feat:优化增量执行、工具调用容错与评测体系

This commit is contained in:
2026-10-11 04:40:25 -07:00
parent cd5a2b4d8f
commit dadd797391
22 changed files with 2198 additions and 578 deletions
+1 -1
View File
@@ -30,7 +30,7 @@ MONGODB_CHECKPOINT_TTL=604800
# 上传的数据文件:存放目录(相对路径相对项目根目录,容器部署时挂载出来)、单个文件大小上限(MB)、保留时长(秒,默认 7 天) # 上传的数据文件:存放目录(相对路径相对项目根目录,容器部署时挂载出来)、单个文件大小上限(MB)、保留时长(秒,默认 7 天)
DATA_FILE_DIR=data/files DATA_FILE_DIR=data/files
DATA_FILE_MAX_MB=50 DATA_FILE_MAX_MB=500
DATA_FILE_TTL=604800 DATA_FILE_TTL=604800
# 上传后用 LLM 分析文件(一句话说明、图表推荐、列的含义),false 关闭;没有配置 LLM_API_KEY 时同样不分析 # 上传后用 LLM 分析文件(一句话说明、图表推荐、列的含义),false 关闭;没有配置 LLM_API_KEY 时同样不分析
DATA_FILE_INSIGHTS=true DATA_FILE_INSIGHTS=true
+39 -2
View File
@@ -7,7 +7,7 @@ from copy import deepcopy
from app.actions.layout import EXPANDERS from app.actions.layout import EXPANDERS
from app.actions.registry import apply_action from app.actions.registry import apply_action
from app.actions.update import fit_existing_path from app.actions.update import fit_existing_path, same_value
from app.resources.material_fields import check_update_field, find_material from app.resources.material_fields import check_update_field, find_material
from app.schemas.actions import Action, AddAction, AgentAction, UpdateAction from app.schemas.actions import Action, AddAction, AgentAction, UpdateAction
from app.schemas.capabilities import DATA_BINDING_ERROR, DATA_BINDING_SUPPORTED from app.schemas.capabilities import DATA_BINDING_ERROR, DATA_BINDING_SUPPORTED
@@ -60,12 +60,47 @@ def complete_add(action: AddAction, catalog: ComponentCatalog) -> AddAction:
return action.model_copy(update={"node": node, "sample_data": sample_data}) return action.model_copy(update={"node": node, "sample_data": sample_data})
def misplaced_props(actions: list[AgentAction], catalog: ComponentCatalog) -> list[str]:
"""新增组件的 props 里物料没有的键:前端按物料的结构读取,多出来的键不会生效,多半是写错了层级
评测里模型把 series 写在 props 下而不是 props.option 下,物料默认的 series 留着,图表引用的列对不上;
它看到的报错说 encode 用了别的列,和自己写的不一样,反复原样提交。返回给模型的说明,每项一个组件
"""
notes = []
for action in actions:
if not isinstance(action, AddAction):
continue
try:
known = (
find_material(catalog, action.node.type).template().get("props") or {}
)
except ValueError: # 物料不存在,补全时会报错
continue
written = (
action.node.model_dump(mode="json", exclude_unset=True).get("props") or {}
)
for key in sorted(set(written) - set(known)):
parent = next(
(
name
for name, value in known.items()
if isinstance(value, dict) and key in value
),
None,
)
where = f"(物料里在 props.{parent}.{key})" if parent else ""
notes.append(f"{action.node.id} 的 props.{key}{where}")
return notes
def run_batch( def run_batch(
screen: Screen, catalog: ComponentCatalog, actions: list[AgentAction] screen: Screen, catalog: ComponentCatalog, actions: list[AgentAction]
) -> tuple[list[Action], Screen]: ) -> tuple[list[Action], Screen]:
"""从原始大屏开始按顺序模拟执行,返回前端要执行的动作和执行后的大屏;第一个错误抛 BatchError """从传入的大屏开始按顺序模拟执行,返回前端要执行的动作和执行后的大屏;第一个错误抛 BatchError
布局意图(arrange、align、distribute、set_layout)在这里按当时的大屏展开成 move 和 resize 布局意图(arrange、align、distribute、set_layout)在这里按当时的大屏展开成 move 和 resize
update 写入的值和字段现在的值完全一样时没有变化,跳过:不经过字段契约、也不交付;
模型常把已经正确的值再写一遍(例如已经是 category 的轴类型),字段不在契约里时整批都会失败
""" """
current = screen.model_copy(deep=True) current = screen.model_copy(deep=True)
completed: list[Action] = [] completed: list[Action] = []
@@ -86,6 +121,8 @@ def run_batch(
None, None,
) )
if node is not None: if node is not None:
if same_value(node, action):
continue
check_update_field(catalog, node.type, action) check_update_field(catalog, node.type, action)
# 字段定义允许、但组件中还没有的中间对象,折叠进值里再执行 # 字段定义允许、但组件中还没有的中间对象,折叠进值里再执行
action = fit_existing_path(node, action) action = fit_existing_path(node, action)
+35 -2
View File
@@ -1,5 +1,6 @@
"""在大屏 JSON 副本上按路径更新组件属性,用于后端校验动作能否执行""" """在大屏 JSON 副本上按路径更新组件属性,用于后端校验动作能否执行"""
import json
from copy import deepcopy from copy import deepcopy
from app.actions.base import BaseActionHandler from app.actions.base import BaseActionHandler
@@ -7,6 +8,30 @@ from app.schemas.actions import UpdateAction
from app.schemas.screen import Screen, ScreenNode from app.schemas.screen import Screen, ScreenNode
def same_value(node: ScreenNode, action: UpdateAction) -> bool:
"""字段已经存在、并且值和要写入的完全一样(按 JSON 比较,1 和 1.0、true 和 1 不算一样)"""
if action.section == "name":
current = node.name
else:
current = getattr(node, action.section)
for segment in action.path:
if (
isinstance(current, dict)
and isinstance(segment, str)
and segment in current
) or (
isinstance(current, list)
and type(segment) is int
and segment < len(current)
):
current = current[segment]
else:
return False
return json.dumps(current, sort_keys=True) == json.dumps(
action.value, sort_keys=True
)
def fit_existing_path(node: ScreenNode, action: UpdateAction) -> UpdateAction: def fit_existing_path(node: ScreenNode, action: UpdateAction) -> UpdateAction:
"""字段定义允许的路径在组件中可能还没有中间对象(例如未开启过的数值标签) """字段定义允许的路径在组件中可能还没有中间对象(例如未开启过的数值标签)
@@ -17,7 +42,11 @@ def fit_existing_path(node: ScreenNode, action: UpdateAction) -> UpdateAction:
return action return action
current = getattr(node, action.section) or {} current = getattr(node, action.section) or {}
for index, segment in enumerate(action.path[:-1]): for index, segment in enumerate(action.path[:-1]):
if isinstance(current, dict) and isinstance(segment, str) and segment not in current: if (
isinstance(current, dict)
and isinstance(segment, str)
and segment not in current
):
rest = action.path[index + 1 :] rest = action.path[index + 1 :]
if not all(isinstance(item, str) for item in rest): if not all(isinstance(item, str) for item in rest):
return action return action
@@ -29,7 +58,11 @@ def fit_existing_path(node: ScreenNode, action: UpdateAction) -> UpdateAction:
) )
if isinstance(current, dict) and isinstance(segment, str): if isinstance(current, dict) and isinstance(segment, str):
current = current[segment] current = current[segment]
elif isinstance(current, list) and type(segment) is int and segment < len(current): elif (
isinstance(current, list)
and type(segment) is int
and segment < len(current)
):
current = current[segment] current = current[segment]
else: else:
return action return action
+3
View File
@@ -10,6 +10,7 @@ from langchain.agents import create_agent
from langchain.agents.middleware import ModelCallLimitMiddleware from langchain.agents.middleware import ModelCallLimitMiddleware
from app.agent.data_tools import DATA_TOOLS, DataContext from app.agent.data_tools import DATA_TOOLS, DataContext
from app.agent.tool_args import normalize_tool_args
from app.clients.llm import get_llm from app.clients.llm import get_llm
from app.config import settings from app.config import settings
from app.prompt.qa import QA_DATA_RULES, QA_SYSTEM_PROMPT from app.prompt.qa import QA_DATA_RULES, QA_SYSTEM_PROMPT
@@ -37,6 +38,8 @@ def get_qa_agent(with_data: bool = False):
), ),
context_schema=QAContext, context_schema=QAContext,
middleware=[ middleware=[
# 有工具时才挂参数兼容:挂了 wrap_tool_call 中间件,没有工具也会多出一个工具节点
*([normalize_tool_args] if with_data else []),
# 超过上限时抛 ModelCallLimitExceededError,由节点按失败处理; # 超过上限时抛 ModelCallLimitExceededError,由节点按失败处理;
# 写成 end 时中间件只插入一条英文说明,节点分不清是正常结束还是用尽了次数 # 写成 end 时中间件只插入一条英文说明,节点分不清是正常结束还是用尽了次数
ModelCallLimitMiddleware( ModelCallLimitMiddleware(
+28 -6
View File
@@ -16,6 +16,8 @@ from app.schemas.actions import AddAction, AgentAction, UpdateAction
from app.schemas.screen import Screen from app.schemas.screen import Screen
QUERY_KEY = "$query" QUERY_KEY = "$query"
# 图表物料的数据都在这里,series 的 encode 按它的列名引用;整个换成查询结果时数据就不再是示例数据
SOURCE_PATH = "props.option.dataset.source"
@dataclass(frozen=True) @dataclass(frozen=True)
@@ -124,10 +126,30 @@ def _join(path: str, key: str | int) -> str:
return f"{path}.{key}" if path else str(key) return f"{path}.{key}" if path else str(key)
def encode_problems(screen: Screen, node_ids: set[str]) -> list[str]: def _binding(node) -> tuple[Any, Any]:
"""用了 $query 的组件,series 的 encode 引用的列必须在 dataset.source 里 option = node.props.get("option")
if not isinstance(option, dict):
return None, None
return option.get("dataset"), option.get("series")
数据换成查询结果后列名变了,encode 还引用物料默认的列名时图表是空的,前端也不报错,所以在这里拦下
def rebound_nodes(before: Screen, after: Screen) -> set[str]:
"""这一批改过数据或映射(option.dataset、option.series)的组件,新增的也算
只按前后状态比较,不看是哪种动作改的:$query、直接写数据、单独改 encode 都会让两者对不上;没动过的旧组件不核对
"""
old = {node.id: _binding(node) for node in before.nodes}
return {
node.id
for node in after.nodes
if node.id not in old or old[node.id] != _binding(node)
}
def encode_problems(screen: Screen, node_ids: set[str]) -> list[str]:
"""指定组件的 series encode 引用的列必须在 dataset.source 里
数据换了列名、或 encode 改成了不存在的列时图表是空的,前端也不报错,所以在这里拦下
""" """
problems = [] problems = []
for node in screen.nodes: for node in screen.nodes:
@@ -161,11 +183,11 @@ def encode_problems(screen: Screen, node_ids: set[str]) -> list[str]:
] ]
if missing: if missing:
wrong.append( wrong.append(
f"series.{position}.encode.{key} 用了 {'、'.join(missing)}" f"props.option.series.{position}.encode.{key} 用了 {'、'.join(missing)}"
) )
if wrong: if wrong:
problems.append( problems.append(
f"组件 {node.id}:{';'.join(wrong)},但 dataset.source 的列是 {'、'.join(columns)};" f"组件 {node.id}:{';'.join(wrong)},但 props.option.dataset.source 的列是 {'、'.join(columns)};"
"把 encode 改成查询结果的列名,或在 SQL 里用 AS 起对应的别名" "把 encode 改成 dataset.source 里的列名;数据来自 $query 时也可以在 SQL 里用 AS 起对应的别名"
) )
return problems return problems
+96 -38
View File
@@ -5,27 +5,31 @@
组件目录和物料在物料 MCP,各类任务的做法在 skill;系统提示只放每次都要用、别处放不下的原则 组件目录和物料在物料 MCP,各类任务的做法在 skill;系统提示只放每次都要用、别处放不下的原则
""" """
from collections.abc import Sequence
from typing import Any, Literal, Self from typing import Any, Literal, Self
from langchain.agents import create_agent from langchain.agents import create_agent
from langchain.agents.middleware import ModelCallLimitMiddleware, wrap_model_call from langchain.agents.middleware import ModelCallLimitMiddleware
from langchain.agents.structured_output import ToolStrategy from langchain.agents.structured_output import ToolStrategy
from langchain_core.messages import AIMessage, HumanMessage from langchain_core.messages import AIMessage, BaseMessage, SystemMessage
from loguru import logger from loguru import logger
from pydantic import BaseModel, Field, model_validator from pydantic import BaseModel, Field, ValidationError, model_validator
from app.agent.data_tools import DATA_TOOLS from app.agent.data_tools import DATA_TOOLS
from app.agent.skills import skills_prompt from app.agent.skills import skills_prompt
from app.agent.tool_args import normalize_tool_args
from app.agent.tools import ( from app.agent.tools import (
ScreenAgentContext, ScreenAgentContext,
ScreenAgentState, ScreenAgentState,
apply_actions, apply_actions,
get_nodes, get_nodes,
last_write_failed,
) )
from app.clients.llm import get_llm from app.clients.llm import get_llm
from app.clients.mcp import get_mcp_tools from app.clients.mcp import get_mcp_tools
from app.config import settings from app.config import settings
from app.prompt.domain import INPUT_AS_DATA from app.prompt.domain import INPUT_AS_DATA
from app.schemas.actions import Action, AddAction
from app.schemas.clarification import Clarification from app.schemas.clarification import Clarification
SYSTEM_PROMPT = f"""你是大屏编辑助手,按用户的要求新建大屏,新增、修改、删除组件,或调整画布 SYSTEM_PROMPT = f"""你是大屏编辑助手,按用户的要求新建大屏,新增、修改、删除组件,或调整画布
@@ -45,13 +49,15 @@ SYSTEM_PROMPT = f"""你是大屏编辑助手,按用户的要求新建大屏,
class AgentReply(BaseModel): class AgentReply(BaseModel):
kind: Literal["done", "clarify", "unable"] = Field( kind: Literal["done", "clarify", "unable"] = Field(
description="done:最后一次 apply_actions 已成功,或请求不需要修改(例如只是提问);" description="done:要求的修改都已写入草稿,且最后一次 apply_actions 成功,草稿交付给用户;"
"clarify:需要用户回答才能继续;unable:无法完成" "或请求不需要修改(例如只是提问);clarify:需要用户回答才能继续;unable:无法完成。"
"clarify 和 unable 时草稿整个丢弃,用户看到的仍是原始大屏"
) )
message: str = Field( message: str = Field(
default="", default="",
description="给用户看的一两句话。done:说明改了什么,用了默认值或示例数据时一并说明,没有修改时直接回答;" description="给用户看的一两句话,不要提草稿、工具、参数这类内部说法。done:说明改了什么,用了默认值或示例数据时一并说明,没有修改时直接回答;"
"unable:说明原因,有可行的替代做法时一并给出;clarify:不用填", "unable:草稿整个丢弃,用户的大屏和请求前一样,不要说已经删除、新增或修改了什么;"
"说明没能完成什么和原因,有可行的替代做法时一并给出;clarify:不用填",
) )
clarification: Clarification | None = Field( clarification: Clarification | None = Field(
default=None, description="kind=clarify 时填写,其他情况不填" default=None, description="kind=clarify 时填写,其他情况不填"
@@ -64,39 +70,89 @@ class AgentReply(BaseModel):
return self return self
# 模型没有调用工具时最多纠正几次 FINALIZE_RULES = """<finalize>
NO_TOOL_CALL_RETRIES = 2 这次执行已经结束,不能再调用其他工具,也不能再修改大屏,只能调用 AgentReply 给出结论
NO_TOOL_CALL_HINT = ( 按上面的对话和执行情况判断:用户要求的修改没有全部完成时不能回复 done
"每一轮都要调用工具。要结束时调用 AgentReply,给用户的话写在 message 里" 上面 apply_actions 返回的"执行成功""已删除"说的是这次的草稿,不是用户看到的大屏;回复 unable 时草稿整个丢弃,用户的大屏没有任何变化
) </finalize>"""
@wrap_model_call def reached_call_limit(messages: Sequence[BaseMessage]) -> bool:
async def require_tool_call(request, handler): """是不是因为调用次数用尽而结束:没结束的每次模型调用都带工具调用,数到上限就是被截断的"""
"""ToolStrategy 要求每一轮都调用工具,结束时调用 AgentReply;模型偶尔只回文本(想直接回复用户)或空回复 calls = sum(
isinstance(item, AIMessage) and bool(item.tool_calls) for item in messages
)
return calls >= settings.max_agent_steps
原样重试往往得到同样的结果,或者模型转而重复提交 apply_actions 直到调用次数用尽,
所以把它的回复和纠正说明加在下一次调用里;只用于这一次调用,不写进 agent 的消息 def _describe(action: Action) -> str:
if isinstance(action, AddAction):
return f"add {action.node.id}"
return f"{action.type} {getattr(action, 'target_id', '画布')}"
def _run_state(result: dict, limit_hit: bool, texts: list[str]) -> str:
"""收尾时告诉模型执行到了哪一步;对话里有每次工具结果,这里只给整体情况"""
lines = [
"调用次数已用完,执行被截断"
if limit_hit
else "上一轮没有调用 AgentReply 给出结论"
]
if actions := result.get("actions"):
described = ";".join(_describe(item) for item in actions[:40])
# 只写"已成功提交"时,模型回复 unable 也会说"已删除了某某",用户以为修改已经生效
lines.append(
f"草稿里的动作 {len(actions)} 个,还没有交付:{described};"
"只有回复 done 才会交付,回复 unable 或 clarify 时草稿整个丢弃,用户的大屏保持原样"
)
else:
lines.append("还没有成功提交过修改")
if last_write_failed(result.get("messages", [])):
lines.append("最后一次 apply_actions 失败,那一批没有写入草稿")
if texts:
lines.append(
f"上一轮只输出了文本,没有调用 AgentReply:{' '.join(texts)[:1000]}"
)
return "\n".join(lines)
async def finalize_reply(result: dict, *, limit_hit: bool) -> AgentReply | None:
"""运行正常结束却没有结论时补一次结论:调用次数用尽,或模型只回了文本、空回复
只调用一次,只绑定 AgentReply 并强制调用,没有能修改大屏的工具,不会再进入修改循环;
说明放在系统消息里,不追加用户消息,否则模型会把内部要求当成用户的话来回复
强制调用也可能给出不合法的参数,校验不过就返回 None,由节点按失败处理,不重试
""" """
response = await handler(request) messages = list(result["messages"])
for _ in range(NO_TOOL_CALL_RETRIES): # 末尾没有工具调用的助手消息(模型的文本、空回复、调用上限的说明)不放进对话:
message = response.result[0] if response.result else None # 以助手消息结尾时部分接口会把它当成续写的开头;模型的文本写进说明里
if not isinstance(message, AIMessage) or message.tool_calls: trailing: list[AIMessage] = []
break while (
logger.warning( messages and isinstance(messages[-1], AIMessage) and not messages[-1].tool_calls
f"模型没有调用工具,纠正后重新调用: {str(message.content)[:200]}" ):
) trailing.insert(0, messages.pop())
# 空回复不放回去,部分接口不接受内容为空的助手消息 texts = (
replied = [message] if message.content else [] [] if limit_hit else [str(item.content) for item in trailing if item.content]
request = request.override( )
messages=[ run_state = _run_state(result, limit_hit, texts)
*request.messages, system = SystemMessage(
*replied, f"{SYSTEM_PROMPT}\n\n{FINALIZE_RULES}\n\n<run_state>\n{run_state}\n</run_state>"
HumanMessage(content=NO_TOOL_CALL_HINT), )
] model = get_llm().bind_tools([AgentReply], tool_choice="AgentReply")
) response = await model.ainvoke(
response = await handler(request) [system, *messages], config={"run_name": "finalize_reply", "tags": ["finalize"]}
return response )
call = next(
(item for item in response.tool_calls if item["name"] == "AgentReply"), None
)
if call is None:
logger.warning(f"收尾没有调用 AgentReply: {str(response.content)[:200]}")
return None
try:
return AgentReply.model_validate(call["args"])
except ValidationError as error:
logger.warning(f"收尾给出的结论不合法: {error}")
return None
# 是否挂数据工具 → agent # 是否挂数据工具 → agent
@@ -120,9 +176,11 @@ async def get_screen_agent(with_data: bool = False):
response_format=ToolStrategy(AgentReply), response_format=ToolStrategy(AgentReply),
state_schema=ScreenAgentState, state_schema=ScreenAgentState,
context_schema=ScreenAgentContext, context_schema=ScreenAgentContext,
# 超过上限时插入说明并结束;和模型只回文本、空回复一样没有结论,由节点调用 finalize_reply 收尾
# 模型没有调用工具时 agent 直接结束,不在循环里纠正:纠正说明会被当成用户的话,或让模型回到修改循环
middleware=[ middleware=[
require_tool_call, # 列表、对象参数被写成 JSON 字符串时先解析回来,见 tool_args.py
# 超过上限时插入说明并结束,没有结构化结果,由节点按失败处理 normalize_tool_args,
ModelCallLimitMiddleware( ModelCallLimitMiddleware(
run_limit=settings.max_agent_steps, exit_behavior="end" run_limit=settings.max_agent_steps, exit_behavior="end"
), ),
+264
View File
@@ -0,0 +1,264 @@
"""工具参数兼容:模型把列表、对象参数写成 JSON 字符串时,按工具的参数 schema 解析回来
部分模型(例如 claude-haiku-5.5)在嵌套较深、内容较长时,常把 apply_actions 的 actions 整个写成 JSON 字符串,
或把列表里的每一项写成字符串;参数校验报 "Input should be a valid list" 后它往往原样再交,评测里出现过
连续 10 次直到调用次数用尽。这类问题和哪个工具无关,所以在工具调用入口统一处理,对本地工具和 MCP 工具都生效
- 合法的输入不动:schema 里任何一处接受字符串(包括联合类型的某个成员、数组某个成员的 items)时不解析;
布尔 schema、解析不了的 $ref 这类看不准的情况按"什么都接受"处理,交给原有的参数校验
- 只处理两层:参数本身,以及列表参数里的每一项
- 写成字符串的长 JSON 多数只是结尾的括号多写或少写了一个:原文以括号结尾、模型输出没有因为长度被截断时,
按还没闭合的括号重新补齐结尾;截断时丢掉的内容补不回来,补齐反而会得到一个意思不同的合法动作,所以不补
- 仍然解析不出来时不执行,用中文说明错在哪里:参数校验的英文报错只说"应该是列表",模型会照旧写成字符串
- 解析或补齐后执行时,在工具结果末尾说明:补齐只能猜结尾,评测里模型把中间的 option 提前闭合、结尾又少一个括号,
补齐后 series 落到了 props 下,执行结果对不上而它看不出原因,反复原样提交
"""
import json
from dataclasses import replace
from typing import Any
from langchain.agents.middleware import wrap_tool_call
from langchain_core.messages import AIMessage, ToolMessage
from langchain_core.tools import BaseTool
from langgraph.types import Command
from loguru import logger
STRUCTURED = {"array", "object"}
ANY = frozenset({"string", "number", "integer", "boolean", "null", "array", "object"})
CLOSERS = {"[": "]", "{": "}"}
# 各家接口表示"输出达到长度上限被截断"的写法
TRUNCATED = {"length", "max_tokens"}
def _resolve(schema: Any, defs: dict) -> Any:
"""展开 $ref;指向的定义找不到时按 True(什么都接受)处理"""
while isinstance(schema, dict) and "$ref" in schema:
name = schema["$ref"].rsplit("/", 1)[-1]
if name not in defs:
return True
schema = defs[name]
return schema
def _types(schema: Any, defs: dict) -> set[str]:
"""schema 接受的 JSON 类型:anyOf、oneOf 取并集,allOf 取交集,没有类型约束时什么都接受"""
schema = _resolve(schema, defs)
if schema is False:
return set()
if not isinstance(schema, dict):
return set(ANY)
kind = schema.get("type")
if kind:
types = set(kind) if isinstance(kind, list) else {kind}
elif "properties" in schema:
types = {"object"}
elif "items" in schema:
types = {"array"}
else:
types = set(ANY)
for member in schema.get("allOf", []):
types &= _types(member, defs)
if members := [*schema.get("anyOf", []), *schema.get("oneOf", [])]:
types &= set().union(*(_types(member, defs) for member in members))
return types
def _items(schema: Any, defs: dict) -> list:
"""各个数组成员的 items schema;有数组成员没有限定 items 时,其中有一个 True(什么都接受)"""
schema = _resolve(schema, defs)
if not isinstance(schema, dict):
return [True] if schema is True else []
kind = schema.get("type")
own = set(kind) if isinstance(kind, list) else {kind}
found = []
if "prefixItems" in schema:
# 每个位置各有 schema 的元组,不逐项解析
found.append(True)
elif "items" in schema:
found.append(schema["items"])
elif "array" in own:
found.append(True)
for key in ("anyOf", "oneOf", "allOf"):
for member in schema.get(key, []):
found.extend(_items(member, defs))
return found
def _reclose(text: str) -> str | None:
"""去掉结尾连续的 ] 和 },按还没闭合的括号重新补齐;结尾停在字符串中间时补不了"""
body = text.rstrip().rstrip("]} \t\r\n")
stack: list[str] = []
in_string = escaped = False
for char in body:
if in_string:
if escaped:
escaped = False
elif char == "\\":
escaped = True
elif char == '"':
in_string = False
elif char == '"':
in_string = True
elif char in CLOSERS:
stack.append(char)
elif char in "]}" and (not stack or CLOSERS[stack.pop()] != char):
# 中间就有括号对不上,不是结尾的问题,补不了
return None
if in_string or not stack:
return None
return body + "".join(CLOSERS[item] for item in reversed(stack))
def _matches(value: Any, types: set[str]) -> bool:
return (isinstance(value, list) and "array" in types) or (
isinstance(value, dict) and "object" in types
)
def _parse(value: Any, types: set[str], repair: bool) -> tuple[Any, str | None, bool]:
"""schema 只接受列表或对象、却收到字符串时解析成 JSON;返回 (值, 解析不了时的说明, 是否补了结尾括号)"""
if not isinstance(value, str) or "string" in types or not types & STRUCTURED:
return value, None, False
expected = "数组" if "array" in types else "对象"
patched = False
try:
parsed = json.loads(value)
except json.JSONDecodeError as error:
# 原文不以括号结尾多半是被截断了,补齐结尾会丢掉后面的内容,不补
ends_closed = value.rstrip().endswith(("]", "}"))
repaired = _reclose(value) if repair and ends_closed else None
try:
parsed, patched = (json.loads(repaired) if repaired else None), True
except json.JSONDecodeError:
parsed = None
if not _matches(parsed, types):
problem = (
f"应该直接写成 JSON {expected},不要写成字符串;"
f"这个字符串也不是合法的 JSON(第 {error.pos} 个字符附近:{error.msg})"
)
return value, problem, False
if not _matches(parsed, types):
return value, f"应该直接写成 JSON {expected},不要写成字符串", False
return parsed, None, patched
def normalize_args(
tool: BaseTool, args: dict, *, repair: bool = True
) -> tuple[dict, list[str], list[str]]:
"""按工具的参数 schema 把写成字符串的列表、对象参数解析回来
返回新的参数(不修改传入的)、解析不了的参数说明、补了结尾括号的参数名;有说明时不该执行这次调用
repair=False 时不补括号,用于模型输出被截断的情况
"""
schema = tool.tool_call_schema
schema = schema if isinstance(schema, dict) else schema.model_json_schema()
defs = schema.get("$defs", {})
properties = schema.get("properties", {})
fixed, problems, repaired = {}, [], []
for name, value in args.items():
prop = properties.get(name)
if prop is None:
fixed[name] = value
continue
value, problem, patched = _parse(value, _types(prop, defs), repair)
if problem:
problems.append(f"{name} {problem}")
item_schemas = _items(prop, defs)
if isinstance(value, list) and item_schemas:
item_types = set().union(*(_types(item, defs) for item in item_schemas))
parsed_items = []
for index, item in enumerate(value):
item, problem, item_patched = _parse(item, item_types, repair)
if problem:
problems.append(f"{name} 的第 {index + 1} 项{problem}")
patched = patched or item_patched
parsed_items.append(item)
value = parsed_items
if patched:
repaired.append(name)
fixed[name] = value
return fixed, problems, repaired
def _truncated(state: Any, call_id: str) -> bool:
"""发起这次调用的模型输出是不是因为长度上限被截断"""
messages = state.get("messages", []) if isinstance(state, dict) else []
message = next(
(
item
for item in reversed(messages)
if isinstance(item, AIMessage)
and any(call["id"] == call_id for call in item.tool_calls)
),
None,
)
metadata = message.response_metadata if message is not None else {}
return bool(
{metadata.get("finish_reason"), metadata.get("stop_reason")} & TRUNCATED
)
@wrap_tool_call
async def normalize_tool_args(request, handler):
"""在参数校验前统一修正参数;没有改动时原样执行,解析不了时返回中文说明,不执行"""
tool = request.tool
args = request.tool_call.get("args")
if tool is None or not isinstance(args, dict):
return await handler(request)
truncated = _truncated(request.state, request.tool_call["id"])
fixed, problems, repaired = normalize_args(tool, args, repair=not truncated)
if problems:
logger.info(
f"工具 {tool.name} 的参数写成了不合法的 JSON 字符串"
f"{'(模型输出被截断)' if truncated else ''}: {problems}"
)
return ToolMessage(
"参数格式错误,这次调用没有执行:\n"
+ "\n".join(f"- {item}" for item in problems)
+ ("\n- 输出达到了长度上限被截断,分成几批提交" if truncated else ""),
tool_call_id=request.tool_call["id"],
name=tool.name,
status="error",
)
if fixed == args:
return await handler(request)
changed = [name for name in args if fixed[name] != args[name]]
strict = [name for name in changed if name not in repaired]
notes = []
if strict:
logger.info(
f"工具 {tool.name} 的参数 {'、'.join(strict)} 写成了 JSON 字符串,已解析"
)
notes.append(f"{'、'.join(strict)} 写成了 JSON 字符串,已解析后执行")
if repaired:
logger.info(
f"工具 {tool.name} 的参数 {'、'.join(repaired)} 写成了 JSON 字符串,结尾括号对不上,补齐后解析"
)
# 补的是结尾,括号多写或少写在中间时,补齐后层级是错的,但它是合法的 JSON,执行时不一定报错
notes.append(
f"{'、'.join(repaired)} 写成了 JSON 字符串,结尾的括号数量不对,已按还没闭合的括号补齐后执行;"
"如果中间的括号写错了,补齐后的层级也是错的,对照执行结果核对"
)
result = await handler(
request.override(tool_call={**request.tool_call, "args": fixed})
)
# 不告诉模型时它会一直写成字符串,补齐猜错层级时它也看不出原因
return _with_note(
result, "参数说明:" + ";".join(notes) + ";以后直接写成 JSON,不要写成字符串"
)
def _with_note(result: Any, note: str) -> Any:
"""在工具返回的消息末尾加一句说明;工具返回 Command 时改其中的工具消息"""
if isinstance(result, ToolMessage):
return result.model_copy(update={"content": f"{result.content}\n{note}"})
if isinstance(result, Command) and isinstance(result.update, dict):
messages = [
item.model_copy(update={"content": f"{item.content}\n{note}"})
if isinstance(item, ToolMessage)
else item
for item in result.update.get("messages", [])
]
return replace(result, update={**result.update, "messages": messages})
return result
+221 -69
View File
@@ -1,6 +1,7 @@
"""screen agent 的本地工具:读取本次请求的大屏、提交整批动作 """screen agent 的本地工具:读取当前大屏、在当前大屏上提交一批修改
它们读写的是本次请求的大屏和 agent 状态,所以放在本地;物料等能力由 MCP 提供 它们读写的是本次请求的大屏和 agent 状态,所以放在本地;物料等能力由 MCP 提供
读取、修改和校验都基于同一份草稿:原始大屏加上之前写入草稿的修改;回复 done 时草稿交付给用户
""" """
from collections.abc import Sequence from collections.abc import Sequence
@@ -9,44 +10,57 @@ from typing import Annotated, NotRequired
from langchain.agents import AgentState from langchain.agents import AgentState
from langchain.tools import ToolRuntime, tool from langchain.tools import ToolRuntime, tool
from langchain_core.messages import ToolMessage from langchain_core.messages import AIMessage, BaseMessage, ToolMessage
from langgraph.types import Command from langgraph.types import Command
from pydantic import Field from pydantic import Field
from app.actions.batch import BatchError, run_batch from app.actions.batch import BatchError, misplaced_props, run_batch
from app.agent.data_tools import DataContext from app.agent.data_tools import DataContext
from app.agent.query_values import EmbeddedQuery, encode_problems, fill_queries from app.agent.query_values import (
from app.schemas.actions import Action, AddAction, AgentAction SOURCE_PATH,
EmbeddedQuery,
encode_problems,
fill_queries,
rebound_nodes,
)
from app.schemas.actions import (
Action,
AddAction,
AgentAction,
DeleteAction,
UpdateAction,
)
from app.schemas.component_resource import ComponentCatalog from app.schemas.component_resource import ComponentCatalog
from app.schemas.screen import Screen from app.schemas.screen import Screen
@dataclass(frozen=True) @dataclass(frozen=True)
class ScreenAgentContext: class ScreenAgentContext:
screen: Screen # 本次请求的原始大屏,只读 screen: Screen # 本次请求的原始大屏,只读,交付的动作从它开始执行
catalog: ComponentCatalog # 当前物料目录,只用于补全和校验 catalog: ComponentCatalog # 当前物料目录,只用于补全和校验
data: DataContext = field(default_factory=DataContext) # 本会话可用的数据表 data: DataContext = field(default_factory=DataContext) # 本会话可用的数据表
def _latest(_, new): # 两个字段只由 apply_actions 写,同一轮只执行第一次提交,所以一步里最多一次写入,不需要合并规则
# 同一步里有多次提交时取最后一个,避免并行工具调用写同一字段时报错 # 提交是否失败不另存字段,由 last_write_failed 从消息推出:参数解析、参数校验、查询、执行各处的失败都只体现在消息上
return new
class ScreenAgentState(AgentState): class ScreenAgentState(AgentState):
# 最近一次 apply_actions 成功时补全后的动作;失败时清空,以最后一次提交为准 # 草稿:最近一次写入后的大屏;还没有写入时不存在,用原始大屏
actions: NotRequired[Annotated[list[Action] | None, _latest]] working_screen: NotRequired[Screen]
# 最近一次 apply_actions 失败的原因,成功时清空 # 交付记录:各次写入的动作(补全、展开后)按顺序拼接,在原始大屏上依次执行得到草稿
apply_error: NotRequired[Annotated[str | None, _latest]] actions: NotRequired[list[Action]]
def current_screen(runtime: ToolRuntime[ScreenAgentContext]) -> Screen:
return runtime.state.get("working_screen") or runtime.context.screen
@tool @tool
def get_nodes(node_ids: list[str], runtime: ToolRuntime[ScreenAgentContext]) -> dict: def get_nodes(node_ids: list[str], runtime: ToolRuntime[ScreenAgentContext]) -> dict:
"""读取组件的完整配置(layout、props、style 等)。修改属性或迁移数据前先读取 """读取组件的完整配置(layout、props、style 等)。修改属性或迁移数据前先读取
读取的是原始大屏,不包含本次已提交的修改 读取的是草稿:原始大屏加上之前写入草稿的修改;和 apply_actions 在同一轮调用时读不到这次提交的结果,要核对就在下一轮读取
""" """
nodes = {node.id: node for node in runtime.context.screen.nodes} nodes = {node.id: node for node in current_screen(runtime).nodes}
return { return {
"nodes": [ "nodes": [
nodes[item].model_dump(mode="json") for item in node_ids if item in nodes nodes[item].model_dump(mode="json") for item in node_ids if item in nodes
@@ -61,13 +75,38 @@ def _result_text(
actions: list[Action], actions: list[Action],
undeclared_sample: frozenset[str] = frozenset(), undeclared_sample: frozenset[str] = frozenset(),
queries: Sequence[EmbeddedQuery] = (), queries: Sequence[EmbeddedQuery] = (),
unchanged: bool = False,
repeated: bool = False,
relabeled: Sequence[str] = (),
) -> str: ) -> str:
"""执行结果:除了动作数量,还给出执行后的画布和每个组件的布局 """执行结果:除了动作数量,还给出执行后的画布和每个组件的布局
模型提交后才第一次看到结果,能自己发现位置超出画布、互相重叠、没排整齐等问题再修正, 模型提交后才第一次看到结果,能自己发现位置超出画布、互相重叠、没排整齐等问题再修正;
不用靠重复提交来试;本次动过布局或新增的组件前面标 *,已删除的单独列出 这一批动过布局或新增的组件前面标 *,这一批删除的单独列出
没有带来变化时也照样给出布局:只说"已经写入"会让模型以为达到了目标,可能当前大屏本来就是这样、并不符合要求
统一说"写入草稿"、"回复 done 时交付":说"已生效"时模型回复 unable 也会说已经改了什么,
不说怎么交付时它会再交一批空的或重复的来"确认"
""" """
lines = [f"执行成功,共 {len(actions)} 个动作已生效"] if unchanged and repeated:
# 之前的提交已经改过:模型不确定是否写入才会重复提交,要写明已经在草稿里、不要重复,并说怎么结束,
# 只说"没有变化"时它以为修改没成功,会换着字段反复提交直到调用次数用尽
lines = [
(
"执行成功,但这一批没有改变草稿(0 个动作):提交的内容和草稿里已有的一样,之前的提交已经改过。"
"对照要求和下面的布局:已经达到就调用 AgentReply 回复 done,草稿会一起交付;"
"没达到就提交不同的修改,不要原样重复提交"
)
]
elif unchanged:
# 这次运行还没有改过大屏:原始大屏本来就是这样,模型的方案并没有改动任何东西,只陈述事实;
# 加上"完成就回复"之类的提示时,大屏并不符合要求它也会当成完成
lines = [
"执行成功,共 0 个动作写入草稿:提交的修改和原始大屏已有的一样,大屏原本就是这样"
]
else:
lines = [
f"执行成功,共 {len(actions)} 个动作已写入草稿;回复 done 时草稿交付给用户,不需要再提交确认"
]
sample = [ sample = [
action.node.id action.node.id
for action in actions for action in actions
@@ -78,10 +117,15 @@ def _result_text(
if guessed := [item for item in sample if item in undeclared_sample]: if guessed := [item for item in sample if item in undeclared_sample]:
lines.append( lines.append(
f"未写 sample_data,按物料默认标成了示例数据:{'、'.join(guessed)};" f"未写 sample_data,按物料默认标成了示例数据:{'、'.join(guessed)};"
"如果这些组件展示的数据都来自大屏上已有的组件或上传的数据文件,把 sample_data 设为 false 重新提交整批,否则不用处理" "如果这些组件展示的数据都来自大屏上已有的组件或上传的数据文件,"
"用 $query 整体替换它们的 option.dataset.source,或再提交一批:先 delete,再用 sample_data=false 重新 add,否则不用处理"
) )
if declared := [item for item in sample if item not in undeclared_sample]: if declared := [item for item in sample if item not in undeclared_sample]:
lines.append(f"标为示例数据的新增组件:{'、'.join(declared)}") lines.append(f"标为示例数据的新增组件:{'、'.join(declared)}")
if relabeled:
lines.append(
f"数据整体换成了 $query 的查询结果,已改标为非示例数据:{'、'.join(relabeled)}"
)
# 模型提交时只写了 SQL,看不到结果;给出列名和行数,用来核对 encode 和数据量 # 模型提交时只写了 SQL,看不到结果;给出列名和行数,用来核对 encode 和数据量
if queries: if queries:
lines.append("$query 已替换为查询结果:") lines.append("$query 已替换为查询结果:")
@@ -110,74 +154,182 @@ def _result_text(
return "\n".join(lines) return "\n".join(lines)
def _reply(
runtime: ToolRuntime, text: str, *, error: bool = False, **update
) -> Command:
"""返回工具消息;状态字段由调用方显式传入,没传的保持不变"""
status = "error" if error else "success"
message = ToolMessage(text, tool_call_id=runtime.tool_call_id, status=status)
return Command(update={**update, "messages": [message]})
def _first_write(message: AIMessage) -> str | None:
"""一条模型消息里第一个 apply_actions 调用的 ID;同一轮只执行它"""
return next(
(call["id"] for call in message.tool_calls if call["name"] == "apply_actions"),
None,
)
def _is_first_submission(runtime: ToolRuntime) -> bool:
"""按模型消息里 tool_calls 的顺序,这次调用是不是第一个 apply_actions
同一轮的多个调用都从同一份草稿出发,都执行会互相覆盖;只执行第一个,与并发谁先跑完无关
直接调用工具、找不到发起调用的模型消息时,视为第一个
"""
message = next(
(
item
for item in reversed(runtime.state.get("messages", []))
if isinstance(item, AIMessage)
and any(call["id"] == runtime.tool_call_id for call in item.tool_calls)
),
None,
)
return message is None or _first_write(message) == runtime.tool_call_id
def last_write_failed(messages: Sequence[BaseMessage]) -> bool:
"""最后一次被接受的提交是否失败:最后一条带 apply_actions 的模型消息里,第一个调用的结果是错误
参数解析、参数校验、查询、执行任何一处失败,结果都是 status="error" 的工具消息,都在这里判断,
不靠各处各自记状态,以后多一种失败也不会漏;同一轮被拒绝的后续调用不是第一个,不算
没有结果时也算失败:这一批没有执行,最后想做的修改没有写入草稿
"""
for position in range(len(messages) - 1, -1, -1):
message = messages[position]
if isinstance(message, AIMessage) and (call_id := _first_write(message)):
reply = next(
(
item
for item in messages[position + 1 :]
if isinstance(item, ToolMessage) and item.tool_call_id == call_id
),
None,
)
return reply is None or reply.status == "error"
return False
@tool @tool
def apply_actions( def apply_actions(
actions: Annotated[ actions: Annotated[
list[AgentAction], list[AgentAction],
# 空的一批会清掉之前成功的结果;参数校验失败不改状态,之前的结果保留 # 不限制至少一个:以前整批替换时空的一批会清掉之前的结果,增量执行下空的一批只是没有变化,
Field(min_length=1, description="至少一个动作;不需要修改时不要调用,直接回复"), # 按普通的无变化回复,比参数校验的英文报错更能让模型知道该结束了
Field(description="这一批要执行的动作;不需要修改时不要调用,直接回复"),
], ],
runtime: ToolRuntime[ScreenAgentContext], runtime: ToolRuntime[ScreenAgentContext],
) -> Command: ) -> Command:
"""提交本次要对大屏做的全部修改 """把一批修改写进这次的草稿;回复 done 时草稿交付给用户,不需要再提交确认
每次调用都是完整的一批,从原始大屏开始按顺序模拟执行,不在上一次调用的结果上累加 第一次提交从原始大屏开始,之后每一批都在草稿上继续修改;回复 unable 或 clarify 时草稿整个丢弃
成功时返回执行后的画布和各组件布局,用它核对结果,不对再重新提交整批 成功时返回这一批执行后的画布和各组件布局,用它核对结果;结果不对再提交一批修正,已经写入的动作不用重复提交
失败时整批不生效,返回带动作序号的错误,修正后重新提交整批;以最后一次调用的结果为准 一批里任何一个动作失败,这一批都不写入草稿,之前写入的保留;按带序号的错误修正这一批后再提交
同一轮只执行第一次调用,同一轮的 get_nodes 也读不到这次的结果;看到执行结果后再调用 AgentReply
""" """
if not _is_first_submission(runtime):
# 被拒绝的调用不写任何状态字段,不会覆盖第一次提交的结果
return _reply(
runtime,
"同一轮只执行第一次 apply_actions,这一批没有执行;需要的话在下一轮根据执行结果再提交",
error=True,
)
before = current_screen(runtime)
misplaced = misplaced_props(actions, runtime.context.catalog)
# 写错层级的属性不生效,失败和成功时都说明;它常是 encode 对不上的真正原因
hint = (
"\n新增组件写了物料没有的属性,前端不会使用:"
+ ";".join(misplaced)
+ ";检查是不是写错了层级"
if misplaced
else ""
)
try: try:
actions, queries = fill_queries(actions, runtime.context.data) actions, queries = fill_queries(actions, runtime.context.data)
completed, after = run_batch( completed, after = run_batch(before, runtime.context.catalog, actions)
runtime.context.screen, runtime.context.catalog, actions # 这一批改过数据或映射的组件都核对,不论数据来自这一批的 $query、之前的提交还是直接写入
) problems = encode_problems(after, rebound_nodes(before, after))
problems = encode_problems(after, {item.node_id for item in queries})
except BatchError as error: except BatchError as error:
problems = [str(error)] problems = [str(error)]
if problems: if problems:
message = "执行失败,整批未生效:\n" + "\n".join( # 失败只体现在工具消息上,由 last_write_failed 判断,不写状态字段
message = "执行失败,这一批未写入草稿,之前写入的修改保留:\n" + "\n".join(
f"- {item}" for item in problems f"- {item}" for item in problems
) )
return Command( return _reply(runtime, message + hint, error=True)
update={
"actions": None,
"apply_error": message,
"messages": [
ToolMessage(
message, tool_call_id=runtime.tool_call_id, status="error"
)
],
}
)
# 和上一次成功的提交完全一样:结果早已生效,再提交不会有变化,直接告诉模型结束,不再空转
if completed == runtime.state.get("actions"):
return Command(
update={
"messages": [
ToolMessage(
"这一批的结果和上一次提交完全相同,已经生效,不要重复提交;"
"任务完成就调用 AgentReply 回复用户,没完成就改成不同的内容",
tool_call_id=runtime.tool_call_id,
)
],
}
)
undeclared = frozenset( undeclared = frozenset(
action.node.id action.node.id
for action in actions for action in actions
if isinstance(action, AddAction) if isinstance(action, AddAction)
and "sample_data" not in action.model_fields_set and "sample_data" not in action.model_fields_set
) )
return Command( log = runtime.state.get("actions", [])
update={ delivered, relabeled = _relabel_real_data(
"actions": completed, [*log, *completed], _replaced_sources(actions, queries)
"apply_error": None,
"messages": [
ToolMessage(
_result_text(
runtime.context.screen, after, completed, undeclared, queries
),
tool_call_id=runtime.tool_call_id,
)
],
}
) )
# 按交付结果判断有没有变化,不比较参数:同样的 move 再执行一次会继续移动,同样的 update 再执行一次没有变化
# 交付结果除了大屏还有只在动作上的示例数据标记,只比大屏会把"重新 add 改 sample_data"当成没变化而丢掉
unchanged = after == before and _sample_marks(delivered) == _sample_marks(log)
# 这次运行之前已有成功提交时,没有变化多半是在重复自己的修改;还没有时说明原始大屏本来就是这样
text = _result_text(
before,
after,
delivered[len(log) :],
undeclared,
queries,
unchanged,
repeated=bool(log),
relabeled=relabeled,
)
if unchanged:
return _reply(runtime, text + hint)
return _reply(runtime, text + hint, working_screen=after, actions=delivered)
def _replaced_sources(
actions: Sequence[AgentAction], queries: Sequence[EmbeddedQuery]
) -> set[str]:
"""这一批用 update 把整个 dataset.source 换成 $query 结果的组件;只替换其中一部分时不算"""
return {
item.node_id
for item in queries
if item.path == SOURCE_PATH and isinstance(actions[item.index], UpdateAction)
}
def _relabel_real_data(
actions: list[Action], node_ids: set[str]
) -> tuple[list[Action], list[str]]:
"""数据整体换成查询结果的组件,把它最近一次 add 的示例数据标记改为 false
标记只在 add 动作上、不影响大屏,改它不影响"交付记录依次执行得到草稿";
原始大屏上的组件没有 add,跨请求的标记要等大屏和前端契约支持
"""
result = list(actions)
relabeled = []
for node_id in sorted(node_ids):
index = next(
(
position
for position in range(len(result) - 1, -1, -1)
if isinstance(result[position], AddAction)
and result[position].node.id == node_id
),
None,
)
if index is not None and result[index].sample_data:
result[index] = result[index].model_copy(update={"sample_data": False})
relabeled.append(node_id)
return result, relabeled
def _sample_marks(actions: Sequence[Action]) -> dict[str, bool]:
"""按交付记录,当前每个新增组件是否标为示例数据;这个标记只在 add 动作上,不在大屏里"""
marks: dict[str, bool] = {}
for action in actions:
if isinstance(action, AddAction):
marks[action.node.id] = action.sample_data
elif isinstance(action, DeleteAction):
marks.pop(action.target_id, None)
return marks
+3 -2
View File
@@ -55,7 +55,8 @@ class Settings:
default_factory=lambda: float(os.getenv("INTENT_QA_MIN_PROBABILITY", "0.5")) default_factory=lambda: float(os.getenv("INTENT_QA_MIN_PROBABILITY", "0.5"))
) )
# screen agent 单次运行的模型调用上限,包括读取 skill、查询物料、提交修正和最终回复 # screen agent 单次运行的模型调用上限,包括读取 skill、查询物料、提交修正和最终回复;
# 用尽或没有给出结论时节点再收尾 1 次,所以一次请求最多 MAX_AGENT_STEPS + 1 次
max_agent_steps: int = field( max_agent_steps: int = field(
default_factory=lambda: int(os.getenv("MAX_AGENT_STEPS", "12")) default_factory=lambda: int(os.getenv("MAX_AGENT_STEPS", "12"))
) )
@@ -101,7 +102,7 @@ class Settings:
) )
# 单个数据文件的大小上限(MB) # 单个数据文件的大小上限(MB)
data_file_max_mb: int = field( data_file_max_mb: int = field(
default_factory=lambda: int(os.getenv("DATA_FILE_MAX_MB", "50")) default_factory=lambda: int(os.getenv("DATA_FILE_MAX_MB", "500"))
) )
# 数据文件保留时长(秒),从上传时算起,默认 7 天 # 数据文件保留时长(秒),从上传时算起,默认 7 天
data_file_ttl: int = field( data_file_ttl: int = field(
+516
View File
@@ -0,0 +1,516 @@
"""单元格的类型规则:单个值用 classify 判断,一段 Arrow 数组用 classify_piece 批量判断
classify 是规则的唯一定义:一个值是什么类型、转换成什么值。大表不能逐个调用它,classify_piece 对常见写法
(空值记号、布尔、ASCII 数字和百分数、编号、ISO 日期和常见的日期时间)用正则和 Arrow 转换批量得出结论,
其余的值去重后逐个交给 classify。批量规则只认它有把握的写法,认不出的一律交给 classify,所以结论和逐个
调用 classify 相同;测试里逐值对比两者
"""
import json
import math
import re
from dataclasses import dataclass
from datetime import UTC, date, datetime, time
from typing import Any
import numpy as np
import pyarrow as pa
import pyarrow.compute as pc
NULL_TOKENS = {"", "-", "--", "—", "——", "N/A", "n/a", "NULL", "null", "#N/A"}
EXCEL_ERROR = re.compile(r"^#(?:NULL!|DIV/0!|VALUE!|REF!|NAME\?|NUM!|N/A)$")
CURRENCY_SYMBOLS = "¥¥$€"
NUMBER = re.compile(r"^[+-]?(?:\d+|\d{1,3}(?:,\d{3})+)(?:\.\d+)?(?:[eE][+-]?\d+)?$")
INTEGER = re.compile(r"^[+-]?\d+$")
# 16 位及以上的纯数字是证件号、卡号一类的编号,转成数值会丢失精度
MIN_ID_DIGITS = 16
# 按年份后的第一个字符选格式;%m、%d 也接受一位数。没有对应格式的直接判为不是日期:
# strptime 每次失败都要抛异常,一个值试完所有格式约 16 μs,还要抢一把模块级的锁,多线程时更慢
DATE_FORMATS = {
"-": ("%Y-%m-%d", "%Y-%m"),
"/": ("%Y/%m/%d", "%Y/%m"),
".": ("%Y.%m.%d",),
"年": ("%Y年%m月%d日", "%Y年%m月"),
}
DATETIME_FORMATS = {
"-": ("%Y-%m-%d %H:%M:%S", "%Y-%m-%d %H:%M"),
"/": ("%Y/%m/%d %H:%M:%S", "%Y/%m/%d %H:%M"),
}
INT64_RANGE = range(-(2**63), 2**63)
# 超过这个范围的整数存成 float 会丢精度
EXACT_FLOAT_INT = 2**53
# str.isspace() 为真的全部字符,也就是 str.strip() 去掉的字符;Arrow 自带的去空白函数按另一套定义,
# 读取时显式传入这些字符,结果才和 Python 一致
WHITESPACE = "\t\n\x0b\x0c\r\x1c\x1d\x1e\x1f \x85\xa0            

   "
# 正则里匹配上面这些字符的写法(RE2 的 \s 只认 ASCII 空白)
WHITESPACE_CLASS = "".join(f"\\x{{{ord(char):x}}}" for char in WHITESPACE)
# 类型代码,classify 返回的类型名按这个顺序编号
KINDS = ("null", "bool", "int", "float", "percent", "date", "datetime", "text")
NULL, BOOL, INT, FLOAT, PERCENT, DATE, DATETIME, TEXT = range(len(KINDS))
KIND_CODES = {name: code for code, name in enumerate(KINDS)}
MICROS_PER_DAY = 86_400_000_000
EPOCH = datetime(1970, 1, 1)
def classify(value: Any) -> tuple[str, Any]:
"""单个值的类型和转换后的值;类型是 KINDS 之一"""
if value is None:
return "null", None
if isinstance(value, bool):
return "bool", value
if isinstance(value, int):
return ("int", value) if value in INT64_RANGE else ("text", str(value))
if isinstance(value, float):
return (
("null", None)
if math.isnan(value) or math.isinf(value)
else ("float", value)
)
if isinstance(value, datetime):
# Excel 只设了日期格式的单元格读出来是零点的 datetime
if value.time() == time() and value.tzinfo is None:
return "date", value.date()
return "datetime", _naive(value)
if isinstance(value, date):
return "date", value
if not isinstance(value, str):
return "text", as_text(value)
text = value.strip()
if text in NULL_TOKENS or EXCEL_ERROR.match(text):
return "null", None
if text.lower() in ("true", "false"):
return "bool", text.lower() == "true"
number = parse_number(text)
if number is not None:
if isinstance(number, int):
# 带千分位的超长整数和超出范围的原生整数一样保留文本
return ("int", number) if number in INT64_RANGE else ("text", text)
return "float", number
if text.endswith("%"):
number = parse_number(text[:-1].rstrip())
if number is not None:
return "percent", number / 100
moment = parse_datetime(text)
if moment is not None:
return ("datetime" if isinstance(moment, datetime) else "date"), moment
return "text", text
def parse_number(text: str) -> int | float | None:
"""千分位和货币符号可以有;有前导零或 16 位以上的数字串是编号,不算数值
判断数字一律用 isdecimal:isdigit 对 ²、① 也返回 True,但 int() 解析不了它们
"""
if text and text[0] in CURRENCY_SYMBOLS:
text = text[1:].lstrip()
# 最常见的纯数字不用正则
plain = text.isdecimal()
if not plain and not NUMBER.match(text):
return None
digits = text.lstrip("+-")
if len(digits) > 1 and digits[0] == "0" and digits[1].isdecimal():
return None
if digits.isdecimal() and len(digits) >= MIN_ID_DIGITS:
return None
if plain:
return int(text)
text = text.replace(",", "")
return int(text) if INTEGER.match(text) else float(text)
def parse_datetime(text: str) -> date | datetime | None:
if len(text) < 6 or not text[:4].isdecimal():
return None
# 标准 ISO 写法先用 fromisoformat,比 strptime 快一个数量级
if ":" in text:
try:
return _naive(datetime.fromisoformat(text))
except ValueError:
pass
for pattern in DATETIME_FORMATS.get(text[4], ()):
try:
return datetime.strptime(text, pattern)
except ValueError:
continue
return None
if len(text) == 10 and text[4] == "-":
try:
return date.fromisoformat(text)
except ValueError:
pass
for pattern in DATE_FORMATS.get(text[4], ()):
try:
return datetime.strptime(text, pattern).date()
except ValueError:
continue
return None
def as_text(value: Any) -> str:
"""单元格转成文本:列名、保留为文本的值、警告里引用的内容都用它;Excel 里年份等数字读出来是 2025.0"""
if value is None:
return ""
if isinstance(value, str):
return value.strip()
if isinstance(value, bool):
return "true" if value else "false"
if isinstance(value, float) and value.is_integer():
return str(int(value))
if isinstance(value, datetime):
return value.isoformat(sep=" ")
if isinstance(value, (date, time)):
return value.isoformat()
if isinstance(value, (dict, list)):
# Parquet 的嵌套值里可能有日期、Decimal、bytes,JSON 表示不了的写成字符串
return json.dumps(value, ensure_ascii=False, default=str)
return str(value)
def _naive(value: datetime) -> datetime:
# 带时区的时间统一换成 UTC 再去掉时区,同一列里的值才能比较
if value.tzinfo is None:
return value
return value.astimezone(UTC).replace(tzinfo=None)
# ---- 批量判断 ----
# classify 只会把这些开头的字符串判成文本以外的类型(空值记号和布尔另外判断):
# 货币符号,或可选的正负号加数字(数值、百分数、日期都以数字开头)
CANDIDATE = r"^(?:[¥¥$€]|[+-]?\p{Nd})"
# 最多 15 位整数部分的 ASCII 数值,保证整数在 int64 内、百分数除以 100 与 Python 结果相同;
# 有前导零的(编号)、货币符号后带空格的等写法不在这里,交给 classify
_DIGITS = r"[¥¥$€]?[+-]?(?:0|[1-9][0-9]{0,14}|[1-9][0-9]{0,2}(?:,[0-9]{3}){1,4})(?:\.[0-9]+)?(?:[eE][+-]?[0-9]{1,3})?"
FAST_NUMBER = f"^{_DIGITS}$"
FAST_PERCENT = f"^{_DIGITS}%$"
# 有前导零或 16 位以上的纯数字是编号:classify 判成文本
FAST_ID = r"^0[0-9]+$|^[0-9]{16,}$"
FAST_DATE = r"^[1-9][0-9]{3}-[0-9]{2}-[0-9]{2}$"
# fromisoformat 能解析的写法,以及 DATETIME_FORMATS 里分隔符为 - 或 /、各部分一两位数字的写法
FAST_DATETIMES = (
r"^[1-9][0-9]{3}-[0-9]{2}-[0-9]{2}[T ][0-9]{2}:[0-9]{2}(?::[0-9]{2}(?:\.[0-9]{3}(?:[0-9]{3})?)?)?$",
r"^[1-9][0-9]{3}-[0-9]{1,2}-[0-9]{1,2} [0-9]{1,2}:[0-9]{1,2}(?::[0-9]{1,2})?$",
r"^[1-9][0-9]{3}/[0-9]{1,2}/[0-9]{1,2} [0-9]{1,2}:[0-9]{1,2}(?::[0-9]{1,2})?$",
)
NULL_SET = pa.array(sorted(NULL_TOKENS))
BOOL_SET = pa.array(["true", "false"])
NULL_TEXT = pa.scalar(None, pa.string())
@dataclass
class Classified:
"""一段单元格逐个分类的结果,和对每个值调用 classify 的结论相同
kinds 是每个值的类型代码;ints、reals、moments、flags 和这段等长,只在对应类型的位置上有意义,
这段里没有这种类型时为空
"""
piece: pa.Array
kinds: np.ndarray
# int
ints: np.ndarray | None = None
# float、percent(已除以 100)
reals: np.ndarray | None = None
# date、datetime:自 1970 年起的微秒数,日期按零点
moments: np.ndarray | None = None
# bool
flags: np.ndarray | None = None
def values(self, name: str) -> np.ndarray:
"""取某类转换后的值,没有时补零"""
array = getattr(self, name)
if array is None:
array = np.zeros(len(self.kinds), _DTYPES[name])
setattr(self, name, array)
return array
def count(self) -> np.ndarray:
return np.bincount(self.kinds, minlength=len(KINDS))
_DTYPES = {"ints": np.int64, "reals": np.float64, "moments": np.int64, "flags": bool}
def union_layout(piece: pa.UnionArray) -> tuple[np.ndarray, np.ndarray]:
"""dense union 每个位置属于第几个子数组、是子数组里的第几个
切片后的 union,pyarrow 的 type_codes、offsets 不计切片的偏移,所以直接读缓冲区
"""
start, stop = piece.offset, piece.offset + len(piece)
_, codes, offsets = piece.buffers()[:3]
return (
np.frombuffer(codes, np.int8)[start:stop],
np.frombuffer(offsets, np.int32)[start:stop],
)
def mask(values: pa.Array | pa.ChunkedArray) -> np.ndarray:
"""Arrow 布尔数组转成 numpy,空值当作 False"""
return pc.fill_null(values, False).to_numpy(zero_copy_only=False)
def classify_piece(piece: pa.Array, order: list | None = None) -> Classified:
"""批量判断一段单元格;段的类型是读取时产生的几种之一,见 grid.py
order 是字符串段依次尝试的批量规则,同一列的各段传同一个列表:这些规则认的写法互不重叠,先试哪个
结论都一样,所以每段之后把认出最多的规则挪到前面,后面的段大多一次就认完
"""
kind = piece.type
if pa.types.is_union(kind):
return _classify_union(piece, list(FAST_RULES) if order is None else order)
valid = mask(piece.is_valid())
if pa.types.is_string(kind):
return _classify_strings(
piece, valid, list(FAST_RULES) if order is None else order
)
result = Classified(piece, np.where(valid, TEXT, NULL).astype(np.int8))
if pa.types.is_null(kind):
return result
if pa.types.is_boolean(kind):
result.kinds[valid] = BOOL
result.flags = mask(piece)
elif pa.types.is_integer(kind):
result.kinds[valid] = INT
result.ints = pc.fill_null(piece, 0).to_numpy()
elif pa.types.is_floating(kind):
reals = pc.fill_null(piece, 0.0).to_numpy()
result.kinds[valid] = np.where(np.isfinite(reals[valid]), FLOAT, NULL)
result.reals = reals
elif pa.types.is_date32(kind):
result.kinds[valid] = DATE
result.moments = (
pc.fill_null(piece, 0).cast(pa.int32()).to_numpy().astype(np.int64)
* MICROS_PER_DAY
)
elif pa.types.is_timestamp(kind):
moments = pc.fill_null(piece.cast(pa.int64()), 0).to_numpy()
midnight = moments % MICROS_PER_DAY == 0
result.kinds[valid] = np.where(midnight[valid], DATE, DATETIME)
result.moments = moments
else:
raise TypeError(f"不支持的单元格类型:{kind}")
return result
def _classify_union(piece: pa.UnionArray, order: list) -> Classified:
"""类型混杂的段(见 grid.py):每个子数组各判断一次,再按位置拼回"""
result = Classified(piece, np.full(len(piece), NULL, np.int8))
codes, offsets = union_layout(piece)
for code in range(piece.type.num_fields):
at = np.flatnonzero(codes == code)
if not len(at):
continue
child = classify_piece(piece.field(code), order)
rows = offsets[at]
result.kinds[at] = child.kinds[rows]
for name in _DTYPES:
values = getattr(child, name)
if values is not None:
result.values(name)[at] = values[rows]
return result
def _classify_strings(piece: pa.Array, valid: np.ndarray, order: list) -> Classified:
"""字符串段:按 classify 的顺序(空值记号 → 布尔 → 数值 → 百分数 → 日期时间 → 文本)批量判断"""
result = Classified(piece, np.where(valid, TEXT, NULL).astype(np.int8))
# 空值记号、Excel 错误值、布尔都不以货币符号或数字开头,只在其余的值里找;
# 这些值剩下的就是文本,类型已经记好。数值、日期列整段都是候选,省掉这几步
pending = valid & mask(pc.match_substring_regex(piece, CANDIDATE))
others = valid & ~pending
if others.any():
hit = others & mask(
pc.or_(
pc.is_in(piece, value_set=NULL_SET),
pc.match_substring_regex(piece, EXCEL_ERROR.pattern),
)
)
result.kinds[hit] = NULL
others &= ~hit
if others.any():
lower = pc.utf8_lower(piece)
hit = others & mask(pc.is_in(lower, value_set=BOOL_SET))
if hit.any():
result.kinds[hit] = BOOL
result.values("flags")[hit] = mask(pc.equal(lower, "true"))[hit]
claimed = {}
for rule in order:
left = int(pending.sum())
if not left:
break
rule(piece, result, pending)
claimed[rule] = left - int(pending.sum())
order.sort(key=lambda rule: -claimed.get(rule, 0))
if pending.any():
_delegate(piece, result, pending)
return result
def _matched(piece: pa.Array, pattern: str, pending: np.ndarray) -> np.ndarray:
"""还没有结论、并且符合 pattern 的值的位置"""
return np.flatnonzero(pending & mask(pc.match_substring_regex(piece, pattern)))
def _plain_number(strings: pa.Array) -> pa.Array:
"""去掉货币符号、千分位和正号,剩下 Arrow 能直接转换的数字
只用于符合 FAST_NUMBER 的值:货币符号只在开头,正号在开头或指数里(1e+5 去掉后是 1e5,值不变)
"""
return pc.replace_substring_regex(strings, r"[¥¥$€,+]", "")
def _fast_numbers(piece: pa.Array, result: Classified, pending: np.ndarray) -> None:
rows = _matched(piece, FAST_NUMBER, pending)
if not len(rows):
return
plain = _plain_number(piece.take(rows))
integer = mask(pc.match_substring_regex(plain, r"^-?[0-9]+$"))
if integer.any():
at = rows[integer]
result.kinds[at] = INT
result.values("ints")[at] = pc.cast(
plain.filter(integer), pa.int64()
).to_numpy()
if not integer.all():
at = rows[~integer]
result.kinds[at] = FLOAT
result.values("reals")[at] = pc.cast(
plain.filter(~integer), pa.float64()
).to_numpy()
pending[rows] = False
def _fast_percents(piece: pa.Array, result: Classified, pending: np.ndarray) -> None:
rows = _matched(piece, FAST_PERCENT, pending)
if not len(rows):
return
plain = _plain_number(pc.utf8_slice_codeunits(piece.take(rows), 0, -1))
result.kinds[rows] = PERCENT
result.values("reals")[rows] = pc.cast(plain, pa.float64()).to_numpy() / 100
pending[rows] = False
def _fast_ids(piece: pa.Array, result: Classified, pending: np.ndarray) -> None:
# 类型已经记为文本,只需标记为已判断
pending[_matched(piece, FAST_ID, pending)] = False
def _fast_dates(piece: pa.Array, result: Classified, pending: np.ndarray) -> None:
rows = _matched(piece, FAST_DATE, pending)
if len(rows):
_put_moments(result, pending, rows, piece.take(rows), pa.date32(), DATE)
def _fast_datetimes(piece: pa.Array, result: Classified, pending: np.ndarray) -> None:
for pattern in FAST_DATETIMES:
rows = _matched(piece, pattern, pending)
if not len(rows):
continue
# 统一成 Arrow 认识的 ISO 写法:/ 换成 -,一位数补零(相邻的两位数各占一次,所以替换两遍)
strings = pc.replace_substring(piece.take(rows), "/", "-")
for _ in range(2):
strings = pc.replace_substring_regex(
strings, r"([-: ])([0-9])([-: ]|$)", r"\10\2\3"
)
_put_moments(result, pending, rows, strings, pa.timestamp("us"), DATETIME)
def _put_moments(
result: Classified,
pending: np.ndarray,
rows: np.ndarray,
strings: pa.Array,
arrow_type: pa.DataType,
kind: int,
) -> None:
"""转换成日期或时间并填回;转换不了的(2 月 30 日、25 点)留给 classify"""
ok = np.zeros(len(rows), bool)
moments = np.zeros(len(rows), np.int64)
_cast_parts(strings, arrow_type, 0, ok, moments)
at = rows[ok]
result.kinds[at] = kind
result.values("moments")[at] = moments[ok]
pending[at] = False
# 二分到这么短还转换失败,就整段交给 classify
MIN_CAST_PART = 64
def _cast_parts(
strings: pa.Array,
arrow_type: pa.DataType,
offset: int,
ok: np.ndarray,
moments: np.ndarray,
) -> None:
"""Arrow 的转换遇到一个坏值就整体失败;失败时对半分开重试,只把含坏值的小段留下"""
try:
values = pc.cast(strings, arrow_type)
except pa.ArrowInvalid:
if len(strings) <= MIN_CAST_PART:
return
half = len(strings) // 2
_cast_parts(strings.slice(0, half), arrow_type, offset, ok, moments)
_cast_parts(strings.slice(half), arrow_type, offset + half, ok, moments)
return
end = offset + len(strings)
ok[offset:end] = True
if pa.types.is_date32(arrow_type):
moments[offset:end] = (
values.cast(pa.int32()).to_numpy().astype(np.int64) * MICROS_PER_DAY
)
else:
moments[offset:end] = values.cast(pa.int64()).to_numpy()
# 批量规则;认的写法互不重叠,顺序不影响结论
FAST_RULES = (_fast_numbers, _fast_percents, _fast_ids, _fast_dates, _fast_datetimes)
def _delegate(piece: pa.Array, result: Classified, pending: np.ndarray) -> None:
"""批量规则认不出的值去重后逐个交给 classify,再按位置填回"""
rows = np.flatnonzero(pending)
values = piece.take(rows)
unique = pc.unique(values)
codes = pc.index_in(values, value_set=unique).to_numpy()
outcomes = [classify(value) for value in unique.to_pylist()]
kinds = np.array([KIND_CODES[kind] for kind, _ in outcomes], np.int8)
result.kinds[rows] = kinds[codes]
for name, wanted, convert in _FIELDS:
chosen = np.isin(kinds, wanted)
if not chosen.any():
continue
converted = np.array(
[
convert(value) if pick else 0
for (_, value), pick in zip(outcomes, chosen)
],
_DTYPES[name],
)
at = chosen[codes]
result.values(name)[rows[at]] = converted[codes[at]]
pending[rows] = False
def _micros(value: date) -> int:
if isinstance(value, datetime):
return (value - EPOCH) // _MICROSECOND
return (value - EPOCH.date()).days * MICROS_PER_DAY
_MICROSECOND = datetime(1970, 1, 1, microsecond=1) - EPOCH
# classify 的结果填进 Classified 的哪个字段:字段名、对应的类型代码、怎样转换
_FIELDS = (
("ints", (INT,), int),
("reals", (FLOAT, PERCENT), float),
("moments", (DATE, DATETIME), _micros),
("flags", (BOOL,), bool),
)
+125 -211
View File
@@ -1,52 +1,46 @@
"""列名清理和类型推断,输出 Arrow 列 """列名清理和类型推断,输出 Arrow 列
CSV 的单元格全是字符串;Excel 的数值、日期单元格直接采用,文本单元格再按规则推断 单元格的类型规则在 cells.py;这里按列汇总:一列的非空值必须全部符合同一类型才转换,否则保留文本,
一列的非空值必须全部符合同一类型才转换,否则保留文本:宁可让模型在 SQL 里 TRY_CAST,也不悄悄丢掉值 宁可让模型在 SQL 里 TRY_CAST,也不悄悄丢掉值。各列互不相关,用线程并行推断(Arrow 的计算不占 GIL)
""" """
import json import os
import math
import re import re
from collections import Counter
from collections.abc import Sequence from collections.abc import Sequence
from concurrent.futures import ThreadPoolExecutor
from dataclasses import dataclass from dataclasses import dataclass
from datetime import UTC, date, datetime, time
from typing import Any from typing import Any
import numpy as np
import pyarrow as pa import pyarrow as pa
import pyarrow.compute as pc
from app.datafiles.base import Reporter, brief, no_report, where from app.datafiles.base import Reporter, brief, no_report, where
from app.datafiles.cells import (
BOOL,
DATE,
DATETIME,
EXACT_FLOAT_INT,
FAST_RULES,
FLOAT,
INT,
KINDS,
MICROS_PER_DAY,
NULL,
NULL_TEXT,
PERCENT,
TEXT,
Classified,
as_text,
classify_piece,
)
from app.datafiles.grid import Column
from app.schemas.data_file import DataColumn, ParseWarning from app.schemas.data_file import DataColumn, ParseWarning
NULL_TOKENS = {"", "-", "--", "—", "——", "N/A", "n/a", "NULL", "null", "#N/A"} NUMERIC_KINDS = (INT, FLOAT, PERCENT)
EXCEL_ERROR = re.compile(r"^#(?:NULL!|DIV/0!|VALUE!|REF!|NAME\?|NUM!|N/A)$") DATE_KINDS = (DATE, DATETIME)
# 同时推断的列数
CURRENCY_SYMBOLS = "¥¥$€" INFER_THREADS = min(8, os.cpu_count() or 1)
NUMBER = re.compile(r"^[+-]?(?:\d+|\d{1,3}(?:,\d{3})+)(?:\.\d+)?(?:[eE][+-]?\d+)?$")
INTEGER = re.compile(r"^[+-]?\d+$")
# 16 位及以上的纯数字是证件号、卡号一类的编号,转成数值会丢失精度
MIN_ID_DIGITS = 16
# 按年份后的第一个字符选格式;%m、%d 也接受一位数
DATE_FORMATS = {
"-": ("%Y-%m-%d", "%Y-%m"),
"/": ("%Y/%m/%d", "%Y/%m"),
".": ("%Y.%m.%d",),
"年": ("%Y年%m月%d日", "%Y年%m月"),
}
DATETIME_FORMATS = (
"%Y-%m-%d %H:%M:%S",
"%Y-%m-%d %H:%M",
"%Y/%m/%d %H:%M:%S",
"%Y/%m/%d %H:%M",
)
INT64_RANGE = range(-(2**63), 2**63)
# 超过这个范围的整数存成 float 会丢精度
EXACT_FLOAT_INT = 2**53
NUMERIC_KINDS = ("int", "float", "percent")
DATE_KINDS = ("date", "datetime")
@dataclass @dataclass
@@ -84,19 +78,28 @@ def clean_names(
def build_columns( def build_columns(
names: list[str], names: list[str],
values: list[Sequence[Any]], columns: list[Column],
sheet: str | None, sheet: str | None,
report: Reporter = no_report, report: Reporter = no_report,
) -> tuple[TypedColumns, ParseWarning | None]: ) -> tuple[TypedColumns, ParseWarning | None]:
"""按列推断类型;有列因少数无法解析的值保留为文本时,返回一条 kept_as_text 警告""" """按列推断类型;有列因少数无法解析的值保留为文本时,返回一条 kept_as_text 警告
columns, arrays, notes = [], [], {}
for number, (name, column_values) in enumerate(zip(names, values), start=1): columns 会被清空:每列推断完就释放它的原始单元格,大表不必同时留着原始和转换后的两份
column, array, note = infer_column(name, column_values) """
report(number / len(names)) described, arrays, notes = [], [], {}
columns.append(column) with ThreadPoolExecutor(INFER_THREADS) as pool:
arrays.append(array) futures = [
if note: pool.submit(infer_column, name, column)
notes[name] = note for name, column in zip(names, columns)
]
columns.clear()
for number, (name, future) in enumerate(zip(names, futures), start=1):
column, array, note = future.result()
report(number / len(names))
described.append(column)
arrays.append(array)
if note:
notes[name] = note
warning = None warning = None
if notes: if notes:
listed = "、".join(f"「{name}」" for name in notes) listed = "、".join(f"「{name}」" for name in notes)
@@ -105,7 +108,7 @@ def build_columns(
message=f"{where(sheet)}{listed}有无法解析的值,整列按文本处理", message=f"{where(sheet)}{listed}有无法解析的值,整列按文本处理",
) )
return TypedColumns( return TypedColumns(
columns, pa.Table.from_arrays(arrays, names=names), notes described, pa.Table.from_arrays(arrays, names=names), notes
), warning ), warning
@@ -154,194 +157,105 @@ def _arrow_kind(arrow_type: pa.DataType) -> str | None:
def infer_column( def infer_column(
name: str, values: Sequence[Any] name: str, column: Column
) -> tuple[DataColumn, pa.Array, str | None]: ) -> tuple[DataColumn, pa.ChunkedArray, str | None]:
value_kinds, parsed = [], [] order = list(FAST_RULES)
for value in values: parts = [classify_piece(piece, order) for piece in column.pieces]
kind, result = classify(value) counts = sum((part.count() for part in parts), np.zeros(len(KINDS), np.int64))
value_kinds.append(kind) present = {kind for kind in range(len(KINDS)) if counts[kind] and kind != NULL}
parsed.append(result)
kinds = Counter(kind for kind in value_kinds if kind != "null")
present = set(kinds)
if present and present <= {"bool"}: def build(make, arrow_type: pa.DataType) -> pa.ChunkedArray:
return DataColumn(name=name, type="boolean"), pa.array(parsed, pa.bool_()), None return pa.chunked_array([make(part) for part in parts], arrow_type)
if present and present <= {"int", "float"}:
if all( if present and present <= {BOOL}:
value is None array = build(
or isinstance(value, int) lambda part: _pick(part, "flags", (BOOL,), pa.bool_()), pa.bool_()
or (value.is_integer() and abs(value) < EXACT_FLOAT_INT)
for value in parsed
):
ints = [None if value is None else int(value) for value in parsed]
return (
DataColumn(name=name, type="number"),
pa.array(ints, pa.int64()),
None,
)
return (
DataColumn(name=name, type="number"),
pa.array(parsed, pa.float64()),
None,
) )
if present == {"percent"}: return DataColumn(name=name, type="boolean"), array, None
column = DataColumn(name=name, type="number", format="percent") if present and present <= {INT, FLOAT}:
return column, pa.array(parsed, pa.float64()), None if all(_integral(part) for part in parts):
if present and present <= {"date"}: array = build(lambda part: _numbers(part, np.int64), pa.int64())
return DataColumn(name=name, type="date"), pa.array(parsed, pa.date32()), None else:
array = build(lambda part: _numbers(part, np.float64), pa.float64())
return DataColumn(name=name, type="number"), array, None
if present == {PERCENT}:
array = build(
lambda part: _pick(part, "reals", (PERCENT,), pa.float64()), pa.float64()
)
return DataColumn(name=name, type="number", format="percent"), array, None
if present and present <= {DATE}:
array = build(_dates, pa.date32())
return DataColumn(name=name, type="date"), array, None
if present and present <= set(DATE_KINDS): if present and present <= set(DATE_KINDS):
stamps = [ array = build(
datetime.combine(value, time()) if type(value) is date else value lambda part: _pick(part, "moments", DATE_KINDS, pa.timestamp("us")),
for value in parsed pa.timestamp("us"),
]
return (
DataColumn(name=name, type="datetime"),
pa.array(stamps, pa.timestamp("us")),
None,
) )
return DataColumn(name=name, type="datetime"), array, None
texts = [ array = build(_texts, pa.string())
None if result is None else as_text(value) return DataColumn(name=name, type="text"), array, _text_note(counts, parts)
for value, result in zip(values, parsed)
]
note = _text_note(texts, value_kinds, kinds)
return DataColumn(name=name, type="text"), pa.array(texts, pa.string()), note
def classify(value: Any) -> tuple[str, Any]: def _pick(
"""单个值的类型和转换后的值;类型是 null、bool、int、float、percent、date、datetime、text 之一""" part: Classified, field: str, kinds: tuple[int, ...], arrow_type: pa.DataType
if value is None: ) -> pa.Array:
return "null", None """这段里属于 kinds 的值,其余位置为空"""
if isinstance(value, bool): keep = np.isin(part.kinds, kinds)
return "bool", value return pa.array(part.values(field), arrow_type, mask=~keep)
if isinstance(value, int):
return ("int", value) if value in INT64_RANGE else ("text", str(value))
if isinstance(value, float):
return (
("null", None)
if math.isnan(value) or math.isinf(value)
else ("float", value)
)
if isinstance(value, datetime):
# Excel 只设了日期格式的单元格读出来是零点的 datetime
if value.time() == time() and value.tzinfo is None:
return "date", value.date()
return "datetime", _naive(value)
if isinstance(value, date):
return "date", value
if not isinstance(value, str):
return "text", as_text(value)
text = value.strip()
if text in NULL_TOKENS or EXCEL_ERROR.match(text):
return "null", None
if text.lower() in ("true", "false"):
return "bool", text.lower() == "true"
number = parse_number(text)
if number is not None:
return ("int" if isinstance(number, int) else "float"), number
if text.endswith("%"):
number = parse_number(text[:-1].rstrip())
if number is not None:
return "percent", number / 100
moment = parse_datetime(text)
if moment is not None:
return ("datetime" if isinstance(moment, datetime) else "date"), moment
return "text", text
def parse_number(text: str) -> int | float | None: def _integral(part: Classified) -> bool:
"""千分位和货币符号可以有;有前导零或 16 位以上的数字串是编号,不算数值 """小数都是整数值并且在 float 能精确表示的范围内(Excel 的整数读出来是 100.0)"""
reals = part.values("reals")[part.kinds == FLOAT]
判断数字一律用 isdecimal:isdigit 对 ²、① 也返回 True,但 int() 解析不了它们 return bool(np.all((np.floor(reals) == reals) & (np.abs(reals) < EXACT_FLOAT_INT)))
"""
if text and text[0] in CURRENCY_SYMBOLS:
text = text[1:].lstrip()
# 最常见的纯数字不用正则
plain = text.isdecimal()
if not plain and not NUMBER.match(text):
return None
digits = text.lstrip("+-")
if len(digits) > 1 and digits[0] == "0" and digits[1].isdecimal():
return None
if digits.isdecimal() and len(digits) >= MIN_ID_DIGITS:
return None
if plain:
return int(text)
text = text.replace(",", "")
return int(text) if INTEGER.match(text) else float(text)
def parse_datetime(text: str) -> date | datetime | None: def _numbers(part: Classified, dtype) -> pa.Array:
if len(text) < 6 or not text[:4].isdecimal(): is_int, is_float = part.kinds == INT, part.kinds == FLOAT
return None # 其他位置可能是 NaN(float 段里的空值),先清零再转换
# 标准 ISO 写法先用 fromisoformat,比 strptime 快一个数量级 reals = np.where(is_float, part.values("reals"), 0.0)
if ":" in text: values = np.where(is_int, part.values("ints").astype(dtype), reals.astype(dtype))
try: keep = is_int | is_float
return _naive(datetime.fromisoformat(text)) return pa.array(values, pa.from_numpy_dtype(dtype), mask=~keep)
except ValueError:
pass
for pattern in DATETIME_FORMATS:
try:
return datetime.strptime(text, pattern)
except ValueError:
continue
return None
if len(text) == 10 and text[4] == "-":
try:
return date.fromisoformat(text)
except ValueError:
pass
for pattern in DATE_FORMATS.get(text[4], ()):
try:
return datetime.strptime(text, pattern).date()
except ValueError:
continue
return None
def as_text(value: Any) -> str: def _dates(part: Classified) -> pa.Array:
"""单元格转成文本:列名、保留为文本的值、警告里引用的内容都用它;Excel 里年份等数字读出来是 2025.0""" days = (part.values("moments") // MICROS_PER_DAY).astype(np.int32)
if value is None: return pa.array(days, pa.date32(), mask=part.kinds != DATE)
return ""
if isinstance(value, str):
return value.strip()
if isinstance(value, bool):
return "true" if value else "false"
if isinstance(value, float) and value.is_integer():
return str(int(value))
if isinstance(value, datetime):
return value.isoformat(sep=" ")
if isinstance(value, (date, time)):
return value.isoformat()
if isinstance(value, (dict, list)):
# Parquet 的嵌套值里可能有日期、Decimal、bytes,JSON 表示不了的写成字符串
return json.dumps(value, ensure_ascii=False, default=str)
return str(value)
def _naive(value: datetime) -> datetime: def _texts(part: Classified) -> pa.Array:
# 带时区的时间统一换成 UTC 再去掉时区,同一列里的值才能比较 """保留为文本:字符串是去掉空白后的原文,其他类型按 as_text 转换;空值记号仍为空"""
if value.tzinfo is None: keep = part.kinds != NULL
return value if pa.types.is_string(part.piece.type):
return value.astimezone(UTC).replace(tzinfo=None) return pc.if_else(pa.array(keep), part.piece, NULL_TEXT)
return pa.array(
[
as_text(value) if kept else None
for value, kept in zip(part.piece.to_pylist(), keep)
],
pa.string(),
)
def _text_note( def _text_note(counts: np.ndarray, parts: list[Classified]) -> str | None:
texts: list[str | None], value_kinds: list[str], kinds: Counter[str]
) -> str | None:
"""大部分值是数值或日期、少数无法解析时,说明原因,模型可以在 SQL 里用 TRY_CAST""" """大部分值是数值或日期、少数无法解析时,说明原因,模型可以在 SQL 里用 TRY_CAST"""
failed = kinds["text"] failed = int(counts[TEXT])
numeric = sum(kinds[kind] for kind in NUMERIC_KINDS) numeric = int(sum(counts[kind] for kind in NUMERIC_KINDS))
dates = sum(kinds[kind] for kind in DATE_KINDS) dates = int(sum(counts[kind] for kind in DATE_KINDS))
if not failed or max(numeric, dates) < failed: if not failed or max(numeric, dates) < failed:
return None return None
examples: list[str] = [] examples: list[str] = []
for text, kind in zip(texts, value_kinds): for part in parts:
if kind == "text" and text not in examples: failures = _texts(part).filter(pa.array(part.kinds == TEXT))
examples.append(text) for text in pc.unique(failures).to_pylist():
if text not in examples:
examples.append(text)
if len(examples) == 2: if len(examples) == 2:
break break
if len(examples) == 2:
break
shown = "、".join(f"「{example}」" for example in examples) shown = "、".join(f"「{example}」" for example in examples)
label = "数值" if numeric >= dates else "日期" label = "数值" if numeric >= dates else "日期"
return f"看起来是{label},有 {failed} 个值无法解析,例如{shown}" return f"看起来是{label},有 {failed} 个值无法解析,例如{shown}"
+122 -34
View File
@@ -1,4 +1,7 @@
"""按文件内容识别格式和文本编码,不看扩展名""" """按文件内容识别格式和文本编码,不看扩展名
文本文件整个解码一遍确认编码,但只留开头一段用来判断格式和分隔符;读取时再按确认的编码边读边解码
"""
import codecs import codecs
import csv import csv
@@ -15,11 +18,16 @@ from app.schemas.data_file import DataFormat
# xlsx/xlsb 解压后的总大小上限,防止解压炸弹 # xlsx/xlsb 解压后的总大小上限,防止解压炸弹
MAX_UNZIPPED_BYTES = 1024 * 1024 * 1024 MAX_UNZIPPED_BYTES = 1024 * 1024 * 1024
# 带 BOM 的文本直接按 BOM 解码;UTF-32 的 BOM 以 UTF-16 的开头,解出来含 NUL,按二进制拒绝 # 整体是一个 JSON 值的文件只能整个读进内存解析,峰值内存约为文件大小的 6~7 倍;
# 这个上限让它和最大的 CSV 一样峰值约 1 GB。JSONL 边读边解析,不受这个限制
MAX_JSON_BYTES = 150 * 1024 * 1024
# 带 BOM 的文本直接按 BOM 解码:编码名、跳过 BOM 后的解码方式、读取时的编码(会自己去掉 BOM);
# UTF-32 的 BOM 以 UTF-16 的开头,解出来含 NUL,按二进制拒绝
BOMS = ( BOMS = (
(codecs.BOM_UTF8, "utf-8", "utf-8"), (codecs.BOM_UTF8, "utf-8", "utf-8", "utf-8-sig"),
(codecs.BOM_UTF16_LE, "utf-16-le", "utf-16"), (codecs.BOM_UTF16_LE, "utf-16", "utf-16-le", "utf-16"),
(codecs.BOM_UTF16_BE, "utf-16-be", "utf-16"), (codecs.BOM_UTF16_BE, "utf-16", "utf-16-be", "utf-16"),
) )
# 没有 BOM 时依次尝试;GB18030 兼容 GBK、GB2312,中文 Windows 上 Excel 另存的 CSV 多是这种 # 没有 BOM 时依次尝试;GB18030 兼容 GBK、GB2312,中文 Windows 上 Excel 另存的 CSV 多是这种
@@ -30,23 +38,42 @@ DELIMITERS = (",", "\t", ";", "|")
# 判断分隔符只看开头这么多行 # 判断分隔符只看开头这么多行
SNIFF_LINES = 50 SNIFF_LINES = 50
# 判断格式和分隔符用的开头文本长度;每次解码读入的字节数
HEAD_CHARS = 1024 * 1024
READ_BYTES = 1024 * 1024
# 开头像 JSON 才按 JSON 处理:对象的 { 后面只能是 " 或 },数组的 [ 后面是一个值的开头; # 开头像 JSON 才按 JSON 处理:对象的 { 后面只能是 " 或 },数组的 [ 后面是一个值的开头;
# CSV 的列名也可能以 [ 开头,例如 [日期],[销售额] # CSV 的列名也可能以 [ 开头,例如 [日期],[销售额]
JSON_START = re.compile(r'\{\s*["}]|\[\s*(?:[\[{"\]\d-]|(?:true|false|null)\s*[,\]])') JSON_START = re.compile(r'\{\s*["}]|\[\s*(?:[\[{"\]\d-]|(?:true|false|null)\s*[,\]])')
UNKNOWN_ENCODING = "无法识别文件编码,请另存为 UTF-8 后重新上传"
@dataclass(frozen=True) @dataclass(frozen=True)
class Detected: class Detected:
format: DataFormat format: DataFormat
# 文本格式的编码名(返回给调用方)和读取时传给 open 的编码
encoding: str | None = None encoding: str | None = None
# 文本格式解码后的内容 codec: str | None = None
text: str | None = None # 文本格式开头的一段
head: str | None = None
# CSV/TSV 的分隔符 # CSV/TSV 的分隔符
delimiter: str | None = None delimiter: str | None = None
# JSON 已解析的内容,避免读取时再解析一遍 # JSON 已解析的内容,避免读取时再解析一遍
value: Any = None value: Any = None
@dataclass(frozen=True)
class _Text:
encoding: str
codec: str
head: str
# 开头一段之后还有内容
truncated: bool
# 整个文件只有空白
blank: bool
def detect(path: Path) -> Detected: def detect(path: Path) -> Detected:
with path.open("rb") as file: with path.open("rb") as file:
head = file.read(8) head = file.read(8)
@@ -57,14 +84,19 @@ def detect(path: Path) -> Detected:
if head.startswith(b"PAR1"): if head.startswith(b"PAR1"):
return Detected("parquet") return Detected("parquet")
text, encoding = decode_text(path.read_bytes()) text = scan_text(path)
stripped = text.lstrip() if text.blank:
if not stripped:
raise DataFileError(422, "文件是空的") raise DataFileError(422, "文件是空的")
if JSON_START.match(stripped): if JSON_START.match(text.head.lstrip()):
return _json_format(text, encoding) return _json_format(path, text)
delimiter = sniff_delimiter(text) delimiter = sniff_delimiter(text.head)
return Detected("tsv" if delimiter == "\t" else "csv", encoding, text, delimiter) return Detected(
"tsv" if delimiter == "\t" else "csv",
text.encoding,
text.codec,
text.head,
delimiter,
)
def _zip_format(path: Path) -> DataFormat: def _zip_format(path: Path) -> DataFormat:
@@ -83,41 +115,76 @@ def _zip_format(path: Path) -> DataFormat:
raise DataFileError(415, UNSUPPORTED_FORMAT) raise DataFileError(415, UNSUPPORTED_FORMAT)
def decode_text(data: bytes) -> tuple[str, str]: def scan_text(path: Path) -> _Text:
"""返回解码后的文本和编码名""" """确认文本编码:BOM 优先,否则依次严格尝试 FALLBACK_ENCODINGS"""
for bom, codec, name in BOMS: with path.open("rb") as file:
if data.startswith(bom): start = file.read(4)
for bom, name, decoding, codec in BOMS:
if start.startswith(bom):
try: try:
text = data[len(bom) :].decode(codec) return _decode(path, name, decoding, codec, len(bom))
except UnicodeDecodeError: except UnicodeDecodeError:
raise DataFileError( raise DataFileError(422, UNKNOWN_ENCODING) from None
422, "无法识别文件编码,请另存为 UTF-8 后重新上传"
) from None
return _reject_binary(text), name
for encoding in FALLBACK_ENCODINGS: for encoding in FALLBACK_ENCODINGS:
try: try:
return _reject_binary(data.decode(encoding)), encoding return _decode(path, encoding, encoding, encoding, 0)
except UnicodeDecodeError: except UnicodeDecodeError:
continue continue
raise DataFileError(422, "无法识别文件编码,请另存为 UTF-8 后重新上传") raise DataFileError(422, UNKNOWN_ENCODING)
def _reject_binary(text: str) -> str: def _decode(path: Path, name: str, decoding: str, codec: str, skip: int) -> _Text:
# 文本文件不会有 NUL;图片、PDF 等二进制文件偶尔能按 GB18030 解码成功,在这里拦下 """严格解码整个文件,留下开头一段;含 NUL 时按二进制文件拒绝"""
if "\x00" in text: decoder = codecs.getincrementaldecoder(decoding)("strict")
raise DataFileError(415, UNSUPPORTED_FORMAT) parts: list[str] = []
return text kept, blank, truncated = 0, True, False
with path.open("rb") as file:
file.seek(skip)
while True:
chunk = file.read(READ_BYTES)
text = decoder.decode(chunk, final=not chunk)
# 文本文件不会有 NUL;图片、PDF 等二进制文件偶尔能按 GB18030 解码成功,在这里拦下
if "\x00" in text:
raise DataFileError(415, UNSUPPORTED_FORMAT)
blank = blank and not text.strip()
if kept < HEAD_CHARS:
parts.append(text)
kept += len(text)
elif text:
truncated = True
if not chunk:
break
head = "".join(parts)
if len(head) > HEAD_CHARS:
head, truncated = head[:HEAD_CHARS], True
cut = head.rfind("\n")
if truncated and cut >= 0:
# 截断处可能在一行中间,丢掉最后半行
head = head[: cut + 1]
return _Text(name, codec, head, truncated, blank)
def _json_format(text: str, encoding: str) -> Detected: def _json_format(path: Path, text: _Text) -> Detected:
if _first_line_is_record(text):
return Detected("jsonl", text.encoding, text.codec, text.head)
# 整体是一个 JSON 值,只能整个读进来解析
if path.stat().st_size > MAX_JSON_BYTES:
raise DataFileError(
413,
f"JSON 文件不能超过 {MAX_JSON_BYTES // 1024 // 1024} MB,"
"请转成 JSONL(每行一个 JSON 对象)后上传",
)
content = path.read_bytes().decode(text.codec)
try: try:
return Detected("json", encoding, text, value=json.loads(text)) return Detected(
"json", text.encoding, text.codec, text.head, value=json.loads(content)
)
except ValueError as error: except ValueError as error:
json_error = error json_error = error
first = next(line for line in text.splitlines() if line.strip()) first = next(line for line in content.splitlines() if line.strip())
try: try:
if isinstance(json.loads(first), dict): if isinstance(json.loads(first), dict):
return Detected("jsonl", encoding, text) return Detected("jsonl", text.encoding, text.codec, text.head)
except ValueError: except ValueError:
pass pass
# 开头像 JSON 却不是合法 JSON,按 JSON 报错比当成 CSV 解析出一堆乱码更有用 # 开头像 JSON 却不是合法 JSON,按 JSON 报错比当成 CSV 解析出一堆乱码更有用
@@ -128,6 +195,27 @@ def _json_format(text: str, encoding: str) -> Detected:
) )
def _first_line_is_record(text: _Text) -> bool:
"""第一行是完整的 JSON 对象、后面还有内容:整体不可能是一个 JSON 值,是 JSONL
不用整个解析就能判断,大 JSONL 文件不必整个读进内存
"""
lines = text.head.splitlines(keepends=True)
index = next((i for i, line in enumerate(lines) if line.strip()), None)
if index is None:
return False
first, rest = lines[index], lines[index + 1 :]
if first == first.rstrip("\r\n") and (rest or text.truncated):
# 第一行在开头一段里没有结束
return False
if not (text.truncated or any(line.strip() for line in rest)):
return False
try:
return isinstance(json.loads(first), dict)
except ValueError:
return False
def sniff_delimiter(text: str) -> str: def sniff_delimiter(text: str) -> str:
"""选出让开头各行列数最一致、列数最多的分隔符;都只有一列时按逗号处理""" """选出让开头各行列数最一致、列数最多的分隔符;都只有一列时按逗号处理"""
lines = [] lines = []
+306
View File
@@ -0,0 +1,306 @@
"""按列存放的单元格,读取、拆表和类型推断共用
每列是若干段 Arrow 数组,每段一种类型:字符串(已去掉首尾空白,空串记为空值)、布尔、int64、float64、
date32、timestamp[us],或全是空值。读取时每 CHUNK_ROWS 行转一次,同一时间只有一个块是 Python 对象;
块里同一列的类型混杂时(表头和数值在同一列、数值列里夹着"待定"),这一段存成 Arrow 的 dense union:
每种类型一个子数组,原值和原类型都保留。Arrow 的 is_valid 等函数不认 union 的子数组里的空值,
所以按段处理的地方都经过 each,对子数组分别计算再按位置拼回
"""
from collections.abc import Iterable, Sequence
from dataclasses import dataclass
from datetime import date, datetime, time
from itertools import zip_longest
from typing import Any
import numpy as np
import pyarrow as pa
import pyarrow.compute as pc
from app.datafiles.cells import (
EXACT_FLOAT_INT,
INT64_RANGE,
NULL_TEXT,
WHITESPACE,
as_text,
mask,
union_layout,
)
CHUNK_ROWS = 65_536
NoneType = type(None)
class Column:
def __init__(self, pieces: Sequence[pa.Array]):
self.pieces = [piece for piece in pieces if len(piece)]
self.offsets = np.cumsum([0, *map(len, self.pieces)])
@classmethod
def from_values(cls, values: Sequence[Any]) -> "Column":
return cls([_convert(values)])
def __len__(self) -> int:
return int(self.offsets[-1])
def present(self) -> np.ndarray:
"""每个单元格是否有值"""
return _join(
[
each(piece, lambda part: mask(part.is_valid()), bool)
for piece in self.pieces
],
bool,
)
def matches(self, test) -> np.ndarray:
"""字符串单元格中 test(Arrow 字符串数组 → 布尔数组)为真的位置;其他类型的单元格都不算"""
def check(part: pa.Array) -> np.ndarray:
if pa.types.is_string(part.type):
return mask(test(part))
return np.zeros(len(part), bool)
return _join([each(piece, check, bool) for piece in self.pieces], bool)
def slice(self, start: int, stop: int) -> "Column":
pieces = []
for piece, offset in zip(self.pieces, self.offsets):
begin, end = max(start - offset, 0), min(stop - offset, len(piece))
if begin < end:
pieces.append(piece.slice(begin, end - begin))
return Column(pieces)
def take(self, rows: np.ndarray) -> "Column":
"""按行号(升序)取出若干行;行号连续时不复制数据"""
if len(rows) and rows[-1] - rows[0] + 1 == len(rows):
return self.slice(int(rows[0]), int(rows[-1]) + 1)
bounds = np.searchsorted(rows, self.offsets)
return Column(
[
piece.take(pa.array(rows[bounds[i] : bounds[i + 1]] - self.offsets[i]))
for i, piece in enumerate(self.pieces)
]
)
def to_pylist(self) -> list[Any]:
return [value for piece in self.pieces for value in piece.to_pylist()]
@dataclass
class Grid:
"""不带表头的原始网格;第 i 行是文件或 sheet 中的第 i + 1 行,各列一样长"""
sheet: str | None
columns: list[Column]
height: int
@classmethod
def from_rows(cls, sheet: str | None, rows: Iterable[Sequence[Any]]) -> "Grid":
builder = GridBuilder()
for row in rows:
builder.add(row)
return builder.finish(sheet)
def rows(
self, start: int, stop: int, columns: Sequence[int] | None = None
) -> list[list[Any]]:
"""[start, stop) 行的 Python 值,只取 columns 这些列"""
return self._rows(
[column.slice(start, stop) for column in self._pick(columns)], stop - start
)
def rows_at(
self, rows: np.ndarray, columns: Sequence[int] | None = None
) -> list[list[Any]]:
return self._rows(
[column.take(rows) for column in self._pick(columns)], len(rows)
)
def _pick(self, columns: Sequence[int] | None) -> list[Column]:
return self.columns if columns is None else [self.columns[c] for c in columns]
@staticmethod
def _rows(columns: list[Column], count: int) -> list[list[Any]]:
if not columns:
return [[] for _ in range(count)]
return [list(row) for row in zip(*(column.to_pylist() for column in columns))]
class GridBuilder:
"""逐行接收单元格,每 CHUNK_ROWS 行转成列;行长短不一时补空值
strings 为真表示单元格全是字符串(CSV),跳过类型检查;已经按列读好的字符串(Arrow 读 CSV)用 add_strings
"""
def __init__(self, strings: bool = False):
self._strings = strings
self._columns: list[list[pa.Array]] = []
self._pending: list[Sequence[Any]] = []
self._done = 0
@property
def height(self) -> int:
return self._done + len(self._pending)
def add(self, row: Sequence[Any]) -> None:
self._pending.append(row)
if len(self._pending) >= CHUNK_ROWS:
self._flush()
def add_rows(self, rows: Sequence[Sequence[Any]]) -> None:
self._pending.extend(rows)
if len(self._pending) >= CHUNK_ROWS:
self._flush()
def add_strings(self, columns: Sequence[pa.Array]) -> None:
"""一批按列存放的字符串单元格,各列等长"""
self._flush()
self._append([[_normalize(column)] for column in columns], len(columns[0]))
def finish(self, sheet: str | None) -> Grid:
self._flush()
return Grid(sheet, [Column(pieces) for pieces in self._columns], self._done)
def _flush(self) -> None:
rows, count = self._pending, len(self._pending)
if not count:
return
width = max(map(len, rows))
self._append(
[
[
_normalize(pa.array(values, pa.string()))
if self._strings
else _convert(values)
]
for values in (zip_longest(*rows) if width else ())
],
count,
)
self._pending = []
def _append(self, columns: list[list[pa.Array]], count: int) -> None:
for position, pieces in enumerate(columns):
if position == len(self._columns):
# 这一列在前面的块里没有出现过
self._columns.append([pa.nulls(self._done)])
self._columns[position].extend(pieces)
for pieces in self._columns[len(columns) :]:
pieces.append(pa.nulls(count))
self._done += count
def each(piece: pa.Array, compute, dtype) -> np.ndarray:
"""对一段计算 numpy 结果;union 段对每个子数组分别计算,再按位置拼回"""
if not pa.types.is_union(piece.type):
return compute(piece)
codes, offsets = union_layout(piece)
result = np.zeros(len(piece), dtype)
for code in range(piece.type.num_fields):
at = np.flatnonzero(codes == code)
if len(at):
result[at] = compute(piece.field(code))[offsets[at]]
return result
def _join(parts: list[np.ndarray], dtype) -> np.ndarray:
return np.concatenate(parts) if parts else np.zeros(0, dtype)
def _normalize(piece: pa.Array) -> pa.Array:
"""字符串去掉首尾空白(和 str.strip 相同),空串记为空值"""
if not pa.types.is_string(piece.type):
return piece
piece = pc.utf8_trim(piece, WHITESPACE)
return pc.if_else(pc.equal(piece, ""), NULL_TEXT, piece)
def _convert(values: Sequence[Any]) -> pa.Array:
"""一段单元格转成 Arrow:类型一致时转成一个数组,混杂时转成 dense union"""
types = set(map(type, values))
types.discard(NoneType)
if str in types and len(types) > 1:
# 空白单元格(Excel 读出来是 "")不算一种类型
values = [
None if type(v) is str and (not v or v.isspace()) else v for v in values
]
types = set(map(type, values))
types.discard(NoneType)
piece = _uniform(values, types)
if piece is not None:
return _normalize(piece)
families: dict[type, int] = {}
children: list[list[Any]] = []
codes = np.empty(len(values), np.int8)
offsets = np.empty(len(values), np.int32)
for index, value in enumerate(values):
family = _family(value)
code = families.get(family)
if code is None:
code = families[family] = len(children)
children.append([])
codes[index] = code
offsets[index] = len(children[code])
children[code].append(value)
return pa.UnionArray.from_dense(
pa.array(codes),
pa.array(offsets),
[_normalize(_child(child)) for child in children],
)
def _uniform(values: Sequence[Any], types: set[type]) -> pa.Array | None:
"""类型一致(或能无损地合成一种)时转成一个 Arrow 数组,否则返回空"""
if not types:
return pa.nulls(len(values))
if types == {str}:
return pa.array(values, pa.string())
if types == {bool}:
return pa.array(values, pa.bool_())
if types <= {int, float}:
if float not in types:
try:
return pa.array(values, pa.int64())
except (OverflowError, pa.ArrowException):
return None
# 整数和小数混在一起时存成 float64;整数超出 float 能精确表示的范围就不合并
if all(
-EXACT_FLOAT_INT < v < EXACT_FLOAT_INT for v in values if type(v) is int
):
return pa.array(values, pa.float64())
return None
if types <= {date, datetime}:
if any(type(v) is datetime and v.tzinfo is not None for v in values):
return None
if types == {date}:
return pa.array(values, pa.date32())
# 日期按零点存,类型推断时零点的时间仍算日期,和 classify 一致
return pa.array(
[datetime.combine(v, time()) if type(v) is date else v for v in values],
pa.timestamp("us"),
)
return None
def _family(value: Any) -> type:
"""union 里归为同一个子数组的类型:日期和不带时区的时间合在一起,超出 int64 的整数等存成文本,其余按 Python 类型分开"""
kind = type(value)
if kind is date or (kind is datetime and value.tzinfo is None):
return date
if kind is int and value not in INT64_RANGE:
return object
if kind in (NoneType, str, bool, int, float):
return kind
return object
def _child(values: list[Any]) -> pa.Array:
types = set(map(type, values))
types.discard(NoneType)
piece = _uniform(values, types)
if piece is not None:
return piece
# 超出 int64 的整数、带时区的时间、时长等:存成文本,classify 对它们的结论与文本相同
return pa.array([None if v is None else as_text(v) for v in values], pa.string())
+46 -22
View File
@@ -3,6 +3,7 @@
只读写本地文件,不碰数据库;上传流程在子进程里调用(worker.py),超时直接终止子进程 只读写本地文件,不碰数据库;上传流程在子进程里调用(worker.py),超时直接终止子进程
""" """
import gc
from collections.abc import Callable from collections.abc import Callable
from pathlib import Path from pathlib import Path
from typing import Any, Literal from typing import Any, Literal
@@ -18,9 +19,9 @@ from app.datafiles.columns import (
clean_names, clean_names,
) )
from app.datafiles.detect import detect from app.datafiles.detect import detect
from app.datafiles.grid import Grid
from app.datafiles.profile import profile_table from app.datafiles.profile import profile_table
from app.datafiles.readers import ( from app.datafiles.readers import (
Grid,
Records, Records,
read_delimited, read_delimited,
read_excel, read_excel,
@@ -37,23 +38,28 @@ ParseStage = Literal["detecting", "reading", "building"]
Progress = Callable[[ParseStage, float], None] Progress = Callable[[ParseStage, float], None]
# 各格式识别、读取两个环节结束时的整体进度,剩下的是建表(拆表、推断类型、写 Parquet、算列概况) # 各格式识别、读取两个环节结束时的整体进度,剩下的是建表(拆表、推断类型、写 Parquet、算列概况)
# 数字来自 scripts/measure_parse_progress.py 的实测:各环节真实耗时和进度大致成正比,预计时间才准 # 数字来自 scripts/measure_parse_progress.py 在 100 万行上的实测:各环节真实耗时和进度大致成正比,预计时间才准
STAGE_ENDS: dict[str, tuple[float, float]] = { STAGE_ENDS: dict[str, tuple[float, float]] = {
"csv": (0.01, 0.08), "csv": (0.06, 0.28),
"tsv": (0.01, 0.08), "tsv": (0.06, 0.28),
"xlsx": (0.01, 0.30), "xlsx": (0.01, 0.84),
"xls": (0.01, 0.30), "xls": (0.01, 0.84),
"xlsb": (0.01, 0.30), "xlsb": (0.01, 0.84),
"json": (0.10, 0.35), "json": (0.27, 0.84),
"jsonl": (0.03, 0.47), "jsonl": (0.03, 0.90),
"parquet": (0.02, 0.25), "parquet": (0.01, 0.75),
} }
# 建表各步骤完成时占其所在范围的比例:拆表(只有网格要)占网格的这么多, # 建表各步骤完成时占其所在范围的比例:拆表(只有网格要)占网格的这么多,
# 然后每张表推断完列类型、写完 Parquet,剩下的是算列概况 # 然后每张表推断完列类型、写完 Parquet,剩下的是算列概况
SPLIT_AT = 0.35 SPLIT_AT = 0.2
TYPED_AT = 0.8 TYPED_AT = 0.4
WRITTEN_AT = 0.85 WRITTEN_AT = 0.8
# Parquet 每个行组的行数。DuckDB 按行组并行扫描,pyarrow 默认一个行组放 100 多万行,整张表只能单线程读;
# 行组越小并行度越高,但字典编码按行组重复,文件变大、写得也慢,这个大小两头折中
ROW_GROUP_ROWS = 128 * 1024
def no_progress(stage: ParseStage, value: float) -> None: def no_progress(stage: ParseStage, value: float) -> None:
@@ -89,6 +95,18 @@ def parse_file(
progress 在识别、读取、建表过程中被调用,用来向调用方报告进度 progress 在识别、读取、建表过程中被调用,用来向调用方报告进度
""" """
# 读取时成批创建短命的行列表,分代垃圾回收反复扫描它们却回收不到什么,关掉后读取快约三成;
# 解析不产生循环引用,结束后照常回收
enabled = gc.isenabled()
gc.disable()
try:
return _parse_file(path, out_dir, progress)
finally:
if enabled:
gc.enable()
def _parse_file(path: Path, out_dir: Path, progress: Progress) -> ParsedFile:
progress("detecting", 0.0) progress("detecting", 0.0)
detected = detect(path) detected = detect(path)
detect_end, read_end = STAGE_ENDS[detected.format] detect_end, read_end = STAGE_ENDS[detected.format]
@@ -111,7 +129,7 @@ def parse_file(
report: Reporter, report: Reporter,
) -> None: ) -> None:
target = out_dir / f"{len(tables) + 1}.parquet" target = out_dir / f"{len(tables) + 1}.parquet"
pq.write_table(typed.table, target) pq.write_table(typed.table, target, row_group_size=ROW_GROUP_ROWS)
report(WRITTEN_AT) report(WRITTEN_AT)
profile, preview = profile_table(target, typed.columns, typed.notes) profile, preview = profile_table(target, typed.columns, typed.notes)
tables.append( tables.append(
@@ -130,6 +148,8 @@ def parse_file(
def add_grid(grid: Grid, report: Reporter) -> None: def add_grid(grid: Grid, report: Reporter) -> None:
found, dropped = split_tables(grid, scaled(report, 0.0, SPLIT_AT)) found, dropped = split_tables(grid, scaled(report, 0.0, SPLIT_AT))
# 拆出的表只引用用到的行;build_columns 推断完一列就释放一列,网格不能再留着
del grid
warnings.extend(dropped) warnings.extend(dropped)
tables = scaled(report, SPLIT_AT, 1.0) tables = scaled(report, SPLIT_AT, 1.0)
for number, table in enumerate(found): for number, table in enumerate(found):
@@ -138,7 +158,7 @@ def parse_file(
names, renamed = clean_names(table.header, table.sheet) names, renamed = clean_names(table.header, table.sheet)
typed, kept = build_columns( typed, kept = build_columns(
names, names,
list(zip(*table.rows)), table.columns,
table.sheet, table.sheet,
scaled(share, 0.0, TYPED_AT), scaled(share, 0.0, TYPED_AT),
) )
@@ -163,18 +183,22 @@ def parse_file(
match detected.format: match detected.format:
case "csv" | "tsv": case "csv" | "tsv":
grid = read_delimited(detected.text, detected.delimiter, reading) # 不留局部变量,网格才能在建表时用完即释放
add_grid(grid, building) add_grid(
read_delimited(path, detected.codec, detected.delimiter, reading),
building,
)
case "xlsx" | "xls" | "xlsb": case "xlsx" | "xls" | "xlsb":
grids, sheet_warnings = read_excel(path, reading) grids, sheet_warnings = read_excel(path, reading)
warnings.extend(sheet_warnings) warnings.extend(sheet_warnings)
# 各 sheet 按行数分摊建表的进度 # 各 sheet 按行数分摊建表的进度
total_rows = sum(len(grid.rows) for grid in grids) total_rows = sum(grid.height for grid in grids)
done_rows = 0 done_rows = 0
for grid in grids: while grids:
start = done_rows / total_rows start = done_rows / total_rows
done_rows += len(grid.rows) done_rows += grids[0].height
add_grid(grid, scaled(building, start, done_rows / total_rows)) # 取出后不再引用,这个 sheet 的网格才能在建表时用完即释放
add_grid(grids.pop(0), scaled(building, start, done_rows / total_rows))
case "json": case "json":
result = read_json(detected.value, reading) result = read_json(detected.value, reading)
if isinstance(result, Grid): if isinstance(result, Grid):
@@ -182,7 +206,7 @@ def parse_file(
else: else:
add_records(result, building) add_records(result, building)
case "jsonl": case "jsonl":
add_records(read_jsonl(detected.text, reading), building) add_records(read_jsonl(path, detected.codec, reading), building)
case "parquet": case "parquet":
typed, renamed = arrow_columns(read_parquet(path)) typed, renamed = arrow_columns(read_parquet(path))
if typed.table.num_rows: if typed.table.num_rows:
+55 -20
View File
@@ -1,12 +1,13 @@
"""列概况和预览:上传时在 Parquet 上用 DuckDB 计算,存进文件记录,供 describe_data、上传后的分析和接口返回""" """列概况和预览:上传时在 Parquet 上用 DuckDB 计算,存进文件记录,供 describe_data、上传后的分析和接口返回"""
import math import math
import os
from pathlib import Path from pathlib import Path
from typing import Any from typing import Any
import duckdb import duckdb
from app.datafiles.columns import EXACT_FLOAT_INT from app.datafiles.cells import EXACT_FLOAT_INT
from app.datafiles.query import connect, json_value, quote_identifier from app.datafiles.query import connect, json_value, quote_identifier
from app.schemas.data_file import ColumnProfile, DataColumn, TableProfile, TopValue from app.schemas.data_file import ColumnProfile, DataColumn, TableProfile, TopValue
@@ -20,28 +21,30 @@ TOP_VALUES_MAX_DISTINCT = 50
EXAMPLE_SCAN_ROWS = 1000 EXAMPLE_SCAN_ROWS = 1000
PROFILE_CELL_CHARS = 50 PROFILE_CELL_CHARS = 50
PREVIEW_CELL_CHARS = 200 PREVIEW_CELL_CHARS = 200
# 算列概况的线程数:在解析子进程里、只读自己刚写的文件,不受模型查询的线程限制;
# Parquet 按行组并行扫描,见 parse.py 的 ROW_GROUP_ROWS
PROFILE_THREADS = min(8, os.cpu_count() or 1)
def profile_table( def profile_table(
path: Path, columns: list[DataColumn], notes: dict[str, str] path: Path, columns: list[DataColumn], notes: dict[str, str]
) -> tuple[TableProfile, list[list[Any]]]: ) -> tuple[TableProfile, list[list[Any]]]:
"""返回列概况和接口用的预览(前 PREVIEW_ROWS 行)""" """返回列概况和接口用的预览(前 PREVIEW_ROWS 行)"""
connection = connect({"t": path}) connection = connect({"t": path}, threads=PROFILE_THREADS)
try: try:
connection.execute("SUMMARIZE t") total, statistics = _statistics(connection, columns)
fields = [column[0] for column in connection.description]
summary = {row[0]: dict(zip(fields, row)) for row in connection.fetchall()}
head = connection.execute(f"FROM t LIMIT {EXAMPLE_SCAN_ROWS}").fetchall() head = connection.execute(f"FROM t LIMIT {EXAMPLE_SCAN_ROWS}").fetchall()
profiled = [ profiled = [
_profile_column( _profile_column(
connection, connection,
column, column,
summary[column.name], stats,
total,
head, head,
position, position,
notes.get(column.name), notes.get(column.name),
) )
for position, column in enumerate(columns) for position, (column, stats) in enumerate(zip(columns, statistics))
] ]
sample = connection.execute( sample = connection.execute(
f"FROM t USING SAMPLE reservoir({SAMPLE_ROWS} ROWS) REPEATABLE (42)" f"FROM t USING SAMPLE reservoir({SAMPLE_ROWS} ROWS) REPEATABLE (42)"
@@ -84,24 +87,57 @@ def read_rows(
return [[json_value(value, PREVIEW_CELL_CHARS) for value in row] for row in rows] return [[json_value(value, PREVIEW_CELL_CHARS) for value in row] for row in rows]
def _statistics(
connection: duckdb.DuckDBPyConnection, columns: list[DataColumn]
) -> tuple[int, list[dict[str, Any]]]:
"""总行数和每列的统计值,一次扫描算完:非空数、去重数(估算),数值列的最小、最大、平均和四分位数,
日期列的最小、最大
四分位数用 quantile_disc 精确计算:近似算法在多线程下每次结果略有不同,同一个文件重新上传概况会变
"""
selects, layout = ["count(*)"], []
for column in columns:
name = quote_identifier(column.name)
fields = {
"count": f"count({name})",
"approx_unique": f"approx_count_distinct({name})",
}
if column.type == "number":
fields |= {
"min": f"min({name})",
"max": f"max({name})",
"avg": f"avg({name})",
"quartiles": f"quantile_disc({name}, [0.25, 0.5, 0.75])",
}
elif column.type in ("date", "datetime"):
fields |= {"min": f"min({name})", "max": f"max({name})"}
layout.append(list(fields))
selects.extend(fields.values())
row = connection.execute(f"SELECT {', '.join(selects)} FROM t").fetchone()
values = iter(row[1:])
return row[0], [dict(zip(fields, values)) for fields in layout]
def _profile_column( def _profile_column(
connection: duckdb.DuckDBPyConnection, connection: duckdb.DuckDBPyConnection,
column: DataColumn, column: DataColumn,
stats: dict[str, Any], stats: dict[str, Any],
total: int,
head: list[tuple], head: list[tuple],
position: int, position: int,
note: str | None, note: str | None,
) -> ColumnProfile: ) -> ColumnProfile:
item: dict[str, Any] = {} item: dict[str, Any] = {}
if column.type == "number": if column.type == "number":
item = { quartiles = stats["quartiles"] or [None] * 3
key: _number(stats[key]) item = {key: _number(stats[key]) for key in ("min", "max", "avg")}
for key in ("min", "max", "avg", "q25", "q50", "q75") item |= {
key: _number(value) for key, value in zip(("q25", "q50", "q75"), quartiles)
} }
elif column.type in ("date", "datetime"): elif column.type in ("date", "datetime"):
# SUMMARIZE 的时间用空格分隔,和预览里的 ISO 写法统一 # 和预览里的时间写法一致
item = { item = {
key: None if stats[key] is None else stats[key].replace(" ", "T", 1) key: None if stats[key] is None else stats[key].isoformat()
for key in ("min", "max") for key in ("min", "max")
} }
if ( if (
@@ -124,23 +160,22 @@ def _profile_column(
examples.append(value) examples.append(value)
if len(examples) == EXAMPLES: if len(examples) == EXAMPLES:
break break
null_ratio = float(stats["null_percentage"]) / 100
return ColumnProfile( return ColumnProfile(
**column.model_dump(), **column.model_dump(),
null_ratio=round(null_ratio, 4), null_ratio=round(1 - stats["count"] / total, 4),
# approx_unique 是估算值,小表上可能比行数还大 # approx_unique 是估算值,小表上可能比非空值的个数还大
distinct=min(stats["approx_unique"], round(stats["count"] * (1 - null_ratio))), distinct=min(stats["approx_unique"], stats["count"]),
examples=examples, examples=examples,
note=note, note=note,
**item, **item,
) )
def _number(text: str | None) -> int | float | None: def _number(raw: Any) -> int | float | None:
"""SUMMARIZE 的统计值是字符串;整数不带小数点,其余保留 6 位小数;Parquet 里的 inf 返回空""" """整数不带小数点,其余保留 6 位小数;Parquet 里的 inf、NaN 返回空"""
if text is None: if raw is None:
return None return None
value = float(text) value = float(raw)
if not math.isfinite(value): if not math.isfinite(value):
return None return None
if value.is_integer() and abs(value) < EXACT_FLOAT_INT: if value.is_integer() and abs(value) < EXACT_FLOAT_INT:
+8 -3
View File
@@ -87,8 +87,13 @@ def run_query(
) )
def connect(tables: Mapping[str, Path]) -> duckdb.DuckDBPyConnection: def connect(
"""建好视图并锁定的连接;列概况也用它,和模型查询看到的数据一致""" tables: Mapping[str, Path], *, threads: int = THREADS
) -> duckdb.DuckDBPyConnection:
"""建好视图并锁定的连接;列概况也用它,和模型查询看到的数据一致
threads 只给解析子进程算列概况时放宽,模型写的查询一律用 THREADS
"""
connection = duckdb.connect( connection = duckdb.connect(
":memory:", ":memory:",
config={ config={
@@ -96,7 +101,7 @@ def connect(tables: Mapping[str, Path]) -> duckdb.DuckDBPyConnection:
"autoload_known_extensions": False, "autoload_known_extensions": False,
"python_enable_replacements": False, "python_enable_replacements": False,
"memory_limit": MEMORY_LIMIT, "memory_limit": MEMORY_LIMIT,
"threads": THREADS, "threads": threads,
}, },
) )
try: try:
+182 -63
View File
@@ -1,16 +1,20 @@
"""把各种格式读成原始网格(CSV、Excel)或按列的记录(JSON、JSONL),Parquet 直接读成 Arrow 表 """把各种格式读成原始网格(CSV、Excel)或按列的记录(JSON、JSONL),Parquet 直接读成 Arrow 表
CSV 和 Excel 可能有标题行、多行表头、合计行,所以先读成不带表头的网格,再由 tables.py 识别表区域 CSV 和 Excel 可能有标题行、多行表头、合计行,所以先读成不带表头的网格,再由 tables.py 识别表区域
文本格式边读边转成列(grid.py),不把整个文件解码进内存
""" """
import codecs
import csv import csv
import io import io
import json import json
from dataclasses import dataclass, field from dataclasses import dataclass, field
from itertools import islice
from pathlib import Path from pathlib import Path
from typing import Any from typing import Any
import pyarrow as pa import pyarrow as pa
import pyarrow.csv as pacsv
import pyarrow.parquet as pq import pyarrow.parquet as pq
from python_calamine import ( from python_calamine import (
CalamineWorkbook, CalamineWorkbook,
@@ -28,7 +32,8 @@ from app.datafiles.base import (
rows_exceeded, rows_exceeded,
scaled, scaled,
) )
from app.datafiles.columns import as_text from app.datafiles.cells import as_text
from app.datafiles.grid import CHUNK_ROWS, Column, Grid, GridBuilder
from app.schemas.data_file import ParseWarning from app.schemas.data_file import ParseWarning
# 原始网格比数据多出的标题、表头、空行,读取时先粗略拦住明显超限的文件,精确的行数在拆表后检查 # 原始网格比数据多出的标题、表头、空行,读取时先粗略拦住明显超限的文件,精确的行数在拆表后检查
@@ -43,37 +48,120 @@ MAX_BAD_LINE_RATIO = 0.01
# 逐行读取时每隔这么多行报告一次进度 # 逐行读取时每隔这么多行报告一次进度
PROGRESS_ROWS = 20_000 PROGRESS_ROWS = 20_000
# Arrow 读 CSV 每次处理的字节数;单行超过它时 Arrow 报错,改用标准库
ARROW_BLOCK_BYTES = 4 * 1024 * 1024
# 单元格最长的字符数;csv 模块默认 128KB,长文本单元格会报错 # 单元格最长的字符数;csv 模块默认 128KB,长文本单元格会报错
csv.field_size_limit(16 * 1024 * 1024) csv.field_size_limit(16 * 1024 * 1024)
@dataclass
class Grid:
"""不带表头的原始网格;rows[i] 是文件或 sheet 中的第 i + 1 行"""
sheet: str | None
rows: list[list[Any]]
@dataclass @dataclass
class Records: class Records:
"""已有列结构的数据,按列存放""" """已有列结构的数据,按列存放"""
names: list[str] names: list[str]
columns: list[list[Any]] columns: list[Column]
warnings: list[ParseWarning] = field(default_factory=list) warnings: list[ParseWarning] = field(default_factory=list)
def read_delimited(text: str, delimiter: str, report: Reporter = no_report) -> Grid: def read_delimited(
rows = [] path: Path, codec: str, delimiter: str, report: Reporter = no_report
stream = io.StringIO(text, newline="") ) -> Grid:
for row in csv.reader(stream, delimiter=delimiter): """codec 是 detect 确认过的编码,整个文件都能按它解码
rows.append(row)
if len(rows) > MAX_TABLE_ROWS + ROW_SLACK: 先用 Arrow 读(C++ 实现,快一个数量级);它要求每行列数相同,遇到标题行、说明行这类长短不齐的行,
raise rows_exceeded(None) 或引号不成对、单行过长,就报错。这时改用标准库 csv 从头读一遍。Arrow 读得出来时两者结果相同,测试里对比
if len(rows) % PROGRESS_ROWS == 0: """
report(stream.tell() / len(text)) report = _increasing(report)
return Grid(None, rows) try:
return _read_delimited_arrow(path, codec, delimiter, report)
except pa.ArrowInvalid:
return _read_delimited_python(path, codec, delimiter, report)
def _read_delimited_arrow(
path: Path, codec: str, delimiter: str, report: Reporter
) -> Grid:
with path.open("rb") as raw, io.TextIOWrapper(raw, codec, newline="") as text:
first = next(csv.reader(text, delimiter=delimiter), [])
if not first:
raise pa.ArrowInvalid("第一行是空行,列数要从后面的行推断")
names = [f"c{i}" for i in range(len(first))]
builder = GridBuilder(strings=True)
size = max(path.stat().st_size, 1)
with path.open("rb") as raw:
if codec == "utf-8-sig" and raw.read(3) != codecs.BOM_UTF8:
raw.seek(0)
reader = pacsv.open_csv(
raw,
read_options=pacsv.ReadOptions(
column_names=names,
block_size=ARROW_BLOCK_BYTES,
encoding="utf8" if codec.startswith("utf-8") else codec,
),
parse_options=pacsv.ParseOptions(
delimiter=delimiter, newlines_in_values=True, ignore_empty_lines=False
),
convert_options=pacsv.ConvertOptions(
column_types={name: pa.string() for name in names},
strings_can_be_null=False,
quoted_strings_can_be_null=False,
),
)
batches: list[pa.RecordBatch] = []
waiting = 0
for batch in reader:
batches.append(batch)
waiting += len(batch)
# 攒到一个块再转:每段太短时,类型推断的固定开销会变大
if waiting >= CHUNK_ROWS:
_add_batches(builder, batches)
batches, waiting = [], 0
if builder.height > MAX_TABLE_ROWS + ROW_SLACK:
raise rows_exceeded(None)
report(raw.tell() / size)
if batches:
_add_batches(builder, batches)
if builder.height > MAX_TABLE_ROWS + ROW_SLACK:
raise rows_exceeded(None)
return builder.finish(None)
def _add_batches(builder: GridBuilder, batches: list[pa.RecordBatch]) -> None:
builder.add_strings(
[
pa.concat_arrays([batch.column(i) for batch in batches])
for i in range(batches[0].num_columns)
]
)
def _read_delimited_python(
path: Path, codec: str, delimiter: str, report: Reporter
) -> Grid:
builder = GridBuilder(strings=True)
size = max(path.stat().st_size, 1)
with path.open("rb") as raw, io.TextIOWrapper(raw, codec, newline="") as text:
reader = csv.reader(text, delimiter=delimiter)
while rows := list(islice(reader, PROGRESS_ROWS)):
builder.add_rows(rows)
if builder.height > MAX_TABLE_ROWS + ROW_SLACK:
raise rows_exceeded(None)
report(raw.tell() / size)
return builder.finish(None)
def _increasing(report: Reporter) -> Reporter:
"""Arrow 读到一半失败、换标准库从头读时,进度不往回走"""
reached = 0.0
def inner(fraction: float) -> None:
nonlocal reached
if fraction > reached:
reached = fraction
report(fraction)
return inner
def read_excel( def read_excel(
@@ -104,9 +192,19 @@ def read_excel(
sheet = book.get_sheet_by_name(meta.name) sheet = book.get_sheet_by_name(meta.name)
if sheet.total_height > MAX_TABLE_ROWS + ROW_SLACK: if sheet.total_height > MAX_TABLE_ROWS + ROW_SLACK:
raise rows_exceeded(meta.name) raise rows_exceeded(meta.name)
# 保留左上角的空白区域,行号和 Excel 里看到的一致 # 解析 sheet 的 XML 占读取的大部分时间,逐行转成列占剩下的
rows = sheet.to_python(skip_empty_area=False) rows = scaled(
if not any(cell != "" for row in rows for cell in row): report, (number + 0.8) / len(sheets), (number + 1) / len(sheets)
)
# iter_rows 保留顶部的空行,行号和 Excel 里看到的一致;左侧的空列不影响识别表区域;
# 没有单元格的 sheet 上 iter_rows 会崩溃,不调用
builder = GridBuilder()
cells = sheet.iter_rows() if sheet.height else iter(())
while batch := list(islice(cells, PROGRESS_ROWS)):
builder.add_rows(batch)
rows(builder.height / sheet.height)
grid = builder.finish(meta.name)
if not any(column.present().any() for column in grid.columns):
warnings.append( warnings.append(
ParseWarning( ParseWarning(
code="empty_sheet_skipped", code="empty_sheet_skipped",
@@ -114,7 +212,7 @@ def read_excel(
) )
) )
continue continue
grids.append(Grid(meta.name, rows)) grids.append(grid)
finally: finally:
book.close() book.close()
return grids, warnings return grids, warnings
@@ -137,42 +235,53 @@ def read_json(value: Any, report: Reporter = no_report) -> Grid | Records:
if not records: if not records:
raise DataFileError(422, "JSON 里没有找到数据") raise DataFileError(422, "JSON 里没有找到数据")
if all(isinstance(item, list) for item in records): if all(isinstance(item, list) for item in records):
return Grid(None, [[_scalar(cell) for cell in row] for row in records]) if len(records) > MAX_TABLE_ROWS + ROW_SLACK:
raise rows_exceeded(None)
return Grid.from_rows(
None, ([_scalar(cell) for cell in row] for row in records)
)
if all(isinstance(item, dict) for item in records): if all(isinstance(item, dict) for item in records):
result = _flatten_records(records, report) flattener = _Flattener()
for number, record in enumerate(records):
if number % PROGRESS_ROWS == 0:
report(number / len(records))
flattener.add(record)
result = flattener.finish()
result.warnings = warnings result.warnings = warnings
return result return result
if not any(isinstance(item, (dict, list)) for item in records): if not any(isinstance(item, (dict, list)) for item in records):
return Records(["值"], [list(records)], warnings) if len(records) > MAX_TABLE_ROWS:
raise rows_exceeded(None)
return Records(["值"], [Column.from_values(records)], warnings)
raise DataFileError(422, "JSON 数组里的元素类型不一致,需要都是对象或都是数组") raise DataFileError(422, "JSON 数组里的元素类型不一致,需要都是对象或都是数组")
def read_jsonl(text: str, report: Reporter = no_report) -> Records: def read_jsonl(path: Path, codec: str, report: Reporter = no_report) -> Records:
records, bad, total = [], [], 0 """逐行读取;codec 是 detect 确认过的编码"""
lines = text.splitlines() flattener, bad, total = _Flattener(), [], 0
# 解析每行 JSON 和展开成列各占一半 size = max(path.stat().st_size, 1)
parsing = scaled(report, 0.0, 0.5) with path.open("rb") as raw, io.TextIOWrapper(raw, codec) as text:
for number, line in enumerate(lines, start=1): for number, line in enumerate(text, start=1):
if number % PROGRESS_ROWS == 0: if number % PROGRESS_ROWS == 0:
parsing(number / len(lines)) report(raw.tell() / size)
if not line.strip(): if not line.strip():
continue continue
total += 1 total += 1
try: try:
record = json.loads(line) record = json.loads(line)
except ValueError: except ValueError:
record = None record = None
if isinstance(record, dict): if isinstance(record, dict):
records.append(record) flattener.add(record)
else: else:
bad.append(number) bad.append(number)
if not records: if not flattener.count:
raise DataFileError(422, "JSONL 里没有找到数据") raise DataFileError(422, "JSONL 里没有找到数据")
if len(bad) > total * MAX_BAD_LINE_RATIO: if len(bad) > total * MAX_BAD_LINE_RATIO:
raise DataFileError( raise DataFileError(
422, f"JSONL 有 {len(bad)} 行无法解析(例如第 {bad[0]} 行),超过 1%" 422, f"JSONL 有 {len(bad)} 行无法解析(例如第 {bad[0]} 行),超过 1%"
) )
result = _flatten_records(records, scaled(report, 0.5, 1.0)) result = flattener.finish()
if bad: if bad:
result.warnings.append( result.warnings.append(
ParseWarning( ParseWarning(
@@ -226,24 +335,34 @@ def _is_record_list(value: Any) -> bool:
) )
def _flatten_records(records: list[dict], report: Reporter = no_report) -> Records: class _Flattener:
# 列按首次出现的顺序排列,记录里没有的字段为空 """逐条展开记录,按列存放:列按首次出现的顺序排列,记录里没有的字段为空"""
index: dict[str, int] = {}
columns: list[list[Any]] = [] def __init__(self) -> None:
for row_number, record in enumerate(records): self._index: dict[str, int] = {}
if row_number % PROGRESS_ROWS == 0: self._builder = GridBuilder()
report(row_number / len(records))
@property
def count(self) -> int:
return self._builder.height
def add(self, record: dict) -> None:
if self._builder.height >= MAX_TABLE_ROWS:
raise rows_exceeded(None)
flat: dict[str, Any] = {} flat: dict[str, Any] = {}
_flatten(record, "", 0, flat) _flatten(record, "", 0, flat)
row: list[Any] = [None] * len(self._index)
for name, value in flat.items(): for name, value in flat.items():
if name not in index: position = self._index.setdefault(name, len(self._index))
index[name] = len(columns) if position < len(row):
columns.append([None] * row_number) row[position] = value
columns[index[name]].append(value) else:
for column in columns: row.append(value)
if len(column) == row_number: self._builder.add(row)
column.append(None)
return Records(list(index), columns) def finish(self) -> Records:
grid = self._builder.finish(None)
return Records(list(self._index), grid.columns)
def _flatten(record: dict, prefix: str, level: int, out: dict[str, Any]) -> None: def _flatten(record: dict, prefix: str, level: int, out: dict[str, Any]) -> None:
+109 -82
View File
@@ -1,13 +1,17 @@
"""在原始网格里识别表:按空行切块,找表头,合并两级表头,去掉标题、说明和合计行 """在原始网格里识别表:按空行切块,找表头,合并两级表头,去掉标题、说明和合计行
只用规则;判断不了的地方照常产出表,同时记一条警告,调用方据此提示用户 只用规则;判断不了的地方照常产出表,同时记一条警告,调用方据此提示用户
要看每一行的规则(空行切块、合计行)对整列批量计算;表头、标题、说明只看每块开头或末尾的几行,取成 Python 值判断
""" """
import math import math
import re
from dataclasses import dataclass, field from dataclasses import dataclass, field
from typing import Any from typing import Any
import numpy as np
import pyarrow as pa
import pyarrow.compute as pc
from app.datafiles.base import ( from app.datafiles.base import (
MAX_TABLE_ROWS, MAX_TABLE_ROWS,
Reporter, Reporter,
@@ -16,8 +20,8 @@ from app.datafiles.base import (
rows_exceeded, rows_exceeded,
where, where,
) )
from app.datafiles.columns import as_text, classify from app.datafiles.cells import WHITESPACE_CLASS, as_text, classify, mask
from app.datafiles.readers import Grid from app.datafiles.grid import Column, Grid
from app.schemas.data_file import ParseWarning from app.schemas.data_file import ParseWarning
# 只在每块的前几行里找表头 # 只在每块的前几行里找表头
@@ -40,6 +44,11 @@ TOTAL_LABELS = {
} }
# 例如"华东小计""一季度合计" # 例如"华东小计""一季度合计"
TOTAL_SUFFIXES = ("小计", "合计", "总计") TOTAL_SUFFIXES = ("小计", "合计", "总计")
_TOTAL_SET = pa.array(sorted(TOTAL_LABELS))
_TOTAL_STRIP = f"[{WHITESPACE_CLASS}::]"
_TOTAL_SUFFIX = f"(?:{'|'.join(TOTAL_SUFFIXES)})$"
# 每个合计字样(含后缀)至少有其中一个字符;英文按转小写之前的大小写都列出
_TOTAL_HINT = "[计和均tTmMeE]"
# 警告里引用的单元格文本最长显示这么多字 # 警告里引用的单元格文本最长显示这么多字
QUOTE_CHARS = 30 QUOTE_CHARS = 30
@@ -51,8 +60,8 @@ class RawTable:
# 同一 sheet 切出多张表时的序号,从 1 开始;只有一张表时为空 # 同一 sheet 切出多张表时的序号,从 1 开始;只有一张表时为空
part: int | None part: int | None
header: list[Any] header: list[Any]
# 数据行,宽度与表头相同 # 数据列,和表头一一对应
rows: list[list[Any]] columns: list[Column]
warnings: list[ParseWarning] = field(default_factory=list) warnings: list[ParseWarning] = field(default_factory=list)
@@ -62,8 +71,8 @@ class _Piece:
# 表头所在行在网格中的下标 # 表头所在行在网格中的下标
header_row: int header_row: int
header: list[Any] header: list[Any]
# 数据行在网格中的下标 # 数据行在网格中的范围 [start, stop)
data: list[int] data: list[tuple[int, int]]
warnings: list[ParseWarning] warnings: list[ParseWarning]
@@ -74,40 +83,41 @@ def split_tables(
几步都要扫一遍所有行,每步结束报告一次进度,比例来自大表上的实测耗时 几步都要扫一遍所有行,每步结束报告一次进度,比例来自大表上的实测耗时
""" """
rows = _normalize(grid.rows)
report(0.25)
sheet = grid.sheet sheet = grid.sheet
present = [column.present() for column in grid.columns]
filled = np.logical_or.reduce(present) if present else np.zeros(grid.height, bool)
blocks = _blocks(filled)
used = _used_columns(present, blocks)
firsts = grid.rows_at(np.array([start for start, _ in blocks], np.int64))
report(0.3)
pieces: list[_Piece] = [] pieces: list[_Piece] = []
dropped: list[ParseWarning] = [] dropped: list[ParseWarning] = []
blocks = _blocks(rows) for number, ((start, stop), columns, first) in enumerate(
report(0.3) zip(blocks, used, firsts), start=1
for number, block in enumerate(blocks, start=1): ):
columns = [
c for c in range(len(rows[0])) if any(rows[r][c] is not None for r in block)
]
last = pieces[-1] if pieces else None last = pieces[-1] if pieces else None
# 表中间夹着空行时,后面的块以数据行开头、列也落在上一张表内,接在上一张表后面 # 表中间夹着空行时,后面的块以数据行开头、列也落在上一张表内,接在上一张表后面
if ( if (
last last
and not _is_label_row([rows[block[0]][c] for c in columns]) and not _is_label_row([first[c] for c in columns])
and set(columns) <= set(last.columns) and set(columns) <= set(last.columns)
): ):
last.data.extend(block) last.data.append((start, stop))
continue continue
if len(block) < 2: if stop - start < 2:
dropped.append( dropped.append(
ParseWarning( ParseWarning(
code="block_dropped", code="block_dropped",
message=f"{where(sheet)}第 {block[0] + 1} 行的{_quote(rows[block[0]])}不是表格,已跳过", message=f"{where(sheet)}第 {start + 1} 行的{_quote(first)}不是表格,已跳过",
) )
) )
continue continue
pieces.append(_read_block(rows, block, columns, sheet)) pieces.append(_read_block(grid, start, stop, columns, sheet))
report(0.3 + 0.25 * number / len(blocks)) report(0.3 + 0.25 * number / len(blocks))
tables = [] tables = []
for number, piece in enumerate(pieces, start=1): for number, piece in enumerate(pieces, start=1):
table = _finish(rows, piece, sheet) table = _finish(grid, piece, sheet)
report(0.55 + 0.45 * number / len(pieces)) report(0.55 + 0.45 * number / len(pieces))
if table is None: if table is None:
dropped.append( dropped.append(
@@ -125,41 +135,28 @@ def split_tables(
return tables, dropped return tables, dropped
def _normalize(grid: list[list[Any]]) -> list[list[Any]]: def _blocks(filled: np.ndarray) -> list[tuple[int, int]]:
"""去掉字符串两端空白,空字符串记为 None,补齐成同样宽度""" """连续的非空行组成的块 [start, stop)"""
width = max((len(row) for row in grid), default=0) edges = np.flatnonzero(np.diff(np.concatenate(([0], filled, [0])).astype(np.int8)))
rows = [] return list(zip(edges[0::2].tolist(), edges[1::2].tolist()))
for row in grid:
cells = [_cell(value) for value in row]
cells.extend([None] * (width - len(cells)))
rows.append(cells)
return rows
def _cell(value: Any) -> Any: def _used_columns(
if isinstance(value, str): present: list[np.ndarray], blocks: list[tuple[int, int]]
value = value.strip() ) -> list[list[int]]:
return value or None """每块里有值的列;块之间只隔着空行,按块的起点分段求或即可"""
return value if not blocks:
return []
starts = np.array([start for start, _ in blocks])
def _blocks(rows: list[list[Any]]) -> list[list[int]]: used = np.array([np.logical_or.reduceat(column, starts) for column in present])
blocks, current = [], [] return [np.flatnonzero(used[:, index]).tolist() for index in range(len(blocks))]
for index, row in enumerate(rows):
if any(value is not None for value in row):
current.append(index)
elif current:
blocks.append(current)
current = []
if current:
blocks.append(current)
return blocks
def _read_block( def _read_block(
rows: list[list[Any]], block: list[int], columns: list[int], sheet: str | None grid: Grid, start: int, stop: int, columns: list[int], sheet: str | None
) -> _Piece: ) -> _Piece:
cells = [[rows[r][c] for c in columns] for r in block] # 表头最多在第 HEADER_SEARCH_ROWS 行,两级表头还要看它下面两行
cells = grid.rows(start, min(stop, start + HEADER_SEARCH_ROWS + 2), columns)
need = max(1, math.ceil(len(columns) / 2)) need = max(1, math.ceil(len(columns) / 2))
dense = [sum(value is not None for value in row) >= need for row in cells] dense = [sum(value is not None for value in row) >= need for row in cells]
search = range(min(HEADER_SEARCH_ROWS, len(cells))) search = range(min(HEADER_SEARCH_ROWS, len(cells)))
@@ -173,7 +170,7 @@ def _read_block(
warnings.append( warnings.append(
ParseWarning( ParseWarning(
code="header_uncertain", code="header_uncertain",
message=f"{where(sheet)}没有找到明确的表头,用第 {block[header_at] + 1} 行作表头,可能识别错误", message=f"{where(sheet)}没有找到明确的表头,用第 {start + header_at + 1} 行作表头,可能识别错误",
) )
) )
if header_at > 0: if header_at > 0:
@@ -183,7 +180,7 @@ def _read_block(
warnings.append( warnings.append(
ParseWarning( ParseWarning(
code="title_skipped", code="title_skipped",
message=f"{where(sheet)}{_row_range(block[0], block[header_at - 1])}是标题,已跳过:{skipped}", message=f"{where(sheet)}{_row_range(start, start + header_at - 1)}是标题,已跳过:{skipped}",
) )
) )
@@ -191,7 +188,7 @@ def _read_block(
data_at = header_at + 1 data_at = header_at + 1
if ( if (
confident confident
and header_at + 2 < len(cells) and header_at + 2 < stop - start
and _is_subheader(cells[header_at], cells[header_at + 1], cells[header_at + 2]) and _is_subheader(cells[header_at], cells[header_at + 1], cells[header_at + 2])
): ):
header = _merge_header(cells[header_at], cells[header_at + 1]) header = _merge_header(cells[header_at], cells[header_at + 1])
@@ -199,20 +196,28 @@ def _read_block(
warnings.append( warnings.append(
ParseWarning( ParseWarning(
code="header_merged", code="header_merged",
message=f"{where(sheet)}{_row_range(block[header_at], block[header_at + 1])}合并为表头", message=f"{where(sheet)}{_row_range(start + header_at, start + header_at + 1)}合并为表头",
) )
) )
return _Piece(columns, block[header_at], header, block[data_at:], warnings) return _Piece(
columns, start + header_at, header, [(start + data_at, stop)], warnings
)
def _finish(rows: list[list[Any]], piece: _Piece, sheet: str | None) -> RawTable | None: def _finish(grid: Grid, piece: _Piece, sheet: str | None) -> RawTable | None:
data = [(index, [rows[index][c] for c in piece.columns]) for index in piece.data] ranges = [np.arange(start, stop) for start, stop in piece.data if start < stop]
rows = np.concatenate(ranges) if ranges else np.zeros(0, np.int64)
warnings = piece.warnings warnings = piece.warnings
# 表格下方只有一个文字单元格的行是说明,例如"注:数据来源于……" # 表格下方只有一个文字单元格的行是说明,例如"注:数据来源于……"
notes = [] notes = []
while data and len(piece.columns) >= 3 and _is_note(data[-1][1]): while len(rows) and len(piece.columns) >= 3:
notes.append(data.pop()) last = int(rows[-1])
row = grid.rows(last, last + 1, piece.columns)[0]
if not _is_note(row):
break
notes.append((last, row))
rows = rows[:-1]
if notes: if notes:
notes.reverse() notes.reverse()
shown = "".join(_quote(row) for _, row in notes) shown = "".join(_quote(row) for _, row in notes)
@@ -224,22 +229,59 @@ def _finish(rows: list[list[Any]], piece: _Piece, sheet: str | None) -> RawTable
) )
# 合计、小计行不是明细数据,留着会让求和重复计算 # 合计、小计行不是明细数据,留着会让求和重复计算
totals = [index for index, row in data if _is_total(row)] totals = _totals(grid, piece.columns, rows)
if totals: if totals.any():
data = [(index, row) for index, row in data if not _is_total(row)] listed = brief([str(index + 1) for index in rows[totals].tolist()])
listed = brief([str(index + 1) for index in totals])
warnings.append( warnings.append(
ParseWarning( ParseWarning(
code="total_rows_removed", code="total_rows_removed",
message=f"{where(sheet)}已排除 {len(totals)} 行合计、小计(第 {listed} 行)", message=f"{where(sheet)}已排除 {int(totals.sum())} 行合计、小计(第 {listed} 行)",
) )
) )
rows = rows[~totals]
if not data: if not len(rows):
return None return None
if len(data) > MAX_TABLE_ROWS: if len(rows) > MAX_TABLE_ROWS:
raise rows_exceeded(sheet) raise rows_exceeded(sheet)
return RawTable(sheet, None, piece.header, [row for _, row in data], warnings) columns = [grid.columns[c].take(rows) for c in piece.columns]
return RawTable(sheet, None, piece.header, columns, warnings)
def _totals(grid: Grid, columns: list[int], rows: np.ndarray) -> np.ndarray:
"""rows 里哪些是合计行:每行只看前两个非空单元格,其中有合计字样的文字就是"""
if not len(rows):
return np.zeros(0, bool)
low, high = int(rows[0]), int(rows[-1]) + 1
seen = np.zeros(high - low, np.int64)
total = np.zeros(high - low, bool)
for c in columns:
column = grid.columns[c].slice(low, high)
present = column.present()
open_ = present & (seen < 2)
if open_.any():
total |= open_ & column.matches(_is_total_label)
seen += present
if (seen >= 2).all():
break
return total[rows - low]
def _is_total_label(strings: pa.Array) -> pa.Array:
"""去掉空白和冒号、转小写后是合计字样;合计字样都含 _TOTAL_HINT 里的字符,先用它筛掉绝大多数值"""
found = np.zeros(len(strings), bool)
rows = np.flatnonzero(mask(pc.match_substring_regex(strings, _TOTAL_HINT)))
if len(rows):
key = pc.utf8_lower(
pc.replace_substring_regex(strings.take(rows), _TOTAL_STRIP, "")
)
found[rows] = mask(
pc.or_(
pc.is_in(key, value_set=_TOTAL_SET),
pc.match_substring_regex(key, _TOTAL_SUFFIX),
)
)
return pa.array(found)
def _is_label(value: Any) -> bool: def _is_label(value: Any) -> bool:
@@ -276,21 +318,6 @@ def _is_note(row: list[Any]) -> bool:
return len(values) == 1 and _is_label(values[0]) return len(values) == 1 and _is_label(values[0])
def _is_total(row: list[Any]) -> bool:
checked = 0
for value in row:
if value is None:
continue
if isinstance(value, str):
key = re.sub(r"[\s::]", "", value).lower()
if key in TOTAL_LABELS or key.endswith(TOTAL_SUFFIXES):
return True
checked += 1
if checked == 2:
return False
return False
def _quote(row: list[Any]) -> str: def _quote(row: list[Any]) -> str:
"""一行里第一个非空单元格,截断后加引号""" """一行里第一个非空单元格,截断后加引号"""
value = next((value for value in row if value is not None), None) value = next((value for value in row if value is not None), None)
+32 -18
View File
@@ -1,39 +1,49 @@
"""调用 screen agent,按它的回复和最后一次提交的结果写回状态 """调用 screen agent,按它的回复和成功提交的动作写回状态
kind 由模型决定;代码只在"回复 done 但最后一次提交失败"时不采纳,因为此时没有能交付的动作 kind 由模型决定;代码只在"回复 done 但最后一次提交失败"时不采纳:最后想做的修改没有写入草稿,
模型却说已完成。这是保守做法,last_write_failed 只看工具结果,任务是否完成仍以模型的结论和评测为准
agent 正常结束却没有结论时(调用次数用尽、只回了文本或空回复)补一次只读收尾;取消和异常照常抛出,不收尾
""" """
import json import json
from langchain_core.messages import HumanMessage from langchain_core.messages import HumanMessage
from app.agent.screen_agent import AgentReply, get_screen_agent from app.agent.screen_agent import (
from app.agent.tools import ScreenAgentContext AgentReply,
finalize_reply,
get_screen_agent,
reached_call_limit,
)
from app.agent.tools import ScreenAgentContext, last_write_failed
from app.processor.base import BaseNode from app.processor.base import BaseNode
from app.processor.history import history_messages from app.processor.history import history_messages
from app.processor.state import QueryState, run_data from app.processor.state import QueryState, run_data
from app.processor.summary import request_summary from app.processor.summary import request_summary
from app.resources.component import get_catalog_cache from app.resources.component import get_catalog_cache
# 评测按这句话识别"草稿因最后一次提交失败没有交付",单独核对草稿本身是否满足要求
DRAFT_REJECTED = "修改没有通过校验,这次没有生效,可以换个说法再试"
def reply_update(result: dict) -> QueryState:
def reply_update(result: dict, *, limit_hit: bool = False) -> QueryState:
"""把 agent 的运行结果转换成 QueryState 的更新""" """把 agent 的运行结果转换成 QueryState 的更新"""
reply: AgentReply | None = result.get("structured_response") reply: AgentReply | None = result.get("structured_response")
if reply is None: if reply is None:
# 调用次数用尽等情况,agent 没有给出结论 # 收尾也没有给出合法的结论;按真实原因说明,不再一律说超出上限
return { answer = (
"status": "failed", "处理步骤超出上限,这次没有完成,修改没有生效,可以换个说法再试"
"answer": "处理步骤超出上限,这次没有完成,可以换个说法再试", if limit_hit
} else "这次没有得出结论,修改没有生效,可以换个说法再试"
)
return {"status": "failed", "answer": answer}
if reply.kind == "clarify": if reply.kind == "clarify":
return {"status": "waiting_for_input", "clarification": reply.clarification} return {"status": "waiting_for_input", "clarification": reply.clarification}
if reply.kind == "unable": if reply.kind == "unable":
return {"status": "failed", "answer": reply.message} return {"status": "failed", "answer": reply.message}
if result.get("apply_error"): if last_write_failed(result.get("messages", [])):
return { return {"status": "failed", "answer": DRAFT_REJECTED}
"status": "failed", # 各次成功提交的动作按顺序拼接;clarify、unable 不交付,下次请求从原始大屏重新执行
"answer": "修改没有通过校验,这次没有生效,可以换个说法再试",
}
return { return {
"status": "completed", "status": "completed",
"actions": result.get("actions") or [], "actions": result.get("actions") or [],
@@ -70,12 +80,16 @@ class ScreenAgentNode(BaseNode):
data=run_data(), data=run_data(),
), ),
) )
limit_hit = False
if result.get("structured_response") is None: if result.get("structured_response") is None:
limit_hit = reached_call_limit(result["messages"])
last = result["messages"][-1] last = result["messages"][-1]
self.logger.warning( self.logger.warning(
f"agent 没有给出结论: {last.type} {str(last.content)[:500]} " f"agent 没有给出结论({'调用次数用尽' if limit_hit else '没有调用 AgentReply'}),收尾一次: "
f"tool_calls={getattr(last, 'tool_calls', None)} " f"{last.type} {str(last.content)[:500]} "
f"metadata={getattr(last, 'response_metadata', None)}" f"metadata={getattr(last, 'response_metadata', None)}"
) )
updated_state.update(reply_update(result)) reply = await finalize_reply(result, limit_hit=limit_hit)
result = {**result, "structured_response": reply}
updated_state.update(reply_update(result, limit_hit=limit_hit))
return updated_state return updated_state
+3 -1
View File
@@ -20,7 +20,9 @@ class QueryState(TypedDict):
input_context: NotRequired[InputContext] # 校验后的上下文,包含最新大屏和澄清历史 input_context: NotRequired[InputContext] # 校验后的上下文,包含最新大屏和澄清历史
intent_decision: NotRequired[DecisionResponse | None] # Jev 原始意图判断及分数 intent_decision: NotRequired[DecisionResponse | None] # Jev 原始意图判断及分数
intent: NotRequired[Route | None] # 门控后的去向 intent: NotRequired[Route | None] # 门控后的去向
actions: NotRequired[list[Action] | None] # screen agent 最后一次成功提交的动作 actions: NotRequired[
list[Action] | None
] # screen agent 各次成功提交的动作,按执行顺序拼接
clarification: NotRequired[Clarification | None] # 当前澄清问题及选项 clarification: NotRequired[Clarification | None] # 当前澄清问题及选项
clarification_reply: NotRequired[ clarification_reply: NotRequired[
ClarificationReply | None ClarificationReply | None
+1 -1
View File
@@ -15,7 +15,7 @@ class AddAction(ActionModel):
node: ScreenNode = Field( node: ScreenNode = Field(
description="新组件。id、type、layout 必填,其余只写需要覆盖物料默认配置的字段;" description="新组件。id、type、layout 必填,其余只写需要覆盖物料默认配置的字段;"
"后端按 type 找到物料,用默认配置补全:对象逐层覆盖,数组整体替换,null 明确覆盖。" "后端按 type 找到物料,用默认配置补全:对象逐层覆盖,数组整体替换,null 明确覆盖。"
"id 不能与现有组件重复,同一批里先 delete 的 ID 可以再用;type 用 list_components 返回的 type。" "id 不能与当前大屏上的组件重复(包括之前提交里新增的),先 delete 的 ID 可以再用;type 用 list_components 返回的 type。"
'任何字段值都可以写成 {"$query": "SELECT ..."},后端替换成上传数据的查询结果(对象数组,键是结果的列名)' '任何字段值都可以写成 {"$query": "SELECT ..."},后端替换成上传数据的查询结果(对象数组,键是结果的列名)'
) )
sample_data: bool = Field( sample_data: bool = Field(
+3 -3
View File
@@ -1,7 +1,7 @@
"""异步执行与恢复大屏任务,当前使用单进程内存检查点""" """异步执行与恢复大屏任务,当前使用单进程内存检查点"""
from contextlib import aclosing from contextlib import aclosing
from datetime import datetime, timezone from datetime import UTC, datetime
from uuid import uuid4 from uuid import uuid4
from langgraph.types import Command from langgraph.types import Command
@@ -201,7 +201,7 @@ class AgentService:
prompt=state["request"]["prompt"], prompt=state["request"]["prompt"],
outcome=result.status, outcome=result.status,
summary=self._summarize(result), summary=self._summarize(result),
created_at=datetime.now(timezone.utc), created_at=datetime.now(UTC),
) )
) )
except Exception: except Exception:
@@ -210,7 +210,7 @@ class AgentService:
@staticmethod @staticmethod
def _result(state, thread_id): def _result(state, thread_id):
waiting = bool(state.get("__interrupt__")) waiting = bool(state.get("__interrupt__"))
# 只交付完成状态下的动作;actions 是 agent 最后一次成功提交、已通过模拟执行的整批动作 # 只交付完成状态下的动作;actions 是 agent 各次成功提交的动作按顺序拼接,在原始大屏上依次执行得到最终大屏
deliver = state["status"] == "completed" deliver = state["status"] == "completed"
return AgentResult( return AgentResult(