fix: only record db content if it is the last chunk in stream (#1091)

For copilot chat, we are actually receiving multiple streams, and
we were recording entries and alerts for each one, repeating those.
So detect if we are in the last chunk and propagate it to the
pipeline, so we can record it only in this case, avoiding dupes.
In the case of other providers, we only receive one request, so
always force saving it

Closes: #936

Author: yrobla

src/codegate/db/connection.py

@@ -129,6 +129,7 @@ async def record_request(self, prompt_params: Optional[Prompt] = None) -> Option
         active_workspace = await DbReader().get_active_workspace()
         workspace_id = active_workspace.id if active_workspace else "1"
         prompt_params.workspace_id = workspace_id
+
         sql = text(
             """
                 INSERT INTO prompts (id, timestamp, provider, request, type, workspace_id)
@@ -302,7 +303,7 @@ async def record_context(self, context: Optional[PipelineContext]) -> None:
                 await self.record_outputs(context.output_responses, initial_id)
                 await self.record_alerts(context.alerts_raised, initial_id)
                 logger.info(
-                    f"Recorded context in DB. Output chunks: {len(context.output_responses)}. "
+                    f"Updated context in DB. Output chunks: {len(context.output_responses)}. "
                     f"Alerts: {len(context.alerts_raised)}."
                 )
         except Exception as e:

src/codegate/pipeline/output.py

@@ -127,7 +127,10 @@ def _record_to_db(self) -> None:
         loop.create_task(self._db_recorder.record_context(self._input_context))
 
     async def process_stream(
-        self, stream: AsyncIterator[ModelResponse], cleanup_sensitive: bool = True
+        self,
+        stream: AsyncIterator[ModelResponse],
+        cleanup_sensitive: bool = True,
+        finish_stream: bool = True,
     ) -> AsyncIterator[ModelResponse]:
         """
         Process a stream through all pipeline steps
@@ -167,7 +170,7 @@ async def process_stream(
         finally:
             # NOTE: Don't use await in finally block, it will break the stream
             # Don't flush the buffer if we assume we'll call the pipeline again
-            if cleanup_sensitive is False:
+            if cleanup_sensitive is False and finish_stream:
                 self._record_to_db()
                 return
 
@@ -194,7 +197,8 @@ async def process_stream(
                 yield chunk
                 self._context.buffer.clear()
 
-            self._record_to_db()
+            if finish_stream:
+                self._record_to_db()
             # Cleanup sensitive data through the input context
             if cleanup_sensitive and self._input_context and self._input_context.sensitive:
                 self._input_context.sensitive.secure_cleanup()

src/codegate/providers/copilot/provider.py

@@ -905,8 +905,16 @@ async def stream_iterator():
                     )
                     yield mr
 
+            # needs to be set as the flag gets reset on finish_data
+            finish_stream_flag = any(
+                choice.get("finish_reason") == "stop"
+                for record in list(self.stream_queue._queue)
+                for choice in record.get("content", {}).get("choices", [])
+            )
             async for record in self.output_pipeline_instance.process_stream(
-                stream_iterator(), cleanup_sensitive=False
+                stream_iterator(),
+                cleanup_sensitive=False,
+                finish_stream=finish_stream_flag,
             ):
                 chunk = record.model_dump_json(exclude_none=True, exclude_unset=True)
                 sse_data = f"data: {chunk}\n\n".encode("utf-8")