SupplyAgent/api/services/agent_catalog.py

"""Read-only catalog for the business Agents shown in the web console."""
from __future__ import annotations

from importlib import import_module
from pathlib import Path
from typing import Any

from supply_infra.agent_prompt_injection import (
    get_agent_document_injection,
    save_agent_document_injection,
)

_PROJECT_ROOT = Path(__file__).resolve().parents[2]

_AGENTS: dict[str, dict[str, str]] = {
    "demand_belong_category_agent": {
        "display_name": "分配 Agent",
        "subtitle": "语义归属与分类挂载",
        "prompt_path": "agents/demand_belong_category_agent/prompt/system_prompt.md",
        "tools_module": "agents.demand_belong_category_agent.tools",
    },
    "demand_grade_agent": {
        "display_name": "分等级 Agent",
        "subtitle": "证据融合与优先级判断",
        "prompt_path": "agents/demand_grade_agent/prompt/system_prompt.md",
        "tools_module": "agents.demand_grade_agent.tools",
    },
    "demand_video_expand_agent": {
        "display_name": "拓展 Agent",
        "subtitle": "视频点位与需求拓展",
        "prompt_path": "agents/demand_video_expand_agent/prompt/system_prompt.md",
        "tools_module": "agents.demand_video_expand_agent.tools",
    },
}


def get_agent_detail(agent_name: str) -> dict[str, Any] | None:
    """Return the current system prompt and registered tool definitions."""
    spec = _AGENTS.get(agent_name)
    if spec is None:
        return None

    prompt_path = _PROJECT_ROOT / spec["prompt_path"]
    prompt = prompt_path.read_text(encoding="utf-8")

    tools_module = import_module(spec["tools_module"])
    tools = []
    for tool_fn in tools_module.ALL_TOOLS:
        tools.append(
            {
                "name": getattr(tool_fn, "_tool_name", tool_fn.__name__),
                "description": getattr(
                    tool_fn,
                    "_tool_description",
                    (tool_fn.__doc__ or "").strip(),
                ),
                "parameters": getattr(tool_fn, "_tool_parameters", {}),
            }
        )

    return {
        "name": agent_name,
        "display_name": spec["display_name"],
        "subtitle": spec["subtitle"],
        "system_prompt": prompt,
        "document_injection": get_agent_document_injection(agent_name),
        "tools": tools,
    }


def update_agent_document_injection(
    agent_name: str,
    *,
    content: str,
    enabled: bool,
) -> dict[str, Any] | None:
    """Persist the editable document injection for an allow-listed Agent."""
    if agent_name not in _AGENTS:
        return None
    return save_agent_document_injection(
        agent_name,
        content=content,
        enabled=enabled,
    )