Spaces:
Running
Running
File size: 5,842 Bytes
5ea3240 b0af996 bb5d2bb 5ea3240 bb5d2bb 5ea3240 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 | from __future__ import annotations
import re
from typing import Any
import duckdb
import pandas as pd
from .llm import GeminiGateway
from .security import validate_readonly_sql
class SQLWorkspace:
def __init__(self):
self.conn = duckdb.connect(database=":memory:")
self.tables: list[str] = []
@staticmethod
def _safe_table_name(name: str) -> str:
cleaned = re.sub(r"[^A-Za-z0-9_]", "_", name).strip("_")
if not cleaned or cleaned[0].isdigit():
cleaned = "t_" + cleaned
return cleaned[:80].lower()
def add_dataframe(self, name: str, df: pd.DataFrame) -> str:
table = self._safe_table_name(name)
base = table
suffix = 2
while table in self.tables:
table = f"{base}_{suffix}"
suffix += 1
view = f"_df_{len(self.tables)}"
self.conn.register(view, df)
self.conn.execute(f'CREATE TABLE "{table}" AS SELECT * FROM "{view}"')
self.conn.unregister(view)
self.tables.append(table)
return table
def schema_text(self) -> str:
pieces = []
for table in self.tables:
rows = self.conn.execute(f'DESCRIBE "{table}"').fetchall()
cols = ", ".join(f"{r[0]} {r[1]}" for r in rows)
pieces.append(f"{table}({cols})")
return "\n".join(pieces)
def analytics_context(self, max_rows: int = 20) -> tuple[str, list[dict[str, Any]]]:
"""Build deterministic table evidence for corpus-level analytical synthesis.
This uses no LLM call. It exposes schema, bounded rows, numeric ranges and
categorical/boolean distributions so the normal grounded generation step
can synthesize trends across documents and structured data together.
"""
if not self.tables:
return "", []
blocks: list[str] = []
sources: list[dict[str, Any]] = []
for idx, table in enumerate(self.tables, start=1):
df = self.conn.execute(f'SELECT * FROM "{table}" LIMIT {max(1, int(max_rows))}').fetchdf()
desc_rows = self.conn.execute(f'DESCRIBE "{table}"').fetchall()
schema = ", ".join(f"{row[0]} {row[1]}" for row in desc_rows)
insights: list[str] = []
for col in df.columns:
series = df[col].dropna()
if series.empty:
continue
if pd.api.types.is_bool_dtype(series):
counts = series.astype(str).value_counts().to_dict()
insights.append(f"{col}: values={counts}")
elif pd.api.types.is_numeric_dtype(series):
insights.append(
f"{col}: min={series.min()}, max={series.max()}, mean={round(float(series.mean()), 3)}"
)
elif series.nunique(dropna=True) <= 8:
counts = series.astype(str).value_counts().to_dict()
insights.append(f"{col}: values={counts}")
preview = df.to_markdown(index=False) if len(df) else "(no rows)"
summary = "; ".join(insights[:12]) or "No compact descriptive statistics available."
block = (
f"[T{idx}] TABLE: {table}\n"
f"SCHEMA: {schema}\n"
f"DESCRIPTIVE SIGNALS: {summary}\n"
f"ROWS (bounded preview):\n{preview}"
)
blocks.append(block)
sources.append(
{
"id": f"T{idx}",
"type": "table",
"title": table,
"rows": int(self.conn.execute(f'SELECT COUNT(*) FROM "{table}"').fetchone()[0]),
"schema": schema,
"snippet": f"{summary}\n{preview}"[:1600],
}
)
return "\n\n".join(blocks), sources
def generate_sql(self, question: str, gateway: GeminiGateway) -> str:
if not self.tables:
raise ValueError("No CSV/XLSX tables are loaded in this session")
prompt = f"""You write DuckDB SQL for a read-only analytics assistant.
Available tables:\n{self.schema_text()}
Question: {question}
Return JSON with keys sql and rationale. The SQL must be a single SELECT or WITH query. Never modify data."""
data = gateway.complete_json(prompt, {"sql": "", "rationale": ""})
return validate_readonly_sql(str(data.get("sql", "")))
def execute_sql(self, sql: str) -> pd.DataFrame:
validated = validate_readonly_sql(sql)
return self.conn.execute(validated).fetchdf()
def benchmark_query(self, question: str, gateway: GeminiGateway) -> tuple[str, pd.DataFrame]:
"""Generate and execute SQL with one LLM call for component evaluation.
Routing is evaluated separately by the semantic-planner benchmark. The
Text2SQL component benchmark therefore avoids an extra planner call and
a second natural-language answer-generation call.
"""
sql = self.generate_sql(question, gateway)
return sql, self.execute_sql(sql)
def ask(self, question: str, gateway: GeminiGateway) -> tuple[str, str, list[dict[str, Any]]]:
sql = self.generate_sql(question, gateway)
result = self.execute_sql(sql)
preview = result.head(200)
result_md = preview.to_markdown(index=False) if len(preview) else "(no rows)"
answer_prompt = f"""Answer the user's data question using the SQL result below.
Question: {question}\nSQL: {sql}\nResult:\n{result_md}
Mention the computed result clearly. Do not invent values outside the table."""
answer = gateway.complete(answer_prompt)
sources = [{"id": "SQL1", "type": "sql", "title": "DuckDB query", "sql": sql, "rows": len(result)}]
return answer, sql, sources
|