aigc-home/calliope-backend/tests/test_agent_roles.py

"""Named-agent sessions (HEL-613): the five AIGC roles. agent_key is stamped
at creation, filters the session list, injects a per-agent persona + agent-chat
mode section into the assembled system prompt, and routes the main loop's LLM
through the agent's bound role (models page 智能体默认模型 panel)."""
from __future__ import annotations

import asyncio

from calliope.agent.harness.registry import ToolContext

# ─────────────────────────────────────────────────────────────────────────
# Registry
# ─────────────────────────────────────────────────────────────────────────


def test_agents_endpoint_lists_five_roles(client):
    agents = client.get("/api/agent/agents").json()
    assert [a["key"] for a in agents] == [
        "script",
        "style",
        "storyboard",
        "assets",
        "shot",
    ]
    names = {a["name"] for a in agents}
    assert names == {"剧本创作", "视觉设计", "分镜设计", "素材资产", "镜头脚本"}
    # LLM roles reuse the models-page five-role binding surface.
    assert {a["llm_role"] for a in agents} == {"story", "style", "script", "assets", "video"}
    # Personas are server-side — never shipped to the browser.
    assert all("persona" not in a for a in agents)
    assert all(isinstance(a.get("skill_count"), int) for a in agents)


# ─────────────────────────────────────────────────────────────────────────
# Session create / list semantics
# ─────────────────────────────────────────────────────────────────────────


def test_create_session_with_agent_key(client):
    s = client.post("/api/agent/sessions", json={"agent_key": "style"}).json()
    assert s["agent_key"] == "style"
    assert s["origin"] == "chat"


def test_agent_key_rejects_unknown_value(client):
    r = client.post("/api/agent/sessions", json={"agent_key": "director"})
    assert r.status_code == 422


def test_agent_sessions_hidden_from_generic_list(client):
    plain = client.post("/api/agent/sessions", json={"title": "canvas chat"}).json()
    named = client.post("/api/agent/sessions", json={"agent_key": "script"}).json()
    listed = client.get("/api/agent/sessions").json()
    ids = {s["id"] for s in listed}
    assert plain["id"] in ids
    assert named["id"] not in ids


def test_agent_key_filter_returns_only_that_agent(client):
    a = client.post("/api/agent/sessions", json={"agent_key": "script"}).json()
    client.post("/api/agent/sessions", json={"agent_key": "style"})
    client.post("/api/agent/sessions", json={"title": "plain"})
    listed = client.get("/api/agent/sessions?agent_key=script").json()
    assert {s["id"] for s in listed} == {a["id"]}
    assert all(s["agent_key"] == "script" for s in listed)


def test_plain_session_has_no_agent_key(client):
    s = client.post("/api/agent/sessions", json={}).json()
    assert s["agent_key"] is None


# ─────────────────────────────────────────────────────────────────────────
# Prompt assembly: persona + mode per agent_key
# ─────────────────────────────────────────────────────────────────────────


def _assemble(agent_key: str | None) -> str:
    from calliope.agent.harness import get_prompts

    ctx = ToolContext(session_id=1, project_id=None, origin="chat", agent_key=agent_key)
    return asyncio.run(get_prompts().assemble(ctx))


def test_persona_section_injected_per_agent(client):
    prompt = _assemble("script")
    assert "「剧本创作」" in prompt
    assert "编剧" in prompt
    assert "SESSION MODE: AGENT CHAT" in prompt
    # The generic sandbox mode text must be replaced for named agents.
    assert "GENERATING IN SANDBOX" not in prompt


def test_each_agent_gets_its_own_persona(client):
    for key, marker in (
        ("style", "视觉"),
        ("storyboard", "分镜"),
        ("assets", "资产"),
        ("shot", "镜头脚本"),
    ):
        assert marker in _assemble(key), key


def test_plain_prompt_unchanged(client):
    prompt = _assemble(None)
    assert "Calliope's production agent" in prompt
    assert "SESSION MODE: SANDBOX" in prompt
    assert "AGENT CHAT" not in prompt


def test_agent_chat_mentions_section_suppressed(client):
    assert "Tagged workflows" not in _assemble("assets")
    assert "Tagged workflows" in _assemble(None)


# ─────────────────────────────────────────────────────────────────────────
# LLM routing: the main loop uses the agent's bound role
# ─────────────────────────────────────────────────────────────────────────


class _RoleSpyClient:
    instances: list["_RoleSpyClient"] = []

    def __init__(self, role: str | None = None) -> None:
        self.role = role
        self.requests: list[list[dict]] = []
        _RoleSpyClient.instances.append(self)

    @classmethod
    def for_role(cls, role: str, *, timeout: float = 120.0) -> "_RoleSpyClient":
        return cls(role=role)

    def chat_stream(self, messages, temperature=0.4, tools=None):
        self.requests.append(list(messages))
        return self

    def __aiter__(self):
        return self

    async def __anext__(self):
        raise StopAsyncIteration

    async def close(self):
        pass


def test_run_turn_uses_agent_llm_role_and_prompt(client):
    """End-to-end through run_turn with a scripted LLM: the client must be
    built for the agent's role (story ← 剧本创作) and receive the per-agent
    persona in the system message."""
    import calliope.agent.harness.loop as loop_mod

    session = client.post(
        "/api/agent/sessions", json={"agent_key": "script", "title": "剧本会话"}
    ).json()

    from calliope.agent.harness import log as session_log

    session_log.append_event(
        session["id"], session_log.USER_MESSAGE, {"content": "写个三幕式大纲"}
    )

    orig = loop_mod.LLMClient
    _RoleSpyClient.instances = []
    loop_mod.LLMClient = _RoleSpyClient
    try:
        ctx = ToolContext(
            session_id=session["id"], project_id=None, origin="chat", agent_key="script"
        )
        asyncio.run(loop_mod.run_turn(ctx, []))
    finally:
        loop_mod.LLMClient = orig

    assert _RoleSpyClient.instances, "no LLM client was built"
    spy = _RoleSpyClient.instances[0]
    assert spy.role == "story"
    system = next(
        (m for m in spy.requests[0] if m.get("role") == "system"), None
    )
    assert system is not None
    assert "「剧本创作」" in system["content"]
    assert "SESSION MODE: AGENT CHAT" in system["content"]


def test_plain_run_turn_keeps_main_role(client):
    import calliope.agent.harness.loop as loop_mod

    session = client.post("/api/agent/sessions", json={}).json()
    from calliope.agent.harness import log as session_log

    session_log.append_event(session["id"], session_log.USER_MESSAGE, {"content": "hi"})

    orig = loop_mod.LLMClient
    _RoleSpyClient.instances = []
    loop_mod.LLMClient = _RoleSpyClient
    try:
        ctx = ToolContext(session_id=session["id"], project_id=None, origin="chat")
        asyncio.run(loop_mod.run_turn(ctx, []))
    finally:
        loop_mod.LLMClient = orig

    assert _RoleSpyClient.instances[0].role == "main"