use serde::{Deserialize, Serialize};
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
#[serde(tag = "kind")]
pub enum ModelKind {
#[serde(rename = "chat")]
Chat {
#[serde(rename = "contextWindow", default)]
context_window: Option<u32>,
#[serde(rename = "maxOutput", default)]
max_output: Option<u32>,
#[serde(rename = "supportsStream", default)]
supports_stream: bool,
#[serde(rename = "capabilities", default)]
capabilities: Vec<ChatCapability>,
},
#[serde(rename = "embedding")]
Embedding {
#[serde(rename = "contextWindow", default)]
context_window: Option<u32>,
#[serde(rename = "dimensions", default)]
dimensions: Option<u32>,
#[serde(rename = "similarity")]
#[serde(default)]
similarity: ModelSimilarity,
#[serde(rename = "capabilities", default)]
capabilities: Vec<EmbeddingCapability>,
},
}
#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
pub enum ChatCapability {
#[serde(rename = "reasoning")]
Reasoning,
}
#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
pub enum EmbeddingCapability {
#[serde(rename = "outputDimensions")]
OutputDimensions,
}
#[derive(Debug, Clone, Copy, PartialEq, Eq, Default, Serialize, Deserialize)]
pub enum ModelSimilarity {
#[default]
#[serde(rename = "cosine")]
Cosine,
#[serde(rename = "dotProduct")]
DotProduct,
#[serde(rename = "euclidean")]
Euclidean,
}
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct ModelInfo {
#[serde(rename = "modelId")]
pub model_id: String,
#[serde(rename = "name")]
pub name: String,
#[serde(flatten)]
pub kind: ModelKind,
#[serde(rename = "provider")]
pub provider: String,
}
#[derive(Debug, Clone, thiserror::Error)]
pub enum ModelError {
#[error("模型提供方不支持该操作")]
NotSupported,
#[error("模型不存在: {0}")]
ModelNotFound(String),
#[error("模型提供方不可用: {0}")]
ProviderUnavailable(String),
#[error("请求参数非法: {0}")]
InvalidRequest(String),
#[error("模型传输失败: {0}")]
Transport(String),
#[error("模型内部错误: {0}")]
Internal(String),
}
#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
pub enum ModelChatRole {
#[serde(rename = "system")]
System,
#[serde(rename = "user")]
User,
#[serde(rename = "assistant")]
Assistant,
}
impl Default for ModelChatRole {
fn default() -> Self {
Self::User
}
}
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct ModelChatMessage {
#[serde(rename = "role", default)]
pub role: ModelChatRole,
#[serde(rename = "content", default)]
pub content: String,
}
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct ModelChatRequest {
#[serde(rename = "modelId", default)]
pub model_id: String,
#[serde(rename = "messages", default)]
pub messages: Vec<ModelChatMessage>,
#[serde(rename = "temperature", default)]
pub temperature: Option<f32>,
#[serde(rename = "maxTokens", default)]
pub max_tokens: Option<u32>,
#[serde(rename = "reasoningEffort", default)]
pub reasoning_effort: Option<String>,
}
#[derive(Debug, Clone, Default, Serialize, Deserialize)]
pub struct ModelTokenUsage {
#[serde(rename = "inputTokens", default)]
pub input_tokens: u32,
#[serde(rename = "outputTokens", default)]
pub output_tokens: u32,
#[serde(rename = "totalTokens", default)]
pub total_tokens: u32,
}
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct ModelChatResponse {
#[serde(rename = "modelId", default)]
pub model_id: String,
#[serde(rename = "content", default)]
pub content: String,
#[serde(rename = "finishReason", default)]
pub finish_reason: Option<String>,
#[serde(rename = "usage", default)]
pub usage: Option<ModelTokenUsage>,
}
#[derive(Debug, Clone, Default, PartialEq, Eq, Serialize, Deserialize)]
pub struct EmbeddingTemplateArgs {
#[serde(rename = "title", default)]
pub title: Option<String>,
}
#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, Serialize, Deserialize)]
pub enum SemanticTask {
#[serde(rename = "retrieval_document")]
RetrievalDocument,
#[serde(rename = "retrieval_query")]
RetrievalQuery,
#[serde(rename = "semantic_similarity")]
SemanticSimilarity,
#[serde(rename = "classification")]
Classification,
#[serde(rename = "clustering")]
Clustering,
#[serde(rename = "plain_text")]
PlainText,
}
impl SemanticTask {
pub fn as_str(self) -> &'static str {
match self {
SemanticTask::RetrievalDocument => "retrieval_document",
SemanticTask::RetrievalQuery => "retrieval_query",
SemanticTask::SemanticSimilarity => "semantic_similarity",
SemanticTask::Classification => "classification",
SemanticTask::Clustering => "clustering",
SemanticTask::PlainText => "plain_text",
}
}
}
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct ModelEmbeddingRequest {
#[serde(rename = "modelId", default)]
pub model_id: String,
#[serde(rename = "input", default)]
pub input: Vec<String>,
#[serde(rename = "templateArgs", default)]
pub template_args: Option<Vec<EmbeddingTemplateArgs>>,
#[serde(rename = "taskType")]
pub task_type: SemanticTask,
#[serde(rename = "dimensions", default)]
pub dimensions: Option<u32>,
}
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct ModelEmbeddingResponse {
#[serde(rename = "modelId", default)]
pub model_id: String,
#[serde(rename = "dimensions", default)]
pub dimensions: u32,
#[serde(rename = "vectors", default)]
pub vectors: Vec<Vec<f32>>,
}
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct ModelSimilarityRequest {
#[serde(rename = "modelId", default)]
pub model_id: String,
#[serde(rename = "query", default)]
pub query: Vec<f32>,
#[serde(rename = "targets", default)]
pub targets: Vec<Vec<f32>>,
}
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct ModelSimilarityResponse {
#[serde(rename = "modelId", default)]
pub model_id: String,
#[serde(rename = "similarities", default)]
pub similarities: Vec<f32>,
}
#[derive(Debug, Clone)]
pub enum ModelStreamChunk {
Delta(String),
Done { finish_reason: Option<String> },
Error(String),
}