Spaces:
Sleeping
Sleeping
| """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] | |
| 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 | |
| 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, | |
| } | |
| 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 {} | |