feat: 开源混合流 + Ultrafast 推理框架（v5.2.1 提取版）

c55968e verified 1 day ago

2.57 kB

//! 核心类型：Provider 抽象 / 流式 chunk / 对话请求。
//!
//! 这一层刻意做到最小——你可以把任何 LLM 后端（Ollama / OpenAI 兼容接口 /
//! 自建推理服务）适配成 `Provider`，混合流与 Ultrafast 规则就能直接复用。

use std::future::Future;
use std::pin::Pin;

use futures::Stream;
use tokio_util::sync::CancellationToken;

/// 一个可取消的流式响应流。
pub type ProviderStream = Pin<Box<dyn Stream<Item = StreamChunk> + Send>>;

/// 装箱 future（避免 async-trait 依赖）。
pub type BoxFuture<T> = Pin<Box<dyn Future<Output = T> + Send>>;

/// 流式响应分片。
#[derive(Debug, Clone)]
pub enum StreamChunk {
    /// 正文增量。
    Delta { text: String },
    /// 思考过程增量（推理模型）。
    Reasoning { text: String },
    /// 接棒信号：混合流从"本机草稿"切到"云端正式"时下发，
    /// 前端应清空已显示的草稿内容再继续渲染后续 delta。
    Replaced { message: String },
    /// 正常结束（带用量统计）。
    Done {
        prompt_tokens: Option<u64>,
        completion_tokens: Option<u64>,
        eval_ms: Option<u64>,
        cached_tokens: Option<u64>,
    },
    /// 出错。
    Error { message: String },
    /// 被截断（达到 max_tokens 上限等）。
    Truncated { message: String },
}

/// 思考强度档位。
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum ThinkLevel {
    Off,
    Low,
    Medium,
    High,
    Frenzy,
}

impl Default for ThinkLevel {
    fn default() -> Self {
        ThinkLevel::Medium
    }
}

/// 一条对话消息。
#[derive(Debug, Clone)]
pub struct ChatMessage {
    pub role: MessageRole,
    pub content: String,
}

#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum MessageRole {
    System,
    User,
    Assistant,
}

/// 一次对话请求。`model` 是"模型锁"的落点：加速开关**不允许**改它。
#[derive(Debug, Clone, Default)]
pub struct ChatRequest {
    pub model: String,
    pub messages: Vec<ChatMessage>,
    pub think: Option<ThinkLevel>,
    pub max_tokens: Option<u32>,
    /// 工具定义等由调用方自行透传；框架不解释其语义。
    pub tools: Option<serde_json::Value>,
}

/// 任何能流式输出 LLM 回复的后端。
pub trait Provider: Send + Sync {
    /// 发起一次流式请求。`cancel` 被触发时应尽快结束流并返回。
    fn chat_stream(
        &self,
        request: ChatRequest,
        cancel: CancellationToken,
    ) -> BoxFuture<Result<ProviderStream, String>>;
}

	//! 核心类型：Provider 抽象 / 流式 chunk / 对话请求。
	//!
	//! 这一层刻意做到最小——你可以把任何 LLM 后端（Ollama / OpenAI 兼容接口 /
	//! 自建推理服务）适配成 `Provider`，混合流与 Ultrafast 规则就能直接复用。

	use std::future::Future;
	use std::pin::Pin;

	use futures::Stream;
	use tokio_util::sync::CancellationToken;

	/// 一个可取消的流式响应流。
	pub type ProviderStream = Pin<Box<dyn Stream<Item = StreamChunk> + Send>>;

	/// 装箱 future（避免 async-trait 依赖）。
	pub type BoxFuture<T> = Pin<Box<dyn Future<Output = T> + Send>>;

	/// 流式响应分片。
	#[derive(Debug, Clone)]
	pub enum StreamChunk {
	/// 正文增量。
	Delta { text: String },
	/// 思考过程增量（推理模型）。
	Reasoning { text: String },
	/// 接棒信号：混合流从"本机草稿"切到"云端正式"时下发，
	/// 前端应清空已显示的草稿内容再继续渲染后续 delta。
	Replaced { message: String },
	/// 正常结束（带用量统计）。
	Done {
	prompt_tokens: Option<u64>,
	completion_tokens: Option<u64>,
	eval_ms: Option<u64>,
	cached_tokens: Option<u64>,
	},
	/// 出错。
	Error { message: String },
	/// 被截断（达到 max_tokens 上限等）。
	Truncated { message: String },
	}

	/// 思考强度档位。
	#[derive(Debug, Clone, Copy, PartialEq, Eq)]
	pub enum ThinkLevel {
	Off,
	Low,
	Medium,
	High,
	Frenzy,
	}

	impl Default for ThinkLevel {
	fn default() -> Self {
	ThinkLevel::Medium
	}
	}

	/// 一条对话消息。
	#[derive(Debug, Clone)]
	pub struct ChatMessage {
	pub role: MessageRole,
	pub content: String,
	}

	#[derive(Debug, Clone, Copy, PartialEq, Eq)]
	pub enum MessageRole {
	System,
	User,
	Assistant,
	}

	/// 一次对话请求。`model` 是"模型锁"的落点：加速开关不允许改它。
	#[derive(Debug, Clone, Default)]
	pub struct ChatRequest {
	pub model: String,
	pub messages: Vec<ChatMessage>,
	pub think: Option<ThinkLevel>,
	pub max_tokens: Option<u32>,
	/// 工具定义等由调用方自行透传；框架不解释其语义。
	pub tools: Option<serde_json::Value>,
	}

	/// 任何能流式输出 LLM 回复的后端。
	pub trait Provider: Send + Sync {
	/// 发起一次流式请求。`cancel` 被触发时应尽快结束流并返回。
	fn chat_stream(
	&self,
	request: ChatRequest,
	cancel: CancellationToken,
	) -> BoxFuture<Result<ProviderStream, String>>;
	}