Merge pull request #149 from stacklok/normalize-vllm-output

Respond with JSON if the request is non-stream

Author: aponcedeleonch

src/codegate/providers/anthropic/provider.py

@@ -48,4 +48,4 @@ async def create_message(
 
             is_fim_request = self._is_fim_request(request, data)
             stream = await self.complete(data, x_api_key, is_fim_request)
-            return self._completion_handler.create_streaming_response(stream)
+            return self._completion_handler.create_response(stream)

src/codegate/providers/completion/base.py

@@ -1,7 +1,8 @@
 from abc import ABC, abstractmethod
+from collections.abc import Iterator
 from typing import Any, AsyncIterator, Optional, Union
 
-from fastapi.responses import StreamingResponse
+from fastapi.responses import JSONResponse, StreamingResponse
 from litellm import ChatCompletionRequest, ModelResponse
 
 
@@ -23,5 +24,17 @@ async def execute_completion(
         pass
 
     @abstractmethod
-    def create_streaming_response(self, stream: AsyncIterator[Any]) -> StreamingResponse:
+    def _create_streaming_response(self, stream: AsyncIterator[Any]) -> StreamingResponse:
         pass
+
+    @abstractmethod
+    def _create_json_response(self, response: Any) -> JSONResponse:
+        pass
+
+    def create_response(self, response: Any) -> Union[JSONResponse, StreamingResponse]:
+        """
+        Create a FastAPI response from the completion response.
+        """
+        if isinstance(response, Iterator):
+            return self._create_streaming_response(response)
+        return self._create_json_response(response)

src/codegate/providers/litellmshim/litellmshim.py

@@ -1,10 +1,17 @@
 from typing import Any, AsyncIterator, Callable, Optional, Union
 
-from fastapi.responses import StreamingResponse
-from litellm import ChatCompletionRequest, ModelResponse, acompletion
+import structlog
+from fastapi.responses import JSONResponse, StreamingResponse
+from litellm import (
+    ChatCompletionRequest,
+    ModelResponse,
+    acompletion,
+)
 
 from codegate.providers.base import BaseCompletionHandler, StreamGenerator
 
+logger = structlog.get_logger("codegate")
+
 
 class LiteLLmShim(BaseCompletionHandler):
     """
@@ -42,7 +49,7 @@ async def execute_completion(
             return await self._fim_completion_func(**request)
         return await self._completion_func(**request)
 
-    def create_streaming_response(self, stream: AsyncIterator[Any]) -> StreamingResponse:
+    def _create_streaming_response(self, stream: AsyncIterator[Any]) -> StreamingResponse:
         """
         Create a streaming response from a stream generator. The StreamingResponse
         is the format that FastAPI expects for streaming responses.
@@ -56,3 +63,14 @@ def create_streaming_response(self, stream: AsyncIterator[Any]) -> StreamingResp
             },
             status_code=200,
         )
+
+    def _create_json_response(self, response: ModelResponse) -> JSONResponse:
+        """
+        Create a JSON FastAPI response from a ModelResponse object.
+        ModelResponse is obtained when the request is not streaming.
+        """
+        # ModelResponse is not a Pydantic object but has a json method we can use to serialize
+        if isinstance(response, ModelResponse):
+            return JSONResponse(status_code=200, content=response.json())
+        # Most of others objects in LiteLLM are Pydantic, we can use the model_dump method
+        return JSONResponse(status_code=200, content=response.model_dump())

src/codegate/providers/llamacpp/completion_handler.py

@@ -2,7 +2,7 @@
 import json
 from typing import Any, AsyncIterator, Iterator, Optional, Union
 
-from fastapi.responses import StreamingResponse
+from fastapi.responses import JSONResponse, StreamingResponse
 from litellm import ChatCompletionRequest, ModelResponse
 from llama_cpp.llama_types import (
     CreateChatCompletionStreamResponse,
@@ -75,7 +75,7 @@ async def execute_completion(
 
         return convert_to_async_iterator(response) if stream else response
 
-    def create_streaming_response(self, stream: AsyncIterator[Any]) -> StreamingResponse:
+    def _create_streaming_response(self, stream: AsyncIterator[Any]) -> StreamingResponse:
         """
         Create a streaming response from a stream generator. The StreamingResponse
         is the format that FastAPI expects for streaming responses.
@@ -89,3 +89,6 @@ def create_streaming_response(self, stream: AsyncIterator[Any]) -> StreamingResp
             },
             status_code=200,
         )
+
+    def _create_json_response(self, response: Any) -> JSONResponse:
+        raise NotImplementedError("JSON Reponse in LlamaCPP not implemented yet.")

src/codegate/providers/llamacpp/provider.py

@@ -43,4 +43,4 @@ async def create_completion(
 
             is_fim_request = self._is_fim_request(request, data)
             stream = await self.complete(data, None, is_fim_request=is_fim_request)
-            return self._completion_handler.create_streaming_response(stream)
+            return self._completion_handler.create_response(stream)

src/codegate/providers/openai/provider.py

@@ -49,4 +49,4 @@ async def create_completion(
 
             is_fim_request = self._is_fim_request(request, data)
             stream = await self.complete(data, api_key, is_fim_request=is_fim_request)
-            return self._completion_handler.create_streaming_response(stream)
+            return self._completion_handler.create_response(stream)

src/codegate/providers/vllm/provider.py

@@ -57,4 +57,4 @@ async def create_completion(
 
             is_fim_request = self._is_fim_request(request, data)
             stream = await self.complete(data, api_key, is_fim_request=is_fim_request)
-            return self._completion_handler.create_streaming_response(stream)
+            return self._completion_handler.create_response(stream)

tests/providers/litellmshim/test_litellmshim.py

@@ -117,7 +117,7 @@ async def mock_stream_gen():
     generator = mock_stream_gen()
 
     litellm_shim = LiteLLmShim(stream_generator=sse_stream_generator)
-    response = litellm_shim.create_streaming_response(generator)
+    response = litellm_shim._create_streaming_response(generator)
 
     # Verify response metadata
     assert isinstance(response, StreamingResponse)

tests/providers/test_registry.py

@@ -43,12 +43,15 @@ def execute_completion(
     ) -> Any:
         pass
 
-    def create_streaming_response(
+    def _create_streaming_response(
         self,
         stream: AsyncIterator[Any],
     ) -> StreamingResponse:
         return StreamingResponse(stream)
 
+    def _create_json_response(self, response: Any) -> Any:
+        raise NotImplementedError
+
 
 class MockInputNormalizer(ModelInputNormalizer):
     def normalize(self, data: Dict) -> Dict: