349 lines
11 KiB
Python
349 lines
11 KiB
Python
#!/usr/bin/env python3
|
||
"""
|
||
API data models and response schemas.
|
||
"""
|
||
|
||
from typing import Dict, List, Optional, Any, AsyncGenerator, Union
|
||
from pydantic import BaseModel, Field, field_validator, ConfigDict
|
||
|
||
class Message(BaseModel):
|
||
role: str
|
||
# content can be a plain string, or a list of content blocks for multimodal
|
||
# input (e.g. text + image). Both OpenAI-style ({"type": "image_url", ...})
|
||
# and LangChain standard blocks ({"type": "image", ...}) are accepted; they
|
||
# are normalized later in process_messages.
|
||
content: Union[str, List[Dict[str, Any]]]
|
||
|
||
|
||
class DatasetRequest(BaseModel):
|
||
system_prompt: Optional[str] = None
|
||
mcp_settings: Optional[List[Dict]] = None
|
||
files: Optional[Dict[str, List[str]]] = Field(default=None, description="Files organized by key groups. Each key maps to a list of file paths (supports zip files)")
|
||
unique_id: Optional[str] = None
|
||
|
||
@field_validator('files', mode='before')
|
||
@classmethod
|
||
def validate_files(cls, v):
|
||
"""Validate dict format with key-grouped files"""
|
||
if v is None:
|
||
return None
|
||
if isinstance(v, dict):
|
||
# Validate dict format
|
||
for key, value in v.items():
|
||
if not isinstance(key, str):
|
||
raise ValueError(f"Key in files dict must be string, got {type(key)}")
|
||
if not isinstance(value, list):
|
||
raise ValueError(f"Value in files dict must be list, got {type(value)} for key '{key}'")
|
||
for item in value:
|
||
if not isinstance(item, str):
|
||
raise ValueError(f"File paths must be strings, got {type(item)} in key '{key}'")
|
||
return v
|
||
else:
|
||
raise ValueError(f"Files must be a dict with key groups, got {type(v)}")
|
||
|
||
|
||
class ChatRequest(BaseModel):
|
||
messages: List[Message]
|
||
model: str = "qwen3-next"
|
||
model_server: str = ""
|
||
dataset_ids: Optional[List[str]] = None
|
||
bot_id: str
|
||
stream: Optional[bool] = False
|
||
language: Optional[str] = "zh"
|
||
tool_response: Optional[bool] = True
|
||
system_prompt: Optional[str] = ""
|
||
mcp_settings: Optional[List[Dict]] = None
|
||
user_identifier: Optional[str] = ""
|
||
session_id: Optional[str] = None
|
||
enable_thinking: Optional[bool] = False
|
||
skills: Optional[List[str]] = None
|
||
enable_memory: Optional[bool] = False
|
||
enable_self_knowledge: Optional[bool] = False
|
||
shell_env: Optional[Dict[str, str]] = None
|
||
|
||
model_config = ConfigDict(extra='allow')
|
||
|
||
|
||
class ChatRequestV2(BaseModel):
|
||
messages: List[Message]
|
||
stream: Optional[bool] = False
|
||
tool_response: Optional[bool] = True
|
||
bot_id: str
|
||
language: Optional[str] = "zh"
|
||
user_identifier: Optional[str] = ""
|
||
session_id: Optional[str] = None
|
||
|
||
model_config = ConfigDict(extra='allow')
|
||
|
||
|
||
class ChatRequestV3(BaseModel):
|
||
"""
|
||
v3 API 请求模型 - 从数据库读取配置
|
||
|
||
所有配置参数从数据库读取，前端只需传递：
|
||
- bot_id: Bot 的用户ID（用于从数据库查找配置）
|
||
- messages: 对话消息列表
|
||
- session_id: 可选的会话ID
|
||
- user_identifier: 当前登录用户的用户名，用于标识用户身份
|
||
"""
|
||
messages: List[Message]
|
||
bot_id: str
|
||
stream: Optional[bool] = False
|
||
session_id: Optional[str] = None
|
||
user_identifier: Optional[str] = None
|
||
|
||
|
||
class VisionMessage(BaseModel):
|
||
"""Message supporting multimodal content (text + images).
|
||
|
||
content accepts either:
|
||
- a plain string (text only), or
|
||
- a list of OpenAI-style content parts, e.g.:
|
||
[
|
||
{"type": "text", "text": "what is in this image?"},
|
||
{"type": "image_url", "image_url": {"url": "https://... or data:image/png;base64,..."}}
|
||
]
|
||
"""
|
||
role: str
|
||
content: Union[str, List[Dict[str, Any]]]
|
||
|
||
|
||
class LLMPassthroughRequest(BaseModel):
|
||
"""LLM passthrough request model - supports vision/multimodal content.
|
||
|
||
Only bot_id and messages are required. Config (model/api_key/model_server)
|
||
is resolved from the database via bot_id.
|
||
"""
|
||
messages: List[VisionMessage]
|
||
bot_id: str
|
||
stream: Optional[bool] = False
|
||
user_identifier: Optional[str] = None
|
||
|
||
model_config = ConfigDict(extra='allow')
|
||
|
||
|
||
class FileProcessRequest(BaseModel):
|
||
unique_id: str
|
||
files: Optional[Dict[str, List[str]]] = Field(default=None, description="Files organized by key groups. Each key maps to a list of file paths (supports zip files)")
|
||
system_prompt: Optional[str] = None
|
||
mcp_settings: Optional[List[Dict]] = None
|
||
|
||
model_config = ConfigDict(extra='allow')
|
||
|
||
@field_validator('files', mode='before')
|
||
@classmethod
|
||
def validate_files(cls, v):
|
||
"""Validate dict format with key-grouped files"""
|
||
if v is None:
|
||
return None
|
||
if isinstance(v, dict):
|
||
# Validate dict format
|
||
for key, value in v.items():
|
||
if not isinstance(key, str):
|
||
raise ValueError(f"Key in files dict must be string, got {type(key)}")
|
||
if not isinstance(value, list):
|
||
raise ValueError(f"Value in files dict must be list, got {type(value)} for key '{key}'")
|
||
for item in value:
|
||
if not isinstance(item, str):
|
||
raise ValueError(f"File paths must be strings, got {type(item)} in key '{key}'")
|
||
return v
|
||
else:
|
||
raise ValueError(f"Files must be a dict with key groups, got {type(v)}")
|
||
|
||
|
||
class DatasetResponse(BaseModel):
|
||
success: bool
|
||
message: str
|
||
unique_id: Optional[str] = None
|
||
dataset_structure: Optional[str] = None
|
||
|
||
|
||
class ChatCompletionResponse(BaseModel):
|
||
id: str
|
||
object: str = "chat.completion"
|
||
created: int
|
||
model: str
|
||
choices: List[Dict[str, Any]]
|
||
usage: Optional[Dict[str, int]] = None
|
||
|
||
|
||
class ChatResponse(BaseModel):
|
||
choices: List[Dict]
|
||
usage: Optional[Dict] = None
|
||
|
||
|
||
class FileProcessResponse(BaseModel):
|
||
success: bool
|
||
message: str
|
||
unique_id: str
|
||
processed_files: List[str]
|
||
|
||
|
||
class ErrorResponse(BaseModel):
|
||
error: Dict[str, Any]
|
||
|
||
@classmethod
|
||
def create(cls, message: str, error_type: str = "invalid_request_error", code: Optional[str] = None):
|
||
error_data = {
|
||
"message": message,
|
||
"type": error_type
|
||
}
|
||
if code:
|
||
error_data["code"] = code
|
||
return cls(error=error_data)
|
||
|
||
|
||
class HealthCheckResponse(BaseModel):
|
||
status: str = "healthy"
|
||
timestamp: str
|
||
version: str = "1.0.0"
|
||
|
||
|
||
class SystemStatusResponse(BaseModel):
|
||
status: str
|
||
projects_count: int
|
||
total_projects: List[str]
|
||
active_projects: List[str]
|
||
system_info: Dict[str, Any]
|
||
|
||
|
||
class CacheStatusResponse(BaseModel):
|
||
cached_projects: List[str]
|
||
cache_info: Dict[str, Any]
|
||
|
||
|
||
class ProjectStatusResponse(BaseModel):
|
||
unique_id: str
|
||
project_exists: bool
|
||
project_path: Optional[str] = None
|
||
processed_files_count: int
|
||
processed_files: Dict[str, Dict]
|
||
document_files_count: int
|
||
document_files: List[str]
|
||
has_system_prompt: bool
|
||
has_mcp_settings: bool
|
||
readme_exists: bool
|
||
log_file_exists: bool
|
||
dataset_structure: Optional[str] = None
|
||
error: Optional[str] = None
|
||
|
||
|
||
class ProjectListResponse(BaseModel):
|
||
projects: List[str]
|
||
count: int
|
||
|
||
|
||
class ProjectStatsResponse(BaseModel):
|
||
unique_id: str
|
||
total_processed_files: int
|
||
total_document_files: int
|
||
total_document_size: int
|
||
total_document_size_mb: float
|
||
has_system_prompt: bool
|
||
has_mcp_settings: bool
|
||
has_readme: bool
|
||
document_files_detail: List[Dict[str, Any]]
|
||
embedding_files_count: int
|
||
embedding_files_detail: List[Dict[str, Any]]
|
||
|
||
|
||
class ProjectActionResponse(BaseModel):
|
||
success: bool
|
||
message: str
|
||
unique_id: str
|
||
action: str
|
||
|
||
|
||
# Utility functions for creating responses
|
||
def create_success_response(message: str, **kwargs) -> Dict[str, Any]:
|
||
"""Create a standardized success response"""
|
||
return {
|
||
"success": True,
|
||
"message": message,
|
||
**kwargs
|
||
}
|
||
|
||
|
||
def create_error_response(message: str, error_type: str = "error", **kwargs) -> Dict[str, Any]:
|
||
"""Create a standardized error response"""
|
||
return {
|
||
"success": False,
|
||
"error": error_type,
|
||
"message": message,
|
||
**kwargs
|
||
}
|
||
|
||
|
||
def create_chat_response(
|
||
messages: List[Message],
|
||
model: str,
|
||
content: str,
|
||
usage: Optional[Dict[str, int]] = None
|
||
) -> Dict[str, Any]:
|
||
"""Create a chat completion response"""
|
||
import time
|
||
import uuid
|
||
|
||
return {
|
||
"id": f"chatcmpl-{uuid.uuid4().hex[:8]}",
|
||
"object": "chat.completion",
|
||
"created": int(time.time()),
|
||
"model": model,
|
||
"choices": [
|
||
{
|
||
"index": 0,
|
||
"message": {
|
||
"role": "assistant",
|
||
"content": content
|
||
},
|
||
"finish_reason": "stop"
|
||
}
|
||
],
|
||
"usage": usage or {
|
||
"prompt_tokens": 0,
|
||
"completion_tokens": 0,
|
||
"total_tokens": 0
|
||
}
|
||
}
|
||
|
||
|
||
# ============================================================================
|
||
# Models related to chat history queries
|
||
# ============================================================================
|
||
|
||
class ChatHistoryRequest(BaseModel):
|
||
"""Chat history query request"""
|
||
session_id: str = Field(..., description="Session ID (thread_id)")
|
||
last_message_id: Optional[str] = Field(None, description="ID of the previous message, used to paginate and fetch earlier messages")
|
||
limit: int = Field(20, ge=1, le=100, description="Maximum number of messages to return per query")
|
||
|
||
|
||
class ChatHistoryMessage(BaseModel):
|
||
"""Chat history message"""
|
||
id: str = Field(..., description="Unique message ID")
|
||
role: str = Field(..., description="Message role: user or assistant")
|
||
content: str = Field(..., description="Message content")
|
||
timestamp: Optional[str] = Field(None, description="Message timestamp (ISO 8601)")
|
||
|
||
|
||
class ChatHistoryResponse(BaseModel):
|
||
"""Chat history query response"""
|
||
messages: List[ChatHistoryMessage] = Field(..., description="List of messages in reverse chronological order")
|
||
has_more: bool = Field(..., description="Whether more history messages are available")
|
||
|
||
|
||
class BatchSaveChatRequest(BaseModel):
|
||
"""Batch save chat records request"""
|
||
session_id: str = Field(..., description="Session ID (thread_id)")
|
||
messages: List[Message] = Field(..., description="List of messages to save, supporting user and assistant roles")
|
||
bot_id: Optional[str] = Field(None, description="Bot ID")
|
||
|
||
|
||
class BatchSaveChatResponse(BaseModel):
|
||
"""Batch save chat records response"""
|
||
success: bool = Field(..., description="Whether the operation succeeded")
|
||
message: str = Field(..., description="Response message")
|
||
session_id: str = Field(..., description="Session ID")
|
||
saved_count: int = Field(..., description="Number of messages saved successfully")
|
||
message_ids: List[str] = Field(..., description="List of saved message IDs")
|