fastapi_main.py

from typing import List, Dict, Any
from dataclasses import dataclass
import re
from datetime import datetime
from fastapi import FastAPI, Request, Form #, UploadFile, HTTPException, File
# from fastapi.staticfiles import StaticFiles
from fastapi.templating import Jinja2Templates
from fastapi.responses import HTMLResponse
from starlette.status import HTTP_429_TOO_MANY_REQUESTS #, HTTP_303_SEE_OTHER, 
from starlette.responses import JSONResponse
from dotenv import load_dotenv
from function_tools import handle_db_operation_from_fastapi, classify_query_intent
from utils.logging_base import setup_logger

from slowapi import Limiter
from slowapi.util import get_remote_address
from slowapi.errors import RateLimitExceeded
from slowapi.middleware import SlowAPIMiddleware

DATE_STAMP, TIME_STAMP = datetime.now().strftime('%Y-%m-%d'), datetime.now().strftime('%H_%M_%S')

@dataclass
class OperationResult:
    final_output_message: str
    rows = List[Dict[str, Any]]

# Configure logging
logger = setup_logger(DATE_STAMP, TIME_STAMP)

# GPT_MODEL = 'gpt-4.1-nano'
GPT_MODELS = [
    'gpt-4.1-nano',
    'gpt-o4-mini',
    'gpt-4.1-mini',
    'gpt-o3-mini',
    'gpt-4o-mini-audio',
    'gpt-o1-mini',
    'gpt-5-mini'
    ]
load_dotenv("utils/.env")
load_dotenv(".env")

templates = Jinja2Templates(directory="templates")

ai_agent_app = FastAPI()
limiter = Limiter(key_func = get_remote_address)
ai_agent_app.state.limiter = limiter
ai_agent_app.add_middleware(SlowAPIMiddleware)


@ai_agent_app.exception_handler(RateLimitExceeded)
async def rate_limit_handler(request: Request, exc: RateLimitExceeded):
    return JSONResponse(
        status_code = HTTP_429_TOO_MANY_REQUESTS,
        content = {
            "detail": "Rate limit exceeded bro. See you later",
            "request_details" : f"Rate limit hit at {request.url} from IP: {request.client.host}",
            "exc_details" : exc.detail
            }
    )

@ai_agent_app.get("/", response_class = HTMLResponse)
def form_page(request: Request):
    return templates.TemplateResponse("form.html", {
        "request" : request,
        "models" : GPT_MODELS
        })

@ai_agent_app.post('/message')
def create(message: str):
    """
    Post a message
    """
    return {"message" : message}

def is_dml_query(query: str) -> bool:
    return bool(re.search(r"\b(delete|update|insert|drop|modify|change|remove)\b", query, re.IGNORECASE))


@ai_agent_app.post('/NL-query')
@limiter.limit("2/minute")
async def handle_nl_query(request: Request,
                          nl_query: str = Form(...),
                          model: str = Form(...),
                          confirm: str = Form(default = None)):
    """
    Process the NL query and return results.
    """
    print(f"Model received: {model}")
    intent = await classify_query_intent(nl_query, model)
    if intent == "DML" and confirm != 'yes':
        return templates.TemplateResponse(
            "confirm.html",{
                "request" : request,
                "query" : nl_query,
                "model" : model
            }
        )
    if intent == "DML" and confirm == 'no':
        return templates.TemplateResponse(
            "form.html", {
            "request" : request,
            "models" : GPT_MODELS,
            "message" : "Query cancelled"
        })
    try:
        result = await handle_db_operation_from_fastapi(nl_query, model = model)
        return templates.TemplateResponse(
            "results.html",
            {
                "request" : request,
                "message" : result.final_output_message,
                "rows" : result.rows,
                "sql_query" : result.sql_query
            },
        )

    except Exception as e:
        return templates.TemplateResponse(
            "results.html",
            {
                "request" : request,
                "message" : f"Error {e}",
                "rows" : [],
                "sql_query" : None
            },
            status_code = 500
        )


@ai_agent_app.post("/confirm-query")
async def confirm_query(request: Request, nl_query: str = Form(...), model: str = Form(...)):
    """"""
    if is_dml_query(nl_query):
        return templates.TemplateResponse(
            "confirm.html",{
                "request" : request,
                "query" : nl_query,
                "model" : model
            }
        )
    # Non-DML query - process directly
    return await handle_nl_query(request, nl_query, model = model)