forked from zovos/vk_hackathon
60 lines
1.7 KiB
Python
60 lines
1.7 KiB
Python
import asyncio
|
|
import logging
|
|
import os
|
|
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 .schemas import IndexAPIRequest, IndexAPIResponse, SparseEmbeddingRequest
|
|
from .sparse import embed_sparse_texts
|
|
|
|
HOST = os.getenv("HOST", "0.0.0.0")
|
|
PORT = int(os.getenv("PORT", "8004"))
|
|
UVICORN_WORKERS = 8
|
|
|
|
logging.basicConfig(level=os.getenv("LOG_LEVEL", "INFO"))
|
|
logger = logging.getLogger("index-service")
|
|
|
|
app = FastAPI(title="Index Service", version="0.2.0")
|
|
|
|
|
|
@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.to_thread(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()
|