feat(llm): 引入 rig-core 并建立新 LLM 门面（Ollama 首个打通）
- 依赖改用 rig-core 0.41（仅 reqwest+rustls，不引入 agent/derive） - 新门面：provider 枚举 + 统一配置 + 单一生成入口 + 凭据校验 + 错误映射 - Ollama 走 rig /api/chat，参数映射与错误文案与旧实现一致 - 用 rig 请求录制 mock 断言请求形状；8 个新单元测试
This commit is contained in:
@@ -11,6 +11,7 @@ pub mod openai;
|
||||
pub mod openrouter;
|
||||
pub mod parsing;
|
||||
pub mod prompts;
|
||||
pub mod rig;
|
||||
pub mod thinking;
|
||||
|
||||
pub use anthropic::AnthropicClient;
|
||||
|
||||
378
src/llm/rig/mod.rs
Normal file
378
src/llm/rig/mod.rs
Normal file
@@ -0,0 +1,378 @@
|
||||
//! 基于 rig-core 的新 LLM 门面。
|
||||
//!
|
||||
//! 迁移扩张期：本模块与旧的手写实现并存，按 ticket 逐个接入 provider。
|
||||
//! 全部 provider 接入并切换后，旧实现将被删除。
|
||||
|
||||
use crate::config::manager::ConfigManager;
|
||||
use anyhow::{Context, Result, bail};
|
||||
use rig_core::{
|
||||
client::{CompletionClient, Nothing, VerifyClient},
|
||||
completion::{AssistantContent, CompletionError, CompletionModel},
|
||||
http_client::{HttpClientExt, ReqwestClient},
|
||||
providers::ollama,
|
||||
wasm_compat::{WasmCompatSend, WasmCompatSync},
|
||||
};
|
||||
use std::time::Duration;
|
||||
|
||||
/// LLM 客户端运行时配置（与用户配置解耦的参数）。
|
||||
#[derive(Debug, Clone)]
|
||||
pub struct LlmClientConfig {
|
||||
pub max_tokens: u64,
|
||||
pub temperature: f64,
|
||||
pub timeout: Duration,
|
||||
}
|
||||
|
||||
impl Default for LlmClientConfig {
|
||||
fn default() -> Self {
|
||||
Self {
|
||||
max_tokens: 500,
|
||||
temperature: 0.7,
|
||||
timeout: Duration::from_secs(30),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// Provider 后端枚举。
|
||||
///
|
||||
/// 泛型 `H` 是 HTTP 后端（生产环境为 reqwest 客户端），作为测试注入点：
|
||||
/// 单测注入 rig 提供的 mock 后端以断言请求形状与响应处理。
|
||||
#[derive(Clone)]
|
||||
pub enum Backend<H = ReqwestClient> {
|
||||
Ollama(ollama::Client<H>),
|
||||
}
|
||||
|
||||
/// 基于 rig 的 LLM 客户端门面。
|
||||
#[derive(Clone)]
|
||||
pub struct LlmClient<H = ReqwestClient> {
|
||||
backend: Backend<H>,
|
||||
model: String,
|
||||
provider: String,
|
||||
config: LlmClientConfig,
|
||||
thinking_enabled: bool,
|
||||
}
|
||||
|
||||
impl LlmClient {
|
||||
/// 从配置管理器构建（默认 reqwest 后端）。
|
||||
pub async fn from_config(manager: &ConfigManager) -> Result<Self> {
|
||||
Self::from_config_with_think(manager, manager.config().llm.thinking_enabled).await
|
||||
}
|
||||
|
||||
/// 从配置构建，thinking 由参数显式指定。
|
||||
pub async fn from_config_with_think(
|
||||
manager: &ConfigManager,
|
||||
thinking_enabled: bool,
|
||||
) -> Result<Self> {
|
||||
let config = manager.config();
|
||||
let cfg = LlmClientConfig {
|
||||
max_tokens: config.llm.max_tokens as u64,
|
||||
temperature: config.llm.temperature as f64,
|
||||
timeout: Duration::from_secs(config.llm.timeout),
|
||||
};
|
||||
let provider = manager.llm_provider().to_string();
|
||||
let model = manager.llm_model().to_string();
|
||||
let base_url = manager.llm_base_url();
|
||||
|
||||
let http = ReqwestClient::builder()
|
||||
.timeout(cfg.timeout)
|
||||
.build()
|
||||
.context("Failed to create HTTP client")?;
|
||||
|
||||
let backend = match provider.as_str() {
|
||||
"ollama" => Backend::Ollama(build_ollama_client(&base_url, http)?),
|
||||
// 其余 provider 在后续 ticket 接入，扩张期内新门面尚未被应用使用。
|
||||
"openai" | "anthropic" | "kimi" | "deepseek" | "openrouter" => {
|
||||
bail!("Provider '{}' is not available in the new LLM backend yet", provider)
|
||||
}
|
||||
_ => bail!("Unknown LLM provider: {}", provider),
|
||||
};
|
||||
|
||||
Ok(Self {
|
||||
backend,
|
||||
model,
|
||||
provider,
|
||||
config: cfg,
|
||||
thinking_enabled,
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
impl<H> LlmClient<H> {
|
||||
/// 内部与测试构造入口。
|
||||
pub(crate) fn new(
|
||||
backend: Backend<H>,
|
||||
model: impl Into<String>,
|
||||
provider: impl Into<String>,
|
||||
config: LlmClientConfig,
|
||||
thinking_enabled: bool,
|
||||
) -> Self {
|
||||
Self {
|
||||
backend,
|
||||
model: model.into(),
|
||||
provider: provider.into(),
|
||||
config,
|
||||
thinking_enabled,
|
||||
}
|
||||
}
|
||||
|
||||
/// 当前是否处于 thinking 模式。
|
||||
pub(crate) fn thinking_enabled(&self) -> bool {
|
||||
self.thinking_enabled
|
||||
}
|
||||
}
|
||||
|
||||
impl<H> LlmClient<H>
|
||||
where
|
||||
H: HttpClientExt
|
||||
+ Clone
|
||||
+ Default
|
||||
+ std::fmt::Debug
|
||||
+ Send
|
||||
+ Sync
|
||||
+ WasmCompatSend
|
||||
+ WasmCompatSync
|
||||
+ 'static,
|
||||
{
|
||||
/// 生成文本（thinking=false 的非流式路径）。
|
||||
pub async fn generate(&self, system: Option<&str>, user: &str) -> Result<String> {
|
||||
match &self.backend {
|
||||
Backend::Ollama(client) => {
|
||||
let model = client.completion_model(self.model.as_str());
|
||||
complete(&model, &self.provider, system, user, &self.config).await
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// 检查 provider 是否可用（走 rig 凭据校验接口）。
|
||||
pub async fn is_available(&self) -> bool {
|
||||
match &self.backend {
|
||||
Backend::Ollama(client) => client.verify().await.is_ok(),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// 构建 Ollama 客户端（无 API key；支持自定义 base_url 与注入的 HTTP 后端）。
|
||||
fn build_ollama_client(base_url: &str, http: ReqwestClient) -> Result<ollama::Client> {
|
||||
ollama::Client::builder()
|
||||
.api_key(Nothing)
|
||||
.base_url(base_url)
|
||||
.http_client(http)
|
||||
.build()
|
||||
.map_err(|e| anyhow::anyhow!("Failed to build Ollama client: {}", e))
|
||||
}
|
||||
|
||||
/// 非流式补全：系统提示、温度、max_tokens 映射到统一请求，聚合正文文本。
|
||||
async fn complete<M>(
|
||||
model: &M,
|
||||
provider: &str,
|
||||
system: Option<&str>,
|
||||
user: &str,
|
||||
config: &LlmClientConfig,
|
||||
) -> Result<String>
|
||||
where
|
||||
M: CompletionModel,
|
||||
{
|
||||
let mut builder = model.completion_request(user);
|
||||
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
|
||||
.completion(request)
|
||||
.await
|
||||
.map_err(|e| map_completion_error(provider, e))?;
|
||||
|
||||
let text = collect_text(response.choice.iter());
|
||||
let text = text.trim().to_string();
|
||||
if text.is_empty() {
|
||||
bail!("No response from {}", provider_display_name(provider));
|
||||
}
|
||||
Ok(text)
|
||||
}
|
||||
|
||||
/// 从响应内容中聚合正文文本（忽略 reasoning/tool call 等块）。
|
||||
fn collect_text<'a>(items: impl Iterator<Item = &'a AssistantContent>) -> String {
|
||||
items
|
||||
.filter_map(|item| match item {
|
||||
AssistantContent::Text(text) => Some(text.text.as_str()),
|
||||
_ => None,
|
||||
})
|
||||
.collect()
|
||||
}
|
||||
|
||||
/// provider 用户可见展示名（与旧实现文案一致，首字母大写）。
|
||||
pub(crate) fn provider_display_name(provider: &str) -> &str {
|
||||
match provider {
|
||||
"ollama" => "Ollama",
|
||||
"openai" => "OpenAI",
|
||||
"anthropic" => "Anthropic",
|
||||
"kimi" => "Kimi",
|
||||
"deepseek" => "DeepSeek",
|
||||
"openrouter" => "OpenRouter",
|
||||
other => other,
|
||||
}
|
||||
}
|
||||
|
||||
/// 把 rig 的类型化补全错误映射为与旧实现一致的用户可见文案。
|
||||
pub(crate) fn map_completion_error(provider: &str, e: CompletionError) -> anyhow::Error {
|
||||
// rig 在 provider 响应没有任何正文内容时的统一错误 → 映射为旧文案风格。
|
||||
if matches!(&e, CompletionError::ResponseError(msg) if msg == "No content provided") {
|
||||
return anyhow::anyhow!("No response from {}", provider_display_name(provider));
|
||||
}
|
||||
let name = provider_display_name(provider);
|
||||
match (e.provider_response_status(), e.provider_response_body()) {
|
||||
(Some(status), Some(body)) => anyhow::anyhow!("{} API error: {} - {}", name, status, body),
|
||||
(Some(status), None) => anyhow::anyhow!("{} API error: {}", name, status),
|
||||
(None, Some(body)) => anyhow::anyhow!("{} API error: {}", name, body),
|
||||
(None, None) => anyhow::anyhow!("{} API request failed: {}", name, e),
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
use rig_core::test_utils::{MockHttpResponse, RecordingHttpClient};
|
||||
|
||||
const OLLAMA_OK: &str = r#"{
|
||||
"model": "llama3.2",
|
||||
"created_at": "2024-01-01T00:00:00Z",
|
||||
"message": {"role": "assistant", "content": "hello from ollama"},
|
||||
"done": true
|
||||
}"#;
|
||||
|
||||
/// 用注入的 HTTP 后端构建 Ollama 客户端（与生产构造路径一致）。
|
||||
fn build_backend(recorder: RecordingHttpClient) -> Backend<RecordingHttpClient> {
|
||||
Backend::Ollama(
|
||||
ollama::Client::builder()
|
||||
.api_key(Nothing)
|
||||
.base_url("http://localhost:11434")
|
||||
.http_client(recorder)
|
||||
.build()
|
||||
.expect("build ollama client with mock backend"),
|
||||
)
|
||||
}
|
||||
|
||||
fn client_with(
|
||||
recorder: RecordingHttpClient,
|
||||
) -> (LlmClient<RecordingHttpClient>, RecordingHttpClient) {
|
||||
let backend = build_backend(recorder.clone());
|
||||
let client = LlmClient::new(
|
||||
backend,
|
||||
"llama3.2",
|
||||
"ollama",
|
||||
LlmClientConfig {
|
||||
max_tokens: 123,
|
||||
temperature: 0.5,
|
||||
timeout: Duration::from_secs(30),
|
||||
},
|
||||
false,
|
||||
);
|
||||
(client, recorder)
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn generate_maps_request_params_and_returns_text() {
|
||||
let recorder = RecordingHttpClient::new(OLLAMA_OK);
|
||||
let (client, recorder) = client_with(recorder);
|
||||
|
||||
let text = client.generate(Some("be helpful"), "say hi").await.unwrap();
|
||||
assert_eq!(text, "hello from ollama");
|
||||
|
||||
let captured = recorder.requests();
|
||||
assert_eq!(captured.len(), 1);
|
||||
assert_eq!(captured[0].uri, "http://localhost:11434/api/chat");
|
||||
|
||||
let body: serde_json::Value = serde_json::from_slice(&captured[0].body).unwrap();
|
||||
assert_eq!(body["model"], "llama3.2");
|
||||
assert_eq!(body["stream"], false);
|
||||
assert_eq!(body["messages"][0]["role"], "system");
|
||||
assert_eq!(body["messages"][0]["content"], "be helpful");
|
||||
assert_eq!(body["messages"][1]["role"], "user");
|
||||
assert_eq!(body["messages"][1]["content"], "say hi");
|
||||
assert_eq!(body["options"]["temperature"], 0.5);
|
||||
assert_eq!(body["options"]["num_predict"], 123);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn generate_without_system_omits_preamble() {
|
||||
let recorder = RecordingHttpClient::new(OLLAMA_OK);
|
||||
let (client, recorder) = client_with(recorder);
|
||||
|
||||
client.generate(None, "say hi").await.unwrap();
|
||||
|
||||
let captured = recorder.requests();
|
||||
let body: serde_json::Value = serde_json::from_slice(&captured[0].body).unwrap();
|
||||
assert_eq!(body["messages"][0]["role"], "user");
|
||||
assert_eq!(body["messages"].as_array().unwrap().len(), 1);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn http_error_maps_to_provider_error_message() {
|
||||
let recorder = RecordingHttpClient::new("ignored");
|
||||
recorder.set_response(MockHttpResponse::ErrorResponse(
|
||||
http::StatusCode::BAD_GATEWAY,
|
||||
"upstream exploded".into(),
|
||||
));
|
||||
let (client, _) = client_with(recorder);
|
||||
|
||||
let err = client.generate(None, "hi").await.unwrap_err();
|
||||
let msg = err.to_string();
|
||||
assert!(msg.contains("Ollama API error:"), "got: {msg}");
|
||||
assert!(msg.contains("502"), "got: {msg}");
|
||||
assert!(msg.contains("upstream exploded"), "got: {msg}");
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn empty_response_errors() {
|
||||
let recorder = RecordingHttpClient::new(
|
||||
r#"{"model":"llama3.2","created_at":"x","message":{"role":"assistant","content":""},"done":true}"#,
|
||||
);
|
||||
let (client, _) = client_with(recorder);
|
||||
|
||||
let err = client.generate(None, "hi").await.unwrap_err();
|
||||
assert_eq!(err.to_string(), "No response from Ollama");
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn is_available_uses_verify_endpoint() {
|
||||
let recorder = RecordingHttpClient::new("");
|
||||
let (client, recorder) = client_with(recorder);
|
||||
|
||||
assert!(client.is_available().await);
|
||||
let captured = recorder.requests();
|
||||
assert_eq!(captured.len(), 1);
|
||||
assert_eq!(captured[0].uri, "http://localhost:11434/api/tags");
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn is_available_false_on_authentication_error() {
|
||||
let recorder = RecordingHttpClient::new("");
|
||||
recorder.set_response(MockHttpResponse::ErrorResponse(
|
||||
http::StatusCode::UNAUTHORIZED,
|
||||
"nope".into(),
|
||||
));
|
||||
let (client, _) = client_with(recorder);
|
||||
|
||||
assert!(!client.is_available().await);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn map_error_prefers_status_and_body() {
|
||||
let err = map_completion_error(
|
||||
"Ollama",
|
||||
CompletionError::from_http_response(http::StatusCode::TOO_MANY_REQUESTS, "slow down"),
|
||||
);
|
||||
let msg = err.to_string();
|
||||
assert!(msg.contains("Ollama API error: 429"), "got: {msg}");
|
||||
assert!(msg.contains("slow down"), "got: {msg}");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn map_error_without_body_keeps_provider_prefix() {
|
||||
let err = map_completion_error("Ollama", CompletionError::ProviderError("boom".into()));
|
||||
assert!(err.to_string().contains("Ollama API request failed: ProviderError: boom"));
|
||||
}
|
||||
}
|
||||
Reference in New Issue
Block a user