diff --git a/README.md b/README.md index 7ecc7f3..803e176 100644 --- a/README.md +++ b/README.md @@ -11,7 +11,7 @@ Welcome to the `api` repository! This repository hosts the core API infrastructu - **FastAPI**: A modern, fast (high-performance) web framework for building APIs with Python 3.13 based on standard Python type hints. - **MongoDB**: A NoSQL database used for storing application data, offering flexibility and scalability. - **Qdrant**: A vector database used for implementing Retrieval-Augmented Generation (RAG) features, enhancing information retrieval capabilities. -- **LLM Integration**: Routes for calling the MistralAI API to leverage Large Language Models (LLMs) for various NLP tasks. +- **LLM Integration**: Routes for calling the MistralAI, OpenAI, or Anthropic APIs to leverage Large Language Models (LLMs) for various NLP tasks. - **Versioning**: A systematic versioning approach with endpoints like `/v1`, `/v2`, etc., to manage API changes and ensure backward compatibility. ## 🚀 Getting Started @@ -21,7 +21,7 @@ Welcome to the `api` repository! This repository hosts the core API infrastructu - Python 3.13 - MongoDB instance - Qdrant instance -- MistralAI API key +- API keys for MistralAI, OpenAI, or Anthropic (depending on the provider you wish to use) ### Installation @@ -46,8 +46,10 @@ Welcome to the `api` repository! This repository hosts the core API infrastructu Create a `.env` file in the root directory and add the following variables: ``` MONGODB_URI=your_mongodb_connection_string - QDRANT_URI=your_qdrant_connection_string + QDRANT_URL=your_qdrant_connection_string MISTRAL_API_KEY=your_mistral_api_key + OPENAI_API_KEY=your_openai_api_key + ANTHROPIC_API_KEY=your_anthropic_api_key ``` ### Running the API @@ -72,7 +74,7 @@ The API will be accessible at `http://127.0.0.1:3001`. - `GET /v1/health`: Check the status of the API. - **LLM Routes**: - - `POST /v1/chat/completions`: Generate text completions using the MistralAI API. + - `POST /v1/chat/completions`: Generate text completions using the configured AI provider. - **Data Routes**: - `GET /v1/data`: Retrieve data from MongoDB. diff --git a/pyproject.toml b/pyproject.toml index bf37022..4a75c7f 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -13,6 +13,8 @@ dependencies = [ "pymongo>=4.3.3", "uvicorn>=0.34.0", "mistralai>=1.5.0", + "openai>=1.3.8", + "anthropic>=0.21.3", "PyJWT>=2.10.1" ] diff --git a/src/api/classes/embeddings.py b/src/api/classes/embeddings.py index 0761227..456bb46 100644 --- a/src/api/classes/embeddings.py +++ b/src/api/classes/embeddings.py @@ -1,11 +1,12 @@ +from api.providers import AIProvider from api.utils import CustomLogger, InferenceUtils logger = CustomLogger.get_logger(__name__) class Embeddings: - def __init__(self, mistralai_service, inference_utils: InferenceUtils): - self.mistralai_service = mistralai_service + def __init__(self, provider: AIProvider, inference_utils: InferenceUtils): + self.provider = provider self.inference_utils = inference_utils async def generate_embeddings( @@ -15,14 +16,14 @@ async def generate_embeddings( job_id: str, output_format: str = "dict", ): - response = await self.mistralai_service.generate_embeddings( + response = await self.provider.generate_embeddings( inputs=inputs, model=model, ) data = response.get("data", None) if not data: - raise ValueError("Invalid response from Mistral API") + raise ValueError("Invalid response from AI provider") if output_format == "points": return self.inference_utils.embedding_to_points(inputs, data) diff --git a/src/api/classes/text_generation.py b/src/api/classes/text_generation.py index 3d3d6cc..08be34a 100644 --- a/src/api/classes/text_generation.py +++ b/src/api/classes/text_generation.py @@ -1,9 +1,11 @@ from typing import Any, AsyncGenerator, Dict +from api.providers import AIProvider + class TextGeneration: - def __init__(self, mistralai_service, inference_utils): - self.mistralai_service = mistralai_service + def __init__(self, provider: AIProvider, inference_utils): + self.provider = provider self.inference_utils = inference_utils async def generate_stream_response( @@ -29,7 +31,7 @@ async def generate_stream_response( Returns: Un générateur asynchrone de dictionnaires contenant les chunks et métadonnées """ - async for response, finish_reason in self.mistralai_service.stream( + async for response, finish_reason in self.provider.stream( model=model, messages=messages, temperature=temperature, @@ -67,7 +69,7 @@ async def complete( Returns: Un dictionnaire contenant la réponse et l'identifiant de la tâche """ - response = await self.mistralai_service.complete( + response = await self.provider.complete( model=model, messages=messages, temperature=temperature, diff --git a/src/api/databases/qdrant_connector.py b/src/api/databases/qdrant_connector.py index 4d8c244..08fb4ad 100644 --- a/src/api/databases/qdrant_connector.py +++ b/src/api/databases/qdrant_connector.py @@ -159,7 +159,9 @@ async def insert_vectors(self, collection_name, vectors, payloads=None): points = [ models.PointStruct(id=idx, vector=vector, payload=payload) - for idx, (vector, payload) in enumerate(zip(vectors, payloads)) + for idx, (vector, payload) in enumerate( + zip(vectors, payloads, strict=False) + ) ] await self.client.upsert(collection_name=collection_name, points=points) diff --git a/src/api/providers/__init__.py b/src/api/providers/__init__.py new file mode 100644 index 0000000..1e8da14 --- /dev/null +++ b/src/api/providers/__init__.py @@ -0,0 +1,15 @@ +from .ai_provider import ( + AIProvider, + AnthropicProvider, + MistralAIProvider, + OpenAIProvider, + get_provider, +) + +__all__ = [ + "AIProvider", + "MistralAIProvider", + "OpenAIProvider", + "AnthropicProvider", + "get_provider", +] diff --git a/src/api/providers/ai_provider.py b/src/api/providers/ai_provider.py new file mode 100644 index 0000000..cc90d96 --- /dev/null +++ b/src/api/providers/ai_provider.py @@ -0,0 +1,233 @@ +from __future__ import annotations + +import json +import os +from abc import ABC, abstractmethod +from json.decoder import JSONDecodeError +from typing import AsyncGenerator, Tuple + +from anthropic import AsyncAnthropic +from fastapi import HTTPException +from mistralai import Mistral +from openai import AsyncOpenAI + +from api.utils import CustomLogger + +logger = CustomLogger.get_logger(__name__) + + +class AIProvider(ABC): + @abstractmethod + async def complete( + self, + model: str, + messages: list, + temperature: float, + max_tokens: int, + top_p: float, + ) -> str: + pass + + @abstractmethod + async def stream( + self, + model: str, + messages: list, + temperature: float, + max_tokens: int, + top_p: float, + ) -> AsyncGenerator[Tuple[str, str | None], None]: + pass + + @abstractmethod + async def check_model(self, model: str): + pass + + @abstractmethod + async def list_models(self): + pass + + @abstractmethod + async def generate_embeddings(self, inputs: list[str], model: str): + pass + + +class MistralAIProvider(AIProvider): + def __init__(self) -> None: + self.api_key = os.getenv("MISTRAL_API_KEY") + if not self.api_key: + raise ValueError("MISTRAL_API_KEY environment variable is not set") + self.client = Mistral(api_key=self.api_key) + + async def complete(self, model, messages, temperature, max_tokens, top_p): + response = await self.client.chat.complete_async( + model=model, + messages=messages, + temperature=temperature, + max_tokens=max_tokens, + top_p=top_p, + ) + return response.choices[0].message.content + + async def stream(self, model, messages, temperature, max_tokens, top_p): + response = await self.client.chat.stream_async( + model=model, + messages=messages, + temperature=temperature, + max_tokens=max_tokens, + top_p=top_p, + ) + async for chunk in response: + if chunk.data.choices[0].delta.content is not None: + yield ( + chunk.data.choices[0].delta.content, + chunk.data.choices[0].finish_reason, + ) + + async def check_model(self, model): + try: + model = await self.client.models.retrieve_async(model_id=model) + if model is None: + raise HTTPException(status_code=404, detail="Model not found") + return model + except Exception as e: + if "status 404" in str(e).lower(): + raise HTTPException(status_code=404, detail="Model not found") from e + logger.error(f"An error occurred while checking model : {e}") + raise HTTPException(status_code=503, detail="Service unavailable") from e + + async def list_models(self): + return await self.client.models.list_async() + + async def generate_embeddings( + self, inputs: list[str], model: str = "mistral-embed" + ): + try: + response = await self.client.embeddings.create_async( + model=model, inputs=inputs + ) + return json.loads(response.model_dump_json()) + except JSONDecodeError as e: + logger.error(f"An error occurred while generating embeddings : {e}") + raise ConnectionError( + f"An error occurred while generating embeddings : {e}" + ) from e + except Exception as e: + logger.error(f"An error occurred while generating embeddings : {e}") + raise ConnectionError( + f"An error occurred while generating embeddings : {e}" + ) from e + + +class OpenAIProvider(AIProvider): + def __init__(self) -> None: + self.api_key = os.getenv("OPENAI_API_KEY") + if not self.api_key: + raise ValueError("OPENAI_API_KEY environment variable is not set") + self.client = AsyncOpenAI(api_key=self.api_key) + + async def complete(self, model, messages, temperature, max_tokens, top_p): + resp = await self.client.chat.completions.create( + model=model, + messages=messages, + temperature=temperature, + max_tokens=max_tokens, + top_p=top_p, + ) + return resp.choices[0].message.content + + async def stream(self, model, messages, temperature, max_tokens, top_p): + stream = await self.client.chat.completions.create( + model=model, + messages=messages, + temperature=temperature, + max_tokens=max_tokens, + top_p=top_p, + stream=True, + ) + async for chunk in stream: + delta = chunk.choices[0].delta + if delta and delta.content: + yield delta.content, chunk.choices[0].finish_reason + + async def check_model(self, model): + try: + return await self.client.models.retrieve(model) + except Exception as e: + logger.error(f"OpenAI model check failed: {e}") + raise HTTPException(status_code=404, detail="Model not found") from e + + async def list_models(self): + return await self.client.models.list() + + async def generate_embeddings( + self, inputs: list[str], model: str = "text-embedding-3-small" + ): + try: + resp = await self.client.embeddings.create(model=model, input=inputs) + return json.loads(resp.model_dump_json()) + except Exception as e: + logger.error(f"OpenAI embeddings error: {e}") + raise ConnectionError(f"OpenAI embeddings error: {e}") from e + + +class AnthropicProvider(AIProvider): + def __init__(self) -> None: + self.api_key = os.getenv("ANTHROPIC_API_KEY") + if not self.api_key: + raise ValueError("ANTHROPIC_API_KEY environment variable is not set") + self.client = AsyncAnthropic(api_key=self.api_key) + + async def complete(self, model, messages, temperature, max_tokens, top_p): + resp = await self.client.messages.create( + model=model, + messages=messages, + temperature=temperature, + max_tokens=max_tokens, + top_p=top_p, + ) + return resp.content[0].text if resp.content else "" + + async def stream(self, model, messages, temperature, max_tokens, top_p): + stream = await self.client.messages.create( + model=model, + messages=messages, + temperature=temperature, + max_tokens=max_tokens, + top_p=top_p, + stream=True, + ) + async for event in stream: + if event.type == "content_block_delta" and event.delta.text: + yield event.delta.text, None + + async def check_model(self, model): + try: + return await self.client.models.retrieve(model) + except Exception as e: + logger.error(f"Anthropic model check failed: {e}") + raise HTTPException(status_code=404, detail="Model not found") from e + + async def list_models(self): + return await self.client.models.list() + + async def generate_embeddings( + self, inputs: list[str], model: str = "claude-3-sonnet-20240229" + ): + try: + resp = await self.client.embeddings.create(model=model, input=inputs) + return json.loads(resp.model_dump_json()) + except Exception as e: + logger.error(f"Anthropic embeddings error: {e}") + raise ConnectionError(f"Anthropic embeddings error: {e}") from e + + +def get_provider(name: str) -> AIProvider: + normalized = name.lower() + if normalized in {"mistral", "mistralai"}: + return MistralAIProvider() + if normalized == "openai": + return OpenAIProvider() + if normalized == "anthropic": + return AnthropicProvider() + raise ValueError(f"Unknown provider: {name}") diff --git a/src/api/v1/routes/auth/auth_routes.py b/src/api/v1/routes/auth/auth_routes.py index e69e1ff..42e4ab8 100644 --- a/src/api/v1/routes/auth/auth_routes.py +++ b/src/api/v1/routes/auth/auth_routes.py @@ -1,6 +1,8 @@ from datetime import timedelta from typing import List +from fastapi import APIRouter, Depends, Form, HTTPException, status + from api.v1.security import ( APIAuth, APIKeyNotFoundError, @@ -13,7 +15,6 @@ get_current_user_with_token, oauth2_scheme, ) -from fastapi import APIRouter, Depends, Form, HTTPException, status from .auth_models import ( ApiKeyEntry, diff --git a/src/api/v1/routes/chat/chat_models.py b/src/api/v1/routes/chat/chat_models.py index 907244e..f49af64 100644 --- a/src/api/v1/routes/chat/chat_models.py +++ b/src/api/v1/routes/chat/chat_models.py @@ -16,6 +16,9 @@ class ChatCompletionsRequest(BaseModel): """Modèle pour une requête de complétion de chat""" model: str = Field(..., description="Identifiant du modèle à utiliser") + provider: str = Field( + "mistral", description="Provider à utiliser (mistral, openai, anthropic)" + ) prompt: Optional[str] = Field( None, description="Prompt à utiliser pour la génération" ) diff --git a/src/api/v1/routes/chat/chat_routes.py b/src/api/v1/routes/chat/chat_routes.py index c705e1a..8108f60 100644 --- a/src/api/v1/routes/chat/chat_routes.py +++ b/src/api/v1/routes/chat/chat_routes.py @@ -2,16 +2,18 @@ import uuid from datetime import datetime +from fastapi import APIRouter, Depends, HTTPException +from fastapi.responses import StreamingResponse + from api.classes import TextGeneration from api.databases import MongoDBConnector -from api.utils import CustomLogger +from api.providers import get_provider +from api.utils import CustomLogger, InferenceUtils from api.v1.security import ( ensure_valid_api_key_or_token, get_current_user_with_api_key_or_token, ) -from api.v1.services import get_mongo_client, get_text_generation -from fastapi import APIRouter, Depends, HTTPException -from fastapi.responses import StreamingResponse +from api.v1.services import get_mongo_client from .chat_models import ChatCompletionResponse, ChatCompletionsRequest @@ -75,12 +77,13 @@ async def sse_stream_generator(generator, job_id): ) async def completions( chat_request: ChatCompletionsRequest, - text_generation: TextGeneration = Depends(get_text_generation), user: dict = Depends(get_current_user_with_api_key_or_token), mongodb_client: MongoDBConnector = Depends(get_mongo_client), ): + provider_instance = get_provider(chat_request.provider) + text_generation = TextGeneration(provider_instance, InferenceUtils()) # Validation du modèle - await text_generation.mistralai_service.check_model(chat_request.model) + await text_generation.provider.check_model(chat_request.model) # Formatage des messages messages = text_generation.inference_utils.format_messages( diff --git a/src/api/v1/routes/embeddings/embeddings_models.py b/src/api/v1/routes/embeddings/embeddings_models.py index ddf0fd7..e25b440 100644 --- a/src/api/v1/routes/embeddings/embeddings_models.py +++ b/src/api/v1/routes/embeddings/embeddings_models.py @@ -6,6 +6,7 @@ class EmbeddingsRequest(BaseModel): chunks: List[str] model: str + provider: str = "mistral" class Embedding(BaseModel): diff --git a/src/api/v1/routes/embeddings/embeddings_routes.py b/src/api/v1/routes/embeddings/embeddings_routes.py index d55b262..67c0883 100644 --- a/src/api/v1/routes/embeddings/embeddings_routes.py +++ b/src/api/v1/routes/embeddings/embeddings_routes.py @@ -1,13 +1,17 @@ import json import uuid +from fastapi import APIRouter, Depends, HTTPException + +from api.classes import Embeddings from api.databases import MongoDBConnector +from api.providers import get_provider +from api.utils import InferenceUtils from api.v1.security import ( ensure_valid_api_key_or_token, get_current_user_with_api_key_or_token, ) -from api.v1.services import get_embeddings, get_mongo_client -from fastapi import APIRouter, Depends, HTTPException +from api.v1.services import get_mongo_client from .embeddings_models import EmbeddingsRequest, EmbeddingsResponse @@ -22,10 +26,11 @@ ) async def embeddings( body: EmbeddingsRequest, - embeddings=Depends(get_embeddings), user: dict = Depends(get_current_user_with_api_key_or_token), mongodb_client: MongoDBConnector = Depends(get_mongo_client), ): + provider_instance = get_provider(body.provider) + embeddings = Embeddings(provider_instance, InferenceUtils()) """Get embeddings for the input text.""" # Validation de l'entrée if not body.chunks: diff --git a/src/api/v1/routes/models/models_models.py b/src/api/v1/routes/models/models_models.py index e057b1f..5b66786 100644 --- a/src/api/v1/routes/models/models_models.py +++ b/src/api/v1/routes/models/models_models.py @@ -1,13 +1,9 @@ -from mistralai.models import ( - ModelList, - RetrieveModelV1ModelsModelIDGetResponseRetrieveModelV1ModelsModelIDGet, -) from pydantic import BaseModel class GetModelResponse(BaseModel): - model: RetrieveModelV1ModelsModelIDGetResponseRetrieveModelV1ModelsModelIDGet + model: dict class ListModelsResponse(BaseModel): - models: ModelList + models: list diff --git a/src/api/v1/routes/models/models_routes.py b/src/api/v1/routes/models/models_routes.py index 526077f..47fe5a3 100644 --- a/src/api/v1/routes/models/models_routes.py +++ b/src/api/v1/routes/models/models_routes.py @@ -1,5 +1,6 @@ -from api.v1.services import get_mistral_service -from fastapi import APIRouter, Depends, Path +from fastapi import APIRouter, Path, Query + +from api.providers import get_provider from .models_models import GetModelResponse, ListModelsResponse @@ -13,12 +14,14 @@ ) async def read_model( model_id: str = Path(..., description="The ID of the model to retrieve"), - mistralai_service=Depends(get_mistral_service), + provider: str = Query("mistral", description="AI provider"), ): - return {"model": await mistralai_service.check_model(model_id)} + service = get_provider(provider) + return {"model": await service.check_model(model_id)} @router.get("/", response_model=ListModelsResponse, summary="List all models") -async def list_models(mistralai_service=Depends(get_mistral_service)): - models = await mistralai_service.list_models() +async def list_models(provider: str = Query("mistral", description="AI provider")): + service = get_provider(provider) + models = await service.list_models() return {"models": models} diff --git a/src/api/v1/routes/rag/rag_models.py b/src/api/v1/routes/rag/rag_models.py index 029a3e0..9aba353 100644 --- a/src/api/v1/routes/rag/rag_models.py +++ b/src/api/v1/routes/rag/rag_models.py @@ -6,6 +6,7 @@ class RagEncodeRequest(BaseModel): chunks: List[str] model: str = "mistral-embed" + provider: str = "mistral" class RagEncodeResponse(BaseModel): @@ -17,6 +18,7 @@ class RagRetrieveRequest(BaseModel): query: str model: str = "mistral-embed" limit: int = 5 + provider: str = "mistral" class RetrieveResult(BaseModel): diff --git a/src/api/v1/routes/rag/rag_routes.py b/src/api/v1/routes/rag/rag_routes.py index 57d21df..4797c85 100644 --- a/src/api/v1/routes/rag/rag_routes.py +++ b/src/api/v1/routes/rag/rag_routes.py @@ -2,20 +2,23 @@ import uuid from datetime import datetime +from fastapi import APIRouter, Depends, HTTPException + +from api.classes import Embeddings from api.databases import MongoDBConnector -from api.utils import CustomLogger +from api.providers import get_provider +from api.utils import CustomLogger, InferenceUtils from api.v1.security import ( ensure_valid_api_key_or_token, get_current_user_with_api_key_or_token, ) from api.v1.services import ( - RagService, check_collection_non_existence, check_collection_ownership, get_mongo_client, - get_rag_service, + get_qdrant_client, ) -from fastapi import APIRouter, Depends, HTTPException +from api.v1.services.rag_service import RagService from .rag_models import ( RagEncodeRequest, @@ -41,8 +44,8 @@ async def encode( collection_name: str, body: RagEncodeRequest, - rag_service: RagService = Depends(get_rag_service), mongodb_client: MongoDBConnector = Depends(get_mongo_client), + qdrant_client=Depends(get_qdrant_client), user: dict = Depends(get_current_user_with_api_key_or_token), ): """Encode text to a collection.""" @@ -52,6 +55,10 @@ async def encode( job_id: str = str(uuid.uuid4()) + provider_instance = get_provider(body.provider) + embeddings = Embeddings(provider_instance, InferenceUtils()) + rag_service = RagService(embeddings, qdrant_client, mongodb_client) + await rag_service.encode_to_collection( collection_name, body.chunks, @@ -93,17 +100,21 @@ async def encode( async def retrieve( collection_name: str, body: RagRetrieveRequest, - rag_service: RagService = Depends(get_rag_service), mongodb_client: MongoDBConnector = Depends(get_mongo_client), + qdrant_client=Depends(get_qdrant_client), user: dict = Depends(get_current_user_with_api_key_or_token), ): """Retrieve text from a collection.""" # Validation de l'entrée - if not input: + if not body.query: raise HTTPException(status_code=400, detail="Aucune entrée fournie") job_id: str = str(uuid.uuid4()) + provider_instance = get_provider(body.provider) + embeddings = Embeddings(provider_instance, InferenceUtils()) + rag_service = RagService(embeddings, qdrant_client, mongodb_client) + results = await rag_service.retrieve_in_collection( collection_name, body.query, diff --git a/src/api/v1/routes/vector_db/vector_db_routes.py b/src/api/v1/routes/vector_db/vector_db_routes.py index f1c1cde..c8c3b20 100644 --- a/src/api/v1/routes/vector_db/vector_db_routes.py +++ b/src/api/v1/routes/vector_db/vector_db_routes.py @@ -1,5 +1,9 @@ import json +from bson import ObjectId +from fastapi import APIRouter, Depends +from qdrant_client.models import CollectionInfo + from api.databases import MongoDBConnector, QdrantConnector from api.v1.security import ( ensure_valid_api_key_or_token, @@ -10,9 +14,6 @@ get_mongo_client, get_qdrant_client, ) -from bson import ObjectId -from fastapi import APIRouter, Depends -from qdrant_client.models import CollectionInfo from .vector_db_models import CollectionsResponse diff --git a/src/api/v1/services/__init__.py b/src/api/v1/services/__init__.py index f4856ce..68582b3 100644 --- a/src/api/v1/services/__init__.py +++ b/src/api/v1/services/__init__.py @@ -1,23 +1,32 @@ +from api.providers import ( + AIProvider, + AnthropicProvider, + MistralAIProvider, + OpenAIProvider, + get_provider, +) + from .check import check_collection_non_existence, check_collection_ownership from .get_classes import ( get_embeddings, - get_mistral_service, get_rag_service, get_text_generation, ) from .get_databases import get_mongo_client, get_qdrant_client -from .mistralai_service import MistralAIService from .rag_service import RagService __all__ = [ - "MistralAIService", + "AIProvider", + "MistralAIProvider", + "OpenAIProvider", + "AnthropicProvider", + "get_provider", "get_mongo_client", "get_qdrant_client", "get_embeddings", "RagService", "get_rag_service", "get_text_generation", - "get_mistral_service", "check_collection_ownership", "check_collection_non_existence", ] diff --git a/src/api/v1/services/get_classes.py b/src/api/v1/services/get_classes.py index f36c660..2bf781a 100644 --- a/src/api/v1/services/get_classes.py +++ b/src/api/v1/services/get_classes.py @@ -1,18 +1,18 @@ from fastapi import Depends from api.classes import Embeddings, TextGeneration +from api.providers import get_provider from api.utils import InferenceUtils from .get_databases import get_mongo_client, get_qdrant_client -from .mistralai_service import MistralAIService from .rag_service import RagService def get_embeddings( - mistralai_service: MistralAIService = Depends(), + provider: str = "mistral", inference_utils: InferenceUtils = Depends(), ) -> Embeddings: - return Embeddings(mistralai_service, inference_utils) + return Embeddings(get_provider(provider), inference_utils) def get_rag_service( @@ -24,13 +24,11 @@ def get_rag_service( def get_text_generation( - mistralai_service: MistralAIService = Depends(), + provider: str = "mistral", inference_utils: InferenceUtils = Depends(), ) -> TextGeneration: - return TextGeneration(mistralai_service, inference_utils) + return TextGeneration(get_provider(provider), inference_utils) -def get_mistral_service( - mistralai_service: MistralAIService = Depends(), -) -> MistralAIService: - return mistralai_service +def get_mistral_service(): + return get_provider("mistral") diff --git a/src/api/v1/services/mistralai_service.py b/src/api/v1/services/mistralai_service.py deleted file mode 100644 index bb17c34..0000000 --- a/src/api/v1/services/mistralai_service.py +++ /dev/null @@ -1,92 +0,0 @@ -import json -import os -from json.decoder import JSONDecodeError -from typing import Tuple - -from fastapi import HTTPException -from mistralai import Mistral - -from api.utils import CustomLogger - -logger = CustomLogger.get_logger(__name__) - - -class MistralAIService: - def __init__(self): - self.api_key = os.getenv("MISTRAL_API_KEY", None) - - if self.api_key is None: - raise ValueError("MISTRAL_API_KEY environment variable is not set") - - self.mistral_client = Mistral(api_key=self.api_key) - - async def complete(self, model, messages, temperature, max_tokens, top_p): - response = await self.mistral_client.chat.complete_async( - model=model, - messages=messages, - temperature=temperature, - max_tokens=max_tokens, - top_p=top_p, - ) - return response.choices[0].message.content - - async def stream( - self, model, messages, temperature, max_tokens, top_p - ) -> Tuple[str, str]: - response = await self.mistral_client.chat.stream_async( - model=model, - messages=messages, - temperature=temperature, - max_tokens=max_tokens, - top_p=top_p, - ) - async for chunk in response: - if chunk.data.choices[0].delta.content is not None: - yield ( - chunk.data.choices[0].delta.content, - chunk.data.choices[0].finish_reason, - ) - - async def check_model(self, model): - try: - model = await self.mistral_client.models.retrieve_async(model_id=model) - - if model is None: - raise HTTPException(status_code=404, detail="Model not found") from None - - return model - except Exception as e: - if "status 404" in str(e).lower(): - raise HTTPException(status_code=404, detail="Model not found") from e - else: - logger.error(f"An error occurred while checking model : {e}") - raise HTTPException( - status_code=503, detail="Service unavailable" - ) from e - - async def list_models(self): - return await self.mistral_client.models.list_async() - - async def generate_embeddings( - self, - inputs: list[str], - model: str = "mistral-embed", - ): - try: - response = await self.mistral_client.embeddings.create_async( - model=model, inputs=inputs - ) - - return json.loads(response.model_dump_json()) - - except JSONDecodeError as e: - logger.error(f"An error occurred while generating embeddings : {e}") - raise ConnectionError( - f"An error occurred while generating embeddings : {e}" - ) from e - - except Exception as e: - logger.error(f"An error occurred while generating embeddings : {e}") - raise ConnectionError( - f"An error occurred while generating embeddings : {e}" - ) from e diff --git a/tests/conftest.py b/tests/conftest.py new file mode 100644 index 0000000..0c0701b --- /dev/null +++ b/tests/conftest.py @@ -0,0 +1,9 @@ +import sys +import types +from pathlib import Path + +module = types.ModuleType("mistralai") +module.Mistral = object +sys.modules.setdefault("mistralai", module) + +sys.path.insert(0, str(Path(__file__).resolve().parents[1] / "src")) diff --git a/tests/test_api_auth.py b/tests/test_api_auth.py new file mode 100644 index 0000000..bfecd68 --- /dev/null +++ b/tests/test_api_auth.py @@ -0,0 +1,43 @@ +import os +os.environ["SECRET_KEY"]="test" +from datetime import timedelta + +import pytest + +from api.v1.security.api_auth import APIAuth + + +class DummyMongo: + def __init__(self): + self.user = {"_id": "u", "username": "me", "hashed_password": ""} + + async def find_one(self, collection, query, projection=None): + if collection == "users" and query.get("username") == "me": + return self.user + if collection == "users" and query.get("_id") == "u": + return self.user + return None + + async def insert_one(self, collection, doc): + return True + + def serialize(self, doc): + return doc + + +@pytest.mark.asyncio +async def test_password_hash_and_verify(): + auth = APIAuth() + pw = "secret" + hashed = auth.hash_password(pw) + assert auth.verify_password(pw, hashed) + + +@pytest.mark.asyncio +async def test_token_roundtrip(): + os.environ["SECRET_KEY"] = "s" + auth = APIAuth() + auth.set_mongo_client(DummyMongo()) + token = auth.create_access_token({"sub": "me"}, timedelta(minutes=1)) + user = await auth.verify_token(token) + assert user["username"] == "me" diff --git a/tests/test_classes.py b/tests/test_classes.py new file mode 100644 index 0000000..85d9d59 --- /dev/null +++ b/tests/test_classes.py @@ -0,0 +1,71 @@ +import asyncio +from types import SimpleNamespace +from unittest.mock import AsyncMock + +import pytest + +from api.classes.text_generation import TextGeneration +from api.classes.embeddings import Embeddings + + +class DummyProvider: + def __init__(self): + self.complete = AsyncMock(return_value="result") + self.generate_embeddings = AsyncMock( + return_value={"data": [{"object": "embedding", "embedding": [0.1, 0.2]}]} + ) + self.check_model = AsyncMock() + self.list_models = AsyncMock(return_value=["model1"]) + + async def stream(self, *args, **kwargs): + yield "chunk1", None + yield "chunk2", "stop" + + +@pytest.mark.asyncio +async def test_text_generation_complete(): + provider = DummyProvider() + tg = TextGeneration(provider, SimpleNamespace(format_response=lambda r, j: {"formatted": r, "job_id": j})) + result = await tg.complete("model", [], 0.7, 10, 0.9, "job") + assert result == {"formatted": "result", "job_id": "job"} + provider.complete.assert_awaited_once() + + +@pytest.mark.asyncio +async def test_text_generation_stream(): + provider = DummyProvider() + tg = TextGeneration(provider, SimpleNamespace()) + chunks = [] + async for c in tg.generate_stream_response("model", [], 0.7, 10, 0.9, "job"): + chunks.append(c) + assert chunks[-1]["finish_reason"] == "stop" + + +@pytest.mark.asyncio +async def test_embeddings_generate_embeddings_dict(): + provider = DummyProvider() + embed = Embeddings(provider, SimpleNamespace( + embedding_to_points=lambda i, d: "points", + embedding_to_tuple=lambda i, d: ([0], [[0.1, 0.2]], [{}]), + format_embeddings=lambda i, d, j: {"job_id": j, "embeddings": d}, + )) + data = await embed.generate_embeddings("model", ["a"], "job") + assert data["job_id"] == "job" + provider.generate_embeddings.assert_awaited_once() + + +@pytest.mark.asyncio +async def test_embeddings_output_tuple(): + provider = DummyProvider() + utils = SimpleNamespace( + embedding_to_points=lambda i, d: "points", + embedding_to_tuple=lambda i, d: ([0], [[0.1, 0.2]], [{}]), + format_embeddings=lambda i, d, j: {"job_id": j, "embeddings": d}, + ) + embed = Embeddings(provider, utils) + ids, vectors, payloads = await embed.generate_embeddings( + "model", ["a"], "job", output_format="tuple" + ) + assert ids == [0] + assert vectors == [[0.1, 0.2]] + assert payloads == [{}] diff --git a/tests/test_connectors.py b/tests/test_connectors.py new file mode 100644 index 0000000..bf4341a --- /dev/null +++ b/tests/test_connectors.py @@ -0,0 +1,31 @@ +import os +from unittest.mock import AsyncMock, MagicMock, patch + +import pytest + +from api.databases.mongo_db_connector import MongoDBConnector +from api.databases.qdrant_connector import QdrantConnector + + +class DummyLogger: + def __getattr__(self, name): + return lambda *args, **kwargs: None + + +def test_mongo_serialize(): + connector = MongoDBConnector(DummyLogger()) + doc = {"_id": connector.object_id("64b8a3000000000000000000"), "value": 1} + serialized = connector.serialize(doc) + assert isinstance(serialized["_id"], str) + + +@pytest.mark.asyncio +async def test_qdrant_create_collection(): + with patch.dict(os.environ, {"QDRANT_API_KEY": "k", "QDRANT_URL": "u"}): + with patch("api.databases.qdrant_connector.AsyncQdrantClient") as client_cls: + client = MagicMock() + client.create_collection = AsyncMock(return_value=True) + client_cls.return_value = client + qc = QdrantConnector(DummyLogger()) + await qc.create_collection("col") + client.create_collection.assert_awaited() diff --git a/tests/test_rag_service.py b/tests/test_rag_service.py new file mode 100644 index 0000000..dd86748 --- /dev/null +++ b/tests/test_rag_service.py @@ -0,0 +1,40 @@ +import os +from unittest.mock import AsyncMock, MagicMock, patch + +import pytest + +from api.v1.services.rag_service import RagService + + +class DummyEmbeddings: + def __init__(self): + self.generate_embeddings = AsyncMock(return_value=([0], [[0.1, 0.2]], [{}])) + + +@pytest.mark.asyncio +async def test_encode_to_collection(): + with patch.dict(os.environ, {"QDRANT_API_KEY": "k", "QDRANT_URL": "u"}): + qdrant = MagicMock() + qdrant.get_collection = AsyncMock(side_effect=[None, MagicMock(model_dump_json=lambda: "{}")]) + qdrant.create_collection = AsyncMock(return_value=True) + qdrant.batch_upsert = AsyncMock() + + mongo = MagicMock() + service = RagService(DummyEmbeddings(), qdrant, mongo) + res = await service.encode_to_collection("col", ["text"], "model") + qdrant.create_collection.assert_awaited() + qdrant.batch_upsert.assert_awaited() + assert res["collection_name"] == "col" + + +@pytest.mark.asyncio +async def test_retrieve_in_collection(): + with patch.dict(os.environ, {"QDRANT_API_KEY": "k", "QDRANT_URL": "u"}): + qdrant = MagicMock() + qdrant.search_in_collection = AsyncMock(return_value=[{"payload": {}, "score": 0.9}]) + embeddings = DummyEmbeddings() + service = RagService(embeddings, qdrant, MagicMock()) + results = await service.retrieve_in_collection("col", "q", "model") + embeddings.generate_embeddings.assert_awaited() + qdrant.search_in_collection.assert_awaited_with(collection_name="col", query_vector=[0.1,0.2], limit=5) + assert results[0]["score"] == 0.9 diff --git a/tests/test_routes.py b/tests/test_routes.py new file mode 100644 index 0000000..1c4ad52 --- /dev/null +++ b/tests/test_routes.py @@ -0,0 +1,23 @@ +import pytest +from unittest.mock import AsyncMock, MagicMock, patch + +from api.v1.endpoints import health_check + + +class DummyDB: + def __init__(self): + self.find_one = AsyncMock() + self.insert_one = AsyncMock() + self.get_client = lambda: MagicMock(close=lambda: None) + + +@pytest.fixture(autouse=True) +def startup_patch(): + with patch("api.main.ensure_database_connection", AsyncMock(return_value=(DummyDB(), DummyDB()))): + yield + + +@pytest.mark.asyncio +async def test_health_endpoint(startup_patch): + result = await health_check() + assert result["status"] == "ok" diff --git a/tests/test_utils.py b/tests/test_utils.py new file mode 100644 index 0000000..f30a072 --- /dev/null +++ b/tests/test_utils.py @@ -0,0 +1,38 @@ +import pytest + +from api.utils.inference import InferenceUtils, Message + + +def test_format_messages(): + history = [Message(role="user", content="hi"), Message(role="assistant", content="hello")] + result = InferenceUtils.format_messages("sys", "question", history) + assert result == [ + {"role": "system", "content": "sys"}, + {"role": "user", "content": "hi"}, + {"role": "assistant", "content": "hello"}, + {"role": "user", "content": "question"}, + ] + + +def test_format_response(): + resp = InferenceUtils.format_response("answer", "job") + assert resp == {"result": "answer", "job_id": "job"} + + +def test_embedding_helpers(): + inputs = ["a", "b"] + data = [ + {"object": "embedding", "embedding": [1.0, 0.0]}, + {"object": "embedding", "embedding": [0.0, 1.0]}, + ] + points = InferenceUtils.embedding_to_points(inputs, data) + assert points[0]["payload"]["source_text"] == "a" + + ids, vectors, payloads = InferenceUtils.embedding_to_tuple(inputs, data) + assert ids == [0, 1] + assert vectors == [[1.0, 0.0], [0.0, 1.0]] + assert payloads[1]["source_text"] == "b" + + formatted = InferenceUtils.format_embeddings(inputs, data, "job") + assert formatted["job_id"] == "job" + assert len(formatted["embeddings"]) == 2