diff --git a/.gitignore b/.gitignore index 10ef763..f0168ce 100644 --- a/.gitignore +++ b/.gitignore @@ -138,4 +138,11 @@ local/ pylate/* pylate/ -.byaldi/ \ No newline at end of file +.byaldi/ + +docker/RAG_DB/storage/* +docker/RAG_DB_Milvus/volumes/* +!docker/RAG_DB/storage/.gitkeep +!docker/RAG_DB_Milvus/volumes/etcd/.gitkeep +!docker/RAG_DB_Milvus/volumes/minio/.gitkeep +!docker/RAG_DB_Milvus/volumes/milvus/.gitkeep \ No newline at end of file diff --git a/docker/RAG_DB/.env.example b/docker/RAG_DB/.env.example new file mode 100644 index 0000000..df82a90 --- /dev/null +++ b/docker/RAG_DB/.env.example @@ -0,0 +1,4 @@ +QDRANT_CONTAINER_NAME=rag-db +QDRANT_HTTP_PORT=6333 +QDRANT_GRPC_PORT=6334 +QDRANT_STORAGE_PATH=./storage diff --git a/docker/RAG_DB/README.md b/docker/RAG_DB/README.md new file mode 100644 index 0000000..decaa7a --- /dev/null +++ b/docker/RAG_DB/README.md @@ -0,0 +1,129 @@ +# RAG_DB + +Déploiement dédié de Qdrant pour la base vectorielle. + +## Contenu + +- `docker-compose.yml` : démarrage du service Qdrant +- `.env.example` : configuration des ports et du stockage +- `storage/` : volume local par défaut + +## Variables importantes + +- `QDRANT_HTTP_PORT` : port HTTP exposé pour les clients +- `QDRANT_GRPC_PORT` : port gRPC exposé +- `QDRANT_STORAGE_PATH` : chemin local du stockage persistant +- `QDRANT_CONTAINER_NAME` : nom du container + +## Démarrage + +```bash +cd RAG_DB +cp .env.example .env +docker compose up -d +``` + +## Accès + +- HTTP : `http://:${QDRANT_HTTP_PORT}` +- gRPC : `:${QDRANT_GRPC_PORT}` + +## Exemples Python + +Pour tester Qdrant directement en Python : + +```bash +python3 -m pip install qdrant-client numpy +``` + +### Stocker des multi-vecteurs ColQwen + +Qdrant sait stocker nativement un multi-vecteur par point. C'est le cas le plus simple si vous voulez garder un comportement de type late-interaction avec `MAX_SIM`. + +```python +from qdrant_client import QdrantClient, models + +client = QdrantClient(url="http://localhost:6333") + +collection_name = "colqwen_pages" +vector_size = 128 + +if client.collection_exists(collection_name): + client.delete_collection(collection_name) + +client.create_collection( + collection_name=collection_name, + vectors_config=models.VectorParams( + size=vector_size, + distance=models.Distance.COSINE, + multivector_config=models.MultiVectorConfig( + comparator=models.MultiVectorComparator.MAX_SIM, + ), + ), +) + +page_vectors = [ + [0.12, 0.44, -0.05, 0.91], + [0.18, 0.39, -0.07, 0.88], +] + +client.upsert( + collection_name=collection_name, + points=[ + models.PointStruct( + id="page-1", + vector=page_vectors, + payload={ + "filename": "manuel.pdf", + "page_number": 1, + "chunk_text": "Exemple de page indexee", + }, + ) + ], + wait=True, +) +``` + +### Rechercher avec `MAX_SIM` + +Avec un multi-vecteur de query, le score est calculé dans Qdrant. En pratique, chaque vecteur de la query prend la meilleure similarité contre les vecteurs du document, puis ces meilleurs scores sont additionnés. + +```python +query_vectors = [ + [0.10, 0.41, -0.02, 0.93], + [0.16, 0.37, -0.08, 0.86], +] + +response = client.query_points( + collection_name=collection_name, + query=query_vectors, + limit=5, + with_payload=True, +) + +for hit in response.points: + print(hit.id, hit.score, hit.payload) +``` + +### Calcul `MaxSim` en Python pur + +Si vous voulez vérifier le score côté application : + +```python +import numpy as np + +def maxsim(query_vectors: list[list[float]], doc_vectors: list[list[float]]) -> float: + q = np.asarray(query_vectors, dtype=np.float32) + d = np.asarray(doc_vectors, dtype=np.float32) + similarities = q @ d.T + return float(similarities.max(axis=1).sum()) + +score = maxsim(query_vectors, page_vectors) +print(score) +``` + +Ce calcul correspond à l'idée de `MultiVectorComparator.MAX_SIM` utilisée par Qdrant. + +## Remarques + +Ce module peut être déployé seul sur un serveur distinct. L'orchestrateur devra simplement connaître l'URL HTTP Qdrant, par exemple `http://10.0.0.20:6333`. diff --git a/docker/RAG_DB/docker-compose.yml b/docker/RAG_DB/docker-compose.yml new file mode 100644 index 0000000..23e3c75 --- /dev/null +++ b/docker/RAG_DB/docker-compose.yml @@ -0,0 +1,15 @@ +services: + qdrant: + image: qdrant/qdrant:latest + container_name: ${QDRANT_CONTAINER_NAME:-rag-db} + ports: + - "${QDRANT_HTTP_PORT:-6333}:6333" + - "${QDRANT_GRPC_PORT:-6334}:6334" + volumes: + - ${QDRANT_STORAGE_PATH:-./storage}:/qdrant/storage + networks: + - rag-network + +networks: + rag-network: + external: true diff --git a/docker/RAG_DB/storage/.gitkeep b/docker/RAG_DB/storage/.gitkeep new file mode 100644 index 0000000..8b13789 --- /dev/null +++ b/docker/RAG_DB/storage/.gitkeep @@ -0,0 +1 @@ + diff --git a/docker/RAG_DB_Milvus/.env.example b/docker/RAG_DB_Milvus/.env.example new file mode 100644 index 0000000..042ceef --- /dev/null +++ b/docker/RAG_DB_Milvus/.env.example @@ -0,0 +1,18 @@ +MILVUS_STANDALONE_CONTAINER_NAME=rag-db-milvus +MILVUS_ETCD_CONTAINER_NAME=rag-db-milvus-etcd +MILVUS_MINIO_CONTAINER_NAME=rag-db-milvus-minio + +MILVUS_VERSION=v2.6.15 + +MILVUS_PORT=19530 +MILVUS_HEALTH_PORT=9091 +MILVUS_MINIO_API_PORT=9000 +MILVUS_MINIO_CONSOLE_PORT=9001 + +MILVUS_STORAGE_PATH=./volumes/milvus +MILVUS_ETCD_PATH=./volumes/etcd +MILVUS_MINIO_PATH=./volumes/minio + +MINIO_ROOT_USER=minioadmin +MINIO_ROOT_PASSWORD=minioadmin +MINIO_REGION=us-east-1 diff --git a/docker/RAG_DB_Milvus/README.md b/docker/RAG_DB_Milvus/README.md new file mode 100644 index 0000000..e49d213 --- /dev/null +++ b/docker/RAG_DB_Milvus/README.md @@ -0,0 +1,190 @@ +# RAG_DB_Milvus + +Déploiement dédié de Milvus en mode standalone pour la base vectorielle. + +## Contenu + +- `docker-compose.yml` : démarrage de Milvus et de ses dépendances +- `.env.example` : configuration des ports, des volumes et des noms de conteneurs +- `volumes/` : volumes locaux persistants pour `etcd`, `minio` et `milvus` + +## Variables importantes + +- `MILVUS_PORT` : port Milvus exposé pour les clients +- `MILVUS_HEALTH_PORT` : port HTTP de healthcheck et de Web UI Milvus +- `MILVUS_MINIO_API_PORT` : port API MinIO +- `MILVUS_MINIO_CONSOLE_PORT` : port console MinIO +- `MILVUS_STORAGE_PATH` : chemin local du stockage Milvus +- `MILVUS_ETCD_PATH` : chemin local du stockage etcd +- `MILVUS_MINIO_PATH` : chemin local du stockage MinIO +- `MILVUS_VERSION` : version de l'image Milvus +- `MINIO_ROOT_USER` / `MINIO_ROOT_PASSWORD` : credentials MinIO + +## Démarrage + +```bash +docker network create rag-network +cd RAG_DB_Milvus +cp .env.example .env +docker compose up -d +``` + +## Accès + +- Milvus : `localhost:${MILVUS_PORT}` pour les clients +- Health / Web UI : `http://localhost:${MILVUS_HEALTH_PORT}/healthz` +- MinIO API : `http://localhost:${MILVUS_MINIO_API_PORT}` +- MinIO Console : `http://localhost:${MILVUS_MINIO_CONSOLE_PORT}` + +## Exemples Python + +Pour tester Milvus directement en Python : + +```bash +python3 -m pip install pymilvus==2.6.11 numpy +``` + +### Stocker des vecteurs de page + +Milvus ne gère pas nativement le multi-vecteur `MAX_SIM` comme Qdrant. Dans ce repo, on utilise donc deux niveaux : + +- une collection de pages pour récupérer des candidats rapidement +- une collection de tokens pour faire le reranking late-interaction ensuite + +Exemple minimal pour stocker des vecteurs de page : + +```python +from pymilvus import MilvusClient, DataType + +client = MilvusClient(uri="http://localhost:19530") + +collection_name = "colqwen_pages" +vector_size = 4 + +if collection_name in client.list_collections(): + client.drop_collection(collection_name=collection_name) + +schema = MilvusClient.create_schema(auto_id=False, enable_dynamic_field=False) +schema.add_field(field_name="id", datatype=DataType.VARCHAR, is_primary=True, max_length=64) +schema.add_field(field_name="page_vector", datatype=DataType.FLOAT_VECTOR, dim=vector_size) +schema.add_field(field_name="filename", datatype=DataType.VARCHAR, max_length=256) +schema.add_field(field_name="page_number", datatype=DataType.INT64) + +index_params = client.prepare_index_params() +index_params.add_index( + field_name="page_vector", + index_type="AUTOINDEX", + metric_type="COSINE", +) + +client.create_collection( + collection_name=collection_name, + schema=schema, + index_params=index_params, +) + +client.insert( + collection_name=collection_name, + data=[ + { + "id": "page-1", + "page_vector": [0.14, 0.42, -0.04, 0.90], + "filename": "manuel.pdf", + "page_number": 1, + } + ], +) +``` + +### Rechercher des pages candidates + +```python +results = client.search( + collection_name=collection_name, + data=[[0.11, 0.40, -0.03, 0.92]], + anns_field="page_vector", + limit=5, + output_fields=["filename", "page_number"], + search_params={"metric_type": "COSINE"}, +) + +for hit in results[0]: + print(hit["id"], hit["distance"], hit["entity"]) +``` + +### Stocker les vecteurs token-level pour le reranking + +Pour approcher `MAX_SIM`, vous pouvez stocker un vecteur par token avec un `page_id` commun : + +```python +token_collection = "colqwen_tokens" + +if token_collection in client.list_collections(): + client.drop_collection(collection_name=token_collection) + +schema = MilvusClient.create_schema(auto_id=False, enable_dynamic_field=False) +schema.add_field(field_name="id", datatype=DataType.VARCHAR, is_primary=True, max_length=64) +schema.add_field(field_name="page_id", datatype=DataType.VARCHAR, max_length=64) +schema.add_field(field_name="token_vector", datatype=DataType.FLOAT_VECTOR, dim=vector_size) + +index_params = client.prepare_index_params() +index_params.add_index( + field_name="token_vector", + index_type="AUTOINDEX", + metric_type="COSINE", +) + +client.create_collection( + collection_name=token_collection, + schema=schema, + index_params=index_params, +) + +client.insert( + collection_name=token_collection, + data=[ + {"id": "tok-1", "page_id": "page-1", "token_vector": [0.12, 0.44, -0.05, 0.91]}, + {"id": "tok-2", "page_id": "page-1", "token_vector": [0.18, 0.39, -0.07, 0.88]}, + ], +) +``` + +### Calcul `MaxSim` en Python pour un candidat Milvus + +Une fois les tokens d'une page candidate récupérés, le `MaxSim` se fait côté application : + +```python +import numpy as np + +query_vectors = [ + [0.10, 0.41, -0.02, 0.93], + [0.16, 0.37, -0.08, 0.86], +] + +doc_vectors = [ + [0.12, 0.44, -0.05, 0.91], + [0.18, 0.39, -0.07, 0.88], +] + +def maxsim(query_vectors: list[list[float]], doc_vectors: list[list[float]]) -> float: + q = np.asarray(query_vectors, dtype=np.float32) + d = np.asarray(doc_vectors, dtype=np.float32) + similarities = q @ d.T + return float(similarities.max(axis=1).sum()) + +score = maxsim(query_vectors, doc_vectors) +print(score) +``` + +Dans `RAG_Orch`, c'est cette idée qui est utilisée avec Milvus : recherche de candidats sur les pages, puis reranking token-level. + +## Remarques + +Ce module suit la structure de `RAG_DB`, mais Milvus nécessite aussi `etcd` et `minio`. + +`RAG_Orch` peut maintenant l'utiliser directement avec une configuration du type : + +```env +VECTOR_DB_BACKEND=milvus +MILVUS_URL=http://rag-db-milvus:19530 +``` diff --git a/docker/RAG_DB_Milvus/docker-compose.yml b/docker/RAG_DB_Milvus/docker-compose.yml new file mode 100644 index 0000000..3171855 --- /dev/null +++ b/docker/RAG_DB_Milvus/docker-compose.yml @@ -0,0 +1,83 @@ +services: + milvus-etcd: + image: quay.io/coreos/etcd:v3.5.25 + container_name: ${MILVUS_ETCD_CONTAINER_NAME:-rag-db-milvus-etcd} + environment: + ETCD_AUTO_COMPACTION_MODE: revision + ETCD_AUTO_COMPACTION_RETENTION: 1000 + ETCD_QUOTA_BACKEND_BYTES: 4294967296 + ETCD_SNAPSHOT_COUNT: 50000 + command: + - etcd + - -advertise-client-urls=http://milvus-etcd:2379 + - -listen-client-urls + - http://0.0.0.0:2379 + - --data-dir + - /etcd + volumes: + - ${MILVUS_ETCD_PATH:-./volumes/etcd}:/etcd + healthcheck: + test: ["CMD", "etcdctl", "endpoint", "health"] + interval: 30s + timeout: 20s + retries: 3 + networks: + - rag-network + + milvus-minio: + image: minio/minio:RELEASE.2024-05-28T17-19-04Z + container_name: ${MILVUS_MINIO_CONTAINER_NAME:-rag-db-milvus-minio} + environment: + MINIO_ROOT_USER: ${MINIO_ROOT_USER:-minioadmin} + MINIO_ROOT_PASSWORD: ${MINIO_ROOT_PASSWORD:-minioadmin} + command: + - minio + - server + - /minio_data + - --console-address + - :9001 + ports: + - "${MILVUS_MINIO_API_PORT:-9000}:9000" + - "${MILVUS_MINIO_CONSOLE_PORT:-9001}:9001" + volumes: + - ${MILVUS_MINIO_PATH:-./volumes/minio}:/minio_data + healthcheck: + test: ["CMD", "mc", "ready", "local"] + interval: 30s + timeout: 20s + retries: 3 + networks: + - rag-network + + milvus-standalone: + image: milvusdb/milvus:${MILVUS_VERSION:-v2.6.15} + container_name: ${MILVUS_STANDALONE_CONTAINER_NAME:-rag-db-milvus} + command: ["milvus", "run", "standalone"] + security_opt: + - seccomp:unconfined + environment: + MINIO_REGION: ${MINIO_REGION:-us-east-1} + ETCD_ENDPOINTS: milvus-etcd:2379 + MINIO_ADDRESS: milvus-minio:9000 + ports: + - "${MILVUS_PORT:-19530}:19530" + - "${MILVUS_HEALTH_PORT:-9091}:9091" + volumes: + - ${MILVUS_STORAGE_PATH:-./volumes/milvus}:/var/lib/milvus + depends_on: + milvus-etcd: + condition: service_healthy + milvus-minio: + condition: service_healthy + healthcheck: + test: ["CMD", "curl", "-f", "http://localhost:9091/healthz"] + interval: 30s + start_period: 90s + timeout: 20s + retries: 3 + networks: + - rag-network + +networks: + rag-network: + external: true diff --git a/docker/RAG_Orch/.env.milvus b/docker/RAG_Orch/.env.milvus new file mode 100644 index 0000000..61fa343 --- /dev/null +++ b/docker/RAG_Orch/.env.milvus @@ -0,0 +1,11 @@ +ORCH_CONTAINER_NAME=rag-orch +ORCH_HOST_PORT=8000 + +EMBEDDING_SERVICE_URL=http://host.docker.internal:8001 + +VECTOR_DB_BACKEND=milvus +MILVUS_URL=http://rag-db-milvus:19530 +MILVUS_CANDIDATE_LIMIT=64 +VECTOR_DB_UPSERT_BATCH_SIZE=16 + +DEFAULT_PDF_DPI=150 diff --git a/docker/RAG_Orch/.env.qdrant b/docker/RAG_Orch/.env.qdrant new file mode 100644 index 0000000..35269fe --- /dev/null +++ b/docker/RAG_Orch/.env.qdrant @@ -0,0 +1,10 @@ +ORCH_CONTAINER_NAME=rag-orch +ORCH_HOST_PORT=8000 + +EMBEDDING_SERVICE_URL=http://host.docker.internal:8001 + +VECTOR_DB_BACKEND=qdrant +QDRANT_URL=http://rag-db:6333 +VECTOR_DB_UPSERT_BATCH_SIZE=16 + +DEFAULT_PDF_DPI=150 diff --git a/docker/RAG_Orch/Dockerfile b/docker/RAG_Orch/Dockerfile new file mode 100644 index 0000000..27df97c --- /dev/null +++ b/docker/RAG_Orch/Dockerfile @@ -0,0 +1,17 @@ +FROM python:3.11-slim + +ENV PYTHONDONTWRITEBYTECODE=1 +ENV PYTHONUNBUFFERED=1 + +WORKDIR /app + +RUN apt-get update \ + && apt-get install -y --no-install-recommends libgl1 libglib2.0-0 \ + && rm -rf /var/lib/apt/lists/* + +COPY requirements.txt . +RUN pip install --no-cache-dir -r requirements.txt + +COPY app ./app + +CMD ["uvicorn", "app.main:app", "--host", "0.0.0.0", "--port", "8000"] diff --git a/docker/RAG_Orch/README.md b/docker/RAG_Orch/README.md new file mode 100644 index 0000000..ad77c1e --- /dev/null +++ b/docker/RAG_Orch/README.md @@ -0,0 +1,67 @@ +# RAG_Orch + +Service d'orchestration du pipeline RAG. Il ne porte pas le modèle et ne porte pas la base : il appelle `RAG_Model` pour les embeddings et un backend vectoriel externe pour l'indexation et le retrieval. + +## Contenu + +- `Dockerfile` +- `docker-compose.yml` +- `.env.example` +- `requirements.txt` +- `app/main.py` + +## Endpoints + +- `GET /health` +- `POST /v1/index` +- `POST /v1/retrieve` + +## Variables importantes + +- `ORCH_HOST_PORT` : port HTTP exposé +- `EMBEDDING_SERVICE_URL` : URL du service `RAG_Model` +- `VECTOR_DB_BACKEND` : backend vectoriel à utiliser, `milvus` par défaut, `qdrant` en option +- `MILVUS_URL` : URL du service Milvus +- `MILVUS_CANDIDATE_LIMIT` : nombre de pages candidates retenues avant reranking late-interaction +- `VECTOR_DB_UPSERT_BATCH_SIZE` : taille des batchs d'upsert côté backend vectoriel +- `QDRANT_URL` : URL HTTP du service Qdrant si `VECTOR_DB_BACKEND=qdrant` +- `DEFAULT_PDF_DPI` : DPI de rasterisation + +## Backends pris en charge + +`RAG_Orch` supporte maintenant deux backends : + +- `milvus` : backend par défaut +- `qdrant` : backend conservé pour compatibilité + +Avec Milvus, l'orchestrateur crée deux collections techniques par collection logique : + +- une collection `__pages` pour les vecteurs agrégés par page +- une collection `__tokens` pour les vecteurs token-level utilisés au reranking + +Le retrieval Milvus fonctionne en deux temps : + +1. récupération d'un ensemble de pages candidates via un vecteur de page agrégé +2. reranking late-interaction sur les vecteurs token-level pour approcher le comportement multivector précédemment assuré par Qdrant + +## Démarrage + +### Qdrant +```bash +cd RAG_Orch +cp .env.qdrant .env +docker compose up --build +``` + +### Milvus +```bash +cd RAG_Orch +cp .env.milvus .env +docker compose up --build +``` + +## Exemple de test rapide + +```bash +curl http://localhost:8000/health +``` diff --git a/docker/RAG_Orch/app/__init__.py b/docker/RAG_Orch/app/__init__.py new file mode 100644 index 0000000..8b13789 --- /dev/null +++ b/docker/RAG_Orch/app/__init__.py @@ -0,0 +1 @@ + diff --git a/docker/RAG_Orch/app/config.py b/docker/RAG_Orch/app/config.py new file mode 100644 index 0000000..be5d832 --- /dev/null +++ b/docker/RAG_Orch/app/config.py @@ -0,0 +1,56 @@ +import os +from dataclasses import dataclass +from functools import lru_cache + + +PAYLOAD_JSON_MAX_LENGTH = 65535 +MILVUS_PRIMARY_KEY_MAX_LENGTH = 64 +MILVUS_COLLECTION_SUFFIX_PAGES = "__pages" +MILVUS_COLLECTION_SUFFIX_TOKENS = "__tokens" +MILVUS_PAGE_VECTOR_FIELD = "page_vector" +MILVUS_TOKEN_VECTOR_FIELD = "token_vector" +MILVUS_TOKEN_GROUP_FIELD = "page_id" + + +@dataclass(frozen=True) +class Settings: + embedding_service_url: str + vector_db_backend: str + qdrant_url: str + milvus_url: str + milvus_token: str | None + default_pdf_dpi: int + upsert_batch_size: int + milvus_candidate_limit: int + http_timeout_seconds: float + + +@lru_cache(maxsize=1) +def get_settings() -> Settings: + # Read and cache the runtime configuration exposed through environment variables. + return Settings( + embedding_service_url=os.getenv("EMBEDDING_SERVICE_URL", "http://localhost:8001"), + vector_db_backend=os.getenv("VECTOR_DB_BACKEND", "milvus").strip().lower(), + qdrant_url=os.getenv("QDRANT_URL", "http://localhost:6333"), + milvus_url=os.getenv("MILVUS_URL", "http://localhost:19530"), + milvus_token=os.getenv("MILVUS_TOKEN") or None, + default_pdf_dpi=int(os.getenv("DEFAULT_PDF_DPI", "150")), + upsert_batch_size=int( + os.getenv( + "VECTOR_DB_UPSERT_BATCH_SIZE", + os.getenv("QDRANT_UPSERT_BATCH_SIZE", "16"), + ) + ), + milvus_candidate_limit=int(os.getenv("MILVUS_CANDIDATE_LIMIT", "64")), + http_timeout_seconds=float(os.getenv("HTTP_TIMEOUT_SECONDS", "300")), + ) + + +def get_vector_db_backend() -> str: + # Validate and return the configured vector database backend name. + backend = get_settings().vector_db_backend + if backend not in {"milvus", "qdrant"}: + raise ValueError( + f"Unsupported VECTOR_DB_BACKEND={backend!r}. Expected 'milvus' or 'qdrant'." + ) + return backend diff --git a/docker/RAG_Orch/app/embeddings.py b/docker/RAG_Orch/app/embeddings.py new file mode 100644 index 0000000..4e5f674 --- /dev/null +++ b/docker/RAG_Orch/app/embeddings.py @@ -0,0 +1,47 @@ +import io +from typing import Any + +import httpx +from PIL import Image + +from app.config import get_settings + + +async def call_embedding_api_for_query(query: str) -> dict[str, Any]: + # Request query embeddings from the external embedding service. + timeout = httpx.Timeout(get_settings().http_timeout_seconds) + async with httpx.AsyncClient(timeout=timeout) as client: + response = await client.post( + f"{get_settings().embedding_service_url}/v1/embed/query", + json={"query": query}, + ) + response.raise_for_status() + return response.json() + + +async def call_embedding_api_for_page(image: Image.Image, filename: str) -> dict[str, Any]: + # Request page embeddings from the external embedding service using a PNG upload. + buffer = io.BytesIO() + image.save(buffer, format="PNG") + files = {"file": (filename, buffer.getvalue(), "image/png")} + + timeout = httpx.Timeout(get_settings().http_timeout_seconds) + async with httpx.AsyncClient(timeout=timeout) as client: + response = await client.post( + f"{get_settings().embedding_service_url}/v1/embed/page", + files=files, + ) + response.raise_for_status() + return response.json() + + +async def check_embedding_api_health() -> str: + # Probe the embedding service health endpoint and return a compact status string. + try: + timeout = httpx.Timeout(5.0) + async with httpx.AsyncClient(timeout=timeout) as client: + response = await client.get(f"{get_settings().embedding_service_url}/health") + response.raise_for_status() + except Exception as exc: + return f"error: {exc.__class__.__name__}" + return "ok" diff --git a/docker/RAG_Orch/app/helpers.py b/docker/RAG_Orch/app/helpers.py new file mode 100644 index 0000000..98018f3 --- /dev/null +++ b/docker/RAG_Orch/app/helpers.py @@ -0,0 +1,118 @@ +import json +import math +import re +import uuid +from pathlib import Path + +import fitz +from fastapi import UploadFile + +from app.config import ( + MILVUS_COLLECTION_SUFFIX_PAGES, + MILVUS_COLLECTION_SUFFIX_TOKENS, + PAYLOAD_JSON_MAX_LENGTH, +) + + +def sanitize_identifier(value: str, fallback: str) -> str: + # Normalize an arbitrary name into a safe identifier fragment for documents and collections. + cleaned = re.sub(r"[^a-zA-Z0-9_-]+", "-", value.strip().lower()).strip("-") + return cleaned or fallback + + +def default_collection_name(files: list[UploadFile]) -> str: + # Build the default logical collection name from the uploaded file list. + if len(files) == 1 and files[0].filename: + return f"{Path(files[0].filename).stem}_colqwen35" + return "corpus_colqwen35" + + +def extract_page_text(page: fitz.Page) -> str: + # Extract and compact textual content from a PDF page for payload storage. + text = page.get_text("text") + return re.sub(r"\s+", " ", text).strip() + + +def page_identifier(document_id: str, page_number: int) -> str: + # Generate a deterministic page identifier so Milvus page and token rows stay linked. + return str(uuid.uuid5(uuid.NAMESPACE_URL, f"{document_id}:{page_number}")) + + +def qdrant_point_id() -> str: + # Generate a random point identifier for a Qdrant record. + return str(uuid.uuid4()) + + +def milvus_safe_collection_base_name(logical_name: str) -> str: + # Convert a logical collection name into a Milvus-compatible collection prefix. + cleaned = re.sub(r"[^a-zA-Z0-9_]+", "_", logical_name.strip().lower()).strip("_") + cleaned = cleaned or "corpus_colqwen35" + if cleaned[0].isdigit(): + cleaned = f"c_{cleaned}" + return cleaned[:180] + + +def milvus_collection_names(logical_name: str) -> tuple[str, str]: + # Derive the page and token collection names used by the Milvus late-interaction layout. + base_name = milvus_safe_collection_base_name(logical_name) + return ( + f"{base_name}{MILVUS_COLLECTION_SUFFIX_PAGES}", + f"{base_name}{MILVUS_COLLECTION_SUFFIX_TOKENS}", + ) + + +def serialize_payload(payload: dict[str, object]) -> str: + # Serialize payload data into a bounded JSON string accepted by the Milvus schema. + serialized = json.dumps(payload, ensure_ascii=False) + if len(serialized) > PAYLOAD_JSON_MAX_LENGTH: + serialized = json.dumps( + { + **payload, + "chunk_text": str(payload.get("chunk_text", ""))[:60000], + "chunk_preview": str(payload.get("chunk_preview", ""))[:500], + "_payload_truncated": True, + }, + ensure_ascii=False, + ) + if len(serialized) > PAYLOAD_JSON_MAX_LENGTH: + serialized = serialized[:PAYLOAD_JSON_MAX_LENGTH] + return serialized + + +def deserialize_payload(payload_json: str | None) -> dict[str, object]: + # Decode a Milvus payload JSON string back into a Python dictionary. + if not payload_json: + return {} + try: + payload = json.loads(payload_json) + if isinstance(payload, dict): + return payload + except json.JSONDecodeError: + pass + return {"payload_json": payload_json} + + +def mean_pool_vectors(vectors: list[list[float]]) -> list[float]: + # Compute a normalized mean-pooled vector used for the Milvus page-level candidate search. + if not vectors: + raise ValueError("Cannot pool an empty vector list.") + + vector_size = len(vectors[0]) + sums = [0.0] * vector_size + for vector in vectors: + if len(vector) != vector_size: + raise ValueError("Inconsistent vector sizes in embedding payload.") + for index, value in enumerate(vector): + sums[index] += float(value) + + pooled = [value / len(vectors) for value in sums] + norm = math.sqrt(sum(value * value for value in pooled)) + if norm > 0: + return [value / norm for value in pooled] + return pooled + + +def milvus_quote_string(value: str) -> str: + # Escape and quote a string value so it can be embedded safely in a Milvus filter expression. + escaped = value.replace("\\", "\\\\").replace('"', '\\"') + return f'"{escaped}"' diff --git a/docker/RAG_Orch/app/indexing.py b/docker/RAG_Orch/app/indexing.py new file mode 100644 index 0000000..d19ec26 --- /dev/null +++ b/docker/RAG_Orch/app/indexing.py @@ -0,0 +1,353 @@ +from pathlib import Path +from typing import Any + +import fitz +import httpx +from fastapi import HTTPException, UploadFile +from PIL import Image + +from app.config import get_settings, get_vector_db_backend +from app.embeddings import call_embedding_api_for_page +from app.helpers import default_collection_name, extract_page_text, sanitize_identifier +from app.progress import set_index_progress +from app.schemas import IndexingFileResult, IndexingResponse +from app.vector_store import ( + PendingIndexBatch, + add_index_record, + ensure_collection, + flush_index_batch, + pending_page_count, +) + + +def build_page_payload( + document_id: str, + filename: str, + page_number: int, + machine_modele: str, + dpi: int, + embedding: dict[str, Any], + page_text: str, +) -> dict[str, Any]: + # Assemble the payload stored next to each indexed page. + return { + "document_id": document_id, + "filename": filename, + "page_number": page_number, + "machine_modele": machine_modele, + "dpi": dpi, + "model_name": embedding["model_name"], + "chunk_text": page_text, + "chunk_preview": page_text[:500], + **embedding.get("metadata", {}), + } + + +def initialize_progress( + progress_id: str | None, + collection_name: str, + total_files: int, +) -> None: + # Seed the progress tracker with the initial state of a new indexing request. + if progress_id: + set_index_progress( + progress_id, + status="preparing", + collection_name=collection_name, + total_files=total_files, + files_processed=0, + pages_indexed=0, + total_pages_estimate=0, + points_upserted=0, + error=None, + ) + + +def update_file_start_progress( + progress_id: str | None, + filename: str, + file_page_count: int, + total_pages_estimate: int, +) -> None: + # Update the progress tracker when a new file starts indexing. + if progress_id: + set_index_progress( + progress_id, + status="running", + current_filename=filename, + current_page=0, + current_file_page_count=file_page_count, + total_pages_estimate=total_pages_estimate, + ) + + +def update_page_progress( + progress_id: str | None, + filename: str, + page_number: int, + file_page_count: int, + file_index: int, + total_pages_indexed: int, + total_pages_estimate: int, + points_upserted: int, +) -> None: + # Update the progress tracker after one page has been prepared locally. + if progress_id: + set_index_progress( + progress_id, + status="running", + current_filename=filename, + current_page=page_number, + current_file_page_count=file_page_count, + files_processed=file_index - 1, + pages_indexed=total_pages_indexed, + total_pages_estimate=total_pages_estimate, + points_upserted=points_upserted, + ) + + +def update_flush_progress( + progress_id: str | None, + filename: str, + page_number: int, + file_page_count: int, + file_index: int, + total_pages_indexed: int, + total_pages_estimate: int, + points_upserted: int, +) -> None: + # Update the progress tracker after a backend flush has completed. + if progress_id: + set_index_progress( + progress_id, + status="running", + current_filename=filename, + current_page=page_number, + current_file_page_count=file_page_count, + files_processed=file_index - 1, + pages_indexed=total_pages_indexed, + total_pages_estimate=total_pages_estimate, + points_upserted=points_upserted, + ) + + +def update_file_done_progress( + progress_id: str | None, + filename: str, + file_page_count: int, + file_index: int, + total_pages_indexed: int, + total_pages_estimate: int, + points_upserted: int, +) -> None: + # Update the progress tracker once a file has been fully processed. + if progress_id: + set_index_progress( + progress_id, + status="running", + current_filename=filename, + current_page=file_page_count, + current_file_page_count=file_page_count, + files_processed=file_index, + pages_indexed=total_pages_indexed, + total_pages_estimate=total_pages_estimate, + points_upserted=points_upserted, + ) + + +def update_failure_progress(progress_id: str | None, error: str) -> None: + # Mark the current indexing request as failed in the progress tracker. + if progress_id: + set_index_progress(progress_id, status="failed", error=error) + + +def update_completion_progress( + progress_id: str | None, + total_files: int, + total_pages_indexed: int, + total_pages_estimate: int, + points_upserted: int, +) -> None: + # Mark the current indexing request as completed in the progress tracker. + if progress_id: + set_index_progress( + progress_id, + status="completed", + current_page=0, + current_file_page_count=0, + files_processed=total_files, + pages_indexed=total_pages_indexed, + total_pages_estimate=total_pages_estimate, + points_upserted=points_upserted, + ) + + +async def index_documents_service( + files: list[UploadFile], + collection_name: str | None, + machine_modele: str, + dpi: int | None, + recreate_collection: bool, + progress_id: str | None, +) -> IndexingResponse: + # Index the uploaded PDF files into the configured vector backend and return the API response. + if not files: + raise HTTPException(status_code=400, detail="At least one PDF file is required.") + + backend = get_vector_db_backend() + effective_dpi = dpi or get_settings().default_pdf_dpi + resolved_collection_name = collection_name or default_collection_name(files) + file_results: list[IndexingFileResult] = [] + pending_batch = PendingIndexBatch() + total_pages_indexed = 0 + points_upserted = 0 + recreate_pending = recreate_collection + collection_ready = False + total_pages_estimate = 0 + + initialize_progress(progress_id, resolved_collection_name, len(files)) + + for file_index, upload in enumerate(files, start=1): + filename = upload.filename or f"document-{file_index}.pdf" + if not filename.lower().endswith(".pdf"): + raise HTTPException(status_code=400, detail=f"{filename} is not a PDF file.") + + file_bytes = await upload.read() + document_id = f"{file_index}-{sanitize_identifier(Path(filename).stem, 'document')}" + + try: + document = fitz.open(stream=file_bytes, filetype="pdf") + except Exception as exc: + raise HTTPException( + status_code=400, + detail=f"Unable to open {filename} as a PDF: {exc}", + ) from exc + + file_page_count = len(document) + total_pages_estimate += file_page_count + pages_indexed = 0 + update_file_start_progress( + progress_id, + filename, + file_page_count, + total_pages_estimate, + ) + + try: + for page_offset in range(file_page_count): + page_number = page_offset + 1 + page = document[page_offset] + pix = page.get_pixmap(dpi=effective_dpi) + image = Image.frombytes("RGB", [pix.width, pix.height], pix.samples) + page_text = extract_page_text(page) + + embedding = await call_embedding_api_for_page( + image=image, + filename=f"{Path(filename).stem}-page-{page_number}.png", + ) + vectors = embedding["vectors"] + if not vectors: + raise HTTPException( + status_code=500, + detail=f"Embedding API returned no vectors for {filename} page {page_number}.", + ) + + if not collection_ready: + ensure_collection( + collection_name=resolved_collection_name, + vector_size=len(vectors[0]), + recreate=recreate_pending, + ) + collection_ready = True + recreate_pending = False + + payload = build_page_payload( + document_id=document_id, + filename=filename, + page_number=page_number, + machine_modele=machine_modele, + dpi=effective_dpi, + embedding=embedding, + page_text=page_text, + ) + add_index_record( + batch=pending_batch, + backend=backend, + document_id=document_id, + page_number=page_number, + vectors=vectors, + payload=payload, + ) + + pages_indexed += 1 + total_pages_indexed += 1 + + update_page_progress( + progress_id=progress_id, + filename=filename, + page_number=page_number, + file_page_count=file_page_count, + file_index=file_index, + total_pages_indexed=total_pages_indexed, + total_pages_estimate=total_pages_estimate, + points_upserted=points_upserted + pending_page_count(pending_batch, backend), + ) + + if pending_page_count(pending_batch, backend) >= get_settings().upsert_batch_size: + points_upserted += flush_index_batch( + resolved_collection_name, + pending_batch, + backend, + ) + update_flush_progress( + progress_id=progress_id, + filename=filename, + page_number=page_number, + file_page_count=file_page_count, + file_index=file_index, + total_pages_indexed=total_pages_indexed, + total_pages_estimate=total_pages_estimate, + points_upserted=points_upserted, + ) + except httpx.HTTPStatusError as exc: + error_message = f"Embedding API error while indexing {filename}: {exc.response.text}" + update_failure_progress(progress_id, error_message) + raise HTTPException(status_code=502, detail=error_message) from exc + except Exception as exc: + update_failure_progress(progress_id, str(exc)) + raise + finally: + document.close() + + file_results.append( + IndexingFileResult( + filename=filename, + document_id=document_id, + pages_indexed=pages_indexed, + ) + ) + update_file_done_progress( + progress_id=progress_id, + filename=filename, + file_page_count=file_page_count, + file_index=file_index, + total_pages_indexed=total_pages_indexed, + total_pages_estimate=total_pages_estimate, + points_upserted=points_upserted + pending_page_count(pending_batch, backend), + ) + + points_upserted += flush_index_batch(resolved_collection_name, pending_batch, backend) + update_completion_progress( + progress_id=progress_id, + total_files=len(files), + total_pages_indexed=total_pages_indexed, + total_pages_estimate=total_pages_estimate, + points_upserted=points_upserted, + ) + + return IndexingResponse( + collection_name=resolved_collection_name, + files=file_results, + total_pages_indexed=total_pages_indexed, + points_upserted=points_upserted, + ) diff --git a/docker/RAG_Orch/app/main.py b/docker/RAG_Orch/app/main.py new file mode 100644 index 0000000..13e7b8c --- /dev/null +++ b/docker/RAG_Orch/app/main.py @@ -0,0 +1,97 @@ +import httpx +from fastapi import FastAPI, File, Form, HTTPException, UploadFile + +from app.config import get_settings, get_vector_db_backend +from app.embeddings import call_embedding_api_for_query, check_embedding_api_health +from app.indexing import index_documents_service +from app.progress import get_index_progress +from app.schemas import ( + HealthResponse, + IndexProgressResponse, + IndexingResponse, + RetrievalHit, + RetrievalRequest, + RetrievalResponse, +) +from app.vector_store import check_vector_db_health, search_points + + +app = FastAPI(title="RAG Orchestrator API", version="0.2.0") + + +@app.get("/health", response_model=HealthResponse) +async def health() -> HealthResponse: + # Report the health of the embedding service and the configured vector database backend. + embedding_status = await check_embedding_api_health() + vector_db_status = check_vector_db_health() + status = "ok" if embedding_status == "ok" and vector_db_status == "ok" else "degraded" + return HealthResponse( + status=status, + embedding_api=embedding_status, + vector_db_backend=get_vector_db_backend(), + vector_db=vector_db_status, + ) + + +@app.post("/v1/index", response_model=IndexingResponse) +async def index_documents( + files: list[UploadFile] = File(...), + collection_name: str | None = Form(default=None), + machine_modele: str = Form(default="inconnu"), + dpi: int = Form(default=get_settings().default_pdf_dpi), + recreate_collection: bool = Form(default=False), + progress_id: str | None = Form(default=None), +) -> IndexingResponse: + # Index uploaded PDF documents into the currently selected vector backend. + return await index_documents_service( + files=files, + collection_name=collection_name, + machine_modele=machine_modele, + dpi=dpi, + recreate_collection=recreate_collection, + progress_id=progress_id, + ) + + +@app.get("/v1/index/progress/{progress_id}", response_model=IndexProgressResponse) +async def get_index_progress_endpoint(progress_id: str) -> IndexProgressResponse: + # Return the latest progress snapshot for a given indexing job. + progress = get_index_progress(progress_id) + if progress is None: + raise HTTPException(status_code=404, detail="Unknown progress_id.") + return progress + + +@app.post("/v1/retrieve", response_model=RetrievalResponse) +async def retrieve(payload: RetrievalRequest) -> RetrievalResponse: + # Embed the query, search the backend, and normalize the results into the public API model. + try: + query_embedding = await call_embedding_api_for_query(payload.query) + results = search_points( + collection_name=payload.collection_name, + query_vectors=query_embedding["vectors"], + limit=payload.limit, + ) + except httpx.HTTPStatusError as exc: + raise HTTPException( + status_code=502, + detail=f"Embedding API error: {exc.response.text}", + ) from exc + except HTTPException: + raise + except Exception as exc: + raise HTTPException(status_code=500, detail=str(exc)) from exc + + return RetrievalResponse( + collection_name=payload.collection_name, + query=payload.query, + limit=payload.limit, + hits=[ + RetrievalHit( + point_id=result.point_id, + score=result.score, + payload=result.payload, + ) + for result in results + ], + ) diff --git a/docker/RAG_Orch/app/progress.py b/docker/RAG_Orch/app/progress.py new file mode 100644 index 0000000..dabaffc --- /dev/null +++ b/docker/RAG_Orch/app/progress.py @@ -0,0 +1,36 @@ +import threading +import time +from typing import Any + +from app.schemas import IndexProgressResponse + + +INDEX_PROGRESS: dict[str, IndexProgressResponse] = {} +INDEX_PROGRESS_LOCK = threading.Lock() + + +def set_index_progress(progress_id: str, **updates: Any) -> None: + # Merge new progress information into the in-memory tracker for an index job. + now = time.time() + with INDEX_PROGRESS_LOCK: + current = INDEX_PROGRESS.get(progress_id) + if current is None: + current = IndexProgressResponse( + progress_id=progress_id, + status="pending", + started_at=now, + updated_at=now, + ) + data = current.model_dump() + data.update(updates) + data["progress_id"] = progress_id + data["updated_at"] = now + if data.get("started_at") is None: + data["started_at"] = now + INDEX_PROGRESS[progress_id] = IndexProgressResponse(**data) + + +def get_index_progress(progress_id: str) -> IndexProgressResponse | None: + # Return the latest known state for an index job, if it exists. + with INDEX_PROGRESS_LOCK: + return INDEX_PROGRESS.get(progress_id) diff --git a/docker/RAG_Orch/app/schemas.py b/docker/RAG_Orch/app/schemas.py new file mode 100644 index 0000000..e9d9d4e --- /dev/null +++ b/docker/RAG_Orch/app/schemas.py @@ -0,0 +1,59 @@ +from typing import Any + +from pydantic import BaseModel, Field + + +class IndexingFileResult(BaseModel): + filename: str + document_id: str + pages_indexed: int + + +class IndexingResponse(BaseModel): + collection_name: str + files: list[IndexingFileResult] + total_pages_indexed: int + points_upserted: int + + +class RetrievalRequest(BaseModel): + collection_name: str + query: str + limit: int = Field(default=5, ge=1, le=100) + + +class RetrievalHit(BaseModel): + point_id: str + score: float + payload: dict[str, Any] = Field(default_factory=dict) + + +class RetrievalResponse(BaseModel): + collection_name: str + query: str + limit: int + hits: list[RetrievalHit] + + +class HealthResponse(BaseModel): + status: str + embedding_api: str + vector_db_backend: str + vector_db: str + + +class IndexProgressResponse(BaseModel): + progress_id: str + status: str + collection_name: str | None = None + current_filename: str | None = None + current_page: int = 0 + current_file_page_count: int = 0 + files_processed: int = 0 + total_files: int = 0 + pages_indexed: int = 0 + total_pages_estimate: int = 0 + points_upserted: int = 0 + started_at: float | None = None + updated_at: float | None = None + error: str | None = None diff --git a/docker/RAG_Orch/app/vector_store.py b/docker/RAG_Orch/app/vector_store.py new file mode 100644 index 0000000..f4cd0f2 --- /dev/null +++ b/docker/RAG_Orch/app/vector_store.py @@ -0,0 +1,429 @@ +import uuid +from dataclasses import dataclass, field +from functools import lru_cache +from typing import Any + +from pymilvus import DataType, MilvusClient +from qdrant_client import QdrantClient +from qdrant_client.models import ( + Distance, + MultiVectorComparator, + MultiVectorConfig, + PointStruct, + VectorParams, +) + +from app.config import ( + MILVUS_PAGE_VECTOR_FIELD, + MILVUS_PRIMARY_KEY_MAX_LENGTH, + MILVUS_TOKEN_GROUP_FIELD, + MILVUS_TOKEN_VECTOR_FIELD, + PAYLOAD_JSON_MAX_LENGTH, + get_settings, + get_vector_db_backend, +) +from app.helpers import ( + deserialize_payload, + mean_pool_vectors, + milvus_collection_names, + milvus_quote_string, + page_identifier, + qdrant_point_id, + serialize_payload, +) + + +@dataclass +class PendingIndexBatch: + qdrant_points: list[PointStruct] = field(default_factory=list) + milvus_pages: list[dict[str, Any]] = field(default_factory=list) + milvus_tokens: list[dict[str, Any]] = field(default_factory=list) + + +@dataclass +class SearchHit: + point_id: str + score: float + payload: dict[str, Any] + + +@lru_cache(maxsize=1) +def get_qdrant_client() -> QdrantClient: + # Create and cache the Qdrant client used by the orchestrator process. + return QdrantClient(url=get_settings().qdrant_url) + + +@lru_cache(maxsize=1) +def get_milvus_client() -> MilvusClient: + # Create and cache the Milvus client used by the orchestrator process. + client_kwargs: dict[str, Any] = {"uri": get_settings().milvus_url} + if get_settings().milvus_token: + client_kwargs["token"] = get_settings().milvus_token + return MilvusClient(**client_kwargs) + + +def ensure_qdrant_collection(collection_name: str, vector_size: int, recreate: bool) -> None: + # Create or recreate a Qdrant multivector collection for ColQwen page storage. + client = get_qdrant_client() + + if recreate and client.collection_exists(collection_name): + client.delete_collection(collection_name) + + if not client.collection_exists(collection_name): + client.create_collection( + collection_name=collection_name, + vectors_config=VectorParams( + size=vector_size, + distance=Distance.COSINE, + multivector_config=MultiVectorConfig( + comparator=MultiVectorComparator.MAX_SIM, + ), + ), + ) + + +def milvus_collection_exists(client: MilvusClient, collection_name: str) -> bool: + # Check whether a Milvus collection already exists. + return collection_name in client.list_collections() + + +def ensure_milvus_collection_loaded(collection_name: str) -> None: + # Load a Milvus collection in memory before inserts or searches. + client = get_milvus_client() + client.load_collection(collection_name=collection_name, replica_number=1) + + +def create_milvus_page_collection(collection_name: str, vector_size: int) -> None: + # Create the Milvus page-level candidate collection used before reranking. + client = get_milvus_client() + + schema = MilvusClient.create_schema(auto_id=False, enable_dynamic_field=False) + schema.add_field( + field_name="id", + datatype=DataType.VARCHAR, + is_primary=True, + max_length=MILVUS_PRIMARY_KEY_MAX_LENGTH, + ) + schema.add_field( + field_name=MILVUS_PAGE_VECTOR_FIELD, + datatype=DataType.FLOAT_VECTOR, + dim=vector_size, + ) + schema.add_field( + field_name="payload_json", + datatype=DataType.VARCHAR, + max_length=PAYLOAD_JSON_MAX_LENGTH, + ) + + index_params = client.prepare_index_params() + index_params.add_index( + field_name=MILVUS_PAGE_VECTOR_FIELD, + index_type="AUTOINDEX", + metric_type="COSINE", + ) + + client.create_collection( + collection_name=collection_name, + schema=schema, + index_params=index_params, + ) + + +def create_milvus_token_collection(collection_name: str, vector_size: int) -> None: + # Create the Milvus token-level collection used for late-interaction reranking. + client = get_milvus_client() + + schema = MilvusClient.create_schema(auto_id=False, enable_dynamic_field=False) + schema.add_field( + field_name="id", + datatype=DataType.VARCHAR, + is_primary=True, + max_length=MILVUS_PRIMARY_KEY_MAX_LENGTH, + ) + schema.add_field( + field_name=MILVUS_TOKEN_GROUP_FIELD, + datatype=DataType.VARCHAR, + max_length=MILVUS_PRIMARY_KEY_MAX_LENGTH, + ) + schema.add_field( + field_name=MILVUS_TOKEN_VECTOR_FIELD, + datatype=DataType.FLOAT_VECTOR, + dim=vector_size, + ) + + index_params = client.prepare_index_params() + index_params.add_index( + field_name=MILVUS_TOKEN_VECTOR_FIELD, + index_type="AUTOINDEX", + metric_type="COSINE", + ) + + client.create_collection( + collection_name=collection_name, + schema=schema, + index_params=index_params, + ) + + +def ensure_milvus_collections(collection_name: str, vector_size: int, recreate: bool) -> None: + # Create or recreate the pair of Milvus collections required by the late-interaction layout. + client = get_milvus_client() + page_collection, token_collection = milvus_collection_names(collection_name) + + if recreate: + if milvus_collection_exists(client, token_collection): + client.drop_collection(collection_name=token_collection) + if milvus_collection_exists(client, page_collection): + client.drop_collection(collection_name=page_collection) + + if not milvus_collection_exists(client, page_collection): + create_milvus_page_collection(page_collection, vector_size) + if not milvus_collection_exists(client, token_collection): + create_milvus_token_collection(token_collection, vector_size) + + ensure_milvus_collection_loaded(page_collection) + ensure_milvus_collection_loaded(token_collection) + + +def ensure_collection(collection_name: str, vector_size: int, recreate: bool) -> None: + # Ensure the logical collection exists on the configured backend before indexing begins. + if get_vector_db_backend() == "qdrant": + ensure_qdrant_collection(collection_name, vector_size, recreate) + return + ensure_milvus_collections(collection_name, vector_size, recreate) + + +def add_index_record( + batch: PendingIndexBatch, + backend: str, + document_id: str, + page_number: int, + vectors: list[list[float]], + payload: dict[str, Any], +) -> None: + # Append one indexed page to the backend-specific in-memory batch structure. + if backend == "qdrant": + batch.qdrant_points.append( + PointStruct( + id=qdrant_point_id(), + vector=vectors, + payload=payload, + ) + ) + return + + page_id = page_identifier(document_id, page_number) + batch.milvus_pages.append( + { + "id": page_id, + MILVUS_PAGE_VECTOR_FIELD: mean_pool_vectors(vectors), + "payload_json": serialize_payload(payload), + } + ) + batch.milvus_tokens.extend( + { + "id": str(uuid.uuid4()), + MILVUS_TOKEN_GROUP_FIELD: page_id, + MILVUS_TOKEN_VECTOR_FIELD: vector, + } + for vector in vectors + ) + + +def pending_page_count(batch: PendingIndexBatch, backend: str) -> int: + # Return the number of pending page records regardless of the configured backend. + if backend == "qdrant": + return len(batch.qdrant_points) + return len(batch.milvus_pages) + + +def flush_qdrant_points(collection_name: str, points: list[PointStruct]) -> int: + # Send the accumulated Qdrant points to the server and clear the local batch. + if not points: + return 0 + + client = get_qdrant_client() + client.upsert(collection_name=collection_name, points=points, wait=True) + count = len(points) + points.clear() + return count + + +def flush_milvus_rows( + collection_name: str, + page_rows: list[dict[str, Any]], + token_rows: list[dict[str, Any]], +) -> int: + # Send the accumulated Milvus page and token rows to the server and clear the local batch. + if not page_rows: + return 0 + + client = get_milvus_client() + page_collection, token_collection = milvus_collection_names(collection_name) + ensure_milvus_collection_loaded(page_collection) + ensure_milvus_collection_loaded(token_collection) + + client.insert(collection_name=page_collection, data=page_rows) + if token_rows: + client.insert(collection_name=token_collection, data=token_rows) + + count = len(page_rows) + page_rows.clear() + token_rows.clear() + return count + + +def flush_index_batch(collection_name: str, batch: PendingIndexBatch, backend: str) -> int: + # Flush the backend-specific batch and return the number of indexed pages committed. + if backend == "qdrant": + return flush_qdrant_points(collection_name, batch.qdrant_points) + return flush_milvus_rows(collection_name, batch.milvus_pages, batch.milvus_tokens) + + +def qdrant_search_points( + collection_name: str, + query_vectors: list[list[float]], + limit: int, +) -> list[SearchHit]: + # Execute a Qdrant multivector search and normalize the response into search hits. + client = get_qdrant_client() + + if hasattr(client, "query_points"): + response = client.query_points( + collection_name=collection_name, + query=query_vectors, + limit=limit, + with_payload=True, + ) + results = list(getattr(response, "points", response)) + else: + results = list( + client.search( + collection_name=collection_name, + query_vector=query_vectors, + limit=limit, + with_payload=True, + ) + ) + + return [ + SearchHit( + point_id=str(hit.id), + score=float(hit.score), + payload=hit.payload or {}, + ) + for hit in results + ] + + +def milvus_page_candidates( + collection_name: str, + query_vectors: list[list[float]], + limit: int, +) -> tuple[list[str], dict[str, dict[str, Any]]]: + # Search the Milvus page collection to build the candidate set for reranking. + client = get_milvus_client() + page_collection, _token_collection = milvus_collection_names(collection_name) + ensure_milvus_collection_loaded(page_collection) + + candidate_limit = max(limit * 10, get_settings().milvus_candidate_limit) + pooled_query = mean_pool_vectors(query_vectors) + results = client.search( + collection_name=page_collection, + data=[pooled_query], + anns_field=MILVUS_PAGE_VECTOR_FIELD, + limit=candidate_limit, + output_fields=["payload_json"], + search_params={"metric_type": "COSINE"}, + ) + + rows = results[0] if results else [] + candidate_ids: list[str] = [] + candidate_payloads: dict[str, dict[str, Any]] = {} + for row in rows: + page_id = str(row.get("id", "")) + entity = row.get("entity") or {} + if not page_id: + continue + candidate_ids.append(page_id) + candidate_payloads[page_id] = deserialize_payload(entity.get("payload_json")) + return candidate_ids, candidate_payloads + + +def milvus_search_points( + collection_name: str, + query_vectors: list[list[float]], + limit: int, +) -> list[SearchHit]: + # Execute the Milvus two-stage retrieval flow and normalize the response into search hits. + candidate_ids, candidate_payloads = milvus_page_candidates(collection_name, query_vectors, limit) + if not candidate_ids: + return [] + + client = get_milvus_client() + _page_collection, token_collection = milvus_collection_names(collection_name) + ensure_milvus_collection_loaded(token_collection) + + filter_expression = ( + f'{MILVUS_TOKEN_GROUP_FIELD} in [{", ".join(milvus_quote_string(page_id) for page_id in candidate_ids)}]' + ) + + aggregated_scores = {page_id: 0.0 for page_id in candidate_ids} + for query_vector in query_vectors: + grouped_results = client.search( + collection_name=token_collection, + data=[query_vector], + anns_field=MILVUS_TOKEN_VECTOR_FIELD, + filter=filter_expression, + limit=min(len(candidate_ids), 16384), + output_fields=[MILVUS_TOKEN_GROUP_FIELD], + search_params={"metric_type": "COSINE"}, + group_by_field=MILVUS_TOKEN_GROUP_FIELD, + ) + + for row in (grouped_results[0] if grouped_results else []): + entity = row.get("entity") or {} + page_id = str(entity.get(MILVUS_TOKEN_GROUP_FIELD, "")) + if not page_id: + continue + aggregated_scores[page_id] = aggregated_scores.get(page_id, 0.0) + float( + row.get("distance", 0.0) + ) + + ranked_page_ids = sorted( + aggregated_scores.items(), + key=lambda item: item[1], + reverse=True, + )[:limit] + + return [ + SearchHit( + point_id=page_id, + score=float(score), + payload=candidate_payloads.get(page_id, {}), + ) + for page_id, score in ranked_page_ids + ] + + +def search_points( + collection_name: str, + query_vectors: list[list[float]], + limit: int, +) -> list[SearchHit]: + # Dispatch the search request to the configured backend implementation. + if get_vector_db_backend() == "qdrant": + return qdrant_search_points(collection_name, query_vectors, limit) + return milvus_search_points(collection_name, query_vectors, limit) + + +def check_vector_db_health() -> str: + # Probe the configured vector backend and return a compact status string. + try: + if get_vector_db_backend() == "qdrant": + get_qdrant_client().get_collections() + else: + get_milvus_client().list_collections() + except Exception as exc: + return f"error: {exc.__class__.__name__}" + return "ok" diff --git a/docker/RAG_Orch/docker-compose.yml b/docker/RAG_Orch/docker-compose.yml new file mode 100644 index 0000000..dca8318 --- /dev/null +++ b/docker/RAG_Orch/docker-compose.yml @@ -0,0 +1,21 @@ +services: + rag-orch: + build: + context: . + container_name: ${ORCH_CONTAINER_NAME:-rag-orch} + ports: + - "${ORCH_HOST_PORT:-8000}:8000" + environment: + EMBEDDING_SERVICE_URL: ${EMBEDDING_SERVICE_URL:-http://localhost:8001} + VECTOR_DB_BACKEND: ${VECTOR_DB_BACKEND:-milvus} + MILVUS_URL: ${MILVUS_URL:-http://rag-db-milvus:19530} + MILVUS_CANDIDATE_LIMIT: ${MILVUS_CANDIDATE_LIMIT:-64} + QDRANT_URL: ${QDRANT_URL:-http://localhost:6333} + DEFAULT_PDF_DPI: ${DEFAULT_PDF_DPI:-150} + VECTOR_DB_UPSERT_BATCH_SIZE: ${VECTOR_DB_UPSERT_BATCH_SIZE:-16} + networks: + - rag-network + +networks: + rag-network: + external: true diff --git a/docker/RAG_Orch/requirements.txt b/docker/RAG_Orch/requirements.txt new file mode 100644 index 0000000..35ba450 --- /dev/null +++ b/docker/RAG_Orch/requirements.txt @@ -0,0 +1,8 @@ +fastapi +httpx +pillow +pymilvus==2.6.11 +pymupdf +python-multipart +qdrant-client +uvicorn[standard]