Removed no longer used code.

Author: blkt

src/codegate/llm_utils/__init__.py

@@ -11,6 +11,8 @@
 from codegate.providers.fim_analyzer import FIMAnalyzer
 from codegate.providers.registry import ProviderRegistry
 from codegate.workspaces.crud import WorkspaceCrud
+from codegate.types.openai import ChatCompletionRequest
+
 
 logger = structlog.get_logger("codegate")
 

src/codegate/llm_utils/llmclient.py

@@ -12,7 +12,7 @@
 from codegate.db.models import Alert, AlertSeverity, Output, Prompt
 from codegate.extract_snippets.message_extractor import CodeSnippet
 from codegate.pipeline.secrets.manager import SecretsManager
-from codegate.types.common import ChatCompletionRequest, ModelResponse
+
 
 logger = structlog.get_logger("codegate")
 
@@ -90,7 +90,7 @@ def add_alert(
         # logger.debug(f"Added alert to context: {self.alerts_raised[-1]}")
 
     def add_input_request(
-        self, normalized_request: ChatCompletionRequest, is_fim_request: bool, provider: str
+        self, normalized_request: Any, is_fim_request: bool, provider: str
     ) -> None:
         try:
             if self.prompt_id is None:
@@ -109,7 +109,7 @@ def add_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:
+    def add_output(self, model_response: Any) -> None:
         try:
             if self.prompt_id is None:
                 logger.warning(f"Tried to record output without response: {model_response}")
@@ -152,7 +152,7 @@ class PipelineResult:
     or a response to return to the client.
     """
 
-    request: Optional[ChatCompletionRequest] = None
+    request: Optional[Any] = None
     response: Optional[PipelineResponse] = None
     context: Optional[PipelineContext] = None
     error_message: Optional[str] = None
@@ -183,13 +183,13 @@ def name(self) -> str:
 
     @staticmethod
     def get_last_user_message(
-        request: ChatCompletionRequest,
+        request: Any,
     ) -> Optional[tuple[str, int]]:
         """
         Get the last user message and its index from the request.
 
         Args:
-            request (ChatCompletionRequest): The chat completion request to process
+            request (Any): The chat completion request to process
 
         Returns:
             Optional[tuple[str, int]]: A tuple containing the message content and
@@ -207,13 +207,13 @@ def get_last_user_message(
 
     @staticmethod
     def get_last_user_message_block(
-        request: ChatCompletionRequest,
+        request: Any,
     ) -> Optional[tuple[str, int]]:
         """
         Get the last block of consecutive 'user' messages from the request.
 
         Args:
-            request (ChatCompletionRequest): The chat completion request to process
+            request (Any): The chat completion request to process
 
         Returns:
             Optional[str, int]: A string containing all consecutive user messages in the
@@ -237,7 +237,7 @@ def get_last_user_message_block(
 
     @abstractmethod
     async def process(
-        self, request: ChatCompletionRequest, context: PipelineContext
+        self, request: Any, context: PipelineContext
     ) -> PipelineResult:
         """Process a request and return either modified request or response stream"""
         pass
@@ -266,7 +266,7 @@ def __init__(
 
     async def process_request(
         self,
-        request: ChatCompletionRequest,
+        request: Any,
         provider: str,
         model: str,
         api_key: Optional[str] = None,
@@ -336,7 +336,7 @@ def _create_instance(self, client_type: ClientType) -> InputPipelineInstance:
 
     async def process_request(
         self,
-        request: ChatCompletionRequest,
+        request: Any,
         provider: str,
         model: str,
         api_key: Optional[str] = None,

src/codegate/muxing/router.py

@@ -1,6 +1,6 @@
 import re
 import shlex
-from typing import Optional
+from typing import Any, Optional
 
 
 from codegate.clients.clients import ClientType
@@ -11,7 +11,6 @@
     PipelineStep,
 )
 from codegate.pipeline.cli.commands import CustomInstructions, Version, Workspace
-from codegate.types.common import ChatCompletionRequest
 
 codegate_regex = re.compile(r"^codegate(?:\s+(.*))?", re.IGNORECASE)
 
@@ -109,14 +108,14 @@ def name(self) -> str:
         return "codegate-cli"
 
     async def process(
-        self, request: ChatCompletionRequest, context: PipelineContext
+        self, request: Any, context: PipelineContext
     ) -> PipelineResult:
         """
         Checks if the last user message contains "codegate" and process the command.
         This short-circuits the pipeline if the message is found.
 
         Args:
-            request (ChatCompletionRequest): The chat completion request to process
+            request (Any): The chat completion request to process
             context (PipelineContext): The current pipeline context
 
         Returns:

src/codegate/pipeline/base.py

@@ -1,6 +1,7 @@
 import itertools
 import json
 import re
+from typing import Any
 
 import structlog
 
@@ -13,7 +14,6 @@
     PipelineStep,
 )
 from codegate.storage.storage_engine import StorageEngine
-from codegate.types.common import ChatCompletionRequest
 from codegate.utils.package_extractor import PackageExtractor
 from codegate.utils.utils import generate_vector_string
 
@@ -71,7 +71,7 @@ def generate_context_str(self, objects: list[object], context: PipelineContext)
         return context_str
 
     async def process(  # noqa: C901
-        self, request: ChatCompletionRequest, context: PipelineContext
+        self, request: Any, context: PipelineContext
     ) -> PipelineResult:
         """
         Use RAG DB to add context to the user request

src/codegate/pipeline/cli/cli.py

@@ -1,4 +1,4 @@
-from typing import Optional
+from typing import Any, Optional
 from urllib.parse import quote
 
 import structlog
@@ -14,7 +14,7 @@
 # from codegate.pipeline.suspicious_commands.suspicious_commands import check_suspicious_code
 from codegate.storage import StorageEngine
 from codegate.utils.package_extractor import PackageExtractor
-from codegate.types.common import ModelResponse, Delta, StreamingChoices
+
 
 logger = structlog.get_logger("codegate")
 
@@ -29,7 +29,7 @@ def __init__(self):
     def name(self) -> str:
         return "code-comment"
 
-    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
         """
@@ -109,10 +109,10 @@ def _split_chunk_at_code_end(self, content: str) -> tuple[str, 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"""
         for content in chunk.get_content():
             # Get current content plus this new chunk