Add LM Studio provider

src/codegate/providers/lm_studio/provider.py

@@ -0,0 +1,56 @@
+import json
+
+from fastapi import Header, HTTPException, Request
+from fastapi.responses import JSONResponse
+
+from codegate.config import Config
+from codegate.pipeline.factory import PipelineFactory
+from codegate.providers.openai.provider import OpenAIProvider
+
+
+class LmStudioProvider(OpenAIProvider):
+    def __init__(
+        self,
+        pipeline_factory: PipelineFactory,
+    ):
+        config = Config.get_config()
+        if config is not None:
+            provided_urls = config.provider_urls
+            self.lm_studio_url = provided_urls.get("lm_studio", "http://localhost:11434/")
+
+        super().__init__(pipeline_factory)
+
+    @property
+    def provider_route_name(self) -> str:
+        return "lm_studio"
+
+    def _setup_routes(self):
+        """
+        Sets up the /chat/completions route for the provider as expected by the
+        LM Studio API. Extracts the API key from the "Authorization" header and
+        passes it to the completion handler.
+        """
+
+        @self.router.get(f"/{self.provider_route_name}/models")
+        @self.router.get(f"/{self.provider_route_name}/v1/models")
+        async def get_models():
+            # dummy method for lm studio
+            return JSONResponse(status_code=200, content=[])
+
+        @self.router.post(f"/{self.provider_route_name}/chat/completions")
+        @self.router.post(f"/{self.provider_route_name}/completions")
+        @self.router.post(f"/{self.provider_route_name}/v1/chat/completions")
+        async def create_completion(
+            request: Request,
+            authorization: str = Header(..., description="Bearer token"),
+        ):
+            if not authorization.startswith("Bearer "):
+                raise HTTPException(status_code=401, detail="Invalid authorization header")
+
+            api_key = authorization.split(" ")[1]
+            body = await request.body()
+            data = json.loads(body)
+
+            data["base_url"] = self.lm_studio_url + "/v1/"
+
+            return await self.process_request(data, api_key, request)

src/codegate/providers/openai/provider.py

@@ -4,9 +4,7 @@
 import httpx
 import structlog
 from fastapi import Header, HTTPException, Request
-from fastapi.responses import JSONResponse
 
-from codegate.config import Config
 from codegate.pipeline.factory import PipelineFactory
 from codegate.providers.base import BaseProvider, ModelFetchError
 from codegate.providers.litellmshim import LiteLLmShim, sse_stream_generator
@@ -19,11 +17,6 @@ def __init__(
         pipeline_factory: PipelineFactory,
     ):
         completion_handler = LiteLLmShim(stream_generator=sse_stream_generator)
-        config = Config.get_config()
-        if config is not None:
-            provided_urls = config.provider_urls
-            self.lm_studio_url = provided_urls.get("lm_studio", "http://localhost:11434/")
-
         super().__init__(
             OpenAIInputNormalizer(),
             OpenAIOutputNormalizer(),
@@ -39,8 +32,6 @@ def models(self, endpoint: str = None, api_key: str = None) -> List[str]:
         headers = {}
         if api_key:
             headers["Authorization"] = f"Bearer {api_key}"
-        if not endpoint:
-            endpoint = "https://api.openai.com"
 
         resp = httpx.get(f"{endpoint}/v1/models", headers=headers)
 
@@ -51,19 +42,32 @@ def models(self, endpoint: str = None, api_key: str = None) -> List[str]:
 
         return [model["id"] for model in jsonresp.get("data", [])]
 
+    async def process_request(self, data: dict, api_key: str, request: Request):
+        """
+        Process the request and return the completion stream
+        """
+        is_fim_request = self._is_fim_request(request, data)
+        try:
+            stream = await self.complete(data, api_key, is_fim_request=is_fim_request)
+        except Exception as e:
+            #  check if we have an status code there
+            if hasattr(e, "status_code"):
+                logger = structlog.get_logger("codegate")
+                logger.error("Error in OpenAIProvider completion", error=str(e))
+
+                raise HTTPException(status_code=e.status_code, detail=str(e))  # type: ignore
+            else:
+                # just continue raising the exception
+                raise e
+        return self._completion_handler.create_response(stream)
+
     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.get(f"/{self.provider_route_name}/models")
-        @self.router.get(f"/{self.provider_route_name}/v1/models")
-        async def get_models():
-            # dummy method for lm studio
-            return JSONResponse(status_code=200, content=[])
-
         @self.router.post(f"/{self.provider_route_name}/chat/completions")
         @self.router.post(f"/{self.provider_route_name}/completions")
         @self.router.post(f"/{self.provider_route_name}/v1/chat/completions")
@@ -78,20 +82,4 @@ async def create_completion(
             body = await request.body()
             data = json.loads(body)
 
-            # if model starts with lm_studio, propagate it
-            if data.get("model", "").startswith("lm_studio"):
-                data["base_url"] = self.lm_studio_url + "/v1/"
-            is_fim_request = self._is_fim_request(request, data)
-            try:
-                stream = await self.complete(data, api_key, is_fim_request=is_fim_request)
-            except Exception as e:
-                #  check if we have an status code there
-                if hasattr(e, "status_code"):
-                    logger = structlog.get_logger("codegate")
-                    logger.error("Error in OpenAIProvider completion", error=str(e))
-
-                    raise HTTPException(status_code=e.status_code, detail=str(e))  # type: ignore
-                else:
-                    # just continue raising the exception
-                    raise e
-            return self._completion_handler.create_response(stream)
+            return await self.process_request(data, api_key, request)

src/codegate/server.py

@@ -13,6 +13,7 @@
 from codegate.pipeline.factory import PipelineFactory
 from codegate.providers.anthropic.provider import AnthropicProvider
 from codegate.providers.llamacpp.provider import LlamaCppProvider
+from codegate.providers.lm_studio.provider import LmStudioProvider
 from codegate.providers.ollama.provider import OllamaProvider
 from codegate.providers.openai.provider import OpenAIProvider
 from codegate.providers.registry import ProviderRegistry, get_provider_registry
@@ -96,6 +97,12 @@ async def log_user_agent(request: Request, call_next):
             pipeline_factory,
         ),
     )
+    registry.add_provider(
+        "lm_studio",
+        LmStudioProvider(
+            pipeline_factory,
+        ),
+    )
 
     # Create and add system routes
     system_router = APIRouter(tags=["System"])

tests/test_server.py

@@ -108,8 +108,8 @@ def test_provider_registration(mock_registry, mock_secrets_mgr, mock_pipeline_fa
     # Verify all providers were registered
     registry_instance = mock_registry.return_value
     assert (
-        registry_instance.add_provider.call_count == 5
-    )  # openai, anthropic, llamacpp, vllm, ollama
+        registry_instance.add_provider.call_count == 6
+    )  # openai, anthropic, llamacpp, vllm, ollama, lm_studio
 
     # Verify specific providers were registered
     provider_names = [call.args[0] for call in registry_instance.add_provider.call_args_list]