148 lines
5.9 KiB
Python
148 lines
5.9 KiB
Python
"""数据文件的本地工具：describe_data 查看列概况，query_data 执行只读 SQL；问答和 screen agent 共用
|
||
|
||
apply_actions 里的 $query 也通过 DataContext.query 执行，见 tools.py
|
||
|
||
工具只能看到 DataContext 里的表，也就是当前用户在本会话里未过期的文件，查询引擎也只为这些表建视图
|
||
工具通过 runtime.context.data 读取 DataContext，用到它们的 agent 上下文都要有 data 字段；
|
||
上下文类型因 agent 而异，所以写 ToolRuntime[Any]（不写类型参数时按 None 处理，序列化会告警）
|
||
"""
|
||
|
||
from dataclasses import dataclass, field
|
||
from pathlib import Path
|
||
from typing import Annotated, Any
|
||
|
||
from langchain.tools import ToolRuntime, tool
|
||
from langchain_core.messages import ToolMessage
|
||
from pydantic import Field
|
||
|
||
from app.config import settings
|
||
from app.datafiles.query import MAX_ROWS, QueryError, QueryResult, run_query
|
||
from app.datafiles.storage import table_path
|
||
from app.schemas.data_file import DataFileRecord, DataTable
|
||
|
||
# 列超过这么多时，不指定 columns 只返回列名和类型，避免一次塞进太多概况
|
||
BRIEF_COLUMNS = 50
|
||
|
||
|
||
@dataclass(frozen=True)
|
||
class DataContext:
|
||
"""本次运行可用的数据文件；root 是数据文件的根目录，Parquet 路径由它和文件记录算出"""
|
||
|
||
files: tuple[DataFileRecord, ...] = ()
|
||
root: Path = field(default_factory=lambda: settings.data_file_dir)
|
||
|
||
def find(self, name: str) -> tuple[DataFileRecord, DataTable] | None:
|
||
for file in self.files:
|
||
for table in file.tables:
|
||
if table.table == name:
|
||
return file, table
|
||
return None
|
||
|
||
def paths(self) -> dict[str, Path]:
|
||
return {
|
||
table.table: table_path(self.root, file.file_id, table)
|
||
for file in self.files
|
||
for table in file.tables
|
||
}
|
||
|
||
def query(self, sql: str, *, max_rows: int = MAX_ROWS) -> QueryResult:
|
||
"""执行一条只读 SQL；失败时抛 QueryError，原因可以直接返回给模型修正"""
|
||
paths = self.paths()
|
||
# 任务执行期间被删除或清理的文件跳过，不影响查询其他表
|
||
existing = {name: path for name, path in paths.items() if path.exists()}
|
||
try:
|
||
return run_query(
|
||
sql, existing, timeout=settings.data_query_timeout, max_rows=max_rows
|
||
)
|
||
except QueryError as error:
|
||
if gone := [name for name in paths if name not in existing]:
|
||
raise QueryError(
|
||
f"{error}；{'、'.join(gone)} 的文件已删除或过期，不能再查询"
|
||
) from None
|
||
raise
|
||
|
||
|
||
def _error(runtime: ToolRuntime, message: str) -> ToolMessage:
|
||
# 返回给模型修正，不中断 agent
|
||
return ToolMessage(message, tool_call_id=runtime.tool_call_id, status="error")
|
||
|
||
|
||
def _table_list(data: DataContext) -> str:
|
||
return "、".join(data.paths()) or "无"
|
||
|
||
|
||
@tool
|
||
def describe_data(
|
||
table: str,
|
||
runtime: ToolRuntime[Any],
|
||
columns: Annotated[
|
||
list[str] | None,
|
||
Field(description="只看这些列；不传时看全部列"),
|
||
] = None,
|
||
) -> dict[str, Any] | ToolMessage:
|
||
"""查看一张数据表的列概况（类型、空值率、去重值个数、数值范围和四分位、常见值、样例）、随机样例行，以及解析时做的调整
|
||
|
||
分析或取数前先看要用的表；表名见请求 data_files 里的 table。统计值来自全表，样例只是其中几行
|
||
列很多时只返回列名和类型，再用 columns 指定要看的列
|
||
"""
|
||
data: DataContext = runtime.context.data
|
||
found = data.find(table)
|
||
if found is None:
|
||
return _error(runtime, f"没有表 {table}，可用的表：{_table_list(data)}")
|
||
file, item = found
|
||
profiled = item.profile.columns
|
||
result: dict[str, Any] = {
|
||
"table": item.table,
|
||
"label": item.label,
|
||
"rows": item.rows,
|
||
}
|
||
if columns is None and len(profiled) > BRIEF_COLUMNS:
|
||
result["columns"] = [
|
||
column.model_dump(
|
||
mode="json", include={"name", "type", "format"}, exclude_none=True
|
||
)
|
||
for column in profiled
|
||
]
|
||
result["note"] = (
|
||
f"共 {len(profiled)} 列，只列出列名和类型；用 columns 指定列查看详细概况和样例"
|
||
)
|
||
else:
|
||
index = {column.name: position for position, column in enumerate(profiled)}
|
||
names = list(index) if columns is None else list(dict.fromkeys(columns))
|
||
positions = [index[name] for name in names if name in index]
|
||
result["columns"] = [
|
||
profiled[position].model_dump(mode="json", exclude_none=True)
|
||
for position in positions
|
||
]
|
||
# 样例行的值和 columns 一一对应
|
||
result["sample_rows"] = [
|
||
[row[position] for position in positions] for row in item.profile.sample
|
||
]
|
||
if missing := [name for name in names if name not in index]:
|
||
result["missing_columns"] = missing
|
||
# 针对这张表的警告，以及针对整个文件的警告（例如跳过的坏行）
|
||
result["warnings"] = [
|
||
warning.message
|
||
for warning in file.warnings
|
||
if warning.table in (item.table, None)
|
||
]
|
||
return result
|
||
|
||
|
||
@tool
|
||
def query_data(sql: str, runtime: ToolRuntime[Any]) -> dict[str, Any] | ToolMessage:
|
||
"""对上传的数据表执行一条只读 SQL（DuckDB 语法），表名是请求 data_files 里的 table（t1、t2……）
|
||
|
||
中文或带符号的列名用双引号括起来；宽表转长表用 UNPIVOT，PIVOT 要用 IN 写明取值
|
||
最多返回 200 行，truncated 为 true 时结果不完整，改用聚合或 LIMIT
|
||
回答和图表里的数字以查询结果为准，不要自己心算
|
||
"""
|
||
data: DataContext = runtime.context.data
|
||
try:
|
||
return data.query(sql).to_dict()
|
||
except QueryError as error:
|
||
return _error(runtime, str(error))
|
||
|
||
|
||
DATA_TOOLS = [describe_data, query_data]
|