From 3561a064d5d8842e4cb6b50e3115b322d8c56c04 Mon Sep 17 00:00:00 2001 From: q Date: Sat, 18 Apr 2026 23:02:05 +0300 Subject: [PATCH] =?UTF-8?q?Revert=20to=20v1.0-working=20+=20RERANK=5FLIMIT?= =?UTF-8?q?=2015=E2=86=9217?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Co-Authored-By: Claude Sonnet 4.6 --- index/Dockerfile | 2 +- index/main.py | 220 ++++++++++++++++++++++---- search/Dockerfile | 2 +- search/main.py | 385 ++++++++++++++++++++++++++++++++++++++++------ 4 files changed, 535 insertions(+), 74 deletions(-) diff --git a/index/Dockerfile b/index/Dockerfile index f36ed37..4ec50f7 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 *.py . +COPY main.py . ENV HOST=0.0.0.0 ENV PORT=8000 diff --git a/index/main.py b/index/main.py index eb1497c..d6083ea 100644 --- a/index/main.py +++ b/index/main.py @@ -1,44 +1,192 @@ import asyncio 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 - -from chunking import build_chunks -from index_schemas import ( - IndexAPIRequest, - IndexAPIResponse, - SparseEmbeddingRequest, -) -from sparse import embed_sparse_texts, get_sparse_model +from pydantic import BaseModel HOST = os.getenv("HOST", "0.0.0.0") PORT = int(os.getenv("PORT", "8000")) -UVICORN_WORKERS = 4 +UVICORN_WORKERS = 8 logging.basicConfig(level=os.getenv("LOG_LEVEL", "INFO")) logger = logging.getLogger("index-service") -_thread_pool = ThreadPoolExecutor(max_workers=4) + +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 -@asynccontextmanager -async def lifespan(app: FastAPI): - await asyncio.to_thread(get_sparse_model) - logger.info("BM25 model preloaded") - yield +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 -app = FastAPI( - title="Index Service", - version="0.1.0", - lifespan=lifespan, -) +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" + + +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") @app.get("/health") @@ -50,18 +198,38 @@ 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, ) ) +@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) + + +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": [{"indices": v.indices, "values": v.values} for v in vectors]} + vectors = await asyncio.to_thread(embed_sparse_texts, payload.texts) + return {"vectors": vectors} @app.exception_handler(Exception) diff --git a/search/Dockerfile b/search/Dockerfile index e3b5814..1954c62 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 *.py . +COPY main.py . ENV HOST=0.0.0.0 ENV PORT=8000 diff --git a/search/main.py b/search/main.py index d2553a2..32a9b2a 100644 --- a/search/main.py +++ b/search/main.py @@ -2,47 +2,157 @@ 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 -from qdrant_client import AsyncQdrantClient +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 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): - 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.http = httpx.AsyncClient() + 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: @@ -50,11 +160,179 @@ async def lifespan(app: FastAPI): await app.state.qdrant.close() -app = FastAPI( - title="Search Service", - version="0.1.0", - lifespan=lifespan, -) +app = FastAPI(title="Search Service", version="0.1.0", lifespan=lifespan) + +DENSE_PREFETCH_K = 80 +SPARSE_PREFETCH_K = 200 +RETRIEVE_K = 150 +RERANK_LIMIT = 17 + + +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") @@ -72,28 +350,43 @@ async def search(payload: SearchAPIRequest) -> SearchAPIResponse: client: httpx.AsyncClient = app.state.http qdrant: AsyncQdrantClient = app.state.qdrant - primary_query = build_primary_query(question) + dense_query = build_dense_query(question) sparse_query = build_sparse_query(question) - extra_queries = build_extra_dense_queries(question) - 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_task = embed_dense(client, dense_query) + sparse_task = asyncio.to_thread(lambda: embed_sparse_sync(sparse_query)) + dense_vector, sparse_vector = await asyncio.gather(dense_task, sparse_task) - extra_dense: list[list[float]] = [] - if extra_queries: + dense_vectors = [dense_vector] + if question.hyde and len(question.hyde) > 0: try: - extra_dense = await embed_dense_batch(client, extra_queries[:2]) + hyde_texts = question.hyde[:2] + hyde_vectors = await embed_dense_batch(client, hyde_texts) + dense_vectors.extend(hyde_vectors) except Exception as e: - logger.warning("Extra dense embedding failed: %s", e) + logger.warning(f"HyDE embedding failed: {e}") - all_points = await qdrant_search(qdrant, primary_dense, extra_dense, sparse_vector, question) + all_points = await qdrant_search(qdrant, dense_vectors, sparse_vector) - if not all_points: + if all_points is None: return SearchAPIResponse(results=[]) - reranked_head, tail = await rerank_points(client, query, all_points) - message_ids = aggregate_message_ids(reranked_head, tail) + 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] return SearchAPIResponse(results=[SearchAPIItem(message_ids=message_ids)])