agent.py

"""Agent loop — OpenAI Responses API with tool use + self-healing.

Loop:
    response = openai.responses.create(input, tools, [previous_response_id])
    while function_calls in response.output:
        for fc in function_calls:
            result = run(fc)
        response = openai.responses.create(
            input=function_call_outputs,
            previous_response_id=response.id,
            tools=tools,
        )

Tools exposed to the model:
  - web_search                    OpenAI-hosted live web search (built-in)
  - list_tables()
  - get_schema(tables: list[str])
  - run_sql(sql: str)             validated, read-only
  - submit_answer(...)            structured DB-answer tool (chart/SQL bundle)

For DB questions the model runs the DB tools and finalizes via `submit_answer`.
For general / current-events questions it can call `web_search` (executed by
OpenAI) and reply naturally; submit_answer is not required in that case.
If a tool fails, the error message is sent back verbatim so the model can
self-heal on the next turn.
"""
from __future__ import annotations

import json
import time
from dataclasses import dataclass, field
from typing import Any, Literal

from openai import OpenAI
from pydantic import BaseModel, Field, ValidationError

from src.config import Settings
from src.db import ReadOnlyDB
from src.pii import redact_dict_rows, redact_rows
from src.validator import InvalidSQL, validate_and_normalize


SYSTEM_PROMPT = """You are a helpful assistant with access to:

  • A read-only SQLite SALES database (tools: `list_tables`, `get_schema`, `run_sql`).
  • A `submit_answer` tool that finalizes a structured DB answer (SQL + chart spec).
  • A built-in `web_search` tool (executed by OpenAI) for live external information.

Pick the right path based on what the user asks:

A) DATA QUESTIONS (anything answerable from the SALES database)
   1. Call `list_tables` (unless you already know them).
   2. Call `get_schema` on the relevant tables — never invent column names.
   3. Write ONE `SELECT` statement. Call `run_sql` to execute it.
   4. If `run_sql` returns an error, READ IT, fix the SQL, and try again.
   5. Once you have rows that answer the question, call `submit_answer` with:
        - `sql`     : the final SQL.
        - `answer`  : a short natural-language answer in the user's language.
        - `reasoning`: one sentence on how you got there.
        - `chart`   : a chart spec (see chart guidance below).

   Hard rules:
     - SELECT only. The validator rejects INSERT/UPDATE/DELETE/DDL.
     - Always qualify columns when joining tables.
     - Prefer aggregates (SUM, AVG, COUNT) for quantitative questions.
     - Use date('now', '-30 days') style for "last N days".
     - If the question is ambiguous, make the most reasonable assumption and
       state it in your final answer.
     - Never call `submit_answer` before you have actually run a successful SELECT.

B) GENERAL or CURRENT-EVENTS QUESTIONS (not answerable from the database)
   - Use `web_search` whenever the question needs fresh, factual, or external
     information you are not certain of (news, prices, dates, definitions
     beyond your training cutoff).
   - Otherwise just answer directly from your own knowledge.
   - Reply naturally in the user's language using clear Markdown when helpful.
   - DO NOT call `submit_answer` for these — the UI only renders SQL/chart
     when submit_answer is used.

C) GREETINGS / SMALL TALK / META QUESTIONS ABOUT YOU
   - Reply briefly and warmly in the user's language. No tools needed.
   - Optionally hint at what you can do (analyze the sales DB, search the web).

D) MIXED QUESTIONS (DB metric + external fact)
   - You may use both: query the DB AND call `web_search`.
   - Finalize with `submit_answer`; include the web-derived facts inside the
     `answer` field. The chart still describes the DB result.

Chart guidance (only for `submit_answer`):
  - "bar"     — ranking / comparison across discrete categories (top-N, group totals).
  - "line"    — trend over an ordered axis (dates, months, sequential periods).
  - "pie"     — part-to-whole when there are ≤6 categories that sum meaningfully.
  - "scatter" — relationship between two numeric columns.
  - "none"    — single value, single row, or non-visualizable result.
  - `x` MUST be a column name from the executed query; `y` is the numeric column(s).
  - `title` is short (≤8 words) in the user's language.
"""


# ---------- Tool schema (OpenAI Responses API format) ----------
#
# Note the differences vs. Chat Completions:
#   - Function tools are FLAT: {"type": "function", "name": ..., "parameters": ...}
#     (no nested "function" key).
#   - Built-in tools like `web_search` are referenced just by their type.

TOOLS: list[dict[str, Any]] = [
    {"type": "web_search"},
    {
        "type": "function",
        "name": "list_tables",
        "description": "Return the list of tables available in the database.",
        "parameters": {"type": "object", "properties": {}, "required": []},
    },
    {
        "type": "function",
        "name": "get_schema",
        "description": (
            "Return columns, foreign keys and a few sample rows for the given tables. "
            "Call this before writing SQL so you know the exact column names and types."
        ),
        "parameters": {
            "type": "object",
            "properties": {
                "tables": {
                    "type": "array",
                    "items": {"type": "string"},
                    "description": "Table names to inspect.",
                }
            },
            "required": ["tables"],
        },
    },
    {
        "type": "function",
        "name": "run_sql",
        "description": (
            "Execute a single read-only SELECT statement and return the rows. "
            "If the SQL is invalid or fails at runtime, the error is returned "
            "so you can correct it and try again."
        ),
        "parameters": {
            "type": "object",
            "properties": {
                "sql": {"type": "string", "description": "A single SELECT statement."}
            },
            "required": ["sql"],
        },
    },
    {
        "type": "function",
        "name": "submit_answer",
        "description": (
            "Finalize a structured DB answer. Call this ONLY for data questions, "
            "after a successful run_sql. The UI uses this to render the SQL block "
            "and the chart. Do NOT call this for general / web / small-talk replies."
        ),
        "parameters": {
            "type": "object",
            "properties": {
                "sql": {"type": "string", "description": "The final SQL that produced the answer."},
                "answer": {"type": "string", "description": "Natural-language answer to the user."},
                "reasoning": {
                    "type": "string",
                    "description": "One sentence on how you arrived at the answer.",
                },
                "chart": {
                    "type": "object",
                    "description": (
                        "Visualization spec for the result. Use type='none' if the result "
                        "cannot meaningfully be charted (single value, single row, etc.)."
                    ),
                    "properties": {
                        "type": {
                            "type": "string",
                            "enum": ["bar", "line", "pie", "scatter", "none"],
                        },
                        "title": {"type": "string"},
                        "x": {
                            "type": "string",
                            "description": "Column name for the x / category / label axis.",
                        },
                        "y": {
                            "type": "array",
                            "items": {"type": "string"},
                            "description": "Numeric column name(s) to plot.",
                        },
                    },
                    "required": ["type"],
                },
            },
            "required": ["sql", "answer", "chart"],
        },
    },
]


# ---------- Output schema ----------

class TraceStep(BaseModel):
    turn: int
    tool: str
    arguments: dict[str, Any]
    result_preview: str
    is_error: bool = False
    duration_ms: int


class AgentResult(BaseModel):
    answer: str
    sql: str = ""
    reasoning: str = ""
    columns: list[str] = Field(default_factory=list)
    rows: list[list[Any]] = Field(default_factory=list)
    row_count_total: int = 0
    truncated: bool = False
    turns_used: int
    trace: list[TraceStep] = Field(default_factory=list)
    chart: dict[str, Any] | None = None
    response_id: str | None = None
    web_sources: list[dict[str, str]] = Field(default_factory=list)


class _ChartSpec(BaseModel):
    type: Literal["bar", "line", "pie", "scatter", "none"]
    title: str = ""
    x: str | None = None
    y: list[str] = Field(default_factory=list)


class _Submission(BaseModel):
    sql: str
    answer: str
    reasoning: str = ""
    chart: _ChartSpec | None = None


# ---------- Agent ----------

@dataclass
class _RunState:
    last_query_result: dict[str, Any] | None = None
    trace: list[TraceStep] = field(default_factory=list)


class TextToSqlAgent:
    def __init__(self, settings: Settings, db: ReadOnlyDB) -> None:
        self._settings = settings
        self._client = OpenAI(api_key=settings.openai_api_key)
        self._model = settings.openai_model
        self._db = db

    # ---------- Tool dispatch ----------

    def _exec_tool(self, name: str, args: dict[str, Any], state: _RunState) -> str:
        """Run a tool and return its content as a string for the model."""
        if name == "list_tables":
            return json.dumps({"tables": self._db.list_tables()}, ensure_ascii=False)

        if name == "get_schema":
            tables = args.get("tables") or []
            if not isinstance(tables, list) or not tables:
                raise ValueError("`tables` must be a non-empty array of table names.")
            schemas = []
            for t in tables:
                described = self._db.describe_table(t)
                if self._settings.enable_pii_redaction:
                    described["sample_rows"] = redact_dict_rows(described["sample_rows"])
                schemas.append(described)
            return json.dumps({"schemas": schemas}, ensure_ascii=False, default=str)

        if name == "run_sql":
            raw_sql = args.get("sql") or ""
            normalized = validate_and_normalize(
                raw_sql, max_limit=self._settings.max_query_limit
            )
            result = self._db.run_select(
                normalized, max_rows=self._settings.max_rows_to_model
            )
            if self._settings.enable_pii_redaction:
                result["rows"] = redact_rows(result["columns"], result["rows"])
            state.last_query_result = {"sql": normalized, **result}
            return json.dumps(
                {"executed_sql": normalized, **result},
                ensure_ascii=False,
                default=str,
            )

        raise ValueError(f"Unknown tool: {name}")

    # ---------- Main loop ----------

    def ask(
        self,
        question: str,
        previous_response_id: str | None = None,
    ) -> AgentResult:
        state = _RunState()
        submission: _Submission | None = None
        web_sources: list[dict[str, str]] = []

        # First call: send the user message + system instructions.
        next_input: list[dict[str, Any]] | str = [
            {
                "role": "user",
                "content": [{"type": "input_text", "text": question}],
            }
        ]
        current_response_id: str | None = previous_response_id
        last_response_id: str | None = previous_response_id

        for turn in range(1, self._settings.max_agent_turns + 1):
            kwargs: dict[str, Any] = {
                "model": self._model,
                "input": next_input,
                "tools": TOOLS,
                # Always re-send instructions so the rules apply on every turn.
                "instructions": SYSTEM_PROMPT,
            }
            if current_response_id:
                kwargs["previous_response_id"] = current_response_id

            response = self._client.responses.create(**kwargs)
            current_response_id = response.id
            last_response_id = response.id

            output_items = list(response.output or [])

            # Observe web_search calls (executed by OpenAI) for the trace.
            for item in output_items:
                if getattr(item, "type", None) == "web_search_call":
                    action = getattr(item, "action", None)
                    query = ""
                    if isinstance(action, dict):
                        query = str(action.get("query", "") or "")
                    elif action is not None:
                        query = str(getattr(action, "query", "") or "")
                    state.trace.append(
                        TraceStep(
                            turn=turn,
                            tool="web_search",
                            arguments={"query": query} if query else {},
                            result_preview=f"(executed by OpenAI · {getattr(item, 'status', 'completed')})",
                            duration_ms=0,
                        )
                    )

            # Collect URL citations from any assistant message in this response.
            for item in output_items:
                if getattr(item, "type", None) != "message":
                    continue
                for part in getattr(item, "content", []) or []:
                    if getattr(part, "type", None) != "output_text":
                        continue
                    for ann in getattr(part, "annotations", []) or []:
                        if getattr(ann, "type", None) == "url_citation":
                            url = getattr(ann, "url", "") or ""
                            if not url:
                                continue
                            if any(s["url"] == url for s in web_sources):
                                continue
                            web_sources.append(
                                {
                                    "url": url,
                                    "title": (getattr(ann, "title", "") or url)[:200],
                                }
                            )

            # Function calls the model wants us to execute.
            function_calls = [
                item for item in output_items
                if getattr(item, "type", None) == "function_call"
            ]

            if not function_calls:
                # Model finished without further tool calls.
                if submission is not None:
                    return self._finalize_submission(
                        submission, state, turn, last_response_id, web_sources
                    )

                # Natural-language reply (general / web / small talk).
                final_text = (response.output_text or "").strip()
                if not final_text:
                    final_text = (
                        "I can help with questions about the sales database "
                        "or general questions (with live web search when needed). "
                        "What would you like to know?"
                    )
                state.trace.append(
                    TraceStep(
                        turn=turn,
                        tool="(final_text)",
                        arguments={},
                        result_preview=_truncate(final_text, 500),
                        duration_ms=0,
                    )
                )
                last = state.last_query_result or {}
                return AgentResult(
                    answer=final_text,
                    sql="",
                    reasoning="",
                    columns=last.get("columns", []),
                    rows=last.get("rows", []),
                    row_count_total=last.get("row_count_total", 0),
                    truncated=last.get("truncated", False),
                    turns_used=turn,
                    trace=state.trace,
                    chart={"type": "none", "title": ""},
                    response_id=last_response_id,
                    web_sources=web_sources,
                )

            # Execute each function call locally and stage their outputs for the next turn.
            tool_outputs: list[dict[str, Any]] = []
            for fc in function_calls:
                name = getattr(fc, "name", "")
                call_id = getattr(fc, "call_id", "")
                raw_args = getattr(fc, "arguments", "") or "{}"
                t0 = time.perf_counter()
                try:
                    args = json.loads(raw_args)
                except json.JSONDecodeError as e:
                    content = f"Invalid JSON arguments: {e}"
                    state.trace.append(
                        TraceStep(
                            turn=turn,
                            tool=name,
                            arguments={},
                            result_preview=content,
                            is_error=True,
                            duration_ms=int((time.perf_counter() - t0) * 1000),
                        )
                    )
                    tool_outputs.append(
                        {"type": "function_call_output", "call_id": call_id, "output": content}
                    )
                    continue

                if name == "submit_answer":
                    try:
                        submission = _Submission.model_validate(args)
                        ack = "ok"
                        is_error = False
                    except ValidationError as e:
                        ack = f"Validation error on submit_answer: {e}"
                        is_error = True
                    state.trace.append(
                        TraceStep(
                            turn=turn,
                            tool=name,
                            arguments=args,
                            result_preview="(final answer submitted)" if not is_error else ack,
                            is_error=is_error,
                            duration_ms=int((time.perf_counter() - t0) * 1000),
                        )
                    )
                    tool_outputs.append(
                        {"type": "function_call_output", "call_id": call_id, "output": ack}
                    )
                    continue

                # Regular DB tool — surface errors back to the model so it can self-heal.
                try:
                    content = self._exec_tool(name, args, state)
                    is_error = False
                except InvalidSQL as e:
                    content = f"SQL rejected by validator: {e}"
                    is_error = True
                except Exception as e:
                    content = f"Tool error: {type(e).__name__}: {e}"
                    is_error = True

                state.trace.append(
                    TraceStep(
                        turn=turn,
                        tool=name,
                        arguments=args,
                        result_preview=_truncate(content, 500),
                        is_error=is_error,
                        duration_ms=int((time.perf_counter() - t0) * 1000),
                    )
                )
                tool_outputs.append(
                    {"type": "function_call_output", "call_id": call_id, "output": content}
                )

            # Feed tool outputs back to the model on the next turn.
            next_input = tool_outputs

        # Exceeded max turns. If submit_answer was already received, prefer it.
        if submission is not None:
            return self._finalize_submission(
                submission, state, self._settings.max_agent_turns, last_response_id, web_sources
            )
        raise RuntimeError(
            f"Agent exceeded {self._settings.max_agent_turns} turns without finishing."
        )

    # ---------- Finalizers ----------

    def _finalize_submission(
        self,
        submission: "_Submission",
        state: "_RunState",
        turn: int,
        response_id: str | None,
        web_sources: list[dict[str, str]],
    ) -> AgentResult:
        last = state.last_query_result or {}
        chart_spec = (
            _sanitize_chart(submission.chart, last.get("columns", []))
            if submission.chart is not None
            else None
        )
        return AgentResult(
            answer=submission.answer,
            sql=submission.sql,
            reasoning=submission.reasoning,
            columns=last.get("columns", []),
            rows=last.get("rows", []),
            row_count_total=last.get("row_count_total", 0),
            truncated=last.get("truncated", False),
            turns_used=turn,
            trace=state.trace,
            chart=chart_spec,
            response_id=response_id,
            web_sources=web_sources,
        )


# ---------- helpers ----------

def _truncate(s: str, n: int) -> str:
    return s if len(s) <= n else s[: n - 3] + "..."


def _sanitize_chart(spec: "_ChartSpec", result_columns: list[str]) -> dict[str, Any] | None:
    """Drop the chart spec if it references columns the result does not have.

    Keeps `type="none"` so the frontend knows the model intentionally skipped charting.
    """
    if spec.type == "none":
        return {"type": "none", "title": spec.title}
    cols = set(result_columns)
    if spec.x is not None and spec.x not in cols:
        return {"type": "none", "title": spec.title}
    if any(y not in cols for y in spec.y):
        return {"type": "none", "title": spec.title}
    return {
        "type": spec.type,
        "title": spec.title,
        "x": spec.x,
        "y": spec.y,
    }