Maas-Model-platform/backend/app/schemas/llm_schema.py

"""
LLM对话相关的数据传输对象定义

定义AI对话API请求和响应的数据结构，包含联网搜索功能扩展
需求: 4.1, 4.2, 5.1, 5.2, 5.3, 6.1, 6.2, 6.4
搜索需求: 2.4, 5.3, 6.2, 10.3
"""
from typing import List, Literal, Optional, Union, Dict
from pydantic import BaseModel, Field, field_validator


class ImageContent(BaseModel):
    """图像内容"""
    type: Literal["image_url"] = "image_url"
    image_url: dict = Field(..., description="图像URL信息，包含url字段（支持base64格式）")


class TextContent(BaseModel):
    """文本内容"""
    type: Literal["text"] = "text"
    text: str = Field(..., description="文本内容")


class ChatMessage(BaseModel):
    """对话消息"""
    role: Literal["system", "user", "assistant"]
    content: Union[str, List[Union[TextContent, ImageContent]]] = Field(..., description="消息内容，可以是字符串或多模态内容列表")


class ChatRequest(BaseModel):
    """对话请求"""
    model: str = Field(..., description="模型名称")
    messages: List[ChatMessage] = Field(..., min_length=1, description="对话消息列表")
    stream: bool = Field(default=True, description="是否流式输出，默认启用")
    temperature: Optional[float] = Field(default=None, ge=0, lt=2, description="采样温度 [0, 2)")
    top_p: Optional[float] = Field(default=None, gt=0, le=1.0, description="核采样概率 (0, 1.0]")
    max_tokens: Optional[int] = Field(default=None, gt=0, description="最大输出token数")
    conversation_id: Optional[int] = Field(default=None, description="会话ID，用于记录对话历史")
    enable_thinking: bool = Field(default=False, description="是否启用思考模式（仅Qwen3系列支持）")
    thinking_budget: Optional[int] = Field(default=None, gt=0, description="思考过程的最大Token数")

    @field_validator('messages')
    @classmethod
    def validate_messages(cls, v):
        if not v:
            raise ValueError('消息列表不能为空')
        return v


class UsageInfo(BaseModel):
    """Token使用统计"""
    input_tokens: int
    output_tokens: int
    total_tokens: int


class ChatResponse(BaseModel):
    """对话响应"""
    content: str
    finish_reason: str
    usage: UsageInfo
    reasoning_content: Optional[str] = Field(default=None, description="思考过程内容")


class StreamChunk(BaseModel):
    """流式响应块"""
    content: str
    finish_reason: Optional[str] = None
    usage: Optional[UsageInfo] = None
    reasoning_content: Optional[str] = Field(default=None, description="思考过程内容（增量）")


# ==================== 联网搜索功能扩展 ====================

class SearchOptions(BaseModel):
    """搜索配置选项"""
    enable_search: bool = Field(default=False, description="是否启用联网搜索")
    search_strategy: Literal["turbo", "max", "agent"] = Field(default="turbo", description="搜索策略")
    forced_search: bool = Field(default=False, description="是否强制搜索")
    enable_search_extension: bool = Field(default=False, description="是否启用垂域搜索")
    freshness: Optional[Literal[7, 30, 180, 365]] = Field(default=None, description="搜索时效性（天）")
    enable_source: bool = Field(default=False, description="是否返回搜索来源")
    enable_citation: bool = Field(default=False, description="是否启用角标引用")
    citation_format: Literal["[<number>]", "[ref_<number>]"] = Field(default="[<number>]", description="角标格式")
    prepend_search_result: bool = Field(default=False, description="是否提前返回搜索结果")
    intention_options: Optional[Dict[str, str]] = Field(default=None, description="自然语言搜索控制")

    @field_validator('freshness')
    @classmethod
    def validate_freshness_with_strategy(cls, v, info):
        """验证时效性参数仅对turbo策略生效"""
        if v is not None:
            # 注意：这里无法直接访问search_strategy，因为Pydantic的验证顺序
            # 实际的策略兼容性验证将在SearchOptionsValidator中进行
            pass
        return v


class SearchResult(BaseModel):
    """搜索结果"""
    index: int = Field(..., description="搜索结果索引")
    title: str = Field(..., description="页面标题")
    url: str = Field(..., description="页面URL")
    snippet: Optional[str] = Field(default=None, description="内容摘要")


class EnhancedChatRequest(ChatRequest):
    """增强的对话请求，支持搜索选项"""
    search_options: Optional[SearchOptions] = Field(default=None, description="搜索配置选项")


class EnhancedChatResponse(ChatResponse):
    """增强的对话响应，包含搜索信息"""
    search_info: Optional[Dict] = Field(default=None, description="搜索信息")
    search_results: Optional[List[SearchResult]] = Field(default=None, description="搜索结果列表")


class EnhancedStreamChunk(StreamChunk):
    """增强的流式响应块，包含搜索信息"""
    search_info: Optional[Dict] = Field(default=None, description="搜索信息")
    search_results: Optional[List[SearchResult]] = Field(default=None, description="搜索结果列表")