app.py

from flask import Flask, request, jsonify
from flask_cors import CORS
import json
import requests
import chromadb
import logging
import time
from db import run_query
from config import (
    OLLAMA_URL, VECTOR_DB_PATH, API_HOST, API_PORT, API_DEBUG,
    LOG_LEVEL, LOG_FORMAT, OLLAMA_TIMEOUT, OLLAMA_MODEL, MAX_QUESTION_LENGTH
)

# Configure logging
logging.basicConfig(
    level=getattr(logging, LOG_LEVEL),
    format=LOG_FORMAT
)
logger = logging.getLogger(__name__)

app = Flask(__name__)

# Enable CORS for all origins and methods
CORS(app, origins=["http://localhost:3000", "http://localhost:3001", "http://127.0.0.1:3000", "http://127.0.0.1:3001"])

# Alternative: Enable CORS for all origins (less secure, use only for development)
# CORS(app, origins="*")

# ----------------- LOAD VECTOR DB -----------------
try:
    chroma_client = chromadb.PersistentClient(path=VECTOR_DB_PATH)
    collection = chroma_client.get_or_create_collection("schema_docs")
    logger.info("Vector database loaded successfully")
except Exception as e:
    logger.error(f"Failed to load vector database: {e}")
    raise

# ----------------- INPUT VALIDATION -----------------
def is_database_related_question(question):
    """Check if the question is related to database/company data"""
    question_lower = question.lower()
    
    # Database-related keywords
    db_keywords = [
        # Data query terms
        'employee', 'employees', 'worker', 'staff', 'person', 'people',
        'department', 'dept', 'division', 'team',
        'project', 'projects', 'task', 'assignment',
        'salary', 'pay', 'payment', 'wage', 'income',
        'attendance', 'present', 'absent', 'working',
        'leave', 'vacation', 'holiday', 'time off',
        'address', 'location', 'city', 'country',
        'role', 'position', 'job', 'designation',
        'login', 'user', 'username', 'password',
        
        # Query action terms
        'show', 'list', 'display', 'get', 'find', 'search',
        'how many', 'count', 'total', 'sum', 'average',
        'who', 'what', 'where', 'when', 'which',
        
        # Database terms
        'table', 'record', 'data', 'database', 'information'
    ]
    
    # Check if question contains database-related keywords
    for keyword in db_keywords:
        if keyword in question_lower:
            return True
    
    # If no database keywords found, it's likely not database-related
    return False

def validate_input(data):
    """Validate request payload"""
    if not data:
        return {"error": "Request body is required", "code": "MISSING_BODY"}, 400
    
    question = data.get("question", "").strip()
    if not question:
        return {"error": "Question field is required and cannot be empty", "code": "MISSING_QUESTION"}, 400
    
    if len(question) > MAX_QUESTION_LENGTH:
        return {"error": f"Question exceeds maximum length of {MAX_QUESTION_LENGTH} characters", "code": "QUESTION_TOO_LONG"}, 400
    
    # Basic SQL injection prevention
    dangerous_keywords = ['DROP', 'DELETE', 'INSERT', 'UPDATE', 'CREATE', 'ALTER', 'TRUNCATE']
    question_upper = question.upper()
    if any(keyword in question_upper for keyword in dangerous_keywords):
        return {"error": "Question contains potentially dangerous SQL keywords", "code": "DANGEROUS_KEYWORDS"}, 400
    
    # Check if question is database-related
    if not is_database_related_question(question):
        return {
            "error": "I can only help with questions about your company database. Please ask about employees, departments, projects, salaries, or other company data.",
            "code": "NON_DATABASE_QUESTION",
            "suggestion": "Try asking something like: 'How many employees do we have?' or 'Show me all departments'"
        }, 400
    
    return None, None

# ----------------- RAG: FETCH RELEVANT SCHEMA -----------------
def fetch_schema_context(question):
    """Fetch relevant schema context using vector search"""
    try:
        logger.info(f"Fetching schema context for question: {question[:50]}...")
        start_time = time.time()
        
        results = collection.query(query_texts=[question], n_results=5)
        docs = results["documents"][0]
        
        elapsed_time = time.time() - start_time
        logger.info(f"Schema context fetched in {elapsed_time:.2f} seconds, found {len(docs)} relevant documents")
        
        return "\n\n".join(docs)
    except Exception as e:
        logger.error(f"Vector DB fetch error: {e}")
        return ""

# ----------------- CALL OLLAMA -----------------
def ask_ollama(prompt, temperature=0.1):
    """Call Ollama API with timeout and error handling"""
    try:
        logger.info("Calling Ollama API")
        start_time = time.time()
        
        payload = {
            "model": OLLAMA_MODEL,
            "prompt": prompt,
            "temperature": temperature
        }
        
        response = requests.post(
            OLLAMA_URL, 
            json=payload, 
            stream=True, 
            timeout=OLLAMA_TIMEOUT
        )
        response.raise_for_status()
        
        result = ""
        for line in response.iter_lines():
            if line:
                data = line.decode("utf-8")
                try:
                    json_obj = json.loads(data)
                    result += json_obj.get("response", "")
                except json.JSONDecodeError:
                    continue
        
        elapsed_time = time.time() - start_time
        logger.info(f"Ollama API call completed in {elapsed_time:.2f} seconds")
        
        return result.strip()
        
    except requests.exceptions.Timeout:
        logger.error("Ollama API request timed out")
        raise Exception("AI service timeout - please try again")
    except requests.exceptions.ConnectionError:
        logger.error("Failed to connect to Ollama API")
        raise Exception("AI service unavailable")
    except requests.exceptions.RequestException as e:
        logger.error(f"Ollama API request failed: {e}")
        raise Exception("AI service error")
    except Exception as e:
        logger.error(f"Unexpected error calling Ollama: {e}")
        raise Exception("AI service error")

def generate_result_summary(question, sql, result):
    """Generate a human-friendly summary of the SQL results"""
    try:
        row_count = len(result) if result else 0
        question_lower = question.lower()
        
        # Rule-based summary generation (no LLM needed - fast!)
        
        # Handle count queries
        if ('how many' in question_lower or 'count' in question_lower) and result and len(result) == 1 and len(result[0]) == 1:
            count_value = result[0][0]
            if 'employee' in question_lower:
                return f"There are {count_value} employees in the company."
            elif 'department' in question_lower:
                return f"There are {count_value} departments."
            elif 'project' in question_lower:
                return f"There are {count_value} projects."
            else:
                return f"The count is {count_value}."
        
        # Handle search queries for names
        if any(word in question_lower for word in ['named', 'name', 'called']) and 'john' in question_lower:
            if row_count == 0:
                return "No employees found with names starting with 'John'."
            elif row_count == 1:
                if result and len(result[0]) > 0:
                    full_name = result[0][0] if result[0][0] else "Unknown"
                    return f"Found 1 employee: {full_name}"
                else:
                    return "Found 1 employee matching your criteria."
            else:
                return f"Found {row_count} employees with names starting with 'John'."
        
        # Handle general search queries
        if any(word in question_lower for word in ['find', 'search', 'show', 'list', 'get']):
            if row_count == 0:
                return "No records found matching your criteria."
            elif row_count == 1:
                return "Found 1 record matching your criteria."
            else:
                return f"Found {row_count} records matching your criteria."
        
        # Handle salary/financial queries
        if 'salary' in question_lower:
            if row_count == 0:
                return "No salary information found."
            else:
                return f"Retrieved salary information for {row_count} record(s)."
        
        # Default summary
        if row_count == 0:
            return "No results found for your query."
        elif row_count == 1:
            return "Found 1 result for your query."
        else:
            return f"Found {row_count} results for your query."
                
    except Exception as e:
        logger.error(f"Error generating summary: {e}")
        # Emergency fallback
        row_count = len(result) if result else 0
        return f"Query completed. Found {row_count} record(s)." if row_count > 0 else "No results found."

# ----------------- SQL PROMPT CREATION -----------------
def build_prompt(question, context):
    """Build prompt for SQL generation"""
    return f"""
You are an expert MySQL SQL generator.
Convert the natural language question into a valid MySQL SELECT query using ONLY the provided schema.

### SCHEMA
{context}

### RULES
1. Output ONLY the SQL query (no explanations, no markdown, no comments)
2. Use only tables and columns that exist in the schema above
3. Use proper MySQL syntax
4. For search queries with names, use LIKE with wildcards: WHERE column_name LIKE 'pattern%'
5. Be careful with JOIN conditions and WHERE clauses
6. Avoid GROUP BY unless explicitly needed for aggregation
7. Use proper table aliases if joining tables
8. Always end with semicolon

### EXAMPLES
Question: "Find employees named John"
SQL: SELECT * FROM employee_master WHERE full_name LIKE 'John%';

Question: "How many employees are there?"
SQL: SELECT COUNT(*) FROM employee_master WHERE active_status = 1;

### USER QUESTION
{question}

SQL:
"""
# ----------------- API ROUTES -----------------
@app.route("/nl2sql", methods=["POST"])
def nl2sql():
    """Convert natural language to SQL and execute query"""
    try:
        # Validate input
        data = request.get_json()
        validation_error, status_code = validate_input(data)
        if validation_error:
            logger.warning(f"Input validation failed: {validation_error}")
            return jsonify(validation_error), status_code
        
        question = data.get("question", "").strip()
        logger.info(f"Processing question: {question}")
        
        start_time = time.time()
        
        # Step 1 - Retrieve relevant schema from vector DB
        context = fetch_schema_context(question)
        if not context:
            logger.warning("No schema context found for question")
            return jsonify({
                "error": "Unable to find relevant database schema for your question",
                "code": "NO_SCHEMA_CONTEXT"
            }), 400

        # Step 2 - Generate SQL via LLM
        sql = ask_ollama(build_prompt(question, context))
        if not sql:
            logger.error("Ollama returned empty SQL")
            return jsonify({
                "error": "Failed to generate SQL query",
                "code": "SQL_GENERATION_FAILED"
            }), 500

        # Step 3 - Execute SQL
        try:
            result = run_query(sql)
        except Exception as db_error:
            logger.error(f"Database execution error: {db_error}")
            # Convert technical database errors to user-friendly messages
            error_str = str(db_error).lower()
            
            if "unknown column" in error_str:
                user_message = "I couldn't find the requested information in the database. The data you're looking for might not be available or might be stored differently."
            elif "syntax error" in error_str or "sql syntax" in error_str:
                user_message = "I had trouble understanding your request. Could you please rephrase your question?"
            elif "table doesn't exist" in error_str:
                user_message = "The information you're looking for is not available in our database."
            elif "group by" in error_str:
                user_message = "I had trouble organizing the data for your request. Could you try asking in a different way?"
            else:
                user_message = "I encountered an issue while searching the database. Please try rephrasing your question or ask something else."
            
            return jsonify({
                "error": user_message,
                "code": "DATABASE_ERROR",
                "suggestion": "Try asking about employees, departments, projects, or salaries in a simpler way."
            }), 400
        
        # Step 4 - Generate human-friendly summary
        summary = generate_result_summary(question, sql, result)
        
        total_time = time.time() - start_time
        logger.info(f"Request completed successfully in {total_time:.2f} seconds")
        
        return jsonify({
            "sql": sql,
            "result": result,
            "summary": summary,
            "execution_time": round(total_time, 2),
            "row_count": len(result)
        })
        
    except Exception as e:
        logger.error(f"Error processing request: {e}")
        # Handle different types of errors with user-friendly messages
        error_str = str(e).lower()
        
        if "timeout" in error_str or "ai service timeout" in error_str:
            user_message = "I'm taking longer than usual to process your request. Please try again with a simpler question."
        elif "connection" in error_str or "service unavailable" in error_str:
            user_message = "I'm having trouble connecting to our AI service. Please try again in a moment."
        else:
            user_message = "I encountered an unexpected issue. Please try asking your question in a different way."
        
        return jsonify({
            "error": user_message,
            "code": "PROCESSING_ERROR",
            "suggestion": "Try asking simpler questions about employees, departments, or projects."
        }), 500

@app.route("/", methods=["GET"])
def home():
    """Health check endpoint"""
    try:
        # Check vector DB connectivity
        collection.peek()
        
        return jsonify({
            "status": "NL2SQL API running",
            "version": "1.0.0",
            "services": {
                "vector_db": "connected",
                "ollama_url": OLLAMA_URL,
                "model": OLLAMA_MODEL
            }
        })
    except Exception as e:
        logger.error(f"Health check failed: {e}")
        return jsonify({
            "status": "NL2SQL API running with warnings",
            "error": str(e)
        }), 200

@app.errorhandler(404)
def not_found(error):
    """Handle 404 errors"""
    return jsonify({
        "error": "Endpoint not found",
        "code": "NOT_FOUND"
    }), 404

@app.errorhandler(405)
def method_not_allowed(error):
    """Handle 405 errors"""
    return jsonify({
        "error": "Method not allowed",
        "code": "METHOD_NOT_ALLOWED"
    }), 405

@app.errorhandler(500)
def internal_error(error):
    """Handle 500 errors"""
    logger.error(f"Internal server error: {error}")
    return jsonify({
        "error": "Internal server error",
        "code": "INTERNAL_ERROR"
    }), 500

if __name__ == "__main__":
    logger.info(f"Starting NL2SQL API on {API_HOST}:{API_PORT}")
    app.run(host=API_HOST, port=API_PORT, debug=API_DEBUG)