Files
ai-video-fullstack/backend/migrations/versions/20260813_0021_organize_tool_resources.py
T

511 lines
19 KiB
Python
Raw Normal View History

"""separate built-ins, resources, and per-use tool bindings
Revision ID: 20260813_0021
Revises: 20260812_0020
"""
from __future__ import annotations
from collections.abc import Mapping, Sequence
from copy import deepcopy
import json
import re
from alembic import op
import sqlalchemy as sa
from sqlalchemy.dialects import postgresql
revision: str = "20260813_0021"
down_revision: str | Sequence[str] | None = "20260812_0020"
branch_labels: str | Sequence[str] | None = None
depends_on: str | Sequence[str] | None = None
PLACEHOLDER = re.compile(r"{{\s*([A-Za-z][A-Za-z0-9_]*)\s*}}")
FULL_PLACEHOLDER = re.compile(r"^{{\s*([A-Za-z][A-Za-z0-9_]*)\s*}}$")
def _source_for(value) -> dict:
raw = json.dumps(value, ensure_ascii=False)
return (
{"source": "template", "value": value}
if PLACEHOLDER.search(raw)
else {"source": "literal", "value": value}
)
def _llm_rule(name: str, definition: Mapping | None) -> dict:
value_type = str((definition or {}).get("type") or "string")
return {
"variableName": name,
"expression": {
"type": "llm",
"parameterName": name,
"valueType": value_type,
"description": f"从当前对话中提取并确认 {name}",
},
}
def _builtin(
kind: str,
tool: Mapping,
definitions: Mapping,
allowed=None,
*,
camel_case: bool = False,
) -> dict:
definition = tool.get("definition") or {}
config = deepcopy(definition.get("config") or {})
config.pop("kind", None)
config["description"] = str(tool.get("description") or "")
if kind == "update_state":
names = list(allowed) if allowed is not None else list(definitions)
config["updates"] = [
_llm_rule(str(name), definitions.get(str(name)))
for name in names
if str(name) in definitions
]
if camel_case:
aliases = {
"message_type": "messageType",
"custom_message": "customMessage",
"capture_reason": "captureReason",
"allow_end_interruption": "allowEndInterruption",
}
config = {aliases.get(key, key): value for key, value in config.items()}
return config
def _builtin_key(kind: str) -> str:
return {
"end_conversation": "endConversation",
"update_state": "updateState",
"skip_turn": "skipTurn",
"request_human_handoff": "requestHumanHandoff",
}.get(kind, kind)
def _binding(
tool_id: str,
tool: Mapping,
arguments=None,
assignments=None,
inherited_assignments: Mapping | None = None,
) -> dict:
config = ((tool.get("definition") or {}).get("config") or {})
if tool.get("type") == "mcp":
schema = config.get("input_schema") or {}
names = list((schema.get("properties") or {}).keys())
else:
names = [
str(parameter.get("name"))
for parameter in config.get("parameters") or []
if parameter.get("name")
]
if arguments is None:
argument_bindings = {name: {"source": "model"} for name in names}
else:
argument_bindings = {
str(name): _source_for(value) for name, value in arguments.items()
}
return {
"toolId": tool_id,
"argumentBindings": argument_bindings,
"templateBindings": {},
"resultAssignments": dict(
(inherited_assignments or {}) if assignments is None else assignments
),
}
def _state_expression(value) -> dict:
"""Translate deterministic v3 values without silently changing semantics."""
if isinstance(value, str):
match = FULL_PLACEHOLDER.fullmatch(value)
if match:
return {"type": "variable", "variableName": match.group(1)}
if PLACEHOLDER.search(value):
raise RuntimeError(
"Update State 节点中的混合字符串模板无法无损迁移;"
"请在升级前改成完整变量占位符"
)
if not isinstance(value, (str, int, float, bool)):
raise RuntimeError("Update State 节点只支持字符串、数字、布尔值或变量")
return {"type": "constant", "value": value}
def _rewrite_templates(value, inputs: set[str]):
if isinstance(value, str):
def replace(match):
name = match.group(1)
if name.startswith(("system__", "secret__")):
return match.group(0)
if name.startswith("input__"):
inputs.add(name.removeprefix("input__"))
return match.group(0)
inputs.add(name)
return "{{input__" + name + "}}"
return PLACEHOLDER.sub(replace, value)
if isinstance(value, list):
return [_rewrite_templates(item, inputs) for item in value]
if isinstance(value, dict):
return {key: _rewrite_templates(item, inputs) for key, item in value.items()}
return value
def _add_template_bindings(binding: dict, inputs: set[str]) -> None:
binding.setdefault("templateBindings", {}).update(
{
name: {"source": "variable", "variableName": name}
for name in sorted(inputs)
}
)
def _freeze_inherited_agent_for_update_state(
data: dict,
settings: Mapping,
definitions: Mapping,
allowed: list,
) -> None:
"""Preserve the old per-Agent Update State write scope.
Before v4 an Agent could inherit the global Update State resource while its
writable variable names still lived on the node. A v4 built-in owns its
complete configuration, so this one exceptional case must snapshot the
current global stage configuration onto the Agent.
"""
runtime_mode = str(settings.get("runtimeMode") or "pipeline")
if runtime_mode != "realtime":
for source, target in (
("defaultLlmResourceId", "llmResourceId"),
("defaultAsrResourceId", "asrResourceId"),
("defaultTtsResourceId", "ttsResourceId"),
):
data[target] = settings.get(source)
for key in ("enableInterrupt", "turnConfig"):
if key in settings:
data[key] = deepcopy(settings[key])
for key in (
"visionEnabled",
"visionModelResourceId",
"visionPromptMode",
"visionPrompt",
"knowledgeBaseId",
"knowledgeMode",
"knowledgeTopN",
"knowledgeScoreThreshold",
):
if key in settings:
data[key] = deepcopy(settings[key])
data["toolBindings"] = deepcopy(settings.get("toolBindings") or [])
data["builtInTools"] = deepcopy(settings.get("builtInTools") or {})
data["builtInTools"]["updateState"]["updates"] = [
_llm_rule(str(name), definitions.get(str(name)))
for name in allowed
if str(name) in definitions
]
global_prompt = str(settings.get("globalPrompt") or "").strip()
node_prompt = str(data.get("prompt") or "").strip()
if global_prompt:
data["prompt"] = "\n\n".join(part for part in (global_prompt, node_prompt) if part)
data["inheritGlobalConfig"] = False
def upgrade() -> None:
connection = op.get_bind()
op.add_column(
"assistants",
sa.Column(
"built_in_tools",
postgresql.JSONB(astext_type=sa.Text()),
server_default=sa.text("'{}'::jsonb"),
nullable=False,
),
)
op.add_column(
"assistant_tool_bindings",
sa.Column(
"config",
postgresql.JSONB(astext_type=sa.Text()),
server_default=sa.text("'{}'::jsonb"),
nullable=False,
),
)
tool_rows = connection.execute(
sa.text("SELECT id, type, description, definition, mcp_server_id FROM tools")
).mappings().all()
tools = {str(row["id"]): dict(row) for row in tool_rows}
template_inputs: dict[str, set[str]] = {}
resource_assignments: dict[str, dict] = {}
server_inputs: dict[str, set[str]] = {}
server_rows = connection.execute(
sa.text("SELECT id, url, config FROM mcp_servers")
).mappings().all()
for row in server_rows:
inputs: set[str] = set()
config = deepcopy(row["config"] or {})
url = _rewrite_templates(str(row["url"] or ""), inputs)
config["headers"] = _rewrite_templates(config.get("headers") or {}, inputs)
server_id = str(row["id"])
server_inputs[server_id] = inputs
connection.execute(
sa.text(
"UPDATE mcp_servers SET url=:url, config=CAST(:config AS jsonb) "
"WHERE id=:id"
),
{"id": server_id, "url": url, "config": json.dumps(config)},
)
for tool_id, tool in tools.items():
if tool.get("type") == "system":
continue
definition = deepcopy(tool.get("definition") or {})
config = definition.get("config") or {}
definition["config"] = config
resource_assignments[tool_id] = dict(
config.get("dynamic_variable_assignments") or {}
)
config.pop("dynamic_variable_assignments", None)
inputs: set[str] = set()
if tool.get("type") == "http":
for key in ("url", "headers", "body"):
config[key] = _rewrite_templates(config.get(key), inputs)
config["template_inputs"] = [
{"name": name, "type": "string", "location": "body", "required": True}
for name in sorted(inputs)
]
elif tool.get("type") == "mcp":
inputs.update(server_inputs.get(str(tool.get("mcp_server_id") or ""), set()))
config["template_inputs"] = [
{
"name": name,
"type": "string",
"location": "body",
"required": True,
}
for name in sorted(inputs)
]
definition["schema_version"] = 2
template_inputs[tool_id] = inputs
connection.execute(
sa.text("UPDATE tools SET definition=CAST(:definition AS jsonb) WHERE id=:id"),
{"id": tool_id, "definition": json.dumps(definition)},
)
tool["definition"] = definition
assistant_rows = connection.execute(
sa.text(
"SELECT id, type, runtime_mode, graph, startup, "
"dynamic_variable_definitions FROM assistants"
)
).mappings().all()
for assistant in assistant_rows:
assistant_id = str(assistant["id"])
definitions = assistant["dynamic_variable_definitions"] or {}
binding_rows = connection.execute(
sa.text(
"SELECT tool_id FROM assistant_tool_bindings "
"WHERE assistant_id=:assistant_id ORDER BY created_at, tool_id"
),
{"assistant_id": assistant_id},
).mappings().all()
selected_ids = [str(row["tool_id"]) for row in binding_rows]
system_by_kind: dict[str, Mapping] = {}
for tool_id in selected_ids:
tool = tools.get(tool_id) or {}
kind = ((tool.get("definition") or {}).get("config") or {}).get("kind")
if tool.get("type") == "system" and kind:
if kind in system_by_kind:
raise RuntimeError(
f"Assistant {assistant_id} 同时绑定多个 {kind} System Tool"
)
system_by_kind[str(kind)] = tool
builtins = (
{
kind: _builtin(kind, tool, definitions)
for kind, tool in system_by_kind.items()
}
if (
assistant["type"] == "prompt"
and assistant["runtime_mode"] != "realtime"
)
else {}
)
for tool_id in selected_ids:
tool = tools.get(tool_id) or {}
if tool.get("type") == "system":
continue
binding = _binding(
tool_id,
tool,
inherited_assignments=resource_assignments.get(tool_id),
)
_add_template_bindings(binding, template_inputs.get(tool_id, set()))
connection.execute(
sa.text(
"UPDATE assistant_tool_bindings SET config=CAST(:config AS jsonb) "
"WHERE assistant_id=:assistant_id AND tool_id=:tool_id"
),
{
"assistant_id": assistant_id,
"tool_id": tool_id,
"config": json.dumps(
{
"argument_bindings": binding["argumentBindings"],
"template_bindings": binding["templateBindings"],
"result_assignments": binding["resultAssignments"],
}
),
},
)
startup = deepcopy(assistant["startup"] or {})
for action in startup.get("actions") or []:
tool_id = str(action.pop("toolId", action.pop("tool_id", "")) or "")
arguments = action.pop("arguments", {}) or {}
tool = tools.get(tool_id) or {}
binding = _binding(
tool_id,
tool,
arguments=arguments,
inherited_assignments=resource_assignments.get(tool_id),
)
_add_template_bindings(binding, template_inputs.get(tool_id, set()))
action["toolBinding"] = binding
graph = deepcopy(assistant["graph"] or {})
if assistant["type"] == "workflow":
graph["specVersion"] = 4
settings = graph.setdefault("settings", {})
global_ids = list(settings.pop("toolIds", []) or [])
settings["toolBindings"] = []
settings["builtInTools"] = {}
for tool_id in global_ids:
tool_id = str(tool_id)
tool = tools.get(tool_id) or {}
if tool.get("type") == "system":
kind = ((tool.get("definition") or {}).get("config") or {}).get("kind")
if kind:
settings["builtInTools"][_builtin_key(str(kind))] = _builtin(
str(kind), tool, definitions, camel_case=True
)
continue
binding = _binding(
tool_id,
tool,
inherited_assignments=resource_assignments.get(tool_id),
)
_add_template_bindings(binding, template_inputs.get(tool_id, set()))
settings["toolBindings"].append(binding)
for node in graph.get("nodes") or []:
data = node.setdefault("data", {})
if node.get("type") == "agent":
inherits_global = bool(data.get("inheritGlobalConfig", True))
ids = list(data.pop("toolIds", []) or [])
allowed = list(data.pop("stateVariableNames", []) or [])
data["toolBindings"] = []
data["builtInTools"] = {}
if inherits_global:
if "updateState" in settings["builtInTools"]:
_freeze_inherited_agent_for_update_state(
data, settings, definitions, allowed
)
else:
for tool_id in ids:
tool_id = str(tool_id)
tool = tools.get(tool_id) or {}
if tool.get("type") == "system":
kind = ((tool.get("definition") or {}).get("config") or {}).get("kind")
if kind:
data["builtInTools"][_builtin_key(str(kind))] = _builtin(
str(kind),
tool,
definitions,
allowed if kind == "update_state" else None,
camel_case=True,
)
continue
binding = _binding(
tool_id,
tool,
inherited_assignments=resource_assignments.get(tool_id),
)
_add_template_bindings(
binding, template_inputs.get(tool_id, set())
)
data["toolBindings"].append(binding)
elif node.get("type") == "action":
tool_id = str(data.pop("toolId", "") or "")
arguments = data.pop("arguments", {}) or {}
mode = str(data.pop("resultAssignmentMode", "none") or "none")
overrides = data.pop("resultAssignments", {}) or {}
assignments = (
None if mode == "inherit" else overrides if mode == "override" else {}
)
binding = _binding(
tool_id,
tools.get(tool_id) or {},
arguments,
assignments,
resource_assignments.get(tool_id),
)
_add_template_bindings(binding, template_inputs.get(tool_id, set()))
data["toolBinding"] = binding
elif node.get("type") == "update_state":
assignments = data.pop("assignments", {}) or {}
data["updates"] = [
{
"variableName": str(name),
"expression": _state_expression(value),
}
for name, value in assignments.items()
]
connection.execute(
sa.text(
"UPDATE assistants SET built_in_tools=CAST(:builtins AS jsonb), "
"startup=CAST(:startup AS json), graph=CAST(:graph AS json) WHERE id=:id"
),
{
"id": assistant_id,
"builtins": json.dumps(builtins),
"startup": json.dumps(startup),
"graph": json.dumps(graph),
},
)
if assistant["type"] == "workflow":
connection.execute(
sa.text(
"UPDATE assistant_tool_bindings SET config='{}'::jsonb "
"WHERE assistant_id=:assistant_id"
),
{"assistant_id": assistant_id},
)
connection.execute(
sa.text(
"DELETE FROM assistant_tool_bindings "
"WHERE tool_id IN (SELECT id FROM tools WHERE type='system')"
)
)
connection.execute(
sa.text("DELETE FROM tools WHERE type='system'")
)
op.create_check_constraint(
"ck_tools_external_type",
"tools",
"type IN ('http', 'mcp', 'client')",
)
def downgrade() -> None:
raise RuntimeError(
"Tool Resource v4 是一次性边界迁移;请从部署前数据库备份回滚。"
)