main.py

"""FastAPI Backend for OpenNL2SQL with Groq AI Integration
Author: Amal SP  
Created: December 2025
"""

from fastapi import FastAPI, HTTPException
from fastapi.middleware.cors import CORSMiddleware
from pydantic import BaseModel
from typing import Optional, List, Dict, Any
import os
import logging
from groq import Groq
import json

# Configure logging
logging.basicConfig(
    level=logging.INFO,
    format='%(asctime)s - %(name)s - %(levelname)s - %(message)s'
)
logger = logging.getLogger(__name__)

# Initialize FastAPI app
app = FastAPI(
    title="OpenNL2SQL API",
    description="AI-powered Natural Language to SQL Analytics System",
    version="1.0.0"
)

# CORS configuration
app.add_middleware(
    CORSMiddleware,
    allow_origins=["*"],
    allow_credentials=True,
    allow_methods=["*"],
    allow_headers=["*"],
)

# Initialize Groq client
GROQ_API_KEY = os.getenv("GROQ_API_KEY")
if GROQ_API_KEY:
    groq_client = Groq(api_key=GROQ_API_KEY)
    logger.info("Groq client initialized successfully")
else:
    groq_client = None
    logger.warning("GROQ_API_KEY not found - running in demo mode")

# Request/Response Models
class QueryRequest(BaseModel):
    question: str
    session_id: Optional[str] = None

class QueryResponse(BaseModel):
    success: bool
    sql: Optional[str] = None
    results: Optional[List[Dict[str, Any]]] = None
    sql_explanation: Optional[str] = None
    results_explanation: Optional[str] = None
    error: Optional[str] = None
    session_id: str

def generate_sql_with_groq(question: str) -> tuple:
    """Generate SQL using Groq AI"""
    try:
        # Sample database schema
        schema = """
        Database Schema:
        - customers (id, name, email, created_at)
        - orders (id, customer_id, total, status, created_at)
        - products (id, name, price, category)
        - order_items (id, order_id, product_id, quantity, price)
        """
        
        prompt = f"""{schema}

Convert this natural language question to a SQL query:
Question: {question}

Generate ONLY a valid SELECT SQL query. No explanations.
SQL Query:"""

        response = groq_client.chat.completions.create(
            model="mixtral-8x7b-32768",
            messages=[
                {"role": "system", "content": "You are a SQL expert. Generate only valid SQL SELECT queries without any explanations or markdown formatting."},
                {"role": "user", "content": prompt}
            ],
            temperature=0.2,
            max_tokens=500
        )
        
        sql = response.choices[0].message.content.strip()
        # Clean up the SQL
        sql = sql.replace("```sql", "").replace("```", "").strip()
        
        return sql, None
    except Exception as e:
        logger.error(f"Error generating SQL: {str(e)}")
        return None, str(e)

def explain_sql_with_groq(sql: str, question: str) -> str:
    """Generate explanation for SQL query"""
    try:
        prompt = f"""Explain this SQL query in simple terms:
        
Original Question: {question}
SQL Query: {sql}

Provide a brief, clear explanation:"""

        response = groq_client.chat.completions.create(
            model="mixtral-8x7b-32768",
            messages=[
                {"role": "system", "content": "You are a helpful assistant that explains SQL queries in simple terms."},
                {"role": "user", "content": prompt}
            ],
            temperature=0.3,
            max_tokens=300
        )
        
        return response.choices[0].message.content.strip()
    except Exception as e:
        logger.error(f"Error explaining SQL: {str(e)}")
        return "SQL query generated successfully."

@app.get("/")
async def root():
    """Health check endpoint"""
    return {
        "status": "healthy",
        "service": "OpenNL2SQL API",
        "version": "1.0.0",
        "message": "FastAPI backend with Groq AI integration running on Hugging Face Spaces!",
        "groq_enabled": groq_client is not None
    }

@app.get("/health")
async def health_check():
    """Detailed health check"""
    return {
        "status": "healthy",
        "groq_api_configured": groq_client is not None,
        "service": "OpenNL2SQL API"
    }

@app.post("/query", response_model=QueryResponse)
async def process_query(request: QueryRequest):
    """Process natural language query with Groq AI"""
    session_id = request.session_id or "demo-session"
    
    # Check if Groq is available
    if not groq_client:
        return QueryResponse(
            success=False,
            error="GROQ_API_KEY not configured. Please add it in HF Spaces Settings > Variables.",
            session_id=session_id
        )
    
    try:
        # Generate SQL using Groq
        sql, error = generate_sql_with_groq(request.question)
        
        if error:
            return QueryResponse(
                success=False,
                error=f"Failed to generate SQL: {error}",
                session_id=session_id
            )
        
        # Generate explanation
        explanation = explain_sql_with_groq(sql, request.question)
        
        # For demo: return mock results
        # In production, you'd execute the SQL against a real database
        results = [
            {"info": "SQL generated successfully! In production, this would execute against your database."},
            {"note": "Connect your database to see real query results."}
        ]
        
        return QueryResponse(
            success=True,
            sql=sql,
            results=results,
            sql_explanation=explanation,
            results_explanation=f"Generated SQL query for: '{request.question}'. Ready to execute against your database.",
            session_id=session_id
        )
        
    except Exception as e:
        logger.error(f"Error processing query: {str(e)}")
        return QueryResponse(
            success=False,
            error=f"Error: {str(e)}",
            session_id=session_id
        )

if __name__ == "__main__":
    import uvicorn
    uvicorn.run(app, host="0.0.0.0", port=7860)