Spaces:
Sleeping
Sleeping
| """`run_sql` — execute SQL against the active dataset's DuckDB. | |
| Identical behavior across all three data sources (bundled PKDD, uploaded | |
| file, S3 connector views) — the dataset layer normalises everything to a | |
| DuckDB substrate. | |
| Results are returned as a pandas DataFrame in `payload["df"]`. The | |
| `summary` field gives the LLM a compact preview (shape + head) so it | |
| doesn't have to ingest the whole result set. | |
| """ | |
| from __future__ import annotations | |
| import logging | |
| import pandas as pd | |
| from pydantic import BaseModel, Field | |
| from lexsi_ds.agent.context import AgentContext | |
| from lexsi_ds.agent.tools.base import Tool, ToolResult, ToolSpec | |
| log = logging.getLogger(__name__) | |
| # Preview cap for the LLM observation. Full DataFrame is in payload. | |
| _PREVIEW_ROWS = 20 | |
| _PREVIEW_CHARS = 4000 | |
| class RunSqlArgs(BaseModel): | |
| sql: str = Field(..., description="DuckDB-compatible SQL to execute.") | |
| label: str | None = Field( | |
| default=None, | |
| description=( | |
| "Optional human label for this query (e.g. 'context_df' or " | |
| "'top10_risky_loans'). Used by the UI/trace; the planner can " | |
| "pass it through so later tools can reference the result by name." | |
| ), | |
| ) | |
| def _run(args: RunSqlArgs, ctx: AgentContext) -> ToolResult: | |
| sql = (args.sql or "").strip().rstrip(";") | |
| if not sql: | |
| return ToolResult(ok=False, summary="run_sql: empty SQL", error="empty_sql") | |
| con = ctx.duck() | |
| try: | |
| df = con.execute(sql).df() | |
| except Exception as e: # noqa: BLE001 — surface the SQL error to the LLM | |
| return ToolResult( | |
| ok=False, | |
| summary=f"run_sql: execution failed: {type(e).__name__}: {str(e)[:300]}", | |
| error=str(e), | |
| ) | |
| finally: | |
| con.close() | |
| # Stash on context so explainability tools can find the row set by label. | |
| if args.label: | |
| ctx.cache[f"sql_result:{args.label}"] = df | |
| ctx.cache["last_sql_result"] = df | |
| ctx.cache["last_sql"] = sql | |
| preview = _preview(df) | |
| summary = ( | |
| f"run_sql: {len(df):,} rows × {len(df.columns)} cols" | |
| + (f" (label: {args.label})" if args.label else "") | |
| + f"\ncolumns: {list(df.columns)}\n" | |
| + f"head:\n{preview}" | |
| ) | |
| return ToolResult( | |
| ok=True, | |
| summary=summary, | |
| payload={"df": df, "n_rows": len(df), "columns": list(df.columns), | |
| "label": args.label}, | |
| ) | |
| def _preview(df: pd.DataFrame) -> str: | |
| if df.empty: | |
| return "(empty)" | |
| head = df.head(_PREVIEW_ROWS).to_string(index=False, max_cols=12) | |
| if len(head) > _PREVIEW_CHARS: | |
| head = head[:_PREVIEW_CHARS] + "\n…[truncated]" | |
| return head | |
| class _RunSqlTool: | |
| spec = ToolSpec( | |
| name="run_sql", | |
| description=( | |
| "Execute DuckDB-compatible SQL against the active dataset and " | |
| "return the result as a DataFrame. " | |
| "PREFER `text_to_sql` for any SQL involving JOINs, conditional " | |
| "aggregates, or string-value filters — that tool sees the full " | |
| "KG and is far less likely to pick the wrong join path. Use " | |
| "`run_sql` inline ONLY for trivial single-table queries (e.g. " | |
| "`SELECT AVG(amount) FROM fin_loan`) OR to execute SQL that " | |
| "`text_to_sql` returned in a previous step. " | |
| "Pass `label` to give the result a name later tools can " | |
| "reference (e.g. label='context_df' before `train_tabular_model`)." | |
| ), | |
| args_schema=RunSqlArgs, | |
| returns=( | |
| "summary: '<N> rows × <M> cols', columns list, and the head of " | |
| "the result. payload['df'] holds the full pandas DataFrame. " | |
| "0-row results on a string filter mean the value is wrong — " | |
| "call `sample_values` and retry." | |
| ), | |
| ) | |
| def run(self, args: RunSqlArgs, ctx: AgentContext) -> ToolResult: | |
| return _run(args, ctx) | |
| TOOL: Tool = _RunSqlTool() |