Ported llamacpp.

Author: blkt

src/codegate/providers/llamacpp/completion_handler.py

@@ -11,35 +11,46 @@
 from codegate.config import Config
 from codegate.inference.inference_engine import LlamaCppInferenceEngine
 from codegate.providers.base import BaseCompletionHandler
+from codegate.types.openai import (
+    stream_generator,
+    LegacyCompletion,
+    StreamingChatCompletion,
+)
+
+
+# async def llamacpp_stream_generator(
+#     stream: AsyncIterator[CreateChatCompletionStreamResponse],
+# ) -> AsyncIterator[str]:
+#     """OpenAI-style SSE format"""
+#     try:
+#         async for chunk in stream:
+#             chunk = json.dumps(chunk)
+#             try:
+#                 yield f"data:{chunk}\n\n"
+#             except Exception as e:
+#                 yield f"data:{str(e)}\n\n"
+#     except Exception as e:
+#         yield f"data: {str(e)}\n\n"
+#     finally:
+#         yield "data: [DONE]\n\n"
 
 
-async def llamacpp_stream_generator(
-    stream: AsyncIterator[CreateChatCompletionStreamResponse],
-) -> AsyncIterator[str]:
-    """OpenAI-style SSE format"""
-    try:
-        async for chunk in stream:
-            chunk = json.dumps(chunk)
-            try:
-                yield f"data:{chunk}\n\n"
-            except Exception as e:
-                yield f"data:{str(e)}\n\n"
-    except Exception as e:
-        yield f"data: {str(e)}\n\n"
-    finally:
-        yield "data: [DONE]\n\n"
-
-
-async def convert_to_async_iterator(
-    sync_iterator: Iterator[CreateChatCompletionStreamResponse],
-) -> AsyncIterator[CreateChatCompletionStreamResponse]:
+async def completion_to_async_iterator(
+    sync_iterator: Iterator[dict],
+) -> AsyncIterator[LegacyCompletion]:
     """
     Convert a synchronous iterator to an asynchronous iterator. This makes the logic easier
     because both the pipeline and the completion handler can use async iterators.
     """
     for item in sync_iterator:
-        yield item
-        await asyncio.sleep(0)
+        yield LegacyCompletion(**item)
+
+
+async def chat_to_async_iterator(
+        sync_iterator: Iterator[dict],
+) -> AsyncIterator[StreamingChatCompletion]:
+    for item in sync_iterator:
+        yield StreamingChatCompletion(**item)
 
 
 class LlamaCppCompletionHandler(BaseCompletionHandler):
@@ -57,33 +68,60 @@ async def execute_completion(
         """
         Execute the completion request with inference engine API
         """
-        model_path = f"{request['base_url']}/{request['model']}.gguf"
+        model_path = f"{base_url}/{request.get_model()}.gguf"
 
         # Create a copy of the request dict and remove stream_options
         # Reason - Request error as JSON:
         # {'error': "Llama.create_completion() got an unexpected keyword argument 'stream_options'"}
-        request_dict = dict(request)
-        request_dict.pop("stream_options", None)
-        # Remove base_url from the request dict. We use this field as a standard across
-        # all providers to specify the base URL of the model.
-        request_dict.pop("base_url", None)
-
         if is_fim_request:
+            request_dict = request.dict(exclude={
+                "best_of",
+                "frequency_pentalty",
+                "n",
+                "stream_options",
+                "user",
+            })
+
             response = await self.inference_engine.complete(
                 model_path,
                 Config.get_config().chat_model_n_ctx,
                 Config.get_config().chat_model_n_gpu_layers,
                 **request_dict,
             )
+
+            if stream:
+                return completion_to_async_iterator(response)
+            return LegacyCompletion(**response)
         else:
+            request_dict = request.dict(exclude={
+                "audio",
+                "frequency_pentalty",
+                "include_reasoning",
+                "metadata",
+                "max_completion_tokens",
+                "modalities",
+                "n",
+                "parallel_tool_calls",
+                "prediction",
+                "prompt",
+                "reasoning_effort",
+                "service_tier",
+                "store",
+                "stream_options",
+                "user",
+            })
+
             response = await self.inference_engine.chat(
                 model_path,
                 Config.get_config().chat_model_n_ctx,
                 Config.get_config().chat_model_n_gpu_layers,
                 **request_dict,
             )
 
-        return convert_to_async_iterator(response) if stream else response
+            if stream:
+                return chat_to_async_iterator(response)
+            else:
+                return StreamingChatCompletion(**response)
 
     def _create_streaming_response(
         self,
@@ -95,7 +133,7 @@ def _create_streaming_response(
         is the format that FastAPI expects for streaming responses.
         """
         return StreamingResponse(
-            llamacpp_stream_generator(stream),
+            stream_generator(stream),
             headers={
                 "Cache-Control": "no-cache",
                 "Connection": "keep-alive",

src/codegate/providers/llamacpp/provider.py

@@ -12,6 +12,11 @@
 from codegate.providers.base import BaseProvider, ModelFetchError
 from codegate.providers.fim_analyzer import FIMAnalyzer
 from codegate.providers.llamacpp.completion_handler import LlamaCppCompletionHandler
+from codegate.types.openai import (
+    ChatCompletionRequest,
+    LegacyCompletionRequest,
+)
+
 
 logger = structlog.get_logger("codegate")
 
@@ -21,6 +26,10 @@ def __init__(
         self,
         pipeline_factory: PipelineFactory,
     ):
+        if self._get_base_url() != "":
+            self.base_url = self._get_base_url()
+        else:
+            self.base_url = "./codegate_volume/models"
         completion_handler = LlamaCppCompletionHandler()
         super().__init__(
             None,
@@ -83,17 +92,32 @@ def _setup_routes(self):
         """
 
         @self.router.post(f"/{self.provider_route_name}/completions")
+        @DetectClient()
+        async def create_completion(
+            request: Request,
+        ):
+            body = await request.body()
+            print(body)
+            req = LegacyCompletionRequest.model_validate_json(body)
+            is_fim_request = FIMAnalyzer.is_fim_request(request.url.path, req)
+            return await self.process_request(
+                req,
+                None,
+                is_fim_request,
+                request.state.detected_client,
+            )
+
         @self.router.post(f"/{self.provider_route_name}/chat/completions")
         @DetectClient()
         async def create_completion(
             request: Request,
         ):
             body = await request.body()
-            data = json.loads(body)
-            data["base_url"] = Config.get_config().model_base_path
-            is_fim_request = FIMAnalyzer.is_fim_request(request.url.path, data)
+            print(body)
+            req = ChatCompletionRequest.model_validate_json(body)
+            is_fim_request = FIMAnalyzer.is_fim_request(request.url.path, req)
             return await self.process_request(
-                data,
+                req,
                 None,
                 is_fim_request,
                 request.state.detected_client,

src/codegate/types/google/_request_models.py

@@ -0,0 +1,17 @@
+from typing import (
+    Iterable,
+    List,
+)
+
+import pydantic
+
+
+class GenerateContentRequest(pydantic.BaseModel):
+  model: str | None = None
+  contents: List[Content] | None = None
+  config: Config | None = None
+
+  def get_messages(self) -> Iterable[Content]:
+      if self.content is not None:
+          for content in self.content:
+              yield content

src/codegate/types/openai/__init__.py

@@ -22,7 +22,6 @@
     StreamingChatCompletion,
     ToolCall,
     Usage,
-    VllmMessageError,
 )
 
 from ._request_models import (
@@ -54,3 +53,13 @@
 from ._shared_models import (
     ServiceTier,
 )
+
+from ._legacy_models import (
+    LegacyCompletionRequest,
+    LegacyCompletionTokenDetails,
+    LegacyPromptTokenDetails,
+    LegacyUsage,
+    LegacyLogProbs,
+    LegacyMessage,
+    LegacyCompletion,
+)

src/codegate/types/openai/_legacy_models.py

@@ -0,0 +1,136 @@
+from typing import (
+    Any,
+    Iterable,
+    List,
+    Literal,
+)
+
+import pydantic
+
+from ._response_models import (
+    Usage,
+)
+from ._request_models import (
+    StreamOption,
+    Message,
+)
+
+
+class LegacyCompletionRequest(pydantic.BaseModel):
+    prompt: str | None = None
+    model: str
+    best_of: int | None = 1
+    echo: bool | None = False
+    frequency_pentalty: float | None = 0.0
+    logit_bias: dict | None = None
+    logprobs: int | None = None
+    max_tokens: int | None = None
+    n: int | None = None
+    presence_penalty: float | None = 0.0
+    seed: int | None = None
+    stop: str | List[Any] | None = None
+    stream: bool | None = False
+    stream_options: StreamOption | None = None
+    suffix: str | None = None
+    temperature: float | None = 1.0
+    top_p: float | None = 1.0
+    user: str | None = None
+
+    def get_stream(self) -> bool:
+        return self.stream
+
+    def get_model(self) -> str:
+        return self.model
+
+    def get_messages(self) -> Iterable[Message]:
+        yield self
+
+    def get_content(self) -> Iterable[Any]:
+        yield self
+
+    def get_text(self) -> str | None:
+        return self.prompt
+
+    def set_text(self, text) -> None:
+        self.prompt = text
+
+    def first_message(self) -> Message | None:
+        return self
+
+    def last_user_message(self) -> tuple[Message, int] | None:
+        return self, 0
+
+    def last_user_block(self) -> Iterable[tuple[Message, int]]:
+        yield self, 0
+
+    def get_system_prompt(self) -> Iterable[str]:
+        yield self.get_text()
+
+    def set_system_prompt(self, text) -> None:
+        self.set_text(text)
+
+    def add_system_prompt(self, text, sep="\n") -> None:
+        original = self.get_text()
+        self.set_text(f"{original}{sep}{text}")
+
+    def get_prompt(self, default=None):
+        if self.prompt is not None:
+            return self.get_text()
+        return default
+
+
+class LegacyCompletionTokenDetails(pydantic.BaseModel):
+    accepted_prediction_tokens: int
+    audio_tokens: int
+    reasoning_tokens: int
+
+
+class LegacyPromptTokenDetails(pydantic.BaseModel):
+    audio_tokens: int
+    cached_tokens: int
+
+
+class LegacyUsage(pydantic.BaseModel):
+    completion_tokens: int
+    prompt_tokens: int
+    total_tokens: int
+    completion_tokens_details: LegacyCompletionTokenDetails | None = None
+    prompt_tokens_details: LegacyPromptTokenDetails | None = None
+
+
+class LegacyLogProbs(pydantic.BaseModel):
+    text_offset: List[Any]
+    token_logprobs: List[Any]
+    tokens: List[Any]
+    top_logprobs: List[Any]
+
+
+class LegacyMessage(pydantic.BaseModel):
+    text: str
+    finish_reason: str | None = None
+    index: int = 0
+    logprobs: LegacyLogProbs | None = None
+
+    def get_text(self) -> str | None:
+        return self.text
+
+    def set_text(self, text) -> None:
+        self.text = text
+
+
+class LegacyCompletion(pydantic.BaseModel):
+    id: str
+    choices: List[LegacyMessage]
+    created: int
+    model: str
+    system_fingerprint: str | None = None
+    object: Literal["text_completion"] = "text_completion"
+    usage: Usage | None = None
+
+    def get_content(self) -> Iterable[LegacyMessage]:
+        for message in self.choices:
+            yield message
+
+    def set_text(self, text) -> None:
+        if self.choices:
+            self.choices[0].set_text(text)