Recorded Copilot chats in pipeline

Author: aponcedeleonch

src/codegate/pipeline/base.py

@@ -276,6 +276,7 @@ async def process_request(
         api_key: Optional[str] = None,
         api_base: Optional[str] = None,
         extra_headers: Optional[Dict[str, str]] = None,
+        is_copilot: bool = False,
     ) -> PipelineResult:
         """Process a request through all pipeline steps"""
         self.context.sensitive = PipelineSensitiveData(
@@ -289,6 +290,9 @@ async def process_request(
         self.context.metadata["extra_headers"] = extra_headers
         current_request = request
 
+        # For Copilot provider=openai. Use a flag to not clash with other places that may use that.
+        provider_db = "copilot" if is_copilot else provider
+
         for step in self.pipeline_steps:
             result = await step.process(current_request, self.context)
             if result is None:
@@ -297,7 +301,7 @@ async def process_request(
             if result.shortcuts_processing():
                 # Also record the input when shortchutting
                 self.context.add_input_request(
-                    current_request, is_fim_request=self.is_fim, provider=provider
+                    current_request, is_fim_request=self.is_fim, provider=provider_db
                 )
                 return result
 
@@ -309,7 +313,7 @@ async def process_request(
 
         # 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
+            current_request, is_fim_request=self.is_fim, provider=provider_db
         )
         return PipelineResult(request=current_request, context=self.context)
 
@@ -334,9 +338,10 @@ async def process_request(
         api_key: Optional[str] = None,
         api_base: Optional[str] = None,
         extra_headers: Optional[Dict[str, str]] = None,
+        is_copilot: bool = False,
     ) -> PipelineResult:
         """Create a new pipeline instance and process the request"""
         instance = self.create_instance()
         return await instance.process_request(
-            request, provider, model, api_key, api_base, extra_headers
+            request, provider, model, api_key, api_base, extra_headers, is_copilot
         )

src/codegate/pipeline/output.py

@@ -6,6 +6,7 @@
 from litellm import ModelResponse
 from litellm.types.utils import Delta, StreamingChoices
 
+from codegate.db.connection import DbRecorder
 from codegate.pipeline.base import CodeSnippet, PipelineContext
 
 logger = structlog.get_logger("codegate")
@@ -75,6 +76,7 @@ def __init__(
         self,
         pipeline_steps: list[OutputPipelineStep],
         input_context: Optional[PipelineContext] = None,
+        db_recorder: Optional[DbRecorder] = None,
     ):
         self._input_context = input_context
         self._pipeline_steps = pipeline_steps
@@ -83,6 +85,10 @@ def __init__(
         # the remaining content in the buffer when the stream ends, we need
         # to store the parameters like model, timestamp, etc.
         self._buffered_chunk = None
+        if not db_recorder:
+            self._db_recorder = DbRecorder()
+        else:
+            self._db_recorder = db_recorder
 
     def _buffer_chunk(self, chunk: ModelResponse) -> None:
         """
@@ -105,6 +111,11 @@ def _store_chunk_content(self, chunk: ModelResponse) -> None:
             if choice.delta is not None and choice.delta.content is not None:
                 self._context.processed_content.append(choice.delta.content)
 
+    async def _record_to_db(self):
+        if self._input_context and not self._input_context.metadata.get("stored_in_db", False):
+            await self._db_recorder.record_context(self._input_context)
+            self._input_context.metadata["stored_in_db"] = True
+
     async def process_stream(
         self, stream: AsyncIterator[ModelResponse]
     ) -> AsyncIterator[ModelResponse]:
@@ -133,6 +144,13 @@ async def process_stream(
 
                     current_chunks = processed_chunks
 
+                # **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()
+
                 # Yield all processed chunks
                 for c in current_chunks:
                     self._store_chunk_content(c)

src/codegate/providers/base.py

@@ -193,7 +193,9 @@ async def _cleanup_after_streaming(
         finally:
             if context:
                 # Record to DB the objects captured during the stream
-                await self._db_recorder.record_context(context)
+                if not context.metadata.get("stored_in_db", False):
+                    await self._db_recorder.record_context(context)
+                    context.metadata["stored_in_db"] = True
                 # Ensure sensitive data is cleaned up after the stream is consumed
                 if context.sensitive:
                     context.sensitive.secure_cleanup()

src/codegate/providers/copilot/pipeline.py

@@ -1,11 +1,11 @@
 import json
 from abc import ABC, abstractmethod
-from typing import Dict
+from typing import Dict, Tuple
 
 import structlog
 from litellm.types.llms.openai import ChatCompletionRequest
 
-from codegate.pipeline.base import PipelineContext
+from codegate.pipeline.base import PipelineContext, SequentialPipelineProcessor
 from codegate.pipeline.factory import PipelineFactory
 from codegate.providers.normalizer.completion import CompletionNormalizer
 
@@ -30,7 +30,7 @@ def _create_normalizer(self):
         pass
 
     @abstractmethod
-    def create_pipeline(self):
+    def create_pipeline(self) -> SequentialPipelineProcessor:
         """Each strategy defines which pipeline to create"""
         pass
 
@@ -64,7 +64,7 @@ def _get_copilot_headers(headers: Dict[str, str]) -> Dict[str, str]:
 
         return copilot_headers
 
-    async def process_body(self, headers: list[str], body: bytes) -> (bytes, PipelineContext):
+    async def process_body(self, headers: list[str], body: bytes) -> Tuple[bytes, PipelineContext]:
         """Common processing logic for all strategies"""
         try:
             normalized_body = self.normalizer.normalize(body)
@@ -85,6 +85,7 @@ async def process_body(self, headers: list[str], body: bytes) -> (bytes, Pipelin
                 api_key=headers_dict.get("authorization", "").replace("Bearer ", ""),
                 api_base="https://" + headers_dict.get("host", ""),
                 extra_headers=CopilotPipeline._get_copilot_headers(headers_dict),
+                is_copilot=True,
             )
 
             if result.request:
@@ -142,7 +143,7 @@ class CopilotFimPipeline(CopilotPipeline):
     def _create_normalizer(self):
         return CopilotFimNormalizer()
 
-    def create_pipeline(self):
+    def create_pipeline(self) -> SequentialPipelineProcessor:
         return self.pipeline_factory.create_fim_pipeline()
 
 
@@ -155,5 +156,5 @@ class CopilotChatPipeline(CopilotPipeline):
     def _create_normalizer(self):
         return CopilotChatNormalizer()
 
-    def create_pipeline(self):
+    def create_pipeline(self) -> SequentialPipelineProcessor:
         return self.pipeline_factory.create_input_pipeline()

src/codegate/providers/copilot/provider.py

@@ -10,7 +10,6 @@
 
 from codegate.ca.codegate_ca import CertificateAuthority
 from codegate.config import Config
-from codegate.db.connection import DbRecorder
 from codegate.pipeline.base import PipelineContext
 from codegate.pipeline.factory import PipelineFactory
 from codegate.pipeline.output import OutputPipelineInstance
@@ -125,7 +124,6 @@ def __init__(self, loop: asyncio.AbstractEventLoop):
         self._closing = False
         self.pipeline_factory = PipelineFactory(SecretsManager())
         self.context_tracking: Optional[PipelineContext] = None
-        self._db_recorder = DbRecorder()
 
     def _select_pipeline(self, method: str, path: str) -> Optional[CopilotPipeline]:
         if method == "POST" and path == "v1/engines/copilot-codex/completions":
@@ -144,7 +142,7 @@ async def _body_through_pipeline(
         path: str,
         headers: list[str],
         body: bytes,
-    ) -> (bytes, PipelineContext):
+    ) -> Tuple[bytes, PipelineContext]:
         logger.debug(f"Processing body through pipeline: {len(body)} bytes")
         strategy = self._select_pipeline(method, path)
         if strategy is None:
@@ -346,12 +344,6 @@ async def handle_http_request(self) -> None:
         except Exception as e:
             logger.error(f"Error handling HTTP request: {e}")
             self.send_error_response(502, str(e).encode())
-        # 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)
 
     async def _get_target_url(self) -> Optional[str]:
         """Determine target URL based on request path and headers"""

tests/pipeline/test_output.py

@@ -1,4 +1,5 @@
 from typing import List
+from unittest.mock import MagicMock
 
 import pytest
 from litellm import ModelResponse
@@ -94,7 +95,8 @@ async def test_single_step_processing(self):
         """Test processing a stream through a single step"""
         step = MockOutputPipelineStep("test_step", modify_content=True)
         context = MockContext()
-        instance = OutputPipelineInstance([step], context)
+        db_recorder = MagicMock()
+        instance = OutputPipelineInstance([step], context, db_recorder)
 
         async def mock_stream():
             yield create_model_response("Hello")
@@ -118,7 +120,8 @@ async def test_multiple_steps_processing(self):
             MockOutputPipelineStep("step2", modify_content=True),
         ]
         context = MockContext()
-        instance = OutputPipelineInstance(steps, context)
+        db_recorder = MagicMock()
+        instance = OutputPipelineInstance(steps, context, db_recorder)
 
         async def mock_stream():
             yield create_model_response("Hello")
@@ -141,7 +144,8 @@ async def test_step_pausing(self):
             MockOutputPipelineStep("step2", modify_content=True),
         ]
         context = MockContext()
-        instance = OutputPipelineInstance(steps, context)
+        db_recorder = MagicMock()
+        instance = OutputPipelineInstance(steps, context, db_recorder)
 
         async def mock_stream():
             yield create_model_response("he")
@@ -197,7 +201,8 @@ async def process_chunk(
                 return []
 
         context = MockContext()
-        instance = OutputPipelineInstance([ReplacementStep()], context)
+        db_recorder = MagicMock()
+        instance = OutputPipelineInstance([ReplacementStep()], context, db_recorder)
 
         async def mock_stream():
             yield create_model_response("he")
@@ -221,7 +226,8 @@ async def test_buffer_processing(self):
         """Test that content is properly buffered and cleared"""
         step = MockOutputPipelineStep("test_step")
         context = MockContext()
-        instance = OutputPipelineInstance([step], context)
+        db_recorder = MagicMock()
+        instance = OutputPipelineInstance([step], context, db_recorder)
 
         async def mock_stream():
             yield create_model_response("Hello")
@@ -242,7 +248,8 @@ async def test_empty_stream(self):
         """Test handling of an empty stream"""
         step = MockOutputPipelineStep("test_step")
         context = MockContext()
-        instance = OutputPipelineInstance([step], context)
+        db_recorder = MagicMock()
+        instance = OutputPipelineInstance([step], context, db_recorder)
 
         async def mock_stream():
             if False:
@@ -275,7 +282,10 @@ async def process_chunk(
                 assert input_context.metadata["test"] == "value"
                 return [chunk]
 
-        instance = OutputPipelineInstance([ContextCheckingStep()], input_context=input_context)
+        db_recorder = MagicMock()
+        instance = OutputPipelineInstance(
+            [ContextCheckingStep()], input_context=input_context, db_recorder=db_recorder
+        )
 
         async def mock_stream():
             yield create_model_response("test")
@@ -288,7 +298,8 @@ async def test_buffer_flush_on_stream_end(self):
         """Test that buffer is properly flushed when stream ends"""
         step = MockOutputPipelineStep("test_step", should_pause=True)
         context = MockContext()
-        instance = OutputPipelineInstance([step], context)
+        db_recorder = MagicMock()
+        instance = OutputPipelineInstance([step], context, db_recorder)
 
         async def mock_stream():
             yield create_model_response("Hello")