more_dots/core/response.py

"""
统一响应格式模块

提供标准化的 API 响应和流式事件格式
"""

from typing import Any, Dict, Generic, List, Optional, TypeVar, Literal
from pydantic import BaseModel, Field
import time
import uuid


T = TypeVar("T")


class ApiResponse(BaseModel, Generic[T]):
    """
    统一 API 响应格式
    
    所有 API 响应都使用这个格式，提供一致的响应结构
    
    Usage:
        @router.get("/users/{user_id}")
        async def get_user(user_id: str) -> ApiResponse[User]:
            user = await user_service.get(user_id)
            return ApiResponse.success(data=user)
    """
    
    code: str = Field(default="success", description="响应代码")
    message: str = Field(default="", description="响应消息")
    data: Optional[T] = Field(default=None, description="响应数据")
    trace_id: Optional[str] = Field(default=None, description="追踪ID")
    timestamp: int = Field(
        default_factory=lambda: int(time.time() * 1000),
        description="时间戳(毫秒)"
    )
    
    @classmethod
    def success(cls, data: T = None, message: str = "", trace_id: Optional[str] = None) -> "ApiResponse[T]":
        """创建成功响应"""
        return cls(
            code="success",
            message=message,
            data=data,
            trace_id=trace_id or uuid.uuid4().hex,
        )
    
    @classmethod
    def error(
        cls,
        code: str = "error",
        message: str = "",
        data: T = None,
        trace_id: Optional[str] = None,
    ) -> "ApiResponse[T]":
        """创建错误响应"""
        return cls(
            code=code,
            message=message,
            data=data,
            trace_id=trace_id or uuid.uuid4().hex,
        )
    
    @classmethod
    def from_exception(cls, exc: Exception, trace_id: Optional[str] = None) -> "ApiResponse[None]":
        """从异常创建错误响应"""
        return cls.error(
            code="internal_error",
            message=str(exc),
            trace_id=trace_id,
        )
    
    def is_success(self) -> bool:
        """判断是否成功"""
        return self.code == "success"


class PagedResponse(BaseModel, Generic[T]):
    """
    分页响应格式
    
    用于返回分页数据
    """
    
    items: List[T] = Field(default_factory=list, description="数据列表")
    total: int = Field(default=0, description="总数")
    page: int = Field(default=1, description="当前页")
    page_size: int = Field(default=20, description="每页大小")
    total_pages: int = Field(default=0, description="总页数")
    
    @classmethod
    def create(
        cls,
        items: List[T],
        total: int,
        page: int = 1,
        page_size: int = 20,
    ) -> "PagedResponse[T]":
        """创建分页响应"""
        total_pages = (total + page_size - 1) // page_size if page_size > 0 else 0
        return cls(
            items=items,
            total=total,
            page=page,
            page_size=page_size,
            total_pages=total_pages,
        )


class StreamEvent(BaseModel):
    """
    流式响应事件
    
    用于 SSE (Server-Sent Events) 流式响应
    
    Usage:
        async def event_stream():
            yield StreamEvent(event="start", data="Processing started")
            # ... 处理逻辑
            yield StreamEvent(event="result", data=json.dumps(result))
            yield StreamEvent(event="done", data="")
    """
    
    event: str = Field(..., description="事件类型")
    data: str = Field(default="", description="事件数据")
    event_id: Optional[str] = Field(default=None, description="事件ID")
    retry: Optional[int] = Field(default=None, description="重试间隔(毫秒)")
    
    def to_sse(self) -> str:
        """转换为 SSE 格式字符串"""
        lines = [f"event: {self.event}"]
        if self.event_id:
            lines.append(f"id: {self.event_id}")
        if self.retry:
            lines.append(f"retry: {self.retry}")
        lines.append(f"data: {self.data}")
        lines.append("")
        lines.append("")
        return "\n".join(lines)
    
    @classmethod
    def message(cls, data: str, event_id: Optional[str] = None) -> "StreamEvent":
        """创建消息事件"""
        return cls(event="message", data=data, event_id=event_id)
    
    @classmethod
    def done(cls) -> "StreamEvent":
        """创建完成事件"""
        return cls(event="done", data="[DONE]")
    
    @classmethod
    def error(cls, message: str) -> "StreamEvent":
        """创建错误事件"""
        return cls(event="error", data=message)


class WorkflowEvent(BaseModel):
    """
    工作流事件
    
    用于工作流执行过程中的状态通知
    """
    
    workflow_id: str = Field(..., description="工作流ID")
    event_type: Literal[
        "started",
        "node_started",
        "node_completed",
        "node_failed",
        "completed",
        "failed",
    ] = Field(..., description="事件类型")
    node_name: Optional[str] = Field(None, description="节点名称")
    data: Optional[Dict[str, Any]] = Field(None, description="事件数据")
    error: Optional[str] = Field(None, description="错误信息")
    timestamp: int = Field(
        default_factory=lambda: int(time.time() * 1000),
        description="时间戳"
    )
    
    @classmethod
    def started(cls, workflow_id: str) -> "WorkflowEvent":
        """创建开始事件"""
        return cls(workflow_id=workflow_id, event_type="started")
    
    @classmethod
    def node_started(cls, workflow_id: str, node_name: str) -> "WorkflowEvent":
        """创建节点开始事件"""
        return cls(
            workflow_id=workflow_id,
            event_type="node_started",
            node_name=node_name,
        )
    
    @classmethod
    def node_completed(
        cls,
        workflow_id: str,
        node_name: str,
        data: Optional[Dict[str, Any]] = None,
    ) -> "WorkflowEvent":
        """创建节点完成事件"""
        return cls(
            workflow_id=workflow_id,
            event_type="node_completed",
            node_name=node_name,
            data=data,
        )
    
    @classmethod
    def completed(
        cls,
        workflow_id: str,
        data: Optional[Dict[str, Any]] = None,
    ) -> "WorkflowEvent":
        """创建完成事件"""
        return cls(workflow_id=workflow_id, event_type="completed", data=data)
    
    @classmethod
    def failed(
        cls,
        workflow_id: str,
        error: str,
        node_name: Optional[str] = None,
    ) -> "WorkflowEvent":
        """创建失败事件"""
        return cls(
            workflow_id=workflow_id,
            event_type="failed",
            node_name=node_name,
            error=error,
        )


class ErrorCode:
    """错误代码常量"""
    
    SUCCESS = "success"
    UNKNOWN_ERROR = "unknown_error"
    INVALID_REQUEST = "invalid_request"
    INVALID_WORKFLOW_TYPE = "invalid_workflow_type"
    SQL_GENERATION_FAILED = "sql_generation_failed"
    TOOL_NOT_FOUND = "tool_not_found"
    TOOL_EXECUTION_FAILED = "tool_execution_failed"
    INTERNAL_ERROR = "internal_error"
    TIMEOUT = "timeout"
    RATE_LIMITED = "rate_limited"
    UNAUTHORIZED = "unauthorized"