Update src/controllers/_data_extractor.py

src/controllers/_data_extractor.py

@@ -3,7 +3,7 @@ from fastapi import APIRouter, HTTPException
 
 from src.services import DataExtractorService
 from pydantic import BaseModel
-from src.schemas import Message
+from src.schemas import Message, SQLQueryExtractor
 
 from loguru import logger
 
@@ -14,9 +14,9 @@ class UserQuery(BaseModel):
 
 class Response(BaseModel):
     status: str
-    data: Optional[list] = []
+    data: Optional[list[dict]] = []
     sql_query: str
-
+    output: SQLQueryExtractor
 
 class DataExtractorController:
     def __init__(self):
@@ -33,10 +33,10 @@ class DataExtractorController:
         logger.info("Extracting data...")
         try:
             async with DataExtractorService() as service:
-                response, sql_query = await service.extract(
+                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)
+                return Response(status="success", data=response, sql_query=sql_query, output=output)
         except HTTPException as e:
             logger.error(e)
             raise e