"""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 {}