vk_hackathon/search/main.py

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,
    HTTP_MAX_RETRIES,
    HTTP_TIMEOUT,
    PORT,
    QDRANT_URL,
    logger,
    validate_required_env,
)
from query_builder import (
    build_extra_dense_queries,
    build_primary_query,
    build_sparse_query,
    embed_dense,
    embed_dense_multi,
    embed_sparse,
)
from rerank import rerank_points
from retrieval import qdrant_search
from schemas import SearchAPIItem, SearchAPIRequest, SearchAPIResponse, SparseVector


async def _embed_dense_with_retry(client: httpx.AsyncClient, text: str) -> list[float]:
    last_exc: Exception | None = None
    for attempt in range(HTTP_MAX_RETRIES + 1):
        try:
            return await embed_dense(client, text)
        except (httpx.TransportError, httpx.HTTPStatusError) as exc:
            if isinstance(exc, httpx.HTTPStatusError) and exc.response.status_code < 500:
                raise
            last_exc = exc
            if attempt < HTTP_MAX_RETRIES:
                await asyncio.sleep(0.5 * (attempt + 1))
    raise RuntimeError(f"Dense embedding failed after retries: {last_exc}")


@asynccontextmanager
async def lifespan(app: FastAPI):
    validate_required_env()
    app.state.http = httpx.AsyncClient(timeout=HTTP_TIMEOUT)
    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.2.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
    primary_query = build_primary_query(question)
    if not primary_query:
        raise HTTPException(status_code=400, detail="question.text is required")

    client: httpx.AsyncClient = app.state.http
    qdrant: AsyncQdrantClient = app.state.qdrant

    extra_texts = build_extra_dense_queries(question)
    sparse_text = build_sparse_query(question)

    async def _no_extra() -> list:
        return []

    extra_task = embed_dense_multi(client, extra_texts) if extra_texts else _no_extra()
    primary_dense, extra_dense_vecs, sparse_vec = await asyncio.gather(
        _embed_dense_with_retry(client, primary_query),
        extra_task,
        asyncio.to_thread(embed_sparse, sparse_text),
    )

    points = await qdrant_search(
        qdrant,
        primary_dense,
        extra_dense_vecs,
        sparse_vec,
        question,
    )

    if not points:
        return SearchAPIResponse(results=[])

    reranked_head, retrieval_tail = await rerank_points(client, primary_query, points)
    message_ids = aggregate_message_ids(reranked_head, retrieval_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()