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

c55968e verified 1 day ago

7.13 kB

//! # perch-ultrafast
//!
//! **混合流 + Ultrafast 推理框架**——把"用户感觉慢"从架构层面解决。
//!
//! 从 [栖木 Perch](https://huggingface.co/hongcoding/perch)（macOS 本地优先 AI 工作台）
//! 中提取的可复用 Rust  crate。
//!
//! ## 解决的问题
//!
//! 云端大模型首 token 延迟 0.5~2s，用户盯着空屏幕等。本机小模型首 token 0.05~0.2s。
//!
//! **混合流**：两条流同时开——本机先出草稿（用户立刻有东西读），云端一到就接棒
//! （质量升级）。首字延迟 = 本机延迟，最终质量 = 云端质量。
//!
//! ## 快速开始
//!
//! ```rust,no_run
//! use std::sync::Arc;
//! use perch_ultrafast::{hybrid::{run_hybrid, HybridConfig}, types::*};
//! use tokio_util::sync::CancellationToken;
//!
//! # async fn demo(local: Arc<dyn Provider>, cloud: Arc<dyn Provider>, req: ChatRequest) {
//! let cancel = CancellationToken::new();
//! let result = run_hybrid(
//!     local.clone(),
//!     cloud,
//!     ChatRequest { think: Some(ThinkLevel::Off), ..req.clone() }, // 本机：关思考求快
//!     req,                                                        // 云端：原样
//!     cancel,
//!     HybridConfig::default(),
//!     |chunk| match chunk {
//!         StreamChunk::Delta { text } => print!("{text}"),
//!         StreamChunk::Replaced { .. } => println!("\n[云端接棒]"),
//!         StreamChunk::Done { .. } => println!("\n[done]"),
//!         _ => {}
//!     },
//! )
//! .await;
//! println!("最终 {} 字，云端接棒 = {}", result.content.chars().count(), result.used_cloud);
//! # }
//! ```
//!
//! ## 两个模块
//!
//! - [`hybrid`]：混合流 mux（本机草稿 → 云端接棒，含双向降级）
//! - [`ultrafast`]：加速规则（关思考 / 回复预算 / 简洁提示 / **模型锁红线**）
//! - [`pool`]：连接池单例（复用 TCP + TLS，启用压缩）
//!
//! ## 实测数据（栖木 Perch，MiMo 云端 + 本机 qwen3-0.6B）
//!
//! | 场景 | 首字延迟 | 说明 |
//! |---|---|---|
//! | 纯云端 | 0.88~1.39s | 用户盯着空屏幕 |
//! | 本机热身 | **0.04~0.18s** | 混合流草稿 |
//! | 混合流首字 | **< 0.2s** | 本机草稿立刻可读 |
//! | 云端接棒 | ~0.9s | 无缝替换，质量升级 |
//!
//! ## 设计红线
//!
//! - **加速绝不换模型**（不偷偷路由到小模型）
//! - **只改本次下发值**（不动用户持久化档位）
//! - **云端失败保留草稿**（不把用户的答案变成错误）
//! - **本机失败静默降级**（行为等同没用混合流，不打扰用户）

pub mod hybrid;
pub mod pool;
pub mod types;
pub mod ultrafast;

pub use hybrid::{run_hybrid, HybridConfig, HybridResult};
pub use types::{
    BoxFuture, ChatMessage, ChatRequest, MessageRole, Provider, ProviderStream, StreamChunk,
    ThinkLevel,
};
pub use ultrafast::{apply_ultrafast, UltrafastConfig};

#[cfg(test)]
mod tests {
    use super::*;
    use std::sync::Arc;
    use tokio_util::sync::CancellationToken;

    /// 测试用 provider：按给定脚本吐 chunk（真异步，不阻塞 runtime）。
    struct ScriptedProvider {
        chunks: Vec<StreamChunk>,
        delay_ms: u64,
    }

    impl Provider for ScriptedProvider {
        fn chat_stream(
            &self,
            _request: ChatRequest,
            _cancel: CancellationToken,
        ) -> BoxFuture<Result<ProviderStream, String>> {
            let chunks = self.chunks.clone();
            let delay = std::time::Duration::from_millis(self.delay_ms);
            Box::pin(async move {
                let (tx, rx) = tokio::sync::mpsc::channel::<StreamChunk>(16);
                tokio::spawn(async move {
                    for chunk in chunks {
                        tokio::time::sleep(delay).await;
                        if tx.send(chunk).await.is_err() {
                            return;
                        }
                    }
                });
                let s = futures::stream::unfold(rx, |mut rx| async move {
                    rx.recv().await.map(|chunk| (chunk, rx))
                });
                Ok(Box::pin(s) as ProviderStream)
            })
        }
    }

    fn request() -> ChatRequest {
        ChatRequest {
            model: "m".to_string(),
            messages: vec![],
            think: None,
            max_tokens: None,
            tools: None,
        }
    }

    #[tokio::test]
    async fn cloud_takes_over_and_replaces_draft() {
        let local = Arc::new(ScriptedProvider {
            chunks: vec![
                StreamChunk::Delta { text: "draft".into() },
                StreamChunk::Done {
                    prompt_tokens: None,
                    completion_tokens: None,
                    eval_ms: None,
                    cached_tokens: None,
                },
            ],
            delay_ms: 0,
        });
        let cloud = Arc::new(ScriptedProvider {
            chunks: vec![
                StreamChunk::Delta { text: "final".into() },
                StreamChunk::Done {
                    prompt_tokens: None,
                    completion_tokens: None,
                    eval_ms: None,
                    cached_tokens: None,
                },
            ],
            delay_ms: 30,
        });
        let mut events = Vec::new();
        let result = run_hybrid(
            local,
            cloud,
            request(),
            request(),
            CancellationToken::new(),
            HybridConfig::default(),
            |c| events.push(c),
        )
        .await;
        assert!(result.used_cloud, "云端应该接棒");
        assert_eq!(result.content, "final");
        let replaced = events.iter().any(|c| matches!(c, StreamChunk::Replaced { .. }));
        assert!(replaced, "接棒时应下发 Replaced 让前端清草稿");
    }

    #[tokio::test]
    async fn cloud_failure_falls_back_to_local() {
        struct FailingProvider;
        impl Provider for FailingProvider {
            fn chat_stream(
                &self,
                _request: ChatRequest,
                _cancel: CancellationToken,
            ) -> BoxFuture<Result<ProviderStream, String>> {
                Box::pin(async { Err("cloud down".to_string()) })
            }
        }
        let local = Arc::new(ScriptedProvider {
            chunks: vec![
                StreamChunk::Delta { text: "local only".into() },
                StreamChunk::Done {
                    prompt_tokens: None,
                    completion_tokens: None,
                    eval_ms: None,
                    cached_tokens: None,
                },
            ],
            delay_ms: 0,
        });
        let result = run_hybrid(
            local,
            Arc::new(FailingProvider),
            request(),
            request(),
            CancellationToken::new(),
            HybridConfig::default(),
            |_| {},
        )
        .await;
        assert!(!result.used_cloud);
        assert_eq!(result.content, "local only", "云端挂了应保留本地草稿");
        assert!(result.finished);
    }
}

	//! # perch-ultrafast
	//!
	//! 混合流 + Ultrafast 推理框架——把"用户感觉慢"从架构层面解决。
	//!
	//! 从 [栖木 Perch](https://huggingface.co/hongcoding/perch)（macOS 本地优先 AI 工作台）
	//! 中提取的可复用 Rust crate。
	//!
	//! ## 解决的问题
	//!
	//! 云端大模型首 token 延迟 0.5~2s，用户盯着空屏幕等。本机小模型首 token 0.05~0.2s。
	//!
	//! 混合流：两条流同时开——本机先出草稿（用户立刻有东西读），云端一到就接棒
	//! （质量升级）。首字延迟 = 本机延迟，最终质量 = 云端质量。
	//!
	//! ## 快速开始
	//!
	//! ```rust,no_run
	//! use std::sync::Arc;
	//! use perch_ultrafast::{hybrid::{run_hybrid, HybridConfig}, types::*};
	//! use tokio_util::sync::CancellationToken;
	//!
	//! # async fn demo(local: Arc<dyn Provider>, cloud: Arc<dyn Provider>, req: ChatRequest) {
	//! let cancel = CancellationToken::new();
	//! let result = run_hybrid(
	//! local.clone(),
	//! cloud,
	//! ChatRequest { think: Some(ThinkLevel::Off), ..req.clone() }, // 本机：关思考求快
	//! req, // 云端：原样
	//! cancel,
	//! HybridConfig::default(),
	//! \|chunk\| match chunk {
	//! StreamChunk::Delta { text } => print!("{text}"),
	//! StreamChunk::Replaced { .. } => println!("\n[云端接棒]"),
	//! StreamChunk::Done { .. } => println!("\n[done]"),
	//! _ => {}
	//! },
	//! )
	//! .await;
	//! println!("最终 {} 字，云端接棒 = {}", result.content.chars().count(), result.used_cloud);
	//! # }
	//! ```
	//!
	//! ## 两个模块
	//!
	//! - [`hybrid`]：混合流 mux（本机草稿 → 云端接棒，含双向降级）
	//! - [`ultrafast`]：加速规则（关思考 / 回复预算 / 简洁提示 / 模型锁红线）
	//! - [`pool`]：连接池单例（复用 TCP + TLS，启用压缩）
	//!
	//! ## 实测数据（栖木 Perch，MiMo 云端 + 本机 qwen3-0.6B）
	//!
	//! \| 场景 \| 首字延迟 \| 说明 \|
	//! \|---\|---\|---\|
	//! \| 纯云端 \| 0.88~1.39s \| 用户盯着空屏幕 \|
	//! \| 本机热身 \| 0.04~0.18s \| 混合流草稿 \|
	//! \| 混合流首字 \| < 0.2s \| 本机草稿立刻可读 \|
	//! \| 云端接棒 \| ~0.9s \| 无缝替换，质量升级 \|
	//!
	//! ## 设计红线
	//!
	//! - 加速绝不换模型（不偷偷路由到小模型）
	//! - 只改本次下发值（不动用户持久化档位）
	//! - 云端失败保留草稿（不把用户的答案变成错误）
	//! - 本机失败静默降级（行为等同没用混合流，不打扰用户）

	pub mod hybrid;
	pub mod pool;
	pub mod types;
	pub mod ultrafast;

	pub use hybrid::{run_hybrid, HybridConfig, HybridResult};
	pub use types::{
	BoxFuture, ChatMessage, ChatRequest, MessageRole, Provider, ProviderStream, StreamChunk,
	ThinkLevel,
	};
	pub use ultrafast::{apply_ultrafast, UltrafastConfig};

	#[cfg(test)]
	mod tests {
	use super::*;
	use std::sync::Arc;
	use tokio_util::sync::CancellationToken;

	/// 测试用 provider：按给定脚本吐 chunk（真异步，不阻塞 runtime）。
	struct ScriptedProvider {
	chunks: Vec<StreamChunk>,
	delay_ms: u64,
	}

	impl Provider for ScriptedProvider {
	fn chat_stream(
	&self,
	_request: ChatRequest,
	_cancel: CancellationToken,
	) -> BoxFuture<Result<ProviderStream, String>> {
	let chunks = self.chunks.clone();
	let delay = std::time::Duration::from_millis(self.delay_ms);
	Box::pin(async move {
	let (tx, rx) = tokio::sync::mpsc::channel::<StreamChunk>(16);
	tokio::spawn(async move {
	for chunk in chunks {
	tokio::time::sleep(delay).await;
	if tx.send(chunk).await.is_err() {
	return;
	}
	}
	});
	let s = futures::stream::unfold(rx, \|mut rx\| async move {
	rx.recv().await.map(\|chunk\| (chunk, rx))
	});
	Ok(Box::pin(s) as ProviderStream)
	})
	}
	}

	fn request() -> ChatRequest {
	ChatRequest {
	model: "m".to_string(),
	messages: vec![],
	think: None,
	max_tokens: None,
	tools: None,
	}
	}

	#[tokio::test]
	async fn cloud_takes_over_and_replaces_draft() {
	let local = Arc::new(ScriptedProvider {
	chunks: vec![
	StreamChunk::Delta { text: "draft".into() },
	StreamChunk::Done {
	prompt_tokens: None,
	completion_tokens: None,
	eval_ms: None,
	cached_tokens: None,
	},
	],
	delay_ms: 0,
	});
	let cloud = Arc::new(ScriptedProvider {
	chunks: vec![
	StreamChunk::Delta { text: "final".into() },
	StreamChunk::Done {
	prompt_tokens: None,
	completion_tokens: None,
	eval_ms: None,
	cached_tokens: None,
	},
	],
	delay_ms: 30,
	});
	let mut events = Vec::new();
	let result = run_hybrid(
	local,
	cloud,
	request(),
	request(),
	CancellationToken::new(),
	HybridConfig::default(),
	\|c\| events.push(c),
	)
	.await;
	assert!(result.used_cloud, "云端应该接棒");
	assert_eq!(result.content, "final");
	let replaced = events.iter().any(\|c\| matches!(c, StreamChunk::Replaced { .. }));
	assert!(replaced, "接棒时应下发 Replaced 让前端清草稿");
	}

	#[tokio::test]
	async fn cloud_failure_falls_back_to_local() {
	struct FailingProvider;
	impl Provider for FailingProvider {
	fn chat_stream(
	&self,
	_request: ChatRequest,
	_cancel: CancellationToken,
	) -> BoxFuture<Result<ProviderStream, String>> {
	Box::pin(async { Err("cloud down".to_string()) })
	}
	}
	let local = Arc::new(ScriptedProvider {
	chunks: vec![
	StreamChunk::Delta { text: "local only".into() },
	StreamChunk::Done {
	prompt_tokens: None,
	completion_tokens: None,
	eval_ms: None,
	cached_tokens: None,
	},
	],
	delay_ms: 0,
	});
	let result = run_hybrid(
	local,
	Arc::new(FailingProvider),
	request(),
	request(),
	CancellationToken::new(),
	HybridConfig::default(),
	\|_\| {},
	)
	.await;
	assert!(!result.used_cloud);
	assert_eq!(result.content, "local only", "云端挂了应保留本地草稿");
	assert!(result.finished);
	}
	}