agents/validation.py

import openai

def query_validation_and_optimization(state):
    """Validates and optimizes the SQL query."""
    sql_query = state.get("generated_sql", "")
    
    if not sql_query:
        return {"optimized_sql": "", "validation_result": {"error": "No SQL query to validate."}}
    
    # Optimization
    prompt = f"""
    Validate and optimize the following SQL query. Check for:
    1. Syntax correctness
    2. Proper table and column references
    3. Performance optimization opportunities

    ```sql
    {sql_query}
    ```

    **Rules:**
    - DO NOT create or suggest any indexes.
    - DO NOT include any CREATE INDEX, ALTER TABLE, or similar DDL statements.
    - Focus only on query structure optimization without modifying the database schema.
    - Optimize filters and joins.
    - Use appropriate WHERE clauses and join conditions.
    - Output only the SQL query without markdown formatting.
    - The table name must be the full ecommerceaie5.table_name format.
    """
    response = openai.chat.completions.create(
        model="gpt-4o-mini",
        messages=[{"role": "user", "content": prompt}],
        temperature=0.0
    )
    
    optimized_query = response.choices[0].message.content.strip()
    # Remove markdown code block formatting if present
    if optimized_query.startswith("```sql"):
        optimized_query = optimized_query.replace("```sql", "").replace("```", "").strip()
    return {
        "optimized_sql": optimized_query,
        "validation_result": {"valid": True}
    }