src/controllers/_data_extractor.py

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))