File size: 7,133 Bytes

c55968e

//! # 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);
    }
}