feat: iPet → iAs 完整迁移,Rust 版微信 AI 助手

Phase 1  CLI 框架 + 配置系统
  - clap 子命令: login / listen / send / whoami / usage
  - config.json + env var 替换
  - tracing 日志系统
  - state 持久化(auth/runtime 文件存 + PostgreSQL)

Phase 2  微信通道
  - wechat::client — 完整 iLink Bot HTTP API 实现
  - 扫码登录(终端二维码 + 轮询状态)
  - 长轮询 getupdates / 消息收发 / 监听注册

Phase 3  AI 对话(纯文本 + function calling)
  - LlmProvider trait: DeepSeek + LM Studio 实现
  - SSE 流式解析(text / reasoning / tool_calls delta / usage)
  - Conversation: 消息历史 + chat / chat_with_tools 工具循环

Phase 4  PostgreSQL 集成
  - app_state(认证 KV 存储)
  - chat_records(消息收发记录)
  - llm_usage(Token 用量统计缓存命中率)
  - user_memories(长期记忆持久化)
  - pending_approvals(审批确认码)
  - scheduled_tasks(定时任务表)

Phase 5  一切皆 Skill(工具系统)
  - SkillRegistry: 系统 + 用户 skills 双目录合并
  - SKILL.md 解析器 + 子进程执行器(stdin JSON → stdout)
  - 9 个系统 Skills: datetime / weather / search / email /
    shell / schedule / memos / read_memories / read_summaries
  - ApprovalManager: High 风险技能 → 确认码审批(透明模式)
  - High 风险技能:确认码审批(透明模式)

Phase 6  定时任务调度器

上下文管理
  - ChatSession: checkpoint + token budget (28K) + summaries
  - Token 估算器(中英文自适应)
  - 12h 空闲 → trigger_idle_summary(不入会话)
  - Budget 溢出 → trigger_overflow_summary(入会话 + drain 旧消息)
  - Summarizer: LLM 生成自然语言摘要(fallback 简单截断)
  - 长期记忆 / 摘要 通过 read_memories / read_summaries 工具按需读取

工具调用日志 + Token 统计
  - INFO: 工具名 + 参数 + 结果摘要
  - DEBUG: 子进程 exit/stdout/stderr
  - ias usage --since --until --model 查看用量和缓存命中率
This commit is contained in:
2026-06-01 17:21:43 +08:00
parent 3a2d2769b5
commit b9de3665d9
53 changed files with 9874 additions and 2 deletions
+242
View File
@@ -0,0 +1,242 @@
use crate::context::builder;
use crate::context::types::ChatSession;
use crate::llm::provider::{create_provider, BoxedProvider, StreamReceiver};
use crate::llm::types::{ConversationConfig, Message, StreamChunk, ToolCall, Usage};
use std::future::Future;
use std::pin::Pin;
use std::sync::Arc;
use tokio::sync::Mutex;
pub type ToolExecutor =
Arc<dyn Fn(&str, &str) -> Pin<Box<dyn Future<Output = Result<String, String>> + Send>> + Send + Sync>;
/// 摘要生成器:接收要摘要的消息文本 → 返回 LLM 生成的摘要
pub type Summarizer =
Arc<dyn Fn(String) -> Pin<Box<dyn Future<Output = Result<String, String>> + Send>> + Send + Sync>;
pub struct Conversation {
config: ConversationConfig,
provider: BoxedProvider,
session: Arc<Mutex<ChatSession>>,
tool_executor: Option<ToolExecutor>,
}
impl Conversation {
pub fn new(config: ConversationConfig) -> Result<Self, String> {
let provider = create_provider(&config)?;
Ok(Self {
config,
provider,
session: Arc::new(Mutex::new(ChatSession::new())),
tool_executor: None,
})
}
pub fn with_provider(config: ConversationConfig, provider: BoxedProvider) -> Self {
Self {
config,
provider,
session: Arc::new(Mutex::new(ChatSession::new())),
tool_executor: None,
}
}
pub fn set_tool_executor(&mut self, executor: ToolExecutor) { self.tool_executor = Some(executor); }
pub fn session(&self) -> Arc<Mutex<ChatSession>> { Arc::clone(&self.session) }
pub fn model(&self) -> &str { &self.config.model }
pub fn provider_name(&self) -> &str { self.provider.name() }
pub fn spec(&self) -> &ConversationConfig { &self.config }
pub fn summarizer(&self) -> Summarizer { self.create_summarizer() }
/// 创建 LLM 摘要生成器(用当前 provider 做非流式调用)
pub fn create_summarizer(&self) -> Summarizer {
let provider = self.provider.clone();
Arc::new(move |prompt: String| {
let provider = provider.clone();
Box::pin(async move {
let msgs = vec![Message::system("你是一个对话摘要助手,用中文输出简洁摘要。"), Message::user(prompt)];
let mut rx = provider.chat_stream(&ConversationConfig::default(), &msgs).await
.map_err(|e| format!("摘要请求失败: {}", e))?;
let mut text = String::new();
while let Some(chunk) = rx.recv().await {
match chunk {
StreamChunk::Text(t) => text.push_str(&t),
StreamChunk::Done { text: full, .. } => {
if !full.is_empty() { text = full; }
break;
}
StreamChunk::Error(e) => return Err(e),
_ => {}
}
}
Ok(text)
})
})
}
/// 单次对话(无工具)
pub async fn chat(&self, user_message: String) -> Result<ChatHandle, String> {
// 空闲超时检查
{
let s = self.session.lock().await;
if s.is_idle_timeout() {
drop(s);
builder::trigger_idle_summary(&self.session, Some(&self.summarizer())).await;
}
}
self.session.lock().await.add_user(user_message.clone());
let messages = builder::build_context(&self.session, &self.config.system_prompt).await;
let rx = self.provider.chat_stream(&self.config, &messages).await?;
Ok(ChatHandle {
rx: Some(rx), full_text: String::new(), tool_calls: None, usage: None,
session: Arc::clone(&self.session),
})
}
/// 对话 + 工具循环(带上下文管理)
pub async fn chat_with_tools(
&self,
user_message: String,
) -> Result<(String, bool, Option<Usage>), String> {
// ── 空闲超时检查 ──
{
let s = self.session.lock().await;
if s.is_idle_timeout() {
drop(s);
builder::trigger_idle_summary(&self.session, Some(&self.summarizer())).await;
tracing::info!("⏰ 检测到 12h 空闲,已生成摘要");
}
}
self.session.lock().await.add_user(user_message.clone());
let mut used_tools = false;
let mut last_usage = None;
let mut turn_count = 0u32;
loop {
turn_count += 1;
if turn_count > 5 {
return Err("工具调用次数过多(最多5轮)".to_string());
}
let messages = builder::build_context(&self.session, &self.config.system_prompt).await;
let rx = self.provider.chat_stream(&self.config, &messages).await?;
let mut full_text = String::new();
let mut tool_calls: Option<Vec<ToolCall>> = None;
let mut rx = rx;
while let Some(chunk) = rx.recv().await {
match chunk {
StreamChunk::Text(t) => full_text.push_str(&t),
StreamChunk::Done { text, tool_calls: tc, usage, .. } => {
if !text.is_empty() { full_text = text; }
tool_calls = tc;
last_usage = usage;
break;
}
StreamChunk::Error(e) => return Err(e),
_ => {}
}
}
if let Some(calls) = tool_calls {
if !calls.is_empty() {
used_tools = true;
tracing::info!(
"🔧 LLM 请求 {} 个工具: {}",
calls.len(),
calls.iter().map(|c| format!("{}({:.80})", c.function.name, c.function.arguments)).collect::<Vec<_>>().join(", ")
);
self.session.lock().await.add_assistant_tool_calls(calls.clone());
for tc in &calls {
let result = match &self.tool_executor {
Some(exec) => match exec(&tc.function.name, &tc.function.arguments).await {
Ok(r) => r,
Err(e) => format!("执行失败: {}", e),
},
None => "工具执行器未配置".to_string(),
};
tracing::info!("📦 {} → {:.150}", tc.function.name, result);
self.session.lock().await.add_tool_result(&tc.id, &tc.function.name, &result);
}
continue;
}
}
// 无工具调用:记录回复,检查是否需要溢出摘要
if !full_text.is_empty() {
self.session.lock().await.add_assistant(full_text.clone());
}
// 检查 token 是否接近预算,触发溢出摘要
{
let s = self.session.lock().await;
let recent = s.recent_messages();
let estimated: u32 = recent.iter().map(|m| builder::estimate_tokens(m)).sum();
if estimated > s.token_budget {
drop(s);
builder::trigger_overflow_summary(&self.session, Some(&self.summarizer())).await;
tracing::info!("📦 上下文超预算,已触发溢出摘要");
}
}
return Ok((full_text, used_tools, last_usage));
}
}
}
// ─── ChatHandle ───
pub struct ChatHandle {
rx: Option<StreamReceiver>,
full_text: String,
tool_calls: Option<Vec<ToolCall>>,
usage: Option<Usage>,
session: Arc<Mutex<ChatSession>>,
}
impl ChatHandle {
pub async fn next_chunk(&mut self) -> Option<StreamChunk> {
let chunk = self.rx.as_mut()?.recv().await;
match &chunk {
Some(StreamChunk::Text(t)) => self.full_text.push_str(t),
Some(StreamChunk::Done { text, tool_calls, usage, .. }) => {
if !text.is_empty() { self.full_text = text.clone(); }
self.tool_calls = tool_calls.clone();
self.usage = usage.clone();
if !self.full_text.is_empty() {
let mut s = self.session.lock().await;
s.add_assistant(self.full_text.clone());
}
self.rx = None;
}
Some(StreamChunk::Error(_)) => { self.rx = None; }
_ => {}
}
chunk
}
pub async fn consume(mut self) -> Result<String, String> {
while let Some(chunk) = self.next_chunk().await {
if let StreamChunk::Error(e) = chunk { return Err(e); }
}
Ok(self.full_text.clone())
}
pub fn current_text(&self) -> &str { &self.full_text }
pub fn get_tool_calls(&self) -> Option<&Vec<ToolCall>> { self.tool_calls.as_ref() }
pub fn get_usage(&self) -> Option<&Usage> { self.usage.as_ref() }
}
pub const DEFAULT_SYSTEM_PROMPT: &str = "You are a concise and helpful assistant replying in Chinese. \
Keep replies practical and natural. \
Unless the user asks for detail, keep most replies under 120 Chinese characters. \
\
如果需要了解用户的长期记忆或历史对话摘要,使用 read_memories / read_summaries 工具。\
如果用户提到重要的个人信息或约定,使用 write_memory 工具记录下来。";
+195
View File
@@ -0,0 +1,195 @@
use crate::llm::provider::{parse_chat_chunk, LlmProvider, ParsedChunk, StreamReceiver, StreamSender};
use crate::llm::types::{ConversationConfig, Message, StreamChunk, ToolCall, Usage};
use async_trait::async_trait;
use reqwest::Client as HttpClient;
use std::collections::BTreeMap;
use tokio::sync::mpsc;
pub struct DeepSeekProvider {
http: HttpClient,
api_key: String,
base_url: String,
}
impl DeepSeekProvider {
pub fn new() -> Result<Self, String> {
let api_key = std::env::var("DEEPSEEK_API_KEY")
.map_err(|_| "请设置 DEEPSEEK_API_KEY 环境变量".to_string())?;
let base_url = std::env::var("DEEPSEEK_BASE_URL")
.unwrap_or_else(|_| "https://api.deepseek.com/v1".to_string());
Ok(Self {
http: HttpClient::builder()
.timeout(std::time::Duration::from_secs(120))
.build()
.map_err(|e| format!("创建 HTTP client 失败: {}", e))?,
api_key,
base_url,
})
}
}
#[async_trait]
impl LlmProvider for DeepSeekProvider {
fn name(&self) -> &str { "deepseek" }
async fn chat_stream(
&self,
config: &ConversationConfig,
messages: &[Message],
) -> Result<StreamReceiver, String> {
let url = format!("{}/chat/completions", self.base_url);
let mut body = serde_json::json!({
"model": config.model,
"messages": messages,
"temperature": config.temperature,
"max_tokens": config.max_tokens,
"stream": true,
});
// 如果有工具,添加 tools 参数
if let Some(ref tools) = config.tools {
body["tools"] = serde_json::Value::Array(tools.clone());
}
if config.thinking {
body["thinking"] = serde_json::json!({"type": "enabled"});
body["reasoning_effort"] = serde_json::json!("high");
body.as_object_mut().map(|obj| {
obj.remove("temperature");
obj.remove("top_p");
});
}
let (tx, rx) = mpsc::channel::<StreamChunk>(64);
let http = self.http.clone();
let api_key = self.api_key.clone();
tokio::spawn(async move {
if let Err(e) = stream_deepseek(http, &url, &api_key, body, tx.clone()).await {
let _ = tx.send(StreamChunk::Error(e)).await;
}
});
Ok(rx)
}
}
async fn stream_deepseek(
http: HttpClient,
url: &str,
api_key: &str,
body: serde_json::Value,
tx: StreamSender,
) -> Result<(), String> {
let response = http
.post(url)
.header("Authorization", format!("Bearer {}", api_key))
.header("Content-Type", "application/json")
.header("Accept", "text/event-stream")
.json(&body)
.send()
.await
.map_err(|e| format!("请求失败: {}", e))?;
if !response.status().is_success() {
let status = response.status();
let text = response.text().await.unwrap_or_default();
let _ = tx.send(StreamChunk::Error(format!("HTTP {}: {}", status, text))).await;
return Err(format!("HTTP {}: {}", status, text));
}
let mut full_text = String::new();
let mut full_reasoning = String::new();
let mut usage: Option<Usage> = None;
let mut tool_call_builders: BTreeMap<i32, (Option<String>, Option<String>, String)> = BTreeMap::new();
use futures_util::StreamExt;
let mut buf = String::new();
let mut stream = response.bytes_stream();
while let Some(chunk_result) = stream.next().await {
let chunk = chunk_result.map_err(|e| format!("读取流失败: {}", e))?;
let text = String::from_utf8_lossy(&chunk);
buf.push_str(&text);
while let Some(line_end) = buf.find('\n') {
let line = buf[..line_end].to_string();
buf = buf[line_end + 1..].to_string();
let trimmed = line.trim();
if trimmed.is_empty() { continue; }
if let Some(parsed) = parse_chat_chunk(trimmed) {
match parsed {
ParsedChunk::Text(t) => {
full_text.push_str(&t);
let _ = tx.send(StreamChunk::Text(t)).await;
}
ParsedChunk::Reasoning(r) => {
full_reasoning.push_str(&r);
let _ = tx.send(StreamChunk::Reasoning(r)).await;
}
ParsedChunk::ToolCallDelta { index, id, name, arguments } => {
let entry = tool_call_builders.entry(index).or_insert((None, None, String::new()));
if let Some(call_id) = id { entry.0 = Some(call_id); }
if let Some(call_name) = name { entry.1 = Some(call_name); }
entry.2.push_str(&arguments);
}
ParsedChunk::FinishReason(_) => {}
ParsedChunk::Usage(u) => { usage = Some(u); }
}
}
}
}
// 检查最后的 buffer
if !buf.trim().is_empty() {
if let Some(parsed) = parse_chat_chunk(buf.trim()) {
match parsed {
ParsedChunk::Text(t) => full_text.push_str(&t),
ParsedChunk::ToolCallDelta { index, id, name, arguments } => {
let entry = tool_call_builders.entry(index).or_insert((None, None, String::new()));
if let Some(call_id) = id { entry.0 = Some(call_id); }
if let Some(call_name) = name { entry.1 = Some(call_name); }
entry.2.push_str(&arguments);
}
ParsedChunk::Usage(u) => { usage = Some(u); }
_ => {}
}
}
}
// 构建 tool_calls
let tool_calls: Option<Vec<ToolCall>> = if tool_call_builders.is_empty() {
None
} else {
Some(
tool_call_builders
.into_values()
.filter_map(|(id, name, args)| {
let name = name?;
let id = id.unwrap_or_else(|| format!("call_{}", rand::random::<u32>()));
Some(ToolCall {
id,
call_type: "function".to_string(),
function: crate::llm::types::ToolFunctionCall {
name,
arguments: args,
},
})
})
.collect(),
)
};
let _ = tx
.send(StreamChunk::Done {
text: full_text,
reasoning: if full_reasoning.is_empty() { None } else { Some(full_reasoning) },
tool_calls,
usage,
})
.await;
Ok(())
}
+157
View File
@@ -0,0 +1,157 @@
use crate::llm::provider::{parse_chat_chunk, LlmProvider, ParsedChunk, StreamReceiver, StreamSender};
use crate::llm::types::{ConversationConfig, Message, StreamChunk};
use async_trait::async_trait;
use reqwest::Client as HttpClient;
use tokio::sync::mpsc;
/// LM Studio 提供商(OpenAI 兼容接口)
pub struct LmStudioProvider {
http: HttpClient,
base_url: String,
api_key: String,
}
impl LmStudioProvider {
pub fn new() -> Result<Self, String> {
let base_url = std::env::var("LM_STUDIO_BASE_URL")
.unwrap_or_else(|_| "http://localhost:1234/v1".to_string());
// 标准化 base_url: 确保有协议和 /v1 路径
let base_url = normalize_base_url(&base_url);
let api_key = std::env::var("LM_STUDIO_API_KEY").unwrap_or_default();
Ok(Self {
http: HttpClient::builder()
.timeout(std::time::Duration::from_secs(120))
.build()
.map_err(|e| format!("创建 HTTP client 失败: {}", e))?,
base_url,
api_key,
})
}
}
#[async_trait]
impl LlmProvider for LmStudioProvider {
fn name(&self) -> &str {
"lmstudio"
}
async fn chat_stream(
&self,
config: &ConversationConfig,
messages: &[Message],
) -> Result<StreamReceiver, String> {
let url = format!("{}/chat/completions", self.base_url);
let body = serde_json::json!({
"model": config.model,
"messages": messages,
"temperature": config.temperature,
"max_tokens": config.max_tokens,
"stream": true,
});
let (tx, rx) = mpsc::channel::<StreamChunk>(64);
let http = self.http.clone();
let api_key = self.api_key.clone();
tokio::spawn(async move {
if let Err(e) = stream_openai(http, &url, &api_key, body, tx.clone()).await {
let _ = tx.send(StreamChunk::Error(e)).await;
}
});
Ok(rx)
}
}
/// 与 DeepSeek 公用相同的 SSE 流式处理,但去除 thinking 相关字段
async fn stream_openai(
http: HttpClient,
url: &str,
api_key: &str,
body: serde_json::Value,
tx: StreamSender,
) -> Result<(), String> {
let mut req = http.post(url).header("Content-Type", "application/json");
if !api_key.is_empty() {
req = req.header("Authorization", format!("Bearer {}", api_key));
}
let response = req
.json(&body)
.send()
.await
.map_err(|e| format!("请求失败: {}", e))?;
if !response.status().is_success() {
let status = response.status();
let text = response.text().await.unwrap_or_default();
let _ = tx.send(StreamChunk::Error(format!("HTTP {}: {}", status, text))).await;
return Err(format!("HTTP {}: {}", status, text));
}
let mut full_text = String::new();
let mut buf = String::new();
let mut stream = response.bytes_stream();
use futures_util::StreamExt;
while let Some(chunk_result) = stream.next().await {
let chunk = chunk_result.map_err(|e| format!("读取流失败: {}", e))?;
let text = String::from_utf8_lossy(&chunk);
buf.push_str(&text);
while let Some(line_end) = buf.find('\n') {
let line = buf[..line_end].to_string();
buf = buf[line_end + 1..].to_string();
let trimmed = line.trim();
if trimmed.is_empty() {
continue;
}
// LM Studio 使用标准 OpenAI SSE 格式
if let Some(parsed) = parse_chat_chunk(trimmed) {
match parsed {
ParsedChunk::Text(t) => {
full_text.push_str(&t);
let _ = tx.send(StreamChunk::Text(t)).await;
}
ParsedChunk::Reasoning(r) => {
let _ = tx.send(StreamChunk::Reasoning(r)).await;
}
_ => {}
}
}
}
}
let _ = tx
.send(StreamChunk::Done {
text: full_text,
reasoning: None,
tool_calls: None,
usage: None,
})
.await;
Ok(())
}
/// 确保 base_url 有协议前缀和 /v1 路径
fn normalize_base_url(raw: &str) -> String {
let raw = raw.trim();
let with_protocol = if raw.starts_with("http://") || raw.starts_with("https://") {
raw.to_string()
} else {
format!("http://{}", raw)
};
// 去掉末尾 /
let trimmed = with_protocol.trim_end_matches('/');
// 确保路径以 /v1 结尾
if trimmed.ends_with("/v1") {
trimmed.to_string()
} else {
format!("{}/v1", trimmed)
}
}
+8
View File
@@ -0,0 +1,8 @@
pub mod conversation;
pub mod deepseek;
pub mod lmstudio;
pub mod provider;
pub mod types;
pub use conversation::{Conversation, Summarizer, ToolExecutor, DEFAULT_SYSTEM_PROMPT};
pub use types::{ConversationConfig, Message, Role, StreamChunk, Usage};
+187
View File
@@ -0,0 +1,187 @@
use crate::llm::types::{ConversationConfig, Message, StreamChunk};
use serde::Deserialize;
use async_trait::async_trait;
use std::sync::Arc;
use tokio::sync::mpsc;
/// 流式返回:每个 StreamChunk 是 delta 或控制信号
pub type StreamReceiver = mpsc::Receiver<StreamChunk>;
pub type StreamSender = mpsc::Sender<StreamChunk>;
/// LLM 提供商抽象
#[async_trait]
pub trait LlmProvider: Send + Sync {
/// 提供商名称
fn name(&self) -> &str;
/// 发起流式对话
async fn chat_stream(
&self,
config: &ConversationConfig,
messages: &[Message],
) -> Result<StreamReceiver, String>;
}
// ─── 内置提供商注册 ───
pub type BoxedProvider = Arc<dyn LlmProvider>;
/// 从配置创建恰当的提供商
pub fn create_provider(_config: &ConversationConfig) -> Result<BoxedProvider, String> {
// 根据 model 前缀或环境变量选择
let provider = std::env::var("LLM_PROVIDER")
.unwrap_or_else(|_| "deepseek".to_string())
.to_lowercase();
match provider.as_str() {
"deepseek" => Ok(Arc::new(super::deepseek::DeepSeekProvider::new()?)),
"lmstudio" => Ok(Arc::new(super::lmstudio::LmStudioProvider::new()?)),
_ => Err(format!("不支持的 LLM 提供商: {}", provider)),
}
}
// ─── 内部:SSE 解析工具 ───
/// 从字节流中解析 SSE 行
pub(crate) fn parse_sse_line(line: &str) -> Option<(String, String)> {
// 支持两种格式:
// data: {...}
// event: ...
let line = line.trim();
if line.is_empty() || line.starts_with(':') {
return None;
}
if let Some(pos) = line.find(": ") {
let field = line[..pos].trim().to_string();
let value = line[pos + 2..].trim_start().to_string();
Some((field, value))
} else if let Some(pos) = line.find(':') {
let field = line[..pos].trim().to_string();
let value = line[pos + 1..].trim_start().to_string();
Some((field, value))
} else {
None
}
}
/// 流式块解析结果
pub(crate) enum ParsedChunk {
Text(String),
Reasoning(String),
ToolCallDelta {
index: i32,
id: Option<String>,
name: Option<String>,
arguments: String,
},
FinishReason(String),
Usage(super::types::Usage),
}
/// 从 JSON body 中解析 DeepSeek/OpenAI 流式 delta
pub(crate) fn parse_chat_chunk(
line: &str,
) -> Option<ParsedChunk> {
if !line.starts_with("data: ") {
return None;
}
let data = line["data: ".len()..].trim();
if data == "[DONE]" {
return None;
}
#[derive(Deserialize)]
struct ToolCallDelta {
#[serde(default)]
index: Option<i32>,
#[serde(default)]
id: Option<String>,
#[serde(rename = "type", default)]
call_type: Option<String>,
#[serde(default)]
function: Option<ToolFunctionDelta>,
}
#[derive(Deserialize)]
struct ToolFunctionDelta {
#[serde(default)]
name: Option<String>,
#[serde(default)]
arguments: Option<String>,
}
#[derive(Deserialize)]
struct Delta {
#[serde(default)]
content: Option<String>,
#[serde(default)]
reasoning_content: Option<String>,
#[serde(default)]
tool_calls: Option<Vec<ToolCallDelta>>,
}
#[derive(Deserialize)]
struct ChunkChoice {
delta: Delta,
#[serde(default)]
finish_reason: Option<String>,
}
#[derive(Deserialize)]
struct ChunkResponse {
choices: Vec<ChunkChoice>,
#[serde(default)]
usage: Option<super::types::Usage>,
}
let parsed: ChunkResponse = match serde_json::from_str(data) {
Ok(p) => p,
Err(_) => return None,
};
// 提取 usage(可能在最后一个 chunk
if let Some(ref usage) = parsed.usage {
if usage.total_tokens > 0 {
return Some(ParsedChunk::Usage(usage.clone()));
}
}
for choice in parsed.choices {
// 工具调用 delta
if let Some(tool_calls) = &choice.delta.tool_calls {
for tc in tool_calls {
let idx = tc.index.unwrap_or(0);
let args = tc.function.as_ref()
.and_then(|f| f.arguments.clone())
.unwrap_or_default();
let name = tc.function.as_ref().and_then(|f| f.name.clone());
return Some(ParsedChunk::ToolCallDelta {
index: idx,
id: tc.id.clone(),
name,
arguments: args,
});
}
}
if let Some(reasoning) = &choice.delta.reasoning_content {
if !reasoning.is_empty() {
return Some(ParsedChunk::Reasoning(reasoning.clone()));
}
}
if let Some(content) = &choice.delta.content {
if !content.is_empty() {
return Some(ParsedChunk::Text(content.clone()));
}
}
if let Some(reason) = &choice.finish_reason {
if !reason.is_empty() {
return Some(ParsedChunk::FinishReason(reason.clone()));
}
}
}
None
}
+125
View File
@@ -0,0 +1,125 @@
use serde::{Deserialize, Serialize};
// ─── 角色 ───
#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
#[serde(rename_all = "lowercase")]
pub enum Role {
System,
User,
Assistant,
Tool,
}
// ─── 消息 ───
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct Message {
pub role: Role,
#[serde(default)]
pub content: String,
#[serde(skip_serializing_if = "Option::is_none")]
pub tool_call_id: Option<String>,
#[serde(skip_serializing_if = "Option::is_none")]
pub name: Option<String>,
#[serde(skip_serializing_if = "Option::is_none")]
pub tool_calls: Option<Vec<ToolCall>>,
}
impl Message {
pub fn system(content: impl Into<String>) -> Self {
Self { role: Role::System, content: content.into(), tool_call_id: None, name: None, tool_calls: None }
}
pub fn user(content: impl Into<String>) -> Self {
Self { role: Role::User, content: content.into(), tool_call_id: None, name: None, tool_calls: None }
}
pub fn assistant(content: impl Into<String>) -> Self {
Self { role: Role::Assistant, content: content.into(), tool_call_id: None, name: None, tool_calls: None }
}
pub fn assistant_with_tool_calls(tool_calls: Vec<ToolCall>) -> Self {
Self { role: Role::Assistant, content: String::new(), tool_call_id: None, name: None, tool_calls: Some(tool_calls) }
}
pub fn tool_result(tool_call_id: &str, name: &str, result: &str) -> Self {
Self { role: Role::Tool, content: result.to_string(), tool_call_id: Some(tool_call_id.to_string()), name: Some(name.to_string()), tool_calls: None }
}
}
// ─── 工具调用 ───
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct ToolCall {
pub id: String,
#[serde(rename = "type")]
pub call_type: String, // "function"
pub function: ToolFunctionCall,
}
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct ToolFunctionCall {
pub name: String,
pub arguments: String, // JSON string
}
// ─── 流式响应块 ───
#[derive(Debug, Clone)]
pub enum StreamChunk {
/// 文本片段 delta
Text(String),
/// 思考/推理内容(DeepSeek thinking
Reasoning(String),
/// 完成(携带完整文本,以及可选的工具调用)
Done {
text: String,
reasoning: Option<String>,
tool_calls: Option<Vec<ToolCall>>,
usage: Option<Usage>,
},
/// 错误
Error(String),
}
// ─── Token 用量 ───
#[derive(Debug, Clone, Serialize, Deserialize, Default)]
pub struct Usage {
pub prompt_tokens: u32,
pub completion_tokens: u32,
pub total_tokens: u32,
#[serde(default)]
pub prompt_cache_hit_tokens: u32,
#[serde(default)]
pub prompt_cache_miss_tokens: u32,
}
// ─── 对话配置 ───
#[derive(Debug, Clone)]
pub struct ConversationConfig {
pub system_prompt: String,
pub model: String,
pub temperature: f32,
pub max_tokens: u32,
pub thinking: bool,
/// LLM function calling 的工具定义(JSON 数组)
pub tools: Option<Vec<serde_json::Value>>,
}
impl Default for ConversationConfig {
fn default() -> Self {
Self {
system_prompt: "You are a concise and helpful assistant replying in Chinese. \
Keep replies practical and natural."
.to_string(),
model: "deepseek-v4-flash".to_string(),
temperature: 0.7,
max_tokens: 4096,
thinking: true,
tools: None,
}
}
}