diff --git a/api/clients/model/__init__.py b/api/clients/model/__init__.py deleted file mode 100644 index 2f491236f..000000000 --- a/api/clients/model/__init__.py +++ /dev/null @@ -1,8 +0,0 @@ -from ._albertmodelprovider import AlbertModelProvider -from ._basemodelprovider import BaseModelProvider -from ._mistralmodelprovider import MistralModelProvider -from ._openaimodelprovider import OpenaiModelProvider -from ._teimodelprovider import TeiModelProvider -from ._vllmmodelprovider import VllmModelProvider - -__all__ = [AlbertModelProvider, BaseModelProvider, MistralModelProvider, OpenaiModelProvider, TeiModelProvider, VllmModelProvider] diff --git a/api/clients/model/_albertmodelprovider.py b/api/clients/model/_albertmodelprovider.py deleted file mode 100644 index 27824082d..000000000 --- a/api/clients/model/_albertmodelprovider.py +++ /dev/null @@ -1,68 +0,0 @@ -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 - -logger = logging.getLogger(__name__) - - -class AlbertModelProvider(BaseModelProvider): - ENDPOINT_TABLE = ProviderEndpoints( - audio_transcriptions="/v1/audio/transcriptions", - chat_completions="/v1/chat/completions", - embeddings="/v1/embeddings", - models="/v1/models", - ocr="/v1/ocr", - rerank="/v1/rerank", - ) - - def __init__( - self, - url: str, - key: str, - timeout: int, - model_name: str, - model_hosting_zone: str | None, - model_total_params: int | None, - model_active_params: int | None, - ) -> None: - """ - Initialize the Albert model provider and check if the model is available. - """ - super().__init__( - url=url, - key=key, - timeout=timeout, - model_name=model_name, - model_hosting_zone=model_hosting_zone, - model_total_params=model_total_params, - 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 diff --git a/api/clients/model/_mistralmodelprovider.py b/api/clients/model/_mistralmodelprovider.py deleted file mode 100644 index c7c688f58..000000000 --- a/api/clients/model/_mistralmodelprovider.py +++ /dev/null @@ -1,68 +0,0 @@ -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 - -logger = logging.getLogger(__name__) - - -class MistralModelProvider(BaseModelProvider): - ENDPOINT_TABLE = ProviderEndpoints( - audio_transcriptions="/v1/chat/completions", - chat_completions="/v1/chat/completions", - embeddings="/v1/embeddings", - models="/v1/models", - ocr="/v1/ocr", - rerank=None, - ) - - def __init__( - self, - url: str, - key: str, - timeout: int, - model_name: str, - model_hosting_zone: str | None, - model_total_params: int | None, - model_active_params: int | None, - ) -> None: - """ - Initialize the Mistral model provider and check if the model is available. - """ - super().__init__( - model_name=model_name, - model_hosting_zone=model_hosting_zone, - model_total_params=model_total_params, - model_active_params=model_active_params, - url=url, - key=key, - 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 diff --git a/api/clients/model/_openaimodelprovider.py b/api/clients/model/_openaimodelprovider.py deleted file mode 100644 index eb4080bca..000000000 --- a/api/clients/model/_openaimodelprovider.py +++ /dev/null @@ -1,67 +0,0 @@ -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 - -logger = logging.getLogger(__name__) - - -class OpenaiModelProvider(BaseModelProvider): - ENDPOINT_TABLE = ProviderEndpoints( - audio_transcriptions="/v1/audio/transcriptions", - chat_completions="/v1/chat/completions", - embeddings="/v1/embeddings", - models="/v1/models", - ocr="/v1/chat/completions", - rerank=None, - ) - - def __init__( - self, - url: str, - key: str, - timeout: int, - model_name: str, - model_hosting_zone: str | None, - model_total_params: int | None, - model_active_params: int | None, - ) -> None: - """ - Initialize the OpenAI model provider and check if the model is available. - """ - super().__init__( - model_name=model_name, - model_hosting_zone=model_hosting_zone, - model_total_params=model_total_params, - model_active_params=model_active_params, - url=url, - key=key, - 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 diff --git a/api/clients/model/_teimodelprovider.py b/api/clients/model/_teimodelprovider.py deleted file mode 100644 index de2669266..000000000 --- a/api/clients/model/_teimodelprovider.py +++ /dev/null @@ -1,63 +0,0 @@ -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 - -logger = logging.getLogger(__name__) - - -class TeiModelProvider(BaseModelProvider): - ENDPOINT_TABLE = ProviderEndpoints( - audio_transcriptions=None, - chat_completions=None, - embeddings="/v1/embeddings", - models="/info", - ocr=None, - rerank="/rerank", - ) - - def __init__( - self, - url: str, - key: str, - timeout: int, - model_name: str, - model_hosting_zone: str | None, - model_total_params: int | None, - model_active_params: int | None, - ) -> None: - """ - Initialize the TEI model client and check if the model is available. - """ - super().__init__( - url=url, - key=key, - timeout=timeout, - model_name=model_name, - model_hosting_zone=model_hosting_zone, - model_total_params=model_total_params, - 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 diff --git a/api/clients/model/_vllmmodelprovider.py b/api/clients/model/_vllmmodelprovider.py deleted file mode 100644 index ce4e046d5..000000000 --- a/api/clients/model/_vllmmodelprovider.py +++ /dev/null @@ -1,67 +0,0 @@ -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 - -logger = logging.getLogger(__name__) - - -class VllmModelProvider(BaseModelProvider): - ENDPOINT_TABLE = ProviderEndpoints( - audio_transcriptions="/v1/audio/transcriptions", - chat_completions="/v1/chat/completions", - embeddings="/v1/embeddings", - models="/v1/models", - ocr="/v1/chat/completions", - rerank="/v2/rerank", - ) - - def __init__( - self, - url: str, - key: str, - timeout: int, - model_name: str, - model_hosting_zone: str | None, - model_total_params: int | None, - model_active_params: int | None, - ) -> None: - """ - Initialize the vLLM model provider and check if the model is available. - """ - super().__init__( - url=url, - key=key, - timeout=timeout, - model_name=model_name, - model_hosting_zone=model_hosting_zone, - model_total_params=model_total_params, - 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 diff --git a/api/dependencies.py b/api/dependencies.py index 529c1df0a..6a5e0b65b 100644 --- a/api/dependencies.py +++ b/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 @@ -43,6 +44,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() + + # repositories def _router_repository(session: AsyncSession) -> PostgresRouterRepository: return PostgresRouterRepository(postgres_session=session, app_title=configuration.settings.app_title) diff --git a/api/endpoints/chat.py b/api/endpoints/chat.py index f8906f480..5ca2c033f 100644 --- a/api/endpoints/chat.py +++ b/api/endpoints/chat.py @@ -58,6 +58,7 @@ async def chat_completions( request_context: ContextVar[RequestContext] = Depends(get_request_context), ) -> JSONResponse | StreamingResponseWithStatusCode: """Creates a model response for the given chat conversation.""" + model_provider = await model_registry.get_model_provider( model=body.model, endpoint=EndpointRoute.CHAT_COMPLETIONS, diff --git a/api/helpers/_documentmanager.py b/api/helpers/_documentmanager.py index 64b76a10d..6e91a8a4e 100755 --- a/api/helpers/_documentmanager.py +++ b/api/helpers/_documentmanager.py @@ -12,7 +12,6 @@ from sqlalchemy.exc import NoResultFound from sqlalchemy.ext.asyncio import AsyncSession -from api.clients.model import BaseModelProvider as ModelProvider from api.helpers._elasticsearchvectorstore import ElasticsearchVectorStore from api.helpers.models import ModelRegistry from api.schemas.chunks import Chunk, ChunkMetadata, InputChunk @@ -35,6 +34,7 @@ ) from api.utils.variables import EndpointRoute +from ..infrastructure.http.model import ModelHttpClient from ._parsermanager import ParserManager logger = logging.getLogger(__name__) @@ -552,7 +552,7 @@ def _split( return chunks - async def _create_embeddings(self, provider: ModelProvider, input_texts: list[str], redis_client: AsyncRedis) -> list[float]: + async def _create_embeddings(self, provider: ModelHttpClient, input_texts: list[str], redis_client: AsyncRedis) -> list[float]: response = await provider.forward_request( request_content=RequestContent( method="POST", diff --git a/api/helpers/models/_modelregistry.py b/api/helpers/models/_modelregistry.py index a61499b1d..341c8e1cb 100644 --- a/api/helpers/models/_modelregistry.py +++ b/api/helpers/models/_modelregistry.py @@ -3,12 +3,20 @@ from typing import Literal from redis.asyncio import Redis as AsyncRedis -from sqlalchemy import Integer, and_, cast, delete, func, insert, or_, select, text, update +from sqlalchemy import Integer, and_, cast, delete, func, insert, or_, select, text from sqlalchemy.exc import IntegrityError, NoResultFound from sqlalchemy.ext.asyncio import AsyncSession -from api.clients.model import BaseModelProvider as ModelProvider -from api.schemas.admin.providers import Provider, ProviderCarbonFootprintZone, ProviderType +from api.domain.provider import Provider +from api.infrastructure.http.model import ( + AlbertModelHttpClient, + MistralModelHttpClient, + ModelHttpClient, + OpenaiModelHttpClient, + TeiModelHttpClient, + VllmModelHttpClient, +) +from api.schemas.admin.providers import ProviderCarbonFootprintZone, ProviderType from api.schemas.admin.routers import Router, RouterLoadBalancingStrategy from api.schemas.core.configuration import Model as ModelConfiguration from api.schemas.core.context import RequestContext @@ -121,7 +129,7 @@ async def setup(self, models: list[ModelConfiguration], postgres_session: AsyncS ) logger.info(f"Router {model.name} are created (id: {router_id})") except RouterAlreadyExistsException: - continue + pass except RouterAliasAlreadyExistsException: continue except Exception as e: @@ -137,7 +145,7 @@ async def setup(self, models: list[ModelConfiguration], postgres_session: AsyncS provider_id = await self.create_provider( router_id=router.id, user_id=0, # setup as master user - type=provider.type, + provider_type=provider.type, url=provider.url, key=provider.key, timeout=provider.timeout, @@ -235,91 +243,6 @@ async def create_router( return router_id - @staticmethod - async def delete_router(router_id: int, postgres_session: AsyncSession) -> None: - """ - Delete a model router and all its providers. - - Args: - router_id(int): The router ID - postgres_session(AsyncSession): Database postgres_session - """ - # Check if router exists - query = select(RouterTable).where(RouterTable.id == router_id) - result = await postgres_session.execute(query) - try: - result.scalar_one() - except NoResultFound: - raise RouterNotFoundException() - - # Delete will cascade to providers and aliases due to foreign key constraints - await postgres_session.execute(delete(RouterTable).where(RouterTable.id == router_id)) - await postgres_session.commit() - - # TODO: delete queue - - async def update_router( - self, - router_id: int, - name: str | None, - type: ModelType | None, - aliases: list[str] | None, - load_balancing_strategy: RouterLoadBalancingStrategy | None, - cost_prompt_tokens: float | None, - cost_completion_tokens: float | None, - postgres_session: AsyncSession, - ) -> None: - """ - Update a model router. - - Args: - router_id(int): The router ID - name(Optional[str]): Optional new name - type(Optional[ModelType]): Optional new type - aliases(Optional[List[str]]): Optional new aliases list (replaces existing) - load_balancing_strategy(Optional[RouterLoadBalancingStrategy]): Optional new routing strategy - cost_prompt_tokens(Optional[float]): Optional new cost of a million prompt tokens - cost_completion_tokens(Optional[float]): Optional new cost of a million completion tokens - postgres_session(AsyncSession): Database postgres_session - - """ - # Check if model exists - routers = await self.get_routers(router_id=router_id, name=None, postgres_session=postgres_session) - - # Check alias integrity in aliases of other routers - if aliases: - query = select(RouterAliasTable.router_id).where(RouterAliasTable.value.in_(aliases)).where(RouterAliasTable.router_id != router_id) - result = await postgres_session.execute(query) - conflicting_aliases = result.scalars().all() - if conflicting_aliases: - raise RouterAliasAlreadyExistsException() - - # Update router properties - update_values = {} - if type is not None: - update_values["type"] = type.value - if load_balancing_strategy is not None: - update_values["load_balancing_strategy"] = load_balancing_strategy.value - if name is not None: - update_values["name"] = name - if cost_prompt_tokens is not None: - update_values["cost_prompt_tokens"] = cost_prompt_tokens - if cost_completion_tokens is not None: - update_values["cost_completion_tokens"] = cost_completion_tokens - - if update_values: - await postgres_session.execute(update(RouterTable).where(RouterTable.id == router_id).values(**update_values)) - - # Update aliases if provided - if aliases is not None: - query = delete(RouterAliasTable).where(RouterAliasTable.router_id == router_id) - await postgres_session.execute(query) - if aliases: - query = insert(RouterAliasTable).values([{"value": alias, "router_id": router_id} for alias in aliases]) - await postgres_session.execute(query) - - await postgres_session.commit() - @staticmethod async def get_routers( router_id: int | None, @@ -439,7 +362,7 @@ async def create_provider( self, router_id: int, user_id: int, - type: ProviderType, + provider_type: ProviderType, url: str, key: str | None, timeout: int, @@ -457,7 +380,7 @@ async def create_provider( Args: router_id(int): The model router ID user_id(int): The user ID of owner of the provider - type(ProviderType): provider type + provider_type(ProviderType): provider type url(str): provider URL key(str | None): provider API key timeout(int): Request timeout @@ -475,21 +398,32 @@ async def create_provider( routers = await self.get_routers(router_id=router_id, name=None, postgres_session=postgres_session) router = routers[0] - if type.value not in self.MODEL_TYPE_TO_MODEL_PROVIDER_TYPE_MAPPING[router.type]: + if provider_type.value not in self.MODEL_TYPE_TO_MODEL_PROVIDER_TYPE_MAPPING[router.type]: raise InvalidProviderTypeException() # call model to get the vector size, max context length try: - provider = ModelProvider.import_module(type=type)( + provider_class: dict[ProviderType, type[ModelHttpClient]] = { + ProviderType.ALBERT: AlbertModelHttpClient, + ProviderType.MISTRAL: MistralModelHttpClient, + ProviderType.OPENAI: OpenaiModelHttpClient, + ProviderType.TEI: TeiModelHttpClient, + ProviderType.VLLM: VllmModelHttpClient, + } + + provider = provider_class[provider_type]( url=url, key=key, timeout=timeout, model_name=model_name, - model_hosting_zone=model_hosting_zone, - model_total_params=model_total_params, model_active_params=model_active_params, + model_total_params=model_total_params, + model_hosting_zone=model_hosting_zone, ) - max_context_length = await provider.get_max_context_length() + model_info = await provider.get_model_info() + if model_info is None: + raise Exception(f"Provider {provider_type.value} not reachable") + max_context_length = model_info.max_context_length if router.type == ModelType.TEXT_EMBEDDINGS_INFERENCE: vector_size = await provider.get_vector_size() else: @@ -515,7 +449,7 @@ async def create_provider( .values( router_id=router_id, user_id=user_id, - type=type.value, + type=provider_type.value, url=url, key=key, timeout=timeout, @@ -650,80 +584,6 @@ async def get_providers( return providers - async def update_provider( - self, - provider_id: int, - router_id: int | None, - timeout: int | None, - model_hosting_zone: ProviderCarbonFootprintZone | None, - model_total_params: int | None, - model_active_params: int | None, - qos_metric: Metric | None, - qos_limit: float | None, - postgres_session: AsyncSession, - ) -> None: - """ - Update a model provider by ID. - - Args: - provider_id(int): The provider ID - router_id(int | None): The new router ID to assign to the provider - timeout(int | None): The new timeout for the provider requests - model_hosting_zone(ProviderCarbonFootprintZone | None): The new model carbon footprint zone - model_total_params(int | None): The new model carbon footprint total params - model_active_params(int | None): The new model carbon footprint active params - qos_metric(Metric | None): The new QoS metric - qos_limit(float | None): The new QoS limit - postgres_session(AsyncSession): Database postgres_session - """ - # Check if provider exists - providers = await self.get_providers(provider_id=provider_id, router_id=None, postgres_session=postgres_session) - provider = providers[0] - - routers = await self.get_routers(router_id=provider.router_id, name=None, postgres_session=postgres_session) - current_router = routers[0] - - # Update provider - update_value = {} - if router_id is not None: - # check if new router exists - routers = await self.get_routers(router_id=router_id, name=None, postgres_session=postgres_session) - new_router = routers[0] - - # provider.type is a ProviderType enum, mapping contains provider type strings - if provider.type.value not in self.MODEL_TYPE_TO_MODEL_PROVIDER_TYPE_MAPPING[new_router.type]: - raise InvalidProviderTypeException("New router type is not compatible with the provider type.") - - # consistency check - if new_router.providers > 0: - if new_router.vector_size != current_router.vector_size: - raise InconsistentModelVectorSizeException() - if new_router.max_context_length != current_router.max_context_length: - raise InconsistentModelMaxContextLengthException() - - update_value["router_id"] = router_id - - if timeout is not None: - update_value["timeout"] = timeout - if model_hosting_zone is not None: - update_value["model_hosting_zone"] = model_hosting_zone - if model_total_params is not None: - update_value["model_total_params"] = model_total_params - if model_active_params is not None: - update_value["model_active_params"] = model_active_params - if qos_metric is not None: - update_value["qos_metric"] = qos_metric - if qos_limit is not None: - update_value["qos_limit"] = qos_limit - - if update_value: - try: - await postgres_session.execute(update(ProviderTable).where(ProviderTable.id == provider_id).values(**update_value)) - await postgres_session.commit() - except IntegrityError: - await postgres_session.rollback() - raise ProviderAlreadyExistsException("provider already exists for the new router.") - async def get_models(self, name: str | None, user_info: UserInfo, postgres_session: AsyncSession) -> list[Model]: """ Get models for a user. @@ -808,7 +668,7 @@ async def get_model_provider( postgres_session: AsyncSession, redis_client: AsyncRedis, request_context: ContextVar[RequestContext], - ) -> ModelProvider: + ) -> ModelHttpClient: """ Get a model provider for a given model, endpoint, user priority, postgres_session and redis client. @@ -819,7 +679,7 @@ async def get_model_provider( redis_client(AsyncRedis): Redis client request_context(ContextVar[RequestContext]): Request context Returns: - ModelProvider: The chosen provider + ModelHttpClient: The chosen provider """ try: routers = await self.get_routers(router_id=None, name=model, postgres_session=postgres_session) @@ -867,7 +727,15 @@ async def get_model_provider( providers = await self.get_providers(router_id=router.id, provider_id=provider_id, postgres_session=postgres_session) provider = providers[0] - model_provider = ModelProvider.import_module(type=provider.type)( + provider_class: dict[ProviderType, type[ModelHttpClient]] = { + ProviderType.ALBERT: AlbertModelHttpClient, + ProviderType.MISTRAL: MistralModelHttpClient, + ProviderType.OPENAI: OpenaiModelHttpClient, + ProviderType.TEI: TeiModelHttpClient, + ProviderType.VLLM: VllmModelHttpClient, + } + + model_provider = provider_class[provider.type]( url=provider.url, key=provider.key, timeout=provider.timeout, diff --git a/api/infrastructure/http/__init__.py b/api/infrastructure/http/__init__.py new file mode 100644 index 000000000..e69de29bb diff --git a/api/infrastructure/http/model/__init__.py b/api/infrastructure/http/model/__init__.py new file mode 100644 index 000000000..d129c5ffe --- /dev/null +++ b/api/infrastructure/http/model/__init__.py @@ -0,0 +1,16 @@ +from ._albertmodelhttpclient import AlbertModelHttpClient +from ._mistralmodelhttpclient import MistralModelHttpClient +from ._modelhttpclient import ModelHttpClient +from ._openaimodelhttpclient import OpenaiModelHttpClient +from ._teimodelhttpclient import TeiCreateRerankBody, TeiModelHttpClient +from ._vllmmodelhttpclient import VllmModelHttpClient + +__all__ = [ + "AlbertModelHttpClient", + "MistralModelHttpClient", + "TeiCreateRerankBody", + "TeiModelHttpClient", + "VllmModelHttpClient", + "OpenaiModelHttpClient", + "ModelHttpClient", +] diff --git a/api/infrastructure/http/model/_albertmodelhttpclient.py b/api/infrastructure/http/model/_albertmodelhttpclient.py new file mode 100644 index 000000000..606677ef4 --- /dev/null +++ b/api/infrastructure/http/model/_albertmodelhttpclient.py @@ -0,0 +1,16 @@ +from api.infrastructure.http.model._modelhttpclient import ModelHttpClient, ModelHttpClientEndpoints +from api.schemas.audio import AudioTranscriptionResponseFormat +from api.schemas.core.models import RequestContent + + +class AlbertModelHttpClient(ModelHttpClient): + ENDPOINT_TABLE = ModelHttpClientEndpoints() + + # request formatting + @staticmethod + def format_audio_transcription_request(request_content: RequestContent) -> RequestContent: + if request_content.form["response_format"] == AudioTranscriptionResponseFormat.TEXT: + # @TODO: remove this once the Albert API supports the text response format + request_content.form["response_format"] = AudioTranscriptionResponseFormat.JSON.value + + return request_content diff --git a/api/infrastructure/http/model/_mistralmodelhttpclient.py b/api/infrastructure/http/model/_mistralmodelhttpclient.py new file mode 100644 index 000000000..c133e3df7 --- /dev/null +++ b/api/infrastructure/http/model/_mistralmodelhttpclient.py @@ -0,0 +1,83 @@ +import base64 + +from mistralai.client.models import AudioChunk, ChatCompletionRequest, TextChunk, UserMessage + +from api.infrastructure.fastapi.schemas.models import ModelResponse, ModelsResponse +from api.infrastructure.http.model._modelhttpclient import ModelHttpClient, ModelHttpClientEndpoints +from api.schemas.audio import AudioTranscription +from api.schemas.core.models import RequestContent + + +class MistralModelHttpClient(ModelHttpClient): + ENDPOINT_TABLE = ModelHttpClientEndpoints(audio_transcriptions="/v1/chat/completions", rerank=None) + + # request formatting + @staticmethod + def format_chat_completion_request(request_content: RequestContent) -> RequestContent: + # see https://docs.mistral.ai/api#operation-chat_completion_v1_chat_completions_post + request_content.body = { + "frequency_penalty": request_content.body.get("frequency_penalty") or 0.0, + "max_tokens": request_content.body.get("max_tokens"), + "messages": request_content.body.get("messages"), + "model": request_content.body.get("model"), + "n": request_content.body.get("n"), + "parallel_tool_calls": request_content.body.get("parallel_tool_calls") or False, + "prediction": request_content.body.get("prediction") or {}, + "presence_penalty": request_content.body.get("presence_penalty") or 0.0, + "prompt_mode": request_content.body.get("prompt_mode"), + "random_seed": request_content.body.get("random_seed") or request_content.body.get("seed"), + "response_format": request_content.body.get("response_format") or {"type": "text"}, + "safe_prompt": request_content.body.get("safe_prompt") or False, + "stop": request_content.body.get("stop") or [], + "stream": request_content.body.get("stream") or False, + "temperature": request_content.body.get("temperature"), + "tool_choice": request_content.body.get("tool_choice"), + "tools": request_content.body.get("tools"), + "top_p": request_content.body.get("top_p") or 1.0, + } + + return request_content + + @staticmethod + def format_audio_transcription_request(request_content: RequestContent) -> RequestContent: + text = request_content.form.get("prompt") or f"Transcribe this audio in this language : {request_content.form.get('language', 'en')}" + input_audio = base64.b64encode(request_content.files["file"][1]).decode("utf-8") + request_content.body = ChatCompletionRequest( + model=request_content.form["model"], + messages=[ + UserMessage( + role="user", + content=[AudioChunk(type="input_audio", input_audio=input_audio), TextChunk(type="text", text=text)], + ) + ], + temperature=request_content.form.get("temperature"), + ).model_dump() + request_content.files = {} + request_content.form = {} + + return request_content + + # response formatting + @staticmethod + def format_response_to_models_response(request_content: RequestContent, response_data: dict) -> ModelsResponse: + return ModelsResponse( + data=[ + ModelResponse( + id=model.get("id"), + created=model.get("created"), + owned_by=model.get("owned_by"), + max_context_length=model.get("max_context_length"), + aliases=model.get("aliases", []), + ) + for model in response_data.get("data", []) + ] + ) + + @staticmethod + def format_response_to_audio_transcription_response(request_content: RequestContent, response_data: dict) -> AudioTranscription: + return AudioTranscription( + id=response_data["id"], + model=response_data["model"], + text=response_data["choices"][0]["message"]["content"], + usage=response_data["usage"], + ) diff --git a/api/clients/model/_basemodelprovider.py b/api/infrastructure/http/model/_modelhttpclient.py similarity index 58% rename from api/clients/model/_basemodelprovider.py rename to api/infrastructure/http/model/_modelhttpclient.py index 7b7762eea..34109c97c 100644 --- a/api/clients/model/_basemodelprovider.py +++ b/api/infrastructure/http/model/_modelhttpclient.py @@ -1,33 +1,63 @@ -from abc import ABC import ast -import importlib from json import JSONDecodeError, dumps, loads import logging import time import traceback +from typing import Annotated from urllib.parse import urljoin from fastapi import HTTPException import httpx +from pydantic import BaseModel, Field, StringConstraints from redis.asyncio import Redis as AsyncRedis -from api.schemas.admin.providers import ProviderType -from api.schemas.audio import AudioTranscription, CreateAudioTranscription -from api.schemas.chat import ChatCompletionChunk, CreateChatCompletion -from api.schemas.core.models import Metric, ProviderEndpoints, RequestContent -from api.schemas.rerank import CreateRerank, Reranks +from api.domain.provider.entities import ProviderCarbonFootprintZone +from api.infrastructure.fastapi.schemas.models import ModelResponse, ModelsResponse +from api.schemas.audio import AudioTranscription +from api.schemas.chat import ChatCompletionChunk +from api.schemas.core.models import Metric, RequestContent +from api.schemas.rerank import Reranks from api.schemas.usage import Usage from api.utils.carbon import get_carbon_footprint from api.utils.context import generate_request_id, global_context, request_context -from api.utils.exceptions import ModelIsTooBusyException, RequestFormatFailedException, ResponseFormatFailedException +from api.utils.exceptions import ModelIsTooBusyException from api.utils.redis import redis_retry, safe_redis_reset from api.utils.variables import PREFIX__REDIS_METRIC_GAUGE, PREFIX__REDIS_METRIC_TIMESERIE, REDIS__TIMESERIE_RETENTION_SECONDS, EndpointRoute logger = logging.getLogger(__name__) -class BaseModelProvider(ABC): - ENDPOINT_TABLE: ProviderEndpoints = ProviderEndpoints() +class ModelHttpClientEndpoints(BaseModel): + audio_transcriptions: Annotated[str | None, StringConstraints(strip_whitespace=True, min_length=1, pattern=r"^/", to_lower=True), Field(default="/v1/audio/transcriptions")] # fmt: off + chat_completions: Annotated[str | None, StringConstraints(strip_whitespace=True, min_length=1, pattern=r"^/", to_lower=True), Field(default="/v1/chat/completions")] # fmt: off + embeddings: Annotated[str | None, StringConstraints(strip_whitespace=True, min_length=1, pattern=r"^/", to_lower=True), Field(default="/v1/embeddings")] # fmt: off + models: Annotated[str | None, StringConstraints(strip_whitespace=True, min_length=1, pattern=r"^/", to_lower=True), Field(default="/v1/models")] # fmt: off + ocr: Annotated[str | None, StringConstraints(strip_whitespace=True, min_length=1, pattern=r"^/", to_lower=True), Field(default="/v1/ocr")] # fmt: off + rerank: Annotated[str | None, StringConstraints(strip_whitespace=True, min_length=1, pattern=r"^/", to_lower=True), Field(default="/v1/rerank")] # fmt: off + + def get_endpoint(self, endpoint: EndpointRoute) -> str | None: + if endpoint == EndpointRoute.AUDIO_TRANSCRIPTIONS: + endpoint = self.audio_transcriptions + elif endpoint == EndpointRoute.CHAT_COMPLETIONS: + endpoint = self.chat_completions + elif endpoint == EndpointRoute.EMBEDDINGS: + endpoint = self.embeddings + elif endpoint == EndpointRoute.MODELS: + endpoint = self.models + elif endpoint == EndpointRoute.OCR: + endpoint = self.ocr + elif endpoint == EndpointRoute.RERANK: + endpoint = self.rerank + else: + endpoint = None + + endpoint = endpoint if endpoint is None else endpoint.lstrip("/") + + return endpoint + + +class ModelHttpClient: + ENDPOINT_TABLE: ModelHttpClientEndpoints = ModelHttpClientEndpoints() def __init__( self, @@ -35,62 +65,176 @@ def __init__( key: str, timeout: int, model_name: str, - model_hosting_zone: str | None, + model_hosting_zone: ProviderCarbonFootprintZone, model_total_params: int | None, model_active_params: int | None, - ) -> None: + ): + self.url = url + self.key = key + self.timeout = timeout self.model_name = model_name self.model_hosting_zone = model_hosting_zone self.model_total_params = model_total_params self.model_active_params = model_active_params - self.url = url - self.key = key - self.timeout = timeout - self.type: ProviderType | None = None # set by the child ModelProvider class when the provider is created self.id: int | None = None # set by the ModelRegistry when the provider is created self.cost_prompt_tokens: float | None = None # set by the ModelRegistry when the provider is retrieved self.cost_completion_tokens: float | None = None # set by the ModelRegistry when the provider is retrieved + self.endpoint_table = self.ENDPOINT_TABLE self.headers = {"Authorization": f"Bearer {self.key}"} if self.key else {} + # request formatting @staticmethod - def import_module(type: ProviderType) -> "type[BaseModelProvider]": - """ - Static method to import a subclass of BaseModelProvider. + def format_audio_transcription_request(self, request_content: RequestContent) -> RequestContent: + """This method can be overridden by children clients to format the audio transcription request.""" + return request_content - Args: - type(str): The type of model provider to import. + @staticmethod + def format_chat_completion_request(request_content: RequestContent) -> RequestContent: + """This method can be overridden by children clients to format the chat completion request.""" + # @TODO: setup default temperature by model (default=1.0) + # @TODO: catch stream options to avoid double usage computation + return request_content - Returns: - Type[BaseModelProvider]: The subclass of BaseModelProvider. - """ + @staticmethod + def format_embeddings_request(request_content: RequestContent) -> RequestContent: + """This method can be overridden by children clients to format the embeddings request.""" + return request_content + + @staticmethod + def format_ocr_request(request_content: RequestContent) -> RequestContent: + """This method can be overridden by children clients to format the OCR request.""" + return request_content + + @staticmethod + def format_rerank_request(request_content: RequestContent) -> RequestContent: + """This method can be overridden by children clients to format the rerank request.""" + return request_content + + def _format_request(self, request_content: RequestContent) -> RequestContent: + if "model" in request_content.body: + request_content.body["model"] = self.model_name + + if "model" in request_content.form: + request_content.form["model"] = self.model_name + + if request_content.endpoint == EndpointRoute.AUDIO_TRANSCRIPTIONS: + if self.ENDPOINT_TABLE.audio_transcriptions is None: + raise NotImplementedError(f"{type(self).__name__} does not support audio transcriptions") + request_content = self.format_audio_transcription_request(request_content=request_content) - module = importlib.import_module(f"api.clients.model._{type.value}modelprovider") + elif request_content.endpoint == EndpointRoute.CHAT_COMPLETIONS: + if self.ENDPOINT_TABLE.chat_completions is None: + raise NotImplementedError(f"{type(self).__name__} does not support chat completions") + request_content = self.format_chat_completion_request(request_content=request_content) - return getattr(module, f"{type.capitalize()}ModelProvider") + elif request_content.endpoint == EndpointRoute.EMBEDDINGS: + request_content = self.format_embeddings_request(request_content=request_content) + elif request_content.endpoint == EndpointRoute.OCR: + if self.ENDPOINT_TABLE.ocr is None: + raise NotImplementedError(f"{type(self).__name__} does not support OCR") + request_content = self.format_ocr_request(request_content=request_content) + + elif request_content.endpoint == EndpointRoute.RERANK: + if self.ENDPOINT_TABLE.rerank is None: + raise NotImplementedError(f"{type(self).__name__} does not support rerank") + request_content = self.format_rerank_request(request_content=request_content) + + return request_content + + # response formatting @staticmethod - async def get_max_context_length() -> int | None: + def format_response_to_audio_transcription_response(request_content: RequestContent, response_data: dict) -> AudioTranscription: + """This method can be overridden by children clients to format the audio transcription response.""" + return AudioTranscription(**response_data) + + @staticmethod + def format_response_to_rerank_response(request_content: RequestContent, response_data: dict) -> Reranks: + """This method can be overridden by children clients to format the rerank response.""" + return Reranks(**response_data) + + @staticmethod + def format_response_to_models_response(request_content: RequestContent, response_data: dict) -> ModelsResponse: + """This method can be overridden by children clients to format the models response.""" + return ModelsResponse(**response_data) + + def _format_response(self, request_content: RequestContent, response: httpx.Response, request_latency: int | None = None) -> httpx.Response: + # @TODO: move additional data logic into a common method called by the format_response_to_* methods and get_extra_stream_chunk + # @TODO: rename _format_response to _format_unstream_response + # @TODO: unify response_content and response_data and metrics into a single request context object + content_type = response.headers.get("Content-Type", "") + if content_type != "application/json": + return response + + response_data = response.json() + additional_data = request_content.additional_data + + if request_content.model is not None: + additional_data.update({"model": request_content.model}) + + if request_context.get().id is None: + if isinstance(response_data, dict) and "id" in response_data: + request_context.get().id = response_data["id"] + else: + request_context.get().id = generate_request_id() + + additional_data.update({"id": request_context.get().id}) + + usage: Usage | None = self._get_usage(request_content=request_content, response_data=response_data, request_latency=request_latency) + if usage is not None: + additional_data.update({"usage": usage.model_dump()}) + + if request_content.endpoint == EndpointRoute.AUDIO_TRANSCRIPTIONS: + response_data = self.format_response_to_audio_transcription_response( + request_content=request_content, + response_data=response_data, + ).model_dump() + + elif request_content.endpoint == EndpointRoute.MODELS: + response_data = self.format_response_to_models_response(request_content=request_content, response_data=response_data).model_dump() + + elif request_content.endpoint == EndpointRoute.RERANK: + response_data = self.format_response_to_rerank_response(request_content=request_content, response_data=response_data).model_dump() + + response_data.update(additional_data) + response = httpx.Response(status_code=response.status_code, content=dumps(response_data)) + + return response + + async def get_model_info(self) -> ModelResponse | None: """ - Get the max context length of the model provider to store in the database. Useful - to check provider consistency. + Check if the model provider is healthy. """ - pass + request_content = RequestContent(endpoint=EndpointRoute.MODELS, method="GET") + response = await self.forward_request(request_content=request_content) + + if response.status_code != 200: + return None + + data = response.json()["data"] + model = next((ModelResponse(**model) for model in data if model["id"] == self.model_name or self.model_name in model["aliases"]), None) + + return model async def get_vector_size(self) -> int | None: - if self.ENDPOINT_TABLE.embeddings is None: + if self.endpoint_table.embeddings is None: return None - url = urljoin(base=self.url, url=self.ENDPOINT_TABLE.embeddings.lstrip("/")) + request_content = RequestContent( + model=self.model_name, + endpoint=EndpointRoute.EMBEDDINGS, + method="POST", + body={"model": self.model_name, "input": "hello world"}, + ) + response = await self.forward_request(request_content=request_content) - 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"] 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: @@ -105,8 +249,7 @@ def _get_usage(self, request_content: RequestContent, response_data: dict | list Returns: Usage | None: The usage data. """ - - usage = request_context.get().usage + usage: Usage | None = request_context.get().usage # In Celery worker processes the FastAPI app initialization (which sets global_context.tokenizer) # might not have fully run. Accessing global_context.tokenizer directly could raise AttributeError. # We skip usage computation if tokenizer is absent so we still return the provider response. @@ -141,83 +284,6 @@ def _get_usage(self, request_content: RequestContent, response_data: dict | list return usage - def _format_request(self, request_content: RequestContent) -> RequestContent: - """ - Format a request to a provider model. This method can be overridden by a subclass to add additional headers or parameters. This method format the requested endpoint thanks the ENDPOINT_TABLE attribute. - - Args: - request_content(RequestContent): The request content to format. - - Returns: - request_content(RequestContent): The formatted request content. - """ - if "model" in request_content.body: - request_content.body["model"] = self.model_name - - if "model" in request_content.form: - request_content.form["model"] = self.model_name - try: - if request_content.endpoint == EndpointRoute.AUDIO_TRANSCRIPTIONS: - request_content = CreateAudioTranscription.format_request(provider_type=self.type, request_content=request_content) - - if request_content.endpoint == EndpointRoute.CHAT_COMPLETIONS: - request_content = CreateChatCompletion.format_request(provider_type=self.type, request_content=request_content) - - if request_content.endpoint == EndpointRoute.RERANK: - request_content = CreateRerank.format_request(provider_type=self.type, request_content=request_content) - - except Exception as e: - logger.error(f"Failed to format request for {self.model_name}: {e}.", exc_info=True) - raise RequestFormatFailedException() - - return request_content - - def _format_response(self, request_content: RequestContent, response: httpx.Response, request_latency: int | None = None) -> httpx.Response: - """ - Format a response from a TEI model, overridden base class method to convert TEI reranking response to a standard response. - """ - - content_type = response.headers.get("Content-Type", "") - if content_type == "application/json": - response_data = response.json() - - usage = self._get_usage(request_content=request_content, response_data=response_data, request_latency=request_latency) - - if request_context.get().id is None: - if isinstance(response_data, dict) and "id" in response_data: - request_context.get().id = response_data["id"] - else: - request_context.get().id = generate_request_id() - - additional_data = request_content.additional_data - additional_data.update({"model": request_content.model, "id": request_context.get().id, "usage": usage.model_dump()}) - - try: - if request_content.endpoint == EndpointRoute.AUDIO_TRANSCRIPTIONS: - response_data = AudioTranscription.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, - request_content=request_content, - response_data=response_data, - ).model_dump() - - else: - response_data.update(additional_data) - - except Exception as e: - logger.error(f"Failed to build response from {self.model_name}: {e}.", exc_info=True) - raise ResponseFormatFailedException() - - response = httpx.Response(status_code=response.status_code, content=dumps(response_data)) - - return response - @staticmethod async def _ensure_timeseries_exists(redis_client: AsyncRedis, key: str) -> None: """ @@ -235,15 +301,18 @@ async def _ensure_timeseries_exists(redis_client: AsyncRedis, key: str) -> None: except Exception: pass - async def _log_performance_metric(self, redis_client: AsyncRedis, ttft: int | None, latency: int | None) -> None: + async def _log_performance_metric(self, redis_client: AsyncRedis | None = None, ttft: int | None = None, latency: int | None = None) -> None: """ Log performance metrics in redis. Args: - redis_client(AsyncRedis): The redis client to use for the request. + redis_client(AsyncRedis | None): The redis client to use for the request. If None, performance metrics are not logged. ttft(int | None): The time to first token in milliseconds (ms). latency(int | None): The latency in milliseconds (ms). """ + if redis_client is None: + return + request_context.get().ttft = ttft request_context.get().latency = latency @@ -270,25 +339,45 @@ async def _log_performance_metric(self, redis_client: AsyncRedis, ttft: int | No def _elapsed_ms(start_time: float) -> int: return int((time.perf_counter() - start_time) * 1000) # ms - async def forward_request(self, request_content: RequestContent, redis_client: AsyncRedis) -> httpx.Response: + @staticmethod + def _build_inflight_key(provider_id: int) -> str: + return f"{PREFIX__REDIS_METRIC_GAUGE}:{Metric.INFLIGHT.value}:{provider_id}" + + async def _increment_inflight_key(self, redis_client: AsyncRedis | None = None) -> bool: + if redis_client is None: + return False + inflight_key = self._build_inflight_key(provider_id=self.id) + try: + await redis_retry(redis_client.incr, name=inflight_key, max_retries=2) + return True + except Exception: + return False + + async def _decrement_inflight_key(self, inflight_is_incremented: bool, redis_client: AsyncRedis | None = None) -> None: + if redis_client is None or not inflight_is_incremented: + return + inflight_key = self._build_inflight_key(provider_id=self.id) + try: + await redis_retry(redis_client.decr, name=inflight_key, max_retries=2) + except Exception as e: + logger.exception(msg=f"Failed to decrement inflight key {inflight_key} for provider {self.id}: {e}") + + async def forward_request(self, request_content: RequestContent, redis_client: AsyncRedis | None = None) -> httpx.Response: """ Forward a request to a provider model and add model name to the response. Optionally, add additional data to the response. Args: - redis_client(AsyncRedis): The redis client to use for the request. + redis_client(AsyncRedis | None): The redis client to use for the request. If None, performance metrics are not logged. request_content(RequestContent): The request content to use for the request. Returns: httpx.Response: The response from the API. """ - url = urljoin(base=self.url, url=self.ENDPOINT_TABLE.get_endpoint(endpoint=request_content.endpoint).lstrip("/")) + url = urljoin(base=self.url, url=self.endpoint_table.get_endpoint(endpoint=request_content.endpoint)) request_content = self._format_request(request_content=request_content) - - inflight_key = f"{PREFIX__REDIS_METRIC_GAUGE}:{Metric.INFLIGHT.value}:{self.id}" + inflight_is_incremented = await self._increment_inflight_key(redis_client=redis_client) try: - await redis_retry(redis_client.incr, name=inflight_key, max_retries=2) - async with httpx.AsyncClient(timeout=self.timeout) as async_client: try: start_time = time.perf_counter() @@ -329,7 +418,7 @@ async def forward_request(self, request_content: RequestContent, redis_client: A message = response.text raise HTTPException(status_code=response.status_code, detail=message) finally: - await redis_retry(redis_client.decr, name=inflight_key, max_retries=2) + await self._decrement_inflight_key(inflight_is_incremented=inflight_is_incremented, redis_client=redis_client) # add additional data to the response latency = self._elapsed_ms(start_time=start_time) @@ -371,27 +460,19 @@ def _get_extra_stream_chunk(self, request_content: RequestContent, buffer: list[ return extra_chunk - async def forward_stream(self, request_content: RequestContent, redis_client: AsyncRedis): + async def forward_stream(self, request_content: RequestContent, redis_client: AsyncRedis | None = None): """ Forward a stream request to a provider model and add model name to the response. Optionally, add additional data to the response. Args: - redis_client(AsyncRedis): The redis client to use for the request. + redis_client(AsyncRedis | None): The redis client to use for the request. If None, performance metrics are not logged. request_content(RequestContent): The request content to use for the request. """ - url = urljoin(base=self.url, url=self.ENDPOINT_TABLE.get_endpoint(endpoint=request_content.endpoint).lstrip("/")) + url = urljoin(base=self.url, url=self.endpoint_table.get_endpoint(endpoint=request_content.endpoint)) request_content = self._format_request(request_content=request_content) - - inflight_key = f"{PREFIX__REDIS_METRIC_GAUGE}:{Metric.INFLIGHT.value}:{self.id}" - inflight_incremented = False + inflight_is_incremented = await self._increment_inflight_key(redis_client=redis_client) async with httpx.AsyncClient(timeout=self.timeout) as async_client: - try: - await redis_retry(redis_client.incr, name=inflight_key, max_retries=2) - inflight_incremented = True - except Exception: - logger.error("Unable to increment redis requests inflight key") - try: async with async_client.stream( method=request_content.method, @@ -461,8 +542,4 @@ async def forward_stream(self, request_content: RequestContent, redis_client: As logger.exception(msg=f"Failed to forward stream request to {self.model_name}: {e}.") yield dumps({"detail": type(e).__name__}), 500 finally: - if inflight_incremented: - try: - await redis_retry(redis_client.decr, name=inflight_key, max_retries=2) - except Exception: - logger.error("Unable to decrement redis requests inflight key") + await self._decrement_inflight_key(inflight_is_incremented=inflight_is_incremented, redis_client=redis_client) diff --git a/api/infrastructure/http/model/_openaimodelhttpclient.py b/api/infrastructure/http/model/_openaimodelhttpclient.py new file mode 100644 index 000000000..026761e98 --- /dev/null +++ b/api/infrastructure/http/model/_openaimodelhttpclient.py @@ -0,0 +1,5 @@ +from api.infrastructure.http.model._modelhttpclient import ModelHttpClient, ModelHttpClientEndpoints + + +class OpenaiModelHttpClient(ModelHttpClient): + ENDPOINT_TABLE = ModelHttpClientEndpoints(ocr=None, rerank=None) diff --git a/api/infrastructure/http/model/_teimodelhttpclient.py b/api/infrastructure/http/model/_teimodelhttpclient.py new file mode 100644 index 000000000..7eaf8c510 --- /dev/null +++ b/api/infrastructure/http/model/_teimodelhttpclient.py @@ -0,0 +1,48 @@ +from typing import Literal + +from pydantic import BaseModel, Field + +from api.infrastructure.fastapi.schemas.models import ModelResponse, ModelsResponse +from api.infrastructure.http.model._modelhttpclient import ModelHttpClient, ModelHttpClientEndpoints +from api.schemas.core.models import RequestContent +from api.schemas.rerank import RerankResult, Reranks + + +class TeiCreateRerankBody(BaseModel): + query: str = Field(..., examples=["What is Deep Learning?"]) + raw_scores: bool = Field(False, examples=[False]) + return_text: bool = Field(False, examples=[False]) + texts: list[str] = Field(..., examples=[["Deep Learning is ..."]]) + truncate: bool | None = Field(False, examples=[False]) + truncation_direction: Literal["left", "right"] = "right" + + +class TeiModelHttpClient(ModelHttpClient): + ENDPOINT_TABLE = ModelHttpClientEndpoints(audio_transcriptions=None, chat_completions=None, models="/info", ocr=None, rerank="/rerank") + + @staticmethod + def format_rerank_request(request_content: RequestContent) -> RequestContent: + request_content.additional_data["top_n"] = request_content.body.get("top_n") + request_content.body = TeiCreateRerankBody(query=request_content.body["query"], texts=request_content.body["documents"]).model_dump() + return request_content + + @staticmethod + def format_rerank_response(request_content: RequestContent, response_data: dict) -> Reranks: + # @TODO: add model copy (after move additional data logic into format_response_to_* methods) + response_data = sorted(response_data, key=lambda x: x["score"], reverse=True)[: request_content.additional_data.get("top_n")] + request_content.additional_data.pop("top_n") + results = [RerankResult(relevance_score=rank["score"], index=rank["index"]) for rank in response_data] + return Reranks(results=results, **request_content.additional_data) + + # response formatting + @staticmethod + def format_response_to_models_response(request_content: RequestContent, response_data: dict) -> ModelsResponse: + data = [ + ModelResponse( + id=response_data.get("model_id", ""), + created=0, + owned_by="tei", + max_context_length=response_data.get("max_input_length"), + ) + ] + return ModelsResponse(data=data) diff --git a/api/infrastructure/http/model/_vllmmodelhttpclient.py b/api/infrastructure/http/model/_vllmmodelhttpclient.py new file mode 100644 index 000000000..f5447f9d6 --- /dev/null +++ b/api/infrastructure/http/model/_vllmmodelhttpclient.py @@ -0,0 +1,21 @@ +from api.infrastructure.fastapi.schemas.models import ModelResponse, ModelsResponse +from api.infrastructure.http.model._modelhttpclient import ModelHttpClient, ModelHttpClientEndpoints +from api.schemas.core.models import RequestContent + + +class VllmModelHttpClient(ModelHttpClient): + ENDPOINT_TABLE = ModelHttpClientEndpoints(rerank="/v2/rerank") + + # response formatting + @staticmethod + def format_response_to_models_response(request_content: RequestContent, response_data: dict) -> ModelsResponse: + data = [ + ModelResponse( + id=model.get("id"), + created=model.get("created"), + owned_by=model.get("owned_by"), + max_context_length=model.get("max_model_len"), + ) + for model in response_data.get("data", []) + ] + return ModelsResponse(data=data) diff --git a/api/infrastructure/model/_modelprovidergateway.py b/api/infrastructure/model/_modelprovidergateway.py index 4d47c1804..595b5aba7 100644 --- a/api/infrastructure/model/_modelprovidergateway.py +++ b/api/infrastructure/model/_modelprovidergateway.py @@ -1,32 +1,49 @@ -from api.clients.model import BaseModelProvider from api.domain.model import ModelType as RouterType from api.domain.provider import ProviderCapabilities, ProviderGateway, ProviderNotReachableError +from api.domain.provider.entities import ProviderType +from api.infrastructure.http.model import ( + AlbertModelHttpClient, + MistralModelHttpClient, + ModelHttpClient, + OpenaiModelHttpClient, + TeiModelHttpClient, + VllmModelHttpClient, +) class ModelProviderGateway(ProviderGateway): - async def get_capabilities(self, router_type, provider_type, url, key, timeout, model_name): - try: - client = self._build_client(provider_type, url, key, timeout, model_name) - max_context_length = await client.get_max_context_length() - if router_type == RouterType.TEXT_EMBEDDINGS_INFERENCE: - vector_size = await client.get_vector_size() - else: - vector_size = None - return ProviderCapabilities( - max_context_length=max_context_length, - vector_size=vector_size, - ) - except AssertionError as e: + async def get_capabilities(self, router_type, provider_type, url, key, timeout, model_name) -> ProviderCapabilities | ProviderNotReachableError: + client = ModelProviderGateway._build_client(provider_type, url, key, timeout, model_name) + model_info = await client.get_model_info() + + if model_info is None: return ProviderNotReachableError(model_name) - def _build_client(self, provider_type, url, key, timeout, model_name): - cls = BaseModelProvider.import_module(type=provider_type) - return cls( + max_context_length = model_info.max_context_length + + if router_type == RouterType.TEXT_EMBEDDINGS_INFERENCE: + vector_size = await client.get_vector_size() + else: + vector_size = None + + return ProviderCapabilities(max_context_length=max_context_length, vector_size=vector_size) + + @staticmethod + def _build_client(provider_type, url, key, timeout, model_name) -> ModelHttpClient: + provider_class: dict[ProviderType, type[ModelHttpClient]] = { + ProviderType.ALBERT: AlbertModelHttpClient, + ProviderType.MISTRAL: MistralModelHttpClient, + ProviderType.OPENAI: OpenaiModelHttpClient, + ProviderType.TEI: TeiModelHttpClient, + ProviderType.VLLM: VllmModelHttpClient, + } + + return provider_class[provider_type]( url=url, key=key, timeout=timeout, model_name=model_name, + model_active_params=None, + model_total_params=None, model_hosting_zone=None, - model_total_params=0, - model_active_params=0, ) diff --git a/api/schemas/admin/providers.py b/api/schemas/admin/providers.py index e55321ac5..1e374df09 100644 --- a/api/schemas/admin/providers.py +++ b/api/schemas/admin/providers.py @@ -1,5 +1,4 @@ from enum import Enum -from typing import Literal import pycountry from pydantic import Field, constr, model_validator @@ -75,27 +74,3 @@ def validate_model(self): raise ValueError("QoS value is required if QoS metric is provided.") return self - - -class Provider(BaseModel): - object: Literal["provider"] = "provider" - id: int = Field(..., description="Provider ID.") # fmt: off - router_id: int = Field(..., description="ID of the router that owns the provider.") # fmt: off - user_id: int = Field(..., description="ID of the user that owns the provider.") # fmt: off - type: ProviderType = Field(..., description="Provider type.") # fmt: off - url: constr(strip_whitespace=True, min_length=1, to_lower=True) | None = Field(default=None, description="Provider API url. The url must only contain the domain name (without `/v1` suffix for example).") # fmt: off - key: str | None = Field(description="Provider API key.") # fmt: off - timeout: int = Field(..., description="Timeout for the provider requests, after user receive an 500 error (model is too busy).") # fmt: off - model_name: str = Field(..., description="Model name from the model provider.") # fmt: off - model_hosting_zone: ProviderCarbonFootprintZone = Field(default=ProviderCarbonFootprintZone.WOR, description="Model hosting zone using ISO 3166-1 alpha-3 code format (e.g., `WOR` for World, `FRA` for France, `USA` for United States). This determines the electricity mix used for carbon intensity calculations. For more information, see https://ecologits.ai", examples=["WOR"]) # fmt: off - model_total_params: int = Field(default=0, ge=0, description="Total params of the model in billions of parameters for carbon footprint computation. If not provided, the active params will be used if provided, else carbon footprint will not be computed. For more information, see https://ecologits.ai") # fmt: off - model_active_params: int = Field(default=0, ge=0, description="Active params of the model in billions of parameters for carbon footprint computation. If not provided, the total params will be used if provided, else carbon footprint will not be computed. For more information, see https://ecologits.ai") # fmt: off - qos_metric: Metric | None = Field(description="The metric to use for the QoS policy. If not provided, no QoS policy is applied.") # fmt: off - qos_limit: float | None = Field(default=None, ge=0.0, description="The value to use for the quality of service. Depends of the metric, the value can be a percentile, a threshold, etc.") # fmt: off - created: int | None = Field(default=None, description="Time of creation, as Unix timestamp.") # fmt: off - updated: int | None = Field(default=None, description="Time of last update, as Unix timestamp.") # fmt: off - - -class Providers(BaseModel): - object: Literal["list"] = "list" - data: list[Provider] diff --git a/api/schemas/audio.py b/api/schemas/audio.py index 85562d84a..8e8c022f4 100644 --- a/api/schemas/audio.py +++ b/api/schemas/audio.py @@ -1,14 +1,14 @@ +<<<<<<< HEAD import base64 +======= +>>>>>>> 8582c6f3 (feat(models): convert get capabilities of a provider with standard forward request method) from enum import StrEnum from fastapi import File, Form, UploadFile from fastapi.exceptions import RequestValidationError -from mistralai.client.models import AudioChunk, ChatCompletionRequest, TextChunk, UserMessage from pydantic import Field, ValidationError, field_validator from api.schemas import BaseModel -from api.schemas.admin.providers import ProviderType -from api.schemas.core.models import RequestContent from api.schemas.usage import Usage from api.utils.exceptions import FileSizeLimitExceededException from api.utils.variables import SUPPORTED_LANGUAGES @@ -61,57 +61,9 @@ def validate_file(cls, file: UploadFile) -> UploadFile: raise FileSizeLimitExceededException() return file - @staticmethod - def format_request(provider_type: ProviderType, request_content: RequestContent): - match provider_type: - case ProviderType.ALBERT: - request_content.form["response_format"] = "json" - return request_content - - case ProviderType.MISTRAL: - text = request_content.form.get("prompt") or f"Transcribe this audio in this language : {request_content.form.get("language", "en")}" - input_audio = base64.b64encode(request_content.files["file"][1]).decode("utf-8") - request_content.body = ChatCompletionRequest( - model=request_content.form["model"], - messages=[ - UserMessage( - role="user", - content=[AudioChunk(type="input_audio", input_audio=input_audio), TextChunk(type="text", text=text)], - ) - ], - temperature=request_content.form.get("temperature"), - ).model_dump() - request_content.files = {} - request_content.form = {} - return request_content - - case ProviderType.VLLM: - return request_content - - case _: - raise NotImplementedError(f"Provider {provider_type} not implemented") - class AudioTranscription(BaseModel): id: str = Field(default=..., description="A unique identifier for the audio transcription.") text: str = Field(default=..., description="The transcription text.") model: str = Field(default=..., description="The model used to generate the transcription.") usage: Usage = Field(default_factory=Usage, description="Usage information for the request.") - - @classmethod - def build_from(cls, provider_type: ProviderType, request_content: RequestContent, response_data: dict) -> "AudioTranscription": - match provider_type: - case ProviderType.ALBERT: - response_data.update(request_content.additional_data) - return cls(**response_data) - - case ProviderType.MISTRAL: - text = response_data["choices"][0]["message"]["content"] - return cls(text=text, **request_content.additional_data) - - case ProviderType.VLLM: - response_data.update(request_content.additional_data) - return cls(**response_data) - - case _: - raise NotImplementedError(f"Provider {provider_type} not implemented") diff --git a/api/schemas/chat.py b/api/schemas/chat.py index baa54ac41..ab524e3c7 100644 --- a/api/schemas/chat.py +++ b/api/schemas/chat.py @@ -2,13 +2,10 @@ from json import JSONDecodeError from typing import Annotated, Any, Literal -from mistralai.client.models import ChatCompletionRequest from openai.types.chat import ChatCompletion, ChatCompletionChunk from pydantic import Field, field_validator, model_validator from api.schemas import BaseModel -from api.schemas.admin.providers import ProviderType -from api.schemas.core.models import RequestContent from api.schemas.search import Search, SearchArgs from api.schemas.usage import Usage @@ -74,49 +71,6 @@ def convert_search_args_to_tool(self) -> "CreateChatCompletion": return self - @staticmethod - def format_request(provider_type: ProviderType, request_content: RequestContent): - match provider_type: - case ProviderType.ALBERT: - return request_content - - case ProviderType.MISTRAL: - try: - request_content.body = ChatCompletionRequest(**request_content.body).model_dump(by_alias=True) - except Exception: - # apply a minimal formatting and ignore error to let the provider raise the correct 422 error - # see https://docs.mistral.ai/api#operation-chat_completion_v1_chat_completions_post - request_content.body = { - "frequency_penalty": request_content.body.get("frequency_penalty") or 0.0, - "max_tokens": request_content.body.get("max_tokens"), - "messages": request_content.body.get("messages"), - "model": request_content.body.get("model"), - "n": request_content.body.get("n"), - "parallel_tool_calls": request_content.body.get("parallel_tool_calls") or False, - "prediction": request_content.body.get("prediction") or {}, - "presence_penalty": request_content.body.get("presence_penalty") or 0.0, - "prompt_mode": request_content.body.get("prompt_mode"), - "random_seed": request_content.body.get("random_seed") or request_content.body.get("seed"), - "response_format": request_content.body.get("response_format") or {"type": "text"}, - "safe_prompt": request_content.body.get("safe_prompt") or False, - "stop": request_content.body.get("stop") or [], - "stream": request_content.body.get("stream") or False, - "temperature": request_content.body.get("temperature"), - "tool_choice": request_content.body.get("tool_choice"), - "tools": request_content.body.get("tools"), - "top_p": request_content.body.get("top_p") or 1.0, - } - return request_content - - case ProviderType.OPENAI: - return request_content - - case ProviderType.VLLM: - return request_content - - case _: - raise NotImplementedError(f"Provider {provider_type} not implemented") - class ChatCompletion(ChatCompletion): id: str = Field(default=None, description="A unique identifier for the chat completion.") diff --git a/api/schemas/core/models.py b/api/schemas/core/models.py index d77bbe009..33b0e8285 100644 --- a/api/schemas/core/models.py +++ b/api/schemas/core/models.py @@ -1,39 +1,11 @@ -from enum import Enum +from enum import StrEnum from http import HTTPMethod from typing import Annotated from pydantic import BaseModel, Field, StringConstraints -from api.utils import variables from api.utils.variables import EndpointRoute -Endpoint = Enum("Endpoint", {name.upper(): value for name, value in vars(variables).items() if name.startswith("ENDPOINT__")}, type=str) - - -class ProviderEndpoints(BaseModel): - audio_transcriptions: Annotated[str | None, StringConstraints(strip_whitespace=True, min_length=1, pattern=r"^/", to_lower=True), Field(default=None)] # fmt: off - chat_completions: Annotated[str | None, StringConstraints(strip_whitespace=True, min_length=1, pattern=r"^/", to_lower=True), Field(default=None)] - embeddings: Annotated[str | None, StringConstraints(strip_whitespace=True, min_length=1, pattern=r"^/", to_lower=True), Field(default=None)] - models: Annotated[str | None, StringConstraints(strip_whitespace=True, min_length=1, pattern=r"^/", to_lower=True), Field(default=None)] - ocr: Annotated[str | None, StringConstraints(strip_whitespace=True, min_length=1, pattern=r"^/", to_lower=True), Field(default=None)] - rerank: Annotated[str | None, StringConstraints(strip_whitespace=True, min_length=1, pattern=r"^/", to_lower=True), Field(default=None)] - - def get_endpoint(self, endpoint: EndpointRoute) -> str | None: - if endpoint == EndpointRoute.AUDIO_TRANSCRIPTIONS: - return self.audio_transcriptions - elif endpoint == EndpointRoute.CHAT_COMPLETIONS: - return self.chat_completions - elif endpoint == EndpointRoute.EMBEDDINGS: - return self.embeddings - elif endpoint == EndpointRoute.MODELS: - return self.models - elif endpoint == EndpointRoute.OCR: - return self.ocr - elif endpoint == EndpointRoute.RERANK: - return self.rerank - else: - return None - class RequestContent(BaseModel): method: HTTPMethod @@ -48,23 +20,8 @@ class RequestContent(BaseModel): # @TODO: build body with model_fields_set to exclude unset fields -class Metric(str, Enum): +class Metric(StrEnum): TTFT = "ttft" # time to first token LATENCY = "latency" # requests latency INFLIGHT = "inflight" # requests concurrency PERFORMANCE = "performance" # custom performance metric - - -# TEI -class TruncationDirection(Enum): - left = "left" - right = "right" - - -class TEICreateRerank(BaseModel): - query: str = Field(..., examples=["What is Deep Learning?"]) - raw_scores: bool = Field(False, examples=[False]) - return_text: bool = Field(False, examples=[False]) - texts: list[str] = Field(..., examples=[["Deep Learning is ..."]]) - truncate: bool | None = Field(False, examples=[False]) - truncation_direction: TruncationDirection = "right" diff --git a/api/schemas/rerank.py b/api/schemas/rerank.py index 5e80feb30..c8280e8b7 100644 --- a/api/schemas/rerank.py +++ b/api/schemas/rerank.py @@ -3,8 +3,6 @@ from pydantic import Field, StringConstraints from api.schemas import BaseModel -from api.schemas.admin.providers import ProviderType -from api.schemas.core.models import RequestContent, TEICreateRerank from api.schemas.usage import Usage @@ -14,24 +12,6 @@ class CreateRerank(BaseModel): model: Annotated[str, StringConstraints(min_length=1, strip_whitespace=True), Field(default=..., description="The model to use for the reranking, call `/v1/models` endpoint to get the list of available models, only `text-classification` model type is supported.")] # fmt: off top_n: Annotated[int | None, Field(default=None, ge=1, description="The number of top results to return. If set to None, all results will be returned.")] # fmt: off - @staticmethod - def format_request(provider_type: ProviderType, request_content: RequestContent): - match provider_type: - case ProviderType.ALBERT: - return request_content - - case ProviderType.TEI: - request_content.additional_data["top_n"] = request_content.body.get("top_n") - request_content.body = TEICreateRerank(query=request_content.body["query"], texts=request_content.body["documents"]).model_dump() - - return request_content - - case ProviderType.VLLM: - return request_content - - case _: - raise NotImplementedError(f"Provider {provider_type} not implemented") - class Rerank(BaseModel): object: Annotated[Literal["rerank"], Field(default="rerank", description="The type of the object.")] @@ -50,27 +30,3 @@ class Reranks(BaseModel): results: Annotated[list[RerankResult], Field(default=..., description="The list of reranked texts.")] model: Annotated[str, Field(default=..., description="The model used to generate the reranking.")] usage: Annotated[Usage, Field(default_factory=Usage, description="Usage information for the request.")] - - @classmethod - def build_from(cls, provider_type: ProviderType, request_content: RequestContent, response_data: dict): - match provider_type: - case ProviderType.ALBERT: - response_data.update(request_content.additional_data) - return cls(**response_data) - - case ProviderType.TEI: - data = [] - results = [] - if request_content.additional_data["top_n"] is not None: - response_data = sorted(response_data, key=lambda x: x["score"], reverse=True)[: request_content.additional_data.get("top_n")] - request_content.additional_data.pop("top_n") - for rank in response_data: - results.append(RerankResult(relevance_score=rank["score"], index=rank["index"])) - return cls(results=results, **request_content.additional_data) - - case ProviderType.VLLM: - response_data.update(request_content.additional_data) - return cls(**response_data) - - case _: - raise NotImplementedError(f"Provider {provider_type} not implemented") diff --git a/api/tests/integ/test_providers/test_mistral_image_to_text.py b/api/tests/integ/test_providers/test_mistral_image_to_text.py index 84aea95aa..beb09671d 100644 --- a/api/tests/integ/test_providers/test_mistral_image_to_text.py +++ b/api/tests/integ/test_providers/test_mistral_image_to_text.py @@ -35,7 +35,7 @@ def setup_mistral_image_to_text(client: TestClient): ) role_id = create_role(router_id=router_id, client=client) user_id = create_user(role_id=role_id, client=client) - _, key = create_token(user_id=user_id, token_name=f"test-token-{dt.datetime.now().strftime("%Y%m%d%H%M%S")}", client=client) + _, key = create_token(user_id=user_id, token_name=f"test-token-{dt.datetime.now().strftime('%Y%m%d%H%M%S')}", client=client) yield key, process.model_name except Exception as e: diff --git a/api/tests/unit/infrastructure/http/factories.py b/api/tests/unit/infrastructure/http/factories.py new file mode 100644 index 000000000..e09a03dd7 --- /dev/null +++ b/api/tests/unit/infrastructure/http/factories.py @@ -0,0 +1,343 @@ +import random + +import factory +from faker import Faker + +from api.schemas.chat import CreateChatCompletion +from api.schemas.core.models import RequestContent +from api.schemas.usage import Usage +from api.utils.variables import EndpointRoute + +RERANK_DOCUMENTS_COUNT = 10 +RERANK_TOP_N = 2 +fake = Faker() + + +class FormattedRequestContentFactory(factory.DictFactory): + """Correspond to the request content after passed in the ModelHttpClient._format_request method + but before the specific formatting in the ModelHttpClient by children classes + """ + + class Meta: + model = RequestContent + + model = factory.Faker("bothify", text="model-????") + method = factory.Faker("random_element", elements=["GET", "POST", "PUT", "DELETE"]) + endpoint = factory.Faker("random_element", elements=list(EndpointRoute)) + body = factory.LazyFunction(lambda: {}) + form = factory.LazyFunction(lambda: {}) + files = factory.LazyFunction(lambda: {}) + additional_data = factory.LazyFunction(lambda: {}) + + class Params: + audio_transcriptions = factory.Trait( + method="POST", + endpoint=EndpointRoute.AUDIO_TRANSCRIPTIONS, + form=factory.LazyAttribute( + lambda self: { + "model": self.model, + "language": "fr", + "prompt": fake.sentence(), + "temperature": 0.3, + "response_format": "json", + } + ), + files={"file": ("audio.wav", b"test-audio-content", "audio/wav")}, + ) + chat_completions = factory.Trait( + method="POST", + endpoint=EndpointRoute.CHAT_COMPLETIONS, + body=factory.LazyAttribute( + lambda self: CreateChatCompletion( + model=self.model, + messages=[ + {"role": "system", "content": fake.sentence()}, + {"role": "user", "content": fake.sentence()}, + {"role": "assistant", "content": fake.sentence()}, + {"role": "user", "content": fake.sentence()}, + ], + frequency_penalty=None, + response_format={ + "type": "json_schema", + "json_schema": { + "name": "DummyResponseFormat", + "schema": { + "properties": { + "dummy_list": {"items": {"type": "string"}, "title": "Dummy List", "type": "array"}, + "dummy_str": {"title": "Dummy Str", "type": "string"}, + "dummy_optional_bool": {"default": False, "title": "Dummy Optional Bool", "type": "boolean"}, + "dummy_nullable_int": {"anyOf": [{"type": "integer"}, {"type": "null"}], "title": "Dummy Nullable Int"}, + }, + "required": ["dummy_list", "dummy_str", "dummy_nullable_int"], + "title": "DummyResponseFormat", + "type": "object", + }, + }, + }, + tool_choice="required", + seed=10, + stop=None, + stream=False, + ).model_dump() + ), + ) + rerank = factory.Trait( + method="POST", + endpoint=EndpointRoute.RERANK, + body=factory.LazyAttribute( + lambda self: { + "model": self.model, + "query": fake.sentence(), + "documents": fake.sentences(nb=RERANK_DOCUMENTS_COUNT), + "top_n": RERANK_TOP_N, + } + ), + ) + embeddings = factory.Trait(method="POST", endpoint=EndpointRoute.EMBEDDINGS, body={"input": ["test", "test2", "test3"]}) + models = factory.Trait(method="GET", endpoint=EndpointRoute.MODELS) + + +# TEI factories +class TeiFormattedRequestContentFactory(FormattedRequestContentFactory): + """Correspond to the request content after passed in TEI specific formatting in the TEIModelHttpClient""" + + class Params: + rerank = factory.Trait( + method="POST", + endpoint=EndpointRoute.RERANK, + body=factory.LazyAttribute(lambda _: {"query": fake.sentence(), "texts": fake.sentences(nb=RERANK_DOCUMENTS_COUNT)}), + additional_data=factory.LazyAttribute( + lambda self: { + "id": "request-1234567890", + "top_n": RERANK_TOP_N, + "model": self.model, + "usage": Usage(prompt_tokens=10, completion_tokens=10, total_tokens=20).model_dump(), + } + ), + ) + embeddings = factory.Trait(method="POST", endpoint=EndpointRoute.EMBEDDINGS, body={"input": ["test", "test2", "test3"]}) + models = factory.Trait(method="GET", endpoint=EndpointRoute.MODELS) + + +class TeiRerankItemResponseFactory(factory.DictFactory): + class Params: + raw_scores = False + return_text = False + + index = factory.Faker("random_int", min=0, max=2) + score = factory.Maybe( + "raw_scores", + yes_declaration=factory.Faker("pyfloat", right_digits=6, min_value=0, max_value=1), + no_declaration=factory.Faker("pyfloat", right_digits=6, min_value=0, max_value=10), + ) + text = factory.Maybe("return_text", yes_declaration=factory.Faker("sentence"), no_declaration=factory.declarations.SKIP) + + +class TeiRerankResponseFactory: + def __new__(cls, request_content: TeiFormattedRequestContentFactory, raw_scores: bool = False, return_text: bool = False): + len_texts = len(request_content.body["texts"]) + texts_indices = list(range(len_texts)) + random.shuffle(texts_indices) + + return TeiRerankItemResponseFactory.create_batch( + size=len_texts, + raw_scores=raw_scores, + return_text=return_text, + index=factory.Iterator(texts_indices), + ) + + +class TeiModelsResponseFactory(factory.DictFactory): + model_id = "BAAI/bge-m3" + model_sha = None + model_dtype = "float16" + model_type = {"embedding": {"pooling": "cls"}} + max_concurrent_requests = 512 + max_input_length = 8192 + max_batch_tokens = 16384 + max_batch_requests = None + max_client_batch_size = 64 + auto_truncate = False + tokenization_workers = 126 + version = "1.8.3" + sha = "3120a50a84b22bb3cd84152c11d4373faea6d99a" + docker_label = "sha-3120a50" + + class Params: + embedding = factory.Trait( + model_id="BAAI/bge-m3", + model_type={"embedding": {"pooling": "cls"}}, + ) + reranker = factory.Trait( + model_id="BAAI/bge-reranker-v2-m3", + model_type={"reranker": {"id2label": {"0": "LABEL_0"}, "label2id": {"LABEL_0": 0}}}, + ) + + +# Mistral factories +class MistralModelsItemResponseFactory(factory.DictFactory): + id = "mistral-medium-2508" + object = "model" + created = 1773667856 + owned_by = "mistralai" + capabilities = { + "completion_chat": True, + "function_calling": True, + "completion_fim": False, + "fine_tuning": True, + "vision": True, + "ocr": False, + "classification": False, + "moderation": False, + "audio": False, + } + name = "mistral-medium-2508" + description = "Update on Mistral Medium 3 with improved capabilities." + max_context_length = 131072 + aliases = ["mistral-medium-latest"] + deprecation = None + deprecation_replacement_model = None + default_model_temperature = 0.3 + type = "base" + + class Params: + medium = factory.Trait( + id="mistral-medium-2508", + capabilities={ + "completion_chat": True, + "function_calling": True, + "completion_fim": False, + "fine_tuning": True, + "vision": True, + "ocr": False, + "classification": False, + "moderation": False, + "audio": False, + }, + name="mistral-medium-2508", + description="Update on Mistral Medium 3 with improved capabilities.", + max_context_length=131072, + aliases=["mistral-medium-latest"], + default_model_temperature=0.3, + ) + embeddings = factory.Trait( + id="mistral-embed-2312", + capabilities={ + "completion_chat": False, + "function_calling": False, + "completion_fim": False, + "fine_tuning": False, + "vision": False, + "ocr": False, + "classification": False, + "moderation": False, + "audio": False, + }, + name="mistral-embed-2312", + description="Our state-of-the-art semantic for extracting representation of text extracts.", + max_context_length=8192, + aliases=["mistral-embed-2312", "mistral-embed-latest"], + default_model_temperature=None, + ) + ocr = factory.Trait( + id="mistral-ocr-2512", + capabilities={ + "completion_chat": False, + "function_calling": True, + "completion_fim": False, + "fine_tuning": False, + "vision": True, + "ocr": True, + "classification": False, + "moderation": False, + "audio": False, + }, + name="mistral-ocr-2512", + description="Official mistral-ocr-2512 Mistral AI model", + max_context_length=16384, + aliases=["mistral-ocr-latest"], + default_model_temperature=0.0, + ) + + +class MistralModelsResponseFactory(factory.DictFactory): + object = "list" + data = factory.LazyFunction( + lambda: [ + MistralModelsItemResponseFactory(medium=True), + MistralModelsItemResponseFactory(embeddings=True), + MistralModelsItemResponseFactory(ocr=True), + ] + ) + + class Params: + text_generation = factory.Trait( + data=factory.LazyFunction( + lambda: [ + MistralModelsItemResponseFactory(medium=True), + ] + ) + ) + embeddings = factory.Trait( + data=factory.LazyFunction( + lambda: [ + MistralModelsItemResponseFactory(embeddings=True), + ] + ) + ) + ocr = factory.Trait( + data=factory.LazyFunction( + lambda: [ + MistralModelsItemResponseFactory(ocr=True), + ] + ) + ) + + +class MistralAudioTranscriptionResponseFactory(factory.DictFactory): + object = "chat.completion" + id = factory.Faker("bothify", text="chatcmpl-????????") + created = factory.Faker("random_int", min=1, max=9999999999) + model = factory.Faker("bothify", text="mistral-????") + choices = factory.LazyFunction( + lambda: [ + { + "index": 0, + "message": {"role": "assistant", "tool_calls": None, "content": fake.sentence()}, + "finish_reason": "stop", + } + ] + ) + usage = factory.LazyFunction(lambda: {"prompt_tokens": 10, "completion_tokens": 10, "total_tokens": 20}) + + +# vLLM factories +class VllmModelsResponseFactory(factory.DictFactory): + object = factory.Faker("random_element", elements=["list"]) + data = [ + { + "id": "openai/gpt-oss-120b", + "object": "model", + "created": 1773657692, + "owned_by": "vllm", + "root": "openai/gpt-oss-120b", + "parent": None, + "max_model_len": 131072, + "permission": [ + { + "id": "modelperm-aa05efd5693dcdcf", + "object": "model_permission", + "created": 1773657692, + "allow_create_engine": False, + "allow_sampling": True, + "allow_logprobs": True, + "allow_search_indices": False, + "allow_view": True, + "allow_fine_tuning": False, + "organization": "*", + "group": None, + "is_blocking": False, + } + ], + }, + ] diff --git a/api/tests/unit/infrastructure/http/test_mistralmodelhttpclient.py b/api/tests/unit/infrastructure/http/test_mistralmodelhttpclient.py new file mode 100644 index 000000000..feaff3af7 --- /dev/null +++ b/api/tests/unit/infrastructure/http/test_mistralmodelhttpclient.py @@ -0,0 +1,199 @@ +import base64 +from copy import deepcopy + +import pytest + +from api.domain.provider.entities import ProviderCarbonFootprintZone +from api.infrastructure.fastapi.schemas.models import ModelResponse, ModelsResponse +from api.infrastructure.http.model import MistralModelHttpClient +from api.schemas.audio import AudioTranscription +from api.schemas.core.models import RequestContent +from api.schemas.usage import Usage +from api.tests.unit.infrastructure.http.factories import ( + FormattedRequestContentFactory, + MistralAudioTranscriptionResponseFactory, + MistralModelsResponseFactory, +) +from api.utils.variables import EndpointRoute + + +@pytest.fixture +def mistral_model_http_client(): + return MistralModelHttpClient( + url="https://test.com", + key="test-key", + timeout=120, + model_name="mistral-test-model", + model_hosting_zone=ProviderCarbonFootprintZone.WOR, + model_total_params=10, + model_active_params=10, + ) + + +@pytest.fixture +def chat_completion_request_content(mistral_model_http_client): + return FormattedRequestContentFactory(chat_completions=True, model=mistral_model_http_client.model_name) + + +@pytest.fixture +def audio_transcription_request_content(mistral_model_http_client): + return FormattedRequestContentFactory(audio_transcriptions=True, model=mistral_model_http_client.model_name) + + +@pytest.fixture +def mistral_models_request_content(mistral_model_http_client): + return FormattedRequestContentFactory(models=True, model=mistral_model_http_client.model_name) + + +@pytest.fixture +def mistral_models_response_data(): + return MistralModelsResponseFactory() + + +@pytest.fixture +def audio_transcription_response_request_content(mistral_model_http_client): + return RequestContent( + method="POST", + model=mistral_model_http_client.model_name, + endpoint=EndpointRoute.AUDIO_TRANSCRIPTIONS, + additional_data={ + "id": "request-1234567890", + "model": mistral_model_http_client.model_name, + "usage": Usage(prompt_tokens=10, completion_tokens=10, total_tokens=20).model_dump(), + }, + ) + + +@pytest.fixture +def mistral_audio_transcription_response_data(): + return MistralAudioTranscriptionResponseFactory() + + +class TestMistralModelHttpClient: + def test_should_format_valid_chat_completion_request(self, mistral_model_http_client, chat_completion_request_content): + request_content = chat_completion_request_content.model_copy(deep=True) + + result = mistral_model_http_client.format_chat_completion_request(request_content) + + assert result.body == { + "frequency_penalty": 0.0, + "max_tokens": None, + "messages": chat_completion_request_content.body.get("messages"), + "model": chat_completion_request_content.body.get("model"), + "n": 1, + "parallel_tool_calls": False, + "prediction": {}, + "presence_penalty": 0.0, + "prompt_mode": None, + "random_seed": 10, + "response_format": { + "type": "json_schema", + "json_schema": { + "name": "DummyResponseFormat", + "schema": { + "properties": { + "dummy_list": {"items": {"type": "string"}, "title": "Dummy List", "type": "array"}, + "dummy_str": {"title": "Dummy Str", "type": "string"}, + "dummy_optional_bool": {"default": False, "title": "Dummy Optional Bool", "type": "boolean"}, + "dummy_nullable_int": {"anyOf": [{"type": "integer"}, {"type": "null"}], "title": "Dummy Nullable Int"}, + }, + "required": ["dummy_list", "dummy_str", "dummy_nullable_int"], + "title": "DummyResponseFormat", + "type": "object", + }, + }, + }, + "safe_prompt": False, + "stop": [], + "stream": False, + "temperature": None, + "tool_choice": "required", + "tools": None, + "top_p": 1.0, + } + + def test_should_format_audio_transcription_request(self, mistral_model_http_client, audio_transcription_request_content): + request_content = audio_transcription_request_content.model_copy(deep=True) + + result = mistral_model_http_client.format_audio_transcription_request(request_content) + + expected_audio = base64.b64encode(audio_transcription_request_content.files["file"][1]).decode("utf-8") + + assert result.form == {} + assert result.files == {} + assert result.additional_data == {} + assert result.body == { + "model": audio_transcription_request_content.form.get("model"), + "temperature": audio_transcription_request_content.form["temperature"], + "messages": [ + { + "role": "user", + "content": [ + {"type": "input_audio", "input_audio": expected_audio}, + {"type": "text", "text": audio_transcription_request_content.form["prompt"]}, + ], + }, + ], + "stream": False, + } + + def test_should_format_response_to_models_format( + self, + mistral_model_http_client, + mistral_models_request_content, + mistral_models_response_data, + ): + request_content = mistral_models_request_content.model_copy(deep=True) + response_data = deepcopy(mistral_models_response_data) + + result = mistral_model_http_client.format_response_to_models_response(request_content=request_content, response_data=response_data) + + assert result == ModelsResponse( + data=[ + ModelResponse( + id="mistral-medium-2508", + type=None, + aliases=["mistral-medium-latest"], + created=1773667856, + owned_by="mistralai", + max_context_length=131072, + ), + ModelResponse( + id="mistral-embed-2312", + type=None, + aliases=["mistral-embed-2312", "mistral-embed-latest"], + created=1773667856, + owned_by="mistralai", + max_context_length=8192, + ), + ModelResponse( + id="mistral-ocr-2512", + type=None, + aliases=["mistral-ocr-latest"], + created=1773667856, + owned_by="mistralai", + max_context_length=16384, + ), + ] + ) + + def test_should_format_response_to_audio_transcription_format( + self, + mistral_model_http_client, + audio_transcription_response_request_content, + mistral_audio_transcription_response_data, + ): + request_content = audio_transcription_response_request_content.model_copy(deep=True) + response_data = mistral_audio_transcription_response_data + + result = mistral_model_http_client.format_response_to_audio_transcription_response( + request_content=request_content, + response_data=response_data, + ) + + assert result == AudioTranscription( + id=response_data["id"], + model=response_data["model"], + text=response_data["choices"][0]["message"]["content"], + usage={"prompt_tokens": 10, "completion_tokens": 10, "total_tokens": 20}, + ) diff --git a/api/tests/unit/infrastructure/http/test_teimodelhttpclient.py b/api/tests/unit/infrastructure/http/test_teimodelhttpclient.py new file mode 100644 index 000000000..4f32f1425 --- /dev/null +++ b/api/tests/unit/infrastructure/http/test_teimodelhttpclient.py @@ -0,0 +1,129 @@ +from copy import deepcopy + +import pytest + +from api.domain.provider.entities import ProviderCarbonFootprintZone +from api.infrastructure.fastapi.schemas.models import ModelResponse, ModelsResponse +from api.infrastructure.http.model import TeiModelHttpClient +from api.schemas.core.models import RequestContent +from api.schemas.rerank import RerankResult, Reranks +from api.tests.unit.infrastructure.http.factories import ( + FormattedRequestContentFactory, + TeiFormattedRequestContentFactory, + TeiModelsResponseFactory, + TeiRerankResponseFactory, +) +from api.utils.variables import EndpointRoute + + +@pytest.fixture +def tei_model_http_client(): + return TeiModelHttpClient( + url="https://test.com", + key="test-key", + timeout=120, + model_name="test-model", + model_hosting_zone=ProviderCarbonFootprintZone.WOR, + model_total_params=10, + model_active_params=10, + ) + + +@pytest.fixture +def rerank_request_content(): + return FormattedRequestContentFactory(rerank=True) + + +@pytest.fixture +def tei_rerank_request_content(tei_model_http_client): + return TeiFormattedRequestContentFactory(rerank=True, model=tei_model_http_client.model_name) + + +@pytest.fixture +def tei_models_request_content(tei_model_http_client): + return TeiFormattedRequestContentFactory(models=True, model=tei_model_http_client.model_name) + + +@pytest.fixture +def tei_rerank_response_data(tei_rerank_request_content: TeiFormattedRequestContentFactory): + return TeiRerankResponseFactory(request_content=tei_rerank_request_content, raw_scores=False, return_text=False) + + +@pytest.fixture +def tei_models_response_data(): + return TeiModelsResponseFactory(embedding=True) + + +class TestTeiModelHttpClient: + @pytest.mark.asyncio + async def test_should_format_rerank_request(self, tei_model_http_client, rerank_request_content): + # Arrange + request_content = rerank_request_content.model_copy(deep=True) + + # Act + result = tei_model_http_client.format_rerank_request(request_content) + + # Assert + assert result == RequestContent( + method="POST", + model=rerank_request_content.model, + endpoint=EndpointRoute.RERANK, + body={ + "query": rerank_request_content.body["query"], + "texts": rerank_request_content.body["documents"], + "raw_scores": False, + "return_text": False, + "truncate": False, + "truncation_direction": "right", + }, + form={}, + files={}, + additional_data={"top_n": rerank_request_content.body["top_n"]}, + ) + + @pytest.mark.asyncio + async def test_should_format_rerank_response_to_reranks_format( + self, + tei_model_http_client, + tei_rerank_request_content, + tei_rerank_response_data, + ): + # Arrange + request_content = tei_rerank_request_content.model_copy(deep=True) + response_data = deepcopy(tei_rerank_response_data) + + # Act + result = tei_model_http_client.format_rerank_response(request_content=request_content, response_data=response_data) + + # Assert + expected_results = [RerankResult(relevance_score=x["score"], index=x["index"]) for x in tei_rerank_response_data] + expected_results = sorted(expected_results, key=lambda x: x.relevance_score, reverse=True) + expected_results = expected_results[: tei_rerank_request_content.additional_data["top_n"]] + + assert result == Reranks( + id=tei_rerank_request_content.additional_data["id"], + model=tei_model_http_client.model_name, + results=expected_results, + usage=tei_rerank_request_content.additional_data["usage"], + ) + + @pytest.mark.asyncio + async def test_should_format_response_to_models_format(self, tei_model_http_client, tei_models_request_content, tei_models_response_data): + # Arrange + request_content = tei_models_request_content.model_copy(deep=True) + response_data = deepcopy(tei_models_response_data) + + # Act + result = tei_model_http_client.format_response_to_models_response(request_content=request_content, response_data=response_data) + + # Assert + assert result == ModelsResponse( + data=[ + ModelResponse( + id=tei_models_response_data["model_id"], + created=0, + owned_by="tei", + max_context_length=tei_models_response_data["max_input_length"], + ) + ] + ) diff --git a/api/tests/unit/infrastructure/http/test_vllmmodelhttpclient.py b/api/tests/unit/infrastructure/http/test_vllmmodelhttpclient.py new file mode 100644 index 000000000..06ca57980 --- /dev/null +++ b/api/tests/unit/infrastructure/http/test_vllmmodelhttpclient.py @@ -0,0 +1,53 @@ +from copy import deepcopy + +import pytest + +from api.domain.provider.entities import ProviderCarbonFootprintZone +from api.infrastructure.fastapi.schemas.models import ModelResponse, ModelsResponse +from api.infrastructure.http.model import VllmModelHttpClient +from api.tests.unit.infrastructure.http.factories import FormattedRequestContentFactory, VllmModelsResponseFactory + + +@pytest.fixture +def vllm_model_http_client(): + return VllmModelHttpClient( + url="https://test.com", + key="test-key", + timeout=120, + model_name="vllm-test-model", + model_hosting_zone=ProviderCarbonFootprintZone.WOR, + model_total_params=10, + model_active_params=10, + ) + + +@pytest.fixture +def vllm_models_request_content(vllm_model_http_client): + return FormattedRequestContentFactory(models=True, model=vllm_model_http_client.model_name) + + +@pytest.fixture +def vllm_models_response_data(): + return VllmModelsResponseFactory() + + +class TestVllmModelHttpClient: + def test_should_format_response_to_models_format(self, vllm_model_http_client, vllm_models_request_content, vllm_models_response_data): + request_content = vllm_models_request_content.model_copy(deep=True) + response_data = deepcopy(vllm_models_response_data) + + result = vllm_model_http_client.format_response_to_models_response(request_content=request_content, response_data=response_data) + + assert result == ModelsResponse( + data=[ + ModelResponse( + object="model", + type=None, + aliases=[], + id=vllm_models_response_data["data"][0]["id"], + created=vllm_models_response_data["data"][0]["created"], + owned_by=vllm_models_response_data["data"][0]["owned_by"], + max_context_length=vllm_models_response_data["data"][0]["max_model_len"], + ) + ] + ) diff --git a/api/tests/unit/test_helpers/test_modelregistry/test_providers.py b/api/tests/unit/test_helpers/test_modelregistry/test_providers.py deleted file mode 100644 index 9aedf79ca..000000000 --- a/api/tests/unit/test_helpers/test_modelregistry/test_providers.py +++ /dev/null @@ -1,1195 +0,0 @@ -from unittest.mock import AsyncMock, MagicMock, patch - -import pytest -from sqlalchemy.exc import IntegrityError, NoResultFound -from sqlalchemy.ext.asyncio import AsyncSession - -import api.helpers.models._modelregistry as modelregistry_module -from api.helpers.models._modelregistry import ModelRegistry -from api.schemas.admin.providers import Provider, ProviderCarbonFootprintZone, ProviderType -from api.schemas.admin.routers import Router, RouterLoadBalancingStrategy -from api.schemas.core.models import Metric -from api.schemas.models import ModelType -from api.utils.exceptions import ( - InconsistentModelMaxContextLengthException, - InconsistentModelVectorSizeException, - InvalidProviderTypeException, - ProviderAlreadyExistsException, - ProviderNotFoundException, - ProviderNotReachableException, -) - - -class _MappingsResult: - def __init__(self, rows): - self._rows = rows - - def all(self): - return self._rows - - -class _Result: - def __init__(self, scalar_one=None, all_rows=None, iterate_rows=None, mappings=None): - self._scalar_one = scalar_one - self._all_rows = all_rows - self._iterate_rows = iterate_rows - self._mappings = mappings - - def scalar_one(self): - if isinstance(self._scalar_one, Exception): - raise self._scalar_one - return self._scalar_one - - def scalar_one_or_none(self): - if isinstance(self._scalar_one, Exception): - raise self._scalar_one - return self._scalar_one - - def all(self): - return self._all_rows or [] - - def mappings(self): - return _MappingsResult(self._mappings or []) - - def scalars(self): - return self # For scalars().all() - - def __iter__(self): - return iter(self._iterate_rows or []) - - -@pytest.fixture -def postgres_session(): - return AsyncMock(spec=AsyncSession) - - -@pytest.fixture -def model_registry(): - return ModelRegistry( - app_title="TestApp", - queuing_enabled=False, - max_priority=10, - max_retries=3, - retry_countdown=60, - ) - - -@pytest.fixture -def mock_router(): - return Router( - id=1, - name="test-router", - user_id=1, - type=ModelType.TEXT_GENERATION, - aliases=[], - load_balancing_strategy=RouterLoadBalancingStrategy.SHUFFLE, - vector_size=None, - max_context_length=4096, - cost_prompt_tokens=1.0, - cost_completion_tokens=2.0, - providers=0, - created=100, - updated=200, - ) - - -@pytest.fixture -def mock_provider_class(): - """Mock BaseModelProvider class""" - provider_instance = AsyncMock() - provider_instance.get_max_context_length = AsyncMock(return_value=4096) - provider_instance.get_vector_size = AsyncMock(return_value=None) - provider_instance.name = "test-model" - - provider_class = MagicMock(return_value=provider_instance) - return provider_class, provider_instance - - -@pytest.mark.asyncio -async def test_create_provider_success( - postgres_session: AsyncSession, - model_registry: ModelRegistry, - mock_router: Router, - mock_provider_class, -): - provider_class, provider_instance = mock_provider_class - - model_registry.get_routers = AsyncMock(return_value=[mock_router]) - postgres_session.execute = AsyncMock( - side_effect=[ - _Result(scalar_one=123), # insert provider, returning id - ] - ) - - with patch.object(modelregistry_module.ModelProvider, "import_module", return_value=provider_class): - provider_id = await model_registry.create_provider( - router_id=1, - user_id=1, - type=ProviderType.OPENAI, - url="https://api.openai.com/", - key="sk-test", - timeout=30, - model_name="gpt-4", - model_hosting_zone=ProviderCarbonFootprintZone.WOR, - model_total_params=0, - model_active_params=0, - qos_metric=None, - qos_limit=None, - postgres_session=postgres_session, - ) - - assert provider_id == 123 - postgres_session.commit.assert_awaited() - provider_instance.get_max_context_length.assert_awaited_once() - - -@pytest.mark.asyncio -async def test_create_provider_with_vector_size( - postgres_session: AsyncSession, - model_registry: ModelRegistry, - mock_provider_class, -): - provider_class, provider_instance = mock_provider_class - provider_instance.get_max_context_length = AsyncMock(return_value=2048) - provider_instance.get_vector_size = AsyncMock(return_value=768) - - router = Router( - id=1, - name="embedding-router", - user_id=1, - type=ModelType.TEXT_EMBEDDINGS_INFERENCE, - aliases=[], - load_balancing_strategy=RouterLoadBalancingStrategy.SHUFFLE, - vector_size=None, - max_context_length=2048, - cost_prompt_tokens=0.0, - cost_completion_tokens=0.0, - providers=0, - created=100, - updated=200, - ) - - model_registry.get_routers = AsyncMock(return_value=[router]) - postgres_session.execute = AsyncMock(return_value=_Result(scalar_one=456)) - - with patch.object(modelregistry_module.ModelProvider, "import_module", return_value=provider_class): - provider_id = await model_registry.create_provider( - router_id=1, - user_id=1, - type=ProviderType.TEI, - url="https://tei.example.com/", - key=None, - timeout=30, - model_name="embedding-model", - model_hosting_zone=ProviderCarbonFootprintZone.WOR, - model_total_params=0, - model_active_params=0, - qos_metric=None, - qos_limit=None, - postgres_session=postgres_session, - ) - - assert provider_id == 456 - provider_instance.get_max_context_length.assert_awaited_once() - provider_instance.get_vector_size.assert_awaited_once() - - -@pytest.mark.asyncio -async def test_create_provider_invalid_type( - postgres_session: AsyncSession, - model_registry: ModelRegistry, - mock_router: Router, -): - model_registry.get_routers = AsyncMock(return_value=[mock_router]) - - with pytest.raises(InvalidProviderTypeException): - await model_registry.create_provider( - router_id=1, - user_id=1, - type=ProviderType.TEI, # TEI not compatible with TEXT_GENERATION - url="https://tei.example.com/", - key=None, - timeout=30, - model_name="test-model", - model_hosting_zone=ProviderCarbonFootprintZone.WOR, - model_total_params=0, - model_active_params=0, - qos_metric=None, - qos_limit=None, - postgres_session=postgres_session, - ) - - -@pytest.mark.asyncio -async def test_create_provider_not_reachable( - postgres_session: AsyncSession, - model_registry: ModelRegistry, - mock_router: Router, - mock_provider_class, -): - provider_class, provider_instance = mock_provider_class - provider_instance.get_max_context_length = AsyncMock(side_effect=AssertionError("Connection failed")) - - model_registry.get_routers = AsyncMock(return_value=[mock_router]) - - with patch.object(modelregistry_module.ModelProvider, "import_module", return_value=provider_class): - with pytest.raises(ProviderNotReachableException): - await model_registry.create_provider( - router_id=1, - user_id=1, - type=ProviderType.OPENAI, - url="https://invalid-url.com/", - key="sk-test", - timeout=30, - model_name="gpt-4", - model_hosting_zone=ProviderCarbonFootprintZone.WOR, - model_total_params=0, - model_active_params=0, - qos_metric=None, - qos_limit=None, - postgres_session=postgres_session, - ) - - -@pytest.mark.asyncio -async def test_create_provider_inconsistent_vector_size( - postgres_session: AsyncSession, - model_registry: ModelRegistry, - mock_provider_class, -): - provider_class, provider_instance = mock_provider_class - provider_instance.get_max_context_length = AsyncMock(return_value=2048) - provider_instance.get_vector_size = AsyncMock(return_value=512) - - router = Router( - id=1, - name="embedding-router", - user_id=1, - type=ModelType.TEXT_EMBEDDINGS_INFERENCE, - aliases=[], - load_balancing_strategy=RouterLoadBalancingStrategy.SHUFFLE, - vector_size=768, # Different from provider - max_context_length=2048, - cost_prompt_tokens=0.0, - cost_completion_tokens=0.0, - providers=1, # Has existing providers - created=100, - updated=200, - ) - - model_registry.get_routers = AsyncMock(return_value=[router]) - - with patch.object(modelregistry_module.ModelProvider, "import_module", return_value=provider_class): - with pytest.raises(InconsistentModelVectorSizeException): - await model_registry.create_provider( - router_id=1, - user_id=1, - type=ProviderType.TEI, - url="https://tei.example.com/", - key=None, - timeout=30, - model_name="embedding-model", - model_hosting_zone=ProviderCarbonFootprintZone.WOR, - model_total_params=0, - model_active_params=0, - qos_metric=None, - qos_limit=None, - postgres_session=postgres_session, - ) - - -@pytest.mark.asyncio -async def test_create_provider_inconsistent_max_context_length( - postgres_session: AsyncSession, - model_registry: ModelRegistry, - mock_router: Router, - mock_provider_class, -): - provider_class, provider_instance = mock_provider_class - provider_instance.get_max_context_length = AsyncMock(return_value=2048) # Different from router - - router_with_providers = Router( - id=1, - name="test-router", - user_id=1, - type=ModelType.TEXT_GENERATION, - aliases=[], - load_balancing_strategy=RouterLoadBalancingStrategy.SHUFFLE, - vector_size=None, - max_context_length=4096, # Different from provider - cost_prompt_tokens=1.0, - cost_completion_tokens=2.0, - providers=1, # Has existing providers - created=100, - updated=200, - ) - - model_registry.get_routers = AsyncMock(return_value=[router_with_providers]) - - with patch.object(modelregistry_module.ModelProvider, "import_module", return_value=provider_class): - with pytest.raises(InconsistentModelMaxContextLengthException): - await model_registry.create_provider( - router_id=1, - user_id=1, - type=ProviderType.OPENAI, - url="https://api.openai.com/", - key="sk-test", - timeout=30, - model_name="gpt-4", - model_hosting_zone=ProviderCarbonFootprintZone.WOR, - model_total_params=0, - model_active_params=0, - qos_metric=None, - qos_limit=None, - postgres_session=postgres_session, - ) - - -@pytest.mark.asyncio -async def test_create_provider_already_exists( - postgres_session: AsyncSession, - model_registry: ModelRegistry, - mock_router: Router, - mock_provider_class, -): - provider_class, provider_instance = mock_provider_class - - model_registry.get_routers = AsyncMock(return_value=[mock_router]) - postgres_session.execute = AsyncMock(side_effect=IntegrityError("", "", None)) - - with patch.object(modelregistry_module.ModelProvider, "import_module", return_value=provider_class): - with pytest.raises(ProviderAlreadyExistsException): - await model_registry.create_provider( - router_id=1, - user_id=1, - type=ProviderType.OPENAI, - url="https://api.openai.com/", - key="sk-test", - timeout=30, - model_name="gpt-4", - model_hosting_zone=ProviderCarbonFootprintZone.WOR, - model_total_params=0, - model_active_params=0, - qos_metric=None, - qos_limit=None, - postgres_session=postgres_session, - ) - - -@pytest.mark.asyncio -async def test_create_provider_master_user( - postgres_session: AsyncSession, - model_registry: ModelRegistry, - mock_router: Router, - mock_provider_class, -): - provider_class, provider_instance = mock_provider_class - - model_registry.get_routers = AsyncMock(return_value=[mock_router]) - postgres_session.execute = AsyncMock(return_value=_Result(scalar_one=789)) - - with patch.object(modelregistry_module.ModelProvider, "import_module", return_value=provider_class): - provider_id = await model_registry.create_provider( - router_id=1, - user_id=0, # master user - type=ProviderType.OPENAI, - url="https://api.openai.com/", - key="sk-test", - timeout=30, - model_name="gpt-4", - model_hosting_zone=ProviderCarbonFootprintZone.WOR, - model_total_params=0, - model_active_params=0, - qos_metric=None, - qos_limit=None, - postgres_session=postgres_session, - ) - - assert provider_id == 789 - - -@pytest.mark.asyncio -async def test_create_provider_with_qos( - postgres_session: AsyncSession, - model_registry: ModelRegistry, - mock_router: Router, - mock_provider_class, -): - provider_class, provider_instance = mock_provider_class - - model_registry.get_routers = AsyncMock(return_value=[mock_router]) - postgres_session.execute = AsyncMock(return_value=_Result(scalar_one=999)) - - with patch.object(modelregistry_module.ModelProvider, "import_module", return_value=provider_class): - provider_id = await model_registry.create_provider( - router_id=1, - user_id=1, - type=ProviderType.OPENAI, - url="https://api.openai.com/", - key="sk-test", - timeout=30, - model_name="gpt-4", - model_hosting_zone=ProviderCarbonFootprintZone.WOR, - model_total_params=0, - model_active_params=0, - qos_metric=Metric.TTFT, - qos_limit=0.5, - postgres_session=postgres_session, - ) - - assert provider_id == 999 - - -@pytest.mark.asyncio -async def test_delete_provider_not_found(postgres_session: AsyncSession, model_registry: ModelRegistry): - postgres_session.execute = AsyncMock(return_value=_Result(scalar_one=NoResultFound())) - - with pytest.raises(ProviderNotFoundException): - await model_registry.delete_provider(provider_id=999, postgres_session=postgres_session) - - -@pytest.mark.asyncio -async def test_delete_provider_success(postgres_session: AsyncSession, model_registry: ModelRegistry): - postgres_session.execute = AsyncMock(side_effect=[_Result(scalar_one=MagicMock(id=1)), None]) - - await model_registry.delete_provider(provider_id=1, postgres_session=postgres_session) - postgres_session.commit.assert_awaited() - - -@pytest.mark.asyncio -async def test_get_providers_by_router_id(postgres_session: AsyncSession, model_registry: ModelRegistry): - provider_rows = [ - { - "id": 1, - "router_id": 1, - "user_id": 1, - "type": ProviderType.OPENAI.value, - "url": "https://api.openai.com/", - "key": "sk-test", - "timeout": 30, - "model_name": "gpt-4", - "model_hosting_zone": ProviderCarbonFootprintZone.WOR.value, - "model_total_params": 0, - "model_active_params": 0, - "qos_metric": None, - "qos_limit": None, - "created": 100, - "updated": 200, - }, - { - "id": 2, - "router_id": 1, - "user_id": 1, - "type": ProviderType.VLLM.value, - "url": "https://vllm.example.com/", - "key": None, - "timeout": 60, - "model_name": "llama-2", - "model_hosting_zone": ProviderCarbonFootprintZone.WOR.value, - "model_total_params": 0, - "model_active_params": 0, - "qos_metric": Metric.TTFT.value, - "qos_limit": 0.5, - "created": 200, - "updated": 300, - }, - ] - - result = _Result() - result.mappings = lambda: _MappingsResult(provider_rows) - postgres_session.execute = AsyncMock(return_value=result) - - providers = await model_registry.get_providers( - router_id=1, - provider_id=None, - postgres_session=postgres_session, - ) - - assert len(providers) == 2 - assert providers[0].id == 1 - assert providers[0].type == ProviderType.OPENAI.value - assert providers[1].id == 2 - assert providers[1].qos_metric == Metric.TTFT - - -@pytest.mark.asyncio -async def test_get_providers_by_provider_id(postgres_session: AsyncSession, model_registry: ModelRegistry): - provider_row = { - "id": 42, - "router_id": 1, - "user_id": 1, - "type": ProviderType.OPENAI.value, - "url": "https://api.openai.com/", - "key": "sk-test", - "timeout": 30, - "model_name": "gpt-4", - "model_hosting_zone": ProviderCarbonFootprintZone.WOR.value, - "model_total_params": 0, - "model_active_params": 0, - "qos_metric": None, - "qos_limit": None, - "created": 100, - "updated": 200, - } - - result = _Result() - result.mappings = lambda: _MappingsResult([provider_row]) - postgres_session.execute = AsyncMock(return_value=result) - - providers = await model_registry.get_providers( - router_id=None, - provider_id=42, - postgres_session=postgres_session, - ) - - assert len(providers) == 1 - assert providers[0].id == 42 - - -@pytest.mark.asyncio -async def test_get_providers_not_found(postgres_session: AsyncSession, model_registry: ModelRegistry): - result = _Result() - result.mappings = lambda: _MappingsResult([]) - postgres_session.execute = AsyncMock(return_value=result) - - with pytest.raises(ProviderNotFoundException): - await model_registry.get_providers( - router_id=None, - provider_id=999, - postgres_session=postgres_session, - ) - - -@pytest.mark.asyncio -async def test_get_providers_master_user(postgres_session: AsyncSession, model_registry: ModelRegistry): - provider_row = { - "id": 1, - "router_id": 1, - "user_id": None, # master user - "type": ProviderType.OPENAI.value, - "url": "https://api.openai.com/", - "key": "sk-test", - "timeout": 30, - "model_name": "gpt-4", - "model_hosting_zone": ProviderCarbonFootprintZone.WOR.value, - "model_total_params": 0, - "model_active_params": 0, - "qos_metric": None, - "qos_limit": None, - "created": 100, - "updated": 200, - } - - result = _Result() - result.mappings = lambda: _MappingsResult([provider_row]) - postgres_session.execute = AsyncMock(return_value=result) - - providers = await model_registry.get_providers( - router_id=1, - provider_id=None, - postgres_session=postgres_session, - ) - - assert providers[0].user_id == 0 - - -@pytest.mark.asyncio -async def test_get_providers_pagination_and_ordering(postgres_session: AsyncSession, model_registry: ModelRegistry): - provider_rows = [ - { - "id": 1, - "router_id": 1, - "user_id": 1, - "type": ProviderType.OPENAI.value, - "url": "https://api.openai.com/", - "key": "sk-test", - "timeout": 30, - "model_name": "gpt-4", - "model_hosting_zone": ProviderCarbonFootprintZone.WOR.value, - "model_total_params": 0, - "model_active_params": 0, - "qos_metric": None, - "qos_limit": None, - "created": 100, - "updated": 200, - }, - { - "id": 2, - "router_id": 1, - "user_id": 1, - "type": ProviderType.VLLM.value, - "url": "https://vllm.example.com/", - "key": None, - "timeout": 60, - "model_name": "llama-2", - "model_hosting_zone": ProviderCarbonFootprintZone.WOR.value, - "model_total_params": 0, - "model_active_params": 0, - "qos_metric": None, - "qos_limit": None, - "created": 200, - "updated": 300, - }, - ] - - result = _Result() - result.mappings = lambda: _MappingsResult(provider_rows) - postgres_session.execute = AsyncMock(return_value=result) - - providers = await model_registry.get_providers( - router_id=1, - provider_id=None, - postgres_session=postgres_session, - offset=0, - limit=10, - order_by="created", - order_direction="desc", - ) - - assert len(providers) == 2 - - -@pytest.mark.asyncio -async def test_update_provider_not_found(postgres_session: AsyncSession, model_registry: ModelRegistry): - # get_providers with provider_id raises ProviderNotFoundException when empty - result = _Result() - result.mappings = lambda: _MappingsResult([]) - postgres_session.execute = AsyncMock(return_value=result) - - with pytest.raises(ProviderNotFoundException): - await model_registry.update_provider( - provider_id=999, - router_id=None, - timeout=None, - model_hosting_zone=ProviderCarbonFootprintZone.WOR, - model_total_params=0, - model_active_params=0, - qos_metric=None, - qos_limit=None, - postgres_session=postgres_session, - ) - - -@pytest.mark.asyncio -async def test_update_provider_timeout(postgres_session: AsyncSession, model_registry: ModelRegistry): - provider = Provider( - id=1, - router_id=1, - user_id=1, - type=ProviderType.OPENAI.value, - url="https://api.openai.com/", - key="sk-test", - timeout=30, - model_name="gpt-4", - model_hosting_zone=ProviderCarbonFootprintZone.WOR, - model_total_params=0, - model_active_params=0, - qos_metric=None, - qos_limit=None, - created=100, - updated=200, - ) - - router = Router( - id=1, - name="test-router", - user_id=1, - type=ModelType.TEXT_GENERATION, - aliases=[], - load_balancing_strategy=RouterLoadBalancingStrategy.SHUFFLE, - vector_size=None, - max_context_length=4096, - cost_prompt_tokens=0.0, - cost_completion_tokens=0.0, - providers=1, - created=100, - updated=200, - ) - - model_registry.get_providers = AsyncMock(return_value=[provider]) - model_registry.get_routers = AsyncMock(return_value=[router]) - postgres_session.execute = AsyncMock() - - await model_registry.update_provider( - provider_id=1, - router_id=None, - timeout=60, - model_hosting_zone=ProviderCarbonFootprintZone.WOR, - model_total_params=0, - model_active_params=0, - qos_metric=None, - qos_limit=None, - postgres_session=postgres_session, - ) - - postgres_session.commit.assert_awaited() - - -@pytest.mark.asyncio -async def test_update_provider_change_router_invalid_type(postgres_session: AsyncSession, model_registry: ModelRegistry): - current_router = Router( - id=1, - name="text-gen-router", - user_id=1, - type=ModelType.TEXT_GENERATION, - aliases=[], - load_balancing_strategy=RouterLoadBalancingStrategy.SHUFFLE, - vector_size=None, - max_context_length=4096, - cost_prompt_tokens=0.0, - cost_completion_tokens=0.0, - providers=1, - created=100, - updated=200, - ) - - new_router = Router( - id=2, - name="classification-router", - user_id=1, - type=ModelType.TEXT_CLASSIFICATION, - aliases=[], - load_balancing_strategy=RouterLoadBalancingStrategy.SHUFFLE, - vector_size=None, - max_context_length=4096, - cost_prompt_tokens=0.0, - cost_completion_tokens=0.0, - providers=0, - created=200, - updated=300, - ) - - # Mock get_providers to return provider (will be called with router_id=None, provider_id=1) - # Note: get_providers creates provider with type from DB (string), which Pydantic converts to enum - provider_result = _Result() - provider_result.mappings = lambda: _MappingsResult( - [ - { - "id": 1, - "router_id": 1, - "user_id": 1, - "type": ProviderType.MISTRAL.value, # MISTRAL is not compatible with TEXT_CLASSIFICATION - "url": "https://mistral.example.com/", - "key": "sk-test", - "timeout": 30, - "model_name": "mistral-model", - "model_hosting_zone": ProviderCarbonFootprintZone.WOR.value, - "model_total_params": 0, - "model_active_params": 0, - "qos_metric": None, - "qos_limit": None, - "created": 100, - "updated": 200, - } - ] - ) - postgres_session.execute = AsyncMock(return_value=provider_result) - model_registry.get_routers = AsyncMock(side_effect=[[current_router], [new_router]]) - - with pytest.raises(InvalidProviderTypeException): - await model_registry.update_provider( - provider_id=1, - router_id=2, - timeout=None, - model_hosting_zone=ProviderCarbonFootprintZone.WOR, - model_total_params=0, - model_active_params=0, - qos_metric=None, - qos_limit=None, - postgres_session=postgres_session, - ) - - -@pytest.mark.asyncio -async def test_update_provider_change_router_success(postgres_session: AsyncSession, model_registry: ModelRegistry): - provider = Provider( - id=1, - router_id=1, - user_id=1, - type=ProviderType.OPENAI.value, - url="https://api.openai.com/", - key="sk-test", - timeout=30, - model_name="gpt-4o-mini", - model_hosting_zone=ProviderCarbonFootprintZone.WOR, - model_total_params=0, - model_active_params=0, - qos_metric=None, - qos_limit=None, - created=100, - updated=200, - ) - - current_router = Router( - id=1, - name="router-1", - user_id=1, - type=ModelType.TEXT_GENERATION, - aliases=[], - load_balancing_strategy=RouterLoadBalancingStrategy.SHUFFLE, - vector_size=None, - max_context_length=4096, - cost_prompt_tokens=0.0, - cost_completion_tokens=0.0, - providers=1, - created=100, - updated=200, - ) - - new_router = Router( - id=2, - name="router-2", - user_id=1, - type=ModelType.TEXT_EMBEDDINGS_INFERENCE, - aliases=[], - load_balancing_strategy=RouterLoadBalancingStrategy.SHUFFLE, - vector_size=768, - max_context_length=2048, - cost_prompt_tokens=0.0, - cost_completion_tokens=0.0, - providers=0, # skip consistency checks - created=200, - updated=300, - ) - - model_registry.get_providers = AsyncMock(return_value=[provider]) - model_registry.get_routers = AsyncMock(side_effect=[[current_router], [new_router]]) - postgres_session.execute = AsyncMock() - - await model_registry.update_provider( - provider_id=1, - router_id=2, - timeout=None, - model_hosting_zone=ProviderCarbonFootprintZone.WOR, - model_total_params=0, - model_active_params=0, - qos_metric=None, - qos_limit=None, - postgres_session=postgres_session, - ) - - postgres_session.commit.assert_awaited() - - -@pytest.mark.asyncio -async def test_update_provider_change_router_inconsistent_vector_size(postgres_session: AsyncSession, model_registry: ModelRegistry): - provider = Provider( - id=1, - router_id=1, - user_id=1, - type=ProviderType.TEI.value, - url="https://tei.example.com/", - key=None, - timeout=30, - model_name="embedding-model", - model_hosting_zone=ProviderCarbonFootprintZone.WOR, - model_total_params=0, - model_active_params=0, - qos_metric=None, - qos_limit=None, - created=100, - updated=200, - ) - - current_router = Router( - id=1, - name="embedding-router-1", - user_id=1, - type=ModelType.TEXT_EMBEDDINGS_INFERENCE, - aliases=[], - load_balancing_strategy=RouterLoadBalancingStrategy.SHUFFLE, - vector_size=512, - max_context_length=2048, - cost_prompt_tokens=0.0, - cost_completion_tokens=0.0, - providers=1, - created=100, - updated=200, - ) - - new_router = Router( - id=2, - name="embedding-router-2", - user_id=1, - type=ModelType.TEXT_EMBEDDINGS_INFERENCE, - aliases=[], - load_balancing_strategy=RouterLoadBalancingStrategy.SHUFFLE, - vector_size=768, # Different vector size - max_context_length=2048, - cost_prompt_tokens=0.0, - cost_completion_tokens=0.0, - providers=1, # Has existing providers - created=200, - updated=300, - ) - - model_registry.get_providers = AsyncMock(return_value=[provider]) - model_registry.get_routers = AsyncMock(side_effect=[[current_router], [new_router]]) - - with pytest.raises(InconsistentModelVectorSizeException): - await model_registry.update_provider( - provider_id=1, - router_id=2, - timeout=None, - model_hosting_zone=ProviderCarbonFootprintZone.WOR, - model_total_params=0, - model_active_params=0, - qos_metric=None, - qos_limit=None, - postgres_session=postgres_session, - ) - - -@pytest.mark.asyncio -async def test_update_provider_change_router_inconsistent_max_context_length(postgres_session: AsyncSession, model_registry: ModelRegistry): - provider = Provider( - id=1, - router_id=1, - user_id=1, - type=ProviderType.OPENAI.value, - url="https://api.openai.com/", - key="sk-test", - timeout=30, - model_name="gpt-4", - model_hosting_zone=ProviderCarbonFootprintZone.WOR, - model_total_params=0, - model_active_params=0, - qos_metric=None, - qos_limit=None, - created=100, - updated=200, - ) - - current_router = Router( - id=1, - name="router-1", - user_id=1, - type=ModelType.TEXT_GENERATION, - aliases=[], - load_balancing_strategy=RouterLoadBalancingStrategy.SHUFFLE, - vector_size=None, - max_context_length=4096, - cost_prompt_tokens=0.0, - cost_completion_tokens=0.0, - providers=1, - created=100, - updated=200, - ) - - new_router = Router( - id=2, - name="router-2", - user_id=1, - type=ModelType.TEXT_GENERATION, - aliases=[], - load_balancing_strategy=RouterLoadBalancingStrategy.SHUFFLE, - vector_size=None, - max_context_length=2048, # Different max context length - cost_prompt_tokens=0.0, - cost_completion_tokens=0.0, - providers=1, # Has existing providers - created=200, - updated=300, - ) - - model_registry.get_providers = AsyncMock(return_value=[provider]) - model_registry.get_routers = AsyncMock(side_effect=[[current_router], [new_router]]) - - with pytest.raises(InconsistentModelMaxContextLengthException): - await model_registry.update_provider( - provider_id=1, - router_id=2, - timeout=None, - model_hosting_zone=ProviderCarbonFootprintZone.WOR, - model_total_params=0, - model_active_params=0, - qos_metric=None, - qos_limit=None, - postgres_session=postgres_session, - ) - - -@pytest.mark.asyncio -async def test_update_provider_change_router_already_exists(postgres_session: AsyncSession, model_registry: ModelRegistry): - provider = Provider( - id=1, - router_id=1, - user_id=1, - type=ProviderType.OPENAI.value, - url="https://api.openai.com/", - key="sk-test", - timeout=30, - model_name="gpt-4", - model_hosting_zone=ProviderCarbonFootprintZone.WOR, - model_total_params=0, - model_active_params=0, - qos_metric=None, - qos_limit=None, - created=100, - updated=200, - ) - - router = Router( - id=1, - name="test-router", - user_id=1, - type=ModelType.TEXT_GENERATION, - aliases=[], - load_balancing_strategy=RouterLoadBalancingStrategy.SHUFFLE, - vector_size=None, - max_context_length=4096, - cost_prompt_tokens=0.0, - cost_completion_tokens=0.0, - providers=1, - created=100, - updated=200, - ) - - model_registry.get_providers = AsyncMock(return_value=[provider]) - model_registry.get_routers = AsyncMock(return_value=[router]) - postgres_session.execute = AsyncMock(side_effect=IntegrityError("", "", None)) - - with pytest.raises(ProviderAlreadyExistsException): - await model_registry.update_provider( - provider_id=1, - router_id=2, - timeout=None, - model_hosting_zone=None, - model_total_params=None, - model_active_params=None, - qos_metric=None, - qos_limit=None, - postgres_session=postgres_session, - ) - - -@pytest.mark.asyncio -async def test_update_provider_qos_metric(postgres_session: AsyncSession, model_registry: ModelRegistry): - provider = Provider( - id=1, - router_id=1, - user_id=1, - type=ProviderType.OPENAI.value, - url="https://api.openai.com/", - key="sk-test", - timeout=30, - model_name="gpt-4", - model_hosting_zone=ProviderCarbonFootprintZone.WOR, - model_total_params=0, - model_active_params=0, - qos_metric=None, - qos_limit=None, - created=100, - updated=200, - ) - - router = Router( - id=1, - name="test-router", - user_id=1, - type=ModelType.TEXT_GENERATION, - aliases=[], - load_balancing_strategy=RouterLoadBalancingStrategy.SHUFFLE, - vector_size=None, - max_context_length=4096, - cost_prompt_tokens=0.0, - cost_completion_tokens=0.0, - providers=1, - created=100, - updated=200, - ) - - model_registry.get_providers = AsyncMock(return_value=[provider]) - model_registry.get_routers = AsyncMock(return_value=[router]) - postgres_session.execute = AsyncMock() - - await model_registry.update_provider( - provider_id=1, - router_id=None, - timeout=None, - model_hosting_zone=ProviderCarbonFootprintZone.WOR, - model_total_params=0, - model_active_params=0, - qos_metric=Metric.TTFT, - qos_limit=0.5, - postgres_session=postgres_session, - ) - - postgres_session.commit.assert_awaited() - - -@pytest.mark.asyncio -async def test_update_provider_noop(postgres_session: AsyncSession, model_registry: ModelRegistry): - provider = Provider( - id=1, - router_id=1, - user_id=1, - type=ProviderType.OPENAI.value, - url="https://api.openai.com/", - key="sk-test", - timeout=30, - model_name="gpt-4", - model_hosting_zone=ProviderCarbonFootprintZone.WOR.value, - model_total_params=0, - model_active_params=0, - qos_metric=None, - qos_limit=None, - created=100, - updated=200, - ) - - router = Router( - id=1, - name="test-router", - user_id=1, - type=ModelType.TEXT_GENERATION, - aliases=[], - load_balancing_strategy=RouterLoadBalancingStrategy.SHUFFLE, - vector_size=None, - max_context_length=4096, - cost_prompt_tokens=0.0, - cost_completion_tokens=0.0, - providers=1, - created=100, - updated=200, - ) - - result = _Result() - result.mappings = lambda: _MappingsResult( - [ - { - "id": 1, - "router_id": 1, - "user_id": 1, - "type": ProviderType.OPENAI.value, - "url": "https://api.openai.com/", - "key": "sk-test", - "timeout": 30, - "model_name": "gpt-4", - "model_hosting_zone": ProviderCarbonFootprintZone.WOR.value, - "model_total_params": 0, - "model_active_params": 0, - "qos_metric": None, - "qos_limit": None, - "created": 100, - "updated": 200, - } - ] - ) - postgres_session.execute = AsyncMock(return_value=result) - model_registry.get_routers = AsyncMock(return_value=[router]) - - await model_registry.update_provider( - provider_id=1, - router_id=None, - timeout=None, - model_hosting_zone=None, - model_total_params=None, - model_active_params=None, - qos_metric=None, - qos_limit=None, - postgres_session=postgres_session, - ) - - # Should NOT commit if no updates (update_value is empty) - postgres_session.commit.assert_not_awaited() diff --git a/api/tests/unit/test_helpers/test_modelregistry/test_routers.py b/api/tests/unit/test_helpers/test_modelregistry/test_routers.py deleted file mode 100644 index 565cddbdf..000000000 --- a/api/tests/unit/test_helpers/test_modelregistry/test_routers.py +++ /dev/null @@ -1,557 +0,0 @@ -from unittest.mock import AsyncMock, MagicMock - -import pytest -from sqlalchemy.exc import IntegrityError, NoResultFound -from sqlalchemy.ext.asyncio import AsyncSession - -from api.helpers.models._modelregistry import ModelRegistry -from api.schemas.admin.routers import Router, RouterLoadBalancingStrategy -from api.schemas.models import ModelType -from api.utils.exceptions import ( - RouterAliasAlreadyExistsException, - RouterAlreadyExistsException, - RouterNotFoundException, -) - - -class _Result: - def __init__(self, scalar_one=None, all_rows=None, iterate_rows=None, mappings=None, scalars_rows=None): - self._scalar_one = scalar_one - self._all_rows = all_rows - self._iterate_rows = iterate_rows - self._mappings = mappings - self._scalars_rows = scalars_rows - - def scalar_one(self): - if isinstance(self._scalar_one, Exception): - raise self._scalar_one - return self._scalar_one - - def scalar_one_or_none(self): - if isinstance(self._scalar_one, Exception): - raise self._scalar_one - return self._scalar_one - - def all(self): - return self._all_rows or [] - - def mappings(self): - return self._mappings or [] - - def scalars(self): - return self # For scalars().all() - - def __iter__(self): - return iter(self._iterate_rows or []) - - -@pytest.fixture -def postgres_session(): - return AsyncMock(spec=AsyncSession) - - -@pytest.fixture -def model_registry(): - return ModelRegistry( - app_title="TestApp", - queuing_enabled=False, - max_priority=10, - max_retries=3, - retry_countdown=60, - ) - - -@pytest.mark.asyncio -async def test_create_router_success(postgres_session: AsyncSession, model_registry: ModelRegistry): - # Router insert returns id, aliases check returns empty, aliases insert succeeds (one per alias) - postgres_session.execute = AsyncMock( - side_effect=[ - _Result(scalar_one=123), # insert router, returning id - _Result(all_rows=[]), # check aliases integrity - None, # insert alias1 - None, # insert alias2 - ] - ) - - router_id = await model_registry.create_router( - name="test-router", - type=ModelType.TEXT_GENERATION, - aliases=["alias1", "alias2"], - load_balancing_strategy=RouterLoadBalancingStrategy.SHUFFLE, - cost_prompt_tokens=1.0, - cost_completion_tokens=2.0, - user_id=1, - postgres_session=postgres_session, - ) - - assert router_id == 123 - postgres_session.commit.assert_awaited() - - -@pytest.mark.asyncio -async def test_create_router_master_user(postgres_session: AsyncSession, model_registry: ModelRegistry): - # Master user (user_id=0) should be converted to None - postgres_session.execute = AsyncMock( - side_effect=[ - _Result(scalar_one=456), # insert router - _Result(all_rows=[]), # check aliases - None, # insert aliases - ] - ) - - router_id = await model_registry.create_router( - name="master-router", - type=ModelType.TEXT_GENERATION, - aliases=[], - load_balancing_strategy=RouterLoadBalancingStrategy.LEAST_BUSY, - cost_prompt_tokens=0.0, - cost_completion_tokens=0.0, - user_id=0, # master user - postgres_session=postgres_session, - ) - - assert router_id == 456 - postgres_session.commit.assert_awaited() - - -@pytest.mark.asyncio -async def test_create_router_already_exists(postgres_session: AsyncSession, model_registry: ModelRegistry): - postgres_session.execute = AsyncMock(side_effect=IntegrityError("", "", None)) - - with pytest.raises(RouterAlreadyExistsException): - await model_registry.create_router( - name="existing-router", - type=ModelType.TEXT_GENERATION, - aliases=[], - load_balancing_strategy=RouterLoadBalancingStrategy.SHUFFLE, - cost_prompt_tokens=0.0, - cost_completion_tokens=0.0, - user_id=1, - postgres_session=postgres_session, - ) - - -@pytest.mark.asyncio -async def test_create_router_alias_already_exists(postgres_session: AsyncSession, model_registry: ModelRegistry): - postgres_session.execute = AsyncMock( - side_effect=[ - _Result(scalar_one=789), # insert router succeeds - _Result(all_rows=[("existing-alias",)]), # alias already exists - ] - ) - - with pytest.raises(RouterAliasAlreadyExistsException): - await model_registry.create_router( - name="new-router", - type=ModelType.TEXT_GENERATION, - aliases=["existing-alias"], - load_balancing_strategy=RouterLoadBalancingStrategy.SHUFFLE, - cost_prompt_tokens=0.0, - cost_completion_tokens=0.0, - user_id=1, - postgres_session=postgres_session, - ) - - -@pytest.mark.asyncio -async def test_delete_router_not_found(postgres_session: AsyncSession, model_registry: ModelRegistry): - postgres_session.execute = AsyncMock(return_value=_Result(scalar_one=NoResultFound())) - - with pytest.raises(RouterNotFoundException): - await model_registry.delete_router(router_id=999, postgres_session=postgres_session) - - -@pytest.mark.asyncio -async def test_delete_router_success(postgres_session: AsyncSession, model_registry: ModelRegistry): - postgres_session.execute = AsyncMock(side_effect=[_Result(scalar_one=MagicMock(id=1)), None]) - - await model_registry.delete_router(router_id=1, postgres_session=postgres_session) - postgres_session.commit.assert_awaited() - - -@pytest.mark.asyncio -async def test_update_router_not_found(postgres_session: AsyncSession, model_registry: ModelRegistry): - # get_routers is called internally and returns empty list when router_id is provided - # Mock postgres_session.execute to return empty results for get_routers query - postgres_session.execute = AsyncMock(return_value=_Result(all_rows=[])) - - with pytest.raises(RouterNotFoundException): - await model_registry.update_router( - router_id=999, - name="new-name", - type=None, - aliases=None, - load_balancing_strategy=None, - cost_prompt_tokens=None, - cost_completion_tokens=None, - postgres_session=postgres_session, - ) - - -@pytest.mark.asyncio -async def test_update_router_success_all_fields(postgres_session: AsyncSession, model_registry: ModelRegistry): - # Mock get_routers to return a router - router = Router( - id=1, - name="old-name", - user_id=1, - type=ModelType.TEXT_GENERATION, - aliases=["old-alias"], - load_balancing_strategy=RouterLoadBalancingStrategy.SHUFFLE, - vector_size=None, - max_context_length=4096, - cost_prompt_tokens=1.0, - cost_completion_tokens=2.0, - providers=0, - created=100, - updated=200, - ) - - model_registry.get_routers = AsyncMock(return_value=[router]) - # check aliases (scalars().all() returns empty), update router, delete old aliases, insert new aliases - alias_result = _Result(scalars_rows=[]) - alias_result.scalars = lambda: alias_result - alias_result.all = lambda: [] - postgres_session.execute = AsyncMock(side_effect=[alias_result, None, None, None]) - - await model_registry.update_router( - router_id=1, - name="new-name", - type=ModelType.TEXT_EMBEDDINGS_INFERENCE, - aliases=["new-alias1", "new-alias2"], - load_balancing_strategy=RouterLoadBalancingStrategy.LEAST_BUSY, - cost_prompt_tokens=3.0, - cost_completion_tokens=4.0, - postgres_session=postgres_session, - ) - - postgres_session.commit.assert_awaited() - - -@pytest.mark.asyncio -async def test_update_router_alias_conflict(postgres_session: AsyncSession, model_registry: ModelRegistry): - router = Router( - id=1, - name="router", - user_id=1, - type=ModelType.TEXT_GENERATION, - aliases=[], - load_balancing_strategy=RouterLoadBalancingStrategy.SHUFFLE, - vector_size=None, - max_context_length=4096, - cost_prompt_tokens=0.0, - cost_completion_tokens=0.0, - providers=0, - created=100, - updated=200, - ) - - model_registry.get_routers = AsyncMock(return_value=[router]) - # Check aliases returns conflicting router_id (scalars().all() returns [2]) - alias_result = _Result() - alias_result.scalars = lambda: alias_result - alias_result.all = lambda: [2] # alias belongs to router 2 - postgres_session.execute = AsyncMock(return_value=alias_result) - - with pytest.raises(RouterAliasAlreadyExistsException): - await model_registry.update_router( - router_id=1, - name=None, - type=None, - aliases=["conflicting-alias"], - load_balancing_strategy=None, - cost_prompt_tokens=None, - cost_completion_tokens=None, - postgres_session=postgres_session, - ) - - -@pytest.mark.asyncio -async def test_update_router_noop(postgres_session: AsyncSession, model_registry: ModelRegistry): - router = Router( - id=1, - name="router", - user_id=1, - type=ModelType.TEXT_GENERATION, - aliases=[], - load_balancing_strategy=RouterLoadBalancingStrategy.SHUFFLE, - vector_size=None, - max_context_length=4096, - cost_prompt_tokens=0.0, - cost_completion_tokens=0.0, - providers=0, - created=100, - updated=200, - ) - - model_registry.get_routers = AsyncMock(return_value=[router]) - postgres_session.execute = AsyncMock() - - await model_registry.update_router( - router_id=1, - name=None, - type=None, - aliases=None, - load_balancing_strategy=None, - cost_prompt_tokens=None, - cost_completion_tokens=None, - postgres_session=postgres_session, - ) - - postgres_session.commit.assert_awaited() - - -@pytest.mark.asyncio -async def test_get_routers_by_id_success(postgres_session: AsyncSession, model_registry: ModelRegistry): - # Mock router result with provider count - router_row = MagicMock( - _asdict=lambda: { - "id": 1, - "name": "test-router", - "user_id": 1, - "type": ModelType.TEXT_GENERATION.value, - "load_balancing_strategy": RouterLoadBalancingStrategy.SHUFFLE.value, - "cost_prompt_tokens": 1.0, - "cost_completion_tokens": 2.0, - "max_context_length": 4096, - "vector_size": None, - "providers": 2, - "created": 100, - "updated": 200, - } - ) - - postgres_session.execute = AsyncMock( - side_effect=[ - _Result(all_rows=[router_row]), # router query - _Result(all_rows=[]), # aliases query - ] - ) - - routers = await model_registry.get_routers(router_id=1, name=None, postgres_session=postgres_session) - - assert len(routers) == 1 - assert routers[0].id == 1 - assert routers[0].name == "test-router" - assert routers[0].type == ModelType.TEXT_GENERATION - - -@pytest.mark.asyncio -async def test_get_routers_by_id_not_found(postgres_session: AsyncSession, model_registry: ModelRegistry): - postgres_session.execute = AsyncMock(return_value=_Result(all_rows=[])) - - with pytest.raises(RouterNotFoundException): - await model_registry.get_routers(router_id=999, name=None, postgres_session=postgres_session) - - -@pytest.mark.asyncio -async def test_get_routers_by_name_success(postgres_session: AsyncSession, model_registry: ModelRegistry): - router_row = MagicMock( - _asdict=lambda: { - "id": 1, - "name": "test-router", - "user_id": 1, - "type": ModelType.TEXT_GENERATION.value, - "load_balancing_strategy": RouterLoadBalancingStrategy.SHUFFLE.value, - "cost_prompt_tokens": 0.0, - "cost_completion_tokens": 0.0, - "max_context_length": 4096, - "vector_size": None, - "providers": 1, - "created": 100, - "updated": 200, - } - ) - - alias_row = MagicMock(router_id=1, value="alias1") - - postgres_session.execute = AsyncMock( - side_effect=[ - _Result(all_rows=[router_row]), # router query - _Result(all_rows=[alias_row]), # aliases query - ] - ) - - routers = await model_registry.get_routers(router_id=None, name="test-router", postgres_session=postgres_session) - - assert len(routers) == 1 - assert routers[0].name == "test-router" - - -@pytest.mark.asyncio -async def test_get_routers_by_alias_success(postgres_session: AsyncSession, model_registry: ModelRegistry): - router_row = MagicMock( - _asdict=lambda: { - "id": 1, - "name": "test-router", - "user_id": 1, - "type": ModelType.TEXT_GENERATION.value, - "load_balancing_strategy": RouterLoadBalancingStrategy.SHUFFLE.value, - "cost_prompt_tokens": 0.0, - "cost_completion_tokens": 0.0, - "max_context_length": 4096, - "vector_size": None, - "providers": 1, - "created": 100, - "updated": 200, - } - ) - - alias_row = MagicMock(router_id=1, value="my-alias") - - postgres_session.execute = AsyncMock( - side_effect=[ - _Result(all_rows=[router_row]), # router query - _Result(all_rows=[alias_row]), # aliases query - ] - ) - - routers = await model_registry.get_routers(router_id=None, name="my-alias", postgres_session=postgres_session) - - assert len(routers) == 1 - assert "my-alias" in routers[0].aliases - - -@pytest.mark.asyncio -async def test_get_routers_by_name_not_found(postgres_session: AsyncSession, model_registry: ModelRegistry): - router_row = MagicMock( - _asdict=lambda: { - "id": 1, - "name": "other-router", - "user_id": 1, - "type": ModelType.TEXT_GENERATION.value, - "load_balancing_strategy": RouterLoadBalancingStrategy.SHUFFLE.value, - "cost_prompt_tokens": 0.0, - "cost_completion_tokens": 0.0, - "max_context_length": 4096, - "vector_size": None, - "providers": 1, - "created": 100, - "updated": 200, - } - ) - - postgres_session.execute = AsyncMock( - side_effect=[ - _Result(all_rows=[router_row]), # router query - _Result(all_rows=[]), # aliases query - ] - ) - - with pytest.raises(RouterNotFoundException): - await model_registry.get_routers(router_id=None, name="nonexistent", postgres_session=postgres_session) - - -@pytest.mark.asyncio -async def test_get_routers_pagination_and_ordering(postgres_session: AsyncSession, model_registry: ModelRegistry): - router_row1 = MagicMock( - _asdict=lambda: { - "id": 1, - "name": "router-a", - "user_id": 1, - "type": ModelType.TEXT_GENERATION.value, - "load_balancing_strategy": RouterLoadBalancingStrategy.SHUFFLE.value, - "cost_prompt_tokens": 0.0, - "cost_completion_tokens": 0.0, - "max_context_length": 4096, - "vector_size": None, - "providers": 1, - "created": 100, - "updated": 200, - } - ) - - router_row2 = MagicMock( - _asdict=lambda: { - "id": 2, - "name": "router-b", - "user_id": 1, - "type": ModelType.TEXT_EMBEDDINGS_INFERENCE.value, - "load_balancing_strategy": RouterLoadBalancingStrategy.LEAST_BUSY.value, - "cost_prompt_tokens": 0.0, - "cost_completion_tokens": 0.0, - "max_context_length": 2048, - "vector_size": 768, - "providers": 2, - "created": 200, - "updated": 300, - } - ) - - postgres_session.execute = AsyncMock( - side_effect=[ - _Result(all_rows=[router_row1, router_row2]), # router query - _Result(all_rows=[]), # aliases query - ] - ) - - routers = await model_registry.get_routers( - router_id=None, - name=None, - postgres_session=postgres_session, - offset=0, - limit=10, - order_by="name", - order_direction="asc", - ) - - assert len(routers) == 2 - - -@pytest.mark.asyncio -async def test_get_routers_master_user(postgres_session: AsyncSession, model_registry: ModelRegistry): - # Router with user_id=None should be converted to user_id=0 - router_row = MagicMock( - _asdict=lambda: { - "id": 1, - "name": "master-router", - "user_id": None, # master user - "type": ModelType.TEXT_GENERATION.value, - "load_balancing_strategy": RouterLoadBalancingStrategy.SHUFFLE.value, - "cost_prompt_tokens": 0.0, - "cost_completion_tokens": 0.0, - "max_context_length": 4096, - "vector_size": None, - "providers": 0, - "created": 100, - "updated": 200, - } - ) - - postgres_session.execute = AsyncMock( - side_effect=[ - _Result(all_rows=[router_row]), - _Result(all_rows=[]), - ] - ) - - routers = await model_registry.get_routers(router_id=1, name=None, postgres_session=postgres_session) - - assert routers[0].user_id == 0 - - -@pytest.mark.asyncio -async def test_get_router_id_from_model_name_by_name(postgres_session: AsyncSession, model_registry: ModelRegistry): - postgres_session.execute = AsyncMock(return_value=_Result(scalar_one=42)) - - router_id = await model_registry.get_router_id_from_model_name("test-model", postgres_session) - - assert router_id == 42 - - -@pytest.mark.asyncio -async def test_get_router_id_from_model_name_by_alias(postgres_session: AsyncSession, model_registry: ModelRegistry): - postgres_session.execute = AsyncMock(return_value=_Result(scalar_one=99)) - - router_id = await model_registry.get_router_id_from_model_name("alias-name", postgres_session) - - assert router_id == 99 - - -@pytest.mark.asyncio -async def test_get_router_id_from_model_name_not_found(postgres_session: AsyncSession, model_registry: ModelRegistry): - postgres_session.execute = AsyncMock(return_value=_Result(scalar_one=None)) - - router_id = await model_registry.get_router_id_from_model_name("nonexistent", postgres_session) - - assert router_id is None diff --git a/api/use_cases/admin/providers/_createproviderusecase.py b/api/use_cases/admin/providers/_createproviderusecase.py index b87b835e3..8ddf0f5a9 100644 --- a/api/use_cases/admin/providers/_createproviderusecase.py +++ b/api/use_cases/admin/providers/_createproviderusecase.py @@ -78,8 +78,6 @@ async def execute(self, command: CreateProviderCommand) -> CreateProviderUseCase timeout=command.timeout, model_name=command.model_name, ) - # @ TODO: separate health check logic from get_capabilities - match result: case ProviderNotReachableError() as error: return error diff --git a/api/utils/routing.py b/api/utils/routing.py index c37d348e6..c6cc171c3 100644 --- a/api/utils/routing.py +++ b/api/utils/routing.py @@ -4,7 +4,7 @@ from celery.result import AsyncResult from redis.asyncio import Redis as AsyncRedis -from api.schemas.admin.providers import Provider +from api.domain.provider.entities import Provider from api.schemas.admin.routers import RouterLoadBalancingStrategy from api.schemas.core.models import Metric from api.tasks import app, create_model_queue @@ -124,7 +124,7 @@ async def apply_routing_with_queuing( logger.debug(f"Task {task.id}: Result={result}") if result["status_code"] != 200: - logger.error(f"Task {task.id}: Failed with status_code={result["status_code"]}, detail={result.get("body", {}).get("detail", "N/A")}") + logger.error(f"Task {task.id}: Failed with status_code={result['status_code']}, detail={result.get('body', {}).get('detail', 'N/A')}") raise TaskFailedException(status_code=result["status_code"], detail=result["body"]["detail"]) provider_id = result["provider_id"] logger.info(f"Task {task.id}: Successfully returned provider_id={provider_id}")