LQAgentPlatform/foundation/schemas/test_schemas.py

#!/usr/bin/env python
# -*- coding: utf-8 -*-

"""
测试模式定义

提供测试相关的数据模型和配置结构
"""

from typing import Optional, Dict, Any, List
from pydantic import BaseModel, Field


class TestConfig(BaseModel):
    """测试配置"""
    session_id: str = Field(description="会话ID")
    model_type: Optional[str] = Field(default="gemini", description="模型类型")
    temperature: Optional[float] = Field(default=0.7, description="温度参数")
    max_tokens: Optional[int] = Field(default=2000, description="最大token数")


class TestForm(BaseModel):
    """测试表单"""
    input: str = Field(description="输入内容")
    context: Optional[Dict[str, Any]] = Field(default=None, description="上下文信息")
    config: TestConfig = Field(description="配置信息")


class TestResponse(BaseModel):
    """测试响应"""
    output: str = Field(description="输出结果")
    trace_id: Optional[str] = Field(default=None, description="追踪ID")
    processing_time: Optional[float] = Field(default=None, description="处理时间(秒)")


class StreamEvent(BaseModel):
    """流式事件"""
    event: str = Field(description="事件类型")
    data: Dict[str, Any] = Field(description="事件数据")


class TestResult(BaseModel):
    """测试结果"""
    success: bool = Field(description="是否成功")
    message: str = Field(description="消息")
    data: Optional[Dict[str, Any]] = Field(default=None, description="数据")
    error: Optional[str] = Field(default=None, description="错误信息")


# 导出的类
__all__ = [
    "TestConfig",
    "TestForm",
    "TestResponse",
    "StreamEvent",
    "TestResult"
]