fix: large collections issues (#32)

* fix: large collections issues

* feat: resolve some TO DO

* feat: add exception handler

* doc: add collections documentation

---------

Co-authored-by: leoguillaume <[email protected]>

.gitignore

@@ -6,9 +6,6 @@ __pycache__/
 # C extensions
 *.so
 
-# PDF Files
-*.pdf
-
 # Distribution / packaging
 .Python
 build/
@@ -201,4 +198,6 @@ terraform.rc
 .DS_Store
 .idea
 .vscode
-*.json
+.json
+/docs/tutorials/*.json
+/docs/tutorials/*.pdf

.pre-commit-config.yaml

@@ -9,4 +9,4 @@ repos:
         args: [ --fix ]
       # Run the formatter.
       - id: ruff-format
-        types_or: [ python, pyi ]
+        types_or: [ python, pyi ]

CHANGELOG.md

@@ -12,12 +12,28 @@ Tous les changements notables de l'application sont documentés dans ce fichier.
 - 🔄 Refactoring
 - ❌ Deprecated
 
+## [Alpha] - 2024-10-07
+
+- 💣 Création de la notion de Document, objet intermédiaire entre un fichier et une collection de chunks. Ajout des endpoints GET `/v1/documents` et DELETE `/v1/documents` pour solutionner le problème de limite de taille de requête.
+- ❌ Suppression de l'endpoint POST `/v1/chunks` pour récupérer plusieurs chunks de différents documents en une seule requête.
+- ❌ Suppression des endpoint GET `/v1/files` et DELETE `/v1/files`
+- 🎉 Ajout de la possibilité de récupérer les documents d'une collection
+- 🎉 Ajout de la possibilité de supprimer un document d'une collection
+- 🎉 Ajout de la possibilité de récupérer les chunks d'un document
+- 🎉 Ajout d'un chunker "NoChunker" qui permet de considérer le fichier en entier comme un chunk
+- 🐛 Les exceptions sont remontées de manière plus claire dans l'API
+- 🐛 Les modèles d'embeddings font remontées une erreur lorsque le context fourni est trop grand
+- 🔄 Meilleur cloisement des dépendances techniques du projet dans des classes distinctes
+- 🧪 Ajout de tests unitaires pour les endpoints *documents* et *chunks*
+- 🧪 Ajout de la configuration pytest dans le fichier `pyproject.toml`
+
 ## [Alpha] - 2024-10-01
 
 - 💣 Les collections sont appelées dorénavant par leur collection ID et non plus par leur nom
 - 💣 Le endpoint POST `/v1/files` ne créer plus de collection si elle n'existe pas
 - 🎉 Le endpoint POST `/v1/files` accepte maintenant tous les paramètres du chunking
 - 🎉 Ajout de rôles utilisateur et admin pour la création de collection publiques
+- 🎉 Ajout d'un middleware pour limiter la taille des requêtes d'upload de fichiers
 - 🎉 Ajout de la collection "internet" qui permet d'effectuer une recherche sur internet pour compléter la réponse du modèle
 - 🎉 Affichage des sources dans le chat UI
 - 🐛 Les erreurs sont remontées de manière plus claire dans l'upload de fichiers

CONTRIBUTING.md

@@ -28,23 +28,23 @@ feat(collections): collection name retriever
 
 # Tests
 
-Merci avant chaque pull request, de vérifier le bon déploiement de votre API à l'aide en exécutant des tests unitaires.
+Merci, avant chaque pull request, de vérifier le bon déploiement de votre API en exécutant des tests unitaires.
 
-1. Après avoir créer un fichier *config.yml*, lancez l'API en local
+1. Après avoir créé un fichier *config.yml*, lancez l'API en local
 
     ```bash
     uvicorn app.main:app --port 8080 --log-level debug --reload
     ```
 
-2. Executez les tests unitaires
+2. Exécutez les tests unitaires
 
     ```bash
-    PYTHONPATH=. pytest -v --exitfirst app/tests --base-url http://localhost:8080/v1 --api-key-user API_KEY_USER --api-key-admin API_KEY_ADMIN --log-cli-level=INFO
+    PYTHONPATH=. pytest --config-file=pyproject.toml --base-url http://localhost:8080/v1 --api-key-user API_KEY_USER --api-key-admin API_KEY_ADMIN --log-cli-level=INFO
     ```
 
 # Linter
 
-Le linter du projet est [Ruff](https://beta.ruff.rs/docs/configuration/). Les règles de formatages spécifiques au projet sont dans le fichier *[pyproject.toml](./pyproject.toml)*.
+Le linter du projet est [Ruff](https://beta.ruff.rs/docs/configuration/). Les règles de formatage spécifiques au projet sont dans le fichier *[pyproject.toml](./pyproject.toml)*.
 
 ## Configurer Ruff avec pre-commit
 
@@ -61,7 +61,7 @@ Le linter du projet est [Ruff](https://beta.ruff.rs/docs/configuration/). Les r
 1. Installez l'extension *Ruff* (charliermarsh.ruff) dans VSCode
 2. Configurez le linter Ruff dans VSCode pour utiliser le fichier *[pyproject.toml](./pyproject.toml)*
 
-    A l'aide de la commande palette de VSCode (⇧⌘P), recherchez et sélectionnez *Preferences: Open User Settings (JSON)*.
+    À l'aide de la palette de commandes de VSCode (⇧⌘P), recherchez et sélectionnez *Preferences: Open User Settings (JSON)*.
 
     Dans le fichier JSON qui s'ouvre, ajoutez à la fin du fichier les lignes suivantes :
 
@@ -75,6 +75,6 @@ Le linter du projet est [Ruff](https://beta.ruff.rs/docs/configuration/). Les r
     "ruff.nativeServer": "on"
     ```
 
-    ⚠️ **Attention** : Assurez vous que le fichier *[pyproject.toml](./app/pyproject.toml)* est bien spécifié dans la configuration.
+    ⚠️ **Attention** : Assurez-vous que le fichier *[pyproject.toml](./app/pyproject.toml)* est bien spécifié dans la configuration.
 
-3. **Pour exécuter le linter, utilisez la commande palette de VSCode (⇧⌘P) depuis le fichier sur lequel vous voulez l'exécuter, et recherchez et sélectionnez *Ruff: Format document* et *Ruff: Format imports*.**
+3. **Pour exécuter le linter, utilisez la palette de commandes de VSCode (⇧⌘P) depuis le fichier sur lequel vous voulez l'exécuter, puis recherchez et sélectionnez *Ruff: Format document* et *Ruff: Format imports*.**

README.md

@@ -1,15 +1,17 @@
 # Albert API
 ![](https://img.shields.io/badge/python-3.12-green) ![](https://img.shields.io/badge/vLLM-v0.5.5-blue) ![](https://img.shields.io/badge/HuggingFace%20Text%20Embeddings%20Inference-1.5-red)
 
-Albert API est une API open source d'IA générative développée par Etalab. Elle permet d'être un proxy entre des modèles de langage et vos données. Elle aggrège les services suivants :
-- [vLLM](https://github.com/vllm-project/vllm) pour la gestion des modèles de langage
-- [HuggingFace Text Embeddings Inference](https://github.com/huggingface/text-embeddings-inference) pour la génération d'embeddings
-- [Qdrant](https://qdrant.tech/) pour la recherche de similarité
-
-### OpenAI conventions
+Albert API est une API open source d'IA générative développée par Etalab. Elle permet d'être un proxy entre des modèles de langage et vos données. Elle agrège les services suivants :
+- servir des modèles de langage avec [vLLM](https://github.com/vllm-project/vllm)
+- servir des modèles d'embeddings avec [HuggingFace Text Embeddings Inference](https://github.com/huggingface/text-embeddings-inference)
+- accès un *vector store* avec [Qdrant](https://qdrant.tech/) pour la recherche de similarité
 
 En se basant sur les conventions définies par OpenAI, l'API Albert expose des endpoints qui peuvent être appelés avec le [client officiel python d'OpenAI](https://github.com/openai/openai-python/tree/main). Ce formalisme permet d'intégrer facilement l'API Albert avec des bibliothèques tierces comme [Langchain](https://www.langchain.com/) ou [LlamaIndex](https://www.llamaindex.ai/).
 
+## 🚀 Nouveautés
+
+Vous trouverez les changelogs des différentes versions d'Albert API dans le fichier [CHANGELOG.md](./CHANGELOG.md).
+
 ## ⚙️ Fonctionnalités
 
 ### Converser avec un modèle de langage (chat memory)
@@ -30,7 +32,7 @@ L'API Albert permet d'accéder à un ensemble de modèles de langage et d'embedd
 
 ### Interroger vos documents (RAG)
 
-L'API Albert permet d'interroger des documents dans une base vectorielle. Ces documents sont classés dans des collections. Vous pouvez créer vos collections privées et utilisé les collections publiques déjà existantes. Enfin une collection "internet" permet d'effectuer une recherche sur internet pour compléter la réponse du modèle.
+L'API Albert permet d'interroger des documents dans une base vectorielle. Ces documents sont classés dans des collections. Vous pouvez créer vos collections privées et utiliser les collections publiques déjà existantes. Enfin, une collection "internet" permet d'effectuer une recherche sur internet pour compléter la réponse du modèle.
 
 <a target="_blank" href="https://colab.research.google.com/github/etalab-ia/albert-api/blob/main/docs/tutorials/retrival_augmented_generation.ipynb">
   <img src="https://colab.research.google.com/assets/colab-badge.svg" alt="Open In Colab"/>
@@ -46,7 +48,7 @@ L'API Albert permet d'importer sa base de connaissances dans une base vectoriell
 
 ## 🧑‍💻 Contribuez au projet
 
-Albert API est un projet open source, vous pouvez contribuez au projet, veuillez lire notre [guide de contribution](./CONTRIBUTING.md).
+Albert API est un projet open source, vous pouvez contribuer au projet en lisant notre [guide de contribution](./CONTRIBUTING.md).
 
 ## Installation
 

app/endpoints/chat.py

@@ -1,6 +1,6 @@
 from typing import Union
 
-from fastapi import APIRouter, HTTPException, Security
+from fastapi import APIRouter, Security
 from fastapi.responses import StreamingResponse
 import httpx
 
@@ -9,6 +9,7 @@
 from app.utils.lifespan import clients
 from app.utils.security import check_api_key
 from app.utils.variables import LANGUAGE_MODEL_TYPE
+from app.utils.exceptions import WrongModelTypeException, ContextLengthExceededException
 
 router = APIRouter()
 
@@ -22,19 +23,20 @@ async def chat_completions(request: ChatCompletionRequest, user: User = Security
     request = dict(request)
     client = clients.models[request["model"]]
     if client.type != LANGUAGE_MODEL_TYPE:
-        raise HTTPException(status_code=400, detail="Model is not a language model")
+        raise WrongModelTypeException()
 
     url = f"{client.base_url}chat/completions"
     headers = {"Authorization": f"Bearer {client.api_key}"}
 
     if not client.check_context_length(model=request["model"], messages=request["messages"]):
-        raise HTTPException(status_code=400, detail="Context length too large")
+        raise ContextLengthExceededException()
 
     # non stream case
     if not request["stream"]:
         async with httpx.AsyncClient(timeout=20) as async_client:
             response = await async_client.request(method="POST", url=url, headers=headers, json=request)
             response.raise_for_status()
+
             data = response.json()
             return ChatCompletion(**data)
 

app/endpoints/chunks.py

@@ -1,44 +1,29 @@
+from typing import Optional
 from uuid import UUID
 
-from fastapi import APIRouter, HTTPException, Security
-from qdrant_client.http.models import Filter, HasIdCondition
+from fastapi import APIRouter, Security, Query
 
-from app.helpers import VectorStore
-from app.schemas.chunks import Chunk, ChunkRequest, Chunks
+from app.schemas.chunks import Chunks
+from app.schemas.security import User
 from app.utils.lifespan import clients
 from app.utils.security import check_api_key
-from app.schemas.security import User
 
 router = APIRouter()
 
 
-@router.get("/chunks/{collection}/{chunk}")
-async def get_chunk(collection: UUID, chunk: str, user: User = Security(check_api_key)) -> Chunk:
+@router.get("/chunks/{collection}/{document}")
+async def get_chunks(
+    collection: UUID,
+    document: UUID,
+    limit: Optional[int] = Query(default=10, ge=1, le=10),
+    offset: Optional[UUID] = None,
+    user: User = Security(check_api_key),
+) -> Chunks:
     """
     Get a single chunk.
     """
-    collection = str(collection)
-    vectorstore = VectorStore(clients=clients, user=user)
-    ids = [chunk]
-    filter = Filter(must=[HasIdCondition(has_id=ids)])
-    try:
-        chunks = vectorstore.get_chunks(collection_id=collection, filter=filter)
-    except AssertionError as e:
-        raise HTTPException(status_code=400, detail=str(e))
-    return chunks[0]
-
+    collection, document = str(collection), str(document)
+    offset = str(offset) if offset else None
+    data = clients.vectors.get_chunks(collection_id=collection, document_id=document, limit=limit, offset=offset, user=user)
 
-@router.post("/chunks/{collection}")
-async def get_chunks(collection: UUID, request: ChunkRequest, user: User = Security(check_api_key)) -> Chunks:
-    """
-    Get multiple chunks.
-    """
-    collection = str(collection)
-    vectorstore = VectorStore(clients=clients, user=user)
-    ids = request.chunks
-    filter = Filter(must=[HasIdCondition(has_id=ids)])
-    try:
-        chunks = vectorstore.get_chunks(collection_id=collection, filter=filter)
-    except AssertionError as e:
-        raise HTTPException(status_code=400, detail=str(e))
-    return Chunks(data=chunks)
+    return Chunks(data=data)

app/endpoints/collections.py

@@ -1,16 +1,16 @@
-from typing import Literal, Optional, Union
+from typing import Union
 import uuid
 from uuid import UUID
 
-from fastapi import APIRouter, HTTPException, Response, Security
+from fastapi import APIRouter, Response, Security
 from fastapi.responses import JSONResponse
 
-from app.helpers import VectorStore
+
 from app.schemas.collections import Collection, CollectionRequest, Collections
 from app.schemas.security import User
 from app.utils.lifespan import clients
 from app.utils.security import check_api_key
-from app.utils.variables import PUBLIC_COLLECTION_TYPE, INTERNET_COLLECTION_ID
+from app.utils.variables import INTERNET_COLLECTION_ID, PUBLIC_COLLECTION_TYPE
 
 router = APIRouter()
 
@@ -20,23 +20,16 @@ async def create_collection(request: CollectionRequest, user: User = Security(ch
     """
     Create a new collection.
     """
-    vectorstore = VectorStore(clients=clients, user=user)
     collection_id = str(uuid.uuid4())
-    try:
-        vectorstore.create_collection(
-            collection_id=collection_id, collection_name=request.name, collection_model=request.model, collection_type=request.type
-        )
-    except AssertionError as e:
-        raise HTTPException(status_code=400, detail=str(e))
+    clients.vectors.create_collection(
+        collection_id=collection_id, collection_name=request.name, collection_model=request.model, collection_type=request.type, user=user
+    )
 
     return JSONResponse(status_code=201, content={"id": collection_id})
 
 
-@router.get("/collections/{collection}")
 @router.get("/collections")
-async def get_collections(
-    collection: Optional[Union[UUID, Literal["internet"]]] = None, user: User = Security(check_api_key)
-) -> Union[Collection, Collections]:
+async def get_collections(user: User = Security(check_api_key)) -> Union[Collection, Collections]:
     """
     Get list of collections.
     """
@@ -47,21 +40,9 @@ async def get_collections(
         type=PUBLIC_COLLECTION_TYPE,
         description="Use this collection to search on the internet.",
     )
-    if collection == "internet":
-        return internet_collection
-
-    collection_ids = [str(collection)] if collection else []
-    vectorstore = VectorStore(clients=clients, user=user)
-    try:
-        data = vectorstore.get_collection_metadata(collection_ids=collection_ids)
-    except AssertionError as e:
-        # TODO: return a 404 error if collection not found
-        raise HTTPException(status_code=400, detail=str(e))
-
-    if collection:
-        return data[0]
-
+    data = clients.vectors.get_collections(user=user)
     data.append(internet_collection)
+
     return Collections(data=data)
 
 
@@ -71,10 +52,6 @@ async def delete_collections(collection: UUID, user: User = Security(check_api_k
     Delete a collection.
     """
     collection = str(collection)
-    vectorstore = VectorStore(clients=clients, user=user)
-    try:
-        vectorstore.delete_collection(collection_id=collection)
-    except AssertionError as e:
-        raise HTTPException(status_code=400, detail=str(e))
+    clients.vectors.delete_collection(collection_id=collection, user=user)
 
     return Response(status_code=204)

app/endpoints/completions.py

@@ -1,11 +1,12 @@
-from fastapi import APIRouter, HTTPException, Security
+from fastapi import APIRouter, Security
 import httpx
 
 from app.schemas.completions import CompletionRequest, Completions
 from app.schemas.security import User
 from app.utils.lifespan import clients
 from app.utils.security import check_api_key
 from app.utils.variables import LANGUAGE_MODEL_TYPE
+from app.utils.exceptions import WrongModelTypeException, ContextLengthExceededException
 
 router = APIRouter()
 
@@ -21,13 +22,17 @@ async def completions(request: CompletionRequest, user: User = Security(check_ap
     client = clients.models[request["model"]]
 
     if client.type != LANGUAGE_MODEL_TYPE:
-        raise HTTPException(status_code=400, detail="Model is not a language model")
+        raise WrongModelTypeException()
+
+    if not client.check_context_length(model=request["model"], messages=request["messages"]):
+        raise ContextLengthExceededException()
 
     url = f"{client.base_url}completions"
     headers = {"Authorization": f"Bearer {client.api_key}"}
 
     async with httpx.AsyncClient(timeout=20) as async_client:
         response = await async_client.request(method="POST", url=url, headers=headers, json=request)
         response.raise_for_status()
+
         data = response.json()
         return Completions(**data)