diff --git a/src/codegate/providers/anthropic/provider.py b/src/codegate/providers/anthropic/provider.py index efe27225..780cc0b3 100644 --- a/src/codegate/providers/anthropic/provider.py +++ b/src/codegate/providers/anthropic/provider.py @@ -13,6 +13,10 @@ def __init__(self): completion_handler = LiteLLmShim(adapter) super().__init__(completion_handler) + @property + def provider_route_name(self) -> str: + return "anthropic" + def _setup_routes(self): """ Sets up the /messages route for the provider as expected by the Anthropic @@ -20,7 +24,7 @@ def _setup_routes(self): completion handler. """ - @self.router.post("/messages") + @self.router.post(f"/{self.provider_route_name}/messages") async def create_message( request: Request, x_api_key: str = Header(None), diff --git a/src/codegate/providers/base.py b/src/codegate/providers/base.py index bd937873..07c3f042 100644 --- a/src/codegate/providers/base.py +++ b/src/codegate/providers/base.py @@ -39,6 +39,11 @@ def __init__(self, completion_handler: BaseCompletionHandler): def _setup_routes(self) -> None: pass + @property + @abstractmethod + def provider_route_name(self) -> str: + pass + async def complete(self, data: Dict, api_key: str) -> AsyncIterator[Any]: return await self._completion_handler.complete(data, api_key) diff --git a/src/codegate/providers/llamacpp/completion_handler.py b/src/codegate/providers/llamacpp/completion_handler.py index fc0a991d..863c5694 100644 --- a/src/codegate/providers/llamacpp/completion_handler.py +++ b/src/codegate/providers/llamacpp/completion_handler.py @@ -1,7 +1,7 @@ from typing import Any, AsyncIterator, Dict +from litellm import ModelResponse from fastapi.responses import StreamingResponse -from litellm import ModelResponse, acompletion from codegate.providers.base import BaseCompletionHandler from codegate.providers.llamacpp.adapter import BaseAdapter diff --git a/src/codegate/providers/llamacpp/provider.py b/src/codegate/providers/llamacpp/provider.py index d70d370a..5590c622 100644 --- a/src/codegate/providers/llamacpp/provider.py +++ b/src/codegate/providers/llamacpp/provider.py @@ -1,6 +1,6 @@ import json -from fastapi import Header, HTTPException, Request +from fastapi import Request from codegate.providers.base import BaseProvider from codegate.providers.llamacpp.completion_handler import LlamaCppCompletionHandler @@ -13,13 +13,17 @@ def __init__(self): completion_handler = LlamaCppCompletionHandler(adapter) super().__init__(completion_handler) + @property + def provider_route_name(self) -> str: + return "llamacpp" + def _setup_routes(self): """ Sets up the /chat route for the provider as expected by the Llama API. Extracts the API key from the "Authorization" header and passes it to the completion handler. """ - @self.router.post("/completion") + @self.router.post(f"/{self.provider_route_name}/completion") async def create_completion( request: Request, ): diff --git a/src/codegate/providers/openai/provider.py b/src/codegate/providers/openai/provider.py index 1fd7240d..0d118862 100644 --- a/src/codegate/providers/openai/provider.py +++ b/src/codegate/providers/openai/provider.py @@ -13,6 +13,10 @@ def __init__(self): completion_handler = LiteLLmShim(adapter) super().__init__(completion_handler) + @property + def provider_route_name(self) -> str: + return "openai" + def _setup_routes(self): """ Sets up the /chat/completions route for the provider as expected by the @@ -20,7 +24,7 @@ def _setup_routes(self): passes it to the completion handler. """ - @self.router.post("/chat/completions") + @self.router.post(f"/{self.provider_route_name}/chat/completions") async def create_completion( request: Request, authorization: str = Header(..., description="Bearer token"), diff --git a/tests/providers/test_registry.py b/tests/providers/test_registry.py index 24ed382c..dafd4fd7 100644 --- a/tests/providers/test_registry.py +++ b/tests/providers/test_registry.py @@ -20,8 +20,13 @@ def create_streaming_response( class MockProvider(BaseProvider): + + @property + def provider_route_name(self) -> str: + return 'mock_provider' + def _setup_routes(self) -> None: - @self.router.get("/test") + @self.router.get(f"/{self.provider_route_name}/test") def test_route(): return {"message": "test"} @@ -61,5 +66,5 @@ def test_provider_routes_added(app, registry, mock_completion_handler): provider = MockProvider(mock_completion_handler) registry.add_provider("test", provider) - routes = [route for route in app.routes if route.path == "/test"] + routes = [route for route in app.routes if route.path == "/mock_provider/test"] assert len(routes) == 1