harness_zbt/backend/src/zbt/harness/kernel.py

"""统一 Agent 内核：同一编排能力承载客户与运营两种人格。"""

from dataclasses import dataclass
from typing import Any, cast
from uuid import UUID, uuid4

from langchain.agents import create_agent
from langchain_core.messages import AIMessage
from langchain_core.runnables import RunnableConfig
from langsmith import Client, tracing_context

from zbt.core.config import Settings
from zbt.domains.identity.models import AdminUser, H5User
from zbt.harness.model_gateway import ModelGateway
from zbt.harness.policy import HarnessPolicyEngine
from zbt.harness.registries import PersonaRegistry, PromptRegistry, SkillRegistry
from zbt.harness.schemas import (
    AgentAction,
    AgentReply,
    AgentStructuredOutput,
    ProductRecommendationsBlock,
    ProductView,
)
from zbt.harness.tooling import ToolExecutionContext, ToolRegistry


@dataclass(frozen=True)
class ConversationTurn:
    role: str
    content: str


@dataclass(frozen=True)
class AgentInvocation:
    persona: str
    message: str
    h5_user: H5User | None = None
    admin_user: AdminUser | None = None
    history: tuple[ConversationTurn, ...] = ()


class AgentKernel:
    def __init__(
        self,
        *,
        settings: Settings,
        model_gateway: ModelGateway,
        tools: ToolRegistry,
        prompts: PromptRegistry | None = None,
        skills: SkillRegistry | None = None,
        personas: PersonaRegistry | None = None,
        policy: HarnessPolicyEngine | None = None,
    ) -> None:
        self._settings = settings
        self._model_gateway = model_gateway
        self._tools = tools
        self._prompts = prompts or PromptRegistry()
        self._skills = skills or SkillRegistry()
        self._personas = personas or PersonaRegistry()
        self._policy = policy or HarnessPolicyEngine()
        self._langsmith_client = (
            Client(
                api_url=settings.langsmith_endpoint,
                api_key=settings.langsmith_api_key,
            )
            if settings.langsmith_tracing and settings.langsmith_api_key
            else None
        )

    def reply(self, invocation: AgentInvocation) -> AgentReply:
        persona = self._personas.get(invocation.persona)
        system_prompt = self._compose_prompt(persona.prompt_id, persona.skill_ids)
        context = ToolExecutionContext(
            persona=persona.id,
            h5_user=invocation.h5_user,
            admin_user=invocation.admin_user,
        )
        tools = self._tools.build(
            names=persona.tool_names,
            context=context,
            policy_engine=self._policy,
        )
        agent = create_agent(
            model=self._model_gateway.chat_model(persona.id),
            tools=tools,
            system_prompt=system_prompt,
            name=f"zhibaotong-{persona.id}",
        )
        trace_id = uuid4()
        messages = [
            {"role": turn.role, "content": turn.content} for turn in invocation.history[-12:]
        ]
        messages.append({"role": "user", "content": invocation.message})
        with tracing_context(
            project_name=self._settings.langsmith_project,
            enabled=self._langsmith_client is not None,
            client=self._langsmith_client,
            tags=["zhibaotong", "s1", persona.id],
            metadata={
                "persona": persona.id,
                "channel": ("customer_h5" if persona.id == "customer" else "operation_admin"),
                "principal_id": context.principal_id,
            },
        ):
            run_config: RunnableConfig = {"run_id": trace_id}
            result = agent.invoke(
                cast(Any, {"messages": messages}),
                config=run_config,
            )
        structured = self._structured_output(
            cast(dict[str, Any], result),
            context,
        )
        blocks = [block for execution in context.executions for block in execution.result.blocks]
        recommendation = self._recommendation_block(
            context,
            structured.selected_product_codes,
        )
        if recommendation is not None:
            blocks.insert(0, recommendation)
        actions = [
            action for execution in context.executions for action in execution.result.actions
        ]
        if structured.suggested_action is not None:
            actions.append(structured.suggested_action)
        unique_actions = _deduplicate_actions(actions)
        return AgentReply(
            text=structured.message,
            cards=[block.model_dump(mode="json") for block in blocks],
            trace_id=str(trace_id),
            trace_url=self._trace_url(trace_id),
            actions=[action.model_dump(mode="json") for action in unique_actions],
            invoked_tools=tuple(execution.name for execution in context.executions),
        )

    def _compose_prompt(
        self,
        prompt_id: str,
        skill_ids: tuple[str, ...],
    ) -> str:
        prompt = self._prompts.get(prompt_id)
        skill_sections = []
        for skill_id in skill_ids:
            skill = self._skills.get(skill_id)
            skill_sections.append(f"## Skill：{skill.name}\n\n{skill.instructions}")
        return f"{prompt}\n\n# 已加载技能\n\n" + "\n\n".join(skill_sections)

    @staticmethod
    def _structured_output(
        result: dict[str, Any],
        context: ToolExecutionContext,
    ) -> AgentStructuredOutput:
        messages = result.get("messages", [])
        final = messages[-1] if messages else None
        text = (
            str(final.content)
            if isinstance(final, AIMessage)
            else "暂时无法生成有效回答，请稍后重试。"
        )
        selected_codes: list[str] = []
        for execution in context.executions:
            if execution.name != "list_available_products":
                continue
            raw_items = execution.result.data.get("items", [])
            if not isinstance(raw_items, list):
                continue
            for item in raw_items:
                if not isinstance(item, dict):
                    continue
                code = str(item.get("product_code", ""))
                name = str(item.get("name", ""))
                if code and (code in text or name in text):
                    selected_codes.append(code)
        return AgentStructuredOutput(
            message=text,
            selected_product_codes=list(dict.fromkeys(selected_codes))[:2],
        )

    @staticmethod
    def _recommendation_block(
        context: ToolExecutionContext,
        selected_codes: list[str],
    ) -> ProductRecommendationsBlock | None:
        if not selected_codes:
            return None
        product_items: list[dict[str, Any]] = []
        for execution in context.executions:
            if execution.name != "list_available_products":
                continue
            raw_items = execution.result.data.get("items", [])
            if isinstance(raw_items, list):
                product_items.extend(item for item in raw_items if isinstance(item, dict))
        selected = [
            item
            for code in selected_codes
            for item in product_items
            if item.get("product_code") == code
        ][:2]
        if not selected:
            return None
        views = [
            ProductView.model_validate(
                {
                    "product_id": item["product_id"],
                    "product_code": item["product_code"],
                    "name": item["name"],
                    "category": item["category"],
                    "summary": item["summary"],
                    "product_version_id": item["product_version_id"],
                    "plans": item["plans"],
                }
            )
            for item in selected
        ]
        return ProductRecommendationsBlock(items=views)

    def _trace_url(self, trace_id: UUID) -> str | None:
        if self._langsmith_client is None:
            return None
        try:
            run = self._langsmith_client.read_run(trace_id)
            return self._langsmith_client.get_run_url(
                run=run,
                project_name=self._settings.langsmith_project,
            )
        except Exception:
            # Trace 上传可能有短暂延迟，运行详情接口会再次尝试解析。
            return None

    def trace_url(self, trace_id: str) -> str | None:
        try:
            return self._trace_url(UUID(trace_id))
        except ValueError:
            return None


def _deduplicate_actions(actions: list[AgentAction]) -> list[AgentAction]:
    unique: list[AgentAction] = []
    seen: set[tuple[str, str]] = set()
    for action in actions:
        key = (action.type, str(sorted(action.payload.items())))
        if key not in seen:
            seen.add(key)
            unique.append(action)
    return unique