Replace litellm with native API implementations.

Refactors client architecture to use native implementations instead of
`litellm` dependency. Adds support for OpenAPI, Ollama, OpenRouter,
and fixes multiple issues with Anthropic and Copilot
providers. Improves message handling and streaming responses.

Commit message brought you by Anthropic Claude 3.7.

Author: blkt

prompts/default.yaml

@@ -46,7 +46,7 @@ pii_redacted: |
   The context files contain redacted personally identifiable information (PII) that is represented by a UUID encased within <>. For example:
   - <123e4567-e89b-12d3-a456-426614174000>
   - <2d040296-98e9-4350-84be-fda4336057eb>
-  If you encounter any PII redacted with a UUID, DO NOT WARN the user about it. Simplt respond to the user request and keep the PII redacted and intact, using the same UUID.
+  If you encounter any PII redacted with a UUID, DO NOT WARN the user about it. Simply respond to the user request and keep the PII redacted and intact, using the same UUID.
 # Security-focused prompts
 security_audit: "You are a security expert conducting a thorough code review. Identify potential security vulnerabilities, suggest improvements, and explain security best practices."
 
@@ -56,6 +56,6 @@ red_team: "You are a red team member conducting a security assessment. Identify
 # BlueTeam prompts
 blue_team: "You are a blue team member conducting a security assessment. Identify security controls, misconfigurations, and potential vulnerabilities."
 
-# Per client prompts
+# Per client prompts
 client_prompts:
     kodu: "If malicious packages or leaked secrets are found, please end the task, sending the problems found embedded in <attempt_completion><result> tags"

src/codegate/config.py

@@ -16,9 +16,9 @@
 
 # Default provider URLs
 DEFAULT_PROVIDER_URLS = {
-    "openai": "https://api.openai.com/v1",
-    "openrouter": "https://openrouter.ai/api/v1",
-    "anthropic": "https://api.anthropic.com/v1",
+    "openai": "https://api.openai.com",
+    "openrouter": "https://openrouter.ai/api",
+    "anthropic": "https://api.anthropic.com",
     "vllm": "http://localhost:8000",  # Base URL without /v1 path
     "ollama": "http://localhost:11434",  # Default Ollama server URL
     "lm_studio": "http://localhost:1234",

src/codegate/db/connection.py

@@ -121,6 +121,17 @@ def does_db_exist(self):
         return self._db_path.is_file()
 
 
+def row_from_model(model: BaseModel) -> dict:
+    return dict(
+        id=model.id,
+        timestamp=model.timestamp,
+        provider=model.provider,
+        request=model.request.json(exclude_defaults=True, exclude_unset=True),
+        type=model.type,
+        workspace_id=model.workspace_id,
+    )
+
+
 class DbRecorder(DbCodeGate):
     def __init__(self, sqlite_path: Optional[str] = None, *args, **kwargs):
         super().__init__(sqlite_path, *args, **kwargs)
@@ -131,7 +142,10 @@ async def _execute_update_pydantic_model(
         """Execute an update or insert command for a Pydantic model."""
         try:
             async with self._async_db_engine.begin() as conn:
-                result = await conn.execute(sql_command, model.model_dump())
+                row = model
+                if isinstance(model, BaseModel):
+                    row = model.model_dump()
+                result = await conn.execute(sql_command, row)
                 row = result.first()
                 if row is None:
                     return None
@@ -173,7 +187,8 @@ async def record_request(self, prompt_params: Optional[Prompt] = None) -> Option
                 RETURNING *
                 """
         )
-        recorded_request = await self._execute_update_pydantic_model(prompt_params, sql)
+        row = row_from_model(prompt_params)
+        recorded_request = await self._execute_update_pydantic_model(row, sql)
         # Uncomment to debug the recorded request
         # logger.debug(f"Recorded request: {recorded_request}")
         return recorded_request  # type: ignore
@@ -192,7 +207,8 @@ async def update_request(
                 RETURNING *
                 """
         )
-        updated_request = await self._execute_update_pydantic_model(prompt_params, sql)
+        row = row_from_model(prompt_params)
+        updated_request = await self._execute_update_pydantic_model(row, sql)
         # Uncomment to debug the recorded request
         # logger.debug(f"Recorded request: {recorded_request}")
         return updated_request  # type: ignore
@@ -215,7 +231,7 @@ async def record_outputs(
             output=first_output.output,
         )
         full_outputs = []
-        # Just store the model respnses in the list of JSON objects.
+        # Just store the model responses in the list of JSON objects.
         for output in outputs:
             full_outputs.append(output.output)
 
@@ -339,7 +355,7 @@ async def record_context(self, context: Optional[PipelineContext]) -> None:
                     f"Alerts: {len(context.alerts_raised)}."
                 )
         except Exception as e:
-            logger.error(f"Failed to record context: {context}.", error=str(e))
+            logger.error(f"Failed to record context: {context}.", error=str(e), exc_info=e)
 
     async def add_workspace(self, workspace_name: str) -> WorkspaceRow:
         """Add a new workspace to the DB.

src/codegate/db/fim_cache.py

@@ -33,6 +33,18 @@ def __init__(self):
 
     def _extract_message_from_fim_request(self, request: str) -> Optional[str]:
         """Extract the user message from the FIM request"""
+        ### NEW CODE PATH ###
+        if not isinstance(request, str):
+            content_message = None
+            for message in request.get_messages():
+                for content in message.get_content():
+                    if content_message is None:
+                        content_message = content.get_text()
+                    else:
+                        logger.warning("Expected one user message, found multiple.")
+                        return None
+            return content_message
+
         try:
             parsed_request = json.loads(request)
         except Exception as e:

src/codegate/extract_snippets/body_extractor.py

@@ -9,6 +9,7 @@
     KoduCodeSnippetExtractor,
     OpenInterpreterCodeSnippetExtractor,
 )
+from codegate.types.common import MessageTypeFilter
 
 
 class BodyCodeSnippetExtractorError(Exception):
@@ -32,25 +33,22 @@ def _extract_from_user_messages(self, data: dict) -> set[str]:
             raise BodyCodeSnippetExtractorError("Code Extractor not set.")
 
         filenames: List[str] = []
-        for msg in data.get("messages", []):
-            if msg.get("role", "") == "user":
+        for msg in data.get_messages(filters=[MessageTypeFilter.USER]):
+            for content in msg.get_content():
                 extracted_snippets = self._snippet_extractor.extract_unique_snippets(
-                    msg.get("content")
+                    content.get_text(),
                 )
                 filenames.extend(extracted_snippets.keys())
         return set(filenames)
 
     def _extract_from_list_user_messages(self, data: dict) -> set[str]:
         filenames: List[str] = []
-        for msg in data.get("messages", []):
-            if msg.get("role", "") == "user":
-                msgs_content = msg.get("content", [])
-                for msg_content in msgs_content:
-                    if msg_content.get("type", "") == "text":
-                        extracted_snippets = self._snippet_extractor.extract_unique_snippets(
-                            msg_content.get("text")
-                        )
-                        filenames.extend(extracted_snippets.keys())
+        for msg in data.get_messages(filters=[MessageTypeFilter.USER]):
+            for content in msg.get_content():
+                extracted_snippets = self._snippet_extractor.extract_unique_snippets(
+                    content.get_text(),
+                )
+                filenames.extend(extracted_snippets.keys())
         return set(filenames)
 
     @abstractmethod

src/codegate/extract_snippets/message_extractor.py

@@ -279,10 +279,16 @@ def extract_snippets(self, message: str, require_filepath: bool = False) -> List
         """
         regexes = self._choose_regex(require_filepath)
         # Find all code block matches
+        if isinstance(message, str):
+            return [
+                self._get_snippet_for_match(match)
+                for regex in regexes
+                for match in regex.finditer(message)
+            ]
         return [
             self._get_snippet_for_match(match)
             for regex in regexes
-            for match in regex.finditer(message)
+            for match in regex.finditer(message.get_text())
         ]
 
     def extract_unique_snippets(self, message: str) -> Dict[str, CodeSnippet]:

src/codegate/llm_utils/__init__.py

@@ -14,6 +14,10 @@
 from codegate.db import models as db_models
 from codegate.muxing import rulematcher
 from codegate.providers.ollama.adapter import OLlamaToModel
+from codegate.types.ollama import StreamingChatCompletion as OllamaStreamingChatCompletion
+from codegate.types.ollama import StreamingGenerateCompletion as OllamaStreamingGenerateCompletion
+from codegate.muxing.ollama_mappers import openai_chunk_from_ollama_chat, openai_chunk_from_ollama_generate
+from codegate.types.openai import StreamingChatCompletion as OpenAIStreamingChatCompletion
 
 logger = structlog.get_logger("codegate")
 
@@ -41,12 +45,9 @@ def _get_provider_formatted_url(self, model_route: rulematcher.ModelRoute) -> st
             return urljoin(model_route.endpoint.endpoint, "/api/v1")
         return model_route.endpoint.endpoint
 
-    def set_destination_info(self, model_route: rulematcher.ModelRoute, data: dict) -> dict:
+    def get_destination_info(self, model_route: rulematcher.ModelRoute) -> dict:
         """Set the destination provider info."""
-        new_data = copy.deepcopy(data)
-        new_data["model"] = model_route.model.name
-        new_data["base_url"] = self._get_provider_formatted_url(model_route)
-        return new_data
+        return model_route.model.name, self._get_provider_formatted_url(model_route)
 
 
 class OutputFormatter(ABC):
@@ -215,8 +216,8 @@ def _format_ollama(self, chunk: str) -> str:
         """Format the Ollama chunk to OpenAI format."""
         try:
             chunk_dict = json.loads(chunk)
-            ollama_chunk = ChatResponse(**chunk_dict)
-            open_ai_chunk = OLlamaToModel.normalize_chat_chunk(ollama_chunk)
+            ollama_chunk = OllamaStreamingChatCompletion.model_validate(chunk_dict)
+            open_ai_chunk = openai_chunk_from_ollama_chat(ollama_chunk)
             return open_ai_chunk.model_dump_json(exclude_none=True, exclude_unset=True)
         except Exception as e:
             # Sometimes we receive an OpenAI formatted chunk from ollama. Specifically when
@@ -251,10 +252,11 @@ def _format_ollama(self, chunk: str) -> str:
         """Format the Ollama chunk to OpenAI format."""
         try:
             chunk_dict = json.loads(chunk)
-            ollama_chunk = GenerateResponse(**chunk_dict)
-            open_ai_chunk = OLlamaToModel.normalize_fim_chunk(ollama_chunk)
-            return json.dumps(open_ai_chunk, separators=(",", ":"), indent=None)
-        except Exception:
+            ollama_chunk = OllamaStreamingGenerateCompletion.model_validate(chunk_dict)
+            open_ai_chunk = openai_chunk_from_ollama_generate(ollama_chunk)
+            return open_ai_chunk.model_dump_json(exclude_none=True, exclude_unset=True)
+        except Exception as e:
+            print("Error formatting Ollama chunk: ", chunk, e)
             return chunk