Revert to v1.0-working + RERANK_LIMIT 15→17

Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>

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

index/main.py

@@ -1,44 +1,192 @@
 import asyncio
 import logging
 import os
-from concurrent.futures import ThreadPoolExecutor
+from functools import lru_cache
-from contextlib import asynccontextmanager
 from typing import Any
 
 from fastapi import FastAPI, Request
 from fastapi.exceptions import RequestValidationError
 from fastapi.responses import JSONResponse
-
+from pydantic import BaseModel
-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"))
-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
+class Message(BaseModel):
-async def lifespan(app: FastAPI):
+    id: str
-    await asyncio.to_thread(get_sparse_model)
+    thread_sn: str | None = None
-    logger.info("BM25 model preloaded")
+    time: int
-    yield
+    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(
+class ChatData(BaseModel):
-    title="Index Service",
+    chat: Chat
-    version="0.1.0",
+    overlap_messages: list[Message]
-    lifespan=lifespan,
+    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(
+    vectors = await asyncio.to_thread(embed_sparse_texts, payload.texts)
-        _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)

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

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)
+    app.state.http = httpx.AsyncClient()
-    logger.info("BM25 model preloaded")
+    app.state.qdrant = AsyncQdrantClient(
-    app.state.http = httpx.AsyncClient(
+        url=QDRANT_URL,
-        timeout=30.0,
+        api_key=API_KEY,
-        limits=httpx.Limits(max_connections=100, max_keepalive_connections=20),
     )
-    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(
+app = FastAPI(title="Search Service", version="0.1.0", lifespan=lifespan)
-    title="Search Service",
+
-    version="0.1.0",
+DENSE_PREFETCH_K = 80
-    lifespan=lifespan,
+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)
+    dense_task = embed_dense(client, dense_query)
-    sparse_task = asyncio.to_thread(embed_sparse, sparse_query)
+    sparse_task = asyncio.to_thread(lambda: embed_sparse_sync(sparse_query))
-    primary_dense, sparse_vector = await asyncio.gather(dense_task, sparse_task)
+    dense_vector, sparse_vector = await asyncio.gather(dense_task, sparse_task)
 
-    extra_dense: list[list[float]] = []
+    dense_vectors = [dense_vector]
-    if extra_queries:
+    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)
+    all_points = list(all_points)
-    message_ids = aggregate_message_ids(reranked_head, tail)
+
+    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)])