Migrate to multi-file architecture: smarter chunking + fixed RERANK_LIMIT

index: message-based windowed chunking (5 msgs/1h gap), better unicode
cleaning, separate dense (with timestamps)/sparse content renderers,
BM25 preload on startup, ThreadPoolExecutor(4), UVICORN_WORKERS=4,
Dockerfile copies all *.py

search: proper multi-module structure (query_builder, retrieval, rerank,
aggregation), RERANK_LIMIT 60→15 (fixes 429 errors), extra dense vectors
for variants/hyde, date+asker metadata filters, httpx pool (100/20/30s),
BM25 preload on startup, Dockerfile copies all *.py

68/68 unit tests passing

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 main.py .
+COPY *.py .
 
 ENV HOST=0.0.0.0
 ENV PORT=8000

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 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"))
@@ -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)

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

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

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
+DENSE_PREFETCH_K = 80
-SPARSE_PREFETCH_K = 100
+SPARSE_PREFETCH_K = 200
-RETRIEVE_K = 80
+RETRIEVE_K = 150
-RERANK_LIMIT = 60
+RERANK_LIMIT = 15
 TOP_K = 50
 
 HTTP_TIMEOUT = 30.0

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 fastapi.responses import JSONResponse
-from pydantic import BaseModel, Field
+from qdrant_client import AsyncQdrantClient
-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):
-    # 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(
+    app.state.qdrant = AsyncQdrantClient(url=QDRANT_URL, api_key=API_KEY)
-        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)
+    dense_task = embed_dense(client, primary_query)
-    sparse_task = asyncio.to_thread(embed_sparse_sync, sparse_query)
+    sparse_task = asyncio.to_thread(embed_sparse, sparse_query)
-    dense_vector, sparse_vector = await asyncio.gather(dense_task, sparse_task)
+    primary_dense, sparse_vector = await asyncio.gather(dense_task, sparse_task)
 
-    dense_vectors = [dense_vector]
+    extra_dense: list[list[float]] = []
-    if question.hyde and len(question.hyde) > 0:
+    if extra_queries:
         try:
-            hyde_texts = question.hyde[:2]
+            extra_dense = await embed_dense_batch(client, extra_queries[:2])
-            hyde_vectors = await embed_dense_batch(client, hyde_texts)
-            dense_vectors.extend(hyde_vectors)
         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_head, tail = await rerank_points(client, query, 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)])
 

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