7247248b27
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>
82 lines
2.1 KiB
Python
82 lines
2.1 KiB
Python
import asyncio
|
|
import logging
|
|
import os
|
|
from concurrent.futures import ThreadPoolExecutor
|
|
from contextlib import asynccontextmanager
|
|
from typing import Any
|
|
|
|
from fastapi import FastAPI, Request
|
|
from fastapi.exceptions import RequestValidationError
|
|
from fastapi.responses import JSONResponse
|
|
|
|
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"))
|
|
UVICORN_WORKERS = 4
|
|
|
|
logging.basicConfig(level=os.getenv("LOG_LEVEL", "INFO"))
|
|
logger = logging.getLogger("index-service")
|
|
|
|
_thread_pool = ThreadPoolExecutor(max_workers=4)
|
|
|
|
|
|
@asynccontextmanager
|
|
async def lifespan(app: FastAPI):
|
|
await asyncio.to_thread(get_sparse_model)
|
|
logger.info("BM25 model preloaded")
|
|
yield
|
|
|
|
|
|
app = FastAPI(
|
|
title="Index Service",
|
|
version="0.1.0",
|
|
lifespan=lifespan,
|
|
)
|
|
|
|
|
|
@app.get("/health")
|
|
async def health() -> dict[str, str]:
|
|
return {"status": "ok"}
|
|
|
|
|
|
@app.post("/index", response_model=IndexAPIResponse)
|
|
async def index(payload: IndexAPIRequest) -> IndexAPIResponse:
|
|
return IndexAPIResponse(
|
|
results=build_chunks(
|
|
payload.data.overlap_messages,
|
|
payload.data.new_messages,
|
|
)
|
|
)
|
|
|
|
|
|
@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": [{"indices": v.indices, "values": v.values} for v in vectors]}
|
|
|
|
|
|
@app.exception_handler(Exception)
|
|
async def exception_handler(request: Request, exc: Exception) -> JSONResponse:
|
|
logger.exception(exc)
|
|
if isinstance(exc, RequestValidationError):
|
|
return JSONResponse(status_code=422, content={"detail": exc.errors()})
|
|
return JSONResponse(status_code=500, content={"detail": str(exc)})
|
|
|
|
|
|
def main() -> None:
|
|
import uvicorn
|
|
|
|
uvicorn.run("main:app", host=HOST, port=PORT, reload=False, workers=UVICORN_WORKERS)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|