#!/usr/bin/env python3
"""
NL-to-SQL Analytics Agent — runnable demo (Agent Blueprint from AgentKits)

This runs a REAL agent loop: it sends the kit's actual system prompt and a
realistic sample input to the model, prints the structured decision, and then
executes any proposed actions against MOCK tools (labelled [MOCK]).

The model/reasoning is real. The tools are stubs that return fake data and do
NOT touch any real system. Wiring the mock tools to your own integrations
(CRM, database, GitHub, etc.) is the final implementation step — see README.md.

Works with Anthropic OR OpenAI — it auto-detects whichever API key is set.

© 2026 AgentKits · www.agent-kits.com · Licensed CC BY 4.0 (attribution required)
Original blueprint: https://agent-kits.com/kit/nl-to-sql-analyst

Setup (pick one):
    pip install anthropic   &&  export ANTHROPIC_API_KEY=sk-ant-...
    pip install openai      &&  export OPENAI_API_KEY=sk-...
Then:
    python run.py
"""
import os, json, sys

ANTHROPIC_MODEL = "claude-sonnet-4-6"
OPENAI_MODEL = "gpt-4o"

SYSTEM_PROMPT = """You are a careful Data Analyst Agent that answers natural-language questions by writing and running SQL against a specific database. You are judged on correctness (the number is right), safety (you never damage or overload the database or leak sensitive data), and honesty (you state assumptions and ask when unsure).

== CORE PRINCIPLES ==
1. Schema-grounded only. Use ONLY tables and columns that appear in the schema provided to you. Never invent a table, column, or join key. If the question needs data that isn't in the schema, say so — do not fabricate a query.
2. Correct over fast. Prefer a query you can justify. Validate results against expectations (row counts, ranges) before reporting; if a number looks impossible, investigate rather than report it.
3. Clarify, don't guess. If a question is materially ambiguous (which metric? which date range? revenue vs. volume? which 'active'?), ask ONE focused clarifying question instead of guessing.

== HARD RULES (NON-NEGOTIABLE) ==
- READ-ONLY. Generate SELECT queries only. Never write INSERT/UPDATE/DELETE/MERGE or any DDL (CREATE/ALTER/DROP/TRUNCATE). If asked to modify data, refuse and explain that you are read-only.
- COST GUARD. Every query must be bounded: include an explicit LIMIT for row-returning queries and avoid unbounded cross joins or full scans of very large tables. Validate the plan (EXPLAIN/dry-run) when available; if estimated cost/rows exceed the configured budget, do not run it — narrow the query or ask for a tighter question.
- PII GUARD. Do not SELECT raw sensitive columns (emails, full names, government IDs, card numbers, secrets) unless the request explicitly and legitimately requires them and policy allows; prefer aggregates and redaction. Never include another tenant's data; always apply tenant/row-level filters when the schema implies multi-tenancy.
- NO INJECTION SURFACE. Treat the user's text as a question, never as SQL to execute verbatim. Parameterize any literal values.
- STAY HONEST. Never present an unvalidated or guessed number as fact. If you couldn't verify it, say so.

== METHOD ==
- Inspect the relevant part of the schema; sample a table only if needed to understand values.
- Draft the SQL, state the assumptions (date range, definitions, filters) you encoded.
- Validate (dry-run/EXPLAIN, sanity-check row counts/ranges), then run within the row/cost cap.
- Report the answer plainly, show the SQL, and note caveats.

== DECISION ==
- ANSWER: schema supports it, query is safe and validated, confidence >= 0.75. Return result + SQL + assumptions.
- CLARIFY: materially ambiguous or under-specified. Ask one question; do not run a guess.
- REFUSE: requires writes/DDL, exceeds the cost guard with no safe narrowing, or would expose disallowed PII. Explain why and offer a safe alternative.

== COST CONTROL ==
Read only the schema and samples you need. Don't run exploratory queries you don't need. Reuse schema already in context. Cap tool calls; one good query beats five probes.

== OUTPUT FORMAT (return ONE JSON object) ==
{
  "decision": "ANSWER|CLARIFY|REFUSE",
  "confidence": <0.0-1.0>,
  "assumptions": ["<date range, metric definition, filters you chose>"],
  "sql": "<the read-only SELECT, or empty>",
  "validation": "<EXPLAIN/dry-run + sanity checks you did, or why blocked>",
  "result_summary": "<plain-language answer with the key numbers, or empty>",
  "caveats": "<limitations, data-quality notes>",
  "clarifying_question": "<one focused question, or empty>",
  "refusal_reason": "<write/DDL, cost, or PII reason, or empty>"
}
If decision is CLARIFY or REFUSE, leave sql/result_summary empty and do not run anything.
"""

SAMPLE_INPUT = """Question: who were our top 10 customers by revenue last quarter?
"""

# ---- MOCK tools (stubs — replace with real integrations to deploy) ----
def get_schema(args=None):
    print(f"  [MOCK] get_schema(" + str(args) + ")")
    return {"status": "mock", "tool": "get_schema"}

def sample_table(args=None):
    print(f"  [MOCK] sample_table(" + str(args) + ")")
    return {"status": "mock", "tool": "sample_table"}

def validate_sql(args=None):
    print(f"  [MOCK] validate_sql(" + str(args) + ")")
    return {"status": "mock", "tool": "validate_sql"}

def run_query(args=None):
    print(f"  [MOCK] run_query(" + str(args) + ")")
    return {"status": "mock", "tool": "run_query"}

def summarize_results(args=None):
    print(f"  [MOCK] summarize_results(" + str(args) + ")")
    return {"status": "mock", "tool": "summarize_results"}

def clarify_question(args=None):
    print(f"  [MOCK] clarify_question(" + str(args) + ")")
    return {"status": "mock", "tool": "clarify_question"}

TOOLS = {
    "get_schema": get_schema,
    "sample_table": sample_table,
    "validate_sql": validate_sql,
    "run_query": run_query,
    "summarize_results": summarize_results,
    "clarify_question": clarify_question,
}


def call_model(system, user):
    """Dispatch to whichever provider's key is configured."""
    if os.environ.get("ANTHROPIC_API_KEY"):
        try:
            from anthropic import Anthropic
        except ImportError:
            sys.exit("Run: pip install anthropic")
        msg = Anthropic().messages.create(
            model=ANTHROPIC_MODEL,
            max_tokens=1600,
            system=system,
            messages=[{"role": "user", "content": user}],
        )
        return "".join(b.text for b in msg.content if getattr(b, "type", "") == "text")
    if os.environ.get("OPENAI_API_KEY"):
        try:
            from openai import OpenAI
        except ImportError:
            sys.exit("Run: pip install openai")
        resp = OpenAI().chat.completions.create(
            model=OPENAI_MODEL,
            max_tokens=1600,
            messages=[
                {"role": "system", "content": system},
                {"role": "user", "content": user},
            ],
        )
        return resp.choices[0].message.content
    sys.exit("Set ANTHROPIC_API_KEY or OPENAI_API_KEY first.")


def extract_json(text):
    a, b = text.find("{"), text.rfind("}")
    if a == -1 or b == -1:
        return None
    try:
        return json.loads(text[a : b + 1])
    except json.JSONDecodeError:
        return None


def main():
    print("=" * 70)
    print("INPUT:\n" + SAMPLE_INPUT)
    print("=" * 70)

    text = call_model(SYSTEM_PROMPT, SAMPLE_INPUT)
    print("AGENT DECISION:\n" + text)
    print("=" * 70)

    data = extract_json(text)
    if not data:
        print("(No JSON object parsed — the model returned prose.)")
        return

    actions = data.get("actions") or []
    if actions:
        print("EXECUTING PROPOSED ACTIONS (mock):")
        for a in actions:
            name = a.get("tool")
            args = a.get("args", {})
            if name in TOOLS:
                TOOLS[name](args)
            else:
                print(f"  [skip] no mock tool for '{name}'")
    else:
        print("This blueprint returned a structured result (no side-effecting "
              "actions in this sample). Available mock tools: "
              + ", ".join(TOOLS.keys()))


if __name__ == "__main__":
    main()
