feat(llm): 新门面接入 DeepSeek 与 Kimi,实现 thinking 流式机制
- provider 枚举增加 DeepSeek/Kimi 变体(base_url/API key 走配置) - thinking 参数经 additional_params 透传(enabled/disabled 显式发送) - Kimi 温度特例:thinking 恒 1.0、普通恒 0.6;DeepSeek thinking 省略温度 - 流式路径:Reasoning/ReasoningDelta 事件驱动思考状态显示,正文聚合, 仅思考无正文时给出与旧实现一致的错误提示 - 新增 7 个测试:请求形状捕获 + SSE mock 事件序列验证
This commit is contained in:
@@ -4,14 +4,17 @@
|
|||||||
//! 全部 provider 接入并切换后,旧实现将被删除。
|
//! 全部 provider 接入并切换后,旧实现将被删除。
|
||||||
|
|
||||||
use crate::config::manager::ConfigManager;
|
use crate::config::manager::ConfigManager;
|
||||||
|
use crate::llm::thinking::{create_console_thinking_state, ThinkingStateManager};
|
||||||
use anyhow::{Context, Result, bail};
|
use anyhow::{Context, Result, bail};
|
||||||
use rig_core::{
|
use rig_core::{
|
||||||
client::{CompletionClient, Nothing, VerifyClient},
|
client::{CompletionClient, Nothing, VerifyClient},
|
||||||
completion::{AssistantContent, CompletionError, CompletionModel},
|
completion::{AssistantContent, CompletionError, CompletionModel},
|
||||||
http_client::{HttpClientExt, ReqwestClient},
|
http_client::{HttpClientExt, ReqwestClient},
|
||||||
providers::ollama,
|
providers::{deepseek, moonshot, ollama},
|
||||||
|
streaming::StreamedAssistantContent,
|
||||||
wasm_compat::{WasmCompatSend, WasmCompatSync},
|
wasm_compat::{WasmCompatSend, WasmCompatSync},
|
||||||
};
|
};
|
||||||
|
use std::sync::Arc;
|
||||||
use std::time::Duration;
|
use std::time::Duration;
|
||||||
|
|
||||||
/// LLM 客户端运行时配置(与用户配置解耦的参数)。
|
/// LLM 客户端运行时配置(与用户配置解耦的参数)。
|
||||||
@@ -39,6 +42,8 @@ impl Default for LlmClientConfig {
|
|||||||
#[derive(Clone)]
|
#[derive(Clone)]
|
||||||
pub enum Backend<H = ReqwestClient> {
|
pub enum Backend<H = ReqwestClient> {
|
||||||
Ollama(ollama::Client<H>),
|
Ollama(ollama::Client<H>),
|
||||||
|
DeepSeek(deepseek::Client<H>),
|
||||||
|
Kimi(moonshot::Client<H>),
|
||||||
}
|
}
|
||||||
|
|
||||||
/// 基于 rig 的 LLM 客户端门面。
|
/// 基于 rig 的 LLM 客户端门面。
|
||||||
@@ -49,6 +54,19 @@ pub struct LlmClient<H = ReqwestClient> {
|
|||||||
provider: String,
|
provider: String,
|
||||||
config: LlmClientConfig,
|
config: LlmClientConfig,
|
||||||
thinking_enabled: bool,
|
thinking_enabled: bool,
|
||||||
|
thinking_state: Option<Arc<ThinkingStateManager>>,
|
||||||
|
}
|
||||||
|
|
||||||
|
/// 一次生成请求的 provider 相关参数。
|
||||||
|
#[derive(Debug, Clone, PartialEq)]
|
||||||
|
struct RequestOptions {
|
||||||
|
/// `None` 表示省略该参数(如 thinking 模式下不传温度)。
|
||||||
|
temperature: Option<f64>,
|
||||||
|
max_tokens: Option<u64>,
|
||||||
|
/// 透传到请求体的 provider 专属参数(thinking 开关等)。
|
||||||
|
additional_params: Option<serde_json::Value>,
|
||||||
|
/// 是否走流式路径。
|
||||||
|
stream: bool,
|
||||||
}
|
}
|
||||||
|
|
||||||
impl LlmClient {
|
impl LlmClient {
|
||||||
@@ -79,23 +97,47 @@ impl LlmClient {
|
|||||||
|
|
||||||
let backend = match provider.as_str() {
|
let backend = match provider.as_str() {
|
||||||
"ollama" => Backend::Ollama(build_ollama_client(&base_url, http)?),
|
"ollama" => Backend::Ollama(build_ollama_client(&base_url, http)?),
|
||||||
|
"deepseek" => {
|
||||||
|
let key = required_api_key(&provider, manager.get_api_key())?;
|
||||||
|
Backend::DeepSeek(build_deepseek_client(&key, &base_url, http)?)
|
||||||
|
}
|
||||||
|
"kimi" => {
|
||||||
|
let key = required_api_key(&provider, manager.get_api_key())?;
|
||||||
|
Backend::Kimi(build_kimi_client(&key, &base_url, http)?)
|
||||||
|
}
|
||||||
// 其余 provider 在后续 ticket 接入,扩张期内新门面尚未被应用使用。
|
// 其余 provider 在后续 ticket 接入,扩张期内新门面尚未被应用使用。
|
||||||
"openai" | "anthropic" | "kimi" | "deepseek" | "openrouter" => {
|
"openai" | "anthropic" | "openrouter" => {
|
||||||
bail!("Provider '{}' is not available in the new LLM backend yet", provider)
|
bail!("Provider '{}' is not available in the new LLM backend yet", provider)
|
||||||
}
|
}
|
||||||
_ => bail!("Unknown LLM provider: {}", provider),
|
_ => bail!("Unknown LLM provider: {}", provider),
|
||||||
};
|
};
|
||||||
|
|
||||||
|
let thinking_state = if thinking_enabled && supports_thinking(&provider) {
|
||||||
|
Some(create_console_thinking_state())
|
||||||
|
} else {
|
||||||
|
None
|
||||||
|
};
|
||||||
|
|
||||||
Ok(Self {
|
Ok(Self {
|
||||||
backend,
|
backend,
|
||||||
model,
|
model,
|
||||||
provider,
|
provider,
|
||||||
config: cfg,
|
config: cfg,
|
||||||
thinking_enabled,
|
thinking_enabled,
|
||||||
|
thinking_state,
|
||||||
})
|
})
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
|
/// 该 provider 是否支持 thinking 模式(与旧实现的白名单一致)。
|
||||||
|
pub(crate) fn supports_thinking(provider: &str) -> bool {
|
||||||
|
matches!(provider, "deepseek" | "kimi" | "anthropic" | "openai")
|
||||||
|
}
|
||||||
|
|
||||||
|
fn required_api_key(provider: &str, key: Option<String>) -> Result<String> {
|
||||||
|
key.ok_or_else(|| anyhow::anyhow!("{} API key not configured", provider_display_name(provider)))
|
||||||
|
}
|
||||||
|
|
||||||
impl<H> LlmClient<H> {
|
impl<H> LlmClient<H> {
|
||||||
/// 内部与测试构造入口。
|
/// 内部与测试构造入口。
|
||||||
pub(crate) fn new(
|
pub(crate) fn new(
|
||||||
@@ -104,6 +146,7 @@ impl<H> LlmClient<H> {
|
|||||||
provider: impl Into<String>,
|
provider: impl Into<String>,
|
||||||
config: LlmClientConfig,
|
config: LlmClientConfig,
|
||||||
thinking_enabled: bool,
|
thinking_enabled: bool,
|
||||||
|
thinking_state: Option<Arc<ThinkingStateManager>>,
|
||||||
) -> Self {
|
) -> Self {
|
||||||
Self {
|
Self {
|
||||||
backend,
|
backend,
|
||||||
@@ -111,12 +154,42 @@ impl<H> LlmClient<H> {
|
|||||||
provider: provider.into(),
|
provider: provider.into(),
|
||||||
config,
|
config,
|
||||||
thinking_enabled,
|
thinking_enabled,
|
||||||
|
thinking_state,
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
/// 当前是否处于 thinking 模式。
|
/// 按 provider × thinking 计算请求参数。
|
||||||
pub(crate) fn thinking_enabled(&self) -> bool {
|
fn request_options(&self) -> RequestOptions {
|
||||||
self.thinking_enabled
|
match &self.backend {
|
||||||
|
Backend::Ollama(_) => RequestOptions {
|
||||||
|
temperature: Some(self.config.temperature),
|
||||||
|
max_tokens: Some(self.config.max_tokens),
|
||||||
|
additional_params: None,
|
||||||
|
stream: false,
|
||||||
|
},
|
||||||
|
// DeepSeek:thinking 时省略温度;thinking 参数始终显式传(与旧实现一致)。
|
||||||
|
Backend::DeepSeek(_) => RequestOptions {
|
||||||
|
temperature: if self.thinking_enabled {
|
||||||
|
None
|
||||||
|
} else {
|
||||||
|
Some(self.config.temperature)
|
||||||
|
},
|
||||||
|
max_tokens: Some(self.config.max_tokens),
|
||||||
|
additional_params: Some(serde_json::json!({
|
||||||
|
"thinking": { "type": if self.thinking_enabled { "enabled" } else { "disabled" } }
|
||||||
|
})),
|
||||||
|
stream: self.thinking_enabled,
|
||||||
|
},
|
||||||
|
// Kimi:Moonshot API 要求 thinking 恒 1.0、普通恒 0.6(忽略用户配置)。
|
||||||
|
Backend::Kimi(_) => RequestOptions {
|
||||||
|
temperature: Some(if self.thinking_enabled { 1.0 } else { 0.6 }),
|
||||||
|
max_tokens: Some(self.config.max_tokens),
|
||||||
|
additional_params: Some(serde_json::json!({
|
||||||
|
"thinking": { "type": if self.thinking_enabled { "enabled" } else { "disabled" } }
|
||||||
|
})),
|
||||||
|
stream: self.thinking_enabled,
|
||||||
|
},
|
||||||
|
}
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
@@ -132,12 +205,21 @@ where
|
|||||||
+ WasmCompatSync
|
+ WasmCompatSync
|
||||||
+ 'static,
|
+ 'static,
|
||||||
{
|
{
|
||||||
/// 生成文本(thinking=false 的非流式路径)。
|
/// 生成文本;thinking 模式下走流式路径并驱动思考状态显示。
|
||||||
pub async fn generate(&self, system: Option<&str>, user: &str) -> Result<String> {
|
pub async fn generate(&self, system: Option<&str>, user: &str) -> Result<String> {
|
||||||
|
let options = self.request_options();
|
||||||
match &self.backend {
|
match &self.backend {
|
||||||
Backend::Ollama(client) => {
|
Backend::Ollama(client) => {
|
||||||
let model = client.completion_model(self.model.as_str());
|
let model = client.completion_model(self.model.as_str());
|
||||||
complete(&model, &self.provider, system, user, &self.config).await
|
self.run(&model, system, user, &options).await
|
||||||
|
}
|
||||||
|
Backend::DeepSeek(client) => {
|
||||||
|
let model = client.completion_model(self.model.as_str());
|
||||||
|
self.run(&model, system, user, &options).await
|
||||||
|
}
|
||||||
|
Backend::Kimi(client) => {
|
||||||
|
let model = client.completion_model(self.model.as_str());
|
||||||
|
self.run(&model, system, user, &options).await
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
@@ -146,8 +228,92 @@ where
|
|||||||
pub async fn is_available(&self) -> bool {
|
pub async fn is_available(&self) -> bool {
|
||||||
match &self.backend {
|
match &self.backend {
|
||||||
Backend::Ollama(client) => client.verify().await.is_ok(),
|
Backend::Ollama(client) => client.verify().await.is_ok(),
|
||||||
|
Backend::DeepSeek(client) => client.verify().await.is_ok(),
|
||||||
|
Backend::Kimi(client) => client.verify().await.is_ok(),
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
|
/// 按请求参数选择流式/非流式路径。
|
||||||
|
async fn run<M: CompletionModel>(
|
||||||
|
&self,
|
||||||
|
model: &M,
|
||||||
|
system: Option<&str>,
|
||||||
|
user: &str,
|
||||||
|
options: &RequestOptions,
|
||||||
|
) -> Result<String> {
|
||||||
|
if options.stream {
|
||||||
|
self.stream_completion(model, system, user, options).await
|
||||||
|
} else {
|
||||||
|
complete(model, &self.provider, system, user, options).await
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
/// 流式补全:聚合正文,驱动思考状态显示(与旧实现行为一致)。
|
||||||
|
async fn stream_completion<M: CompletionModel>(
|
||||||
|
&self,
|
||||||
|
model: &M,
|
||||||
|
system: Option<&str>,
|
||||||
|
user: &str,
|
||||||
|
options: &RequestOptions,
|
||||||
|
) -> Result<String> {
|
||||||
|
let request = build_request(model, system, user, options);
|
||||||
|
let mut response = model
|
||||||
|
.stream(request)
|
||||||
|
.await
|
||||||
|
.map_err(|e| map_completion_error(&self.provider, e))?;
|
||||||
|
|
||||||
|
use futures_util::StreamExt;
|
||||||
|
|
||||||
|
let state = self.thinking_state.as_deref();
|
||||||
|
let mut text = String::new();
|
||||||
|
let mut has_reasoning = false;
|
||||||
|
let mut has_content = false;
|
||||||
|
|
||||||
|
while let Some(event) = response.next().await {
|
||||||
|
let event = event.map_err(|e| map_completion_error(&self.provider, e))?;
|
||||||
|
match event {
|
||||||
|
StreamedAssistantContent::Text(chunk) => {
|
||||||
|
if has_reasoning && !has_content
|
||||||
|
&& let Some(state) = state
|
||||||
|
{
|
||||||
|
state.end_thinking();
|
||||||
|
}
|
||||||
|
has_content = true;
|
||||||
|
text.push_str(&chunk.text);
|
||||||
|
}
|
||||||
|
StreamedAssistantContent::Reasoning(_)
|
||||||
|
| StreamedAssistantContent::ReasoningDelta { .. } if !has_reasoning => {
|
||||||
|
has_reasoning = true;
|
||||||
|
if let Some(state) = state {
|
||||||
|
state.start_thinking();
|
||||||
|
}
|
||||||
|
}
|
||||||
|
_ => {}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
if let Some(state) = state {
|
||||||
|
state.end_thinking();
|
||||||
|
}
|
||||||
|
|
||||||
|
let text = text.trim().to_string();
|
||||||
|
if text.is_empty() {
|
||||||
|
if has_reasoning && !has_content {
|
||||||
|
bail!(
|
||||||
|
"{} returned reasoning content but no final answer. \
|
||||||
|
The model may have entered an incomplete reasoning state. \
|
||||||
|
Please try again or disable thinking mode.",
|
||||||
|
provider_display_name(&self.provider)
|
||||||
|
);
|
||||||
|
}
|
||||||
|
bail!(
|
||||||
|
"No response from {}. \
|
||||||
|
If thinking mode is enabled, try disabling it or ensure the model supports it.",
|
||||||
|
provider_display_name(&self.provider)
|
||||||
|
);
|
||||||
|
}
|
||||||
|
Ok(text)
|
||||||
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
/// 构建 Ollama 客户端(无 API key;支持自定义 base_url 与注入的 HTTP 后端)。
|
/// 构建 Ollama 客户端(无 API key;支持自定义 base_url 与注入的 HTTP 后端)。
|
||||||
@@ -160,26 +326,58 @@ fn build_ollama_client(base_url: &str, http: ReqwestClient) -> Result<ollama::Cl
|
|||||||
.map_err(|e| anyhow::anyhow!("Failed to build Ollama client: {}", e))
|
.map_err(|e| anyhow::anyhow!("Failed to build Ollama client: {}", e))
|
||||||
}
|
}
|
||||||
|
|
||||||
/// 非流式补全:系统提示、温度、max_tokens 映射到统一请求,聚合正文文本。
|
/// 构建 DeepSeek 客户端(OpenAI 兼容 API)。
|
||||||
|
fn build_deepseek_client(key: &str, base_url: &str, http: ReqwestClient) -> Result<deepseek::Client> {
|
||||||
|
deepseek::Client::builder()
|
||||||
|
.api_key(key)
|
||||||
|
.base_url(base_url)
|
||||||
|
.http_client(http)
|
||||||
|
.build()
|
||||||
|
.map_err(|e| anyhow::anyhow!("Failed to build DeepSeek client: {}", e))
|
||||||
|
}
|
||||||
|
|
||||||
|
/// 构建 Kimi(Moonshot)客户端(OpenAI 兼容 API)。
|
||||||
|
fn build_kimi_client(key: &str, base_url: &str, http: ReqwestClient) -> Result<moonshot::Client> {
|
||||||
|
moonshot::Client::builder()
|
||||||
|
.api_key(key)
|
||||||
|
.base_url(base_url)
|
||||||
|
.http_client(http)
|
||||||
|
.build()
|
||||||
|
.map_err(|e| anyhow::anyhow!("Failed to build Kimi client: {}", e))
|
||||||
|
}
|
||||||
|
|
||||||
|
/// 组装统一的补全请求(preamble/temperature/max_tokens/additional_params)。
|
||||||
|
fn build_request<M: CompletionModel>(
|
||||||
|
model: &M,
|
||||||
|
system: Option<&str>,
|
||||||
|
user: &str,
|
||||||
|
options: &RequestOptions,
|
||||||
|
) -> rig_core::completion::CompletionRequest {
|
||||||
|
let mut builder = model.completion_request(user);
|
||||||
|
if let Some(sys) = system {
|
||||||
|
builder = builder.preamble(sys.to_string());
|
||||||
|
}
|
||||||
|
builder = builder
|
||||||
|
.temperature_opt(options.temperature)
|
||||||
|
.max_tokens_opt(options.max_tokens);
|
||||||
|
if let Some(params) = &options.additional_params {
|
||||||
|
builder = builder.additional_params(params.clone());
|
||||||
|
}
|
||||||
|
builder.build()
|
||||||
|
}
|
||||||
|
|
||||||
|
/// 非流式补全:聚合正文文本。
|
||||||
async fn complete<M>(
|
async fn complete<M>(
|
||||||
model: &M,
|
model: &M,
|
||||||
provider: &str,
|
provider: &str,
|
||||||
system: Option<&str>,
|
system: Option<&str>,
|
||||||
user: &str,
|
user: &str,
|
||||||
config: &LlmClientConfig,
|
options: &RequestOptions,
|
||||||
) -> Result<String>
|
) -> Result<String>
|
||||||
where
|
where
|
||||||
M: CompletionModel,
|
M: CompletionModel,
|
||||||
{
|
{
|
||||||
let mut builder = model.completion_request(user);
|
let request = build_request(model, system, user, options);
|
||||||
if let Some(sys) = system {
|
|
||||||
builder = builder.preamble(sys.to_string());
|
|
||||||
}
|
|
||||||
let request = builder
|
|
||||||
.temperature(config.temperature)
|
|
||||||
.max_tokens(config.max_tokens)
|
|
||||||
.build();
|
|
||||||
|
|
||||||
let response = model
|
let response = model
|
||||||
.completion(request)
|
.completion(request)
|
||||||
.await
|
.await
|
||||||
@@ -234,7 +432,8 @@ pub(crate) fn map_completion_error(provider: &str, e: CompletionError) -> anyhow
|
|||||||
#[cfg(test)]
|
#[cfg(test)]
|
||||||
mod tests {
|
mod tests {
|
||||||
use super::*;
|
use super::*;
|
||||||
use rig_core::test_utils::{MockHttpResponse, RecordingHttpClient};
|
use rig_core::test_utils::{MockHttpResponse, MockStreamingClient, RecordingHttpClient};
|
||||||
|
use std::sync::atomic::{AtomicUsize, Ordering};
|
||||||
|
|
||||||
const OLLAMA_OK: &str = r#"{
|
const OLLAMA_OK: &str = r#"{
|
||||||
"model": "llama3.2",
|
"model": "llama3.2",
|
||||||
@@ -243,40 +442,59 @@ mod tests {
|
|||||||
"done": true
|
"done": true
|
||||||
}"#;
|
}"#;
|
||||||
|
|
||||||
/// 用注入的 HTTP 后端构建 Ollama 客户端(与生产构造路径一致)。
|
const DEEPSEEK_OK: &str = r#"{
|
||||||
fn build_backend(recorder: RecordingHttpClient) -> Backend<RecordingHttpClient> {
|
"id": "cmpl-1",
|
||||||
Backend::Ollama(
|
"model": "deepseek-v4-flash",
|
||||||
ollama::Client::builder()
|
"choices": [{"index": 0, "message": {"role": "assistant", "content": "hello from deepseek"}, "finish_reason": "stop"}],
|
||||||
.api_key(Nothing)
|
"usage": {"completion_tokens": 3, "prompt_tokens": 4, "prompt_cache_hit_tokens": 0, "prompt_cache_miss_tokens": 0, "total_tokens": 7}
|
||||||
.base_url("http://localhost:11434")
|
}"#;
|
||||||
.http_client(recorder)
|
|
||||||
.build()
|
|
||||||
.expect("build ollama client with mock backend"),
|
|
||||||
)
|
|
||||||
}
|
|
||||||
|
|
||||||
fn client_with(
|
fn test_config() -> LlmClientConfig {
|
||||||
recorder: RecordingHttpClient,
|
|
||||||
) -> (LlmClient<RecordingHttpClient>, RecordingHttpClient) {
|
|
||||||
let backend = build_backend(recorder.clone());
|
|
||||||
let client = LlmClient::new(
|
|
||||||
backend,
|
|
||||||
"llama3.2",
|
|
||||||
"ollama",
|
|
||||||
LlmClientConfig {
|
LlmClientConfig {
|
||||||
max_tokens: 123,
|
max_tokens: 123,
|
||||||
temperature: 0.5,
|
temperature: 0.5,
|
||||||
timeout: Duration::from_secs(30),
|
timeout: Duration::from_secs(30),
|
||||||
},
|
}
|
||||||
false,
|
}
|
||||||
|
|
||||||
|
fn ollama_client(
|
||||||
|
recorder: RecordingHttpClient,
|
||||||
|
) -> (LlmClient<RecordingHttpClient>, RecordingHttpClient) {
|
||||||
|
let backend = Backend::Ollama(
|
||||||
|
ollama::Client::builder()
|
||||||
|
.api_key(Nothing)
|
||||||
|
.base_url("http://localhost:11434")
|
||||||
|
.http_client(recorder.clone())
|
||||||
|
.build()
|
||||||
|
.expect("build ollama client with mock backend"),
|
||||||
);
|
);
|
||||||
|
let client = LlmClient::new(backend, "llama3.2", "ollama", test_config(), false, None);
|
||||||
(client, recorder)
|
(client, recorder)
|
||||||
}
|
}
|
||||||
|
|
||||||
|
fn deepseek_client(
|
||||||
|
recorder: RecordingHttpClient,
|
||||||
|
thinking: bool,
|
||||||
|
) -> (LlmClient<RecordingHttpClient>, RecordingHttpClient) {
|
||||||
|
let backend = Backend::DeepSeek(
|
||||||
|
deepseek::Client::builder()
|
||||||
|
.api_key("sk-test")
|
||||||
|
.base_url("https://api.deepseek.com/v1")
|
||||||
|
.http_client(recorder.clone())
|
||||||
|
.build()
|
||||||
|
.expect("build deepseek client with mock backend"),
|
||||||
|
);
|
||||||
|
let client =
|
||||||
|
LlmClient::new(backend, "deepseek-v4-flash", "deepseek", test_config(), thinking, None);
|
||||||
|
(client, recorder)
|
||||||
|
}
|
||||||
|
|
||||||
|
// ---- Ollama 基本路径 ----
|
||||||
|
|
||||||
#[tokio::test]
|
#[tokio::test]
|
||||||
async fn generate_maps_request_params_and_returns_text() {
|
async fn generate_maps_request_params_and_returns_text() {
|
||||||
let recorder = RecordingHttpClient::new(OLLAMA_OK);
|
let recorder = RecordingHttpClient::new(OLLAMA_OK);
|
||||||
let (client, recorder) = client_with(recorder);
|
let (client, recorder) = ollama_client(recorder);
|
||||||
|
|
||||||
let text = client.generate(Some("be helpful"), "say hi").await.unwrap();
|
let text = client.generate(Some("be helpful"), "say hi").await.unwrap();
|
||||||
assert_eq!(text, "hello from ollama");
|
assert_eq!(text, "hello from ollama");
|
||||||
@@ -299,7 +517,7 @@ mod tests {
|
|||||||
#[tokio::test]
|
#[tokio::test]
|
||||||
async fn generate_without_system_omits_preamble() {
|
async fn generate_without_system_omits_preamble() {
|
||||||
let recorder = RecordingHttpClient::new(OLLAMA_OK);
|
let recorder = RecordingHttpClient::new(OLLAMA_OK);
|
||||||
let (client, recorder) = client_with(recorder);
|
let (client, recorder) = ollama_client(recorder);
|
||||||
|
|
||||||
client.generate(None, "say hi").await.unwrap();
|
client.generate(None, "say hi").await.unwrap();
|
||||||
|
|
||||||
@@ -316,7 +534,7 @@ mod tests {
|
|||||||
http::StatusCode::BAD_GATEWAY,
|
http::StatusCode::BAD_GATEWAY,
|
||||||
"upstream exploded".into(),
|
"upstream exploded".into(),
|
||||||
));
|
));
|
||||||
let (client, _) = client_with(recorder);
|
let (client, _) = ollama_client(recorder);
|
||||||
|
|
||||||
let err = client.generate(None, "hi").await.unwrap_err();
|
let err = client.generate(None, "hi").await.unwrap_err();
|
||||||
let msg = err.to_string();
|
let msg = err.to_string();
|
||||||
@@ -330,7 +548,7 @@ mod tests {
|
|||||||
let recorder = RecordingHttpClient::new(
|
let recorder = RecordingHttpClient::new(
|
||||||
r#"{"model":"llama3.2","created_at":"x","message":{"role":"assistant","content":""},"done":true}"#,
|
r#"{"model":"llama3.2","created_at":"x","message":{"role":"assistant","content":""},"done":true}"#,
|
||||||
);
|
);
|
||||||
let (client, _) = client_with(recorder);
|
let (client, _) = ollama_client(recorder);
|
||||||
|
|
||||||
let err = client.generate(None, "hi").await.unwrap_err();
|
let err = client.generate(None, "hi").await.unwrap_err();
|
||||||
assert_eq!(err.to_string(), "No response from Ollama");
|
assert_eq!(err.to_string(), "No response from Ollama");
|
||||||
@@ -339,7 +557,7 @@ mod tests {
|
|||||||
#[tokio::test]
|
#[tokio::test]
|
||||||
async fn is_available_uses_verify_endpoint() {
|
async fn is_available_uses_verify_endpoint() {
|
||||||
let recorder = RecordingHttpClient::new("");
|
let recorder = RecordingHttpClient::new("");
|
||||||
let (client, recorder) = client_with(recorder);
|
let (client, recorder) = ollama_client(recorder);
|
||||||
|
|
||||||
assert!(client.is_available().await);
|
assert!(client.is_available().await);
|
||||||
let captured = recorder.requests();
|
let captured = recorder.requests();
|
||||||
@@ -354,25 +572,189 @@ mod tests {
|
|||||||
http::StatusCode::UNAUTHORIZED,
|
http::StatusCode::UNAUTHORIZED,
|
||||||
"nope".into(),
|
"nope".into(),
|
||||||
));
|
));
|
||||||
let (client, _) = client_with(recorder);
|
let (client, _) = ollama_client(recorder);
|
||||||
|
|
||||||
assert!(!client.is_available().await);
|
assert!(!client.is_available().await);
|
||||||
}
|
}
|
||||||
|
|
||||||
|
// ---- DeepSeek/Kimi 请求参数形状 ----
|
||||||
|
|
||||||
|
#[tokio::test]
|
||||||
|
async fn deepseek_normal_request_carries_thinking_disabled() {
|
||||||
|
let recorder = RecordingHttpClient::new(DEEPSEEK_OK);
|
||||||
|
let (client, recorder) = deepseek_client(recorder, false);
|
||||||
|
|
||||||
|
let text = client.generate(Some("sys"), "hi").await.unwrap();
|
||||||
|
assert_eq!(text, "hello from deepseek");
|
||||||
|
|
||||||
|
let captured = recorder.requests();
|
||||||
|
assert_eq!(captured[0].uri, "https://api.deepseek.com/v1/chat/completions");
|
||||||
|
let body: serde_json::Value = serde_json::from_slice(&captured[0].body).unwrap();
|
||||||
|
assert_eq!(body["model"], "deepseek-v4-flash");
|
||||||
|
// 非流式请求省略 stream 字段(API 默认 false)
|
||||||
|
assert_ne!(body["stream"], true);
|
||||||
|
assert_eq!(body["temperature"], 0.5);
|
||||||
|
assert_eq!(body["thinking"]["type"], "disabled");
|
||||||
|
assert_eq!(body["max_tokens"], 123);
|
||||||
|
}
|
||||||
|
|
||||||
|
#[tokio::test]
|
||||||
|
async fn kimi_normal_request_uses_kimi_temperature() {
|
||||||
|
let recorder = RecordingHttpClient::new(DEEPSEEK_OK);
|
||||||
|
let backend = Backend::Kimi(
|
||||||
|
moonshot::Client::builder()
|
||||||
|
.api_key("sk-test")
|
||||||
|
.base_url("https://api.moonshot.cn/v1")
|
||||||
|
.http_client(recorder.clone())
|
||||||
|
.build()
|
||||||
|
.expect("build kimi client with mock backend"),
|
||||||
|
);
|
||||||
|
let client = LlmClient::new(backend, "kimi-k2.6", "kimi", test_config(), false, None);
|
||||||
|
|
||||||
|
client.generate(None, "hi").await.unwrap();
|
||||||
|
|
||||||
|
let captured = recorder.requests();
|
||||||
|
assert_eq!(captured[0].uri, "https://api.moonshot.cn/v1/chat/completions");
|
||||||
|
let body: serde_json::Value = serde_json::from_slice(&captured[0].body).unwrap();
|
||||||
|
assert_eq!(body["temperature"], 0.6);
|
||||||
|
assert_eq!(body["thinking"]["type"], "disabled");
|
||||||
|
}
|
||||||
|
|
||||||
|
// ---- 流式与 thinking 事件 ----
|
||||||
|
|
||||||
|
/// 构造 thinking 模式下的 DeepSeek 客户端(流式 mock 后端 + 记录回调的思考状态)。
|
||||||
|
fn deepseek_streaming_client(
|
||||||
|
sse: &str,
|
||||||
|
) -> (
|
||||||
|
LlmClient<MockStreamingClient>,
|
||||||
|
Arc<ThinkingStateManager>,
|
||||||
|
Arc<std::sync::atomic::AtomicUsize>,
|
||||||
|
Arc<std::sync::atomic::AtomicUsize>,
|
||||||
|
) {
|
||||||
|
let start_count = Arc::new(AtomicUsize::new(0));
|
||||||
|
let end_count = Arc::new(AtomicUsize::new(0));
|
||||||
|
let state = Arc::new(
|
||||||
|
ThinkingStateManager::new()
|
||||||
|
.on_thinking_start({
|
||||||
|
let c = start_count.clone();
|
||||||
|
move || {
|
||||||
|
c.fetch_add(1, Ordering::SeqCst);
|
||||||
|
}
|
||||||
|
})
|
||||||
|
.on_thinking_end({
|
||||||
|
let c = end_count.clone();
|
||||||
|
move || {
|
||||||
|
c.fetch_add(1, Ordering::SeqCst);
|
||||||
|
}
|
||||||
|
}),
|
||||||
|
);
|
||||||
|
|
||||||
|
let backend = Backend::DeepSeek(
|
||||||
|
deepseek::Client::builder()
|
||||||
|
.api_key("sk-test")
|
||||||
|
.base_url("https://api.deepseek.com/v1")
|
||||||
|
.http_client(MockStreamingClient { sse_bytes: sse.to_string().into() })
|
||||||
|
.build()
|
||||||
|
.expect("build deepseek client with streaming mock"),
|
||||||
|
);
|
||||||
|
let client = LlmClient::new(
|
||||||
|
backend,
|
||||||
|
"deepseek-v4-flash",
|
||||||
|
"deepseek",
|
||||||
|
test_config(),
|
||||||
|
true,
|
||||||
|
Some(state.clone()),
|
||||||
|
);
|
||||||
|
(client, state, start_count, end_count)
|
||||||
|
}
|
||||||
|
|
||||||
|
const SSE_THINK_THEN_TEXT: &str = concat!(
|
||||||
|
"data: {\"choices\":[{\"index\":0,\"delta\":{\"reasoning_content\":\"let me think\"},\"finish_reason\":null}]}\n\n",
|
||||||
|
"data: {\"choices\":[{\"index\":0,\"delta\":{\"content\":\"hello\"},\"finish_reason\":null}]}\n\n",
|
||||||
|
"data: {\"choices\":[{\"index\":0,\"delta\":{},\"finish_reason\":\"stop\"}]}\n\n",
|
||||||
|
"data: [DONE]\n\n",
|
||||||
|
);
|
||||||
|
|
||||||
|
const SSE_REASONING_ONLY: &str = concat!(
|
||||||
|
"data: {\"choices\":[{\"index\":0,\"delta\":{\"reasoning_content\":\"only thoughts\"},\"finish_reason\":null}]}\n\n",
|
||||||
|
"data: {\"choices\":[{\"index\":0,\"delta\":{},\"finish_reason\":\"stop\"}]}\n\n",
|
||||||
|
"data: [DONE]\n\n",
|
||||||
|
);
|
||||||
|
|
||||||
|
#[tokio::test]
|
||||||
|
async fn streaming_aggregates_text_and_drives_thinking_state() {
|
||||||
|
let (client, _, start_count, end_count) = deepseek_streaming_client(SSE_THINK_THEN_TEXT);
|
||||||
|
|
||||||
|
let text = client.generate(None, "hi").await.unwrap();
|
||||||
|
assert_eq!(text, "hello");
|
||||||
|
assert_eq!(start_count.load(Ordering::SeqCst), 1);
|
||||||
|
assert_eq!(end_count.load(Ordering::SeqCst), 1);
|
||||||
|
}
|
||||||
|
|
||||||
|
#[tokio::test]
|
||||||
|
async fn streaming_reasoning_only_errors_with_hint() {
|
||||||
|
let (client, _, _, _) = deepseek_streaming_client(SSE_REASONING_ONLY);
|
||||||
|
|
||||||
|
let err = client.generate(None, "hi").await.unwrap_err();
|
||||||
|
let msg = err.to_string();
|
||||||
|
assert!(
|
||||||
|
msg.contains("DeepSeek returned reasoning content but no final answer"),
|
||||||
|
"got: {msg}"
|
||||||
|
);
|
||||||
|
assert!(msg.contains("disable thinking mode"), "got: {msg}");
|
||||||
|
}
|
||||||
|
|
||||||
|
#[tokio::test]
|
||||||
|
async fn streaming_without_reasoning_skips_state_start() {
|
||||||
|
let (client, _, start_count, _) = deepseek_streaming_client(concat!(
|
||||||
|
"data: {\"choices\":[{\"index\":0,\"delta\":{\"content\":\"plain\"},\"finish_reason\":null}]}\n\n",
|
||||||
|
"data: {\"choices\":[{\"index\":0,\"delta\":{},\"finish_reason\":\"stop\"}]}\n\n",
|
||||||
|
"data: [DONE]\n\n",
|
||||||
|
));
|
||||||
|
|
||||||
|
let text = client.generate(None, "hi").await.unwrap();
|
||||||
|
assert_eq!(text, "plain");
|
||||||
|
assert_eq!(start_count.load(Ordering::SeqCst), 0);
|
||||||
|
}
|
||||||
|
|
||||||
|
// ---- 错误映射 ----
|
||||||
|
|
||||||
#[test]
|
#[test]
|
||||||
fn map_error_prefers_status_and_body() {
|
fn map_error_prefers_status_and_body() {
|
||||||
let err = map_completion_error(
|
let err = map_completion_error(
|
||||||
"Ollama",
|
"DeepSeek",
|
||||||
CompletionError::from_http_response(http::StatusCode::TOO_MANY_REQUESTS, "slow down"),
|
CompletionError::from_http_response(http::StatusCode::TOO_MANY_REQUESTS, "slow down"),
|
||||||
);
|
);
|
||||||
let msg = err.to_string();
|
let msg = err.to_string();
|
||||||
assert!(msg.contains("Ollama API error: 429"), "got: {msg}");
|
assert!(msg.contains("DeepSeek API error: 429"), "got: {msg}");
|
||||||
assert!(msg.contains("slow down"), "got: {msg}");
|
assert!(msg.contains("slow down"), "got: {msg}");
|
||||||
}
|
}
|
||||||
|
|
||||||
#[test]
|
#[test]
|
||||||
fn map_error_without_body_keeps_provider_prefix() {
|
fn map_error_without_body_keeps_provider_prefix() {
|
||||||
let err = map_completion_error("Ollama", CompletionError::ProviderError("boom".into()));
|
let err = map_completion_error("DeepSeek", CompletionError::ProviderError("boom".into()));
|
||||||
assert!(err.to_string().contains("Ollama API request failed: ProviderError: boom"));
|
assert!(err
|
||||||
|
.to_string()
|
||||||
|
.contains("DeepSeek API request failed: ProviderError: boom"));
|
||||||
|
}
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn provider_display_names_match_legacy() {
|
||||||
|
assert_eq!(provider_display_name("ollama"), "Ollama");
|
||||||
|
assert_eq!(provider_display_name("openai"), "OpenAI");
|
||||||
|
assert_eq!(provider_display_name("anthropic"), "Anthropic");
|
||||||
|
assert_eq!(provider_display_name("kimi"), "Kimi");
|
||||||
|
assert_eq!(provider_display_name("deepseek"), "DeepSeek");
|
||||||
|
assert_eq!(provider_display_name("openrouter"), "OpenRouter");
|
||||||
|
}
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn supports_thinking_whitelist_matches_legacy() {
|
||||||
|
for provider in ["deepseek", "kimi", "anthropic", "openai"] {
|
||||||
|
assert!(supports_thinking(provider), "{provider} should support thinking");
|
||||||
|
}
|
||||||
|
for provider in ["ollama", "openrouter"] {
|
||||||
|
assert!(!supports_thinking(provider), "{provider} should not support thinking");
|
||||||
|
}
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
Reference in New Issue
Block a user