harness_zbt/backend/app/domains/agent/service.py

from collections.abc import Callable, Iterator
from dataclasses import dataclass
from datetime import datetime
from secrets import randbelow
from typing import Any

from app.core.errors import AppError
from app.core.identifiers import new_ulid
from app.domains.agent.models import AgentMessage, AgentRun, AgentThread
from app.domains.agent.repository import AgentRepository
from app.domains.agent.runtime import AgentRuntime
from app.domains.identity.models import AdminUser, H5User
from app.harness.events import RuntimeEvent
from app.harness.kernel import AgentInvocation, ConversationTurn
from app.harness.schemas import AgentReply
from app.infrastructure.redis.agent_state import (
    AgentStateStore,
    InMemoryAgentStateStore,
)


@dataclass(frozen=True)
class PreparedMessage:
    thread: AgentThread
    text: str
    run_id: str
    service_no: str
    started_at: datetime
    history: tuple[ConversationTurn, ...]
    h5_user: H5User | None
    admin_user: AdminUser | None


class AgentThreadService:
    """客户与运营人格共用的会话、消息、运行事件和限流服务。"""

    def __init__(
        self,
        repository: AgentRepository,
        clock: Callable[[], datetime],
        runtime: AgentRuntime,
        state_store: AgentStateStore | None = None,
    ) -> None:
        self._repository = repository
        self._clock = clock
        self._runtime = runtime
        self._state_store = state_store or InMemoryAgentStateStore()

    def create_customer_thread(
        self,
        user: H5User,
        *,
        title: str | None,
    ) -> dict[str, Any]:
        return self._create_thread(
            owner_type="H5_USER",
            owner_id=user.id,
            persona="customer",
            title=title or "新的保障咨询",
        )

    def create_operation_thread(
        self,
        user: AdminUser,
        *,
        title: str | None,
    ) -> dict[str, Any]:
        return self._create_thread(
            owner_type="ADMIN_USER",
            owner_id=user.id,
            persona="operation",
            title=title or "新的运营分析",
        )

    def send_customer_message(
        self,
        user: H5User,
        *,
        thread_id: str,
        text: str,
    ) -> dict[str, Any]:
        prepared = self._prepare_message(
            self._owned_thread("H5_USER", user.id, thread_id),
            text,
            h5_user=user,
        )
        return self._execute(prepared)

    def stream_customer_message(
        self,
        user: H5User,
        *,
        thread_id: str,
        text: str,
    ) -> Iterator[RuntimeEvent]:
        prepared = self._prepare_message(
            self._owned_thread("H5_USER", user.id, thread_id),
            text,
            h5_user=user,
        )
        yield self._emit(
            prepared,
            "run.started",
            1,
            {"persona": prepared.thread.persona},
        )
        try:
            reply = self._invoke(prepared)
            result = self._complete(prepared, reply)
        except AppError as error:
            self._save_failed_run(prepared, error)
            yield self._emit(
                prepared,
                "run.failed",
                2,
                {
                    "code": error.code,
                    "message": error.message,
                    "retryable": error.retryable,
                },
            )
            return
        sequence = 2
        for tool_name in reply.invoked_tools:
            yield self._emit(
                prepared,
                "tool.completed",
                sequence,
                {"tool_name": tool_name},
            )
            sequence += 1
        if reply.cards or reply.actions:
            yield self._emit(
                prepared,
                "ui.ready",
                sequence,
                {
                    "cards": reply.cards,
                    "actions": reply.actions,
                },
            )
            sequence += 1
        yield self._emit(
            prepared,
            "run.completed",
            sequence,
            {"result": result},
        )

    def send_operation_message(
        self,
        user: AdminUser,
        *,
        thread_id: str,
        text: str,
    ) -> dict[str, Any]:
        prepared = self._prepare_message(
            self._owned_thread("ADMIN_USER", user.id, thread_id),
            text,
            admin_user=user,
        )
        return self._execute(prepared)

    def list_customer_messages(
        self,
        user: H5User,
        *,
        thread_id: str,
    ) -> dict[str, Any]:
        self._owned_thread("H5_USER", user.id, thread_id)
        return self._message_data(thread_id)

    def list_operation_messages(
        self,
        user: AdminUser,
        *,
        thread_id: str,
    ) -> dict[str, Any]:
        self._owned_thread("ADMIN_USER", user.id, thread_id)
        return self._message_data(thread_id)

    def get_customer_run(self, user: H5User, *, run_id: str) -> dict[str, Any]:
        return self._get_run("H5_USER", user.id, run_id)

    def get_operation_run(
        self,
        user: AdminUser,
        *,
        run_id: str,
    ) -> dict[str, Any]:
        return self._get_run("ADMIN_USER", user.id, run_id)

    def list_operation_runs(
        self,
        user: AdminUser,
        *,
        limit: int = 20,
    ) -> dict[str, Any]:
        threads = self._repository.list_threads("ADMIN_USER", user.id)
        runs = self._repository.list_runs(
            tuple(thread.id for thread in threads),
            limit,
        )
        return {
            "items": [
                {
                    "id": run.id,
                    "thread_id": run.thread_id,
                    "persona": run.persona,
                    "status": run.status,
                    "trace_id": run.trace_id,
                    "trace_url": self._resolve_trace_url(run),
                    "input": run.input,
                    "error_code": run.error_code,
                    "created_at": run.created_at,
                    "completed_at": run.completed_at,
                }
                for run in runs
            ],
            "total": len(runs),
        }

    def list_customer_run_events(
        self,
        user: H5User,
        *,
        run_id: str,
    ) -> list[RuntimeEvent]:
        self._get_run("H5_USER", user.id, run_id)
        return self._state_store.list_events(run_id)

    def _create_thread(
        self,
        *,
        owner_type: str,
        owner_id: str,
        persona: str,
        title: str,
    ) -> dict[str, Any]:
        now = self._clock()
        thread = AgentThread(
            id=new_ulid(),
            owner_type=owner_type,
            owner_id=owner_id,
            persona=persona,
            title=title.strip(),
            status="ACTIVE",
            created_at=now,
            updated_at=now,
        )
        self._repository.save_thread(thread)
        self._state_store.touch_thread(
            thread.id,
            owner_type=thread.owner_type,
            owner_id=thread.owner_id,
            persona=thread.persona,
        )
        return {
            "id": thread.id,
            "persona": thread.persona,
            "title": thread.title,
            "status": thread.status,
        }

    def _prepare_message(
        self,
        thread: AgentThread,
        text: str,
        *,
        h5_user: H5User | None = None,
        admin_user: AdminUser | None = None,
    ) -> PreparedMessage:
        principal = h5_user or admin_user
        if principal is None:
            raise AppError("AGENT_PRINCIPAL_REQUIRED", "Agent缺少登录身份", 401)
        if not self._state_store.allow_request(principal.id):
            raise AppError(
                "AGENT_RATE_LIMITED",
                "请求过于频繁，请稍后再试",
                429,
                retryable=True,
            )
        history = tuple(
            ConversationTurn(
                role="assistant" if message.role == "ASSISTANT" else "user",
                content=message.content,
            )
            for message in self._repository.list_messages(thread.id)[-12:]
        )
        now = self._clock()
        run_id = new_ulid()
        service_no = f"ZBT-{now:%Y%m%d}-{randbelow(10_000):04d}"
        self._repository.save_message(
            AgentMessage(
                id=new_ulid(),
                thread_id=thread.id,
                role="USER",
                content=text,
                metadata={"run_id": run_id, "service_no": service_no},
                created_at=now,
            )
        )
        self._state_store.touch_thread(
            thread.id,
            owner_type=thread.owner_type,
            owner_id=thread.owner_id,
            persona=thread.persona,
        )
        return PreparedMessage(
            thread=thread,
            text=text,
            run_id=run_id,
            service_no=service_no,
            started_at=now,
            history=history,
            h5_user=h5_user,
            admin_user=admin_user,
        )

    def _execute(self, prepared: PreparedMessage) -> dict[str, Any]:
        self._emit(
            prepared,
            "run.started",
            1,
            {"persona": prepared.thread.persona},
        )
        try:
            reply = self._invoke(prepared)
            result = self._complete(prepared, reply)
        except AppError as error:
            self._save_failed_run(prepared, error)
            self._emit(
                prepared,
                "run.failed",
                2,
                {
                    "code": error.code,
                    "message": error.message,
                    "retryable": error.retryable,
                },
            )
            raise
        sequence = 2
        for tool_name in reply.invoked_tools:
            self._emit(
                prepared,
                "tool.completed",
                sequence,
                {"tool_name": tool_name},
            )
            sequence += 1
        if reply.cards or reply.actions:
            self._emit(
                prepared,
                "ui.ready",
                sequence,
                {"cards": reply.cards, "actions": reply.actions},
            )
            sequence += 1
        self._emit(
            prepared,
            "run.completed",
            sequence,
            {"result": result},
        )
        return result

    def _invoke(self, prepared: PreparedMessage) -> AgentReply:
        try:
            return self._runtime.reply(
                AgentInvocation(
                    persona=prepared.thread.persona,
                    message=prepared.text,
                    h5_user=prepared.h5_user,
                    admin_user=prepared.admin_user,
                    history=prepared.history,
                )
            )
        except AppError:
            raise
        except Exception as error:
            raise AppError(
                "AGENT_RUNTIME_FAILED",
                "智能体服务暂不可用，请稍后重试",
                502,
                retryable=True,
            ) from error

    def _complete(
        self,
        prepared: PreparedMessage,
        reply: AgentReply,
    ) -> dict[str, Any]:
        assistant_message = {
            "type": "text",
            "text": reply.text,
            "cards": reply.cards,
            "actions": reply.actions,
            "service_no": prepared.service_no,
        }
        completed_at = self._clock()
        self._repository.save_message(
            AgentMessage(
                id=new_ulid(),
                thread_id=prepared.thread.id,
                role="ASSISTANT",
                content=reply.text,
                metadata={
                    "run_id": prepared.run_id,
                    "service_no": prepared.service_no,
                    "cards": reply.cards,
                    "actions": reply.actions,
                    "invoked_tools": list(reply.invoked_tools),
                    "trace_url": reply.trace_url,
                },
                created_at=completed_at,
            )
        )
        self._repository.save_run(
            AgentRun(
                id=prepared.run_id,
                request_id=f"agent-{new_ulid()}",
                thread_id=prepared.thread.id,
                trace_id=reply.trace_id,
                persona=prepared.thread.persona,
                status="COMPLETED",
                input={"type": "text", "text": prepared.text},
                output={**assistant_message, "trace_url": reply.trace_url},
                error_code=None,
                created_at=prepared.started_at,
                completed_at=completed_at,
            )
        )
        return {
            "run_id": prepared.run_id,
            "thread_id": prepared.thread.id,
            "status": "COMPLETED",
            "assistant_message": assistant_message,
            "service_no": prepared.service_no,
            "trace_id": reply.trace_id,
            "trace_url": reply.trace_url,
            "invoked_tools": list(reply.invoked_tools),
        }

    def _save_failed_run(
        self,
        prepared: PreparedMessage,
        error: AppError,
    ) -> None:
        self._repository.save_run(
            AgentRun(
                id=prepared.run_id,
                request_id=f"agent-{new_ulid()}",
                thread_id=prepared.thread.id,
                trace_id=None,
                persona=prepared.thread.persona,
                status="FAILED",
                input={"type": "text", "text": prepared.text},
                output=None,
                error_code=error.code,
                created_at=prepared.started_at,
                completed_at=self._clock(),
            )
        )

    def _emit(
        self,
        prepared: PreparedMessage,
        event: str,
        sequence: int,
        data: dict[str, Any],
    ) -> RuntimeEvent:
        runtime_event = RuntimeEvent.model_validate(
            {
                "event": event,
                "run_id": prepared.run_id,
                "sequence": sequence,
                "occurred_at": self._clock(),
                "data": data,
            }
        )
        self._state_store.append_event(runtime_event)
        return runtime_event

    def _message_data(self, thread_id: str) -> dict[str, Any]:
        messages = self._repository.list_messages(thread_id)
        return {
            "items": [
                {
                    "id": message.id,
                    "role": message.role,
                    "content": message.content,
                    "metadata": message.metadata,
                    "created_at": message.created_at,
                }
                for message in messages
            ]
        }

    def _get_run(
        self,
        owner_type: str,
        owner_id: str,
        run_id: str,
    ) -> dict[str, Any]:
        run = self._repository.get_run(run_id)
        if run is None:
            raise AppError("AGENT_RUN_NOT_FOUND", "未找到智能体运行", 404)
        self._owned_thread(owner_type, owner_id, run.thread_id)
        trace_url = self._resolve_trace_url(run)
        return {
            "id": run.id,
            "thread_id": run.thread_id,
            "persona": run.persona,
            "status": run.status,
            "trace_id": run.trace_id,
            "trace_url": trace_url,
            "output": run.output,
            "error_code": run.error_code,
            "completed_at": run.completed_at,
        }

    def _resolve_trace_url(self, run: AgentRun) -> str | None:
        trace_url = run.output.get("trace_url") if run.output is not None else None
        if not trace_url and run.trace_id:
            resolver = getattr(self._runtime, "trace_url", None)
            if callable(resolver):
                resolved = resolver(run.trace_id)
                trace_url = str(resolved) if resolved else None
        return str(trace_url) if trace_url else None

    def _owned_thread(
        self,
        owner_type: str,
        owner_id: str,
        thread_id: str,
    ) -> AgentThread:
        thread = self._repository.get_thread(thread_id)
        if thread is None or thread.owner_type != owner_type or thread.owner_id != owner_id:
            raise AppError("AGENT_THREAD_NOT_FOUND", "未找到会话", 404)
        return thread