4ecba7d35a
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>
122 lines
3.4 KiB
Python
122 lines
3.4 KiB
Python
import asyncio
|
|
import logging
|
|
import os
|
|
from contextlib import asynccontextmanager
|
|
|
|
import httpx
|
|
from fastapi import FastAPI, HTTPException, Request
|
|
from fastapi.exceptions import RequestValidationError
|
|
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()
|
|
|
|
|
|
@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.qdrant = AsyncQdrantClient(url=QDRANT_URL, api_key=API_KEY)
|
|
try:
|
|
yield
|
|
finally:
|
|
await app.state.http.aclose()
|
|
await app.state.qdrant.close()
|
|
|
|
|
|
app = FastAPI(
|
|
title="Search Service",
|
|
version="0.1.0",
|
|
lifespan=lifespan,
|
|
)
|
|
|
|
|
|
@app.get("/health")
|
|
async def health() -> dict[str, str]:
|
|
return {"status": "ok"}
|
|
|
|
|
|
@app.post("/search", response_model=SearchAPIResponse)
|
|
async def search(payload: SearchAPIRequest) -> SearchAPIResponse:
|
|
question = payload.question
|
|
query = question.text.strip()
|
|
if not query:
|
|
raise HTTPException(status_code=400, detail="question.text is required")
|
|
|
|
client: httpx.AsyncClient = app.state.http
|
|
qdrant: AsyncQdrantClient = app.state.qdrant
|
|
|
|
primary_query = build_primary_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)
|
|
|
|
extra_dense: list[list[float]] = []
|
|
if extra_queries:
|
|
try:
|
|
extra_dense = await embed_dense_batch(client, extra_queries[:2])
|
|
except Exception as e:
|
|
logger.warning("Extra dense embedding failed: %s", e)
|
|
|
|
all_points = await qdrant_search(qdrant, primary_dense, extra_dense, sparse_vector, question)
|
|
|
|
if not all_points:
|
|
return SearchAPIResponse(results=[])
|
|
|
|
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)])
|
|
|
|
|
|
@app.exception_handler(Exception)
|
|
async def exception_handler(request: Request, exc: Exception) -> JSONResponse:
|
|
logger.exception(exc)
|
|
detail = str(exc) or repr(exc)
|
|
|
|
if isinstance(exc, RequestValidationError):
|
|
return JSONResponse(status_code=422, content={"detail": exc.errors()})
|
|
|
|
if isinstance(exc, HTTPException):
|
|
return JSONResponse(status_code=exc.status_code, content={"detail": exc.detail})
|
|
|
|
return JSONResponse(status_code=500, content={"detail": detail})
|
|
|
|
|
|
def main() -> None:
|
|
import uvicorn
|
|
|
|
uvicorn.run("main:app", host=HOST, port=PORT, reload=False)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|