diff --git a/index/Dockerfile b/index/Dockerfile index 4ec50f7..f36ed37 100644 --- a/index/Dockerfile +++ b/index/Dockerfile @@ -5,7 +5,7 @@ WORKDIR /app COPY requirements.txt . RUN pip install --no-cache-dir -r requirements.txt -COPY main.py . +COPY *.py . ENV HOST=0.0.0.0 ENV PORT=8000 diff --git a/index/main.py b/index/main.py index 0b4bd92..eb1497c 100644 --- a/index/main.py +++ b/index/main.py @@ -3,13 +3,19 @@ import logging import os from concurrent.futures import ThreadPoolExecutor from contextlib import asynccontextmanager -from functools import lru_cache from typing import Any from fastapi import FastAPI, Request from fastapi.exceptions import RequestValidationError -from fastapi.responses import JSONResponse, ORJSONResponse -from pydantic import BaseModel +from fastapi.responses import JSONResponse + +from chunking import build_chunks +from index_schemas import ( + IndexAPIRequest, + IndexAPIResponse, + SparseEmbeddingRequest, +) +from sparse import embed_sparse_texts, get_sparse_model HOST = os.getenv("HOST", "0.0.0.0") PORT = int(os.getenv("PORT", "8000")) @@ -18,199 +24,20 @@ UVICORN_WORKERS = 4 logging.basicConfig(level=os.getenv("LOG_LEVEL", "INFO")) logger = logging.getLogger("index-service") - -class Chat(BaseModel): - id: str - name: str - sn: str - type: str - is_public: bool | None = None - members_count: int | None = None - members: list[dict[str, Any]] | None = None - - -class Message(BaseModel): - id: str - thread_sn: str | None = None - time: int - text: str - sender_id: str - file_snippets: str - parts: list[dict[str, Any]] | None = None - mentions: list[str] | None = None - member_event: dict[str, Any] | None = None - is_system: bool - is_hidden: bool - is_forward: bool - is_quote: bool - - -class ChatData(BaseModel): - chat: Chat - overlap_messages: list[Message] - new_messages: list[Message] - - -class IndexAPIRequest(BaseModel): - data: ChatData - - -class IndexAPIItem(BaseModel): - page_content: str - dense_content: str - sparse_content: str - message_ids: list[str] - - -class IndexAPIResponse(BaseModel): - results: list[IndexAPIItem] - - -class SparseEmbeddingRequest(BaseModel): - texts: list[str] - - -class SparseVector(BaseModel): - indices: list[int] - values: list[float] - - -CHUNK_SIZE = 256 -OVERLAP_SIZE = 128 -SPARSE_MODEL_NAME = "Qdrant/bm25" -FASTEMBED_CACHE_PATH = "/models/fastembed" - _thread_pool = ThreadPoolExecutor(max_workers=4) -@lru_cache(maxsize=1) -def get_sparse_model(): - from fastembed import SparseTextEmbedding - - logger.info("Loading sparse model %s from cache %s", SPARSE_MODEL_NAME, FASTEMBED_CACHE_PATH) - return SparseTextEmbedding(model_name=SPARSE_MODEL_NAME) - - @asynccontextmanager async def lifespan(app: FastAPI): - # Preload BM25 model on startup to avoid cold-start latency await asyncio.to_thread(get_sparse_model) logger.info("BM25 model preloaded") yield -def render_message(message: Message) -> str: - parts_list: list[str] = [] - - if message.sender_id: - sender_name = message.sender_id.split("@")[0].replace(".", " ") - parts_list.append(f"[{sender_name}]:") - - if message.text: - parts_list.append(message.text) - - if message.parts: - for part in message.parts: - media_type = part.get("mediaType", "text") - part_text = part.get("text") - if isinstance(part_text, str) and part_text: - if media_type == "forward": - parts_list.append(f"[Пересланное]: {part_text}") - elif media_type == "quote": - parts_list.append(f"[Цитата]: {part_text}") - else: - parts_list.append(part_text) - - if message.file_snippets: - parts_list.append(f"[Файл]: {message.file_snippets}") - - return " ".join(parts_list).strip() - - -def build_chunks( - chat: Chat, - overlap_messages: list[Message], - new_messages: list[Message], -) -> list[IndexAPIItem]: - new_messages = [m for m in new_messages if not m.is_system and not m.is_hidden] - overlap_messages = [m for m in overlap_messages if not m.is_system and not m.is_hidden] - - result: list[IndexAPIItem] = [] - - def build_text_and_ranges(messages: list[Message]) -> tuple[str, list[tuple[int, int, str]]]: - text_parts: list[str] = [] - message_ranges: list[tuple[int, int, str]] = [] - position = 0 - - for index, message in enumerate(messages): - text = render_message(message) - if not text: - continue - - if index > 0 and text_parts: - text_parts.append("\n") - position += 1 - - start = position - text_parts.append(text) - position += len(text) - message_ranges.append((start, position, message.id)) - - return "".join(text_parts), message_ranges - - def slice_tail(text: str, tail_size: int) -> str: - if tail_size <= 0: - return "" - tail_start = max(0, len(text) - tail_size) - return text[tail_start:] - - overlap_text, _ = build_text_and_ranges(overlap_messages) - previous_chunk_text = slice_tail(overlap_text, OVERLAP_SIZE) - - new_text, new_message_ranges = build_text_and_ranges(new_messages) - - for start in range(0, len(new_text), CHUNK_SIZE): - chunk_body = new_text[start: start + CHUNK_SIZE] - if not chunk_body: - continue - - chunk_body_ranges = [ - ( - max(message_start, start) - start, - min(message_end, start + len(chunk_body)) - start, - message_id, - ) - for message_start, message_end, message_id in new_message_ranges - if message_end > start and message_start < start + len(chunk_body) - ] - - chunk_overlap = previous_chunk_text - chunk_text = chunk_overlap - if chunk_text and chunk_body: - chunk_text += "\n" - chunk_text += chunk_body - - dense_text = f"[{chat.name}] {chunk_text}" - sparse_text = chunk_body - - result.append( - IndexAPIItem( - page_content=chunk_text, - dense_content=dense_text, - sparse_content=sparse_text, - message_ids=[message_id for _, _, message_id in chunk_body_ranges], - ) - ) - previous_chunk_text = slice_tail(chunk_text, OVERLAP_SIZE) - - return result - - app = FastAPI( title="Index Service", version="0.1.0", lifespan=lifespan, - default_response_class=ORJSONResponse, ) @@ -223,32 +50,18 @@ async def health() -> dict[str, str]: async def index(payload: IndexAPIRequest) -> IndexAPIResponse: return IndexAPIResponse( results=build_chunks( - payload.data.chat, payload.data.overlap_messages, payload.data.new_messages, ) ) -def embed_sparse_texts(texts: list[str]) -> list[dict]: - model = get_sparse_model() - vectors = [] - for item in model.embed(texts): - vectors.append( - { - "indices": item.indices.tolist(), - "values": item.values.tolist(), - } - ) - return vectors - - @app.post("/sparse_embedding") async def sparse_embedding(payload: SparseEmbeddingRequest) -> dict[str, Any]: vectors = await asyncio.get_event_loop().run_in_executor( _thread_pool, embed_sparse_texts, payload.texts ) - return {"vectors": vectors} + return {"vectors": [{"indices": v.indices, "values": v.values} for v in vectors]} @app.exception_handler(Exception) diff --git a/index/requirements.txt b/index/requirements.txt index a71d4c8..a8683a5 100644 --- a/index/requirements.txt +++ b/index/requirements.txt @@ -2,4 +2,3 @@ fastapi==0.135.1 uvicorn[standard]==0.42.0 pydantic==2.12.5 fastembed==0.7.4 -orjson==3.10.18 diff --git a/search/Dockerfile b/search/Dockerfile index 1954c62..e3b5814 100644 --- a/search/Dockerfile +++ b/search/Dockerfile @@ -5,7 +5,7 @@ WORKDIR /app COPY requirements.txt . RUN pip install --no-cache-dir -r requirements.txt -COPY main.py . +COPY *.py . ENV HOST=0.0.0.0 ENV PORT=8000 diff --git a/search/config.py b/search/config.py index b95e578..140fd87 100644 --- a/search/config.py +++ b/search/config.py @@ -19,10 +19,10 @@ QDRANT_SPARSE_VECTOR_NAME = os.getenv("QDRANT_SPARSE_VECTOR_NAME", "sparse") OPEN_API_LOGIN = os.getenv("OPEN_API_LOGIN") OPEN_API_PASSWORD = os.getenv("OPEN_API_PASSWORD") -DENSE_PREFETCH_K = 50 -SPARSE_PREFETCH_K = 100 -RETRIEVE_K = 80 -RERANK_LIMIT = 60 +DENSE_PREFETCH_K = 80 +SPARSE_PREFETCH_K = 200 +RETRIEVE_K = 150 +RERANK_LIMIT = 15 TOP_K = 50 HTTP_TIMEOUT = 30.0 diff --git a/search/main.py b/search/main.py index 7fd3847..d2553a2 100644 --- a/search/main.py +++ b/search/main.py @@ -2,163 +2,47 @@ import asyncio import logging import os from contextlib import asynccontextmanager -from functools import lru_cache -from typing import Any import httpx -from fastembed import SparseTextEmbedding from fastapi import FastAPI, HTTPException, Request from fastapi.exceptions import RequestValidationError -from fastapi.responses import JSONResponse, ORJSONResponse -from pydantic import BaseModel, Field -from qdrant_client import AsyncQdrantClient, models - -EMBEDDINGS_DENSE_MODEL = "Qwen/Qwen3-Embedding-0.6B" - -HOST = os.getenv("HOST", "0.0.0.0") -PORT = int(os.getenv("PORT", "8000")) - -API_KEY = os.getenv("API_KEY") -EMBEDDINGS_DENSE_URL = os.getenv("EMBEDDINGS_DENSE_URL") -QDRANT_DENSE_VECTOR_NAME = os.getenv("QDRANT_DENSE_VECTOR_NAME", "dense") -QDRANT_SPARSE_VECTOR_NAME = os.getenv("QDRANT_SPARSE_VECTOR_NAME", "sparse") -SPARSE_MODEL_NAME = "Qdrant/bm25" -RERANKER_MODEL = "nvidia/llama-nemotron-rerank-1b-v2" -RERANKER_URL = os.getenv("RERANKER_URL") -OPEN_API_LOGIN = os.getenv("OPEN_API_LOGIN") -OPEN_API_PASSWORD = os.getenv("OPEN_API_PASSWORD") -QDRANT_URL = os.getenv("QDRANT_URL") -QDRANT_COLLECTION_NAME = os.getenv("QDRANT_COLLECTION_NAME", "evaluation") -REQUIRED_ENV_VARS = [ - "EMBEDDINGS_DENSE_URL", - "RERANKER_URL", - "QDRANT_URL", -] - -logging.basicConfig(level=os.getenv("LOG_LEVEL", "INFO")) -logger = logging.getLogger("search-service") - - -def validate_required_env() -> None: - if bool(OPEN_API_LOGIN) != bool(OPEN_API_PASSWORD): - raise RuntimeError("OPEN_API_LOGIN and OPEN_API_PASSWORD must be set together") - - if not API_KEY and not (OPEN_API_LOGIN and OPEN_API_PASSWORD): - raise RuntimeError("Either API_KEY or OPEN_API_LOGIN and OPEN_API_PASSWORD must be set") - - missing_env_vars = [ - name for name in REQUIRED_ENV_VARS if os.getenv(name) is None or os.getenv(name) == "" - ] - if not missing_env_vars: - return - - logger.error("Empty required env vars: %s", ", ".join(missing_env_vars)) - raise RuntimeError(f"Empty required env vars: {', '.join(missing_env_vars)}") +from fastapi.responses import JSONResponse +from qdrant_client import AsyncQdrantClient +from aggregation import aggregate_message_ids +from config import ( + API_KEY, + HOST, + PORT, + QDRANT_URL, + validate_required_env, + logger, +) +from query_builder import ( + build_primary_query, + build_extra_dense_queries, + build_sparse_query, + embed_dense, + embed_dense_batch, + embed_sparse, + get_sparse_model, +) +from rerank import rerank_points +from retrieval import qdrant_search +from schemas import SearchAPIItem, SearchAPIRequest, SearchAPIResponse validate_required_env() -def get_upstream_request_kwargs() -> dict[str, Any]: - headers = {"Content-Type": "application/json"} - kwargs: dict[str, Any] = {"headers": headers} - - if OPEN_API_LOGIN and OPEN_API_PASSWORD: - kwargs["auth"] = (OPEN_API_LOGIN, OPEN_API_PASSWORD) - return kwargs - - if API_KEY: - headers["Authorization"] = f"Bearer {API_KEY}" - - return kwargs - - -class DateRange(BaseModel): - from_: str = Field(alias="from") - to: str - - -class Entities(BaseModel): - people: list[str] | None = None - emails: list[str] | None = None - documents: list[str] | None = None - names: list[str] | None = None - links: list[str] | None = None - - -class Question(BaseModel): - text: str - asker: str = "" - asked_on: str = "" - variants: list[str] | None = None - hyde: list[str] | None = None - keywords: list[str] | None = None - entities: Entities | None = None - date_mentions: list[str] | None = None - date_range: DateRange | None = None - search_text: str = "" - - -class SearchAPIRequest(BaseModel): - question: Question - - -class SearchAPIItem(BaseModel): - message_ids: list[str] - - -class SearchAPIResponse(BaseModel): - results: list[SearchAPIItem] - - -class DenseEmbeddingItem(BaseModel): - index: int - embedding: list[float] - - -class DenseEmbeddingResponse(BaseModel): - data: list[DenseEmbeddingItem] - - -class SparseVector(BaseModel): - indices: list[int] = Field(default_factory=list) - values: list[float] = Field(default_factory=list) - - -class ChunkMetadata(BaseModel): - chat_name: str - chat_type: str - chat_id: str - chat_sn: str - thread_sn: str | None = None - message_ids: list[str] - start: str - end: str - participants: list[str] = Field(default_factory=list) - mentions: list[str] = Field(default_factory=list) - contains_forward: bool = False - contains_quote: bool = False - - -@lru_cache(maxsize=1) -def get_sparse_model() -> SparseTextEmbedding: - logger.info("Loading local sparse model %s", SPARSE_MODEL_NAME) - return SparseTextEmbedding(model_name=SPARSE_MODEL_NAME) - - @asynccontextmanager async def lifespan(app: FastAPI): - # Preload BM25 and set up HTTP client with connection pooling await asyncio.to_thread(get_sparse_model) logger.info("BM25 model preloaded") app.state.http = httpx.AsyncClient( timeout=30.0, limits=httpx.Limits(max_connections=100, max_keepalive_connections=20), ) - app.state.qdrant = AsyncQdrantClient( - url=QDRANT_URL, - api_key=API_KEY, - ) + app.state.qdrant = AsyncQdrantClient(url=QDRANT_URL, api_key=API_KEY) try: yield finally: @@ -170,181 +54,8 @@ app = FastAPI( title="Search Service", version="0.1.0", lifespan=lifespan, - default_response_class=ORJSONResponse, ) -DENSE_PREFETCH_K = 80 -SPARSE_PREFETCH_K = 200 -RETRIEVE_K = 150 -RERANK_LIMIT = 15 - - -async def embed_dense(client: httpx.AsyncClient, text: str) -> list[float]: - response = await client.post( - EMBEDDINGS_DENSE_URL, - **get_upstream_request_kwargs(), - json={ - "model": os.getenv("EMBEDDINGS_DENSE_MODEL", EMBEDDINGS_DENSE_MODEL), - "input": [text], - }, - ) - response.raise_for_status() - payload = DenseEmbeddingResponse.model_validate(response.json()) - if not payload.data: - raise ValueError("Dense embedding response is empty") - return payload.data[0].embedding - - -async def embed_dense_batch(client: httpx.AsyncClient, texts: list[str]) -> list[list[float]]: - response = await client.post( - EMBEDDINGS_DENSE_URL, - **get_upstream_request_kwargs(), - json={ - "model": os.getenv("EMBEDDINGS_DENSE_MODEL", EMBEDDINGS_DENSE_MODEL), - "input": texts, - }, - ) - response.raise_for_status() - payload = DenseEmbeddingResponse.model_validate(response.json()) - payload.data.sort(key=lambda x: x.index) - return [item.embedding for item in payload.data] - - -def embed_sparse_sync(text: str) -> SparseVector: - vectors = list(get_sparse_model().embed([text])) - if not vectors: - raise ValueError("Sparse embedding response is empty") - item = vectors[0] - return SparseVector( - indices=[int(index) for index in item.indices.tolist()], - values=[float(value) for value in item.values.tolist()], - ) - - -def build_dense_query(question: Question) -> str: - return question.text.strip() - - -def build_sparse_query(question: Question) -> str: - parts = [question.text.strip()] - if question.keywords: - parts.extend(question.keywords) - if question.search_text: - parts = [question.search_text] - return " ".join(parts) - - -async def qdrant_search( - client: AsyncQdrantClient, - dense_vectors: list[list[float]], - sparse_vector: SparseVector, -) -> list[Any] | None: - prefetch_list = [] - for dv in dense_vectors: - prefetch_list.append( - models.Prefetch( - query=dv, - using=QDRANT_DENSE_VECTOR_NAME, - limit=DENSE_PREFETCH_K, - ) - ) - prefetch_list.append( - models.Prefetch( - query=models.SparseVector( - indices=sparse_vector.indices, - values=sparse_vector.values, - ), - using=QDRANT_SPARSE_VECTOR_NAME, - limit=SPARSE_PREFETCH_K, - ) - ) - - response = await client.query_points( - collection_name=QDRANT_COLLECTION_NAME, - prefetch=prefetch_list, - query=models.FusionQuery(fusion=models.Fusion.RRF), - limit=RETRIEVE_K, - with_payload=True, - ) - - if not response.points: - return None - - return response.points - - -def extract_message_ids(point: Any) -> list[str]: - payload = point.payload or {} - metadata = payload.get("metadata") or {} - message_ids = metadata.get("message_ids") or [] - return [str(message_id) for message_id in message_ids] - - -async def get_rerank_scores( - client: httpx.AsyncClient, - label: str, - targets: list[str], -) -> list[float]: - if not targets: - return [] - - for attempt in range(5): - try: - response = await client.post( - RERANKER_URL, - **get_upstream_request_kwargs(), - json={ - "model": RERANKER_MODEL, - "encoding_format": "float", - "text_1": label, - "text_2": targets, - }, - ) - if response.status_code == 429: - wait = 2 ** attempt - logger.warning(f"Rerank 429, retry {attempt+1}/5 in {wait}s") - await asyncio.sleep(wait) - continue - response.raise_for_status() - payload = response.json() - data = payload.get("data") or [] - return [float(sample["score"]) for sample in data] - except Exception as e: - logger.warning(f"Rerank error attempt {attempt+1}/5: {e}") - if attempt < 4: - await asyncio.sleep(2 ** attempt) - continue - logger.error("Rerank failed after 5 attempts, using fallback") - return [] - - logger.error("Rerank 429 after 5 retries, using fallback") - return [] - - -async def rerank_points( - client: httpx.AsyncClient, - query: str, - points: list[Any], -) -> list[Any]: - rerank_candidates = points[:RERANK_LIMIT] - rerank_targets = [point.payload.get("page_content") for point in rerank_candidates] - scores = await get_rerank_scores(client, query, rerank_targets) - - if not scores: - logger.warning("Reranker unavailable, returning RRF order") - return rerank_candidates - - reranked_candidates = [ - point - for _, point in sorted( - zip(scores, rerank_candidates, strict=True), - key=lambda item: item[0], - reverse=True, - ) - ] - - return reranked_candidates - @app.get("/health") async def health() -> dict[str, str]: @@ -361,43 +72,28 @@ async def search(payload: SearchAPIRequest) -> SearchAPIResponse: client: httpx.AsyncClient = app.state.http qdrant: AsyncQdrantClient = app.state.qdrant - dense_query = build_dense_query(question) + primary_query = build_primary_query(question) sparse_query = build_sparse_query(question) + extra_queries = build_extra_dense_queries(question) - dense_task = embed_dense(client, dense_query) - sparse_task = asyncio.to_thread(embed_sparse_sync, sparse_query) - dense_vector, sparse_vector = await asyncio.gather(dense_task, sparse_task) + dense_task = embed_dense(client, primary_query) + sparse_task = asyncio.to_thread(embed_sparse, sparse_query) + primary_dense, sparse_vector = await asyncio.gather(dense_task, sparse_task) - dense_vectors = [dense_vector] - if question.hyde and len(question.hyde) > 0: + extra_dense: list[list[float]] = [] + if extra_queries: try: - hyde_texts = question.hyde[:2] - hyde_vectors = await embed_dense_batch(client, hyde_texts) - dense_vectors.extend(hyde_vectors) + extra_dense = await embed_dense_batch(client, extra_queries[:2]) except Exception as e: - logger.warning(f"HyDE embedding failed: {e}") + logger.warning("Extra dense embedding failed: %s", e) - all_points = await qdrant_search(qdrant, dense_vectors, sparse_vector) + all_points = await qdrant_search(qdrant, primary_dense, extra_dense, sparse_vector, question) - if all_points is None: + if not all_points: return SearchAPIResponse(results=[]) - all_points = list(all_points) - - reranked = await rerank_points(client, query, all_points) - - reranked_ids = {id(p) for p in reranked} - remaining = [p for p in all_points if id(p) not in reranked_ids] - final_points = reranked + remaining - - seen: set[str] = set() - message_ids: list[str] = [] - for point in final_points: - for mid in extract_message_ids(point): - if mid not in seen: - seen.add(mid) - message_ids.append(mid) - message_ids = message_ids[:50] + reranked_head, tail = await rerank_points(client, query, all_points) + message_ids = aggregate_message_ids(reranked_head, tail) return SearchAPIResponse(results=[SearchAPIItem(message_ids=message_ids)]) diff --git a/search/requirements.txt b/search/requirements.txt index 0ca00a6..50b3c1d 100644 --- a/search/requirements.txt +++ b/search/requirements.txt @@ -4,4 +4,3 @@ pydantic==2.12.5 httpx==0.28.1 qdrant-client==1.15.1 fastembed==0.7.4 -orjson==3.10.18