AISqlQueryFastApi / backend /app /api /endpoints.py
Vivek0912's picture
api code changes for schema reader
d779a9b
Raw
History Blame Contribute Delete
4.75 kB
from langchain_openai import ChatOpenAI
from config.settings import settings
# --- LLM-based classifier for user input ---
import pandas as pd
# Use LLMClassifier from services.utility
from app.services.utility import UtilityClass
from fastapi import APIRouter
from pydantic import BaseModel
from app.services.sql_generator import generate_sql_query
from app.services.query_executor import execute_sql_query
from app.services.result_formatter import df_to_chart
from app.services.query_executor import execute_sql_query
from app.db.schema_reader import get_schema
from app.services.query_executor import run_and_handle_sql_query
router = APIRouter()
class QueryRequest(BaseModel):
question: str
from fastapi import HTTPException
@router.post("/process-text")
def process_text(req: QueryRequest):
try:
sql = generate_sql_query(req.question)
# If generate_sql_query returns a chat_message, treat as normal chat or empty SQL result
if isinstance(sql, dict) and "chat_message" in sql:
return {"message": sql["chat_message"]}
# Use utility method to check if the string is a valid SQL query
formattedSqlQuery = UtilityClass.is_valid_sql_query(sql)
print(f"Formatted SQL Query: {formattedSqlQuery}")
if not formattedSqlQuery:
return {"message": str(sql) if sql else "No SQL query could be generated for your question."}
df = run_and_handle_sql_query(formattedSqlQuery, req.question)
# If run_and_handle_sql_query returns a chat_message (for empty SQL results), return it
if isinstance(df, dict) and "chat_message" in df:
return {"message": df["chat_message"]}
chart = None
# if len(df) > 0 and isinstance(df, list) and len(df[0]) > 0 and len(df[0].keys()) >= 2:
# chart = df_to_chart(pd.DataFrame(df))
# Prepare heading and records JSON using LLM
result_json = UtilityClass.prepare_llm_heading_and_records(req.question, df)
# Parse the heading if it's a JSON string containing heading and summary
heading_text = result_json["heading"]
summary_text = ""
if isinstance(heading_text, str):
try:
import json
# Check if it's wrapped in markdown code block
if heading_text.strip().startswith('```json') and heading_text.strip().endswith('```'):
# Extract JSON from markdown code block
json_content = heading_text.strip()
# Remove ```json from start and ``` from end
json_content = json_content[7:-3].strip() # Remove ```json and ```
heading_data = json.loads(json_content)
else:
# Try to parse directly
heading_data = json.loads(heading_text)
if isinstance(heading_data, dict):
heading_text = heading_data.get("heading", heading_text)
summary_text = heading_data.get("summary", summary_text)
except (json.JSONDecodeError, ValueError):
# If not JSON, use as-is
pass
return {
"sql": sql,
"rows": result_json["records"],
"heading": heading_text, # Send just the heading text
"summary": summary_text, # Send just the summary text
"chart": chart
}
except Exception as e:
raise HTTPException(status_code=400, detail=f"An error occurred: {str(e)}")
@router.get("/health")
def health_check():
"""
Health check endpoint to verify API and system status.
Returns:
dict: Health status information including system checks and timestamp
"""
try:
health_status = UtilityClass.get_health_status()
# Return appropriate HTTP status based on health
if health_status["status"] == "healthy":
return health_status
else:
# Return 503 Service Unavailable if any critical component is unhealthy
from fastapi import Response
import json
return Response(
content=json.dumps(health_status),
media_type="application/json",
status_code=503
)
except Exception as e:
# Return 503 if health check itself fails
from datetime import datetime
error_response = {
"status": "unhealthy",
"message": f"Health check failed: {str(e)}",
"timestamp": datetime.utcnow().isoformat()
}
raise HTTPException(status_code=503, detail=error_response)