AstroResearch/src/clients/llm/chat.rs
Asfmq 5f2d2d83f6 refactor: 大文件模块化拆分、自进化Skill管线、前端设计系统统一
后端:
  - runtime: 拆分 AgentConfig/events/duplicate_detector 为独立模块
  - error_recovery: 1499行单体拆为 classification(21种FailoverReason)/overflow/mod
  - executor: 提取 helpers.rs (PreparedCall/ToolExecutionResult/execute_single_tool)
  - skills: 新增 SkillCreator + SelfImprovePipeline(模式检测→自动生成SKILL.md→质量审查)
  - clients/llm: 拆分为 chat/embedding/types 三个子模块
  - services/download: 1548行拆为 mod/headers(反爬+SSRF)/strategies(多级回退)
  - services/batch/asset: 1264行拆为 mod/helpers/process

  前端:
  - 设计系统统一: sky/indigo → blueprint 色系, rounded-xl→lg, shadow-lg→sm
  - 删除 Vite 模板残留 App.css
  - GlobalDialog/PaperDetailModal/UncachedPaperModal 提取公共 BaseModal 组件
2026-06-27 09:56:36 +08:00

759 lines
29 KiB
Rust
Raw Blame History

This file contains ambiguous Unicode characters

This file contains Unicode characters that might be confused with other characters. If you think that this is intentional, you can safely ignore this warning. Use the Escape button to reveal them.

// src/clients/llm/chat.rs
//
// ChatCompleter trait + LlmClient 实现(对话补全、流式对话、图片分析)。
use async_trait::async_trait;
use futures_util::StreamExt;
use reqwest::Client;
use serde::Deserialize;
use std::sync::Arc;
use tracing::error;
use super::types::{
ChatMessage, CompletionResult, FunctionCall, StreamEvent, TokenUsage, ToolCall, ToolDefinition,
};
/// LLM 对话补全的抽象 trait支持真实客户端和测试 mock。
#[async_trait]
pub trait ChatCompleter: Send + Sync {
async fn complete(&self, system_prompt: &str, user_content: &str) -> anyhow::Result<String>;
}
#[derive(Clone, Debug)]
pub struct LlmClient {
api_key: String,
api_base: String,
model: Arc<std::sync::RwLock<String>>,
client: Client,
}
#[async_trait]
impl ChatCompleter for LlmClient {
async fn complete(&self, system_prompt: &str, user_content: &str) -> anyhow::Result<String> {
self.chat_completion(system_prompt, user_content).await
}
}
impl LlmClient {
pub fn new(api_key: String, api_base: String, model: String) -> Self {
LlmClient {
api_key,
api_base,
model: Arc::new(std::sync::RwLock::new(model)),
client: Client::new(),
}
}
pub fn model(&self) -> String {
self.model.read().unwrap_or_else(|e| e.into_inner()).clone()
}
/// 运行时切换模型(用于故障转移)。
pub fn set_model(&self, model: String) {
if let Ok(mut m) = self.model.write() {
*m = model;
tracing::info!("[LlmClient] 模型切换为: {}", m);
}
}
pub fn api_base(&self) -> &str {
&self.api_base
}
pub fn api_key(&self) -> &str {
&self.api_key
}
pub async fn chat_completion(
&self,
system_prompt: &str,
user_content: &str,
) -> anyhow::Result<String> {
let url = format!("{}/chat/completions", self.api_base);
let payload = serde_json::json!({
"model": self.model(),
"messages": [
{
"role": "system",
"content": system_prompt
},
{
"role": "user",
"content": user_content
}
],
"temperature": 0.3
});
let response = self
.client
.post(&url)
.header("Authorization", format!("Bearer {}", self.api_key))
.header("Content-Type", "application/json")
.json(&payload)
.send()
.await?;
if !response.status().is_success() {
let status = response.status();
let body = response.text().await.unwrap_or_default();
error!("LLM 接口调用失败: 状态码={}, 报错={}", status, body);
return Err(anyhow::anyhow!("大模型接口返回错误状态: {}", status));
}
#[derive(Deserialize)]
struct Message {
content: String,
}
#[derive(Deserialize)]
struct Choice {
message: Message,
}
#[derive(Deserialize)]
struct LLMResponse {
choices: Vec<Choice>,
}
let res_data: LLMResponse = response.json().await?;
if let Some(choice) = res_data.choices.first() {
Ok(choice.message.content.clone())
} else {
Err(anyhow::anyhow!("大模型返回空翻译选项集"))
}
}
/// 多轮对话补全(含原生 Tool Calling 支持),非流式
pub async fn chat(
&self,
messages: &[ChatMessage],
tools: &[ToolDefinition],
enable_thinking: bool,
) -> anyhow::Result<CompletionResult> {
let url = format!("{}/chat/completions", self.api_base);
let mut payload = serde_json::json!({
"model": self.model(),
"messages": messages,
"temperature": 0.3
});
// 前端可控的思考模式开关(仅对千问/DashScope 启用 enable_thinking 参数)(仅对千问/DashScope 启用 enable_thinking 参数)
if enable_thinking
&& (self.api_base.contains("dashscope.aliyuncs.com")
|| self.model().to_lowercase().contains("qwen"))
{
if let Some(obj) = payload.as_object_mut() {
obj.insert("enable_thinking".to_string(), serde_json::json!(true));
}
}
if !tools.is_empty() {
payload["tools"] = serde_json::to_value(tools)?;
}
let response = self
.client
.post(&url)
.header("Authorization", format!("Bearer {}", self.api_key))
.header("Content-Type", "application/json")
.json(&payload)
.send()
.await?;
if !response.status().is_success() {
let status = response.status();
let body = response.text().await.unwrap_or_default();
error!("LLM chat 接口调用失败: 状态码={}, 报错={}", status, body);
return Err(anyhow::anyhow!(
"大模型 chat 接口返回错误状态: {} - {}",
status,
body
));
}
#[derive(Deserialize)]
struct ResponseMessage {
content: Option<String>,
tool_calls: Option<Vec<ToolCall>>,
reasoning_content: Option<String>,
}
#[derive(Deserialize)]
struct Choice {
message: ResponseMessage,
}
#[derive(Deserialize)]
struct ChatResponse {
choices: Vec<Choice>,
usage: Option<TokenUsage>,
}
let res_data: ChatResponse = response.json().await?;
if let Some(choice) = res_data.choices.into_iter().next() {
Ok(CompletionResult {
content: choice.message.content,
reasoning_content: choice.message.reasoning_content,
tool_calls: choice.message.tool_calls.unwrap_or_default(),
usage: res_data.usage.unwrap_or_default(),
})
} else {
Err(anyhow::anyhow!("大模型返回空对话选项集"))
}
}
/// 流式分析图片,通过 `on_delta` 回调推送增量文本,返回完整结果。
pub async fn analyze_image_stream(
&self,
system_prompt: &str,
question: &str,
image_base64: &str,
mime_type: &str,
mut on_delta: impl FnMut(&str),
) -> anyhow::Result<String> {
let user_content = serde_json::json!([
{ "type": "text", "text": question },
{ "type": "image_url", "image_url": { "url": format!("data:{};base64,{}", mime_type, image_base64) } }
]);
let payload = serde_json::json!({
"model": self.model(),
"messages": [
{ "role": "system", "content": system_prompt },
{ "role": "user", "content": user_content }
],
"temperature": 0.3,
"stream": true,
"stream_options": { "include_usage": true }
});
let url = format!("{}/chat/completions", self.api_base);
let response = self
.client
.post(&url)
.header("Authorization", format!("Bearer {}", self.api_key))
.header("Content-Type", "application/json")
.json(&payload)
.send()
.await?;
if !response.status().is_success() {
let status = response.status();
let body = response.text().await.unwrap_or_default();
return Err(anyhow::anyhow!(
"视觉模型调用失败 (HTTP {}): {}",
status,
body
));
}
use futures_util::StreamExt;
let mut stream = response.bytes_stream();
let mut accumulated = String::new();
let mut buffer = String::new();
while let Some(chunk) = stream.next().await {
let chunk = chunk?;
buffer.push_str(&String::from_utf8_lossy(&chunk));
while let Some(pos) = buffer.find('\n') {
let line = buffer[..pos].trim().to_string();
buffer = buffer[pos + 1..].to_string();
if line.is_empty() || !line.starts_with("data: ") {
continue;
}
let data = line[6..].trim().to_string();
if data == "[DONE]" {
continue;
}
if let Ok(event) = serde_json::from_str::<serde_json::Value>(&data) {
if let Some(delta) = event["choices"][0]["delta"]["content"].as_str() {
accumulated.push_str(delta);
on_delta(delta);
}
}
}
}
if accumulated.is_empty() {
return Err(anyhow::anyhow!("视觉模型返回空结果"));
}
Ok(accumulated)
}
/// 流式多轮对话补全(含 Tool Call 碎片累积还原),返回 SSE 事件流
pub async fn chat_stream(
&self,
messages: &[ChatMessage],
tools: &[ToolDefinition],
enable_thinking: bool,
) -> anyhow::Result<tokio::sync::mpsc::UnboundedReceiver<StreamEvent>> {
let url = format!("{}/chat/completions", self.api_base);
let mut payload = serde_json::json!({
"model": self.model(),
"messages": messages,
"temperature": 0.3,
"stream": true,
"stream_options": { "include_usage": true }
});
// 前端可控的思考模式开关(仅对千问/DashScope 启用 enable_thinking 参数)
if enable_thinking
&& (self.api_base.contains("dashscope.aliyuncs.com")
|| self.model().to_lowercase().contains("qwen"))
{
if let Some(obj) = payload.as_object_mut() {
obj.insert("enable_thinking".to_string(), serde_json::json!(true));
}
}
if !tools.is_empty() {
payload["tools"] = serde_json::to_value(tools)?;
}
let response = self
.client
.post(&url)
.header("Authorization", format!("Bearer {}", self.api_key))
.header("Content-Type", "application/json")
.json(&payload)
.send()
.await?;
if !response.status().is_success() {
let status = response.status();
let status_code = status.as_u16();
let retry_after = response
.headers()
.get("retry-after")
.and_then(|v| v.to_str().ok())
.and_then(|v| v.parse::<u64>().ok());
let body = response.text().await.unwrap_or_default();
error!("LLM stream 接口调用失败: 状态码={}, 报错={}", status, body);
return Err(anyhow::anyhow!(
"HTTP {}: {} | retry_after={:?}",
status_code,
body,
retry_after
));
}
let (tx, rx) = tokio::sync::mpsc::unbounded_channel();
let mut byte_stream = response.bytes_stream();
tokio::spawn(async move {
// 流式 Tool Call 碎片累积器
let mut tool_call_accumulators: std::collections::HashMap<
usize,
(String, String, String),
> = std::collections::HashMap::new();
let mut buffer = String::new();
while let Some(chunk_result) = byte_stream.next().await {
let chunk = match chunk_result {
Ok(c) => c,
Err(e) => {
let _ = tx.send(StreamEvent::Error(format!("流式读取错误: {}", e)));
break;
}
};
buffer.push_str(&String::from_utf8_lossy(&chunk));
// 按行解析 SSE 事件
while let Some(line_end) = buffer.find('\n') {
let line = buffer[..line_end].trim().to_string();
buffer = buffer[line_end + 1..].to_string();
if line.is_empty() || !line.starts_with("data: ") {
continue;
}
let data = &line[6..];
if data == "[DONE]" {
// 流结束前,将累积的 tool calls 还原并发送
if !tool_call_accumulators.is_empty() {
let mut indices: Vec<usize> =
tool_call_accumulators.keys().cloned().collect();
indices.sort();
let tool_calls: Vec<ToolCall> = indices
.into_iter()
.map(|idx| {
let (id, name, args) =
tool_call_accumulators.remove(&idx).unwrap();
ToolCall {
id: if id.is_empty() {
format!(
"call_{}",
&uuid::Uuid::new_v4().to_string()[..8]
)
} else {
id
},
call_type: "function".to_string(),
function: FunctionCall {
name,
arguments: args,
},
}
})
.collect();
let _ = tx.send(StreamEvent::ToolCallsComplete(tool_calls));
}
let _ = tx.send(StreamEvent::Done);
return;
}
// 解析 JSON delta
let parsed: serde_json::Value = match serde_json::from_str(data) {
Ok(v) => v,
Err(_) => continue,
};
// 提取 usage在最后一条 chunk 中 stream_options 返回)
if let Some(usage_val) = parsed.get("usage") {
if !usage_val.is_null() {
if let Ok(usage) =
serde_json::from_value::<TokenUsage>(usage_val.clone())
{
let _ = tx.send(StreamEvent::Usage(usage));
}
}
}
if let Some(choices) = parsed.get("choices").and_then(|c| c.as_array()) {
if let Some(choice) = choices.first() {
let delta = &choice["delta"];
// 文本增量
if let Some(content) = delta.get("content").and_then(|c| c.as_str()) {
if !content.is_empty() {
let _ = tx.send(StreamEvent::TextDelta(content.to_string()));
}
}
// 推理内容增量
if let Some(reasoning) =
delta.get("reasoning_content").and_then(|c| c.as_str())
{
if !reasoning.is_empty() {
let _ =
tx.send(StreamEvent::ReasoningDelta(reasoning.to_string()));
}
}
// 工具调用增量
if let Some(tool_calls) =
delta.get("tool_calls").and_then(|t| t.as_array())
{
for tc in tool_calls {
let index =
tc.get("index").and_then(|i| i.as_u64()).unwrap_or(0)
as usize;
let id = tc
.get("id")
.and_then(|i| i.as_str())
.map(|s| s.to_string());
let fn_name = tc
.get("function")
.and_then(|f| f.get("name"))
.and_then(|n| n.as_str())
.map(|s| s.to_string());
let fn_args = tc
.get("function")
.and_then(|f| f.get("arguments"))
.and_then(|a| a.as_str())
.unwrap_or("");
let entry =
tool_call_accumulators.entry(index).or_insert_with(|| {
(String::new(), String::new(), String::new())
});
if let Some(ref id_str) = id {
entry.0 = id_str.clone();
}
if let Some(ref name_str) = fn_name {
entry.1 = name_str.clone();
}
entry.2.push_str(fn_args);
let _ = tx.send(StreamEvent::ToolCallDelta {
index,
id,
name: fn_name,
arguments_delta: fn_args.to_string(),
});
}
}
// 检查是否 finish_reason == "tool_calls"
if let Some(finish_reason) =
choice.get("finish_reason").and_then(|f| f.as_str())
{
if finish_reason == "tool_calls"
&& !tool_call_accumulators.is_empty()
{
let mut indices: Vec<usize> =
tool_call_accumulators.keys().cloned().collect();
indices.sort();
let tool_calls: Vec<ToolCall> = indices
.into_iter()
.map(|idx| {
let (id, name, args) =
tool_call_accumulators.remove(&idx).unwrap();
ToolCall {
id,
call_type: "function".to_string(),
function: FunctionCall {
name,
arguments: args,
},
}
})
.collect();
let _ = tx.send(StreamEvent::ToolCallsComplete(tool_calls));
}
}
}
}
}
}
// 异常结束(连接断开等)
if !tool_call_accumulators.is_empty() {
let mut indices: Vec<usize> = tool_call_accumulators.keys().cloned().collect();
indices.sort();
let tool_calls: Vec<ToolCall> = indices
.into_iter()
.map(|idx| {
let (id, name, args) = tool_call_accumulators.remove(&idx).unwrap();
ToolCall {
id: if id.is_empty() {
format!("call_{}", &uuid::Uuid::new_v4().to_string()[..8])
} else {
id
},
call_type: "function".to_string(),
function: FunctionCall {
name,
arguments: args,
},
}
})
.collect();
let _ = tx.send(StreamEvent::ToolCallsComplete(tool_calls));
}
let _ = tx.send(StreamEvent::Done);
});
Ok(rx)
}
}
#[cfg(test)]
mod tests {
use super::*;
use crate::clients::llm::types::{
ChatMessage, FunctionCall, MessageRole, ToolCall, ToolDefinition,
};
#[test]
fn test_llm_client_initialization() {
let client = LlmClient::new("key".to_string(), "base".to_string(), "model".to_string());
assert_eq!(client.api_key(), "key");
assert_eq!(client.api_base(), "base");
assert_eq!(client.model(), "model");
}
#[test]
fn test_chat_message_construction() {
// 测试系统消息构造
let sys = ChatMessage::system("你是一个助手");
assert_eq!(sys.role, MessageRole::System);
assert_eq!(sys.text(), Some("你是一个助手"));
assert!(sys.tool_calls.is_none());
// 测试用户消息构造
let user = ChatMessage::user("你好");
assert_eq!(user.role, MessageRole::User);
assert_eq!(user.text(), Some("你好"));
// 测试助手消息构造
let assistant = ChatMessage::assistant("你好!有什么可以帮助你的吗?");
assert_eq!(assistant.role, MessageRole::Assistant);
assert_eq!(assistant.text(), Some("你好!有什么可以帮助你的吗?"));
// 测试工具结果消息构造
let tool_result = ChatMessage::tool_result("call_123", r#"{"result": 42}"#);
assert_eq!(tool_result.role, MessageRole::Tool);
assert_eq!(tool_result.tool_call_id.as_deref(), Some("call_123"));
assert_eq!(tool_result.text(), Some(r#"{"result": 42}"#));
// 测试带工具调用的助手消息
let tool_calls = vec![ToolCall {
id: "call_abc".to_string(),
call_type: "function".to_string(),
function: FunctionCall {
name: "get_weather".to_string(),
arguments: r#"{"city": ""}"#.to_string(),
},
}];
let assistant_tc = ChatMessage::assistant_with_tool_calls(None, tool_calls.clone());
assert_eq!(assistant_tc.role, MessageRole::Assistant);
assert!(assistant_tc.text().is_none());
assert_eq!(assistant_tc.tool_calls.as_ref().unwrap().len(), 1);
assert_eq!(
assistant_tc.tool_calls.as_ref().unwrap()[0].function.name,
"get_weather"
);
}
#[test]
fn test_tool_definition_construction() {
let tool = ToolDefinition::new(
"search",
"搜索互联网内容",
serde_json::json!({
"type": "object",
"properties": {
"query": {
"type": "string",
"description": "搜索关键词"
}
},
"required": ["query"]
}),
);
assert_eq!(tool.tool_type, "function");
assert_eq!(tool.function.name, "search");
assert_eq!(tool.function.description, "搜索互联网内容");
}
#[test]
fn test_chat_message_serialization() {
// 测试序列化时 skip_serializing_if 生效
let msg = ChatMessage::user("测试消息");
let json = serde_json::to_value(&msg).unwrap();
assert_eq!(json.get("role").unwrap(), "user");
assert_eq!(json.get("content").unwrap(), "测试消息");
// None 字段不应出现在 JSON 中
assert!(json.get("tool_call_id").is_none());
assert!(json.get("tool_calls").is_none());
assert!(json.get("name").is_none());
assert!(json.get("reasoning_content").is_none());
}
#[test]
fn test_chat_message_reasoning_serialization() {
// 测试带有 reasoning_content 的消息序列化和反序列化
let msg = ChatMessage::assistant_with_reasoning(
Some("回答内容".to_string()),
Some("这是思考过程".to_string()),
None,
);
let json = serde_json::to_value(&msg).unwrap();
assert_eq!(json.get("role").unwrap(), "assistant");
assert_eq!(json.get("content").unwrap(), "回答内容");
assert_eq!(json.get("reasoning_content").unwrap(), "这是思考过程");
assert!(json.get("tool_calls").is_none());
// 测试反序列化
let deserialized: ChatMessage = serde_json::from_value(json).unwrap();
assert_eq!(deserialized.role, MessageRole::Assistant);
assert_eq!(deserialized.text(), Some("回答内容"));
assert_eq!(
deserialized.reasoning_content.as_deref(),
Some("这是思考过程")
);
}
#[test]
fn test_message_role_serialization() {
// 测试角色枚举的 rename_all = "lowercase" 序列化
let role_json = serde_json::to_string(&MessageRole::System).unwrap();
assert_eq!(role_json, r#""system""#);
let role_json = serde_json::to_string(&MessageRole::Assistant).unwrap();
assert_eq!(role_json, r#""assistant""#);
let role_json = serde_json::to_string(&MessageRole::Tool).unwrap();
assert_eq!(role_json, r#""tool""#);
// 测试反序列化
let role: MessageRole = serde_json::from_str(r#""user""#).unwrap();
assert_eq!(role, MessageRole::User);
}
/// 真实网络集成测试 —— 需要配置 LLM_API_KEY 和 EMBEDDING_API_KEY
#[tokio::test]
#[ignore]
async fn test_live_llm_and_embedding() -> anyhow::Result<()> {
let config = crate::Config::from_env();
println!("================= 开始大模型与向量模型真实网络集成测试 =================");
// 1. 测试 LlmClient
if config.llm_api_key.is_empty() {
println!("警告: 未在环境配置中检测到 LLM_API_KEY跳过 LlmClient 集成测试。");
} else {
println!(
"测试大模型: {} (API Base: {})",
config.llm_model, config.llm_api_base
);
let llm = LlmClient::new(
config.llm_api_key.clone(),
config.llm_api_base.clone(),
config.llm_model.clone(),
);
match llm
.chat_completion("You are a helpful assistant.", "Say Hello!")
.await
{
Ok(reply) => {
println!("LlmClient 响应成功: {}", reply.trim());
assert!(!reply.trim().is_empty(), "错误: 大模型返回了空响应");
}
Err(e) => panic!("LlmClient 接口调用失败: {}", e),
}
}
// 2. 测试 EmbeddingClient
if config.embedding_api_key.is_empty() {
println!(
"警告: 未在环境配置中检测到 EMBEDDING_API_KEY跳过 EmbeddingClient 集成测试。"
);
} else {
println!(
"测试向量模型: {} (API Base: {})",
config.embedding_model, config.embedding_api_base
);
let embedding_client = crate::clients::llm::EmbeddingClient::new(
config.embedding_api_key.clone(),
config.embedding_api_base.clone(),
config.embedding_model.clone(),
);
let test_text = "active galactic nucleus";
match embedding_client.create_embedding(test_text).await {
Ok(vector) => {
println!("EmbeddingClient 响应成功!向量维度: {}", vector.len());
assert!(!vector.is_empty(), "错误: 向量数据为空");
let preview_len = std::cmp::min(5, vector.len());
println!(
"{} 个向量数值样例: {:?}",
preview_len,
&vector[..preview_len]
);
}
Err(e) => panic!("EmbeddingClient 接口调用失败: {}", e),
}
}
println!("================= 大模型与向量模型真实网络集成测试完成 =================");
Ok(())
}
}