Specify the provider in the URL to properly route traffic

aponcedeleonch · aponcedeleonch · commit af50968bf59f · 2024-11-26T13:39:18.000+02:00
Closes: #93 To route to the appropriate provider in CodeGate
diff --git a/src/codegate/providers/anthropic/provider.py b/src/codegate/providers/anthropic/provider.py
@@ -13,14 +13,18 @@ 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
         API. Extracts the API key from the "x-api-key" header and passes it to the
         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
@@ -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
@@ -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
@@ -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
@@ -13,14 +13,18 @@ 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
         OpenAI API. Extracts the API key from the "Authorization" header and
         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
@@ -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
