x

2026-03-11 23:40:39 +08:00
parent db25d61026
commit e062368ef2
15 changed files with 1592 additions and 229 deletions
@@ -0,0 +1,248 @@
+"""
+统一响应格式模块
+
+提供标准化的 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"