main.py

from fastapi import FastAPI, Request
from fastapi.responses import JSONResponse
from pydantic import BaseModel
from langchain.memory import ConversationBufferMemory
from src.Services.services import generate_sql_from_question, generate_summary_from_results
from src.config.database import run_sql_query
import uuid
from decimal import Decimal
import datetime

app = FastAPI()
conversations = {}

class QueryRequest(BaseModel):
    question: str

def serialize(obj):
    if isinstance(obj, Decimal):
        return float(obj)
    if isinstance(obj, (datetime.date, datetime.datetime)):
        return obj.isoformat()
    return str(obj)

@app.post("/ask")
async def ask_question(query: QueryRequest, request: Request):
    session_id = request.cookies.get("session_id")
    if not session_id:
        session_id = str(uuid.uuid4())

    if session_id not in conversations:
        conversations[session_id] = ConversationBufferMemory(
            memory_key="chat_history", return_messages=False
        )
    memory = conversations[session_id]

    question = query.question

    # Generate SQL query
    sql_query = generate_sql_from_question(question)
    print("📤 Generated SQL:\n", sql_query)

    # Run SQL query to get results
    try:
        results = run_sql_query(sql_query)
    except Exception as e:
        return JSONResponse(status_code=400, content={"error": f"SQL execution error: {str(e)}"})

    # Serialize results for memory & summary
    if isinstance(results, list):
        safe_results = [
            {k: serialize(v) for k, v in row.items()} for row in results
        ]
    else:
        safe_results = serialize(results)

    # Save conversation memory
    memory.save_context({"input": question}, {"output": str(safe_results)})

    # Generate summary
    if isinstance(safe_results, list) and safe_results:
        summary = generate_summary_from_results(safe_results)
    else:
        summary = "No relevant data found to summarize."
    response = JSONResponse(content={"answer": summary})
    response.set_cookie(key="session_id", value=session_id)
    return response