mirror of
https://github.com/tiennm99/DocsGPT.git
synced 2026-10-12 20:15:26 +00:00
Merge pull request #2537 from arc53/agent-exports
feat: agent import / export
This commit is contained in:
24 files changed
+2329
-20
No files matched your search
@@ -0,0 +1,36 @@
|
||||
"""0019 agent slug — stable per-user identifier for YAML export/import.
|
||||
|
||||
Adds ``agents.slug`` (CITEXT) plus a partial unique index on
|
||||
``(user_id, slug)`` where slug is not null, so an exported agent can be
|
||||
re-imported idempotently (and mapped from a repo file for GitOps) while
|
||||
still allowing multiple agents to carry NULL. Idempotent both ways.
|
||||
|
||||
Revision ID: 0019_agent_slug
|
||||
Revises: 0018_tool_attempts_attribution
|
||||
"""
|
||||
|
||||
from typing import Sequence, Union
|
||||
|
||||
from alembic import op
|
||||
|
||||
|
||||
revision: str = "0019_agent_slug"
|
||||
down_revision: Union[str, None] = "0018_tool_attempts_attribution"
|
||||
branch_labels: Union[str, Sequence[str], None] = None
|
||||
depends_on: Union[str, Sequence[str], None] = None
|
||||
|
||||
|
||||
def upgrade() -> None:
|
||||
op.execute("ALTER TABLE agents ADD COLUMN IF NOT EXISTS slug CITEXT;")
|
||||
op.execute(
|
||||
"""
|
||||
CREATE UNIQUE INDEX IF NOT EXISTS ix_agents_user_slug
|
||||
ON agents (user_id, slug)
|
||||
WHERE slug IS NOT NULL;
|
||||
"""
|
||||
)
|
||||
|
||||
|
||||
def downgrade() -> None:
|
||||
op.execute("DROP INDEX IF EXISTS ix_agents_user_slug;")
|
||||
op.execute("ALTER TABLE agents DROP COLUMN IF EXISTS slug;")
|
||||
@@ -4,5 +4,12 @@ from .routes import agents_ns
|
||||
from .sharing import agents_sharing_ns
|
||||
from .webhooks import agents_webhooks_ns
|
||||
from .folders import agents_folders_ns
|
||||
from .portability import agents_portability_ns
|
||||
|
||||
__all__ = ["agents_ns", "agents_sharing_ns", "agents_webhooks_ns", "agents_folders_ns"]
|
||||
__all__ = [
|
||||
"agents_ns",
|
||||
"agents_sharing_ns",
|
||||
"agents_webhooks_ns",
|
||||
"agents_folders_ns",
|
||||
"agents_portability_ns",
|
||||
]
|
||||
@@ -0,0 +1,949 @@
|
||||
"""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.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
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 application.agents.default_tools import (
|
||||
default_tool_id,
|
||||
is_synthesized_tool_id,
|
||||
synthesize_tool_by_name,
|
||||
synthesized_tool_name_for_id,
|
||||
)
|
||||
from application.api import api
|
||||
from application.core.model_utils import validate_model_id
|
||||
from application.core.url_validation import SSRFError, validate_url
|
||||
from application.security.safe_url import UnsafeUserUrlError, validate_user_base_url
|
||||
from application.storage.db.base_repository import looks_like_uuid
|
||||
from application.storage.db.repositories.agents import AgentsRepository
|
||||
from application.storage.db.repositories.prompts import PromptsRepository
|
||||
from application.storage.db.repositories.sources import SourcesRepository
|
||||
from application.storage.db.repositories.user_custom_models import (
|
||||
UserCustomModelsRepository,
|
||||
)
|
||||
from application.storage.db.repositories.user_tools import UserToolsRepository
|
||||
from application.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
|
||||
|
||||
|
||||
class AgentImportError(Exception):
|
||||
"""Raised when an agent YAML document is malformed or unsupported."""
|
||||
|
||||
|
||||
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 application.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. ``api_tool``/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``.
|
||||
"""
|
||||
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
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# 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})
|
||||
continue
|
||||
row = repo.get_any(tid_str, user)
|
||||
if not row:
|
||||
continue
|
||||
requirements = row.get("config_requirements") or {}
|
||||
requires_secrets = _secret_field_names(requirements)
|
||||
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": _safe_export_config(row.get("config") or {}, requirements),
|
||||
"ref": tid_str,
|
||||
}
|
||||
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 serialize_agent(conn, agent: dict, user: str) -> dict:
|
||||
"""Build the portable export document for an agent row."""
|
||||
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": _serialize_prompt(conn, agent, user),
|
||||
"model": _serialize_models(conn, agent, user),
|
||||
"sources": _serialize_sources(conn, agent, user),
|
||||
"tools": _serialize_tools(conn, agent, user),
|
||||
"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)),
|
||||
}
|
||||
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 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":
|
||||
raise AgentImportError("Workflow agents are not supported for import yet")
|
||||
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_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 [],
|
||||
}
|
||||
)
|
||||
|
||||
return {"target": target, "sources": sources, "tools": tools, "prompt": prompt, "models": models}
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# 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) -> list:
|
||||
repo = SourcesRepository(conn)
|
||||
mapping = resolution.get("sources")
|
||||
if not isinstance(mapping, dict):
|
||||
mapping = {}
|
||||
resolved: list[str] = []
|
||||
for src in spec.get("sources") or []:
|
||||
name = src.get("name") or ""
|
||||
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"]))
|
||||
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"]))
|
||||
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
|
||||
|
||||
|
||||
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 application.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 {})
|
||||
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 [])
|
||||
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) -> list:
|
||||
repo = UserToolsRepository(conn)
|
||||
user_tools = repo.list_for_user(user)
|
||||
decisions = resolution.get("tools")
|
||||
if not isinstance(decisions, dict):
|
||||
decisions = {}
|
||||
out: list[str] = []
|
||||
for index, tool in enumerate(spec.get("tools") or []):
|
||||
tool_type = tool.get("type") or ""
|
||||
decision = decisions.get(_tool_key(index))
|
||||
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")
|
||||
continue
|
||||
out.append(default_tool_id(tool_type))
|
||||
continue
|
||||
|
||||
if choice == "skip":
|
||||
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):
|
||||
out.append(tid)
|
||||
else:
|
||||
warnings.append(f"Tool reuse id '{tid}' is not yours; ignored")
|
||||
continue
|
||||
|
||||
custom_name = tool.get("name") or ""
|
||||
match = _find_user_tool(user_tools, tool_type, custom_name)
|
||||
if match and choice != "create":
|
||||
out.append(str(match["id"]))
|
||||
continue
|
||||
|
||||
created_id = _create_tool_from_spec(conn, user, tool, decision.get("secrets") or {}, warnings)
|
||||
if created_id:
|
||||
out.append(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
|
||||
|
||||
|
||||
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] = []
|
||||
|
||||
for entry in model_spec.get("available") or []:
|
||||
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)
|
||||
|
||||
|
||||
def _limit(spec: dict, key: str):
|
||||
return (spec.get("limits") or {}).get(key)
|
||||
|
||||
|
||||
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 = _apply_sources(conn, user, spec, resolution, warnings)
|
||||
tool_ids = _apply_tools(conn, user, spec, resolution, warnings)
|
||||
models, default_model_id = _apply_models(conn, user, spec, resolution, warnings)
|
||||
|
||||
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 2
|
||||
except (TypeError, ValueError):
|
||||
chunks_value = 2
|
||||
|
||||
# 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")),
|
||||
"slug": slug,
|
||||
}
|
||||
# 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):
|
||||
return {
|
||||
"agent_id": str(target["agent_id"]),
|
||||
"action": "updated",
|
||||
"status": target.get("status") or "draft",
|
||||
"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",
|
||||
"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"), {}
|
||||
|
||||
|
||||
@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)
|
||||
agent = repo.get_any(agent_id, user)
|
||||
if not agent:
|
||||
return make_response(
|
||||
jsonify({"success": False, "message": "Agent not found"}), 404
|
||||
)
|
||||
if (agent.get("agent_type") or "") == "workflow":
|
||||
return make_response(
|
||||
jsonify(
|
||||
{"success": False, "message": "Workflow agents can't be exported yet"}
|
||||
),
|
||||
400,
|
||||
)
|
||||
agent["slug"] = ensure_agent_slug(conn, agent, user)
|
||||
export = serialize_agent(conn, agent, user)
|
||||
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:
|
||||
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:
|
||||
result = apply_import(conn, user, doc, resolution)
|
||||
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)
|
||||
@@ -181,6 +181,7 @@ def _format_agent_output(agent: dict, *, pinned: bool = False, include_key_maske
|
||||
out = {
|
||||
"id": str(agent["id"]),
|
||||
"name": agent.get("name", ""),
|
||||
"slug": agent.get("slug", "") or "",
|
||||
"description": agent.get("description", "") or "",
|
||||
"image": (
|
||||
generate_image_url(agent["image"]) if agent.get("image") else ""
|
||||
@@ -1297,6 +1298,10 @@ class AdoptAgent(Resource):
|
||||
|
||||
now = datetime.datetime.now(datetime.timezone.utc)
|
||||
new_key = str(uuid.uuid4())
|
||||
# Copy content columns only. Identity / idempotency columns
|
||||
# (key, slug, shared_token, incoming_webhook_token) are
|
||||
# deliberately NOT copied — the adopted agent is a fresh copy
|
||||
# and must get its own values (slug stays NULL until exported).
|
||||
create_kwargs: dict = {}
|
||||
for col in (
|
||||
"description", "agent_type", "image", "retriever",
|
||||
|
||||
@@ -5,7 +5,13 @@ Main user API routes - registers all namespace modules.
|
||||
from flask import Blueprint
|
||||
|
||||
from application.api import api
|
||||
from .agents import agents_ns, agents_sharing_ns, agents_webhooks_ns, agents_folders_ns
|
||||
from .agents import (
|
||||
agents_folders_ns,
|
||||
agents_ns,
|
||||
agents_portability_ns,
|
||||
agents_sharing_ns,
|
||||
agents_webhooks_ns,
|
||||
)
|
||||
from .analytics import analytics_ns
|
||||
from .attachments import attachments_ns
|
||||
from .conversations import conversations_ns
|
||||
@@ -32,11 +38,12 @@ api.add_namespace(conversations_ns)
|
||||
# Models
|
||||
api.add_namespace(models_ns)
|
||||
|
||||
# Agents (main, sharing, webhooks, folders)
|
||||
# Agents (main, sharing, webhooks, folders, import/export)
|
||||
api.add_namespace(agents_ns)
|
||||
api.add_namespace(agents_sharing_ns)
|
||||
api.add_namespace(agents_webhooks_ns)
|
||||
api.add_namespace(agents_folders_ns)
|
||||
api.add_namespace(agents_portability_ns)
|
||||
|
||||
# Prompts
|
||||
api.add_namespace(prompts_ns)
|
||||
|
||||
@@ -23,6 +23,7 @@ from sqlalchemy import (
|
||||
DateTime,
|
||||
ForeignKey,
|
||||
ForeignKeyConstraint,
|
||||
Index,
|
||||
Integer,
|
||||
MetaData,
|
||||
UniqueConstraint,
|
||||
@@ -204,6 +205,11 @@ agents_table = Table(
|
||||
Column("agent_type", Text),
|
||||
Column("status", Text, nullable=False),
|
||||
Column("key", CITEXT, unique=True),
|
||||
# Stable per-user human identifier used to match an agent across
|
||||
# YAML export/import (idempotent re-import / GitOps). Uniqueness is
|
||||
# enforced by a partial unique index in migration 0019, not here, so
|
||||
# multiple agents may carry NULL.
|
||||
Column("slug", CITEXT),
|
||||
Column("image", Text),
|
||||
Column("source_id", UUID(as_uuid=True), ForeignKey("sources.id", ondelete="SET NULL")),
|
||||
Column("extra_source_ids", ARRAY(UUID(as_uuid=True)), nullable=False, server_default="{}"),
|
||||
@@ -231,6 +237,17 @@ agents_table = Table(
|
||||
Column("legacy_mongo_id", Text),
|
||||
)
|
||||
|
||||
# Per-user uniqueness of the export/import slug. Mirrors the partial unique
|
||||
# index created in migration 0019 so a schema built from this metadata
|
||||
# (e.g. create_all) matches an Alembic-built one.
|
||||
Index(
|
||||
"ix_agents_user_slug",
|
||||
agents_table.c.user_id,
|
||||
agents_table.c.slug,
|
||||
unique=True,
|
||||
postgresql_where=agents_table.c.slug.isnot(None),
|
||||
)
|
||||
|
||||
user_custom_models_table = Table(
|
||||
"user_custom_models",
|
||||
metadata,
|
||||
|
||||
@@ -27,7 +27,7 @@ class AgentsRepository:
|
||||
@staticmethod
|
||||
def _normalize_unique_text(col: str, val):
|
||||
"""Coerce blank strings for nullable unique text columns to NULL."""
|
||||
if col == "key" and val == "":
|
||||
if col in ("key", "slug") and val == "":
|
||||
return None
|
||||
return val
|
||||
|
||||
@@ -35,7 +35,7 @@ class AgentsRepository:
|
||||
values: dict = {"user_id": user_id, "name": name, "status": status}
|
||||
|
||||
_ALLOWED = {
|
||||
"description", "agent_type", "key", "retriever",
|
||||
"description", "agent_type", "key", "slug", "retriever",
|
||||
"default_model_id", "incoming_webhook_token",
|
||||
"source_id", "prompt_id", "folder_id", "workflow_id",
|
||||
"extra_source_ids", "image",
|
||||
@@ -142,6 +142,17 @@ class AgentsRepository:
|
||||
row = result.fetchone()
|
||||
return row_to_dict(row) if row is not None else None
|
||||
|
||||
def find_by_slug(self, user_id: str, slug: str) -> Optional[dict]:
|
||||
"""Resolve one of a user's agents by its stable export/import slug."""
|
||||
if not slug:
|
||||
return None
|
||||
result = self._conn.execute(
|
||||
text("SELECT * FROM agents WHERE user_id = :user_id AND slug = :slug"),
|
||||
{"user_id": user_id, "slug": slug},
|
||||
)
|
||||
row = result.fetchone()
|
||||
return row_to_dict(row) if row is not None else None
|
||||
|
||||
def list_for_user(self, user_id: str) -> list[dict]:
|
||||
result = self._conn.execute(
|
||||
text("SELECT * FROM agents WHERE user_id = :user_id ORDER BY created_at DESC"),
|
||||
@@ -150,14 +161,20 @@ class AgentsRepository:
|
||||
return [row_to_dict(r) for r in result.fetchall()]
|
||||
|
||||
def list_templates(self) -> list[dict]:
|
||||
# Template/system agents are seeded under ``__system__`` (migration
|
||||
# 0001 + the seeder); the legacy ``'system'`` value is kept in the
|
||||
# match so older rows still surface.
|
||||
result = self._conn.execute(
|
||||
text("SELECT * FROM agents WHERE user_id = 'system' ORDER BY name"),
|
||||
text(
|
||||
"SELECT * FROM agents WHERE user_id IN ('system', '__system__') "
|
||||
"ORDER BY name"
|
||||
),
|
||||
)
|
||||
return [row_to_dict(r) for r in result.fetchall()]
|
||||
|
||||
def update(self, agent_id: str, user_id: str, fields: dict) -> bool:
|
||||
allowed = {
|
||||
"name", "description", "agent_type", "status", "key", "source_id",
|
||||
"name", "description", "agent_type", "status", "key", "slug", "source_id",
|
||||
"chunks", "retriever", "prompt_id", "tools", "json_schema", "models",
|
||||
"default_model_id", "folder_id", "workflow_id",
|
||||
"extra_source_ids", "image",
|
||||
|
||||
@@ -161,11 +161,8 @@ class PromptsRepository:
|
||||
)
|
||||
return result.rowcount > 0
|
||||
|
||||
def find_or_create(self, user_id: str, name: str, content: str) -> dict:
|
||||
"""Return existing prompt matching (user, name, content), or create one.
|
||||
|
||||
Used by the seeder to avoid duplicating template prompts.
|
||||
"""
|
||||
def find(self, user_id: str, name: str, content: str) -> Optional[dict]:
|
||||
"""Return a prompt exactly matching (user, name, content), or None."""
|
||||
result = self._conn.execute(
|
||||
text(
|
||||
"SELECT * FROM prompts WHERE user_id = :user_id AND name = :name AND content = :content"
|
||||
@@ -173,6 +170,14 @@ class PromptsRepository:
|
||||
{"user_id": user_id, "name": name, "content": content},
|
||||
)
|
||||
row = result.fetchone()
|
||||
if row is not None:
|
||||
return row_to_dict(row)
|
||||
return row_to_dict(row) if row is not None else None
|
||||
|
||||
def find_or_create(self, user_id: str, name: str, content: str) -> dict:
|
||||
"""Return existing prompt matching (user, name, content), or create one.
|
||||
|
||||
Used by the seeder and agent YAML import to avoid duplicating prompts.
|
||||
"""
|
||||
existing = self.find(user_id, name, content)
|
||||
if existing is not None:
|
||||
return existing
|
||||
return self.create(user_id, name, content)
|
||||
@@ -174,6 +174,25 @@ class SourcesRepository:
|
||||
return row
|
||||
return self.get_by_legacy_id(source_id, user_id)
|
||||
|
||||
def find_by_name(self, user_id: str, name: str) -> Optional[dict]:
|
||||
"""Return a user's source whose name matches ``name`` (case-insensitive).
|
||||
|
||||
Used by agent YAML import to map a portable source name back to a
|
||||
concrete source id. Returns the oldest match when names collide.
|
||||
"""
|
||||
if not name:
|
||||
return None
|
||||
result = self._conn.execute(
|
||||
text(
|
||||
"SELECT * FROM sources "
|
||||
"WHERE user_id = :user_id AND lower(name) = lower(:name) "
|
||||
"ORDER BY created_at LIMIT 1"
|
||||
),
|
||||
{"user_id": user_id, "name": name},
|
||||
)
|
||||
row = result.fetchone()
|
||||
return row_to_dict(row) if row is not None else None
|
||||
|
||||
def list_for_user(
|
||||
self,
|
||||
user_id: str,
|
||||
|
||||
@@ -4,6 +4,7 @@ import { useNavigate } from 'react-router-dom';
|
||||
import { useTranslation } from 'react-i18next';
|
||||
|
||||
import userService from '../api/services/userService';
|
||||
import Download from '../assets/download.svg';
|
||||
import Duplicate from '../assets/duplicate.svg';
|
||||
import Edit from '../assets/edit.svg';
|
||||
import FolderIcon from '../assets/folder.svg';
|
||||
@@ -104,6 +105,17 @@ export default function AgentCard({
|
||||
iconWidth: 14,
|
||||
iconHeight: 14,
|
||||
},
|
||||
{
|
||||
icon: Download,
|
||||
label: t('agents.exportAgent'),
|
||||
onClick: (e: SyntheticEvent) => {
|
||||
e.stopPropagation();
|
||||
handleExport();
|
||||
},
|
||||
variant: 'default',
|
||||
iconWidth: 14,
|
||||
iconHeight: 14,
|
||||
},
|
||||
...(agent.status === 'published'
|
||||
? [
|
||||
{
|
||||
@@ -223,6 +235,25 @@ export default function AgentCard({
|
||||
}
|
||||
};
|
||||
|
||||
const handleExport = async () => {
|
||||
try {
|
||||
const response = await userService.exportAgent(agent.id ?? '', token);
|
||||
if (!response.ok) throw new Error('Failed to export agent');
|
||||
const yamlText = await response.text();
|
||||
const blob = new Blob([yamlText], { type: 'application/x-yaml' });
|
||||
const url = URL.createObjectURL(blob);
|
||||
const link = document.createElement('a');
|
||||
link.href = url;
|
||||
link.download = `${agent.slug || agent.name || 'agent'}.agent.yaml`;
|
||||
document.body.appendChild(link);
|
||||
link.click();
|
||||
document.body.removeChild(link);
|
||||
URL.revokeObjectURL(url);
|
||||
} catch (error) {
|
||||
console.error('Error:', error);
|
||||
}
|
||||
};
|
||||
|
||||
const handleDelete = async () => {
|
||||
try {
|
||||
const response = await userService.deleteAgent(agent.id ?? '', token);
|
||||
|
||||
@@ -12,6 +12,7 @@ import {
|
||||
setConversation,
|
||||
updateConversationId,
|
||||
} from '../conversation/conversationSlice';
|
||||
import ImportAgentModal from '../modals/ImportAgentModal';
|
||||
import {
|
||||
selectAgentFolders,
|
||||
selectSelectedAgent,
|
||||
@@ -284,6 +285,7 @@ function AgentSection({
|
||||
const allAgents = useSelector(config.selectData);
|
||||
const [isCreatingFolder, setIsCreatingFolder] = useState(false);
|
||||
const [newFolderName, setNewFolderName] = useState('');
|
||||
const [showImportModal, setShowImportModal] = useState(false);
|
||||
const newFolderInputRef = useRef<HTMLInputElement>(null);
|
||||
|
||||
const currentFolderId =
|
||||
@@ -540,6 +542,16 @@ function AgentSection({
|
||||
{t('agents.folders.newFolder')}
|
||||
</Button>
|
||||
))}
|
||||
{config.id === 'user' && (
|
||||
<Button
|
||||
type="button"
|
||||
variant="outline"
|
||||
className="bg-card shrink-0 rounded-full whitespace-nowrap"
|
||||
onClick={() => setShowImportModal(true)}
|
||||
>
|
||||
{t('agents.importAgent')}
|
||||
</Button>
|
||||
)}
|
||||
{config.showNewAgentButton && (
|
||||
<Button
|
||||
type="button"
|
||||
@@ -554,6 +566,10 @@ function AgentSection({
|
||||
)}
|
||||
</div>
|
||||
</div>
|
||||
<ImportAgentModal
|
||||
modalState={showImportModal ? 'ACTIVE' : 'INACTIVE'}
|
||||
setModalState={(state) => setShowImportModal(state === 'ACTIVE')}
|
||||
/>
|
||||
|
||||
<div className="flex flex-col gap-4">
|
||||
{isLoading ? (
|
||||
|
||||
@@ -7,6 +7,7 @@ export type ToolSummary = {
|
||||
export type Agent = {
|
||||
id?: string;
|
||||
name: string;
|
||||
slug?: string;
|
||||
description: string;
|
||||
image: string;
|
||||
source: string;
|
||||
|
||||
@@ -26,6 +26,9 @@ const endpoints = {
|
||||
TEMPLATE_AGENTS: '/api/template_agents',
|
||||
ADOPT_AGENT: (id: string) => `/api/adopt_agent?id=${id}`,
|
||||
AGENT_WEBHOOK: (id: string) => `/api/agent_webhook?id=${id}`,
|
||||
EXPORT_AGENT: (id: string) => `/api/export_agent?id=${id}`,
|
||||
IMPORT_AGENT_PLAN: '/api/import_agent/plan',
|
||||
IMPORT_AGENT: '/api/import_agent',
|
||||
PROMPTS: '/api/get_prompts',
|
||||
CREATE_PROMPT: '/api/create_prompt',
|
||||
DELETE_PROMPT: '/api/delete_prompt',
|
||||
|
||||
@@ -37,6 +37,15 @@ const userService = {
|
||||
apiClient.putFormData(endpoints.USER.UPDATE_AGENT(agent_id), data, token),
|
||||
deleteAgent: (id: string, token: string | null): Promise<any> =>
|
||||
apiClient.delete(endpoints.USER.DELETE_AGENT(id), token),
|
||||
exportAgent: (id: string, token: string | null): Promise<Response> =>
|
||||
apiClient.get(endpoints.USER.EXPORT_AGENT(id), token),
|
||||
planImportAgent: (yaml: string, token: string | null): Promise<Response> =>
|
||||
apiClient.post(endpoints.USER.IMPORT_AGENT_PLAN, { yaml }, token),
|
||||
importAgent: (
|
||||
payload: { yaml: string; resolution?: unknown },
|
||||
token: string | null,
|
||||
): Promise<Response> =>
|
||||
apiClient.post(endpoints.USER.IMPORT_AGENT, payload, token),
|
||||
getPinnedAgents: (token: string | null): Promise<any> =>
|
||||
throttledApiClient.get(endpoints.USER.PINNED_AGENTS, token),
|
||||
togglePinAgent: (id: string, token: string | null): Promise<any> =>
|
||||
|
||||
@@ -0,0 +1,5 @@
|
||||
<svg xmlns="http://www.w3.org/2000/svg" width="16" height="16" viewBox="0 0 24 24" fill="none" stroke="#747474" stroke-width="2" stroke-linecap="round" stroke-linejoin="round">
|
||||
<path d="M21 15v4a2 2 0 0 1-2 2H5a2 2 0 0 1-2-2v-4"/>
|
||||
<path d="M7 10l5 5 5-5"/>
|
||||
<path d="M12 15V3"/>
|
||||
</svg>
|
||||
|
After Width: | Height: | Size: 291 B |
@@ -560,6 +560,35 @@
|
||||
"cancel": "Abbrechen",
|
||||
"parse": "Parsen",
|
||||
"import": "Importieren ({{count}})"
|
||||
},
|
||||
"importAgent": {
|
||||
"title": "Import agent from YAML",
|
||||
"description": "Upload an agent YAML file. Sources and tools are matched to your account; anything missing can be mapped or created in the next step. A new agent is imported as a draft; re-importing updates the matching agent and keeps its current status.",
|
||||
"dropzoneText": "Click to upload or drag and drop a .yaml file",
|
||||
"invalidFileType": "Only .yaml or .yml files are supported",
|
||||
"readError": "Could not read this agent file",
|
||||
"importError": "Import failed",
|
||||
"willCreate": "Creates a new agent.",
|
||||
"willUpdate": "Updates an existing agent (matched by {{matchedBy}}).",
|
||||
"sources": "Sources",
|
||||
"tools": "Tools",
|
||||
"models": "Models",
|
||||
"sourceMatched": "{{name}}: using your existing source",
|
||||
"sourceMissing": "{{name}}: not found",
|
||||
"leaveUnattached": "Leave unattached",
|
||||
"toolBuiltin": "{{type}}: built-in",
|
||||
"toolReuse": "{{name}}: reusing your existing tool",
|
||||
"toolUnavailable": "{{type}}: not available on this instance, will be skipped",
|
||||
"toolCreate": "{{name}}: will be created",
|
||||
"secretPlaceholder": "{{field}} (leave blank to skip this tool)",
|
||||
"modelCustom": "{{name}}: custom model",
|
||||
"apiKeyPlaceholder": "API key (leave blank to skip this model)",
|
||||
"warningsTitle": "Agent imported as a draft. Some items need your attention:",
|
||||
"warningsTitlePublished": "Your live agent was updated. Some items need your attention:",
|
||||
"cancel": "Cancel",
|
||||
"review": "Review",
|
||||
"import": "Import",
|
||||
"continueToAgent": "Continue to agent"
|
||||
}
|
||||
},
|
||||
"sharedConv": {
|
||||
@@ -752,7 +781,9 @@
|
||||
"backToRoot": "Zurück",
|
||||
"noSubfolders": "Keine Unterordner",
|
||||
"noFolders": "Noch keine Ordner"
|
||||
}
|
||||
},
|
||||
"importAgent": "Import",
|
||||
"exportAgent": "Export as YAML"
|
||||
},
|
||||
"components": {
|
||||
"fileUpload": {
|
||||
|
||||
@@ -712,6 +712,35 @@
|
||||
"cancel": "Cancel",
|
||||
"parse": "Parse",
|
||||
"import": "Import ({{count}})"
|
||||
},
|
||||
"importAgent": {
|
||||
"title": "Import agent from YAML",
|
||||
"description": "Upload an agent YAML file. Sources and tools are matched to your account; anything missing can be mapped or created in the next step. A new agent is imported as a draft; re-importing updates the matching agent and keeps its current status.",
|
||||
"dropzoneText": "Click to upload or drag and drop a .yaml file",
|
||||
"invalidFileType": "Only .yaml or .yml files are supported",
|
||||
"readError": "Could not read this agent file",
|
||||
"importError": "Import failed",
|
||||
"willCreate": "Creates a new agent.",
|
||||
"willUpdate": "Updates an existing agent (matched by {{matchedBy}}).",
|
||||
"sources": "Sources",
|
||||
"tools": "Tools",
|
||||
"models": "Models",
|
||||
"sourceMatched": "{{name}}: using your existing source",
|
||||
"sourceMissing": "{{name}}: not found",
|
||||
"leaveUnattached": "Leave unattached",
|
||||
"toolBuiltin": "{{type}}: built-in",
|
||||
"toolReuse": "{{name}}: reusing your existing tool",
|
||||
"toolUnavailable": "{{type}}: not available on this instance, will be skipped",
|
||||
"toolCreate": "{{name}}: will be created",
|
||||
"secretPlaceholder": "{{field}} (leave blank to skip this tool)",
|
||||
"modelCustom": "{{name}}: custom model",
|
||||
"apiKeyPlaceholder": "API key (leave blank to skip this model)",
|
||||
"warningsTitle": "Agent imported as a draft. Some items need your attention:",
|
||||
"warningsTitlePublished": "Your live agent was updated. Some items need your attention:",
|
||||
"cancel": "Cancel",
|
||||
"review": "Review",
|
||||
"import": "Import",
|
||||
"continueToAgent": "Continue to agent"
|
||||
}
|
||||
},
|
||||
"sharedConv": {
|
||||
@@ -762,6 +791,8 @@
|
||||
"edit": "Edit",
|
||||
"description": "Discover and create custom versions of DocsGPT that combine instructions, extra knowledge, and any combination of skills",
|
||||
"newAgent": "New Agent",
|
||||
"importAgent": "Import",
|
||||
"exportAgent": "Export as YAML",
|
||||
"backToAll": "Back to all agents",
|
||||
"pageHeader": {
|
||||
"crumbs": {
|
||||
|
||||
@@ -578,6 +578,35 @@
|
||||
"cancel": "Cancelar",
|
||||
"parse": "Analizar",
|
||||
"import": "Importar ({{count}})"
|
||||
},
|
||||
"importAgent": {
|
||||
"title": "Import agent from YAML",
|
||||
"description": "Upload an agent YAML file. Sources and tools are matched to your account; anything missing can be mapped or created in the next step. A new agent is imported as a draft; re-importing updates the matching agent and keeps its current status.",
|
||||
"dropzoneText": "Click to upload or drag and drop a .yaml file",
|
||||
"invalidFileType": "Only .yaml or .yml files are supported",
|
||||
"readError": "Could not read this agent file",
|
||||
"importError": "Import failed",
|
||||
"willCreate": "Creates a new agent.",
|
||||
"willUpdate": "Updates an existing agent (matched by {{matchedBy}}).",
|
||||
"sources": "Sources",
|
||||
"tools": "Tools",
|
||||
"models": "Models",
|
||||
"sourceMatched": "{{name}}: using your existing source",
|
||||
"sourceMissing": "{{name}}: not found",
|
||||
"leaveUnattached": "Leave unattached",
|
||||
"toolBuiltin": "{{type}}: built-in",
|
||||
"toolReuse": "{{name}}: reusing your existing tool",
|
||||
"toolUnavailable": "{{type}}: not available on this instance, will be skipped",
|
||||
"toolCreate": "{{name}}: will be created",
|
||||
"secretPlaceholder": "{{field}} (leave blank to skip this tool)",
|
||||
"modelCustom": "{{name}}: custom model",
|
||||
"apiKeyPlaceholder": "API key (leave blank to skip this model)",
|
||||
"warningsTitle": "Agent imported as a draft. Some items need your attention:",
|
||||
"warningsTitlePublished": "Your live agent was updated. Some items need your attention:",
|
||||
"cancel": "Cancel",
|
||||
"review": "Review",
|
||||
"import": "Import",
|
||||
"continueToAgent": "Continue to agent"
|
||||
}
|
||||
},
|
||||
"sharedConv": {
|
||||
@@ -770,7 +799,9 @@
|
||||
"backToRoot": "Volver",
|
||||
"noSubfolders": "Sin subcarpetas",
|
||||
"noFolders": "No hay carpetas todavía"
|
||||
}
|
||||
},
|
||||
"importAgent": "Import",
|
||||
"exportAgent": "Export as YAML"
|
||||
},
|
||||
"components": {
|
||||
"fileUpload": {
|
||||
|
||||
@@ -578,6 +578,35 @@
|
||||
"cancel": "キャンセル",
|
||||
"parse": "解析",
|
||||
"import": "インポート ({{count}})"
|
||||
},
|
||||
"importAgent": {
|
||||
"title": "Import agent from YAML",
|
||||
"description": "Upload an agent YAML file. Sources and tools are matched to your account; anything missing can be mapped or created in the next step. A new agent is imported as a draft; re-importing updates the matching agent and keeps its current status.",
|
||||
"dropzoneText": "Click to upload or drag and drop a .yaml file",
|
||||
"invalidFileType": "Only .yaml or .yml files are supported",
|
||||
"readError": "Could not read this agent file",
|
||||
"importError": "Import failed",
|
||||
"willCreate": "Creates a new agent.",
|
||||
"willUpdate": "Updates an existing agent (matched by {{matchedBy}}).",
|
||||
"sources": "Sources",
|
||||
"tools": "Tools",
|
||||
"models": "Models",
|
||||
"sourceMatched": "{{name}}: using your existing source",
|
||||
"sourceMissing": "{{name}}: not found",
|
||||
"leaveUnattached": "Leave unattached",
|
||||
"toolBuiltin": "{{type}}: built-in",
|
||||
"toolReuse": "{{name}}: reusing your existing tool",
|
||||
"toolUnavailable": "{{type}}: not available on this instance, will be skipped",
|
||||
"toolCreate": "{{name}}: will be created",
|
||||
"secretPlaceholder": "{{field}} (leave blank to skip this tool)",
|
||||
"modelCustom": "{{name}}: custom model",
|
||||
"apiKeyPlaceholder": "API key (leave blank to skip this model)",
|
||||
"warningsTitle": "Agent imported as a draft. Some items need your attention:",
|
||||
"warningsTitlePublished": "Your live agent was updated. Some items need your attention:",
|
||||
"cancel": "Cancel",
|
||||
"review": "Review",
|
||||
"import": "Import",
|
||||
"continueToAgent": "Continue to agent"
|
||||
}
|
||||
},
|
||||
"sharedConv": {
|
||||
@@ -770,7 +799,9 @@
|
||||
"backToRoot": "戻る",
|
||||
"noSubfolders": "サブフォルダなし",
|
||||
"noFolders": "フォルダがありません"
|
||||
}
|
||||
},
|
||||
"importAgent": "Import",
|
||||
"exportAgent": "Export as YAML"
|
||||
},
|
||||
"components": {
|
||||
"fileUpload": {
|
||||
|
||||
@@ -578,6 +578,35 @@
|
||||
"cancel": "Отмена",
|
||||
"parse": "Разобрать",
|
||||
"import": "Импорт ({{count}})"
|
||||
},
|
||||
"importAgent": {
|
||||
"title": "Import agent from YAML",
|
||||
"description": "Upload an agent YAML file. Sources and tools are matched to your account; anything missing can be mapped or created in the next step. A new agent is imported as a draft; re-importing updates the matching agent and keeps its current status.",
|
||||
"dropzoneText": "Click to upload or drag and drop a .yaml file",
|
||||
"invalidFileType": "Only .yaml or .yml files are supported",
|
||||
"readError": "Could not read this agent file",
|
||||
"importError": "Import failed",
|
||||
"willCreate": "Creates a new agent.",
|
||||
"willUpdate": "Updates an existing agent (matched by {{matchedBy}}).",
|
||||
"sources": "Sources",
|
||||
"tools": "Tools",
|
||||
"models": "Models",
|
||||
"sourceMatched": "{{name}}: using your existing source",
|
||||
"sourceMissing": "{{name}}: not found",
|
||||
"leaveUnattached": "Leave unattached",
|
||||
"toolBuiltin": "{{type}}: built-in",
|
||||
"toolReuse": "{{name}}: reusing your existing tool",
|
||||
"toolUnavailable": "{{type}}: not available on this instance, will be skipped",
|
||||
"toolCreate": "{{name}}: will be created",
|
||||
"secretPlaceholder": "{{field}} (leave blank to skip this tool)",
|
||||
"modelCustom": "{{name}}: custom model",
|
||||
"apiKeyPlaceholder": "API key (leave blank to skip this model)",
|
||||
"warningsTitle": "Agent imported as a draft. Some items need your attention:",
|
||||
"warningsTitlePublished": "Your live agent was updated. Some items need your attention:",
|
||||
"cancel": "Cancel",
|
||||
"review": "Review",
|
||||
"import": "Import",
|
||||
"continueToAgent": "Continue to agent"
|
||||
}
|
||||
},
|
||||
"sharedConv": {
|
||||
@@ -770,7 +799,9 @@
|
||||
"backToRoot": "Назад",
|
||||
"noSubfolders": "Нет подпапок",
|
||||
"noFolders": "Пока нет папок"
|
||||
}
|
||||
},
|
||||
"importAgent": "Import",
|
||||
"exportAgent": "Export as YAML"
|
||||
},
|
||||
"components": {
|
||||
"fileUpload": {
|
||||
|
||||
@@ -578,6 +578,35 @@
|
||||
"cancel": "取消",
|
||||
"parse": "解析",
|
||||
"import": "匯入 ({{count}})"
|
||||
},
|
||||
"importAgent": {
|
||||
"title": "Import agent from YAML",
|
||||
"description": "Upload an agent YAML file. Sources and tools are matched to your account; anything missing can be mapped or created in the next step. A new agent is imported as a draft; re-importing updates the matching agent and keeps its current status.",
|
||||
"dropzoneText": "Click to upload or drag and drop a .yaml file",
|
||||
"invalidFileType": "Only .yaml or .yml files are supported",
|
||||
"readError": "Could not read this agent file",
|
||||
"importError": "Import failed",
|
||||
"willCreate": "Creates a new agent.",
|
||||
"willUpdate": "Updates an existing agent (matched by {{matchedBy}}).",
|
||||
"sources": "Sources",
|
||||
"tools": "Tools",
|
||||
"models": "Models",
|
||||
"sourceMatched": "{{name}}: using your existing source",
|
||||
"sourceMissing": "{{name}}: not found",
|
||||
"leaveUnattached": "Leave unattached",
|
||||
"toolBuiltin": "{{type}}: built-in",
|
||||
"toolReuse": "{{name}}: reusing your existing tool",
|
||||
"toolUnavailable": "{{type}}: not available on this instance, will be skipped",
|
||||
"toolCreate": "{{name}}: will be created",
|
||||
"secretPlaceholder": "{{field}} (leave blank to skip this tool)",
|
||||
"modelCustom": "{{name}}: custom model",
|
||||
"apiKeyPlaceholder": "API key (leave blank to skip this model)",
|
||||
"warningsTitle": "Agent imported as a draft. Some items need your attention:",
|
||||
"warningsTitlePublished": "Your live agent was updated. Some items need your attention:",
|
||||
"cancel": "Cancel",
|
||||
"review": "Review",
|
||||
"import": "Import",
|
||||
"continueToAgent": "Continue to agent"
|
||||
}
|
||||
},
|
||||
"sharedConv": {
|
||||
@@ -770,7 +799,9 @@
|
||||
"backToRoot": "返回",
|
||||
"noSubfolders": "沒有子資料夾",
|
||||
"noFolders": "暫無資料夾"
|
||||
}
|
||||
},
|
||||
"importAgent": "Import",
|
||||
"exportAgent": "Export as YAML"
|
||||
},
|
||||
"components": {
|
||||
"fileUpload": {
|
||||
|
||||
@@ -578,6 +578,35 @@
|
||||
"cancel": "取消",
|
||||
"parse": "解析",
|
||||
"import": "导入 ({{count}})"
|
||||
},
|
||||
"importAgent": {
|
||||
"title": "Import agent from YAML",
|
||||
"description": "Upload an agent YAML file. Sources and tools are matched to your account; anything missing can be mapped or created in the next step. A new agent is imported as a draft; re-importing updates the matching agent and keeps its current status.",
|
||||
"dropzoneText": "Click to upload or drag and drop a .yaml file",
|
||||
"invalidFileType": "Only .yaml or .yml files are supported",
|
||||
"readError": "Could not read this agent file",
|
||||
"importError": "Import failed",
|
||||
"willCreate": "Creates a new agent.",
|
||||
"willUpdate": "Updates an existing agent (matched by {{matchedBy}}).",
|
||||
"sources": "Sources",
|
||||
"tools": "Tools",
|
||||
"models": "Models",
|
||||
"sourceMatched": "{{name}}: using your existing source",
|
||||
"sourceMissing": "{{name}}: not found",
|
||||
"leaveUnattached": "Leave unattached",
|
||||
"toolBuiltin": "{{type}}: built-in",
|
||||
"toolReuse": "{{name}}: reusing your existing tool",
|
||||
"toolUnavailable": "{{type}}: not available on this instance, will be skipped",
|
||||
"toolCreate": "{{name}}: will be created",
|
||||
"secretPlaceholder": "{{field}} (leave blank to skip this tool)",
|
||||
"modelCustom": "{{name}}: custom model",
|
||||
"apiKeyPlaceholder": "API key (leave blank to skip this model)",
|
||||
"warningsTitle": "Agent imported as a draft. Some items need your attention:",
|
||||
"warningsTitlePublished": "Your live agent was updated. Some items need your attention:",
|
||||
"cancel": "Cancel",
|
||||
"review": "Review",
|
||||
"import": "Import",
|
||||
"continueToAgent": "Continue to agent"
|
||||
}
|
||||
},
|
||||
"sharedConv": {
|
||||
@@ -770,7 +799,9 @@
|
||||
"backToRoot": "返回",
|
||||
"noSubfolders": "没有子文件夹",
|
||||
"noFolders": "暂无文件夹"
|
||||
}
|
||||
},
|
||||
"importAgent": "Import",
|
||||
"exportAgent": "Export as YAML"
|
||||
},
|
||||
"components": {
|
||||
"fileUpload": {
|
||||
|
||||
@@ -0,0 +1,485 @@
|
||||
import { AlertTriangle, CheckCircle2 } from 'lucide-react';
|
||||
import { useRef, useState } from 'react';
|
||||
import { useTranslation } from 'react-i18next';
|
||||
import { useSelector } from 'react-redux';
|
||||
import { useNavigate } from 'react-router-dom';
|
||||
|
||||
import userService from '../api/services/userService';
|
||||
import Upload from '../assets/upload.svg';
|
||||
import Spinner from '../components/Spinner';
|
||||
import { Button } from '../components/ui/button';
|
||||
import { Input } from '../components/ui/input';
|
||||
import { Modal } from '../components/ui/modal';
|
||||
import {
|
||||
Select,
|
||||
SelectContent,
|
||||
SelectItem,
|
||||
SelectTrigger,
|
||||
SelectValue,
|
||||
} from '../components/ui/select';
|
||||
import { ActiveState } from '../models/misc';
|
||||
import { selectSourceDocs, selectToken } from '../preferences/preferenceSlice';
|
||||
|
||||
type PlanSource = {
|
||||
name: string;
|
||||
type: string;
|
||||
status: 'matched' | 'missing';
|
||||
target_id: string | null;
|
||||
};
|
||||
|
||||
type PlanTool = {
|
||||
key: string;
|
||||
type: string;
|
||||
name?: string;
|
||||
builtin?: boolean;
|
||||
status: 'builtin' | 'reuse' | 'create' | 'unavailable';
|
||||
target_id?: string | null;
|
||||
requires_secrets?: string[];
|
||||
};
|
||||
|
||||
type PlanModel = {
|
||||
id?: string;
|
||||
display_name?: string;
|
||||
status: 'matched' | 'unavailable' | 'reuse' | 'create';
|
||||
requires_secrets?: string[];
|
||||
};
|
||||
|
||||
type ImportPlan = {
|
||||
target: {
|
||||
action: 'create' | 'update';
|
||||
agent_id: string | null;
|
||||
matched_by: string | null;
|
||||
};
|
||||
sources: PlanSource[];
|
||||
tools: PlanTool[];
|
||||
prompt: { status: string; name?: string };
|
||||
models: PlanModel[];
|
||||
};
|
||||
|
||||
interface ImportAgentModalProps {
|
||||
modalState: ActiveState;
|
||||
setModalState: (state: ActiveState) => void;
|
||||
}
|
||||
|
||||
export default function ImportAgentModal({
|
||||
modalState,
|
||||
setModalState,
|
||||
}: ImportAgentModalProps) {
|
||||
const { t } = useTranslation();
|
||||
const token = useSelector(selectToken);
|
||||
const sourceDocs = useSelector(selectSourceDocs);
|
||||
const navigate = useNavigate();
|
||||
const fileInputRef = useRef<HTMLInputElement>(null);
|
||||
|
||||
const [yamlText, setYamlText] = useState<string>('');
|
||||
const [fileName, setFileName] = useState<string>('');
|
||||
const [loading, setLoading] = useState(false);
|
||||
const [importing, setImporting] = useState(false);
|
||||
const [error, setError] = useState<string | null>(null);
|
||||
const [plan, setPlan] = useState<ImportPlan | null>(null);
|
||||
|
||||
const [sourceMap, setSourceMap] = useState<Record<string, string>>({});
|
||||
const [toolSecrets, setToolSecrets] = useState<
|
||||
Record<string, Record<string, string>>
|
||||
>({});
|
||||
const [modelKeys, setModelKeys] = useState<Record<string, string>>({});
|
||||
const [warnings, setWarnings] = useState<string[] | null>(null);
|
||||
const [importedStatus, setImportedStatus] = useState<string | null>(null);
|
||||
const [goToAgentId, setGoToAgentId] = useState<string | null>(null);
|
||||
|
||||
const reset = () => {
|
||||
setYamlText('');
|
||||
setFileName('');
|
||||
setLoading(false);
|
||||
setImporting(false);
|
||||
setError(null);
|
||||
setPlan(null);
|
||||
setSourceMap({});
|
||||
setToolSecrets({});
|
||||
setModelKeys({});
|
||||
setWarnings(null);
|
||||
setImportedStatus(null);
|
||||
setGoToAgentId(null);
|
||||
};
|
||||
|
||||
const handleClose = () => {
|
||||
setModalState('INACTIVE');
|
||||
reset();
|
||||
};
|
||||
|
||||
const handleFileChange = async (e: React.ChangeEvent<HTMLInputElement>) => {
|
||||
const selectedFile = e.target.files?.[0];
|
||||
if (!selectedFile) return;
|
||||
const name = selectedFile.name.toLowerCase();
|
||||
if (!name.endsWith('.yaml') && !name.endsWith('.yml')) {
|
||||
setError(t('modals.importAgent.invalidFileType'));
|
||||
return;
|
||||
}
|
||||
setError(null);
|
||||
setPlan(null);
|
||||
setFileName(selectedFile.name);
|
||||
setYamlText(await selectedFile.text());
|
||||
};
|
||||
|
||||
const handleAnalyze = async () => {
|
||||
if (!yamlText) return;
|
||||
setLoading(true);
|
||||
setError(null);
|
||||
try {
|
||||
const response = await userService.planImportAgent(yamlText, token);
|
||||
const data = await response.json();
|
||||
if (!response.ok || !data.success) {
|
||||
setError(data.message || t('modals.importAgent.readError'));
|
||||
return;
|
||||
}
|
||||
setPlan(data.plan as ImportPlan);
|
||||
} catch {
|
||||
setError(t('modals.importAgent.readError'));
|
||||
} finally {
|
||||
setLoading(false);
|
||||
}
|
||||
};
|
||||
|
||||
const handleImport = async () => {
|
||||
if (!plan) return;
|
||||
setImporting(true);
|
||||
setError(null);
|
||||
|
||||
const resolution: {
|
||||
sources: Record<string, string>;
|
||||
tools: Record<string, { secrets: Record<string, string> }>;
|
||||
models: Record<string, { api_key: string }>;
|
||||
} = { sources: {}, tools: {}, models: {} };
|
||||
|
||||
plan.sources.forEach((s) => {
|
||||
if (s.status === 'missing' && sourceMap[s.name]) {
|
||||
resolution.sources[s.name] = sourceMap[s.name];
|
||||
}
|
||||
});
|
||||
plan.tools.forEach((tool) => {
|
||||
if (tool.status === 'create') {
|
||||
const secrets = toolSecrets[tool.key] || {};
|
||||
if (Object.values(secrets).some((v) => v)) {
|
||||
resolution.tools[tool.key] = { secrets };
|
||||
}
|
||||
}
|
||||
});
|
||||
plan.models.forEach((model) => {
|
||||
if (
|
||||
model.status === 'create' &&
|
||||
model.display_name &&
|
||||
modelKeys[model.display_name]
|
||||
) {
|
||||
resolution.models[model.display_name] = {
|
||||
api_key: modelKeys[model.display_name],
|
||||
};
|
||||
}
|
||||
});
|
||||
|
||||
try {
|
||||
const response = await userService.importAgent(
|
||||
{ yaml: yamlText, resolution },
|
||||
token,
|
||||
);
|
||||
const data = await response.json();
|
||||
if (!response.ok || !data.success) {
|
||||
setError(data.message || t('modals.importAgent.importError'));
|
||||
return;
|
||||
}
|
||||
const agentId = data.agent_id as string;
|
||||
if (data.warnings && data.warnings.length > 0) {
|
||||
// Keep the modal open so the user sees what was skipped.
|
||||
setGoToAgentId(agentId);
|
||||
setWarnings(data.warnings as string[]);
|
||||
setImportedStatus((data.status as string) || null);
|
||||
setPlan(null);
|
||||
return;
|
||||
}
|
||||
handleClose();
|
||||
navigate(`/agents/edit/${agentId}`);
|
||||
} catch {
|
||||
setError(t('modals.importAgent.importError'));
|
||||
} finally {
|
||||
setImporting(false);
|
||||
}
|
||||
};
|
||||
|
||||
const setToolSecret = (key: string, field: string, value: string) => {
|
||||
setToolSecrets((prev) => ({
|
||||
...prev,
|
||||
[key]: { ...(prev[key] || {}), [field]: value },
|
||||
}));
|
||||
};
|
||||
|
||||
const renderFooter = () => {
|
||||
if (warnings) {
|
||||
return (
|
||||
<Button
|
||||
type="button"
|
||||
onClick={() => {
|
||||
const id = goToAgentId;
|
||||
handleClose();
|
||||
if (id) navigate(`/agents/edit/${id}`);
|
||||
}}
|
||||
className="rounded-3xl px-5"
|
||||
>
|
||||
{t('modals.importAgent.continueToAgent')}
|
||||
</Button>
|
||||
);
|
||||
}
|
||||
return (
|
||||
<>
|
||||
<Button
|
||||
type="button"
|
||||
variant="ghost"
|
||||
onClick={handleClose}
|
||||
className="rounded-3xl px-5"
|
||||
>
|
||||
{t('modals.importAgent.cancel')}
|
||||
</Button>
|
||||
{!plan ? (
|
||||
<Button
|
||||
type="button"
|
||||
onClick={handleAnalyze}
|
||||
disabled={!yamlText || loading}
|
||||
className="w-24 rounded-3xl px-5 disabled:cursor-not-allowed"
|
||||
>
|
||||
{loading ? (
|
||||
<Spinner size="small" />
|
||||
) : (
|
||||
t('modals.importAgent.review')
|
||||
)}
|
||||
</Button>
|
||||
) : (
|
||||
<Button
|
||||
type="button"
|
||||
onClick={handleImport}
|
||||
disabled={importing}
|
||||
className="rounded-3xl px-5 disabled:cursor-not-allowed"
|
||||
>
|
||||
{importing ? (
|
||||
<Spinner size="small" />
|
||||
) : (
|
||||
t('modals.importAgent.import')
|
||||
)}
|
||||
</Button>
|
||||
)}
|
||||
</>
|
||||
);
|
||||
};
|
||||
|
||||
return (
|
||||
<Modal
|
||||
open={modalState === 'ACTIVE'}
|
||||
onOpenChange={(o) => !o && handleClose()}
|
||||
title={t('modals.importAgent.title')}
|
||||
size="lg"
|
||||
contentClassName="max-h-[70vh]"
|
||||
footer={renderFooter()}
|
||||
>
|
||||
<div className="flex flex-col gap-4">
|
||||
{warnings ? (
|
||||
<div className="flex flex-col gap-3">
|
||||
<p className="text-foreground text-sm">
|
||||
{importedStatus === 'published'
|
||||
? t('modals.importAgent.warningsTitlePublished')
|
||||
: t('modals.importAgent.warningsTitle')}
|
||||
</p>
|
||||
<ul className="list-disc space-y-1 pl-5 text-sm text-yellow-700 dark:text-yellow-400">
|
||||
{warnings.map((w, i) => (
|
||||
<li key={i}>{w}</li>
|
||||
))}
|
||||
</ul>
|
||||
</div>
|
||||
) : !plan ? (
|
||||
<div className="flex flex-col gap-4">
|
||||
<p className="text-muted-foreground text-sm">
|
||||
{t('modals.importAgent.description')}
|
||||
</p>
|
||||
<div
|
||||
onClick={() => fileInputRef.current?.click()}
|
||||
className="border-border hover:border-primary flex cursor-pointer flex-col items-center justify-center rounded-xl border-2 border-dashed p-8 transition-colors"
|
||||
>
|
||||
<img
|
||||
src={Upload}
|
||||
alt=""
|
||||
className="mb-3 h-10 w-10 opacity-60 dark:invert"
|
||||
/>
|
||||
<p className="text-foreground text-sm font-medium">
|
||||
{fileName || t('modals.importAgent.dropzoneText')}
|
||||
</p>
|
||||
</div>
|
||||
<input
|
||||
ref={fileInputRef}
|
||||
type="file"
|
||||
accept=".yaml,.yml"
|
||||
onChange={handleFileChange}
|
||||
className="hidden"
|
||||
/>
|
||||
{error && <p className="text-destructive text-sm">{error}</p>}
|
||||
</div>
|
||||
) : (
|
||||
<div className="flex flex-col gap-5">
|
||||
<div className="bg-muted rounded-xl p-3 text-sm">
|
||||
{plan.target.action === 'update'
|
||||
? t('modals.importAgent.willUpdate', {
|
||||
matchedBy: plan.target.matched_by,
|
||||
})
|
||||
: t('modals.importAgent.willCreate')}
|
||||
</div>
|
||||
|
||||
{plan.sources.length > 0 && (
|
||||
<Section title={t('modals.importAgent.sources')}>
|
||||
{plan.sources.map((s) =>
|
||||
s.status === 'matched' ? (
|
||||
<p key={s.name} className="flex items-center gap-2 text-sm">
|
||||
<CheckCircle2 className="h-4 w-4 text-green-600 dark:text-green-400" />
|
||||
{t('modals.importAgent.sourceMatched', { name: s.name })}
|
||||
</p>
|
||||
) : (
|
||||
<div key={s.name} className="flex flex-col gap-1">
|
||||
<p className="flex items-center gap-2 text-sm">
|
||||
<AlertTriangle className="h-4 w-4 text-yellow-600 dark:text-yellow-400" />
|
||||
{t('modals.importAgent.sourceMissing', {
|
||||
name: s.name,
|
||||
})}
|
||||
</p>
|
||||
<Select
|
||||
value={sourceMap[s.name] || ''}
|
||||
onValueChange={(value) =>
|
||||
setSourceMap((prev) => ({ ...prev, [s.name]: value }))
|
||||
}
|
||||
>
|
||||
<SelectTrigger className="w-full">
|
||||
<SelectValue
|
||||
placeholder={t(
|
||||
'modals.importAgent.leaveUnattached',
|
||||
)}
|
||||
/>
|
||||
</SelectTrigger>
|
||||
<SelectContent>
|
||||
{(sourceDocs || []).map(
|
||||
(doc: { id?: string; name?: string }) => (
|
||||
<SelectItem
|
||||
key={String(doc.id)}
|
||||
value={String(doc.id)}
|
||||
>
|
||||
{doc.name}
|
||||
</SelectItem>
|
||||
),
|
||||
)}
|
||||
</SelectContent>
|
||||
</Select>
|
||||
</div>
|
||||
),
|
||||
)}
|
||||
</Section>
|
||||
)}
|
||||
|
||||
{plan.tools.length > 0 && (
|
||||
<Section title={t('modals.importAgent.tools')}>
|
||||
{plan.tools.map((tool) => (
|
||||
<div key={tool.key} className="flex flex-col gap-1">
|
||||
{tool.status === 'builtin' && (
|
||||
<p className="flex items-center gap-2 text-sm">
|
||||
<CheckCircle2 className="h-4 w-4 text-green-600 dark:text-green-400" />
|
||||
{t('modals.importAgent.toolBuiltin', {
|
||||
type: tool.type,
|
||||
})}
|
||||
</p>
|
||||
)}
|
||||
{tool.status === 'reuse' && (
|
||||
<p className="flex items-center gap-2 text-sm">
|
||||
<CheckCircle2 className="h-4 w-4 text-green-600 dark:text-green-400" />
|
||||
{t('modals.importAgent.toolReuse', {
|
||||
name: tool.name || tool.type,
|
||||
})}
|
||||
</p>
|
||||
)}
|
||||
{tool.status === 'unavailable' && (
|
||||
<p className="text-destructive flex items-center gap-2 text-sm">
|
||||
<AlertTriangle className="h-4 w-4" />
|
||||
{t('modals.importAgent.toolUnavailable', {
|
||||
type: tool.type,
|
||||
})}
|
||||
</p>
|
||||
)}
|
||||
{tool.status === 'create' && (
|
||||
<div className="flex flex-col gap-2">
|
||||
<p className="text-sm">
|
||||
{t('modals.importAgent.toolCreate', {
|
||||
name: tool.name || tool.type,
|
||||
})}
|
||||
</p>
|
||||
{(tool.requires_secrets || []).map((field) => (
|
||||
<Input
|
||||
key={field}
|
||||
type="password"
|
||||
placeholder={t(
|
||||
'modals.importAgent.secretPlaceholder',
|
||||
{ field },
|
||||
)}
|
||||
value={toolSecrets[tool.key]?.[field] || ''}
|
||||
onChange={(e) =>
|
||||
setToolSecret(tool.key, field, e.target.value)
|
||||
}
|
||||
className="bg-card h-auto rounded-lg px-3 py-2 text-sm md:text-sm"
|
||||
/>
|
||||
))}
|
||||
</div>
|
||||
)}
|
||||
</div>
|
||||
))}
|
||||
</Section>
|
||||
)}
|
||||
|
||||
{plan.models.some((m) => m.status === 'create') && (
|
||||
<Section title={t('modals.importAgent.models')}>
|
||||
{plan.models
|
||||
.filter((m) => m.status === 'create')
|
||||
.map((m) => (
|
||||
<div key={m.display_name} className="flex flex-col gap-1">
|
||||
<p className="text-sm">
|
||||
{t('modals.importAgent.modelCustom', {
|
||||
name: m.display_name,
|
||||
})}
|
||||
</p>
|
||||
<Input
|
||||
type="password"
|
||||
placeholder={t('modals.importAgent.apiKeyPlaceholder')}
|
||||
value={modelKeys[m.display_name || ''] || ''}
|
||||
onChange={(e) =>
|
||||
setModelKeys((prev) => ({
|
||||
...prev,
|
||||
[m.display_name || '']: e.target.value,
|
||||
}))
|
||||
}
|
||||
className="bg-card h-auto rounded-lg px-3 py-2 text-sm md:text-sm"
|
||||
/>
|
||||
</div>
|
||||
))}
|
||||
</Section>
|
||||
)}
|
||||
|
||||
{error && <p className="text-destructive text-sm">{error}</p>}
|
||||
</div>
|
||||
)}
|
||||
</div>
|
||||
</Modal>
|
||||
);
|
||||
}
|
||||
|
||||
function Section({
|
||||
title,
|
||||
children,
|
||||
}: {
|
||||
title: string;
|
||||
children: React.ReactNode;
|
||||
}) {
|
||||
return (
|
||||
<div className="flex flex-col gap-2">
|
||||
<p className="text-foreground text-sm font-medium">{title}</p>
|
||||
<div className="flex flex-col gap-2">{children}</div>
|
||||
</div>
|
||||
);
|
||||
}
|
||||
@@ -0,0 +1,480 @@
|
||||
"""Tests for agent YAML export/import (application.api.user.agents.portability).
|
||||
|
||||
These exercise real SQL against the ephemeral ``pg_conn`` fixture, calling
|
||||
the dependency-injected serialize/plan/apply functions directly so the
|
||||
logic is covered without a running Flask app.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from unittest.mock import Mock
|
||||
|
||||
import pytest
|
||||
|
||||
from application.agents.default_tools import default_tool_id
|
||||
from application.api.user.agents.portability import (
|
||||
API_VERSION,
|
||||
AgentImportError,
|
||||
agent_to_yaml,
|
||||
apply_import,
|
||||
ensure_agent_slug,
|
||||
parse_agent_yaml,
|
||||
plan_import,
|
||||
serialize_agent,
|
||||
)
|
||||
from application.storage.db.repositories.agents import AgentsRepository
|
||||
from application.storage.db.repositories.prompts import PromptsRepository
|
||||
from application.storage.db.repositories.sources import SourcesRepository
|
||||
from application.storage.db.repositories.user_tools import UserToolsRepository
|
||||
|
||||
|
||||
pytestmark = pytest.mark.integration
|
||||
|
||||
|
||||
def _make_agent(conn, user, **kwargs):
|
||||
defaults = dict(description="d", chunks=2, retriever="classic")
|
||||
defaults.update(kwargs)
|
||||
return AgentsRepository(conn).create(user, kwargs.pop("name", "Bot"), "published", **defaults)
|
||||
|
||||
|
||||
def _doc(**spec):
|
||||
spec.setdefault("name", "Imported")
|
||||
return {
|
||||
"apiVersion": API_VERSION,
|
||||
"kind": "Agent",
|
||||
"metadata": {"slug": spec.pop("_slug", "imported")},
|
||||
"spec": spec,
|
||||
}
|
||||
|
||||
|
||||
def test_serialize_redacts_tool_secret(pg_conn):
|
||||
user = "u_redact"
|
||||
src = SourcesRepository(pg_conn).create("Docs", user_id=user, type="file")
|
||||
prompt = PromptsRepository(pg_conn).create(user, "P", "You are helpful")
|
||||
tool = UserToolsRepository(pg_conn).create(
|
||||
user,
|
||||
"brave",
|
||||
config={"encrypted_credentials": "SECRETBLOB", "region": "US"},
|
||||
custom_name="My Brave",
|
||||
display_name="Brave Search",
|
||||
description="search",
|
||||
config_requirements={"token": {"secret": True, "required": True}, "region": {}},
|
||||
actions=[],
|
||||
)
|
||||
agent = AgentsRepository(pg_conn).create(
|
||||
user,
|
||||
"Support Bot",
|
||||
"published",
|
||||
description="help desk",
|
||||
chunks=3,
|
||||
retriever="classic",
|
||||
prompt_id=str(prompt["id"]),
|
||||
source_id=str(src["id"]),
|
||||
tools=[str(tool["id"])],
|
||||
)
|
||||
|
||||
export = serialize_agent(pg_conn, agent, user)
|
||||
text = agent_to_yaml(export)
|
||||
|
||||
assert "SECRETBLOB" not in text
|
||||
assert "encrypted_credentials" not in text
|
||||
tool_entry = export["spec"]["tools"][0]
|
||||
assert tool_entry["type"] == "brave"
|
||||
assert tool_entry["requires_secrets"] == ["token"]
|
||||
assert "encrypted_credentials" not in tool_entry["config"]
|
||||
assert tool_entry["config"] == {"region": "US"}
|
||||
assert export["spec"]["prompt"] == {"name": "P", "content": "You are helpful"}
|
||||
assert export["spec"]["sources"][0]["name"] == "Docs"
|
||||
|
||||
|
||||
def test_export_assigns_unique_slug(pg_conn):
|
||||
user = "u_slug"
|
||||
a1 = AgentsRepository(pg_conn).create(user, "My Bot", "draft")
|
||||
a2 = AgentsRepository(pg_conn).create(user, "My Bot", "draft")
|
||||
|
||||
s1 = ensure_agent_slug(pg_conn, a1, user)
|
||||
a2_reload = AgentsRepository(pg_conn).get(str(a2["id"]), user)
|
||||
s2 = ensure_agent_slug(pg_conn, a2_reload, user)
|
||||
|
||||
assert s1 == "my-bot"
|
||||
assert s2 == "my-bot-2"
|
||||
|
||||
|
||||
def test_round_trip_same_user_idempotent(pg_conn):
|
||||
user = "u_round"
|
||||
src = SourcesRepository(pg_conn).create("KB", user_id=user, type="file")
|
||||
prompt = PromptsRepository(pg_conn).create(user, "Sys", "Be nice")
|
||||
scheduler_id = default_tool_id("scheduler")
|
||||
agent = AgentsRepository(pg_conn).create(
|
||||
user,
|
||||
"Bot",
|
||||
"published",
|
||||
description="d",
|
||||
chunks=2,
|
||||
retriever="classic",
|
||||
prompt_id=str(prompt["id"]),
|
||||
source_id=str(src["id"]),
|
||||
tools=[scheduler_id],
|
||||
)
|
||||
ensure_agent_slug(pg_conn, agent, user)
|
||||
agent = AgentsRepository(pg_conn).get(str(agent["id"]), user)
|
||||
|
||||
doc = parse_agent_yaml(agent_to_yaml(serialize_agent(pg_conn, agent, user)))
|
||||
r1 = apply_import(pg_conn, user, doc)
|
||||
r2 = apply_import(pg_conn, user, doc)
|
||||
|
||||
assert r1["agent_id"] == r2["agent_id"] == str(agent["id"]) # matched by id
|
||||
assert len(PromptsRepository(pg_conn).list_for_user(user)) == 1 # no dup prompt
|
||||
imported = AgentsRepository(pg_conn).get(r1["agent_id"], user)
|
||||
assert imported["status"] == "published" # update preserves the live status
|
||||
assert [str(s) for s in imported["extra_source_ids"]] == [str(src["id"])]
|
||||
assert imported["tools"] == [scheduler_id] # builtin passthrough
|
||||
|
||||
|
||||
def test_update_preserves_published_status(pg_conn):
|
||||
"""Re-importing over a published agent must not revert it to draft.
|
||||
|
||||
The agent's API key and any active users keep working; only the content
|
||||
is synced. A new agent (no match) still lands as a draft.
|
||||
"""
|
||||
user = "u_pub"
|
||||
agent = AgentsRepository(pg_conn).create(
|
||||
user,
|
||||
"Live Bot",
|
||||
"published",
|
||||
description="old",
|
||||
chunks=2,
|
||||
retriever="classic",
|
||||
)
|
||||
slug = ensure_agent_slug(pg_conn, agent, user)
|
||||
doc = {
|
||||
"apiVersion": API_VERSION,
|
||||
"kind": "Agent",
|
||||
"metadata": {"id": str(agent["id"]), "slug": slug},
|
||||
"spec": {"name": "Live Bot", "description": "new", "retriever": "classic"},
|
||||
}
|
||||
|
||||
result = apply_import(pg_conn, user, doc)
|
||||
|
||||
assert result["action"] == "updated"
|
||||
assert result["status"] == "published" # response reports the preserved status
|
||||
updated = AgentsRepository(pg_conn).get(str(agent["id"]), user)
|
||||
assert updated["status"] == "published" # stayed live
|
||||
assert updated["description"] == "new" # content still synced
|
||||
|
||||
|
||||
def test_import_by_slug_idempotent(pg_conn):
|
||||
user = "u_slug_import"
|
||||
doc = _doc(name="Slug Bot", _slug="slug-bot", retriever="classic")
|
||||
|
||||
r1 = apply_import(pg_conn, user, doc)
|
||||
r2 = apply_import(pg_conn, user, doc)
|
||||
|
||||
assert r1["action"] == "created"
|
||||
assert r2["action"] == "updated"
|
||||
assert r1["agent_id"] == r2["agent_id"]
|
||||
assert len(AgentsRepository(pg_conn).list_for_user(user)) == 1
|
||||
|
||||
|
||||
def test_import_missing_source_drafts_and_warns(pg_conn):
|
||||
user = "u_missing"
|
||||
doc = _doc(
|
||||
name="No Source Bot",
|
||||
_slug="no-source",
|
||||
sources=[{"name": "Nonexistent KB", "type": "file"}],
|
||||
)
|
||||
|
||||
result = apply_import(pg_conn, user, doc)
|
||||
|
||||
assert result["status"] == "draft" # new agents are created as drafts
|
||||
assert any("Nonexistent KB" in w for w in result["warnings"])
|
||||
agent = AgentsRepository(pg_conn).get(result["agent_id"], user)
|
||||
assert agent["status"] == "draft"
|
||||
assert list(agent["extra_source_ids"]) == []
|
||||
|
||||
|
||||
def test_import_matches_existing_source_by_name(pg_conn):
|
||||
user = "u_match"
|
||||
src = SourcesRepository(pg_conn).create("Product Docs", user_id=user, type="file")
|
||||
doc = _doc(
|
||||
name="Matcher",
|
||||
_slug="matcher",
|
||||
sources=[{"name": "product docs", "type": "file"}], # case-insensitive
|
||||
)
|
||||
|
||||
result = apply_import(pg_conn, user, doc)
|
||||
|
||||
agent = AgentsRepository(pg_conn).get(result["agent_id"], user)
|
||||
assert [str(s) for s in agent["extra_source_ids"]] == [str(src["id"])]
|
||||
assert result["warnings"] == []
|
||||
|
||||
|
||||
def test_plan_classifies_references(pg_conn):
|
||||
user = "u_plan"
|
||||
SourcesRepository(pg_conn).create("Known", user_id=user, type="file")
|
||||
tool = UserToolsRepository(pg_conn).create(
|
||||
user, "brave", custom_name="My Brave", display_name="Brave", config={},
|
||||
)
|
||||
doc = _doc(
|
||||
name="Plan Bot",
|
||||
sources=[{"name": "Known"}, {"name": "Unknown"}],
|
||||
tools=[
|
||||
{"type": "scheduler", "builtin": True},
|
||||
{"type": "brave", "name": "My Brave"},
|
||||
],
|
||||
)
|
||||
|
||||
plan = plan_import(pg_conn, user, doc)
|
||||
|
||||
statuses = {s["name"]: s["status"] for s in plan["sources"]}
|
||||
assert statuses == {"Known": "matched", "Unknown": "missing"}
|
||||
tool_statuses = [t["status"] for t in plan["tools"]]
|
||||
assert tool_statuses[0] == "builtin"
|
||||
assert tool_statuses[1] == "reuse"
|
||||
assert plan["tools"][1]["target_id"] == str(tool["id"])
|
||||
assert plan["target"]["action"] == "create"
|
||||
|
||||
|
||||
def test_apply_creates_tool_with_supplied_secret(pg_conn, monkeypatch):
|
||||
user = "u_tool"
|
||||
fake_tool = Mock()
|
||||
fake_tool.get_config_requirements.return_value = {
|
||||
"api_key": {"secret": True, "required": True, "label": "API Key"}
|
||||
}
|
||||
fake_tool.get_actions_metadata.return_value = []
|
||||
monkeypatch.setattr(
|
||||
"application.api.user.agents.portability._tool_instance",
|
||||
lambda tool_type: fake_tool,
|
||||
)
|
||||
doc = _doc(
|
||||
name="Tool Bot",
|
||||
_slug="tool-bot",
|
||||
tools=[
|
||||
{
|
||||
"type": "brave",
|
||||
"name": "My API",
|
||||
"config": {},
|
||||
"requires_secrets": ["api_key"],
|
||||
}
|
||||
],
|
||||
)
|
||||
|
||||
# No secret supplied -> tool skipped, agent still created.
|
||||
res = apply_import(pg_conn, user, doc)
|
||||
agent = AgentsRepository(pg_conn).get(res["agent_id"], user)
|
||||
assert agent["tools"] == []
|
||||
assert any("api_key" in w for w in res["warnings"])
|
||||
|
||||
# Secret supplied via resolution -> tool created and linked, secret encrypted.
|
||||
res2 = apply_import(
|
||||
pg_conn,
|
||||
user,
|
||||
doc,
|
||||
resolution={"tools": {"tool-0": {"secrets": {"api_key": "PLAINTEXT_KEY"}}}},
|
||||
)
|
||||
agent2 = AgentsRepository(pg_conn).get(res2["agent_id"], user)
|
||||
assert len(agent2["tools"]) == 1
|
||||
tool_row = UserToolsRepository(pg_conn).get_any(agent2["tools"][0], user)
|
||||
assert "PLAINTEXT_KEY" not in str(tool_row["config"])
|
||||
assert "encrypted_credentials" in tool_row["config"]
|
||||
|
||||
|
||||
def test_parse_rejects_bad_documents():
|
||||
with pytest.raises(AgentImportError):
|
||||
parse_agent_yaml("kind: NotAgent\napiVersion: docsgpt.arc53.com/v1\nspec:\n name: x\n")
|
||||
with pytest.raises(AgentImportError):
|
||||
parse_agent_yaml("kind: Agent\napiVersion: other/v1\nspec:\n name: x\n")
|
||||
with pytest.raises(AgentImportError):
|
||||
parse_agent_yaml("kind: Agent\napiVersion: docsgpt.arc53.com/v1\nspec: {}\n")
|
||||
with pytest.raises(AgentImportError):
|
||||
parse_agent_yaml(
|
||||
"kind: Agent\napiVersion: docsgpt.arc53.com/v1\n"
|
||||
"spec:\n name: x\n agent_type: workflow\n"
|
||||
)
|
||||
|
||||
|
||||
def test_parse_rejects_yaml_aliases():
|
||||
bomb = (
|
||||
"kind: Agent\napiVersion: docsgpt.arc53.com/v1\n"
|
||||
"spec:\n name: &a x\n description: *a\n"
|
||||
)
|
||||
with pytest.raises(AgentImportError):
|
||||
parse_agent_yaml(bomb)
|
||||
|
||||
|
||||
def test_update_clears_removed_json_schema(pg_conn):
|
||||
"""On update the YAML is authoritative — removing json_schema clears it."""
|
||||
user = "u_clear"
|
||||
agent = AgentsRepository(pg_conn).create(
|
||||
user,
|
||||
"Clear Bot",
|
||||
"published",
|
||||
description="d",
|
||||
chunks=2,
|
||||
retriever="classic",
|
||||
json_schema={"type": "object"},
|
||||
)
|
||||
slug = ensure_agent_slug(pg_conn, agent, user)
|
||||
doc = {
|
||||
"apiVersion": API_VERSION,
|
||||
"kind": "Agent",
|
||||
"metadata": {"slug": slug},
|
||||
"spec": {"name": "Clear Bot", "retriever": "classic", "json_schema": None},
|
||||
}
|
||||
|
||||
result = apply_import(pg_conn, user, doc)
|
||||
|
||||
assert result["action"] == "updated"
|
||||
updated = AgentsRepository(pg_conn).get(str(agent["id"]), user)
|
||||
assert updated["json_schema"] is None
|
||||
|
||||
|
||||
def test_custom_tool_dedup_across_reimport(pg_conn):
|
||||
user = "u_dedup"
|
||||
tool = UserToolsRepository(pg_conn).create(
|
||||
user, "brave", custom_name="My Brave", display_name="Brave", config={},
|
||||
)
|
||||
agent = AgentsRepository(pg_conn).create(
|
||||
user,
|
||||
"Dedup Bot",
|
||||
"published",
|
||||
description="d",
|
||||
chunks=2,
|
||||
retriever="classic",
|
||||
tools=[str(tool["id"])],
|
||||
)
|
||||
ensure_agent_slug(pg_conn, agent, user)
|
||||
agent = AgentsRepository(pg_conn).get(str(agent["id"]), user)
|
||||
doc = parse_agent_yaml(agent_to_yaml(serialize_agent(pg_conn, agent, user)))
|
||||
|
||||
apply_import(pg_conn, user, doc)
|
||||
apply_import(pg_conn, user, doc)
|
||||
|
||||
assert len(UserToolsRepository(pg_conn).list_for_user(user)) == 1 # reused, never duplicated
|
||||
final = AgentsRepository(pg_conn).get(str(agent["id"]), user)
|
||||
assert final["tools"] == [str(tool["id"])]
|
||||
|
||||
|
||||
def test_import_rejects_ssrf_tool_url(pg_conn, monkeypatch):
|
||||
user = "u_ssrf"
|
||||
fake_tool = Mock()
|
||||
fake_tool.get_config_requirements.return_value = {}
|
||||
fake_tool.get_actions_metadata.return_value = []
|
||||
monkeypatch.setattr(
|
||||
"application.api.user.agents.portability._tool_instance",
|
||||
lambda tool_type: fake_tool,
|
||||
)
|
||||
doc = _doc(
|
||||
name="SSRF Bot",
|
||||
_slug="ssrf",
|
||||
tools=[
|
||||
{
|
||||
"type": "mcp_tool",
|
||||
"name": "evil",
|
||||
"config": {"server_url": "http://169.254.169.254/latest/meta-data"},
|
||||
}
|
||||
],
|
||||
)
|
||||
|
||||
result = apply_import(pg_conn, user, doc)
|
||||
|
||||
agent = AgentsRepository(pg_conn).get(result["agent_id"], user)
|
||||
assert agent["tools"] == []
|
||||
assert any("unsafe" in w.lower() for w in result["warnings"])
|
||||
|
||||
|
||||
def test_import_rejects_unowned_source_mapping(pg_conn):
|
||||
owner, attacker = "u_owner", "u_attacker"
|
||||
src = SourcesRepository(pg_conn).create("Secret KB", user_id=owner, type="file")
|
||||
doc = _doc(name="IDOR Bot", _slug="idor", sources=[{"name": "Their Source"}])
|
||||
|
||||
# Attacker maps the missing source to the owner's source id.
|
||||
result = apply_import(
|
||||
pg_conn, attacker, doc, resolution={"sources": {"Their Source": str(src["id"])}}
|
||||
)
|
||||
|
||||
agent = AgentsRepository(pg_conn).get(result["agent_id"], attacker)
|
||||
assert list(agent["extra_source_ids"]) == [] # not linked
|
||||
assert any("not yours" in w for w in result["warnings"])
|
||||
|
||||
|
||||
def test_apply_tolerates_plan_shaped_resolution(pg_conn):
|
||||
"""Echoing the plan's list-shaped resolution back must not crash apply."""
|
||||
user = "u_shape"
|
||||
doc = _doc(name="Shape Bot", _slug="shape", sources=[{"name": "X"}])
|
||||
|
||||
result = apply_import(
|
||||
pg_conn, user, doc, resolution={"sources": [], "tools": [], "models": []}
|
||||
)
|
||||
|
||||
assert result["action"] == "created"
|
||||
|
||||
|
||||
def test_update_clears_models_and_resets_prompt(pg_conn):
|
||||
"""A re-imported file that drops models and prompt clears them on update."""
|
||||
user = "u_clear2"
|
||||
prompt = PromptsRepository(pg_conn).create(user, "Sys", "Be terse")
|
||||
agent = AgentsRepository(pg_conn).create(
|
||||
user,
|
||||
"Model Bot",
|
||||
"published",
|
||||
description="d",
|
||||
chunks=2,
|
||||
retriever="classic",
|
||||
prompt_id=str(prompt["id"]),
|
||||
models=["gpt-4o"],
|
||||
default_model_id="gpt-4o",
|
||||
)
|
||||
slug = ensure_agent_slug(pg_conn, agent, user)
|
||||
doc = {
|
||||
"apiVersion": API_VERSION,
|
||||
"kind": "Agent",
|
||||
"metadata": {"slug": slug},
|
||||
"spec": {
|
||||
"name": "Model Bot",
|
||||
"retriever": "classic",
|
||||
"prompt": "default",
|
||||
"model": {"default": "", "available": []},
|
||||
},
|
||||
}
|
||||
|
||||
apply_import(pg_conn, user, doc)
|
||||
|
||||
updated = AgentsRepository(pg_conn).get(str(agent["id"]), user)
|
||||
assert not updated["models"]
|
||||
assert not updated["default_model_id"]
|
||||
assert updated["prompt_id"] is None # reset to default
|
||||
|
||||
|
||||
def test_tool_skipped_when_encryption_fails(pg_conn, monkeypatch):
|
||||
user = "u_encfail"
|
||||
fake_tool = Mock()
|
||||
fake_tool.get_config_requirements.return_value = {
|
||||
"api_key": {"secret": True, "required": True}
|
||||
}
|
||||
fake_tool.get_actions_metadata.return_value = []
|
||||
monkeypatch.setattr(
|
||||
"application.api.user.agents.portability._tool_instance",
|
||||
lambda tool_type: fake_tool,
|
||||
)
|
||||
# Force credential encryption to fail (returns "").
|
||||
monkeypatch.setattr(
|
||||
"application.api.user.tools.routes.encrypt_credentials",
|
||||
lambda creds, user_id: "",
|
||||
)
|
||||
doc = _doc(
|
||||
name="EncFail Bot",
|
||||
_slug="encfail",
|
||||
tools=[{"type": "brave", "name": "X", "config": {}, "requires_secrets": ["api_key"]}],
|
||||
)
|
||||
|
||||
result = apply_import(
|
||||
pg_conn,
|
||||
user,
|
||||
doc,
|
||||
resolution={"tools": {"tool-0": {"secrets": {"api_key": "K"}}}},
|
||||
)
|
||||
|
||||
agent = AgentsRepository(pg_conn).get(result["agent_id"], user)
|
||||
assert agent["tools"] == []
|
||||
assert any("encryption failed" in w for w in result["warnings"])
|
||||
Reference in new issue
Block a user