Spaces:
Running
Running
| from fastapi import APIRouter, HTTPException | |
| from pydantic import BaseModel | |
| from typing import List, Optional | |
| from src.pipeline.query_pipeline import QueryPipeline | |
| router = APIRouter() | |
| _pipeline = None | |
| def get_pipeline(): | |
| global _pipeline | |
| if _pipeline is None: | |
| _pipeline = QueryPipeline() | |
| return _pipeline | |
| class QueryRequest(BaseModel): | |
| query: str | |
| top_k_retrieval: Optional[int] = 20 | |
| top_k_rerank: Optional[int] = 5 | |
| use_hyde: Optional[bool] = False | |
| class DocResponse(BaseModel): | |
| content: str | |
| score: float | |
| class QueryResponse(BaseModel): | |
| query: str | |
| answer: str | |
| context: List[tuple] | |
| async def chat(request: QueryRequest): | |
| try: | |
| pipe = get_pipeline() | |
| result = pipe.run( | |
| query=request.query, | |
| top_k_retrieval=request.top_k_retrieval, | |
| top_k_rerank=request.top_k_rerank | |
| ) | |
| return result | |
| except Exception as e: | |
| raise HTTPException(status_code=500, detail=str(e)) | |