Add input processing pipeline + codegate-version pipeline step

This adds a pipeline processing before the completion is ran where
the request is either change or can be shortcut. This pipeline consists
of steps, for now we implement a single step `CodegateVersion` that
responds with the codegate version if the verbatim `codegate-version`
string is found in the input.

The pipeline also passes along a context, for now that is unused but I
thought this would be where we store extracted code snippets etc.

To avoid import loops, we also move the `BaseCompletionHandler` class to
a new `completion` package.

Since the shortcut replies are more or less simple strings, we add yet
another package `providers/formatting` whose responsibility is to
convert the string returned by the shortcut response to the format
expected by the client, meaning either a reply or a stream of replies in
the LLM-specific format. We use the `BaseCompletionHandler` as a way to
convert to the LLM-specific format.

Author: jhrozek

scripts/import_packages.py

@@ -1,9 +1,10 @@
 import json
-from utils.embedding_util import generate_embeddings
+
 import weaviate
+from weaviate.classes.config import DataType, Property
 from weaviate.embedded import EmbeddedOptions
-from weaviate.classes.config import Property, DataType
 
+from utils.embedding_util import generate_embeddings
 
 json_files = [
     "data/archived.jsonl",

src/codegate/inference/__init__.py

@@ -1,3 +1,3 @@
 from .inference_engine import LlamaCppInferenceEngine
 
-__all__ = [LlamaCppInferenceEngine]
+__all__ = [LlamaCppInferenceEngine]

src/codegate/pipeline/__init__.py

@@ -0,0 +1,164 @@
+from abc import ABC, abstractmethod
+from dataclasses import dataclass, field
+from typing import Any, Dict, List, Optional
+
+from litellm import ChatCompletionRequest
+
+
+@dataclass
+class CodeSnippet:
+    """
+    Represents a code snippet with its programming language.
+
+    Args:
+        language: The programming language identifier (e.g., 'python', 'javascript')
+        code: The actual code content
+    """
+    language: str
+    code: str
+
+    def __post_init__(self):
+        if not self.language or not self.language.strip():
+            raise ValueError("Language must not be empty")
+        if not self.code or not self.code.strip():
+            raise ValueError("Code must not be empty")
+        self.language = self.language.strip().lower()
+
+@dataclass
+class PipelineContext:
+    code_snippets: List[CodeSnippet] = field(default_factory=list)
+    metadata: Dict[str, Any] = field(default_factory=dict)
+
+    def add_code_snippet(self, snippet: CodeSnippet):
+        self.code_snippets.append(snippet)
+
+    def get_snippets_by_language(self, language: str) -> List[CodeSnippet]:
+        return [s for s in self.code_snippets if s.language.lower() == language.lower()]
+
+@dataclass
+class PipelineResponse:
+    """Response generated by a pipeline step"""
+    content: str
+    step_name: str  # The name of the pipeline step that generated this response
+    model: str  # Taken from the original request's model field
+
+@dataclass
+class PipelineResult:
+    """
+    Represents the result of a pipeline operation.
+    Either contains a modified request to continue processing,
+    or a response to return to the client.
+    """
+    request: Optional[ChatCompletionRequest] = None
+    response: Optional[PipelineResponse] = None
+    error_message: Optional[str] = None
+
+    def shortcuts_processing(self) -> bool:
+        """Returns True if this result should end pipeline processing"""
+        return self.response is not None or self.error_message is not None
+
+    @property
+    def success(self) -> bool:
+        """Returns True if the pipeline step completed without errors"""
+        return self.error_message is None
+
+
+class PipelineStep(ABC):
+    """Base class for all pipeline steps in the processing chain."""
+
+    @property
+    @abstractmethod
+    def name(self) -> str:
+        """
+        Returns the name of the pipeline step.
+
+        Returns:
+            str: A unique identifier for this pipeline step
+        """
+        pass
+
+    @staticmethod
+    def get_last_user_message(
+            request: ChatCompletionRequest,
+        ) -> Optional[tuple[str, int]]:
+        """
+        Get the last user message and its index from the request.
+
+        Args:
+            request (ChatCompletionRequest): The chat completion request to process
+
+        Returns:
+            Optional[tuple[str, int]]: A tuple containing the message content and
+                                       its index, or None if no user message is found
+        """
+        if request.get("messages") is None:
+            return None
+        for i in reversed(range(len(request["messages"]))):
+            if request["messages"][i]["role"] == "user":
+                content = request["messages"][i]["content"]
+
+                # This is really another LiteLLM weirdness. Depending on the
+                # provider inside the ChatCompletionRequest you might either
+                # have a string or a list of Union, one of which is a
+                # ChatCompletionTextObject. We'll handle this better by
+                # either dumping litellm completely or converting to a more sane
+                # format # in our own adapter
+
+                # Handle string content
+                if isinstance(content, str):
+                    return content, i
+
+                # Handle iterable of ChatCompletionTextObject
+                if isinstance(content, (list, tuple)):
+                    # Find first text content
+                    for item in content:
+                        if isinstance(item, dict) and item.get("type") == "text":
+                            return item["text"], i
+
+                    # If no text content found, return None
+                    return None
+
+        return None
+
+    @abstractmethod
+    async def process(
+        self,
+        request: ChatCompletionRequest,
+        context: PipelineContext
+    ) -> PipelineResult:
+        """Process a request and return either modified request or response stream"""
+        pass
+
+
+class PipelineProcessor:
+    def __init__(self, pipeline_steps: List[PipelineStep]):
+        self.pipeline_steps = pipeline_steps
+
+    async def process_request(
+            self,
+            request: ChatCompletionRequest,
+    ) -> PipelineResult:
+        """
+        Process a request through all pipeline steps
+
+        Args:
+            request: The chat completion request to process
+
+        Returns:
+            PipelineResult containing either a modified request or response structure
+        """
+        context = PipelineContext()
+        current_request = request
+
+        for step in self.pipeline_steps:
+            result = await step.process(current_request, context)
+            if result is None:
+                continue
+
+            if result.shortcuts_processing():
+                return result
+
+            if result.request is not None:
+                current_request = result.request
+
+        return PipelineResult(request=current_request)

src/codegate/pipeline/base.py

@@ -0,0 +1,55 @@
+from litellm import ChatCompletionRequest
+
+from codegate import __version__
+from codegate.pipeline.base import (
+    PipelineContext,
+    PipelineResponse,
+    PipelineResult,
+    PipelineStep,
+)
+
+
+class CodegateVersion(PipelineStep):
+    """Pipeline step that handles version information requests."""
+
+    @property
+    def name(self) -> str:
+        """
+        Returns the name of this pipeline step.
+
+        Returns:
+            str: The identifier 'codegate-version'
+        """
+        return "codegate-version"
+
+    async def process(
+            self,
+            request: ChatCompletionRequest,
+            context: PipelineContext
+    ) -> PipelineResult:
+        """
+        Checks if the last user message contains "codegate-version" and
+        responds with the current version.
+        This short-circuits the pipeline if the message is found.
+
+        Args:
+            request (ChatCompletionRequest): The chat completion request to process
+            context (PipelineContext): The current pipeline context
+
+        Returns:
+            PipelineResult: Contains version response if triggered, otherwise continues
+            pipeline
+        """
+        last_user_message = self.get_last_user_message(request)
+
+        if last_user_message is not None and "codegate-version" in last_user_message:
+            return PipelineResult(
+                response=PipelineResponse(
+                    step_name=self.name,
+                    content="Codegate version: {}".format(__version__),
+                    model=request["model"],
+                ),
+            )
+
+        # Fall through
+        return PipelineResult(request=request)

src/codegate/pipeline/version/__init__.py

@@ -8,10 +8,10 @@
 
 
 class AnthropicProvider(BaseProvider):
-    def __init__(self):
+    def __init__(self, pipeline_processor=None):
         adapter = AnthropicAdapter()
         completion_handler = LiteLLmShim(adapter)
-        super().__init__(completion_handler)
+        super().__init__(completion_handler, pipeline_processor)
 
     def _setup_routes(self):
         """

src/codegate/pipeline/version/version.py

@@ -1,46 +1,77 @@
 from abc import ABC, abstractmethod
-from typing import Any, AsyncIterator, Callable, Dict
+from typing import Any, AsyncIterator, Callable, Dict, Optional, Union
 
 from fastapi import APIRouter
-from fastapi.responses import StreamingResponse
+from litellm import ModelResponse
 
-StreamGenerator = Callable[[AsyncIterator[Any]], AsyncIterator[str]]
+from codegate.providers.completion.base import BaseCompletionHandler
+from codegate.providers.formatting.input_pipeline import PipelineResponseFormatter
 
+from ..pipeline.base import PipelineProcessor
 
-class BaseCompletionHandler(ABC):
-    """
-    The completion handler is responsible for executing the completion request
-    and creating the streaming response.
-    """
-
-    @abstractmethod
-    async def complete(self, data: Dict, api_key: str) -> AsyncIterator[Any]:
-        pass
-
-    @abstractmethod
-    def create_streaming_response(
-        self, stream: AsyncIterator[Any]
-    ) -> StreamingResponse:
-        pass
-
+StreamGenerator = Callable[[AsyncIterator[Any]], AsyncIterator[str]]
 
 class BaseProvider(ABC):
     """
     The provider class is responsible for defining the API routes and
     calling the completion method using the completion handler.
     """
 
-    def __init__(self, completion_handler: BaseCompletionHandler):
+    def __init__(
+        self,
+        completion_handler: BaseCompletionHandler,
+        pipeline_processor: Optional[PipelineProcessor] = None
+    ):
         self.router = APIRouter()
         self._completion_handler = completion_handler
+        self._pipeline_processor = pipeline_processor
+        self._pipeline_response_formatter = \
+            PipelineResponseFormatter(completion_handler)
         self._setup_routes()
 
     @abstractmethod
     def _setup_routes(self) -> None:
         pass
 
-    async def complete(self, data: Dict, api_key: str) -> AsyncIterator[Any]:
-        return await self._completion_handler.complete(data, api_key)
+    async def complete(
+            self, data: Dict, api_key: str,
+        ) -> Union[ModelResponse, AsyncIterator[ModelResponse]]:
+        """
+        Main completion flow with pipeline integration
+
+        The flow has three main steps:
+        - Translate the request to the OpenAI API format used internally
+        - Process the request with the pipeline processor. This can modify the request
+          or yield a response. The response can either be returned or streamed back to
+          the client
+        - Execute the completion and translate the response back to the
+          provider-specific format
+        """
+        completion_request = self._completion_handler.translate_request(data, api_key)
+        streaming = data.get("stream", False)
+
+        if self._pipeline_processor is not None:
+            result = await self._pipeline_processor.process_request(completion_request)
+
+            if result.error_message:
+                raise Exception(result.error_message)
+
+            if result.response:
+                return self._pipeline_response_formatter.handle_pipeline_response(
+                    result.response, streaming)
+
+            completion_request = result.request
+
+        # Execute the completion and translate the response
+        # This gives us either a single response or a stream of responses
+        # based on the streaming flag
+        raw_response = await self._completion_handler.execute_completion(
+            completion_request,
+            stream=streaming
+        )
+        if not streaming:
+            return self._completion_handler.translate_response(raw_response)
+        return self._completion_handler.translate_streaming_response(raw_response)
 
     def get_routes(self) -> APIRouter:
         return self.router