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