from typing import Any, Dict, Optional

from pydantic import BaseModel, ConfigDict, Field


class SuperAgentRequest(BaseModel):
    """Super Agent 请求模型"""

    model_config = ConfigDict(extra="forbid")

    query: str = Field(..., desc ription="用户查询")
    conversation_id: Optional[str] = Field(None, description="会话ID")
    user_id: Optional[str] = Field(None, description="用户ID")
    workflow_type: str = Field(default="conversation", description="工作流类型")
    context: Dict[str, str] = Field(default_factory=dict, description="上下文信息")
    timeout_seconds: int = Field(default=30, description="超时时间(秒)")


class SuperAgentResponse(BaseModel):
    """Super Agent 响应模型"""

    conversation_id: str = Field(..., description="会话ID")
    workflow_type: str = Field(..., description="工作流类型")
    status: str = Field(default="success", description="状态: success/error")
    sql: Optional[str] = Field(None, description="生成的SQL")
    result: Optional[str] = Field(None, description="查询结果")
    error: Optional[str] = Field(None, description="错误信息")
    metadata: Dict[str, str] = Field(default_factory=dict, description="元数据")


class SuperAgentStreamEvent(BaseModel):
    """Super Agent 流式响应事件"""

    conversation_id: str = Field(..., description="会话ID")
    event: str = Field(..., description="事件类型: sql_generated/sql_executing/result/error/done")
    data: str = Field(..., description="事件数据")
    timestamp: int = Field(..., description="时间戳(毫秒)")