WIP: Record the prompt at the end of pipeline. Keep DB objects in context

Author: aponcedeleonch

src/codegate/db/connection.py

@@ -18,6 +18,7 @@
     GetAlertsWithPromptAndOutputRow,
     GetPromptWithOutputsRow,
 )
+from codegate.pipeline.base import PipelineContext
 
 logger = structlog.get_logger("codegate")
 alert_queue = asyncio.Queue()
@@ -104,37 +105,20 @@ async def _insert_pydantic_model(
                 return None
 
     async def record_request(
-        self, normalized_request: ChatCompletionRequest, is_fim_request: bool, provider_str: str
+        self, prompt_params: Optional[Prompt] = None
     ) -> Optional[Prompt]:
-        request_str = None
-        if isinstance(normalized_request, BaseModel):
-            request_str = normalized_request.model_dump_json(exclude_none=True, exclude_unset=True)
-        else:
-            try:
-                request_str = json.dumps(normalized_request)
-            except Exception as e:
-                logger.error(f"Failed to serialize output: {normalized_request}", error=str(e))
-
-        if request_str is None:
-            logger.warning("No request found to record.")
-            return
-
-        # Create a new prompt record
-        prompt_params = Prompt(
-            id=str(uuid.uuid4()),  # Generate a new UUID for the prompt
-            timestamp=datetime.datetime.now(datetime.timezone.utc),
-            provider=provider_str,
-            type="fim" if is_fim_request else "chat",
-            request=request_str,
-        )
+        if prompt_params is None:
+            return None
         sql = text(
             """
                 INSERT INTO prompts (id, timestamp, provider, request, type)
                 VALUES (:id, :timestamp, :provider, :request, :type)
                 RETURNING *
                 """
         )
-        return await self._insert_pydantic_model(prompt_params, sql)
+        recorded_request = await self._insert_pydantic_model(prompt_params, sql)
+        logger.info(f"Recorded request: {recorded_request}")
+        return recorded_request
 
     async def _record_output(self, prompt: Prompt, output_str: str) -> Optional[Output]:
         output_params = Output(
@@ -152,6 +136,28 @@ async def _record_output(self, prompt: Prompt, output_str: str) -> Optional[Outp
         )
         return await self._insert_pydantic_model(output_params, sql)
 
+    async def record_outputs(self, outputs: List[Output]) -> List[Output]:
+        if not outputs:
+            return
+        sql = text(
+            """
+                INSERT INTO outputs (id, prompt_id, timestamp, output)
+                VALUES (:id, :prompt_id, :timestamp, :output)
+                RETURNING *
+                """
+        )
+        # We can insert each alert independently in parallel.
+        outputs_tasks = []
+        async with asyncio.TaskGroup() as tg:
+            for output in outputs:
+                try:
+                    outputs_tasks.append(tg.create_task(self._insert_pydantic_model(output, sql)))
+                except Exception as e:
+                    logger.error(f"Failed to record alert: {output}.", error=str(e))
+        recorded_outputs = [output.result() for output in outputs_tasks]
+        logger.info(f"Recorded outputs: {recorded_outputs}")
+        return recorded_outputs
+
     async def record_output_stream(
         self, prompt: Prompt, model_response: AsyncIterator
     ) -> AsyncGenerator:
@@ -193,7 +199,7 @@ async def record_output_non_stream(
 
         return await self._record_output(prompt, output_str)
 
-    async def record_alerts(self, alerts: List[Alert]) -> None:
+    async def record_alerts(self, alerts: List[Alert]) -> List[Alert]:
         if not alerts:
             return
         sql = text(
@@ -208,15 +214,25 @@ async def record_alerts(self, alerts: List[Alert]) -> None:
                 """
         )
         # We can insert each alert independently in parallel.
+        alerts_tasks = []
         async with asyncio.TaskGroup() as tg:
             for alert in alerts:
                 try:
                     result = tg.create_task(self._insert_pydantic_model(alert, sql))
+                    alerts_tasks.append(result)
                     if result and alert.trigger_category == "critical":
                         await alert_queue.put(f"New alert detected: {alert.timestamp}")
                 except Exception as e:
                     logger.error(f"Failed to record alert: {alert}.", error=str(e))
-        return None
+        recorded_alerts = [alert.result() for alert in alerts_tasks]
+        logger.info(f"Recorded alerts: {recorded_alerts}")
+        return recorded_alerts
+
+    async def record_context(self, context: PipelineContext) -> None:
+        logger.info(f"Recording context: {context}")
+        await self.record_request(context.input_request)
+        await self.record_outputs(context.output_responses)
+        await self.record_alerts(context.alerts_raised)
 
 
 class DbReader(DbCodeGate):

src/codegate/pipeline/base.py

@@ -5,12 +5,13 @@
 from abc import ABC, abstractmethod
 from dataclasses import dataclass, field
 from enum import Enum
-from typing import Any, Dict, List, Optional
+from typing import Any, AsyncIterator, Dict, List, Optional
 
 import structlog
-from litellm import ChatCompletionRequest
+from litellm import ChatCompletionRequest, ModelResponse
+from pydantic import BaseModel
 
-from codegate.db.models import Alert
+from codegate.db.models import Alert, Output, Prompt
 from codegate.pipeline.secrets.manager import SecretsManager
 
 logger = structlog.get_logger("codegate")
@@ -73,6 +74,9 @@ class PipelineContext:
     metadata: Dict[str, Any] = field(default_factory=dict)
     sensitive: Optional[PipelineSensitiveData] = field(default_factory=lambda: None)
     alerts_raised: List[Alert] = field(default_factory=list)
+    prompt_id: Optional[str] = field(default_factory=lambda: None)
+    input_request: Optional[Prompt] = field(default_factory=lambda: None)
+    output_responses: List[Output] = field(default_factory=list)
 
     def add_code_snippet(self, snippet: CodeSnippet):
         self.code_snippets.append(snippet)
@@ -90,9 +94,8 @@ def add_alert(
         """
         Add an alert to the pipeline step alerts_raised.
         """
-        if not self.metadata.get("prompt_id"):
-            logger.warning("No prompt_id found in context. Alert will not be created")
-            return
+        if self.prompt_id is None:
+            self.prompt_id = str(uuid.uuid4())
 
         if not code_snippet and not trigger_string:
             logger.warning("No code snippet or trigger string provided for alert. Will not create")
@@ -103,15 +106,57 @@ def add_alert(
         self.alerts_raised.append(
             Alert(
                 id=str(uuid.uuid4()),
-                prompt_id=self.metadata["prompt_id"],
+                prompt_id=self.prompt_id,
                 code_snippet=code_snippet_str,
                 trigger_string=trigger_string,
                 trigger_type=step_name,
                 trigger_category=severity_category.value,
                 timestamp=datetime.datetime.now(datetime.timezone.utc),
             )
         )
+        logger.info(f"Added alert to context: {self.alerts_raised[-1]}")
+
+    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.info(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:
+                self.prompt_id = str(uuid.uuid4())
+
+            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=self.prompt_id,
+                prompt_id=self.input_request.id,
+                timestamp=datetime.datetime.now(datetime.timezone.utc),
+                output=output_str,
+                )
+            )
+            logger.info(f"Added output to context: {self.output_responses[-1]}")
+        except Exception as e:
+            logger.error(f"Failed to serialize output: {model_response}", error=str(e))
+            return
 
 @dataclass
 class PipelineResponse:
@@ -212,16 +257,17 @@ async def process(
 
 
 class InputPipelineInstance:
-    def __init__(self, pipeline_steps: List[PipelineStep], secret_manager: SecretsManager):
+    def __init__(self, pipeline_steps: List[PipelineStep], secret_manager: SecretsManager, is_fim: bool):
         self.pipeline_steps = pipeline_steps
         self.secret_manager = secret_manager
+        self.is_fim = is_fim
         self.context = PipelineContext()
 
     async def process_request(
         self,
         request: ChatCompletionRequest,
         provider: str,
-        prompt_id: str,
+        # prompt_id: str,
         model: str,
         api_key: Optional[str] = None,
         api_base: Optional[str] = None,
@@ -236,7 +282,7 @@ async def process_request(
             provider=provider,
             api_base=api_base,
         )
-        self.context.metadata["prompt_id"] = prompt_id
+        # self.context.metadata["prompt_id"] = prompt_id
         self.context.metadata["extra_headers"] = extra_headers
         current_request = request
 
@@ -254,23 +300,26 @@ async def process_request(
             if result.context is not None:
                 self.context = result.context
 
+        # 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)
         return PipelineResult(request=current_request, context=self.context)
 
 
 class SequentialPipelineProcessor:
-    def __init__(self, pipeline_steps: List[PipelineStep], secret_manager: SecretsManager):
+    def __init__(self, pipeline_steps: List[PipelineStep], secret_manager: SecretsManager, is_fim: bool):
         self.pipeline_steps = pipeline_steps
         self.secret_manager = secret_manager
+        self.is_fim = is_fim
 
     def create_instance(self) -> InputPipelineInstance:
         """Create a new pipeline instance for processing a request"""
-        return InputPipelineInstance(self.pipeline_steps, self.secret_manager)
+        return InputPipelineInstance(self.pipeline_steps, self.secret_manager, self.is_fim) 
 
     async def process_request(
         self,
         request: ChatCompletionRequest,
         provider: str,
-        prompt_id: str,
+        # prompt_id: str,
         model: str,
         api_key: Optional[str] = None,
         api_base: Optional[str] = None,
@@ -279,5 +328,5 @@ async def process_request(
         """Create a new pipeline instance and process the request"""
         instance = self.create_instance()
         return await instance.process_request(
-            request, provider, prompt_id, model, api_key, api_base, extra_headers
+            request, provider, model, api_key, api_base, extra_headers
         )

src/codegate/pipeline/factory.py

@@ -32,13 +32,13 @@ def create_input_pipeline(self) -> SequentialPipelineProcessor:
             SystemPrompt(Config.get_config().prompts.default_chat),
             CodegateContextRetriever(),
         ]
-        return SequentialPipelineProcessor(input_steps, self.secrets_manager)
+        return SequentialPipelineProcessor(input_steps, self.secrets_manager, is_fim=False)
 
     def create_fim_pipeline(self) -> SequentialPipelineProcessor:
         fim_steps: List[PipelineStep] = [
             CodegateSecrets(),
         ]
-        return SequentialPipelineProcessor(fim_steps, self.secrets_manager)
+        return SequentialPipelineProcessor(fim_steps, self.secrets_manager, is_fim=True)
 
     def create_output_pipeline(self) -> OutputPipelineProcessor:
         output_steps: List[OutputPipelineStep] = [

src/codegate/pipeline/output.py

@@ -115,6 +115,7 @@ async def process_stream(
             async for chunk in stream:
                 # Store chunk content in buffer
                 self._buffer_chunk(chunk)
+                self._input_context.add_output(chunk)
 
                 # Process chunk through each step of the pipeline
                 current_chunks = [chunk]

src/codegate/providers/base.py

@@ -109,7 +109,7 @@ async def _run_input_pipeline(
         api_key: Optional[str],
         api_base: Optional[str],
         is_fim_request: bool,
-        prompt_id: str,
+        # prompt_id: str,
     ) -> PipelineResult:
         # Decide which pipeline processor to use
         if is_fim_request:
@@ -125,7 +125,7 @@ async def _run_input_pipeline(
         result = await pipeline_processor.process_request(
             request=normalized_request,
             provider=self.provider_route_name,
-            prompt_id=prompt_id,
+            # prompt_id=prompt_id,
             model=normalized_request.get("model"),
             api_key=api_key,
             api_base=api_base,
@@ -194,10 +194,11 @@ async def _cleanup_after_streaming(
             async for item in stream:
                 yield item
         finally:
-            # Ensure sensitive data is cleaned up after the stream is consumed
-            if context and context.sensitive:
-                context.sensitive.secure_cleanup()
-                await self._db_recorder.record_alerts(context.alerts_raised)
+            if context:
+                await self._db_recorder.record_context(context)
+                # Ensure sensitive data is cleaned up after the stream is consumed
+                if context.sensitive:
+                    context.sensitive.secure_cleanup()
 
     async def complete(
         self, data: Dict, api_key: Optional[str], is_fim_request: bool
@@ -215,22 +216,22 @@ async def complete(
         """
         normalized_request = self._input_normalizer.normalize(data)
         streaming = normalized_request.get("stream", False)
-        prompt_db = await self._db_recorder.record_request(
-            normalized_request, is_fim_request, self.provider_route_name
-        )
+        # prompt_db = await self._db_recorder.record_request(
+        #     normalized_request, is_fim_request, self.provider_route_name
+        # )
 
-        prompt_db_id = prompt_db.id if prompt_db is not None else None
+        # prompt_db_id = prompt_db.id if prompt_db is not None else None
         input_pipeline_result = await self._run_input_pipeline(
             normalized_request,
             api_key,
             data.get("base_url"),
             is_fim_request,
-            prompt_id=prompt_db_id,
+            # prompt_id=prompt_db_id,
         )
         if input_pipeline_result.response:
-            await self._db_recorder.record_alerts(input_pipeline_result.context.alerts_raised)
+            # await self._db_recorder.record_alerts(input_pipeline_result.context.alerts_raised)
             return await self._pipeline_response_formatter.handle_pipeline_response(
-                input_pipeline_result.response, streaming, prompt_db=prompt_db
+                input_pipeline_result.response, streaming, context=input_pipeline_result.context
             )
 
         provider_request = self._input_normalizer.denormalize(input_pipeline_result.request)
@@ -247,17 +248,18 @@ async def complete(
             normalized_response = self._output_normalizer.normalize(model_response)
             pipeline_output = self._run_output_pipeline(normalized_response)
             # Record the output and alerts in the database can be done in parallel
-            async with asyncio.TaskGroup() as tg:
-                tg.create_task(
-                    self._db_recorder.record_output_non_stream(prompt_db, model_response)
-                )
-                if input_pipeline_result and input_pipeline_result.context:
-                    tg.create_task(
-                        self._db_recorder.record_alerts(input_pipeline_result.context.alerts_raised)
-                    )
+            # async with asyncio.TaskGroup() as tg:
+            #     tg.create_task(
+            #         self._db_recorder.record_output_non_stream(prompt_db, model_response)
+            #     )
+            #     if input_pipeline_result and input_pipeline_result.context:
+            #         tg.create_task(
+            #             self._db_recorder.record_alerts(input_pipeline_result.context.alerts_raised)
+            #         )
+            await self._db_recorder.record_context(input_pipeline_result.context)
             return self._output_normalizer.denormalize(pipeline_output)
 
-        model_response = self._db_recorder.record_output_stream(prompt_db, model_response)
+        # model_response = self._db_recorder.record_output_stream(prompt_db, model_response)
         pipeline_output_stream = await self._run_output_stream_pipeline(
             input_pipeline_result.context, model_response, is_fim_request=is_fim_request
         )