feat(models): decouple healthcheck to other capabilities

Author: leoguillaume

api/clients/model/_albertmodelprovider.py

@@ -1,11 +1,7 @@
 import logging
-from urllib.parse import urljoin
-
-import httpx
 
 from api.schemas.admin.providers import ProviderType
 from api.schemas.core.models import ProviderEndpoints
-from api.utils.variables import EndpointRoute
 
 from ._basemodelprovider import BaseModelProvider
 
@@ -45,24 +41,3 @@ def __init__(
             model_active_params=model_active_params,
         )
         self.type = ProviderType.ALBERT
-
-    async def get_max_context_length(self) -> int | None:
-        url = urljoin(base=str(self.url), url=self.ENDPOINT_TABLE.get_endpoint(endpoint=EndpointRoute.MODELS).lstrip("/"))
-
-        try:
-            async with httpx.AsyncClient() as client:
-                response = await client.get(url=url, headers=self.headers, timeout=self.timeout)
-                response.raise_for_status()
-        except Exception as e:
-            # TODO: remove exc_info=True and return error instead of exception
-            logger.error(f"Error getting max context length for {self.model_name}: {e}", exc_info=True)
-            raise AssertionError(f"Model is not reachable ({e}).")
-
-        data = response.json()["data"]
-        models = [model for model in data if model["id"] == self.model_name or self.model_name in model["aliases"]]
-        assert len(models) == 1, f"Model not found ({self.model_name})."
-
-        model = models[0]
-        max_context_length = model.get("max_context_length")
-
-        return max_context_length

api/clients/model/_basemodelprovider.py

@@ -11,6 +11,7 @@
 import httpx
 from redis.asyncio import Redis as AsyncRedis
 
+from api.infrastructure.fastapi.schemas.models import Models
 from api.schemas.admin.providers import ProviderType
 from api.schemas.audio import AudioTranscription, CreateAudioTranscription
 from api.schemas.chat import ChatCompletionChunk, CreateChatCompletion
@@ -70,27 +71,50 @@ def import_module(type: ProviderType) -> "type[BaseModelProvider]":
 
         return getattr(module, f"{type.capitalize()}ModelProvider")
 
-    @staticmethod
-    async def get_max_context_length() -> int | None:
+    async def healthcheck(self, redis_client: AsyncRedis) -> bool:
+        """
+        Check if the model provider is healthy.
+        """
+        request_content = RequestContent(endpoint=EndpointRoute.MODELS, method="GET")
+        response = await self.forward_request(request_content=request_content, redis_client=redis_client)
+
+        if response.status_code != 200:
+            return False
+
+        data = response.json()["data"]
+        models = [model for model in data if model["id"] == self.model_name or self.model_name in model["aliases"]]
+
+        if not models:
+            return False
+
+        return True
+
+    async def get_max_context_length(self, redis_client: AsyncRedis) -> int | None:
         """
         Get the max context length of the model provider to store in the database. Useful
         to check provider consistency.
         """
-        pass
+        request_content = RequestContent(endpoint=EndpointRoute.MODELS, method="GET")
+        response = await self.forward_request(request_content=request_content, redis_client=redis_client)
+
+        if response.status_code != 200:
+            return None
+
+        data = response.json()
+        return data["max_context_length"]
 
-    async def get_vector_size(self) -> int | None:
+    async def get_vector_size(self, redis_client: AsyncRedis) -> int | None:
         if self.ENDPOINT_TABLE.embeddings is None:
             return None
 
-        url = urljoin(base=self.url, url=self.ENDPOINT_TABLE.embeddings.lstrip("/"))
+        request_content = RequestContent(endpoint=EndpointRoute.EMBEDDINGS, method="POST", json={"model": self.model_name, "input": "hello world"})
+        response = await self.forward_request(request_content=request_content, redis_client=redis_client)
 
-        async with httpx.AsyncClient() as client:
-            response = await client.post(url=url, headers=self.headers, json={"model": self.model_name, "input": "hello world"}, timeout=self.timeout)
-            assert response.status_code == 200, f"Model is not reachable ({response.status_code} - {response.text})."
+        if response.status_code != 200:
+            return None
 
-        data = response.json()["data"]
+        data = response.json()
         vector_size = len(data[0]["embedding"])
-
         return vector_size
 
     def _get_usage(self, request_content: RequestContent, response_data: dict | list[dict], request_latency: float | None = 0.0) -> Usage | None:
@@ -203,6 +227,13 @@ def _format_response(self, request_content: RequestContent, response: httpx.Resp
                         response_data=response_data,
                     ).model_dump()
 
+                elif request_content.endpoint == EndpointRoute.MODELS:
+                    response_data = Models.build_from(
+                        provider_type=self.type,
+                        request_content=request_content,
+                        response_data=response_data,
+                    ).model_dump()
+
                 elif request_content.endpoint == EndpointRoute.RERANK:
                     response_data = Reranks.build_from(
                         provider_type=self.type,

api/clients/model/_mistralmodelprovider.py

@@ -1,11 +1,7 @@
 import logging
-from urllib.parse import urljoin
-
-import httpx
 
 from api.schemas.admin.providers import ProviderType
 from api.schemas.core.models import ProviderEndpoints
-from api.utils.variables import EndpointRoute
 
 from ._basemodelprovider import BaseModelProvider
 
@@ -45,24 +41,3 @@ def __init__(
             timeout=timeout,
         )
         self.type = ProviderType.MISTRAL
-
-    async def get_max_context_length(self) -> int | None:
-        url = urljoin(base=str(self.url), url=self.ENDPOINT_TABLE.get_endpoint(endpoint=EndpointRoute.MODELS).lstrip("/"))
-
-        try:
-            async with httpx.AsyncClient() as client:
-                response = await client.get(url=url, headers=self.headers, timeout=self.timeout)
-                response.raise_for_status()
-
-        except Exception as e:
-            logger.error(f"Error getting max context length for {self.model_name}: {e}", exc_info=True)
-            raise AssertionError(f"Model is not reachable ({e}).")
-
-        data = response.json()["data"]
-        models = [model for model in data if model["id"] == self.model_name]
-        assert len(models) == 1, f"Model not found ({self.model_name})."
-
-        model = models[0]
-        max_context_length = model.get("max_context_length")
-
-        return max_context_length

api/clients/model/_openaimodelprovider.py

@@ -1,11 +1,7 @@
 import logging
-from urllib.parse import urljoin
-
-import httpx
 
 from api.schemas.admin.providers import ProviderType
 from api.schemas.core.models import ProviderEndpoints
-from api.utils.variables import EndpointRoute
 
 from ._basemodelprovider import BaseModelProvider
 
@@ -45,23 +41,3 @@ def __init__(
             timeout=timeout,
         )
         self.type = ProviderType.OPENAI
-
-    async def get_max_context_length(self) -> int | None:
-        url = urljoin(base=str(self.url), url=self.ENDPOINT_TABLE.get_endpoint(endpoint=EndpointRoute.MODELS).lstrip("/"))
-
-        try:
-            async with httpx.AsyncClient() as client:
-                response = await client.get(url=url, headers=self.headers, timeout=self.timeout)
-                response.raise_for_status()
-        except Exception as e:
-            logger.error(f"Error getting max context length for {self.model_name}: {e}", exc_info=True)
-            raise AssertionError(f"Model is not reachable ({e}).")
-
-        data = response.json()["data"]
-        models = [model for model in data if model["id"] == self.model_name]
-        assert len(models) == 1, f"Model not found ({self.model_name})."
-
-        model = models[0]
-        max_context_length = model.get("max_context_length")
-
-        return max_context_length

api/clients/model/_teimodelprovider.py

@@ -1,11 +1,7 @@
 import logging
-from urllib.parse import urljoin
-
-import httpx
 
 from api.schemas.admin.providers import ProviderType
 from api.schemas.core.models import ProviderEndpoints
-from api.utils.variables import EndpointRoute
 
 from ._basemodelprovider import BaseModelProvider
 
@@ -45,19 +41,3 @@ def __init__(
             model_active_params=model_active_params,
         )
         self.type = ProviderType.TEI
-
-    async def get_max_context_length(self) -> int | None:
-        url = urljoin(base=self.url, url=self.ENDPOINT_TABLE.get_endpoint(endpoint=EndpointRoute.MODELS).lstrip("/"))
-
-        try:
-            async with httpx.AsyncClient() as client:
-                response = await client.get(url=url, headers=self.headers, timeout=self.timeout)
-                response.raise_for_status()
-        except Exception as e:
-            logger.error(f"Error getting max context length for {self.model_name}: {e}", exc_info=True)
-            raise AssertionError(f"Model is not reachable ({e}).")
-
-        data = response.json()
-        assert self.model_name == data["model_id"], f"Model not found ({self.model_name})."
-        max_context_length = data.get("max_input_length")
-        return max_context_length

api/clients/model/_vllmmodelprovider.py

@@ -1,11 +1,7 @@
 import logging
-from urllib.parse import urljoin
-
-import httpx
 
 from api.schemas.admin.providers import ProviderType
 from api.schemas.core.models import ProviderEndpoints
-from api.utils.variables import EndpointRoute
 
 from ._basemodelprovider import BaseModelProvider
 
@@ -45,23 +41,3 @@ def __init__(
             model_active_params=model_active_params,
         )
         self.type = ProviderType.VLLM
-
-    async def get_max_context_length(self) -> int | None:
-        url = urljoin(base=self.url, url=self.ENDPOINT_TABLE.get_endpoint(endpoint=EndpointRoute.MODELS).lstrip("/"))
-
-        try:
-            async with httpx.AsyncClient() as client:
-                response = await client.get(url=url, headers=self.headers, timeout=self.timeout)
-                response.raise_for_status()
-        except Exception as e:
-            logger.error(f"Error getting max context length for {self.model_name}: {e}", exc_info=True)
-            raise AssertionError(f"Model is not reachable ({e}).")
-
-        data = response.json()
-        models = [model for model in data["data"] if model["id"] == self.model_name]
-        assert len(models) == 1, f"Model not found ({self.model_name})."
-
-        model = models[0]
-        max_context_length = model.get("max_model_len")
-
-        return max_context_length

api/dependencies.py

@@ -2,6 +2,7 @@
 from contextvars import ContextVar
 
 from fastapi import Depends
+from redis.asyncio import Redis as AsyncRedis
 from sqlalchemy.ext.asyncio import AsyncSession
 
 from api.domain.key import KeyRepository
@@ -28,6 +29,14 @@ async def get_postgres_session() -> AsyncGenerator[AsyncSession]:
             raise
 
 
+async def get_redis_client() -> AsyncGenerator[AsyncRedis]:
+    client = AsyncRedis(connection_pool=global_context.redis_pool)
+
+    yield client
+
+    await client.aclose()
+
+
 def get_request_context() -> ContextVar[RequestContext]:
     return request_context
 
@@ -53,11 +62,12 @@ def get_models_use_case(
 
 def create_provider_use_case_factory(
     postgres_session: AsyncSession = Depends(get_postgres_session),
+    redis_client: AsyncRedis = Depends(get_redis_client),
 ) -> CreateProviderUseCase:
     return CreateProviderUseCase(
         router_repository=_router_repository(postgres_session),
         provider_repository=PostgresProviderRepository(postgres_session=postgres_session),
-        provider_gateway=ModelProviderGateway(),
+        provider_gateway=ModelProviderGateway(redis_client=redis_client),
         user_info_repository=_user_info_repository(postgres_session),
     )
 

api/domain/provider/entities.py

@@ -1,4 +1,4 @@
-from enum import Enum
+from enum import Enum, StrEnum
 from typing import Literal
 
 import pycountry
@@ -14,7 +14,7 @@
 ProviderCarbonFootprintZone: type[Enum] = Enum("ProviderCarbonFootprintZone", country_codes_dict, type=str)
 
 
-class ProviderType(str, Enum):
+class ProviderType(StrEnum):
     ALBERT = "albert"
     OPENAI = "openai"
     MISTRAL = "mistral"

api/helpers/models/_modelregistry.py

@@ -492,9 +492,9 @@ async def create_provider(
                 model_total_params=model_total_params,
                 model_active_params=model_active_params,
             )
-            max_context_length = await provider.get_max_context_length()
+            max_context_length = await provider.get_max_context_length(redis_client=self.redis_client)
             if router.type == ModelType.TEXT_EMBEDDINGS_INFERENCE:
-                vector_size = await provider.get_vector_size()
+                vector_size = await provider.get_vector_size(redis_client=self.redis_client)
             else:
                 vector_size = None
 

api/infrastructure/fastapi/schemas/models.py

@@ -1,18 +1,19 @@
-from enum import Enum
-from typing import Literal
+from enum import StrEnum
+from typing import Annotated, Literal
 
 from pydantic import Field
 
-from api.domain.model import Model as ModelEntity
+from api.domain.provider.entities import ProviderType
 from api.schemas import BaseModel
+from api.schemas.core.models import RequestContent
 
 
 class ModelCosts(BaseModel):
     prompt_tokens: float = Field(default=0.0, ge=0.0, description="Cost of a million prompt tokens (decrease user budget)")
     completion_tokens: float = Field(default=0.0, ge=0.0, description="Cost of a million completion tokens (decrease user budget)")
 
 
-class ModelType(str, Enum):
+class ModelType(StrEnum):
     AUTOMATIC_SPEECH_RECOGNITION = "automatic-speech-recognition"
     IMAGE_TEXT_TO_TEXT = "image-text-to-text"
     IMAGE_TO_TEXT = "image-to-text"
@@ -21,10 +22,38 @@ class ModelType(str, Enum):
     TEXT_CLASSIFICATION = "text-classification"
 
 
-class Model(ModelEntity):
-    object: Literal["model"] = "model"
+class Model(BaseModel):
+    object: Annotated[Literal["model"], Field("model", description="Type of the object.")]
+    id: Annotated[str, Field(..., description="The model identifier, which can be referenced in the API endpoints.")]
+    type: Annotated[ModelType | None, Field(default=None, description="The type of the model, which can be used to identify the model type.", examples=["text-generation"])]  # fmt: off
+    aliases: Annotated[list[str], Field(default_factory=list, description="Aliases of the model. It will be used to identify the model by users.", examples=[["model-alias", "model-alias-2"]])]  # fmt: off
+    created: Annotated[int, Field(..., description="Time of creation, as Unix timestamp.")]
+    owned_by: Annotated[str, Field(..., description="The organization that owns the model.")]
+    max_context_length: Annotated[int | None, Field(default=None, description="Maximum amount of tokens a context could contains. Makes sure it is the same for all models.")]  # fmt: off
+    costs: Annotated[ModelCosts, Field(default_factory=ModelCosts, description="Costs of the model.")]
+
+    @classmethod
+    def build_from(cls, provider_type: ProviderType, request_content: RequestContent, response_data: dict) -> "Model":
+        match provider_type:
+            case ProviderType.ALBERT:
+                return cls(**response_data)
+
+            case ProviderType.TEI:
+                return cls(id=response_data["model_id"], created=0, owned_by="tei", max_context_length=response_data["max_input_length"])
+
+            case ProviderType.MISTRAL:
+                return cls(**response_data)
+
+            case ProviderType.OPENAI:
+                return cls(**response_data)
+
+            case ProviderType.VLLM:
+                return cls(max_context_length=response_data["max_model_len"], **response_data)
+
+            case _:
+                raise NotImplementedError(f"Provider {provider_type} not implemented")
 
 
 class Models(BaseModel):
-    object: Literal["list"] = "list"
-    data: list[Model]
+    object: Annotated[Literal["list"], Field("list", description="Type of the object.")]
+    data: Annotated[list[Model], Field(..., description="List of models.")]