Bootstrap provider models on addition

When adding a provider via the API, this bootstrap the available models
for a provider.

Signed-off-by: Juan Antonio Osorio <ozz@stacklok.com>

Author: JAORMX

src/codegate/api/v1.py

@@ -109,20 +109,26 @@ async def get_provider_endpoint(
     status_code=201,
 )
 async def add_provider_endpoint(
-    request: v1_models.ProviderEndpoint,
+    request: v1_models.AddProviderEndpointRequest,
 ) -> v1_models.ProviderEndpoint:
     """Add a provider endpoint."""
     try:
         provend = await pcrud.add_endpoint(request)
     except AlreadyExistsError:
         raise HTTPException(status_code=409, detail="Provider endpoint already exists")
+    except ValueError as e:
+        raise HTTPException(
+            status_code=400,
+            detail=str(e),
+        )
     except ValidationError as e:
         # TODO: This should be more specific
         raise HTTPException(
             status_code=400,
             detail=str(e),
         )
     except Exception:
+        logger.exception("Error while adding provider endpoint")
         raise HTTPException(status_code=500, detail="Internal server error")
 
     return provend
@@ -160,6 +166,10 @@ async def update_provider_endpoint(
     try:
         request.id = provider_id
         provend = await pcrud.update_endpoint(request)
+    except provendcrud.ProviderNotFoundError:
+        raise HTTPException(status_code=404, detail="Provider endpoint not found")
+    except ValueError as e:
+        raise HTTPException(status_code=400, detail=str(e))
     except ValidationError as e:
         # TODO: This should be more specific
         raise HTTPException(

src/codegate/api/v1_models.py

@@ -250,6 +250,14 @@ def get_from_registry(self, registry: ProviderRegistry) -> Optional[BaseProvider
         return registry.get_provider(self.provider_type)
 
 
+class AddProviderEndpointRequest(ProviderEndpoint):
+    """
+    Represents a request to add a provider endpoint.
+    """
+
+    api_key: Optional[str] = None
+
+
 class ConfigureAuthMaterial(pydantic.BaseModel):
     """
     Represents a request to configure auth material for a provider.

src/codegate/providers/anthropic/provider.py

@@ -29,9 +29,18 @@ def __init__(
     def provider_route_name(self) -> str:
         return "anthropic"
 
-    def models(self) -> List[str]:
-        # TODO: This won't work since we need an API Key being set.
-        resp = httpx.get("https://api.anthropic.com/models")
+    def models(self, api_key: str = None) -> List[str]:
+        headers = {
+            "Content-Type": "application/json",
+            "anthropic-version": "2023-06-01",
+        }
+        if api_key:
+            headers["x-api-key"] = api_key
+
+        resp = httpx.get(
+            "https://api.anthropic.com/v1/models",
+            headers=headers,
+        )
         # If Anthropic returned 404, it means it's not accepting our
         # requests. We should throw an error.
         if resp.status_code == 404:

src/codegate/providers/base.py

@@ -55,7 +55,7 @@ def _setup_routes(self) -> None:
         pass
 
     @abstractmethod
-    def models(self) -> List[str]:
+    def models(self, api_key: str = None) -> List[str]:
         pass
 
     @property

src/codegate/providers/crud/crud.py

@@ -11,7 +11,7 @@
 from codegate.db import models as dbmodels
 from codegate.db.connection import DbReader, DbRecorder
 from codegate.providers.base import BaseProvider
-from codegate.providers.registry import ProviderRegistry
+from codegate.providers.registry import ProviderRegistry, get_provider_registry
 
 logger = structlog.get_logger("codegate")
 
@@ -62,22 +62,60 @@ async def get_endpoint_by_name(self, name: str) -> Optional[apimodelsv1.Provider
         return apimodelsv1.ProviderEndpoint.from_db_model(dbendpoint)
 
     async def add_endpoint(
-        self, endpoint: apimodelsv1.ProviderEndpoint
+        self, endpoint: apimodelsv1.AddProviderEndpointRequest
     ) -> apimodelsv1.ProviderEndpoint:
         """Add an endpoint."""
         dbend = endpoint.to_db_model()
+        provider_registry = get_provider_registry()
 
         # We override the ID here, as we want to generate it.
         dbend.id = str(uuid4())
 
-        dbendpoint = await self._db_writer.add_provider_endpoint()
+        prov = endpoint.get_from_registry(provider_registry)
+        if prov is None:
+            raise ValueError("Unknown provider type: {}".format(endpoint.provider_type))
+
+        models = []
+        if endpoint.auth_type == apimodelsv1.ProviderAuthType.api_key and not endpoint.api_key:
+            raise ValueError("API key must be provided for API auth type")
+        if endpoint.auth_type != apimodelsv1.ProviderAuthType.passthrough:
+            try:
+                models = prov.models(endpoint.api_key)
+            except Exception as err:
+                raise ValueError("Unable to get models from provider: {}".format(str(err)))
+
+        dbendpoint = await self._db_writer.add_provider_endpoint(dbend)
+
+        await self._db_writer.push_provider_auth_material(
+            dbmodels.ProviderAuthMaterial(
+                provider_endpoint_id=dbendpoint.id,
+                auth_type=endpoint.auth_type,
+                auth_blob=endpoint.api_key if endpoint.api_key else "",
+            )
+        )
+
+        for model in models:
+            await self._db_writer.add_provider_model(
+                dbmodels.ProviderModel(
+                    provider_endpoint_id=dbendpoint.id,
+                    name=model,
+                )
+            )
         return apimodelsv1.ProviderEndpoint.from_db_model(dbendpoint)
 
     async def update_endpoint(
         self, endpoint: apimodelsv1.ProviderEndpoint
     ) -> apimodelsv1.ProviderEndpoint:
         """Update an endpoint."""
 
+        founddbe = await self._db_reader.get_provider_endpoint_by_id(endpoint.id)
+        if founddbe is None:
+            raise ProviderNotFoundError("Provider not found")
+
+        # TODO: We should probably allow this and reinitialize the provider
+        if endpoint.auth_type != founddbe.auth_type:
+            raise ValueError("Cannot change auth type for provider through this endpoint.")
+
         dbendpoint = await self._db_writer.update_provider_endpoint(endpoint.to_db_model())
         return apimodelsv1.ProviderEndpoint.from_db_model(dbendpoint)
 
@@ -175,6 +213,13 @@ async def initialize_provider_endpoints(preg: ProviderRegistry):
             continue
 
         pimpl = provend.get_from_registry(preg)
+        if pimpl is None:
+            logger.warning(
+                "Provider not found in registry",
+                provider=provend.name,
+                endpoint=provend.endpoint,
+            )
+            continue
         await try_initialize_provider_endpoints(provend, pimpl, db_writer)
 
 

src/codegate/providers/llamacpp/provider.py

@@ -1,4 +1,5 @@
 import json
+from typing import List
 
 import httpx
 import structlog
@@ -27,9 +28,15 @@ def __init__(
     def provider_route_name(self) -> str:
         return "llamacpp"
 
-    def models(self):
+    def models(self, api_key: str = None) -> List[str]:
+        headers = {}
+        if api_key:
+            headers["Authorization"] = f"Bearer {api_key}"
         # HACK: This is using OpenAI's /v1/models endpoint to get the list of models
-        resp = httpx.get(f"{self.base_url}/v1/models")
+        resp = httpx.get(
+            f"{self.base_url}/v1/models",
+            headers=headers,
+        )
         jsonresp = resp.json()
 
         return [model["id"] for model in jsonresp.get("data", [])]

src/codegate/providers/ollama/provider.py

@@ -1,4 +1,5 @@
 import json
+from typing import List
 
 import httpx
 import structlog
@@ -34,8 +35,14 @@ def __init__(
     def provider_route_name(self) -> str:
         return "ollama"
 
-    def models(self):
-        resp = httpx.get(f"{self.base_url}/api/tags")
+    def models(self, api_key: str = None) -> List[str]:
+        headers = {}
+        if api_key:
+            headers["Authorization"] = f"Bearer {api_key}"
+        resp = httpx.get(
+            f"{self.base_url}/api/tags",
+            headers=headers,
+        )
         jsonresp = resp.json()
 
         return [model["name"] for model in jsonresp.get("models", [])]

src/codegate/providers/openai/provider.py

@@ -35,9 +35,12 @@ def __init__(
     def provider_route_name(self) -> str:
         return "openai"
 
-    def models(self) -> List[str]:
-        # NOTE: This won't work since we need an API Key being set.
-        resp = httpx.get(f"{self.lm_studio_url}/v1/models")
+    def models(self, api_key: str = None) -> List[str]:
+        headers = {}
+        if api_key:
+            headers["Authorization"] = f"Bearer {api_key}"
+
+        resp = httpx.get(f"{self.lm_studio_url}/v1/models", headers=headers)
         jsonresp = resp.json()
 
         return [model["id"] for model in jsonresp.get("data", [])]

src/codegate/providers/registry.py

@@ -1,9 +1,26 @@
+from threading import Lock
 from typing import Dict, Optional
 
 from fastapi import FastAPI
 
 from codegate.providers.base import BaseProvider
 
+_provider_registry_lock = Lock()
+_provider_registry_singleton: Optional["ProviderRegistry"] = None
+
+
+def get_provider_registry(app: FastAPI = None) -> "ProviderRegistry":
+    global _provider_registry_singleton
+
+    if _provider_registry_singleton is None:
+        if app is None:
+            raise ValueError("Cannot initialize a ProviderRegistry without an app")
+        with _provider_registry_lock:
+            if _provider_registry_singleton is None:
+                _provider_registry_singleton = ProviderRegistry(app)
+
+    return _provider_registry_singleton
+
 
 class ProviderRegistry:
     def __init__(self, app: FastAPI):

src/codegate/providers/vllm/provider.py

@@ -31,8 +31,15 @@ def __init__(
     def provider_route_name(self) -> str:
         return "vllm"
 
-    def models(self):
-        resp = httpx.get(f"{self.base_url}/v1/models")
+    def models(self, api_key: str = None) -> list[str]:
+        headers = {}
+        if api_key:
+            headers["Authorization"] = f"Bearer {api_key}"
+
+        resp = httpx.get(
+            f"{self.base_url}/v1/models",
+            headers=headers,
+        )
         jsonresp = resp.json()
 
         return [model["id"] for model in jsonresp.get("data", [])]