integration_examples.py

"""
Integration examples: How to use the agentic pipeline in your own code.

Run with: python integration_examples.py
"""

import os
import json
from sqlalchemy import create_engine

# Ensure OPENAI_API_KEY is set
if not os.environ.get("OPENAI_API_KEY"):
    print("⚠️  OPENAI_API_KEY not set. Skipping live examples.")
    print("Set it with: export OPENAI_API_KEY='sk-...'")
else:
    from agentic_pipeline import run_agentic_pipeline, AgenticPipeline
    from agent import verify_sql_query

    DB_URL = "sqlite:////workspaces/nattu22/data.db"

    print("=" * 80)
    print("NLP-to-SQL AGENTIC PIPELINE - INTEGRATION EXAMPLES")
    print("=" * 80)

    # EXAMPLE 1: Simple function call
    print("\n" + "=" * 80)
    print("EXAMPLE 1: Simple Pipeline Run")
    print("=" * 80)
    engine = create_engine(DB_URL)
    result = run_agentic_pipeline("List all customers from USA", engine, dialect="sqlite")
    if result.get("error"):
        print(f"❌ Error: {result['error']}")
    else:
        print(f"✓ Query successful")
        print(f"  SQL: {result['final_result']['sql']}")
        print(f"  Rows: {result['final_result']['rows_returned']}")
        print(f"  Narrative: {result['final_result']['narrative'][:200]}...")

    # EXAMPLE 2: Direct pipeline class for custom logic
    print("\n" + "=" * 80)
    print("EXAMPLE 2: Using AgenticPipeline Class")
    print("=" * 80)

    class CustomPipeline(AgenticPipeline):
        """Custom pipeline that logs each step."""

        def run(self, user_query: str):
            print(f"\n🔍 Starting custom pipeline for: {user_query}")
            result = super().run(user_query)
            print(f"📊 Pipeline completed with {len(result['traces'])} steps")
            return result

    pipeline = CustomPipeline(engine, dialect="sqlite", max_sql_retries=5)
    result = pipeline.run("Count orders by status")
    for i, trace in enumerate(result["traces"], 1):
        status = "✓" if not trace["error"] else "✗"
        print(f"  {i}. {status} {trace['step_name']:<25} {trace['duration_ms']:>8.1f}ms")

    # EXAMPLE 3: Verify SQL before running
    print("\n" + "=" * 80)
    print("EXAMPLE 3: SQL Verification (Safety Checks)")
    print("=" * 80)

    test_queries = [
        ("SELECT * FROM orders", True, "Valid SELECT"),
        ("INSERT INTO orders VALUES (1)", False, "CRUD: INSERT"),
        ("DELETE FROM orders WHERE id=1", False, "CRUD: DELETE"),
        ("DROP TABLE orders", False, "DDL: DROP"),
        ("UPDATE orders SET status='done'", False, "CRUD: UPDATE"),
        ("SELECT * FROM orders; DELETE FROM orders", False, "Multiple statements"),
        ("EXPLAIN SELECT * FROM orders", True, "Valid EXPLAIN"),
        ("DESCRIBE orders", True, "Valid DESCRIBE"),
    ]

    for sql, should_pass, label in test_queries:
        ok, msg = verify_sql_query(sql, dialect="sqlite")
        status = "✓" if ok == should_pass else "✗"
        result_str = "APPROVED" if ok else f"REJECTED: {msg}"
        print(f"  {status} [{label:25}] {result_str}")

    # EXAMPLE 4: Batch queries with aggregated traces
    print("\n" + "=" * 80)
    print("EXAMPLE 4: Batch Processing with Trace Aggregation")
    print("=" * 80)

    queries = [
        "How many customers are there?",
        "What is the average order value?",
        "List top 3 products by price",
    ]

    results = []
    total_time_ms = 0

    print("\nRunning batch...")
    for i, q in enumerate(queries, 1):
        print(f"  {i}. Processing: {q}")
        result = run_agentic_pipeline(q, engine, dialect="sqlite")
        if not result.get("error"):
            total_time = sum(t["duration_ms"] for t in result["traces"])
            total_time_ms += total_time
            results.append(
                {
                    "query": q,
                    "sql": result["final_result"]["sql"],
                    "rows": result["final_result"]["rows_returned"],
                    "duration_ms": total_time,
                }
            )
        else:
            print(f"     ❌ Error: {result['error']}")

    print(f"\n📊 Batch Summary ({len(results)} successful):")
    print(f"  Total time: {total_time_ms:.0f}ms")
    print(f"  Avg per query: {total_time_ms / len(results):.0f}ms")
    print(f"  Total rows: {sum(r['rows'] for r in results)}")

    # EXAMPLE 5: Trace visualization
    print("\n" + "=" * 80)
    print("EXAMPLE 5: Detailed Trace Visualization")
    print("=" * 80)

    result = run_agentic_pipeline("Show distinct countries in customers table", engine, dialect="sqlite")
    if not result.get("error"):
        print(f"\nQuery: {result['user_query']}")
        print(f"SQL: {result['final_result']['sql']}\n")
        print("Timeline:")

        traces = sorted(result["traces"], key=lambda t: sum(s["duration_ms"] for s in result["traces"][:result["traces"].index(t)] if s in result["traces"]))
        cumulative_ms = 0

        for trace in result["traces"]:
            status = "✓" if not trace["error"] else "✗"
            cumulative_ms += trace["duration_ms"]
            bar_width = int(trace["duration_ms"] / 100)
            bar = "█" * bar_width if bar_width > 0 else "·"
            print(f"  {status} {cumulative_ms:>6.0f}ms │{bar:<20} {trace['step_name']:<25} {trace['duration_ms']:>6.1f}ms")

print("\n" + "=" * 80)
print("✅ Integration examples complete!")
print("=" * 80)