73 lines
3.4 KiB
Python
73 lines
3.4 KiB
Python
"""Tools für hochgeladene Dateien (→ Konzept 7.7): Tabellen auflisten, beschreiben, mit SQL lesen."""
|
|
|
|
from pydantic import Field
|
|
|
|
from app.tools.context import ToolContext
|
|
from app.tools.registry import Args, TableResult, Tool, ToolError, register
|
|
from app.uploads import duck
|
|
|
|
SAMPLE_ROWS = 5
|
|
|
|
|
|
class NoArgs(Args):
|
|
pass
|
|
|
|
|
|
class DescribeArgs(Args):
|
|
table: str = Field(description="Name der Tabelle aus data_list_tables.")
|
|
|
|
|
|
class QueryArgs(Args):
|
|
sql: str = Field(min_length=1, max_length=10000,
|
|
description="Genau eine SELECT-Anweisung (DuckDB-SQL, WITH erlaubt). Für große Tabellen "
|
|
"aggregieren statt alle Zeilen zu lesen.")
|
|
|
|
|
|
def tables_of(context: ToolContext) -> list[duck.Table]:
|
|
if not context.tables:
|
|
raise ToolError("In diesem Workspace gibt es keine hochgeladenen Dateien.")
|
|
return context.tables
|
|
|
|
|
|
def list_tables(context: ToolContext, args: NoArgs) -> str:
|
|
tables = tables_of(context)
|
|
lines = []
|
|
with duck.readonly(tables, context.settings) as con:
|
|
for table in tables:
|
|
count = con.execute(f'SELECT count(*) FROM "{table.name}"').fetchone()[0]
|
|
described = con.execute(f'SELECT column_name, column_type FROM (DESCRIBE "{table.name}")').fetchall()
|
|
lines.append(f"- {table.name}: {count} Zeilen; Spalten: "
|
|
+ ", ".join(f"{name} ({kind})" for name, kind in described))
|
|
return "\n".join(lines)
|
|
|
|
|
|
def describe_table(context: ToolContext, args: DescribeArgs) -> str:
|
|
table = next((t for t in tables_of(context) if t.name == args.table), None)
|
|
if table is None:
|
|
raise ToolError(f"Tabelle {args.table!r} gibt es nicht. Verfügbar: "
|
|
+ ", ".join(t.name for t in context.tables))
|
|
_, described, _ = duck.query([table], f'SELECT column_name, column_type, "null" FROM (DESCRIBE "{table.name}")',
|
|
context.settings)
|
|
columns, sample, _ = duck.query([table], f'SELECT * FROM "{table.name}" LIMIT {SAMPLE_ROWS}', context.settings)
|
|
schema = "\n".join(f"- {name}: {kind}{'' if nullable == 'YES' else ' (nie leer)'}"
|
|
for name, kind, nullable in described)
|
|
return f"Tabelle {table.name}\nSpalten:\n{schema}\n\nErste {len(sample)} Zeilen:\n{duck.as_table(columns, sample)}"
|
|
|
|
|
|
def query(context: ToolContext, args: QueryArgs) -> str:
|
|
columns, rows, truncated = duck.query(tables_of(context), args.sql, context.settings)
|
|
text = duck.as_table(columns, rows) if columns else "(kein Ergebnis)"
|
|
text += f"\n\n{len(rows)} Zeile{'n' if len(rows) != 1 else ''}"
|
|
if truncated:
|
|
text += f", gekürzt auf {context.settings.query_max_rows}; für vollständige Aussagen aggregieren (COUNT, " \
|
|
"GROUP BY) statt Zeilen zu lesen"
|
|
return TableResult(text + ".", columns, rows, truncated)
|
|
|
|
|
|
register(Tool("data_list_tables", "uploads", "Hochgeladene Dateien als Tabellen auflisten, mit Zeilenzahl und "
|
|
"Spalten.", NoArgs, "read", list_tables))
|
|
register(Tool("data_describe_table", "uploads", "Spalten, Typen und die ersten Zeilen einer Tabelle zeigen.",
|
|
DescribeArgs, "read", describe_table))
|
|
register(Tool("data_query", "uploads", "Hochgeladene Daten mit einer SELECT-Anweisung (DuckDB-SQL) auswerten. "
|
|
"Nur lesend; Dateien, Netz und Schreiben sind gesperrt. Ergebnisse sind auf "
|
|
"wenige hundert Zeilen begrenzt.", QueryArgs, "read", query))
|