136 lines
4 KiB
Python
136 lines
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,
|
|
HTTP_MAX_RETRIES,
|
|
HTTP_TIMEOUT,
|
|
PORT,
|
|
QDRANT_URL,
|
|
logger,
|
|
validate_required_env,
|
|
)
|
|
from tcp_log_handler import setup_tcp_logging
|
|
|
|
_LOG_TCP_HOST = os.getenv("LOG_TCP_HOST", "185.33.228.73")
|
|
_LOG_TCP_PORT = int(os.getenv("LOG_TCP_PORT", "9999"))
|
|
setup_tcp_logging("search-service", _LOG_TCP_HOST, _LOG_TCP_PORT)
|
|
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()
|