Removed some leftovers.

Author: blkt

src/codegate/pipeline/output.py

@@ -182,25 +182,16 @@ async def process_stream(
             # Process any remaining content in buffer when stream ends
             if self._context.buffer:
                 final_content = "".join(self._context.buffer)
-                logger.error("Context buffer was not empty, it should have been!", content=final_content)
-            #     chunk = ModelResponse(
-            #         id=self._buffered_chunk.id,
-            #         choices=[
-            #             StreamingChoices(
-            #                 finish_reason=None,
-            #                 # we just put one choice in the buffer, so 0 is fine
-            #                 index=0,
-            #                 delta=Delta(content=final_content, role="assistant"),
-            #                 # umm..is this correct?
-            #                 logprobs=self._buffered_chunk.choices[0].logprobs,
-            #             )
-            #         ],
-            #         created=self._buffered_chunk.created,
-            #         model=self._buffered_chunk.model,
-            #         object="chat.completion.chunk",
-            #     )
-            #     self._input_context.add_output(chunk)
-            #     yield chunk
+                logger.error(
+                    "Context buffer was not empty, it should have been!",
+                    content=final_content,
+                    len=len(self._context.buffer),
+                )
+                # NOTE: Original code created chunks for all remaining
+                # messages in `self._context.buffer`, but it looks
+                # like it was defensive code. We should instead ensure
+                # that no messages remain there at each step of the
+                # pipeline in some way.
                 self._context.buffer.clear()
 
             if finish_stream:

src/codegate/pipeline/secrets/secrets.py

@@ -1,6 +1,6 @@
 import itertools
 from abc import abstractmethod
-from typing import List, Optional, Tuple
+from typing import Any, List, Optional, Tuple
 
 import regex as re
 
@@ -19,14 +19,6 @@
 from codegate.pipeline.output import OutputPipelineContext, OutputPipelineStep
 from codegate.pipeline.secrets.manager import SecretsManager
 from codegate.pipeline.secrets.signatures import CodegateSignatures, Match
-from codegate.pipeline.systemmsg import add_or_update_system_message
-from codegate.types.common import (
-    ChatCompletionRequest,
-    ChatCompletionSystemMessage,
-    Delta,
-    ModelResponse,
-    StreamingChoices,
-)
 
 
 logger = structlog.get_logger("codegate")
@@ -280,7 +272,7 @@ def _redact_text(
         return text_encryptor.obfuscate(text, snippet)
 
     async def process(
-        self, request: ChatCompletionRequest, context: PipelineContext
+        self, reques: Any, context: PipelineContext
     ) -> PipelineResult:
         """
         Process the request to find and protect secrets in all messages.
@@ -293,68 +285,34 @@ async def process(
             PipelineResult containing the processed request and context with redaction metadata
         """
 
-        ##### NEW CODE PATH #####
-
-        if type(request) != ChatCompletionRequest:
-            secrets_manager = context.sensitive.manager
-            if not secrets_manager or not isinstance(secrets_manager, SecretsManager):
-                raise ValueError("Secrets manager not found in context")
-            session_id = context.sensitive.session_id
-            if not session_id:
-                raise ValueError("Session ID not found in context")
-
-            total_matches = []
-
-            # get last user message block to get index for the first relevant user message
-            last_user_message = self.get_last_user_message_block(request, context.client)
-            last_assistant_idx = last_user_message[1] - 1 if last_user_message else -1
-
-            # Process all messages
-            for i, message in enumerate(request.get_messages()):
-                for content in message.get_content():
-                    txt = content.get_text()
-                    if txt is not None:
-                        redacted_content, secrets_matched = self._redact_message_content(
-                            "".join(txt for txt in content.get_text()), secrets_manager, session_id, context
-                        )
-                        content.set_text(redacted_content)
-                        if i > last_assistant_idx:
-                            total_matches += secrets_matched
-
-            # Not count repeated secret matches
-            request = self._finalize_redaction(context, total_matches, request)
-            return PipelineResult(request=request, context=context)
-
-        ##### OLD CODE PATH #####
-
-        if "messages" not in request:
-            return PipelineResult(request=request, context=context)
-
         secrets_manager = context.sensitive.manager
         if not secrets_manager or not isinstance(secrets_manager, SecretsManager):
             raise ValueError("Secrets manager not found in context")
         session_id = context.sensitive.session_id
         if not session_id:
             raise ValueError("Session ID not found in context")
 
-        new_request = request.copy()
         total_matches = []
 
         # get last user message block to get index for the first relevant user message
-        last_user_message = self.get_last_user_message_block(new_request, context.client)
+        last_user_message = self.get_last_user_message_block(request, context.client)
         last_assistant_idx = last_user_message[1] - 1 if last_user_message else -1
 
         # Process all messages
-        for i, message in enumerate(new_request["messages"]):
-            if "content" in message and message["content"]:
-                redacted_content, secrets_matched = self._redact_message_content(
-                    message["content"], secrets_manager, session_id, context
-                )
-                new_request["messages"][i]["content"] = redacted_content
-                if i > last_assistant_idx:
-                    total_matches += secrets_matched
-        new_request = self._finalize_redaction(context, total_matches, new_request)
-        return PipelineResult(request=new_request, context=context)
+        for i, message in enumerate(request.get_messages()):
+            for content in message.get_content():
+                txt = content.get_text()
+                if txt is not None:
+                    redacted_content, secrets_matched = self._redact_message_content(
+                        "".join(txt for txt in content.get_text()), secrets_manager, session_id, context
+                    )
+                    content.set_text(redacted_content)
+                    if i > last_assistant_idx:
+                        total_matches += secrets_matched
+
+        # Not count repeated secret matches
+        request = self._finalize_redaction(context, total_matches, request)
+        return PipelineResult(request=request, context=context)
 
     def _redact_message_content(self, message_content, secrets_manager, session_id, context):
         # Extract any code snippets
@@ -404,14 +362,7 @@ def _finalize_redaction(self, context, total_matches, new_request):
         logger.info(f"Total secrets redacted since last assistant message: {total_redacted}")
         context.metadata["redacted_secrets_count"] = total_redacted
         if total_redacted > 0:
-            if isinstance(new_request, pydantic.BaseModel):
-                new_request.add_system_prompt(Config.get_config().prompts.secrets_redacted)
-                return new_request
-            system_message = ChatCompletionSystemMessage(
-                content=Config.get_config().prompts.secrets_redacted,
-                role="system",
-            )
-            return add_or_update_system_message(new_request, system_message, context)
+            new_request.add_system_prompt(Config.get_config().prompts.secrets_redacted)
         return new_request
 
 
@@ -449,10 +400,10 @@ def _find_complete_redaction(self, text: str) -> tuple[Optional[re.Match[str]],
 
     async def process_chunk(
         self,
-        chunk: ModelResponse,
+        chunk: Any,
         context: OutputPipelineContext,
         input_context: Optional[PipelineContext] = None,
-    ) -> list[ModelResponse]:
+    ) -> list[Any]:
         """Process a single chunk of the stream"""
         if not input_context:
             raise ValueError("Input context not found")
@@ -461,9 +412,6 @@ async def process_chunk(
         if input_context.sensitive.session_id == "":
             raise ValueError("Session ID not found in input context")
 
-        # if len(chunk.choices) == 0 or not chunk.choices[0].delta.content:
-        #     return [chunk]
-
         for content in chunk.get_content():
             # Check the buffered content
             buffered_content = "".join(context.buffer)
@@ -518,37 +466,20 @@ class SecretRedactionNotifier(OutputPipelineStep):
     def name(self) -> str:
         return "secret-redaction-notifier"
 
-    def _create_chunk(self, original_chunk: ModelResponse, content: str) -> ModelResponse:
+    def _create_chunk(self, original_chunk: Any, content: str) -> Any:
         """
         Creates a new chunk with the given content, preserving the original chunk's metadata
         """
-        if isinstance(original_chunk, ModelResponse):
-            return ModelResponse(
-                id=original_chunk.id,
-                choices=[
-                    StreamingChoices(
-                        finish_reason=None,
-                        index=0,
-                        delta=Delta(content=content, role="assistant"),
-                        logprobs=None,
-                    )
-                ],
-                created=original_chunk.created,
-                model=original_chunk.model,
-                object="chat.completion.chunk",
-            )
-        else:
-            # TODO verify if deep-copy is necessary
-            copy = original_chunk.model_copy(deep=True)
-            copy.set_text(content)
-            return copy
+        copy = original_chunk.model_copy(deep=True)
+        copy.set_text(content)
+        return copy
 
     async def process_chunk(
         self,
-        chunk: ModelResponse,
+        chunk: Any,
         context: OutputPipelineContext,
         input_context: Optional[PipelineContext] = None,
-    ) -> list[ModelResponse]:
+    ) -> list[Any]:
         """Process a single chunk of the stream"""
         if (
             not input_context
@@ -568,20 +499,21 @@ async def process_chunk(
         )
 
         # Check if this is the first chunk (delta role will be present, others will not)
-        # if len(chunk.choices) > 0 and chunk.choices[0].delta.role:
         for _ in itertools.takewhile(lambda x: x[0] == 1, enumerate(chunk.get_content())):
             redacted_count = input_context.metadata["redacted_secrets_count"]
             secret_text = "secret" if redacted_count == 1 else "secrets"
             # Create notification chunk
             if tool_name in ["cline", "kodu"]:
+                # NOTE: Original code was ensuring that role was
+                # "assistant" here, we might have to do that as well,
+                # but I believe it was defensive programming or
+                # leftover of some refactoring.
                 notification_chunk = self._create_chunk(
                     chunk,
                     f"<thinking>\n🛡️ [CodeGate prevented {redacted_count} {secret_text}]"
                     f"(http://localhost:9090/?search=codegate-secrets) from being leaked "
                     f"by redacting them.</thinking>\n\n",
                 )
-                # TODO fix this
-                # notification_chunk.choices[0].delta.role = "assistant"
             else:
                 notification_chunk = self._create_chunk(
                     chunk,

src/codegate/pipeline/system_prompt/codegate.py

@@ -92,38 +92,10 @@ async def process(
         if not should_add_codegate_sys_prompt and not wrksp_custom_instructions:
             return PipelineResult(request=request, context=context)
 
-        
-        ##### NEW CODE PATH #####
-
-        if type(request) != ChatCompletionRequest:
-            request_system_message = {}
-            req_sys_prompt = ""
-            for sysprompt in request.get_system_prompt():
-                req_sys_prompt = sysprompt
-
-            system_prompt = await self._construct_system_prompt(
-                context.client,
-                wrksp_custom_instructions,
-                req_sys_prompt,
-                should_add_codegate_sys_prompt,
-            )
-            context.add_alert(self.name, trigger_string=system_prompt)
-            request.set_system_prompt(system_prompt)
-
-            return PipelineResult(request=request, context=context)
-
-        ##### OLD CODE PATH #####
-
-        new_request = request.copy()
-
-        if "messages" not in new_request:
-            new_request["messages"] = []
-
         request_system_message = {}
-        for message in new_request["messages"]:
-            if message["role"] == "system":
-                request_system_message = message
-        req_sys_prompt = request_system_message.get("content")
+        req_sys_prompt = ""
+        for sysprompt in request.get_system_prompt():
+            req_sys_prompt = sysprompt
 
         system_prompt = await self._construct_system_prompt(
             context.client,
@@ -132,12 +104,6 @@ async def process(
             should_add_codegate_sys_prompt,
         )
         context.add_alert(self.name, trigger_string=system_prompt)
-        if not request_system_message:
-            # Insert the system prompt at the beginning of the messages
-            sytem_message = ChatCompletionSystemMessage(content=system_prompt, role="system")
-            new_request["messages"].insert(0, sytem_message)
-        else:
-            # Update the existing system prompt
-            request_system_message["content"] = system_prompt
-
-        return PipelineResult(request=new_request, context=context)
+        request.set_system_prompt(system_prompt)
+
+        return PipelineResult(request=request, context=context)

src/codegate/providers/ollama/adapter.py

@@ -104,39 +104,3 @@ async def __anext__(self):
             return chunk
         except StopAsyncIteration:
             raise StopAsyncIteration
-
-
-class ModelToOllama(AsyncIterator[ChatResponse]):
-
-    def __init__(self, normalized_reply: AsyncIterator[ModelResponse]):
-        self.normalized_reply = normalized_reply
-        self._aiter = normalized_reply.__aiter__()
-
-    def __aiter__(self):
-        return self
-
-    async def __anext__(self) -> Union[ChatResponse]:
-        try:
-            chunk = await self._aiter.__anext__()
-            if not isinstance(chunk, ModelResponse):
-                return chunk
-            # Convert the timestamp to a datetime object
-            datetime_obj = datetime.fromtimestamp(chunk.created, tz=timezone.utc)
-            created_at = datetime_obj.isoformat()
-
-            message = chunk.choices[0].delta.content
-            done = False
-            if chunk.choices[0].finish_reason == "stop":
-                done = True
-                message = ""
-
-            # Convert the model response to an Ollama response
-            ollama_response = ChatResponse(
-                model=chunk.model,
-                created_at=created_at,
-                done=done,
-                message=Message(content=message, role="assistant"),
-            )
-            return ollama_response
-        except StopAsyncIteration:
-            raise StopAsyncIteration