agent.py

"""Antern Bot — natural-language-to-SQL agent over the IAmInterviewed_QA DB.

Talks to any OpenAI-compatible LLM endpoint (Groq, ngrok AI Gateway, OpenAI,
self-hosted, ...) in a manual tool-calling loop:
  user question -> model writes T-SQL -> we run it READ-ONLY -> model explains.

The schema is placed in the system instruction so the model knows the tables,
columns, and joins available. Two tools are exposed: `run_sql` (run a read-only
query) and `present` (choose how to display the result: table / chart).
"""
from __future__ import annotations

import json

from openai import OpenAI

import config
import db

# How many result rows to show the MODEL (the full set still goes to the
# frontend). Keeps the prompt small and cheap.
MODEL_ROW_PREVIEW = 12
MAX_TOOL_ITERATIONS = 8

RUN_SQL_TOOL = {
    "type": "function",
    "function": {
        "name": "run_sql",
        "description": (
            "Run a single READ-ONLY T-SQL SELECT query against the "
            "IAmInterviewed_QA SQL Server database and return the rows. Only "
            "SELECT / WITH queries are permitted. Use this whenever you need "
            "data to answer the user."
        ),
        "parameters": {
            "type": "object",
            "properties": {
                "query": {
                    "type": "string",
                    "description": "A single T-SQL SELECT statement. Use TOP "
                    "(not LIMIT) to cap rows. Do not end with a semicolon.",
                }
            },
            "required": ["query"],
        },
    },
}

PRESENT_TOOL = {
    "type": "function",
    "function": {
        "name": "present",
        "description": (
            "Choose how the MOST RECENT run_sql result is displayed. Call this "
            "AFTER run_sql once you have the data. Use 'table' for multi-row "
            "results, a chart ('bar', 'line', 'pie') to visualize an "
            "aggregation, or 'text' for a single value / simple answer."
        ),
        "parameters": {
            "type": "object",
            "properties": {
                "format": {
                    "type": "string",
                    "enum": ["text", "table", "bar", "line", "pie"],
                    "description": "How to render the latest result.",
                },
                "title": {
                    "type": "string",
                    "description": "Short title/caption for the table or chart.",
                },
                "x_field": {
                    "type": "string",
                    "description": "For charts: column name for the category / "
                    "x-axis (also the pie slice labels).",
                },
                "y_fields": {
                    "type": "array",
                    "items": {"type": "string"},
                    "description": "For charts: one or more numeric column names "
                    "to plot on the y-axis (pie uses the first).",
                },
            },
            "required": ["format"],
        },
    },
}

SYSTEM_INTRO = """You are "Antern Bot", a helpful data assistant for the Antern \
recruitment / interview platform. You answer questions about the data in the \
IAmInterviewed_QA database (Microsoft SQL Server 2019).

How you work:
- When a question needs data, call the `run_sql` tool with a single T-SQL SELECT \
query, read the rows, then answer in clear natural language.
- This is SQL Server / T-SQL. Use `TOP n` (never `LIMIT`), `OFFSET ... FETCH` for \
paging, `GETDATE()` for the current time, and square brackets for reserved names.
- Use ONLY tables and columns that appear in the schema below. Do not guess table \
names, and give every table an alias and reference columns by that alias. Watch \
column types when joining (an int id only joins to an int id).
- You have READ-ONLY access. Never attempt INSERT/UPDATE/DELETE/DDL.
- Most tables use soft deletes: rows have IsActive (bit) and DeletedDate. Unless \
the user asks otherwise, filter to active rows (IsActive = 1) for "live" counts.
- Keep result sets small: aggregate (COUNT, SUM, GROUP BY) or use TOP for examples.
- If a query fails, READ the error and fix the query — do not repeat the same \
failing query.
- Explain results conversationally. Never invent data that isn't in the results.
- After you have the data, call the `present` tool to choose how it is shown: \
`table` for multi-row results, `bar`/`line`/`pie` to visualize an aggregation \
(pass x_field and y_fields as exact column names from your query), or `text` for a \
single value. For charts, aggregate in SQL first (GROUP BY / TOP n). When you show \
a table or chart, keep your written answer short — the visual carries the detail.

Note: every table also has standard audit columns not listed below — \
CreatedDate, ModifiedDate, DeletedDate (datetimeoffset) — in addition to the \
IsActive (bit) column shown.

Below is the database schema (each table with its meaningful columns, and the \
foreign-key relationships you can join on):
"""

# Minimal system prompt for the final "summarise the results" call — the schema
# isn't needed there, so we drop it to save tokens.
FINALIZE_SYSTEM = (
    "You are Antern Bot. The user's question and the SQL query results are in the "
    "conversation above. Write a clear, concise, friendly answer based only on "
    "those results. If a table or chart is being shown, keep your text short."
)


class AnternBot:
    def __init__(self) -> None:
        if not config.LLM_API_KEY:
            raise RuntimeError(
                "LLM_API_KEY is not set. Add it to your .env file. "
                "For the default Groq backend, get a free key at "
                "https://console.groq.com/keys"
            )
        self.model = config.LLM_MODEL
        self.client = OpenAI(
            base_url=config.LLM_BASE_URL, api_key=config.LLM_API_KEY, timeout=60
        )
        # Introspect once; the schema goes into the system instruction.
        self.schema = db.introspect_schema()
        self.system_instruction = SYSTEM_INTRO + "\n" + self.schema
        self.tools = [RUN_SQL_TOOL, PRESENT_TOOL]

    @staticmethod
    def _run_sql(sql: str) -> tuple:
        """Execute a query. Returns (record_for_ui, model_response, full_result).
        The model gets only a row preview; the full result is for the frontend."""
        record: dict = {"query": sql}
        try:
            result = db.run_query(sql)
            record["row_count"] = result["row_count"]
            record["truncated"] = result["truncated"]
            preview = result["rows"][:MODEL_ROW_PREVIEW]
            model_view = {
                "columns": result["columns"],
                "rows": preview,
                "row_count": result["row_count"],
                "preview_truncated": len(result["rows"]) > len(preview),
            }
            return record, {"result": model_view}, result
        except db.UnsafeQueryError as exc:
            record["error"] = f"Blocked: {exc}"
            return record, {"error": f"Query rejected by safety guard: {exc}"}, None
        except Exception as exc:  # SQL error, connection, etc.
            record["error"] = str(exc)
            return record, {"error": f"Query failed: {exc}"}, None

    @staticmethod
    def _present(args: dict, last_result) -> tuple:
        """Build a presentation directive paired with the most recent result."""
        fmt = (args.get("format") or "text").lower()
        if fmt == "text" or not last_result or not last_result.get("rows"):
            return None, {"status": "shown as text"}
        presentation = {
            "format": fmt,
            "title": args.get("title"),
            "x_field": args.get("x_field"),
            "y_fields": list(args.get("y_fields") or []),
            "columns": last_result["columns"],
            "rows": last_result["rows"],
            "truncated": last_result.get("truncated", False),
        }
        return presentation, {"status": f"shown as {fmt}"}

    def chat(self, history: list, user_message: str) -> dict:
        """Run one user turn through the tool-calling loop.

        `history` is the prior list of clean {role, content} text turns. Returns
        {answer, queries, presentation, messages}.
        """
        messages = [{"role": "system", "content": self.system_instruction}]
        messages.extend(history)
        messages.append({"role": "user", "content": user_message})

        queries: list[dict] = []
        presentation = None
        last_result = None
        presented = False
        answer = ""
        usage = {"prompt": 0, "completion": 0, "total": 0}
        llm_calls = 0

        for _ in range(MAX_TOOL_ITERATIONS):
            if presented:
                # Finalize call: just summarise the results — schema not needed,
                # so swap in the minimal system prompt to save tokens.
                call_messages = [
                    {"role": "system", "content": FINALIZE_SYSTEM}
                ] + messages[1:]
                kwargs = dict(model=self.model, messages=call_messages, temperature=0)
            else:
                kwargs = dict(
                    model=self.model, messages=messages, temperature=0,
                    tools=self.tools,
                )
            resp = self.client.chat.completions.create(**kwargs)
            llm_calls += 1
            u = getattr(resp, "usage", None)
            if u:
                usage["prompt"] += getattr(u, "prompt_tokens", 0) or 0
                usage["completion"] += getattr(u, "completion_tokens", 0) or 0
                usage["total"] += getattr(u, "total_tokens", 0) or 0
            msg = resp.choices[0].message
            content = (msg.content or "").strip()
            tool_calls = msg.tool_calls or []
            print(
                f"[antern] turn calls={[tc.function.name for tc in tool_calls]} "
                f"text={'yes' if content else 'no'}",
                flush=True,
            )

            # Echo the assistant turn back into the conversation.
            assistant_msg = {"role": "assistant", "content": msg.content or ""}
            if tool_calls:
                assistant_msg["tool_calls"] = [
                    {
                        "id": tc.id,
                        "type": "function",
                        "function": {
                            "name": tc.function.name,
                            "arguments": tc.function.arguments,
                        },
                    }
                    for tc in tool_calls
                ]
            messages.append(assistant_msg)

            if content:
                answer = content
            if not tool_calls:
                break

            # run_sql calls first (so a 'present' in the same batch uses fresh data).
            sql_calls = [tc for tc in tool_calls if tc.function.name == "run_sql"]
            other_calls = [tc for tc in tool_calls if tc.function.name != "run_sql"]
            for tc in sql_calls:
                args = self._args(tc)
                record, model_response, full = self._run_sql(args.get("query", ""))
                if full is not None:
                    last_result = full
                if record.get("error"):
                    print(f"[antern] sql error: {record['error']} | sql={record['query'][:200]}", flush=True)
                queries.append(record)
                messages.append({"role": "tool", "tool_call_id": tc.id,
                                 "content": json.dumps(model_response, default=str)})
            for tc in other_calls:
                if tc.function.name == "present":
                    presentation, model_response = self._present(self._args(tc), last_result)
                    presented = True
                else:
                    model_response = {"error": f"Unknown tool: {tc.function.name}"}
                messages.append({"role": "tool", "tool_call_id": tc.id,
                                 "content": json.dumps(model_response, default=str)})

        if not answer:
            answer = (
                "I found the data but had trouble summarising it. Please try "
                "rephrasing or narrowing your question."
            )

        new_history = list(history)
        new_history.append({"role": "user", "content": user_message})
        new_history.append({"role": "assistant", "content": answer})

        return {
            "answer": answer,
            "queries": queries,
            "presentation": presentation,
            "messages": new_history,
            "usage": usage,
            "llm_calls": llm_calls,
        }

    @staticmethod
    def _args(tool_call) -> dict:
        """Parse a tool call's JSON-string arguments into a dict."""
        try:
            return json.loads(tool_call.function.arguments or "{}")
        except Exception:
            return {}