ForcePilot/backend/package/yuxi/channel/message/models.py

import hashlib
import json
from dataclasses import dataclass, field
from datetime import datetime
from enum import StrEnum
from typing import Any

from yuxi.channel.routing.models import PeerKind


class MessageType(StrEnum):
    TEXT = "text"
    IMAGE = "image"
    VOICE = "voice"
    FILE = "file"
    EVENT = "event"


class ReplyStage(StrEnum):
    TOOL = "tool"
    BLOCK = "block"
    FINAL = "final"


class ChunkType(StrEnum):
    TEXT_DELTA = "text-delta"
    REASONING = "reasoning"
    TOOL_CALL = "tool-call"
    TOOL_RESULT = "tool-result"
    HEARTBEAT = "heartbeat"
    CITATION = "citation"
    STATUS = "status"
    ERROR = "error"


class MentionSource(StrEnum):
    EXPLICIT_BOT = "explicit_bot"
    SUBTEAM = "subteam"
    MENTION_PATTERN = "mention_pattern"
    IMPLICIT_THREAD = "implicit_thread"
    COMMAND_BYPASS = "command_bypass"
    NONE = "none"


@dataclass
class PeerInfo:
    kind: PeerKind
    id: str
    display_name: str | None = None
    username: str | None = None
    tag: str | None = None
    is_bot: bool = False
    is_self: bool = False
    display_label: str | None = None
    roles: list[str] = field(default_factory=list)


@dataclass
class GroupContext:
    id: str | None = None
    name: str | None = None
    guild_id: str | None = None
    team_id: str | None = None
    thread_id: str | None = None
    roles: list[str] = field(default_factory=list)
    kind: str | None = None
    space_id: str | None = None
    parent_id: str | None = None
    native_channel_id: str | None = None
    route_peer_kind: str | None = None
    route_peer_id: str | None = None


@dataclass
class UnifiedMessage:
    msg_id: str
    channel_type: str
    account_id: str
    content: str
    sender: PeerInfo
    channel_config_id: str | None = None
    message_type: MessageType = MessageType.TEXT
    media_urls: list[str] = field(default_factory=list)
    image_base64: str | None = None
    group: GroupContext | None = None
    timestamp: datetime | None = None
    raw_payload: dict[str, Any] = field(default_factory=dict)
    metadata: dict[str, Any] = field(default_factory=dict)

    body_for_agent: str | None = None
    command_body: str | None = None
    body_for_commands: str | None = None
    inbound_history: list[dict[str, Any]] = field(default_factory=list)

    reply_chain: list[dict[str, Any]] = field(default_factory=list)
    reply_to_id: str | None = None
    reply_to_id_full: str | None = None
    root_message_id: str | None = None
    message_thread_id: str | None = None
    reply_to_tag: bool = False
    reply_to_current: bool = False

    forwarded_from: str | None = None
    forwarded_from_type: str | None = None
    forwarded_from_id: str | None = None
    forwarded_from_username: str | None = None
    forwarded_date: float | None = None

    media_path: str | None = None
    media_paths: list[str] = field(default_factory=list)
    media_types: list[str] = field(default_factory=list)

    conversation_label: str | None = None
    group_channel: str | None = None
    group_space: str | None = None
    group_members: str | None = None
    group_system_prompt: str | None = None
    member_role_ids: list[str] = field(default_factory=list)

    was_mentioned: bool = False
    explicitly_mentioned_bot: bool = False
    mentioned_user_ids: list[str] = field(default_factory=list)
    mention_source: MentionSource | None = None

    surface: str | None = None
    originating_channel: str | None = None
    originating_to: str | None = None
    native_channel_id: str | None = None
    native_direct_user_id: str | None = None
    thread_parent_id: str | None = None

    location_lat: float | None = None
    location_lon: float | None = None


@dataclass
class DispatchResult:
    success: bool
    thread_id: str | None = None
    agent_config_id: int | None = None
    error: str | None = None
    session_key: str | None = None
    matched_by: str | None = None
    handled_by: str | None = None
    command_response: str | None = None
    internal_user_id: str | None = None
    reply_sent: bool = False
    explicit_target: str | None = None
    queued_final: bool = False
    counts: dict[str, int] = field(default_factory=dict)
    failed_counts: dict[str, int] = field(default_factory=dict)
    source_reply_delivery_mode: str | None = None
    before_agent_run_blocked: bool = False


@dataclass
class StreamingChunk:
    """Agent 推理流式块"""

    stage: ReplyStage = ReplyStage.BLOCK
    chunk_type: ChunkType = ChunkType.TEXT_DELTA
    content: str | None = None
    tool_name: str | None = None
    tool_input: dict[str, Any] | None = None
    tool_output: str | None = None
    metadata: dict[str, Any] = field(default_factory=dict)
    agent_run_id: str | None = None
    is_final: bool = False
    item_id: str | None = None
    title: str | None = None
    phase: str | None = None
    steps: list[str] | None = None


@dataclass
class ReplyPayload:
    """出站回复载荷（经过去重和合并后）"""

    target_id: str
    content: str
    reply_to_id: str | None = None
    thread_id: str | None = None
    media_urls: list[str] = field(default_factory=list)
    presentation: dict[str, Any] | None = None
    metadata: dict[str, Any] = field(default_factory=dict)
    btw: dict[str, str] | None = None
    audio_as_voice: bool = False
    spoken_text: str | None = None
    is_error: bool = False
    is_reasoning: bool = False
    is_compaction_notice: bool = False
    channel_data: dict[str, Any] | None = None
    trusted_local_media: bool = False
    sensitive_media: bool = False
    delivery: dict[str, Any] | None = None
    reply_to_tag: bool = False
    reply_to_current: bool = False

    @property
    def payload_key(self) -> str:
        key = f"{self.content}|{','.join(sorted(self.media_urls))}|{self.reply_to_id or ''}"
        if self.presentation:
            stable_hash = hashlib.md5(
                json.dumps(self.presentation, sort_keys=True, ensure_ascii=False).encode()
            ).hexdigest()
            key += f"|{stable_hash}"
        return key

    @property
    def content_key(self) -> str:
        return f"{self.content}|{','.join(sorted(self.media_urls))}"


class StreamStatus(StrEnum):
    """流执行终止状态"""

    COMPLETED = "completed"
    CANCELLED = "cancelled"
    ERROR = "error"
    TIMEOUT = "timeout"


@dataclass
class StreamResult:
    """流执行完成后的结构化结果，对应 openclaw 的 MessageReceipt 概念"""

    status: StreamStatus = StreamStatus.COMPLETED
    accumulated_text: str = ""
    chunk_count: int = 0
    text_chunk_count: int = 0
    tool_chunk_count: int = 0
    reasoning_chunk_count: int = 0
    heartbeat_count: int = 0
    citation_count: int = 0
    error_count: int = 0
    start_time: float = 0.0
    end_time: float = 0.0
    token_usage: dict[str, int] | None = None
    error_message: str | None = None

    @property
    def elapsed_ms(self) -> float:
        return (self.end_time - self.start_time) * 1000 if self.start_time > 0 else 0.0

    @property
    def is_success(self) -> bool:
        return self.status == StreamStatus.COMPLETED


class StreamCancellationToken:
    """流取消令牌，对应 openclaw 的 AbortSignal"""

    def __init__(self) -> None:
        self._cancelled = False

    def cancel(self) -> None:
        self._cancelled = True

    @property
    def is_cancelled(self) -> bool:
        return self._cancelled

    def throw_if_cancelled(self) -> None:
        if self._cancelled:
            raise StreamCancelledError()


class StreamCancelledError(Exception):
    pass


class StreamTimeoutError(Exception):
    pass


@dataclass
class MessageReceipt:
    platform_message_id: str
    thread_id: str | None = None
    reply_to_id: str | None = None
    sent_at: float | None = None
    parts: list[dict[str, Any]] = field(default_factory=list)
    raw: list[dict[str, Any]] = field(default_factory=list)