mirror of
https://github.com/tiennm99/DocsGPT.git
synced 2026-10-05 02:13:24 +00:00
arc53-machine
4d001f6a20
A resource attached in a save now ignores any sponsor recorded for it before it was removed: the caller must be able to sponsor it and confirm. YAML import prunes the sponsors of resources it drops, for agents and workflow graphs. Sponsor details, with the resources' names, go only to people who may edit the agent. The workflow read returns the names of every node tool and source, so editors can remove the owner's private ones. An agent image is stored only once the save is known to go ahead.
1868 lines
75 KiB
Python
1868 lines
75 KiB
Python
"""Agent YAML export / import.
|
|
|
|
Agents are mostly references (sources, tools, prompt, models). Export
|
|
translates those into portable identifiers and never emits a secret;
|
|
import resolves them back against the importing user — matching existing
|
|
resources or creating them from the file plus user-supplied tokens.
|
|
|
|
Re-import is idempotent at every level: the agent is matched by
|
|
``metadata.id`` then ``metadata.slug``, tools by ``(type, name)``,
|
|
prompts by ``(name, content)``, custom models by
|
|
``(display_name, upstream_model_id, base_url)`` — so the same file
|
|
applied twice updates rather than duplicating.
|
|
|
|
Workflow agents additionally carry their graph under ``spec.workflow``.
|
|
The graph's agent nodes reference tools, sources and models by raw id in
|
|
their config; export rewrites those into the same portable identities the
|
|
top-level sections use (``tool-N`` keys into ``spec.tools``, source names,
|
|
custom-model display names), and import resolves them back and writes the
|
|
graph through the same validation gate as the workflow API.
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import copy
|
|
import re
|
|
from typing import Any, Optional
|
|
|
|
import yaml
|
|
from flask import current_app, jsonify, make_response, request
|
|
from flask_restx import Namespace, Resource
|
|
|
|
from docsgpt.agents.default_tools import (
|
|
default_tool_id,
|
|
is_synthesized_tool_id,
|
|
synthesize_tool_by_name,
|
|
synthesized_tool_name_for_id,
|
|
)
|
|
from docsgpt.api import api
|
|
from docsgpt.api.pat.rules import allowed_ids
|
|
from docsgpt.api.user.resource_access import AccessDenied, require
|
|
from docsgpt.core.model_utils import validate_model_id
|
|
from docsgpt.core.url_validation import SSRFError, validate_url
|
|
from docsgpt.security.safe_url import UnsafeUserUrlError, validate_user_base_url
|
|
from docsgpt.storage.db.base_repository import looks_like_uuid
|
|
from docsgpt.storage.db.repositories.agents import AgentsRepository
|
|
from docsgpt.storage.db.repositories.prompts import PromptsRepository
|
|
from docsgpt.storage.db.repositories.sources import SourcesRepository
|
|
from docsgpt.storage.db.repositories.user_custom_models import (
|
|
UserCustomModelsRepository,
|
|
)
|
|
from docsgpt.storage.db.repositories.user_tools import UserToolsRepository
|
|
from docsgpt.storage.db.repositories.workflow_edges import WorkflowEdgesRepository
|
|
from docsgpt.storage.db.repositories.workflow_nodes import WorkflowNodesRepository
|
|
from docsgpt.storage.db.repositories.workflows import WorkflowsRepository
|
|
from docsgpt.storage.db.session import db_readonly, db_session
|
|
|
|
|
|
API_VERSION = "docsgpt.arc53.com/v1"
|
|
KIND = "Agent"
|
|
|
|
# Import safety bounds (DoS): cap the document size and the number of
|
|
# child resources a single import may reference / create.
|
|
MAX_IMPORT_BYTES = 512 * 1024
|
|
MAX_LIST_ITEMS = 100
|
|
MAX_WORKFLOW_NODES = 200
|
|
MAX_WORKFLOW_EDGES = 400
|
|
|
|
|
|
class AgentImportError(Exception):
|
|
"""Raised when an agent YAML document is malformed or unsupported."""
|
|
|
|
|
|
class AgentExportError(Exception):
|
|
"""Raised when an agent cannot be exported as a document import would accept."""
|
|
|
|
|
|
class _SafeNoAliasLoader(yaml.SafeLoader):
|
|
"""SafeLoader that also rejects anchors/aliases (billion-laughs guard)."""
|
|
|
|
def compose_node(self, parent, index):
|
|
if self.check_event(yaml.events.AliasEvent):
|
|
raise AgentImportError("YAML anchors/aliases are not allowed")
|
|
return super().compose_node(parent, index)
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Slug helpers
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
def slugify(name: Optional[str]) -> str:
|
|
"""Turn an agent name into a url-safe slug; never empty."""
|
|
s = (name or "").strip().lower()
|
|
s = re.sub(r"[^a-z0-9]+", "-", s)
|
|
s = re.sub(r"-{2,}", "-", s).strip("-")
|
|
return s or "agent"
|
|
|
|
|
|
def _unique_slug(
|
|
repo: AgentsRepository,
|
|
user: str,
|
|
base: str,
|
|
*,
|
|
exclude_id: Optional[str] = None,
|
|
) -> str:
|
|
"""Return ``base`` (slugified) or the first ``-N`` variant free for ``user``."""
|
|
base = slugify(base)
|
|
candidate = base
|
|
n = 2
|
|
while True:
|
|
found = repo.find_by_slug(user, candidate)
|
|
if found is None or (exclude_id and str(found["id"]) == str(exclude_id)):
|
|
return candidate
|
|
candidate = f"{base}-{n}"
|
|
n += 1
|
|
|
|
|
|
def ensure_agent_slug(conn, agent: dict, user: str) -> str:
|
|
"""Return the agent's slug, lazily assigning and persisting one if absent."""
|
|
existing = agent.get("slug")
|
|
if existing:
|
|
return str(existing)
|
|
repo = AgentsRepository(conn)
|
|
slug = _unique_slug(repo, user, agent.get("name"), exclude_id=str(agent["id"]))
|
|
repo.update(str(agent["id"]), user, {"slug": slug})
|
|
return slug
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Tool-manager access (lazy — avoids importing the tool registry at module load)
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
def _tool_manager():
|
|
from docsgpt.api.user.tools.routes import tool_manager
|
|
|
|
return tool_manager
|
|
|
|
|
|
def _tool_instance(tool_type: str):
|
|
return _tool_manager().tools.get(tool_type)
|
|
|
|
|
|
def _secret_field_names(config_requirements: dict, *, required_only: bool = False) -> list:
|
|
out = []
|
|
for key, spec in (config_requirements or {}).items():
|
|
if not isinstance(spec, dict) or not spec.get("secret"):
|
|
continue
|
|
if required_only and not spec.get("required"):
|
|
continue
|
|
out.append(key)
|
|
return out
|
|
|
|
|
|
def _live_requires_secrets(tool_type: str) -> list:
|
|
inst = _tool_instance(tool_type)
|
|
if inst is None:
|
|
return []
|
|
return _secret_field_names(inst.get_config_requirements() or {})
|
|
|
|
|
|
def _safe_export_config(stored_config: dict, config_requirements: dict) -> dict:
|
|
"""Return only config values provably non-secret (allowlist).
|
|
|
|
Export must never leak a credential. A key is emitted ONLY if
|
|
``config_requirements`` declares it and does not mark it secret.
|
|
Anything not described by requirements — e.g. MCP free-form ``headers``,
|
|
``url`` query strings, ``server_url`` that routinely embed tokens, or
|
|
tools with empty ``config_requirements`` — is dropped. Reuse on import
|
|
matches the existing tool; create re-collects config from the user. This
|
|
blocklist-free approach can't be defeated by a stale or empty
|
|
``config_requirements``. (``api_tool`` is the one exception: its
|
|
user-authored ``config["actions"]`` travel via
|
|
``_export_api_tool_actions``, redacted under the same posture.)
|
|
"""
|
|
requirements = config_requirements or {}
|
|
safe: dict[str, Any] = {}
|
|
for key, value in (stored_config or {}).items():
|
|
spec = requirements.get(key)
|
|
if isinstance(spec, dict) and not spec.get("secret"):
|
|
safe[key] = value
|
|
return safe
|
|
|
|
|
|
# Param-type sections an action may carry. Fixed ``value``s are exported only
|
|
# for ``parameters`` (declared function args, treated as non-secret throughout:
|
|
# returned to the client, shown in the actions editor). The free-form HTTP
|
|
# types routinely hold credentials (an ``Authorization`` header fixed as a
|
|
# value), so for those only the ``filled_by_llm`` flag travels — same posture
|
|
# as ``_safe_export_config`` dropping free-form config.
|
|
#
|
|
# Scope note: ``actions``-column rows come from ``transform_actions`` over live
|
|
# tool metadata, which today only ever emits ``parameters`` sections — the
|
|
# HTTP types appear only in ``api_tool``'s ``config["actions"]``, which travel
|
|
# through ``_export_api_tool_actions`` / ``_sanitize_api_tool_actions`` below.
|
|
# The full tuple is kept here so a hand-edited file naming those sections is
|
|
# handled (flags applied, values refused) rather than silently ignored.
|
|
_ACTION_PARAM_TYPES = ("parameters", "query_params", "headers", "body")
|
|
_ACTION_VALUE_SAFE_TYPES = frozenset({"parameters"})
|
|
|
|
# Action keys copied verbatim by ``_export_api_tool_actions`` — everything
|
|
# except ``url`` (query-string-stripped) and the param sections (redacted).
|
|
_API_ACTION_SAFE_KEYS = (
|
|
"name",
|
|
"description",
|
|
"method",
|
|
"active",
|
|
"require_approval",
|
|
"body_content_type",
|
|
"body_encoding_rules",
|
|
)
|
|
|
|
|
|
def _export_api_tool_actions(stored_config: dict) -> dict:
|
|
"""Redacted copy of an ``api_tool``'s ``config["actions"]``.
|
|
|
|
``api_tool`` is the one tool type whose actions are user-authored data in
|
|
``config["actions"]`` (a dict keyed by action name) rather than the
|
|
``actions`` column (the executor reads exactly that — see ToolExecutor).
|
|
The action *shape* — endpoint, method, param schemas, ``filled_by_llm``
|
|
flags, fixed ``parameters`` values — is what makes the tool work, so it
|
|
travels. Fixed values in the free-form HTTP sections and URL query
|
|
strings routinely hold credentials, so those are blanked — same posture
|
|
as ``_safe_export_config``.
|
|
"""
|
|
actions = stored_config.get("actions")
|
|
if not isinstance(actions, dict):
|
|
return {}
|
|
out: dict[str, Any] = {}
|
|
for name, action in actions.items():
|
|
if not isinstance(action, dict):
|
|
continue
|
|
entry: dict[str, Any] = {
|
|
k: copy.deepcopy(action[k]) for k in _API_ACTION_SAFE_KEYS if k in action
|
|
}
|
|
entry["url"] = str(action.get("url") or "").split("?", 1)[0]
|
|
for param_type in _ACTION_PARAM_TYPES:
|
|
section = action.get(param_type)
|
|
if not isinstance(section, dict):
|
|
continue
|
|
section = copy.deepcopy(section)
|
|
props = section.get("properties")
|
|
if param_type not in _ACTION_VALUE_SAFE_TYPES and isinstance(props, dict):
|
|
for details in props.values():
|
|
if isinstance(details, dict) and details.get("value"):
|
|
details["value"] = ""
|
|
entry[param_type] = section
|
|
out[str(name)] = entry
|
|
return out
|
|
|
|
|
|
def _sanitize_api_tool_actions(raw, tool_label: str, warnings: list) -> dict:
|
|
"""Vet a hand-editable ``config["actions"]`` block before storing it.
|
|
|
|
Unlike column actions — whose schema comes from live tool code, so the
|
|
YAML may contribute flags and scalar values only — ``api_tool`` actions
|
|
ARE user-authored schema, and the whole definition is accepted. Each
|
|
action's URL goes through the same SSRF gate as the tool-config routes
|
|
(``_validate_tool_urls``); an action with a missing or unsafe URL is
|
|
dropped with a warning rather than failing the import.
|
|
"""
|
|
if not isinstance(raw, dict):
|
|
return {}
|
|
if len(raw) > MAX_LIST_ITEMS:
|
|
raise AgentImportError(f"Too many api_tool actions (max {MAX_LIST_ITEMS})")
|
|
out: dict[str, Any] = {}
|
|
for name, action in raw.items():
|
|
if not isinstance(action, dict):
|
|
continue
|
|
url = str(action.get("url") or "").strip()
|
|
if not url:
|
|
warnings.append(
|
|
f"Tool '{tool_label}' action '{name}' has no URL; dropped"
|
|
)
|
|
continue
|
|
try:
|
|
validate_url(url)
|
|
except SSRFError:
|
|
warnings.append(
|
|
f"Tool '{tool_label}' action '{name}' has an unsafe URL; dropped"
|
|
)
|
|
continue
|
|
# Export blanks fixed values in the free-form HTTP sections (they
|
|
# routinely hold credentials); a param left pinned with an empty
|
|
# value is silently omitted from the request at run time, so
|
|
# surface it now instead.
|
|
pinned_empty = [
|
|
f"{param_type}.{pname}"
|
|
for param_type in _ACTION_PARAM_TYPES
|
|
if param_type not in _ACTION_VALUE_SAFE_TYPES
|
|
and isinstance(action.get(param_type), dict)
|
|
for pname, details in (action[param_type].get("properties") or {}).items()
|
|
if isinstance(details, dict)
|
|
and details.get("filled_by_llm") is False
|
|
and not details.get("value")
|
|
]
|
|
if pinned_empty:
|
|
warnings.append(
|
|
f"Tool '{tool_label}' action '{name}': fixed values for "
|
|
f"{', '.join(pinned_empty)} don't travel in the file; set them "
|
|
"in the tool's settings or the request will omit them"
|
|
)
|
|
out[str(name)] = action
|
|
return out
|
|
|
|
|
|
def _export_action_overrides(actions) -> list:
|
|
"""Compact per-action customizations that deviate from factory defaults.
|
|
|
|
``transform_actions`` initializes every action ``active`` with every param
|
|
``filled_by_llm: true, value: ""`` — anything else is a user customization
|
|
(a fixed ``chat_id``, a disabled action) that must survive export, or the
|
|
imported tool silently reverts to "let the LLM fill everything".
|
|
"""
|
|
out = []
|
|
for action in actions or []:
|
|
if not isinstance(action, dict) or not action.get("name"):
|
|
continue
|
|
override: dict[str, Any] = {"name": action["name"]}
|
|
customized = False
|
|
if action.get("active", True) is False:
|
|
override["active"] = False
|
|
customized = True
|
|
for param_type in _ACTION_PARAM_TYPES:
|
|
section = action.get(param_type)
|
|
props = section.get("properties") if isinstance(section, dict) else None
|
|
if not isinstance(props, dict):
|
|
continue
|
|
for pname, details in props.items():
|
|
if not isinstance(details, dict):
|
|
continue
|
|
filled = details.get("filled_by_llm", True)
|
|
value = details.get("value", "")
|
|
if filled and not value:
|
|
continue
|
|
entry: dict[str, Any] = {"filled_by_llm": bool(filled)}
|
|
if param_type in _ACTION_VALUE_SAFE_TYPES:
|
|
entry["value"] = value
|
|
override.setdefault(param_type, {})[pname] = entry
|
|
customized = True
|
|
if customized:
|
|
out.append(override)
|
|
return out
|
|
|
|
|
|
def _apply_action_overrides(base_actions: list, overrides, tool_type: str, warnings: list) -> list:
|
|
"""Overlay exported customizations onto freshly transformed live metadata.
|
|
|
|
Matched strictly by (action name, param-type, param name): the YAML
|
|
contributes flags and scalar values only, never schema — a hand-edited
|
|
file can't inject a doctored action definition — and the merge survives
|
|
tool-version drift. Unknown names are skipped with a warning. Values are
|
|
accepted only for the ``parameters`` type, mirroring export.
|
|
"""
|
|
if not isinstance(overrides, list) or not overrides:
|
|
return base_actions
|
|
by_name = {a.get("name"): a for a in base_actions if isinstance(a, dict)}
|
|
for override in overrides:
|
|
if not isinstance(override, dict):
|
|
continue
|
|
name = override.get("name") or ""
|
|
action = by_name.get(name)
|
|
if action is None:
|
|
warnings.append(
|
|
f"Tool '{tool_type}' action '{name}' no longer exists; customization skipped"
|
|
)
|
|
continue
|
|
if override.get("active") is False:
|
|
action["active"] = False
|
|
for param_type in _ACTION_PARAM_TYPES:
|
|
override_props = override.get(param_type)
|
|
if not isinstance(override_props, dict):
|
|
continue
|
|
section = action.get(param_type)
|
|
props = section.get("properties") if isinstance(section, dict) else {}
|
|
for pname, entry in override_props.items():
|
|
details = props.get(pname) if isinstance(props, dict) else None
|
|
if not isinstance(details, dict) or not isinstance(entry, dict):
|
|
warnings.append(
|
|
f"Tool '{tool_type}' action '{name}' has no parameter "
|
|
f"'{pname}'; customization skipped"
|
|
)
|
|
continue
|
|
details["filled_by_llm"] = bool(entry.get("filled_by_llm", True))
|
|
if param_type in _ACTION_VALUE_SAFE_TYPES:
|
|
value = entry.get("value", "")
|
|
if isinstance(value, (str, int, float, bool)) or value is None:
|
|
details["value"] = "" if value is None else value
|
|
elif not details["filled_by_llm"] and not details.get("value"):
|
|
# Pinned but valueless: the executor omits the param from
|
|
# the LLM schema AND skips the empty fixed value, so the
|
|
# request silently goes out without it. Values for the
|
|
# free-form HTTP types never travel in the file, so this
|
|
# can only be repaired in the tool's settings.
|
|
warnings.append(
|
|
f"Tool '{tool_type}' action '{name}': fixed value for "
|
|
f"{param_type} '{pname}' doesn't travel in the file; set "
|
|
"it in the tool's settings or the request will omit it"
|
|
)
|
|
return base_actions
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Export
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
def _serialize_prompt(conn, agent: dict, user: str):
|
|
prompt_id = agent.get("prompt_id")
|
|
if not prompt_id:
|
|
return "default"
|
|
row = PromptsRepository(conn).get(str(prompt_id), user)
|
|
if not row:
|
|
return "default"
|
|
return {"name": row.get("name") or "", "content": row.get("content") or ""}
|
|
|
|
|
|
def _serialize_sources(conn, agent: dict, user: str) -> list:
|
|
ids: list[str] = []
|
|
if agent.get("source_id"):
|
|
ids.append(str(agent["source_id"]))
|
|
for sid in agent.get("extra_source_ids") or []:
|
|
s = str(sid)
|
|
if s and s not in ids:
|
|
ids.append(s)
|
|
repo = SourcesRepository(conn)
|
|
out = []
|
|
for sid in ids:
|
|
row = repo.get(sid, user)
|
|
if not row:
|
|
continue
|
|
out.append(
|
|
{
|
|
"name": row.get("name") or "",
|
|
"type": row.get("type") or "",
|
|
"ref": sid,
|
|
}
|
|
)
|
|
return out
|
|
|
|
|
|
def _serialize_tools(conn, agent: dict, user: str) -> list:
|
|
repo = UserToolsRepository(conn)
|
|
out = []
|
|
for tid in agent.get("tools") or []:
|
|
tid_str = str(tid)
|
|
if is_synthesized_tool_id(tid_str):
|
|
out.append(
|
|
{
|
|
"type": synthesized_tool_name_for_id(tid_str),
|
|
"builtin": True,
|
|
"ref": tid_str,
|
|
}
|
|
)
|
|
continue
|
|
row = repo.get_any(tid_str, user)
|
|
if not row:
|
|
continue
|
|
requirements = row.get("config_requirements") or {}
|
|
requires_secrets = _secret_field_names(requirements)
|
|
config = _safe_export_config(row.get("config") or {}, requirements)
|
|
if (row.get("name") or "") == "api_tool":
|
|
# api_tool's actions live in config, not the actions column — an
|
|
# export without them imports as a dead tool.
|
|
exported_actions = _export_api_tool_actions(row.get("config") or {})
|
|
if exported_actions:
|
|
config["actions"] = exported_actions
|
|
entry: dict[str, Any] = {
|
|
"type": row.get("name") or "",
|
|
# Raw custom_name (may be "") so import matching is symmetric and
|
|
# idempotent — see _find_user_tool.
|
|
"name": row.get("custom_name") or "",
|
|
"display_name": row.get("display_name") or "",
|
|
"description": row.get("description") or "",
|
|
"config": config,
|
|
"ref": tid_str,
|
|
}
|
|
action_overrides = _export_action_overrides(row.get("actions"))
|
|
if action_overrides:
|
|
entry["actions"] = action_overrides
|
|
if requires_secrets:
|
|
entry["requires_secrets"] = requires_secrets
|
|
out.append(entry)
|
|
return out
|
|
|
|
|
|
def _serialize_models(conn, agent: dict, user: str) -> dict:
|
|
repo = UserCustomModelsRepository(conn)
|
|
|
|
def describe(model_id: str):
|
|
mid = str(model_id)
|
|
if looks_like_uuid(mid):
|
|
row = repo.get(mid, user)
|
|
if not row:
|
|
return None
|
|
return {
|
|
"type": "custom",
|
|
"display_name": row.get("display_name") or "",
|
|
"upstream_model_id": row.get("upstream_model_id") or "",
|
|
"base_url": row.get("base_url") or "",
|
|
"capabilities": row.get("capabilities") or {},
|
|
"requires_secrets": ["api_key"],
|
|
"ref": mid,
|
|
}
|
|
return mid
|
|
|
|
available = []
|
|
for m in agent.get("models") or []:
|
|
described = describe(m)
|
|
if described is not None:
|
|
available.append(described)
|
|
|
|
default_value = ""
|
|
raw_default = agent.get("default_model_id") or ""
|
|
if raw_default:
|
|
described = describe(raw_default)
|
|
if isinstance(described, dict):
|
|
default_value = described["display_name"]
|
|
elif isinstance(described, str):
|
|
default_value = described
|
|
return {"default": default_value, "available": available}
|
|
|
|
|
|
def _effective_node_config(config: dict) -> dict:
|
|
"""Return the dict a node's settings actually live in.
|
|
|
|
The builder writes some node configs nested under ``config`` while other
|
|
payloads keep them flat; the engine reads
|
|
``node.config.get("config", node.config)``. Mirror that here so reference
|
|
rewriting touches the dict the engine will read.
|
|
"""
|
|
nested = config.get("config")
|
|
return nested if isinstance(nested, dict) else config
|
|
|
|
|
|
def _strip_shadowed_ref_keys(config: dict) -> None:
|
|
"""Drop flat ``tools``/``sources``/``model_id`` shadowed by a nested config.
|
|
|
|
When a node stores its settings under ``config["config"]``, the engine
|
|
reads only the nested dict — flat leftovers of those keys are inert, but
|
|
on export they would leak the exporter's raw ids into the file, and on
|
|
import they would store a hand-edited file's raw ids verbatim.
|
|
"""
|
|
if isinstance(config.get("config"), dict):
|
|
for key in ("tools", "sources", "model_id"):
|
|
config.pop(key, None)
|
|
|
|
|
|
def _serialize_workflow(conn, agent: dict, user: str):
|
|
"""Build the portable ``spec.workflow`` block plus its reference sections.
|
|
|
|
Agent nodes hold raw tool/source/model ids; those are collected (deduped,
|
|
in first-appearance order) into the same top-level ``tools`` / ``sources``
|
|
/ ``model`` sections a classic export uses, and the node configs are
|
|
rewritten to portable keys: ``tool-N``, the source's name, a custom
|
|
model's display name. Returns ``(workflow_block, tools, sources, model)``;
|
|
the block is None when the agent has no resolvable workflow.
|
|
"""
|
|
empty_model = {"default": "", "available": []}
|
|
workflow_id = agent.get("workflow_id")
|
|
if not workflow_id:
|
|
return None, [], [], empty_model
|
|
row = WorkflowsRepository(conn).get(str(workflow_id), user)
|
|
if row is None:
|
|
return None, [], [], empty_model
|
|
# Lazy import, same as _apply_workflow — keeps route modules decoupled at load.
|
|
from docsgpt.api.user.workflows.routes import get_workflow_graph_version
|
|
|
|
version = get_workflow_graph_version(row)
|
|
node_rows = WorkflowNodesRepository(conn).find_by_version(str(row["id"]), version)
|
|
edge_rows = WorkflowEdgesRepository(conn).find_by_version(str(row["id"]), version)
|
|
|
|
# Pass 1: collect referenced ids across agent nodes, deduped and ordered.
|
|
tool_ids: list[str] = []
|
|
source_ids: list[str] = []
|
|
model_ids: list[str] = []
|
|
for n in node_rows:
|
|
if n.get("node_type") != "agent":
|
|
continue
|
|
cfg = _effective_node_config(n.get("config") or {})
|
|
for tid in cfg.get("tools") or []:
|
|
if str(tid) and str(tid) not in tool_ids:
|
|
tool_ids.append(str(tid))
|
|
for sid in cfg.get("sources") or []:
|
|
if str(sid) and str(sid) not in source_ids:
|
|
source_ids.append(str(sid))
|
|
mid = cfg.get("model_id")
|
|
if mid and str(mid) not in model_ids:
|
|
model_ids.append(str(mid))
|
|
|
|
# Serialize the references through the classic per-resource serializers so
|
|
# secret redaction and identity rules stay in one place. Unresolvable ids
|
|
# (dangling rows) are dropped from both the section and the node config;
|
|
# every emitted entry carries its ``ref``, which keys the maps below.
|
|
tools: list = []
|
|
tool_key_by_id: dict[str, str] = {}
|
|
for entry in _serialize_tools(conn, {"tools": tool_ids}, user):
|
|
tool_key_by_id[str(entry["ref"])] = _tool_key(len(tools))
|
|
tools.append(entry)
|
|
|
|
sources: list = []
|
|
source_name_by_id: dict[str, str] = {}
|
|
for entry in _serialize_sources(conn, {"extra_source_ids": source_ids}, user):
|
|
source_name_by_id[str(entry["ref"])] = entry["name"]
|
|
sources.append(entry)
|
|
|
|
model = _serialize_models(conn, {"models": model_ids, "default_model_id": ""}, user)
|
|
model_name_by_id: dict[str, str] = {}
|
|
for entry in model["available"]:
|
|
if isinstance(entry, dict):
|
|
model_name_by_id[str(entry.get("ref") or "")] = entry.get("display_name") or ""
|
|
elif isinstance(entry, str):
|
|
model_name_by_id[entry] = entry
|
|
|
|
# ``parse_agent_yaml`` bounds these sections at import (DoS cap), so a
|
|
# graph referencing more would export a file that import then rejects —
|
|
# refuse here with a clear message instead.
|
|
if (
|
|
len(tools) > MAX_LIST_ITEMS
|
|
or len(sources) > MAX_LIST_ITEMS
|
|
or len(model["available"]) > MAX_LIST_ITEMS
|
|
):
|
|
raise AgentExportError(
|
|
"This workflow references more than "
|
|
f"{MAX_LIST_ITEMS} tools, sources, or models and cannot be exported"
|
|
)
|
|
|
|
# Pass 2: emit nodes/edges with agent node configs rewritten to the
|
|
# portable keys. (Node settings sit under ``config`` here; the workflow
|
|
# API's wire shape calls that field ``data`` — import converts.)
|
|
nodes_out = []
|
|
for n in node_rows:
|
|
config = copy.deepcopy(n.get("config") or {})
|
|
if n.get("node_type") == "agent":
|
|
cfg = _effective_node_config(config)
|
|
cfg["tools"] = [
|
|
tool_key_by_id[str(t)]
|
|
for t in cfg.get("tools") or []
|
|
if str(t) in tool_key_by_id
|
|
]
|
|
cfg["sources"] = [
|
|
source_name_by_id[str(s)]
|
|
for s in cfg.get("sources") or []
|
|
if str(s) in source_name_by_id
|
|
]
|
|
if cfg.get("model_id"):
|
|
cfg["model_id"] = model_name_by_id.get(str(cfg["model_id"]))
|
|
_strip_shadowed_ref_keys(config)
|
|
nodes_out.append(
|
|
{
|
|
"id": n["node_id"],
|
|
"type": n["node_type"],
|
|
"title": n.get("title") or "",
|
|
"description": n.get("description") or "",
|
|
"position": n.get("position") or {"x": 0, "y": 0},
|
|
"config": config,
|
|
}
|
|
)
|
|
edges_out = [
|
|
{
|
|
"id": e["edge_id"],
|
|
"source": e.get("source_id"),
|
|
"target": e.get("target_id"),
|
|
"sourceHandle": e.get("source_handle"),
|
|
"targetHandle": e.get("target_handle"),
|
|
}
|
|
for e in edge_rows
|
|
]
|
|
block = {
|
|
"name": row.get("name") or "",
|
|
"description": row.get("description") or "",
|
|
"nodes": nodes_out,
|
|
"edges": edges_out,
|
|
}
|
|
return block, tools, sources, model
|
|
|
|
|
|
def _import_config(spec: dict) -> dict:
|
|
"""Validate a spec's ``config`` through the same gate as the API.
|
|
|
|
A YAML is hand-editable, so it must not be a way to install a control the
|
|
write path would have rejected. An invalid block is dropped rather than
|
|
failing the whole import, and the caller surfaces it as a warning.
|
|
"""
|
|
from docsgpt.api.user.agents.routes import normalize_agent_config
|
|
|
|
try:
|
|
return normalize_agent_config(spec.get("config")) or {}
|
|
except ValueError:
|
|
current_app.logger.warning(
|
|
"Dropping invalid guardrails config during agent import"
|
|
)
|
|
return {}
|
|
|
|
|
|
def serialize_agent(conn, agent: dict, user: str) -> dict:
|
|
"""Build the portable export document for an agent row."""
|
|
is_workflow = (agent.get("agent_type") or "") == "workflow"
|
|
if is_workflow:
|
|
# A workflow agent's behavior lives in its graph: the agent-level
|
|
# prompt/source/tool/model slots are unused, so the top-level sections
|
|
# instead carry the references collected from the graph's agent nodes.
|
|
workflow_block, tools, sources, model = _serialize_workflow(conn, agent, user)
|
|
prompt = "default"
|
|
else:
|
|
tools = _serialize_tools(conn, agent, user)
|
|
sources = _serialize_sources(conn, agent, user)
|
|
model = _serialize_models(conn, agent, user)
|
|
prompt = _serialize_prompt(conn, agent, user)
|
|
spec = {
|
|
"name": agent.get("name") or "",
|
|
"description": agent.get("description") or "",
|
|
"agent_type": agent.get("agent_type") or "classic",
|
|
"retriever": agent.get("retriever") or "classic",
|
|
"chunks": int(agent["chunks"]) if agent.get("chunks") is not None else None,
|
|
"prompt": prompt,
|
|
"model": model,
|
|
"sources": sources,
|
|
"tools": tools,
|
|
"limits": {
|
|
"limited_token_mode": bool(agent.get("limited_token_mode", False)),
|
|
"token_limit": agent.get("token_limit"),
|
|
"limited_request_mode": bool(agent.get("limited_request_mode", False)),
|
|
"request_limit": agent.get("request_limit"),
|
|
},
|
|
"json_schema": agent.get("json_schema"),
|
|
"allow_system_prompt_override": bool(agent.get("allow_system_prompt_override", False)),
|
|
"config": agent.get("config") or {},
|
|
}
|
|
if is_workflow:
|
|
spec["workflow"] = workflow_block
|
|
return {
|
|
"apiVersion": API_VERSION,
|
|
"kind": KIND,
|
|
"metadata": {"id": str(agent["id"]), "slug": agent.get("slug") or ""},
|
|
"spec": spec,
|
|
}
|
|
|
|
|
|
class _AgentYamlDumper(yaml.SafeDumper):
|
|
"""SafeDumper that renders multi-line strings as block scalars."""
|
|
|
|
|
|
def _str_representer(dumper, data):
|
|
style = "|" if "\n" in data else None
|
|
return dumper.represent_scalar("tag:yaml.org,2002:str", data, style=style)
|
|
|
|
|
|
_AgentYamlDumper.add_representer(str, _str_representer)
|
|
|
|
|
|
def agent_to_yaml(export: dict) -> str:
|
|
"""Serialize an export document to YAML, preserving key order."""
|
|
return yaml.dump(
|
|
export,
|
|
Dumper=_AgentYamlDumper,
|
|
sort_keys=False,
|
|
default_flow_style=False,
|
|
allow_unicode=True,
|
|
)
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Parse
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
def _validate_workflow_shape(spec: dict) -> None:
|
|
"""Shallow-validate ``spec.workflow`` (deep graph validation runs at apply).
|
|
|
|
A missing key or explicit null is allowed — a draft workflow agent may have
|
|
no graph yet; only a present block must be well-formed and bounded.
|
|
"""
|
|
workflow = spec.get("workflow")
|
|
if workflow is None:
|
|
return
|
|
if not isinstance(workflow, dict):
|
|
raise AgentImportError("spec.workflow must be a mapping")
|
|
nodes = workflow.get("nodes")
|
|
edges = workflow.get("edges")
|
|
if not isinstance(nodes, list) or not isinstance(edges, list):
|
|
raise AgentImportError("spec.workflow.nodes and spec.workflow.edges must be lists")
|
|
if len(nodes) > MAX_WORKFLOW_NODES:
|
|
raise AgentImportError(f"Too many workflow nodes (max {MAX_WORKFLOW_NODES})")
|
|
if len(edges) > MAX_WORKFLOW_EDGES:
|
|
raise AgentImportError(f"Too many workflow edges (max {MAX_WORKFLOW_EDGES})")
|
|
for node in nodes:
|
|
if not isinstance(node, dict) or not node.get("id") or not node.get("type"):
|
|
raise AgentImportError("Every workflow node needs an 'id' and a 'type'")
|
|
for edge in edges:
|
|
if not isinstance(edge, dict) or not edge.get("id"):
|
|
raise AgentImportError("Every workflow edge needs an 'id'")
|
|
|
|
|
|
def parse_agent_yaml(text: str) -> dict:
|
|
"""Parse and shallow-validate an agent YAML document."""
|
|
if text and len(text) > MAX_IMPORT_BYTES:
|
|
raise AgentImportError("Import document too large")
|
|
try:
|
|
doc = yaml.load(text, Loader=_SafeNoAliasLoader)
|
|
except yaml.YAMLError as exc:
|
|
raise AgentImportError(f"Invalid YAML: {exc}") from exc
|
|
if not isinstance(doc, dict):
|
|
raise AgentImportError("Top-level YAML must be a mapping")
|
|
if doc.get("kind") != KIND:
|
|
raise AgentImportError(f"Unsupported kind {doc.get('kind')!r}; expected {KIND!r}")
|
|
if not str(doc.get("apiVersion") or "").startswith("docsgpt."):
|
|
raise AgentImportError(f"Unsupported apiVersion {doc.get('apiVersion')!r}")
|
|
spec = doc.get("spec")
|
|
if not isinstance(spec, dict):
|
|
raise AgentImportError("Missing or invalid 'spec'")
|
|
if not spec.get("name"):
|
|
raise AgentImportError("spec.name is required")
|
|
if spec.get("agent_type") == "workflow":
|
|
_validate_workflow_shape(spec)
|
|
for key in ("sources", "tools"):
|
|
value = spec.get(key)
|
|
if isinstance(value, list) and len(value) > MAX_LIST_ITEMS:
|
|
raise AgentImportError(f"Too many {key} (max {MAX_LIST_ITEMS})")
|
|
available = (spec.get("model") or {}).get("available")
|
|
if isinstance(available, list) and len(available) > MAX_LIST_ITEMS:
|
|
raise AgentImportError(f"Too many models (max {MAX_LIST_ITEMS})")
|
|
return doc
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Reference resolution (shared by plan + apply)
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
def _resolve_target(conn, user: str, metadata: dict) -> dict:
|
|
repo = AgentsRepository(conn)
|
|
metadata = metadata or {}
|
|
agent_id = metadata.get("id")
|
|
if agent_id and looks_like_uuid(str(agent_id)):
|
|
row = repo.get(str(agent_id), user)
|
|
if row:
|
|
return {
|
|
"action": "update",
|
|
"agent_id": str(row["id"]),
|
|
"matched_by": "id",
|
|
"status": row.get("status"),
|
|
}
|
|
slug = metadata.get("slug")
|
|
if slug:
|
|
row = repo.find_by_slug(user, str(slug))
|
|
if row:
|
|
return {
|
|
"action": "update",
|
|
"agent_id": str(row["id"]),
|
|
"matched_by": "slug",
|
|
"status": row.get("status"),
|
|
}
|
|
return {"action": "create", "agent_id": None, "matched_by": None, "status": None}
|
|
|
|
|
|
def _find_user_tool(user_tools: list, tool_type: str, custom_name: str):
|
|
"""Match a stored tool by ``(type, custom_name)`` for idempotent reuse.
|
|
|
|
Compares against the stored ``custom_name`` directly (export emits the
|
|
raw value), so a named tool matches exactly and an unnamed one (``""``)
|
|
reuses the first same-type tool rather than spawning a duplicate.
|
|
"""
|
|
target = custom_name or ""
|
|
for row in user_tools:
|
|
if (row.get("name") or "") != tool_type:
|
|
continue
|
|
if (row.get("custom_name") or "") == target:
|
|
return row
|
|
return None
|
|
|
|
|
|
def _find_custom_model(customs: list, spec_model: dict):
|
|
for row in customs:
|
|
if (
|
|
(row.get("display_name") or "") == (spec_model.get("display_name") or "")
|
|
and (row.get("upstream_model_id") or "") == (spec_model.get("upstream_model_id") or "")
|
|
and (row.get("base_url") or "") == (spec_model.get("base_url") or "")
|
|
):
|
|
return row
|
|
return None
|
|
|
|
|
|
def _tool_key(index: int) -> str:
|
|
return f"tool-{index}"
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Plan (dry run)
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
def _plan_workflow_removal(conn, user: str, target: dict) -> Optional[dict]:
|
|
"""Describe what an explicit ``workflow: null`` would destroy, if anything.
|
|
|
|
Returns a ``delete`` block sized with the current graph, or None when the
|
|
workflow survives — the target is new, published (``_apply_workflow``
|
|
keeps a published agent's graph), has no workflow, or another agent still
|
|
references it.
|
|
"""
|
|
if target.get("action") != "update" or not target.get("agent_id"):
|
|
return None
|
|
if (target.get("status") or "") == "published":
|
|
return None
|
|
agent_row = AgentsRepository(conn).get(str(target["agent_id"]), user) or {}
|
|
existing_id = agent_row.get("workflow_id")
|
|
if not existing_id:
|
|
return None
|
|
row = WorkflowsRepository(conn).get(str(existing_id), user)
|
|
if row is None:
|
|
return None
|
|
if AgentsRepository(conn).count_by_workflow(str(existing_id), user) > 1:
|
|
return None
|
|
# Lazy import, same as _apply_workflow — keeps route modules decoupled at load.
|
|
from docsgpt.api.user.workflows.routes import get_workflow_graph_version
|
|
|
|
version = get_workflow_graph_version(row)
|
|
return {
|
|
"action": "delete",
|
|
"nodes": len(WorkflowNodesRepository(conn).find_by_version(str(row["id"]), version)),
|
|
"edges": len(WorkflowEdgesRepository(conn).find_by_version(str(row["id"]), version)),
|
|
}
|
|
|
|
|
|
def plan_import(conn, user: str, doc: dict) -> dict:
|
|
"""Resolve every reference without writing; returns a resolution report."""
|
|
spec = doc["spec"]
|
|
target = _resolve_target(conn, user, doc.get("metadata") or {})
|
|
|
|
sources_repo = SourcesRepository(conn)
|
|
sources = []
|
|
for src in spec.get("sources") or []:
|
|
name = src.get("name") or ""
|
|
match = sources_repo.find_by_name(user, name)
|
|
sources.append(
|
|
{
|
|
"name": name,
|
|
"type": src.get("type") or "",
|
|
"status": "matched" if match else "missing",
|
|
"target_id": str(match["id"]) if match else None,
|
|
}
|
|
)
|
|
|
|
user_tools = UserToolsRepository(conn).list_for_user(user)
|
|
tools = []
|
|
for index, tool in enumerate(spec.get("tools") or []):
|
|
tool_type = tool.get("type") or ""
|
|
if tool.get("builtin"):
|
|
available = synthesize_tool_by_name(tool_type) is not None
|
|
tools.append(
|
|
{
|
|
"key": _tool_key(index),
|
|
"type": tool_type,
|
|
"builtin": True,
|
|
"status": "builtin" if available else "unavailable",
|
|
"target_id": default_tool_id(tool_type) if available else None,
|
|
}
|
|
)
|
|
continue
|
|
custom_name = tool.get("name") or ""
|
|
match = _find_user_tool(user_tools, tool_type, custom_name)
|
|
if match:
|
|
tools.append(
|
|
{
|
|
"key": _tool_key(index),
|
|
"type": tool_type,
|
|
"name": custom_name,
|
|
"status": "reuse",
|
|
"target_id": str(match["id"]),
|
|
}
|
|
)
|
|
continue
|
|
available = _tool_instance(tool_type) is not None
|
|
tools.append(
|
|
{
|
|
"key": _tool_key(index),
|
|
"type": tool_type,
|
|
"name": custom_name,
|
|
"status": "create" if available else "unavailable",
|
|
"requires_secrets": tool.get("requires_secrets") or _live_requires_secrets(tool_type),
|
|
}
|
|
)
|
|
|
|
prompt_spec = spec.get("prompt")
|
|
if isinstance(prompt_spec, dict):
|
|
existing = PromptsRepository(conn).find(
|
|
user, prompt_spec.get("name") or "", prompt_spec.get("content") or ""
|
|
)
|
|
prompt = {"status": "reuse" if existing else "create", "name": prompt_spec.get("name") or ""}
|
|
else:
|
|
prompt = {"status": "default"}
|
|
|
|
customs = UserCustomModelsRepository(conn).list_for_user(user)
|
|
models = []
|
|
for entry in (spec.get("model") or {}).get("available") or []:
|
|
if isinstance(entry, str):
|
|
models.append(
|
|
{
|
|
"id": entry,
|
|
"status": "matched" if validate_model_id(entry, user) else "unavailable",
|
|
}
|
|
)
|
|
elif isinstance(entry, dict):
|
|
match = _find_custom_model(customs, entry)
|
|
models.append(
|
|
{
|
|
"display_name": entry.get("display_name") or "",
|
|
"status": "reuse" if match else "create",
|
|
"requires_secrets": ["api_key"] if not match else [],
|
|
}
|
|
)
|
|
|
|
workflow_plan = None
|
|
wf_spec = spec.get("workflow")
|
|
if spec.get("agent_type") == "workflow":
|
|
if isinstance(wf_spec, dict):
|
|
action = "create"
|
|
if target.get("action") == "update" and target.get("agent_id"):
|
|
agent_row = AgentsRepository(conn).get(str(target["agent_id"]), user)
|
|
existing_wf = (agent_row or {}).get("workflow_id")
|
|
if existing_wf and WorkflowsRepository(conn).get(str(existing_wf), user):
|
|
action = "update"
|
|
workflow_plan = {
|
|
"nodes": len(wf_spec.get("nodes") or []),
|
|
"edges": len(wf_spec.get("edges") or []),
|
|
"action": action,
|
|
}
|
|
elif "workflow" in spec:
|
|
# Explicit ``workflow: null`` (vs. an omitted key, which changes
|
|
# nothing). Apply clears the link and reaps the row when no other
|
|
# agent references it, destroying the graph, its run history and
|
|
# its artifacts — the dry-run has to say so before the user commits.
|
|
workflow_plan = _plan_workflow_removal(conn, user, target)
|
|
|
|
return {
|
|
"target": target,
|
|
"sources": sources,
|
|
"tools": tools,
|
|
"prompt": prompt,
|
|
"models": models,
|
|
"workflow": workflow_plan,
|
|
}
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Apply
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
def _apply_prompt(conn, user: str, spec: dict) -> Optional[str]:
|
|
prompt_spec = spec.get("prompt")
|
|
if not isinstance(prompt_spec, dict):
|
|
return None
|
|
row = PromptsRepository(conn).find_or_create(
|
|
user, prompt_spec.get("name") or "Imported prompt", prompt_spec.get("content") or ""
|
|
)
|
|
return str(row["id"])
|
|
|
|
|
|
def _apply_sources(conn, user: str, spec: dict, resolution: dict, warnings: list):
|
|
"""Resolve spec sources; returns ``(ids, id_by_name)``.
|
|
|
|
``id_by_name`` maps each spec source name to its resolved id (or None) so
|
|
a workflow graph's node configs — which reference sources by name — can
|
|
be rewritten from the same resolution.
|
|
"""
|
|
repo = SourcesRepository(conn)
|
|
mapping = resolution.get("sources")
|
|
if not isinstance(mapping, dict):
|
|
mapping = {}
|
|
resolved: list[str] = []
|
|
id_by_name: dict[str, Optional[str]] = {}
|
|
entries = [src for src in spec.get("sources") or [] if isinstance(src, dict)]
|
|
# ``id_by_name`` (and a workflow node's source reference) keys on the name
|
|
# alone, but names aren't unique per user — ``find_by_name`` returns the
|
|
# oldest match, so duplicates collapse onto one source. Same hazard the
|
|
# custom-model display names carry below; surface it rather than silently
|
|
# rewiring a node to a different source.
|
|
names = [src.get("name") or "" for src in entries]
|
|
for dup in sorted({n for n in names if n and names.count(n) > 1}):
|
|
warnings.append(
|
|
f"Multiple sources share the name '{dup}'; references to it all "
|
|
"resolve to the oldest one"
|
|
)
|
|
for src in entries:
|
|
name = src.get("name") or ""
|
|
id_by_name.setdefault(name, None)
|
|
mapped = mapping.get(name)
|
|
if mapped:
|
|
# Ownership-check the client-supplied id before linking (IDOR guard).
|
|
owned = repo.get_any(str(mapped), user)
|
|
if owned:
|
|
resolved.append(str(owned["id"]))
|
|
id_by_name[name] = str(owned["id"])
|
|
else:
|
|
warnings.append(f"Source mapping for '{name}' is not yours; ignored")
|
|
continue
|
|
match = repo.find_by_name(user, name)
|
|
if match:
|
|
resolved.append(str(match["id"]))
|
|
id_by_name[name] = str(match["id"])
|
|
else:
|
|
warnings.append(f"Source '{name}' not found; left unattached")
|
|
seen: set = set()
|
|
out = []
|
|
for sid in resolved:
|
|
if sid not in seen:
|
|
seen.add(sid)
|
|
out.append(sid)
|
|
return out, id_by_name
|
|
|
|
|
|
def _validate_tool_urls(tool_type: str, config: dict) -> Optional[str]:
|
|
"""SSRF-guard imported tool config to parity with the create routes.
|
|
|
|
Returns an error message if a URL is unsafe, else None.
|
|
"""
|
|
try:
|
|
if tool_type == "mcp_tool":
|
|
server_url = (config.get("server_url") or "").strip()
|
|
if server_url:
|
|
validate_url(server_url)
|
|
elif tool_type == "api_tool":
|
|
url = (config.get("url") or "").strip()
|
|
if url:
|
|
validate_url(url)
|
|
except SSRFError:
|
|
return f"Tool '{tool_type}' has an unsafe URL; not created"
|
|
return None
|
|
|
|
|
|
def _create_tool_from_spec(conn, user: str, tool: dict, secrets: dict, warnings: list) -> Optional[str]:
|
|
from docsgpt.api.user.tools.routes import _encrypt_secret_fields, transform_actions
|
|
|
|
tool_type = tool.get("type") or ""
|
|
inst = _tool_instance(tool_type)
|
|
if inst is None:
|
|
warnings.append(f"Tool type '{tool_type}' not available on this instance; skipped")
|
|
return None
|
|
config_requirements = inst.get_config_requirements() or {}
|
|
config = dict(tool.get("config") or {})
|
|
config.update(secrets or {})
|
|
if tool_type == "api_tool":
|
|
label = tool.get("display_name") or tool.get("name") or tool_type
|
|
config["actions"] = _sanitize_api_tool_actions(
|
|
config.get("actions"), label, warnings
|
|
)
|
|
if not config["actions"]:
|
|
# Zero actions means a dead tool — the executor resolves api_tool
|
|
# calls exclusively through config["actions"]. Reachable via older
|
|
# exports (which couldn't carry actions) or a fully-dropped block.
|
|
warnings.append(
|
|
f"API tool '{label}' was created without any actions; "
|
|
"configure its actions in the tool's settings before use"
|
|
)
|
|
missing = [k for k in _secret_field_names(config_requirements, required_only=True) if not config.get(k)]
|
|
if missing:
|
|
warnings.append(f"Tool '{tool_type}' needs secret(s) {missing}; not created")
|
|
return None
|
|
url_error = _validate_tool_urls(tool_type, config)
|
|
if url_error:
|
|
warnings.append(url_error)
|
|
return None
|
|
provided_secrets = [k for k in _secret_field_names(config_requirements) if config.get(k)]
|
|
storage_config = _encrypt_secret_fields(config, config_requirements, user)
|
|
if provided_secrets and not storage_config.get("encrypted_credentials"):
|
|
warnings.append(f"Tool '{tool_type}' secret encryption failed; not created")
|
|
return None
|
|
actions = transform_actions(inst.get_actions_metadata() or [])
|
|
actions = _apply_action_overrides(actions, tool.get("actions"), tool_type, warnings)
|
|
created = UserToolsRepository(conn).create(
|
|
user,
|
|
tool_type,
|
|
config=storage_config,
|
|
custom_name=tool.get("name") or "",
|
|
display_name=tool.get("display_name") or tool_type,
|
|
description=tool.get("description") or "",
|
|
config_requirements=config_requirements,
|
|
actions=actions,
|
|
status=True,
|
|
)
|
|
return str(created["id"])
|
|
|
|
|
|
def _apply_tools(conn, user: str, spec: dict, resolution: dict, warnings: list):
|
|
"""Resolve spec tools; returns ``(ids, id_by_key)``.
|
|
|
|
``id_by_key`` maps each spec entry's ``tool-N`` key to its resolved id
|
|
(or None when skipped/unavailable) so a workflow graph's node configs —
|
|
which reference tools by those keys — can be rewritten from the same
|
|
resolution.
|
|
"""
|
|
repo = UserToolsRepository(conn)
|
|
user_tools = repo.list_for_user(user)
|
|
decisions = resolution.get("tools")
|
|
if not isinstance(decisions, dict):
|
|
decisions = {}
|
|
out: list[str] = []
|
|
id_by_key: dict[str, Optional[str]] = {}
|
|
|
|
def _resolve(key: str, tid: Optional[str]) -> None:
|
|
id_by_key[key] = tid
|
|
if tid:
|
|
out.append(tid)
|
|
|
|
for index, tool in enumerate(spec.get("tools") or []):
|
|
key = _tool_key(index)
|
|
tool_type = tool.get("type") or ""
|
|
decision = decisions.get(key)
|
|
if not isinstance(decision, dict):
|
|
decision = {}
|
|
choice = decision.get("decision")
|
|
|
|
if tool.get("builtin"):
|
|
if synthesize_tool_by_name(tool_type) is None:
|
|
warnings.append(f"Built-in tool '{tool_type}' unavailable; skipped")
|
|
_resolve(key, None)
|
|
continue
|
|
_resolve(key, default_tool_id(tool_type))
|
|
continue
|
|
|
|
if choice == "skip":
|
|
_resolve(key, None)
|
|
continue
|
|
if choice == "reuse" and decision.get("tool_id"):
|
|
# Ownership-check the client-supplied id before linking (IDOR guard).
|
|
tid = str(decision["tool_id"])
|
|
if is_synthesized_tool_id(tid) or repo.get_any(tid, user):
|
|
_resolve(key, tid)
|
|
else:
|
|
warnings.append(f"Tool reuse id '{tid}' is not yours; ignored")
|
|
_resolve(key, None)
|
|
continue
|
|
|
|
custom_name = tool.get("name") or ""
|
|
match = _find_user_tool(user_tools, tool_type, custom_name)
|
|
if match and choice != "create":
|
|
_resolve(key, str(match["id"]))
|
|
continue
|
|
|
|
created_id = _create_tool_from_spec(conn, user, tool, decision.get("secrets") or {}, warnings)
|
|
_resolve(key, created_id)
|
|
if created_id:
|
|
# Record so a later identical spec entry reuses rather than re-creates.
|
|
user_tools.append({"id": created_id, "name": tool_type, "custom_name": custom_name})
|
|
return out, id_by_key
|
|
|
|
|
|
def _apply_models(conn, user: str, spec: dict, resolution: dict, warnings: list):
|
|
model_spec = spec.get("model") or {}
|
|
repo = UserCustomModelsRepository(conn)
|
|
customs = repo.list_for_user(user)
|
|
decisions = resolution.get("models")
|
|
if not isinstance(decisions, dict):
|
|
decisions = {}
|
|
name_to_id: dict[str, str] = {}
|
|
models: list[str] = []
|
|
|
|
available = model_spec.get("available") or []
|
|
# ``name_to_id`` (and a workflow node's ``model_id`` reference) keys on
|
|
# display_name alone, while custom-model identity is the full
|
|
# (display_name, upstream_model_id, base_url) triple — duplicates
|
|
# collapse onto the last entry, so surface that instead of silently
|
|
# rewiring nodes.
|
|
custom_names = [
|
|
e.get("display_name") or "" for e in available if isinstance(e, dict)
|
|
]
|
|
for dup in sorted({n for n in custom_names if n and custom_names.count(n) > 1}):
|
|
warnings.append(
|
|
f"Multiple custom models share the display name '{dup}'; "
|
|
"references to it resolve to the last one in the file"
|
|
)
|
|
|
|
for entry in available:
|
|
if isinstance(entry, str):
|
|
if validate_model_id(entry, user):
|
|
models.append(entry)
|
|
name_to_id[entry] = entry
|
|
else:
|
|
warnings.append(f"Model '{entry}' not available; skipped")
|
|
continue
|
|
if not isinstance(entry, dict):
|
|
continue
|
|
display_name = entry.get("display_name") or ""
|
|
match = _find_custom_model(customs, entry)
|
|
if match:
|
|
models.append(str(match["id"]))
|
|
name_to_id[display_name] = str(match["id"])
|
|
continue
|
|
decision = decisions.get(display_name)
|
|
api_key = decision.get("api_key") if isinstance(decision, dict) else None
|
|
if not api_key:
|
|
warnings.append(f"Custom model '{display_name}' needs an API key; skipped")
|
|
continue
|
|
base_url = entry.get("base_url") or ""
|
|
try:
|
|
validate_user_base_url(base_url)
|
|
except UnsafeUserUrlError:
|
|
warnings.append(f"Custom model '{display_name}' has an unsafe base URL; skipped")
|
|
continue
|
|
created = repo.create(
|
|
user,
|
|
entry.get("upstream_model_id") or "",
|
|
display_name,
|
|
base_url,
|
|
api_key,
|
|
capabilities=entry.get("capabilities") or {},
|
|
)
|
|
if not created.get("api_key_encrypted"):
|
|
warnings.append(f"Custom model '{display_name}' key encryption failed; skipped")
|
|
continue
|
|
models.append(str(created["id"]))
|
|
name_to_id[display_name] = str(created["id"])
|
|
|
|
default_id = ""
|
|
raw_default = model_spec.get("default") or ""
|
|
if raw_default:
|
|
if raw_default in name_to_id:
|
|
default_id = name_to_id[raw_default]
|
|
elif validate_model_id(raw_default, user):
|
|
default_id = raw_default
|
|
else:
|
|
# Custom-model default carried as its display_name — resolve against
|
|
# the user's own models even if it wasn't in `available`.
|
|
owned = next(
|
|
(c for c in customs if (c.get("display_name") or "") == raw_default), None
|
|
)
|
|
if owned:
|
|
default_id = str(owned["id"])
|
|
else:
|
|
warnings.append(f"Default model '{raw_default}' not resolved")
|
|
return (models or None), (default_id or None), name_to_id
|
|
|
|
|
|
def _limit(spec: dict, key: str):
|
|
return (spec.get("limits") or {}).get(key)
|
|
|
|
|
|
# Sentinel: the file omitted ``spec.workflow`` entirely, so the agent's
|
|
# existing workflow link is left untouched (vs. explicit null = clear).
|
|
_WORKFLOW_UNCHANGED = object()
|
|
|
|
|
|
def _rewrite_node_refs(
|
|
conn,
|
|
user: str,
|
|
cfg: dict,
|
|
tool_ids_by_key: dict,
|
|
source_ids_by_name: dict,
|
|
model_ids_by_name: dict,
|
|
warnings: list,
|
|
) -> None:
|
|
"""Resolve an agent node's portable references back to the importer's ids.
|
|
|
|
Export writes ``tool-N`` keys / source names / model display names, and a
|
|
hand-edited file may also carry raw ids — those are accepted only when
|
|
they are a builtin synthetic id or ownership-check against the importing
|
|
user (IDOR guard). Anything unresolvable is dropped with a warning so the
|
|
graph imports degraded rather than pointing at foreign resources.
|
|
"""
|
|
tools_repo = UserToolsRepository(conn)
|
|
sources_repo = SourcesRepository(conn)
|
|
|
|
resolved_tools = []
|
|
for ref in cfg.get("tools") or []:
|
|
ref = str(ref)
|
|
if ref in tool_ids_by_key:
|
|
tid = tool_ids_by_key[ref]
|
|
if tid:
|
|
resolved_tools.append(tid)
|
|
else:
|
|
warnings.append(f"Workflow node tool '{ref}' not resolved; removed")
|
|
elif is_synthesized_tool_id(ref):
|
|
resolved_tools.append(ref)
|
|
elif looks_like_uuid(ref) and tools_repo.get_any(ref, user):
|
|
resolved_tools.append(ref)
|
|
else:
|
|
warnings.append(f"Workflow node tool '{ref}' not resolved; removed")
|
|
cfg["tools"] = resolved_tools
|
|
|
|
resolved_sources = []
|
|
for ref in cfg.get("sources") or []:
|
|
ref = str(ref)
|
|
if ref in source_ids_by_name:
|
|
sid = source_ids_by_name[ref]
|
|
if sid:
|
|
resolved_sources.append(sid)
|
|
else:
|
|
warnings.append(f"Workflow node source '{ref}' not resolved; removed")
|
|
elif looks_like_uuid(ref) and sources_repo.get_any(ref, user):
|
|
resolved_sources.append(ref)
|
|
else:
|
|
warnings.append(f"Workflow node source '{ref}' not resolved; removed")
|
|
# Distinct names can collapse onto one source (see ``_apply_sources``), so
|
|
# dedupe rather than storing the same id twice on the node.
|
|
cfg["sources"] = list(dict.fromkeys(resolved_sources))
|
|
|
|
raw_model = cfg.get("model_id")
|
|
if raw_model:
|
|
raw_model = str(raw_model)
|
|
if raw_model in model_ids_by_name:
|
|
cfg["model_id"] = model_ids_by_name[raw_model]
|
|
elif validate_model_id(raw_model, user):
|
|
cfg["model_id"] = raw_model
|
|
else:
|
|
warnings.append(f"Workflow node model '{raw_model}' not resolved; removed")
|
|
cfg["model_id"] = None
|
|
|
|
|
|
def _apply_workflow(
|
|
conn,
|
|
user: str,
|
|
spec: dict,
|
|
target: dict,
|
|
tool_ids_by_key: dict,
|
|
source_ids_by_name: dict,
|
|
model_ids_by_name: dict,
|
|
warnings: list,
|
|
):
|
|
"""Write the imported graph; returns the workflow id, None, or the sentinel.
|
|
|
|
The graph goes through the same validation gate as the workflow API
|
|
(``validate_workflow_structure``) so a hand-edited YAML can't install a
|
|
graph the write path would have rejected. When the matched agent already
|
|
owns a workflow the graph lands as a new version of it (mirroring
|
|
``PUT /api/workflows/<id>``); otherwise a new workflow is created.
|
|
"""
|
|
if "workflow" not in spec:
|
|
return _WORKFLOW_UNCHANGED
|
|
wf_spec = spec.get("workflow")
|
|
if not isinstance(wf_spec, dict):
|
|
return None
|
|
from docsgpt.api.user.workflows.routes import (
|
|
_write_graph,
|
|
get_workflow_graph_version,
|
|
normalize_agent_node_json_schemas,
|
|
validate_workflow_structure,
|
|
)
|
|
|
|
nodes_data = []
|
|
for node in wf_spec.get("nodes") or []:
|
|
config = copy.deepcopy(node.get("config") or {})
|
|
if not isinstance(config, dict):
|
|
config = {}
|
|
if node.get("type") == "agent":
|
|
_rewrite_node_refs(
|
|
conn,
|
|
user,
|
|
_effective_node_config(config),
|
|
tool_ids_by_key,
|
|
source_ids_by_name,
|
|
model_ids_by_name,
|
|
warnings,
|
|
)
|
|
_strip_shadowed_ref_keys(config)
|
|
position = node.get("position")
|
|
nodes_data.append(
|
|
{
|
|
"id": str(node.get("id")),
|
|
"type": str(node.get("type")),
|
|
"title": node.get("title") or "",
|
|
"description": node.get("description") or "",
|
|
"position": position if isinstance(position, dict) else {"x": 0, "y": 0},
|
|
"data": config,
|
|
}
|
|
)
|
|
edges_data = [
|
|
{
|
|
"id": str(edge.get("id")),
|
|
"source": edge.get("source"),
|
|
"target": edge.get("target"),
|
|
"sourceHandle": edge.get("sourceHandle"),
|
|
"targetHandle": edge.get("targetHandle"),
|
|
}
|
|
for edge in wf_spec.get("edges") or []
|
|
]
|
|
|
|
errors = validate_workflow_structure(nodes_data, edges_data, user_id=user)
|
|
if errors:
|
|
raise AgentImportError("Workflow validation failed: " + "; ".join(errors[:5]))
|
|
nodes_data = normalize_agent_node_json_schemas(nodes_data)
|
|
|
|
wf_repo = WorkflowsRepository(conn)
|
|
name = wf_spec.get("name") or spec.get("name") or "Imported workflow"
|
|
description = wf_spec.get("description") or ""
|
|
|
|
existing = None
|
|
if target.get("action") == "update" and target.get("agent_id"):
|
|
agent_row = AgentsRepository(conn).get(str(target["agent_id"]), user)
|
|
existing_id = (agent_row or {}).get("workflow_id")
|
|
if existing_id:
|
|
existing = wf_repo.get(str(existing_id), user)
|
|
if existing is not None:
|
|
from docsgpt.api.user.resource_access import prune_sponsors
|
|
from docsgpt.api.user.workflows.routes import _node_refs
|
|
|
|
pg_workflow_id = str(existing["id"])
|
|
next_version = get_workflow_graph_version(existing) + 1
|
|
_write_graph(conn, pg_workflow_id, next_version, nodes_data, edges_data)
|
|
workflow_fields = {
|
|
"name": name,
|
|
"description": description,
|
|
"current_graph_version": next_version,
|
|
}
|
|
# Sponsors of node resources the file dropped go too, so a stale
|
|
# record can't vouch for the resource if someone adds it back.
|
|
sponsors = existing.get("resource_sponsors") or {}
|
|
pruned = prune_sponsors(sponsors, _node_refs(nodes_data))
|
|
if pruned != sponsors:
|
|
workflow_fields["resource_sponsors"] = pruned
|
|
wf_repo.update(pg_workflow_id, user, workflow_fields)
|
|
WorkflowNodesRepository(conn).delete_other_versions(pg_workflow_id, next_version)
|
|
WorkflowEdgesRepository(conn).delete_other_versions(pg_workflow_id, next_version)
|
|
return pg_workflow_id
|
|
|
|
created = wf_repo.create(user, name, description=description)
|
|
_write_graph(conn, str(created["id"]), 1, nodes_data, edges_data)
|
|
return str(created["id"])
|
|
|
|
|
|
def _prune_agent_sponsors(agents_repo: AgentsRepository, agent_id: str, user: str) -> None:
|
|
"""Drop sponsor records of resources the imported agent no longer references.
|
|
|
|
Args:
|
|
agents_repo: Repository on the import's connection.
|
|
agent_id: The updated agent.
|
|
user: Its owner.
|
|
"""
|
|
from docsgpt.api.user.resource_access import agent_refs, prune_sponsors
|
|
|
|
row = agents_repo.get(agent_id, user)
|
|
if not row:
|
|
return
|
|
sponsors = row.get("resource_sponsors") or {}
|
|
pruned = prune_sponsors(sponsors, agent_refs(row))
|
|
if pruned != sponsors:
|
|
agents_repo.update(agent_id, user, {"resource_sponsors": pruned})
|
|
|
|
|
|
def apply_import(conn, user: str, doc: dict, resolution: Optional[dict] = None) -> dict:
|
|
"""Create or update an agent from a parsed YAML doc.
|
|
|
|
A newly created agent lands as a draft. An update to an existing agent
|
|
(matched by id or slug) preserves that agent's current status, so
|
|
re-importing over a published agent keeps it live — its API key and any
|
|
active users are unaffected — rather than silently reverting it to draft.
|
|
|
|
On update the YAML is authoritative: fields it specifies are written even
|
|
when that clears a value (e.g. removing all models, dropping ``json_schema``,
|
|
or switching ``prompt`` back to default), so re-importing an edited file is
|
|
a true sync rather than an additive merge.
|
|
"""
|
|
if not isinstance(resolution, dict):
|
|
resolution = {}
|
|
spec = doc["spec"]
|
|
metadata = doc.get("metadata") or {}
|
|
warnings: list[str] = []
|
|
|
|
target = _resolve_target(conn, user, metadata)
|
|
|
|
prompt_id = _apply_prompt(conn, user, spec)
|
|
source_ids, source_ids_by_name = _apply_sources(conn, user, spec, resolution, warnings)
|
|
tool_ids, tool_ids_by_key = _apply_tools(conn, user, spec, resolution, warnings)
|
|
models, default_model_id, model_ids_by_name = _apply_models(
|
|
conn, user, spec, resolution, warnings
|
|
)
|
|
|
|
is_workflow = (spec.get("agent_type") or "classic") == "workflow"
|
|
|
|
agents_repo = AgentsRepository(conn)
|
|
is_update = bool(target.get("action") == "update" and target.get("agent_id"))
|
|
exclude_id = str(target["agent_id"]) if is_update else None
|
|
# Slug is recomputed from the file, not preserved from the matched row:
|
|
# a round-tripped export carries metadata.slug so it stays stable, but an
|
|
# update-by-id whose file omits the slug and renames the agent rewrites the
|
|
# slug to follow the new name (excluding the row's own slug from collision
|
|
# checks). Intended — the slug tracks the file.
|
|
slug = _unique_slug(agents_repo, user, metadata.get("slug") or spec.get("name"), exclude_id=exclude_id)
|
|
|
|
try:
|
|
chunks_value = int(spec["chunks"]) if spec.get("chunks") is not None else 6
|
|
except (TypeError, ValueError):
|
|
chunks_value = 6
|
|
|
|
# YAML-authoritative fields — written even when the resolved value is None,
|
|
# so a re-import can CLEAR models / json_schema / prompt-to-default. On
|
|
# create, AgentsRepository.create() skips None kwargs, so Nones become
|
|
# column defaults.
|
|
authoritative: dict[str, Any] = {
|
|
"description": spec.get("description") or "",
|
|
"agent_type": spec.get("agent_type") or "classic",
|
|
"chunks": chunks_value,
|
|
"prompt_id": prompt_id,
|
|
"tools": tool_ids,
|
|
"json_schema": spec.get("json_schema"),
|
|
"models": models,
|
|
"default_model_id": default_model_id,
|
|
"extra_source_ids": source_ids,
|
|
"limited_token_mode": bool(_limit(spec, "limited_token_mode")),
|
|
"limited_request_mode": bool(_limit(spec, "limited_request_mode")),
|
|
"allow_system_prompt_override": bool(spec.get("allow_system_prompt_override")),
|
|
"config": _import_config(spec),
|
|
"slug": slug,
|
|
}
|
|
orphaned_workflow_id: Optional[str] = None
|
|
if is_workflow:
|
|
# A workflow agent's behavior lives in its graph — the agent-level
|
|
# prompt/tool/model/source slots in the file exist only to carry the
|
|
# graph's references through resolution, so they never land on the row.
|
|
authoritative.update(
|
|
{
|
|
"prompt_id": None,
|
|
"tools": [],
|
|
"models": None,
|
|
"default_model_id": None,
|
|
"extra_source_ids": [],
|
|
}
|
|
)
|
|
prior_workflow_id = None
|
|
if is_update:
|
|
prior_workflow_id = (
|
|
agents_repo.get(str(target["agent_id"]), user) or {}
|
|
).get("workflow_id")
|
|
published = (target.get("status") or "") == "published"
|
|
workflow_result = _apply_workflow(
|
|
conn,
|
|
user,
|
|
spec,
|
|
target,
|
|
tool_ids_by_key,
|
|
source_ids_by_name,
|
|
model_ids_by_name,
|
|
warnings,
|
|
)
|
|
# A published workflow agent with no graph can't run — the create and
|
|
# update routes both reject that state, so import must not install it
|
|
# through the back door (e.g. a file that flips ``agent_type`` to
|
|
# workflow over a published classic agent).
|
|
if workflow_result is _WORKFLOW_UNCHANGED:
|
|
if published and not prior_workflow_id:
|
|
raise AgentImportError(
|
|
"A published workflow agent needs a workflow; this file has none"
|
|
)
|
|
elif workflow_result is None and published:
|
|
if not prior_workflow_id:
|
|
raise AgentImportError(
|
|
"A published workflow agent needs a workflow; this file has none"
|
|
)
|
|
warnings.append(
|
|
"File has no workflow; the published agent kept its existing one"
|
|
)
|
|
else:
|
|
authoritative["workflow_id"] = workflow_result
|
|
if workflow_result is None and prior_workflow_id:
|
|
# Explicit ``workflow: null`` on a draft: remember the old
|
|
# link so the row can be reaped after the update — there is
|
|
# no workflow-list API independent of agents, so an
|
|
# unlinked workflow would be unreachable forever.
|
|
orphaned_workflow_id = str(prior_workflow_id)
|
|
|
|
# Optional fields — applied only when present so a partial file doesn't wipe them.
|
|
optional = {
|
|
"retriever": spec.get("retriever") or ("classic" if not source_ids else None),
|
|
"token_limit": _limit(spec, "token_limit"),
|
|
"request_limit": _limit(spec, "request_limit"),
|
|
}
|
|
optional = {k: v for k, v in optional.items() if v is not None}
|
|
|
|
if is_update:
|
|
# Status is intentionally omitted so an update never changes it:
|
|
# re-importing over a published agent keeps it live (its API key and
|
|
# any active users are unaffected) instead of reverting it to draft.
|
|
# Only a brand-new agent (the create path below) starts as a draft.
|
|
fields = {**authoritative, **optional, "name": spec.get("name")}
|
|
if agents_repo.update(str(target["agent_id"]), user, fields):
|
|
_prune_agent_sponsors(agents_repo, str(target["agent_id"]), user)
|
|
if orphaned_workflow_id and not agents_repo.count_by_workflow(
|
|
orphaned_workflow_id, user
|
|
):
|
|
# Nothing references it anymore; delete (cascades nodes/edges
|
|
# and reaps run artifacts) rather than stranding the row.
|
|
WorkflowsRepository(conn).delete(orphaned_workflow_id, user)
|
|
# Destructive and irreversible — the graph, its run history and
|
|
# its artifacts (rows and stored bytes) all go. Say so.
|
|
warnings.append(
|
|
"File has no workflow; the agent's workflow was deleted "
|
|
"along with its run history and artifacts"
|
|
)
|
|
return {
|
|
"agent_id": str(target["agent_id"]),
|
|
"action": "updated",
|
|
"status": target.get("status") or "draft",
|
|
"agent_type": authoritative["agent_type"],
|
|
"slug": slug,
|
|
"warnings": warnings,
|
|
}
|
|
# Row vanished between resolve and write — fall through to create.
|
|
|
|
row = agents_repo.create(user, spec.get("name"), "draft", **{**authoritative, **optional})
|
|
return {
|
|
"agent_id": str(row["id"]),
|
|
"action": "created",
|
|
"status": "draft",
|
|
"agent_type": authoritative["agent_type"],
|
|
"slug": slug,
|
|
"warnings": warnings,
|
|
}
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# HTTP surface
|
|
# ---------------------------------------------------------------------------
|
|
|
|
agents_portability_ns = Namespace(
|
|
"agents", description="Agent import/export operations", path="/api"
|
|
)
|
|
|
|
|
|
def _read_import_payload(req):
|
|
"""Pull (yaml_text, resolution) from a JSON body or uploaded file.
|
|
|
|
Enforces a body-size cap before reading so a giant document can't be
|
|
buffered into memory.
|
|
"""
|
|
if req.content_length and req.content_length > MAX_IMPORT_BYTES:
|
|
raise AgentImportError("Import document too large")
|
|
content_type = req.content_type or ""
|
|
if "application/json" in content_type:
|
|
data = req.get_json(silent=True) or {}
|
|
resolution = data.get("resolution")
|
|
return (
|
|
data.get("yaml") or data.get("content") or "",
|
|
resolution if isinstance(resolution, dict) else {},
|
|
)
|
|
if "file" in req.files:
|
|
raw = req.files["file"].read(MAX_IMPORT_BYTES + 1)
|
|
if len(raw) > MAX_IMPORT_BYTES:
|
|
raise AgentImportError("Import document too large")
|
|
return raw.decode("utf-8", "replace"), {}
|
|
# Raw body: read from the stream with a hard cap so a chunked request
|
|
# (no Content-Length) can't be buffered unbounded into memory.
|
|
raw = req.stream.read(MAX_IMPORT_BYTES + 1)
|
|
if len(raw) > MAX_IMPORT_BYTES:
|
|
raise AgentImportError("Import document too large")
|
|
return raw.decode("utf-8", "replace"), {}
|
|
|
|
|
|
def _restricted_token_denial(conn, user: str, doc: dict) -> Optional[str]:
|
|
"""Why a resource-restricted personal access token may not import ``doc``, if it may not.
|
|
|
|
An import resolves sources, prompts and tools by name and may create them,
|
|
so a token restricted on any of those families cannot be held to its
|
|
allowlist here. A token restricted to specific agents may update exactly
|
|
those; it can never create one.
|
|
"""
|
|
for family in ("sources", "prompts", "tools", "workflows"):
|
|
if allowed_ids(request, family) is not None:
|
|
return f"A token restricted to specific {family} cannot import agents"
|
|
allowed_agents = allowed_ids(request, "agents")
|
|
if allowed_agents is None:
|
|
return None
|
|
target = _resolve_target(conn, user, doc.get("metadata") or {})
|
|
if target["action"] != "update" or target["agent_id"] not in allowed_agents:
|
|
return "This token is restricted to specific agents and may only update those"
|
|
return None
|
|
|
|
|
|
def _token_denied_response(reason: str):
|
|
return make_response(
|
|
jsonify({"success": False, "error": "resource_not_allowed", "message": reason}), 403
|
|
)
|
|
|
|
|
|
@agents_portability_ns.route("/export_agent")
|
|
class ExportAgent(Resource):
|
|
@api.doc(params={"id": "Agent ID"}, description="Export an agent as YAML")
|
|
def get(self):
|
|
if not (decoded_token := request.decoded_token):
|
|
return make_response(jsonify({"success": False}), 401)
|
|
user = decoded_token.get("sub")
|
|
agent_id = request.args.get("id")
|
|
if not agent_id:
|
|
return make_response(jsonify({"success": False, "message": "id is required"}), 400)
|
|
with db_session() as conn:
|
|
repo = AgentsRepository(conn)
|
|
try:
|
|
ra = require(conn, "agent", agent_id, user, "export")
|
|
except AccessDenied as denied:
|
|
return make_response(
|
|
jsonify({"success": False, "message": denied.message}), denied.status
|
|
)
|
|
agent = repo.get_by_id(ra.resource_id)
|
|
if not agent:
|
|
return make_response(
|
|
jsonify({"success": False, "message": "Agent not found"}), 404
|
|
)
|
|
# Serialized as the owner: the agent's prompt, sources, tools and
|
|
# workflow are the owner's (secrets are never exported).
|
|
owner_id = ra.owner_id
|
|
agent["slug"] = ensure_agent_slug(conn, agent, owner_id)
|
|
try:
|
|
export = serialize_agent(conn, agent, owner_id)
|
|
except AgentExportError as exc:
|
|
return make_response(
|
|
jsonify({"success": False, "message": str(exc)}), 400
|
|
)
|
|
body = agent_to_yaml(export)
|
|
filename = f"{export['metadata'].get('slug') or 'agent'}.agent.yaml"
|
|
response = make_response(body, 200)
|
|
response.headers["Content-Type"] = "application/x-yaml; charset=utf-8"
|
|
response.headers["Content-Disposition"] = f'attachment; filename="{filename}"'
|
|
return response
|
|
|
|
|
|
@agents_portability_ns.route("/import_agent/plan")
|
|
class ImportAgentPlan(Resource):
|
|
@api.doc(description="Dry-run an agent YAML import and return the resolution plan")
|
|
def post(self):
|
|
if not (decoded_token := request.decoded_token):
|
|
return make_response(jsonify({"success": False}), 401)
|
|
user = decoded_token.get("sub")
|
|
try:
|
|
yaml_text, _ = _read_import_payload(request)
|
|
doc = parse_agent_yaml(yaml_text)
|
|
except AgentImportError as exc:
|
|
return make_response(jsonify({"success": False, "message": str(exc)}), 400)
|
|
try:
|
|
with db_readonly() as conn:
|
|
if reason := _restricted_token_denial(conn, user, doc):
|
|
return _token_denied_response(reason)
|
|
plan = plan_import(conn, user, doc)
|
|
except Exception:
|
|
current_app.logger.error("Agent import plan failed", exc_info=True)
|
|
return make_response(
|
|
jsonify({"success": False, "message": "Could not analyze the agent file"}), 500
|
|
)
|
|
return make_response(jsonify({"success": True, "plan": plan}), 200)
|
|
|
|
|
|
@agents_portability_ns.route("/import_agent")
|
|
class ImportAgent(Resource):
|
|
@api.doc(description="Import an agent from YAML (created as a draft)")
|
|
def post(self):
|
|
if not (decoded_token := request.decoded_token):
|
|
return make_response(jsonify({"success": False}), 401)
|
|
user = decoded_token.get("sub")
|
|
try:
|
|
yaml_text, resolution = _read_import_payload(request)
|
|
doc = parse_agent_yaml(yaml_text)
|
|
except AgentImportError as exc:
|
|
return make_response(jsonify({"success": False, "message": str(exc)}), 400)
|
|
try:
|
|
with db_session() as conn:
|
|
if reason := _restricted_token_denial(conn, user, doc):
|
|
return _token_denied_response(reason)
|
|
result = apply_import(conn, user, doc, resolution)
|
|
except AgentImportError as exc:
|
|
# Apply-time rejection of the user's document (e.g. the workflow
|
|
# graph failing the API's validation gate) — a client error, with
|
|
# the detail surfaced, matching how PUT /workflows reports it.
|
|
return make_response(jsonify({"success": False, "message": str(exc)}), 400)
|
|
except Exception:
|
|
current_app.logger.error("Agent import failed", exc_info=True)
|
|
return make_response(
|
|
jsonify({"success": False, "message": "Import failed"}), 500
|
|
)
|
|
return make_response(jsonify({"success": True, **result}), 200)
|