src/nodes/sql_executor.py

import sqlite3
import os
from src.state import AgentState

DB_PATH = os.getenv("DB_PATH", "data/olist.db")
MAX_RETRIES = 3

def execute_sql(state: AgentState) -> dict:
    """Execute SQL query with retry logic"""
    # Check if safety checker already flagged an error
    sql_result = state.get("sql_result")
    if sql_result and isinstance(sql_result, dict) and sql_result.get("error"):
        return {}  # safety checker already flagged this

    query = state.get("sql_query")
    if not query:
        return {"sql_result": {"error": "no_query_generated"}}

    retry = state.get("retry_count", 0)
    if retry >= MAX_RETRIES:
        return {
            "sql_result": {"error": "max_retries_exceeded"},
            "error_log": state.get("error_log", []) + ["[sql_executor] Max retries hit, giving up"]
        }

    try:
        conn = sqlite3.connect(DB_PATH)
        conn.row_factory = sqlite3.Row
        cursor = conn.execute(query)
        rows = [dict(row) for row in cursor.fetchall()]
        conn.close()

        if not rows:
            return {"sql_result": {"rows": [], "empty": True}, "retry_count": 0}
        else:
            result = {"sql_result": {"rows": rows, "count": len(rows)}, "retry_count": 0}
            
            # Persist entities from query results for next turn
            row = rows[0]
            session_context_updates = {}
            
            if row.get("order_id"):
                result["last_order_id"] = row["order_id"]
                session_context_updates["order_id"] = row["order_id"]
                print(f"[sql_executor] Persisted order_id: {row['order_id']}")
            
            if row.get("seller_id"):
                result["last_seller_id"] = row["seller_id"]
                session_context_updates["seller_id"] = row["seller_id"]
                print(f"[sql_executor] Persisted seller_id: {row['seller_id']}")
            
            # Handle comma-separated seller_ids (from aggregated queries)
            if row.get("seller_ids"):
                first_seller = row["seller_ids"].split(",")[0].strip()
                result["last_seller_id"] = first_seller
                session_context_updates["seller_id"] = first_seller
                print(f"[sql_executor] Persisted first seller_id: {first_seller}")
            
            if row.get("product_id"):
                result["last_product_id"] = row["product_id"]
                session_context_updates["product_id"] = row["product_id"]
                print(f"[sql_executor] Persisted product_id: {row['product_id']}")
            
            if session_context_updates:
                result["session_context"] = {**state.get("session_context", {}), **session_context_updates}
            
            return result

    except sqlite3.Error as e:
        return {
            "retry_count": retry + 1,
            "sql_result": {"error": str(e), "failed_query": query},
            "error_log": state.get("error_log", []) + [f"[sql_executor] SQL error (attempt {retry+1}): {str(e)}"]
        }
    except Exception as e:
        return {
            "sql_result": {"error": f"unexpected_error: {str(e)}"},
            "error_log": state.get("error_log", []) + [f"[sql_executor] Unexpected error: {str(e)}"]
        }