Fixing tests.

Author: blkt

src/codegate/pipeline/system_prompt/codegate.py

@@ -92,18 +92,15 @@ async def process(
         if not should_add_codegate_sys_prompt and not wrksp_custom_instructions:
             return PipelineResult(request=request, context=context)
 
-        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)
+        # NOTE: this was changed from adding more text to an existing
+        # system prompt to potentially adding a new system prompt.
+        request.add_system_prompt(system_prompt)
 
         return PipelineResult(request=request, context=context)

src/codegate/types/ollama/_request_models.py

@@ -120,7 +120,7 @@ class ChatRequest(pydantic.BaseModel):
     format: dict | None = None
     keep_alive: int | str | None = None
     tools: List[ToolDef] | None = None
-    options: dict
+    options: dict | None = None
 
     def get_stream(self) -> bool:
         return self.stream
@@ -192,7 +192,7 @@ class GenerateRequest(pydantic.BaseModel):
     format: dict | None = None
     keep_alive: int | str | None = None
     images: List[bytes] | None = None
-    options: dict
+    options: dict | None = None
 
     def get_stream(self) -> bool:
         return self.stream

tests/pipeline/extract_snippets/test_extract_snippets.py

@@ -1,12 +1,17 @@
 from unittest.mock import MagicMock, patch
 
 import pytest
-from litellm import ChatCompletionRequest, ModelResponse
-from litellm.types.utils import Delta, StreamingChoices
 
 from codegate.pipeline.base import PipelineContext
 from codegate.pipeline.output import OutputPipelineContext
 from codegate.pipeline.pii.pii import CodegatePii, PiiRedactionNotifier, PiiUnRedactionStep
+from codegate.types.openai import (
+    ChatCompletionRequest,
+    ChoiceDelta,
+    MessageDelta,
+    StreamingChatCompletion,
+    UserMessage,
+)
 
 
 class TestCodegatePii:
@@ -43,7 +48,7 @@ def test_get_redacted_snippet_with_pii(self, pii_step):
 
     @pytest.mark.asyncio
     async def test_process_no_messages(self, pii_step):
-        request = ChatCompletionRequest(model="test-model")
+        request = ChatCompletionRequest(model="test-model", messages=[])
         context = PipelineContext()
 
         result = await pii_step.process(request, context)
@@ -55,7 +60,7 @@ async def test_process_no_messages(self, pii_step):
     async def test_process_with_pii(self, pii_step):
         original_text = "My email is [email protected]"
         request = ChatCompletionRequest(
-            model="test-model", messages=[{"role": "user", "content": original_text}]
+            model="test-model", messages=[UserMessage(role="user", content=original_text)]
         )
         context = PipelineContext()
 
@@ -77,9 +82,10 @@ async def test_process_with_pii(self, pii_step):
         result = await pii_step.process(request, context)
 
         # Verify the user message was anonymized
-        user_messages = [m for m in result.request["messages"] if m["role"] == "user"]
+        user_messages = [m for m in result.request.get_messages() if isinstance(m, UserMessage)]
         assert len(user_messages) == 1
-        assert user_messages[0]["content"] == anonymized_text
+        content = next(user_messages[0].get_content())
+        assert content.get_text() == anonymized_text
 
         # Verify metadata was updated
         assert result.context.metadata["redacted_pii_count"] == 1
@@ -89,9 +95,9 @@ async def test_process_with_pii(self, pii_step):
         assert "pii_manager" in result.context.metadata
 
         # Verify system message was added
-        system_messages = [m for m in result.request["messages"] if m["role"] == "system"]
+        system_messages = [m for m in result.request.get_system_prompt()]
         assert len(system_messages) == 1
-        assert system_messages[0]["content"] == "PII has been redacted"
+        assert system_messages[0] == "PII has been redacted"
 
     def test_restore_pii(self, pii_step):
         anonymized_text = "My email is <test-uuid>"
@@ -121,11 +127,11 @@ def test_is_complete_uuid_invalid(self, unredaction_step):
 
     @pytest.mark.asyncio
     async def test_process_chunk_no_content(self, unredaction_step):
-        chunk = ModelResponse(
+        chunk = StreamingChatCompletion(
             id="test",
             choices=[
-                StreamingChoices(
-                    finish_reason=None, index=0, delta=Delta(content=None), logprobs=None
+                ChoiceDelta(
+                    finish_reason=None, index=0, delta=MessageDelta(content=None), logprobs=None
                 )
             ],
             created=1234567890,
@@ -142,13 +148,13 @@ async def test_process_chunk_no_content(self, unredaction_step):
     @pytest.mark.asyncio
     async def test_process_chunk_with_uuid(self, unredaction_step):
         uuid = "12345678-1234-1234-1234-123456789012"
-        chunk = ModelResponse(
+        chunk = StreamingChatCompletion(
             id="test",
             choices=[
-                StreamingChoices(
+                ChoiceDelta(
                     finish_reason=None,
                     index=0,
-                    delta=Delta(content=f"Text with <{uuid}>"),
+                    delta=MessageDelta(content=f"Text with <{uuid}>"),
                     logprobs=None,
                 )
             ],
@@ -168,6 +174,7 @@ async def test_process_chunk_with_uuid(self, unredaction_step):
 
         result = await unredaction_step.process_chunk(chunk, context, input_context)
 
+        # TODO this should use the abstract interface
         assert result[0].choices[0].delta.content == "Text with [email protected]"
 
 
@@ -199,11 +206,11 @@ def test_format_pii_summary_multiple(self, notifier):
 
     @pytest.mark.asyncio
     async def test_process_chunk_no_pii(self, notifier):
-        chunk = ModelResponse(
+        chunk = StreamingChatCompletion(
             id="test",
             choices=[
-                StreamingChoices(
-                    finish_reason=None, index=0, delta=Delta(content="Hello"), logprobs=None
+                ChoiceDelta(
+                    finish_reason=None, index=0, delta=MessageDelta(content="Hello"), logprobs=None
                 )
             ],
             created=1234567890,
@@ -219,13 +226,13 @@ async def test_process_chunk_no_pii(self, notifier):
 
     @pytest.mark.asyncio
     async def test_process_chunk_with_pii(self, notifier):
-        chunk = ModelResponse(
+        chunk = StreamingChatCompletion(
             id="test",
             choices=[
-                StreamingChoices(
+                ChoiceDelta(
                     finish_reason=None,
                     index=0,
-                    delta=Delta(content="Hello", role="assistant"),
+                    delta=MessageDelta(content="Hello", role="assistant"),
                     logprobs=None,
                 )
             ],
@@ -244,6 +251,7 @@ async def test_process_chunk_with_pii(self, notifier):
         result = await notifier.process_chunk(chunk, context, input_context)
 
         assert len(result) == 2  # Notification chunk + original chunk
+        # TODO this should use the abstract interface
         notification_content = result[0].choices[0].delta.content
         assert "CodeGate protected" in notification_content
         assert "1 email address" in notification_content