This repository was archived by the owner on Jun 5, 2025. It is now read-only.
Copilot DB integration. Keep DB objects in context to record at the end. #331
Conversation
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
aponcedeleonch
commented
Dec 13, 2024
aponcedeleonch
commented
Dec 13, 2024
aponcedeleonch
force-pushed
the
move-db-to-pipeline
branch
from
6e976cd to
73abafe
Compare
aponcedeleonch
changed the title
aponcedeleonch
commented
Dec 14, 2024
src/codegate/pipeline/base.py
Outdated
Comment on lines
310
to
317
| # Create the input request at the end so we make sure the secrets are obfuscated | ||
| self.context.add_input_request( | ||
| current_request, is_fim_request=self.is_fim, provider=provider | ||
| ) |
There was a problem hiding this comment.
This bit creates the request object that is going to be stored in DB at the end of the pipeline steps. Like this we wouldn't store the secrets incoming in the request message
aponcedeleonch
commented
Dec 14, 2024
Comment on lines
+119
to
+156
| def add_input_request( | ||
| self, normalized_request: ChatCompletionRequest, is_fim_request: bool, provider: str | ||
| ) -> None: | ||
| try: | ||
| if self.prompt_id is None: | ||
| self.prompt_id = str(uuid.uuid4()) | ||
|
|
||
| request_str = json.dumps(normalized_request) | ||
|
|
||
| self.input_request = Prompt( | ||
| id=self.prompt_id, | ||
| timestamp=datetime.datetime.now(datetime.timezone.utc), | ||
| provider=provider, | ||
| type="fim" if is_fim_request else "chat", | ||
| request=request_str, | ||
| ) | ||
| logger.debug(f"Added input request to context: {self.input_request}") | ||
| except Exception as e: | ||
| logger.warning(f"Failed to serialize input request: {normalized_request}", error=str(e)) | ||
|
|
||
| def add_output(self, model_response: ModelResponse) -> None: | ||
| try: | ||
| if self.prompt_id is None: | ||
| logger.warning(f"Tried to record output without response: {model_response}") | ||
| return | ||
|
|
||
| if isinstance(model_response, BaseModel): | ||
| output_str = model_response.model_dump_json(exclude_none=True, exclude_unset=True) | ||
| else: | ||
| output_str = json.dumps(model_response) | ||
|
|
||
| self.output_responses.append( | ||
| Output( | ||
| id=str(uuid.uuid4()), | ||
| prompt_id=self.prompt_id, | ||
| timestamp=datetime.datetime.now(datetime.timezone.utc), | ||
| output=output_str, | ||
| ) |
There was a problem hiding this comment.
Moves storing of DB objects to the context. Notably, if an id is not created when adding an alert it will create one. This would enable to create alerts before the actual Prompt object
aponcedeleonch
commented
Dec 14, 2024
Comment on lines
349
to
354
| # This should work for recording FIM and Chat to DB. Now we keep the objects that are going | ||
| # to be written to DB in the pipeline `context`. Copilot also uses pipelines and `context`. | ||
| # For some reason is only working for FIM. Investigatig and enabling on future PR. | ||
| # finally: | ||
| # if self.context_tracking: | ||
| # await self._db_recorder.record_context(self.context_tracking) |
There was a problem hiding this comment.
Since now the pipeline is storing the DB objects Prompt, Output and Alert we only would need to execute the following at the end of the Copilot provider to get everything in DB:
self._db_recorder.record_context(self.context_tracking)However, is not working as expected. Will more look into it later
aponcedeleonch
force-pushed
the
move-db-to-pipeline
branch
2 times, most recently
from
564d287 to
b2bd4bc
Compare
With this change the objects that are going to be stored in DB are kept in the `context` of the pipeline. The pipeline and its `context` are used by all providers, including Copilot. We would need to find a good place in Copilot provider to record the context in DB, e.g. when all the chunks have been transmitted and the stream is about to be closed.
aponcedeleonch
force-pushed
the
move-db-to-pipeline
branch
from
47c3403 to
6a7c00a
Compare
aponcedeleonch
commented
Dec 14, 2024
Comment on lines
+143
to
+153
| # **Needed for Copilot**. This is a hacky way of recording in DB the context | ||
| # when we see the last chunk. Ideally this should be done in a `finally` or | ||
| # `StopAsyncIteration` but Copilot streams in an infite while loop so is not | ||
| # possible | ||
| if len(chunk.choices) > 0 and chunk.choices[0].get("finish_reason", "") == "stop": | ||
| await self._record_to_db() | ||
|
|
There was a problem hiding this comment.
This part is what is allowing recording to DB copilot chats
aponcedeleonch
changed the title
aponcedeleonch
force-pushed
the
move-db-to-pipeline
branch
from
f48b0f9 to
4a06bac
Compare
aponcedeleonch
force-pushed
the
move-db-to-pipeline
branch
from
4a06bac to
dfb5b8f
Compare
This was referenced Dec 16, 2024
Sign up for free
to subscribe to this conversation on GitHub.
Already have an account?
Sign in.
Add this suggestion to a batch that can be applied as a single commit.
This suggestion is invalid because no changes were made to the code.
Suggestions cannot be applied while the pull request is closed.
Suggestions cannot be applied while viewing a subset of changes.
Only one suggestion per line can be applied in a batch.
Add this suggestion to a batch that can be applied as a single commit.
Applying suggestions on deleted lines is not supported.
You must change the existing code in this line in order to create a valid suggestion.
Outdated suggestions cannot be applied.
This suggestion has been applied or marked resolved.
Suggestions cannot be applied from pending reviews.
Suggestions cannot be applied on multi-line comments.
Suggestions cannot be applied while the pull request is queued to merge.
Suggestion cannot be applied right now. Please check back later.
The main purpose of this is to:
requestat the end of the input pipeline. With this we wouldn't store the secrets included in the message requestCloses: #322, #297, #281