Spaces:
Sleeping
Sleeping
File size: 1,370 Bytes
8e073b5 6500c31 ff24e58 6500c31 ff24e58 6500c31 ff24e58 6500c31 dca20c8 6500c31 8e073b5 6500c31 ff24e58 6500c31 dca20c8 6500c31 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 | from typing import Optional
from fastapi import APIRouter, HTTPException
from src.services import DataExtractorService
from pydantic import BaseModel
from src.schemas import Message, SQLQueryExtractor
from loguru import logger
class UserQuery(BaseModel):
user_query: str
class Response(BaseModel):
status: str
data: Optional[list[dict]] = []
sql_query: str
output: SQLQueryExtractor
class DataExtractorController:
def __init__(self):
self.router = APIRouter()
self.router.add_api_route(
"/extract",
self.extrcat,
methods=["POST"],
)
async def extrcat(
self, user_query: UserQuery, message_history: Optional[list[Message]] = []
):
logger.info("Extracting data...")
try:
async with DataExtractorService() as service:
response, sql_query, output = await service.extract(
user_query=user_query.user_query, message_history=message_history
)
return Response(
status="success", data=response, sql_query=sql_query, output=output
)
except HTTPException as e:
logger.error(e)
raise e
except Exception as e:
logger.error(e)
raise HTTPException(status_code=500, detail=str(e))
|