172 lines
6.4 KiB
Python
172 lines
6.4 KiB
Python
"""apply_actions 里的 $query：把字段值写成 {"$query": "SELECT ..."}，提交时替换成上传数据的查询结果
|
||
|
||
模型不用把查询结果逐行抄进动作，又慢又容易抄错；替换发生在 run_batch 之前，替换后照常做字段校验
|
||
结果是对象数组，键是查询结果的列名，正好对应图表物料的 dataset.source；series 的 encode 按列名引用
|
||
"""
|
||
|
||
from dataclasses import dataclass
|
||
from typing import Any
|
||
|
||
from pydantic import JsonValue
|
||
|
||
from app.actions.batch import BatchError
|
||
from app.agent.data_tools import DataContext
|
||
from app.config import settings
|
||
from app.schemas.actions import AddAction, AgentAction, UpdateAction
|
||
from app.schemas.screen import Screen
|
||
|
||
QUERY_KEY = "$query"
|
||
|
||
|
||
@dataclass(frozen=True)
|
||
class EmbeddedQuery:
|
||
"""一处替换：第几个动作、写到哪个字段、结果的列和行数，写进 apply_actions 的结果供模型核对"""
|
||
|
||
index: int
|
||
node_id: str
|
||
path: str
|
||
columns: list[str]
|
||
rows: int
|
||
|
||
|
||
def fill_queries(
|
||
actions: list[AgentAction], data: DataContext
|
||
) -> tuple[list[AgentAction], list[EmbeddedQuery]]:
|
||
"""替换 add 的 node 和 update 的 value 里所有的 $query；查询失败抛 BatchError，整批不生效
|
||
|
||
add 用了 $query 又没写 sample_data 时设为 false：数据来自用户上传的文件，不是示例数据
|
||
"""
|
||
filled: list[AgentAction] = []
|
||
embedded: list[EmbeddedQuery] = []
|
||
for index, action in enumerate(actions):
|
||
found: list[tuple[str, list[str], int]] = []
|
||
node_id = "" # 只有 add 和 update 有字段值，其他动作不会有替换
|
||
try:
|
||
if isinstance(action, AddAction):
|
||
node = _fill(
|
||
action.node.model_dump(mode="json", exclude_unset=True),
|
||
data,
|
||
"",
|
||
found,
|
||
)
|
||
if found:
|
||
fields = action.model_dump(mode="json", exclude_unset=True)
|
||
fields["node"] = node
|
||
fields.setdefault("sample_data", False)
|
||
action = AddAction.model_validate(fields)
|
||
node_id = action.node.id
|
||
elif isinstance(action, UpdateAction):
|
||
value = _fill(action.value, data, action.field_path, found)
|
||
if found:
|
||
action = action.model_copy(update={"value": value})
|
||
node_id = action.target_id
|
||
except ValueError as error: # 含 QueryError
|
||
raise BatchError(index, action, str(error)) from error
|
||
embedded.extend(
|
||
EmbeddedQuery(index, node_id, path, columns, rows)
|
||
for path, columns, rows in found
|
||
)
|
||
filled.append(action)
|
||
return filled, embedded
|
||
|
||
|
||
def _fill(
|
||
value: JsonValue,
|
||
data: DataContext,
|
||
path: str,
|
||
found: list[tuple[str, list[str], int]],
|
||
) -> JsonValue:
|
||
if isinstance(value, dict):
|
||
if QUERY_KEY in value:
|
||
return _query(value, data, path, found)
|
||
return {
|
||
key: _fill(item, data, _join(path, key), found)
|
||
for key, item in value.items()
|
||
}
|
||
if isinstance(value, list):
|
||
return [
|
||
_fill(item, data, _join(path, position), found)
|
||
for position, item in enumerate(value)
|
||
]
|
||
return value
|
||
|
||
|
||
def _query(
|
||
value: dict[str, Any],
|
||
data: DataContext,
|
||
path: str,
|
||
found: list[tuple[str, list[str], int]],
|
||
) -> list[dict[str, Any]]:
|
||
sql = value[QUERY_KEY]
|
||
if len(value) != 1 or not isinstance(sql, str):
|
||
raise ValueError(
|
||
f'{path} 的 $query 要写成 {{"$query": "SELECT ..."}}，对象里不能有其他键'
|
||
)
|
||
if not data.files:
|
||
raise ValueError("会话里没有上传的数据文件，不能使用 $query")
|
||
limit = settings.data_embed_max_rows
|
||
result = data.query(sql, max_rows=limit)
|
||
if result.truncated:
|
||
raise ValueError(
|
||
f"{path} 的 $query 结果超过 {limit} 行，组件放不下；在 SQL 里聚合、筛选或 LIMIT"
|
||
)
|
||
if duplicated := sorted(
|
||
{name for name in result.columns if result.columns.count(name) > 1}
|
||
):
|
||
raise ValueError(
|
||
f"{path} 的 $query 结果有重名的列 {'、'.join(duplicated)}，用 AS 起不同的别名"
|
||
)
|
||
found.append((path, result.columns, len(result.rows)))
|
||
return [dict(zip(result.columns, row, strict=True)) for row in result.rows]
|
||
|
||
|
||
def _join(path: str, key: str | int) -> str:
|
||
return f"{path}.{key}" if path else str(key)
|
||
|
||
|
||
def encode_problems(screen: Screen, node_ids: set[str]) -> list[str]:
|
||
"""用了 $query 的组件，series 的 encode 引用的列必须在 dataset.source 里
|
||
|
||
数据换成查询结果后列名变了，encode 还引用物料默认的列名时图表是空的，前端也不报错，所以在这里拦下
|
||
"""
|
||
problems = []
|
||
for node in screen.nodes:
|
||
if node.id not in node_ids:
|
||
continue
|
||
option = node.props.get("option")
|
||
if not isinstance(option, dict):
|
||
continue
|
||
dataset = option.get("dataset")
|
||
source = dataset.get("source") if isinstance(dataset, dict) else None
|
||
# 空结果或不是对象数组时无从核对
|
||
if not (
|
||
isinstance(source, list)
|
||
and source
|
||
and all(isinstance(row, dict) for row in source)
|
||
):
|
||
continue
|
||
columns = list(dict.fromkeys(key for row in source for key in row))
|
||
series = option.get("series")
|
||
items = series if isinstance(series, list) else [series]
|
||
wrong = []
|
||
for position, item in enumerate(items):
|
||
encode = item.get("encode") if isinstance(item, dict) else None
|
||
if not isinstance(encode, dict):
|
||
continue
|
||
for key, names in encode.items():
|
||
missing = [
|
||
name
|
||
for name in (names if isinstance(names, list) else [names])
|
||
if isinstance(name, str) and name not in columns
|
||
]
|
||
if missing:
|
||
wrong.append(
|
||
f"series.{position}.encode.{key} 用了 {'、'.join(missing)}"
|
||
)
|
||
if wrong:
|
||
problems.append(
|
||
f"组件 {node.id}：{'；'.join(wrong)}，但 dataset.source 的列是 {'、'.join(columns)}；"
|
||
"把 encode 改成查询结果的列名，或在 SQL 里用 AS 起对应的别名"
|
||
)
|
||
return problems
|