Ported OpenAI and OpenRouter providers.

blkt · blkt · commit 17eae23002a0 · 2025-02-20T13:42:03.000+01:00
Also, tests and fixes.
diff --git a/src/codegate/pipeline/comment/output.py b/src/codegate/pipeline/comment/output.py
@@ -130,12 +130,10 @@ async def process_chunk(
         input_context: Optional[PipelineContext] = None,
     ) -> list[ModelResponse]:
         """Process a single chunk of the stream"""
-        # if len(chunk.choices) == 0 or not chunk.choices[0].delta.content:
-        #     return [chunk]
-
         for content in chunk.get_content():
             # Get current content plus this new chunk
-            current_content = "".join(context.processed_content + [txt for txt in content.get_text()])
+            text = content.get_text()
+            current_content = "".join(context.processed_content + [text if text else ""])
 
             # Extract snippets from current content
             snippets = self.extractor.extract_snippets(current_content)
diff --git a/src/codegate/pipeline/secrets/secrets.py b/src/codegate/pipeline/secrets/secrets.py
@@ -502,7 +502,8 @@ async def process_chunk(
                 return []
 
             # No markers or partial markers, let pipeline handle the chunk normally
-            content.set_text(context.prefix_buffer + content.get_text())
+            text = content.get_text()
+            content.set_text(context.prefix_buffer + text if text else "")
             context.prefix_buffer = ""
             return [chunk]
         else:
diff --git a/src/codegate/providers/anthropic/provider.py b/src/codegate/providers/anthropic/provider.py
@@ -28,7 +28,7 @@ def __init__(
         if self._get_base_url() != "":
             self.base_url = self._get_base_url()
         else:
-            self.base_url = "https://api.anthropic.com/"
+            self.base_url = "https://api.anthropic.com"
 
         completion_handler = AnthropicCompletion(stream_generator=stream_generator)
         super().__init__(
diff --git a/src/codegate/providers/litellmshim/litellmshim.py b/src/codegate/providers/litellmshim/litellmshim.py
@@ -50,8 +50,8 @@ async def execute_completion(
         Execute the completion request with LiteLLM's API
         """
         if is_fim_request:
-            return self._fim_completion_func(request, api_key=api_key)
-        return self._completion_func(request, api_key=api_key)
+            return self._fim_completion_func(request, api_key=api_key, base_url=base_url)
+        return self._completion_func(request, api_key=api_key, base_url=base_url)
 
     def _create_streaming_response(
         self,
diff --git a/src/codegate/providers/openai/provider.py b/src/codegate/providers/openai/provider.py
@@ -11,20 +11,32 @@
 from codegate.providers.base import BaseProvider, ModelFetchError
 from codegate.providers.fim_analyzer import FIMAnalyzer
 from codegate.providers.litellmshim import LiteLLmShim
-from codegate.providers.openai.adapter import OpenAIInputNormalizer, OpenAIOutputNormalizer
-from codegate.types.generators import sse_stream_generator
+from codegate.types.openai import (
+    completions_streaming,
+    stream_generator,
+    ChatCompletionRequest,
+)
 
 
 class OpenAIProvider(BaseProvider):
     def __init__(
         self,
         pipeline_factory: PipelineFactory,
         # Enable receiving other completion handlers from childs, i.e. OpenRouter and LM Studio
-        completion_handler: LiteLLmShim = LiteLLmShim(stream_generator=sse_stream_generator),
+        completion_handler: LiteLLmShim = LiteLLmShim(completion_func=completions_streaming, stream_generator=stream_generator),
     ):
+        if self._get_base_url() != "":
+            self.base_url = self._get_base_url()
+        else:
+            self.base_url = "https://api.openai.com/api/v1"
+
+        completion_handler = LiteLLmShim(
+            completion_func=completions_streaming,
+            stream_generator=stream_generator,
+        )
         super().__init__(
-            OpenAIInputNormalizer(),
-            OpenAIOutputNormalizer(),
+            None,
+            None,
             completion_handler,
             pipeline_factory,
         )
@@ -93,11 +105,11 @@ async def create_completion(
 
             api_key = authorization.split(" ")[1]
             body = await request.body()
-            data = json.loads(body)
-            is_fim_request = FIMAnalyzer.is_fim_request(request.url.path, data)
+            req = ChatCompletionRequest.model_validate_json(body)
+            is_fim_request = FIMAnalyzer.is_fim_request(request.url.path, req)
 
             return await self.process_request(
-                data,
+                req,
                 api_key,
                 is_fim_request,
                 request.state.detected_client,
diff --git a/src/codegate/providers/openrouter/provider.py b/src/codegate/providers/openrouter/provider.py
@@ -2,66 +2,24 @@
 from typing import Dict
 
 from fastapi import Header, HTTPException, Request
-from litellm import atext_completion
-from litellm.types.llms.openai import ChatCompletionRequest
 
 from codegate.clients.clients import ClientType
 from codegate.clients.detector import DetectClient
 from codegate.pipeline.factory import PipelineFactory
 from codegate.providers.fim_analyzer import FIMAnalyzer
-from codegate.providers.litellmshim import LiteLLmShim, sse_stream_generator
-from codegate.providers.normalizer.completion import CompletionNormalizer
 from codegate.providers.openai import OpenAIProvider
-
-
-class OpenRouterNormalizer(CompletionNormalizer):
-    def __init__(self):
-        super().__init__()
-
-    def normalize(self, data: Dict) -> ChatCompletionRequest:
-        return super().normalize(data)
-
-    def denormalize(self, data: ChatCompletionRequest) -> Dict:
-        """
-        Denormalize a FIM OpenRouter request. Force it to be an accepted atext_completion format.
-        """
-        denormalized_data = super().denormalize(data)
-        # We are forcing atext_completion which expects to have a "prompt" key in the data
-        # Forcing it in case is not present
-        if "prompt" in data:
-            return denormalized_data
-        custom_prompt = ""
-        for msg_dict in denormalized_data.get("messages", []):
-            content_obj = msg_dict.get("content")
-            if not content_obj:
-                continue
-            if isinstance(content_obj, list):
-                for content_dict in content_obj:
-                    custom_prompt += (
-                        content_dict.get("text", "") if isinstance(content_dict, dict) else ""
-                    )
-            elif isinstance(content_obj, str):
-                custom_prompt += content_obj
-
-        # Erase the original "messages" key. Replace it by "prompt"
-        del denormalized_data["messages"]
-        denormalized_data["prompt"] = custom_prompt
-
-        return denormalized_data
+from codegate.types.openai import (
+        ChatCompletionRequest,
+)
 
 
 class OpenRouterProvider(OpenAIProvider):
     def __init__(self, pipeline_factory: PipelineFactory):
-        super().__init__(
-            pipeline_factory,
-            # We get FIM requests in /completions. LiteLLM is forcing /chat/completions
-            # which returns "choices":[{"delta":{"content":"some text"}}]
-            # instead of "choices":[{"text":"some text"}] expected by the client (Continue)
-            completion_handler=LiteLLmShim(
-                stream_generator=sse_stream_generator, fim_completion_func=atext_completion
-            ),
-        )
-        self._fim_normalizer = OpenRouterNormalizer()
+        super().__init__(pipeline_factory)
+        if self._get_base_url() != "":
+            self.base_url = self._get_base_url()
+        else:
+            self.base_url = "https://openrouter.ai/api/v1"
 
     @property
     def provider_route_name(self) -> str:
@@ -74,12 +32,6 @@ async def process_request(
         is_fim_request: bool,
         client_type: ClientType,
     ):
-        # litellm workaround - add openrouter/ prefix to model name to make it openai-compatible
-        # once we get rid of litellm, this can simply be removed
-        original_model = data.get("model", "")
-        if not original_model.startswith("openrouter/"):
-            data["model"] = f"openrouter/{original_model}"
-
         return await super().process_request(data, api_key, is_fim_request, client_type)
 
     def _setup_routes(self):
@@ -96,14 +48,12 @@ async def create_completion(
 
             api_key = authorization.split(" ")[1]
             body = await request.body()
-            data = json.loads(body)
 
-            base_url = self._get_base_url()
-            data["base_url"] = base_url
-            is_fim_request = FIMAnalyzer.is_fim_request(request.url.path, data)
+            req = ChatCompletionRequest.model_validate_json(body)
+            is_fim_request = FIMAnalyzer.is_fim_request(request.url.path, req)
 
             return await self.process_request(
-                data,
+                req,
                 api_key,
                 is_fim_request,
                 request.state.detected_client,
diff --git a/src/codegate/types/anthropic/_generators.py b/src/codegate/types/anthropic/_generators.py
@@ -65,7 +65,7 @@ async def stream_generator(stream: AsyncIterator[Any]) -> AsyncIterator[str]:
         yield f"event: error\ndata: {body}\n\n"
 
 
-async def acompletion(request, api_key):
+async def acompletion(request, api_key, base_url):
     headers = {
         "anthropic-version": "2023-06-01",
         "x-api-key": api_key,
@@ -79,7 +79,7 @@ async def acompletion(request, api_key):
 
     client = httpx.AsyncClient()
     async with client.stream(
-            "POST", "https://api.anthropic.com/v1/messages",
+            "POST", f"{base_url}/v1/messages",
             headers=headers,
             content=payload,
             timeout=30, # TODO this should not be hardcoded
@@ -90,9 +90,11 @@ async def acompletion(request, api_key):
                 async for event in message_wrapper(resp.aiter_lines()):
                     yield event
             case 400 | 401 | 403 | 404 | 413 | 429:
-                yield MessageError.model_validate_json(resp.text)
+                text = await resp.aread()
+                yield MessageError.model_validate_json(text)
             case 500 | 529:
-                yield MessageError.model_validate_json(resp.text)
+                text = await resp.aread()
+                yield MessageError.model_validate_json(text)
             case _:
                 logger.error(f"unexpected status code {resp.status_code}", provider="anthropic")
                 raise ValueError(f"unexpected status code {resp.status_code}", provider="anthropic")
diff --git a/src/codegate/types/anthropic/_request_models.py b/src/codegate/types/anthropic/_request_models.py
@@ -106,7 +106,7 @@ class AssistantMessage(pydantic.BaseModel):
 
     def get_text(self) -> Iterable[str]:
         if isinstance(self.content, str):
-            yield self.content
+            return self.content
 
     def set_text(self, text) -> None:
         if isinstance(self.content, str):
diff --git a/src/codegate/types/anthropic/_response_models.py b/src/codegate/types/anthropic/_response_models.py
@@ -33,7 +33,7 @@ class ToolUseResponseContent(pydantic.BaseModel):
     name: str
 
     def get_text(self):
-        return iter(()) # empty generator
+        return None
 
     def set_text(self, text):
         pass
@@ -97,7 +97,7 @@ class ToolUse(pydantic.BaseModel):
     input: Dict
 
     def get_text(self) -> str | None:
-        return ""
+        return None
 
     def set_text(self, text):
         pass
diff --git a/src/codegate/types/openai/__init__.py b/src/codegate/types/openai/__init__.py
@@ -3,6 +3,7 @@
 
 from ._generators import (
     completions_streaming,
+    message_wrapper,
     stream_generator,
 )
 
@@ -12,11 +13,13 @@
     Choice,
     ChoiceDelta,
     CompletionTokenDetails,
+    ErrorDetails,
     FunctionCall,
     LogProbs,
     LogProbsContent,
     Message,
     MessageDelta,
+    MessageError,
     PromptTokenDetails,
     RawLogProbsContent,
     StreamingChatCompletion,
diff --git a/src/codegate/types/openai/_generators.py b/src/codegate/types/openai/_generators.py
@@ -7,6 +7,8 @@
 import structlog
 
 from ._response_models import (
+    ErrorDetails,
+    MessageError,
     StreamingChatCompletion,
 )
 
@@ -31,36 +33,52 @@ async def stream_generator(stream: AsyncIterator[StreamingChatCompletion]) -> As
                 yield f"data: {str(e)}\n\n"
     except Exception as e:
         logger.error("failed generating output payloads", exc_info=e)
-        yield f"data: {str(e)}\n\n"
+        err = MessageError(
+            error=ErrorDetails(
+                message=str(e),
+                code=500,
+            ),
+        )
+        data = err.model_dump_json(exclude_none=True, exclude_unset=True)
+        yield f"data: {data}\n\n"
     finally:
+        # Note: I'm not sure this is sent when an error is triggered
+        # during SSE processing.
         yield "data: [DONE]\n\n"
 
 
 async def completions_streaming(request, api_key, base_url):
     if base_url is None:
         base_url = "https://api.openai.com"
-    return streaming(request, api_key, f"{base_url}/v1/chat/completions")
+    async for item in  streaming(request, api_key, f"{base_url}/chat/completions"):
+        yield item
 
 
 async def streaming(request, api_key, url):
+    headers = {
+        "Content-Type": "application/json",
+        "Authorization": f"Bearer {api_key}",
+    }
     payload = request.json(exclude_defaults=True)
     if os.getenv("CODEGATE_DEBUG_OPENAI") is not None:
         print(payload)
+        print(headers)
 
     client = httpx.AsyncClient()
     async with client.stream(
             "POST", url,
+            headers=headers,
             content=payload,
             timeout=30, # TODO this should not be hardcoded
     ) as resp:
         # TODO figure out how to best return failures
         match resp.status_code:
             case 200:
-                async for message in parser(resp.aiter_lines()):
+                async for message in message_wrapper(resp.aiter_lines()):
                     yield message
             case 400 | 401 | 403 | 404 | 413 | 429:
-                logger.error(f"unexpected status code {resp.status_code}: {resp.text}", provider="openai")
-                yield MessageError.model_validate_json(resp.text)
+                text = await resp.aread()
+                yield MessageError.model_validate_json(text)
             # case 500 | 529:
             #     yield MessageError.model_validate_json(resp.text)
             case _:
@@ -73,20 +91,31 @@ async def get_data_lines(lines):
     while True:
         # Get the `data: <type>` line.
         data_line = await anext(lines)
-        # Get the empty line.
-        _ = await anext(lines)
+
+        # As per standard, we ignore comment lines
+        # https://html.spec.whatwg.org/multipage/server-sent-events.html#event-stream-interpretation
+        if data_line.startswith(":"):
+            continue
 
         count = count + 1
 
         if "[DONE]" in data_line:
             break
 
         yield data_line[6:]
-    logger.debug(f"Consumed {count} messages", provider="anthropic", count=count)
 
+        # Get the empty line.
+        _ = await anext(lines)
+    logger.debug(f"Consumed {count} messages", provider="openai", count=count)
 
-async def parser(lines):
+
+async def message_wrapper(lines):
     messages = get_data_lines(lines)
     async for payload in messages:
-        item = StreamingChatCompletion.model_validate_json(payload)
-        yield item
+        try:
+            item = StreamingChatCompletion.model_validate_json(payload)
+            yield item
+        except Exception as e:
+            logger.warn("HTTP error while consuming SSE stream", exc_info=e)
+            item = MessageError.model_validate_json(payload)
+            yield item
diff --git a/src/codegate/types/openai/_response_models.py b/src/codegate/types/openai/_response_models.py
diff --git a/tests/types/openai/streaming_messages.txt b/tests/types/openai/streaming_messages.txt
diff --git a/tests/types/openai/test_openai.py b/tests/types/openai/test_openai.py
