Port index and search logic from working Lotus reference implementation

Both services are now single-file (main.py only), exactly matching
the Lotus solution structure that passes the test stand:
- index: char-based sliding window chunking (256/128), is_system+is_hidden
  filter, render_message consistent with Lotus, UVICORN_WORKERS=8
- search: validate_required_env at module level, embed_dense_batch for
  HyDE, 429 retry on reranker, RRF fusion without per-query filter
- Dockerfiles: COPY main.py . (no extra modules to import)

Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>
This commit is contained in:
q 2026-04-18 18:32:41 +03:00
parent 143efd6531
commit 74c967a34a
4 changed files with 543 additions and 71 deletions

View file

@ -5,7 +5,7 @@ WORKDIR /app
COPY requirements.txt . COPY requirements.txt .
RUN pip install --no-cache-dir -r requirements.txt RUN pip install --no-cache-dir -r requirements.txt
COPY *.py . COPY main.py .
ENV HOST=0.0.0.0 ENV HOST=0.0.0.0
ENV PORT=8000 ENV PORT=8000

View file

@ -1,24 +1,192 @@
import asyncio import asyncio
import logging import logging
import os import os
from functools import lru_cache
from typing import Any from typing import Any
from fastapi import FastAPI, Request from fastapi import FastAPI, Request
from fastapi.exceptions import RequestValidationError from fastapi.exceptions import RequestValidationError
from fastapi.responses import JSONResponse 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
HOST = os.getenv("HOST", "0.0.0.0") HOST = os.getenv("HOST", "0.0.0.0")
PORT = int(os.getenv("PORT", "8004")) PORT = int(os.getenv("PORT", "8000"))
UVICORN_WORKERS = 8 UVICORN_WORKERS = 8
logging.basicConfig(level=os.getenv("LOG_LEVEL", "INFO")) logging.basicConfig(level=os.getenv("LOG_LEVEL", "INFO"))
logger = logging.getLogger("index-service") logger = logging.getLogger("index-service")
app = FastAPI(title="Index Service", version="0.2.0")
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"
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") @app.get("/health")
@ -30,16 +198,38 @@ async def health() -> dict[str, str]:
async def index(payload: IndexAPIRequest) -> IndexAPIResponse: async def index(payload: IndexAPIRequest) -> IndexAPIResponse:
return IndexAPIResponse( return IndexAPIResponse(
results=build_chunks( results=build_chunks(
payload.data.chat,
payload.data.overlap_messages, payload.data.overlap_messages,
payload.data.new_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") @app.post("/sparse_embedding")
async def sparse_embedding(payload: SparseEmbeddingRequest) -> dict[str, Any]: async def sparse_embedding(payload: SparseEmbeddingRequest) -> dict[str, Any]:
vectors = await asyncio.to_thread(embed_sparse_texts, payload.texts) vectors = await asyncio.to_thread(embed_sparse_texts, payload.texts)
return {"vectors": [{"indices": v.indices, "values": v.values} for v in vectors]} return {"vectors": vectors}
@app.exception_handler(Exception) @app.exception_handler(Exception)

View file

@ -5,7 +5,7 @@ WORKDIR /app
COPY requirements.txt . COPY requirements.txt .
RUN pip install --no-cache-dir -r requirements.txt RUN pip install --no-cache-dir -r requirements.txt
COPY *.py . COPY main.py .
ENV HOST=0.0.0.0 ENV HOST=0.0.0.0
ENV PORT=8000 ENV PORT=8000

View file

@ -2,56 +2,157 @@ import asyncio
import logging import logging
import os import os
from contextlib import asynccontextmanager from contextlib import asynccontextmanager
from functools import lru_cache
from typing import Any
import httpx import httpx
from fastembed import SparseTextEmbedding
from fastapi import FastAPI, HTTPException, Request from fastapi import FastAPI, HTTPException, Request
from fastapi.exceptions import RequestValidationError from fastapi.exceptions import RequestValidationError
from fastapi.responses import JSONResponse from fastapi.responses import JSONResponse
from qdrant_client import AsyncQdrantClient from pydantic import BaseModel, Field
from qdrant_client import AsyncQdrantClient, models
from aggregation import aggregate_message_ids EMBEDDINGS_DENSE_MODEL = "Qwen/Qwen3-Embedding-0.6B"
from config import (
API_KEY, HOST = os.getenv("HOST", "0.0.0.0")
HOST, PORT = int(os.getenv("PORT", "8000"))
HTTP_MAX_RETRIES,
HTTP_TIMEOUT, API_KEY = os.getenv("API_KEY")
PORT, EMBEDDINGS_DENSE_URL = os.getenv("EMBEDDINGS_DENSE_URL")
QDRANT_URL, QDRANT_DENSE_VECTOR_NAME = os.getenv("QDRANT_DENSE_VECTOR_NAME", "dense")
logger, QDRANT_SPARSE_VECTOR_NAME = os.getenv("QDRANT_SPARSE_VECTOR_NAME", "sparse")
validate_required_env, SPARSE_MODEL_NAME = "Qdrant/bm25"
) RERANKER_MODEL = "nvidia/llama-nemotron-rerank-1b-v2"
from query_builder import ( RERANKER_URL = os.getenv("RERANKER_URL")
build_extra_dense_queries, OPEN_API_LOGIN = os.getenv("OPEN_API_LOGIN")
build_primary_query, OPEN_API_PASSWORD = os.getenv("OPEN_API_PASSWORD")
build_sparse_query, QDRANT_URL = os.getenv("QDRANT_URL")
embed_dense, QDRANT_COLLECTION_NAME = os.getenv("QDRANT_COLLECTION_NAME", "evaluation")
embed_dense_batch, REQUIRED_ENV_VARS = [
embed_sparse, "EMBEDDINGS_DENSE_URL",
) "RERANKER_URL",
from rerank import rerank_points "QDRANT_URL",
from retrieval import qdrant_search ]
from schemas import SearchAPIItem, SearchAPIRequest, SearchAPIResponse, SparseVector
logging.basicConfig(level=os.getenv("LOG_LEVEL", "INFO"))
logger = logging.getLogger("search-service")
async def _embed_dense_with_retry(client: httpx.AsyncClient, text: str) -> list[float]: def validate_required_env() -> None:
last_exc: Exception | None = None if bool(OPEN_API_LOGIN) != bool(OPEN_API_PASSWORD):
for attempt in range(HTTP_MAX_RETRIES + 1): raise RuntimeError("OPEN_API_LOGIN and OPEN_API_PASSWORD must be set together")
try:
return await embed_dense(client, text) if not API_KEY and not (OPEN_API_LOGIN and OPEN_API_PASSWORD):
except (httpx.TransportError, httpx.HTTPStatusError) as exc: raise RuntimeError("Either API_KEY or OPEN_API_LOGIN and OPEN_API_PASSWORD must be set")
if isinstance(exc, httpx.HTTPStatusError) and exc.response.status_code < 500:
raise missing_env_vars = [
last_exc = exc name for name in REQUIRED_ENV_VARS if os.getenv(name) is None or os.getenv(name) == ""
if attempt < HTTP_MAX_RETRIES: ]
await asyncio.sleep(0.5 * (attempt + 1)) if not missing_env_vars:
raise RuntimeError(f"Dense embedding failed after retries: {last_exc}") return
logger.error("Empty required env vars: %s", ", ".join(missing_env_vars))
raise RuntimeError(f"Empty required env vars: {', '.join(missing_env_vars)}")
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 @asynccontextmanager
async def lifespan(app: FastAPI): async def lifespan(app: FastAPI):
validate_required_env() app.state.http = httpx.AsyncClient()
app.state.http = httpx.AsyncClient(timeout=HTTP_TIMEOUT) app.state.qdrant = AsyncQdrantClient(
app.state.qdrant = AsyncQdrantClient(url=QDRANT_URL, api_key=API_KEY) url=QDRANT_URL,
api_key=API_KEY,
)
try: try:
yield yield
finally: finally:
@ -59,7 +160,179 @@ async def lifespan(app: FastAPI):
await app.state.qdrant.close() await app.state.qdrant.close()
app = FastAPI(title="Search Service", version="0.2.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 = 10
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") @app.get("/health")
@ -70,41 +343,50 @@ async def health() -> dict[str, str]:
@app.post("/search", response_model=SearchAPIResponse) @app.post("/search", response_model=SearchAPIResponse)
async def search(payload: SearchAPIRequest) -> SearchAPIResponse: async def search(payload: SearchAPIRequest) -> SearchAPIResponse:
question = payload.question question = payload.question
primary_query = build_primary_query(question) query = question.text.strip()
if not primary_query: if not query:
raise HTTPException(status_code=400, detail="question.text is required") raise HTTPException(status_code=400, detail="question.text is required")
client: httpx.AsyncClient = app.state.http client: httpx.AsyncClient = app.state.http
qdrant: AsyncQdrantClient = app.state.qdrant qdrant: AsyncQdrantClient = app.state.qdrant
extra_texts = build_extra_dense_queries(question) dense_query = build_dense_query(question)
sparse_text = build_sparse_query(question) sparse_query = build_sparse_query(question)
primary_dense, sparse_vec = await asyncio.gather( dense_task = embed_dense(client, dense_query)
_embed_dense_with_retry(client, primary_query), sparse_task = asyncio.to_thread(lambda: embed_sparse_sync(sparse_query))
asyncio.to_thread(embed_sparse, sparse_text), dense_vector, sparse_vector = await asyncio.gather(dense_task, sparse_task)
)
extra_dense_vecs: list[list[float]] = [] dense_vectors = [dense_vector]
if extra_texts: if question.hyde and len(question.hyde) > 0:
try: try:
extra_dense_vecs = await embed_dense_batch(client, extra_texts[:3]) hyde_texts = question.hyde[:2]
except Exception as exc: hyde_vectors = await embed_dense_batch(client, hyde_texts)
logger.warning("Extra dense embedding failed, skipping: %s", exc) dense_vectors.extend(hyde_vectors)
except Exception as e:
logger.warning(f"HyDE embedding failed: {e}")
points = await qdrant_search( all_points = await qdrant_search(qdrant, dense_vectors, sparse_vector)
qdrant,
primary_dense,
extra_dense_vecs,
sparse_vec,
question,
)
if not points: if all_points is None:
return SearchAPIResponse(results=[]) return SearchAPIResponse(results=[])
reranked_head, retrieval_tail = await rerank_points(client, primary_query, points) all_points = list(all_points)
message_ids = aggregate_message_ids(reranked_head, retrieval_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)]) return SearchAPIResponse(results=[SearchAPIItem(message_ids=message_ids)])