feat(llm): 实现 Phase 0 剩余四个模块
实现 ProviderRegistry、HookExecutor、StreamEvents 和 Auto-compaction 模块，并集成到 LlmCycle 中
This commit is contained in:
徐涛
+216
-1
@@ -1,21 +1,38 @@
|
||||
//! LLM 调用周期控制模块。
|
||||
|
||||
mod retry;
|
||||
pub mod usage;
|
||||
|
||||
pub use retry::RetryConfig;
|
||||
pub use usage::{CostTracker, Usage};
|
||||
|
||||
use std::pin::Pin;
|
||||
use std::sync::Arc;
|
||||
|
||||
use futures_core::stream::Stream;
|
||||
use async_stream::stream;
|
||||
|
||||
use crate::llm::compact::{should_compact, microcompact, CompactConfig, CompactState};
|
||||
use crate::llm::cycle::retry::should_retry;
|
||||
use crate::llm::error::LlmError;
|
||||
use crate::llm::hooks::{HookContext, HookExecutor};
|
||||
use crate::llm::provider::LlmProvider;
|
||||
use crate::llm::stream::StreamEvent;
|
||||
use crate::llm::types::{
|
||||
ChatRequest, ChatResponse, OpenaiChatMessage, OpenaiTool, ToolChoice, ToolDefinition,
|
||||
};
|
||||
|
||||
/// LLM 调用周期配置。
|
||||
pub struct CycleConfig {
|
||||
/// 模型名称。
|
||||
pub model: String,
|
||||
/// 最大输出 token 数。
|
||||
pub max_tokens: Option<u32>,
|
||||
/// 采样温度。
|
||||
pub temperature: Option<f32>,
|
||||
/// 最大对话轮次。
|
||||
pub max_turns: Option<u32>,
|
||||
/// 重试策略配置。
|
||||
pub retry: RetryConfig,
|
||||
}
|
||||
|
||||
@@ -31,15 +48,20 @@ impl Default for CycleConfig {
|
||||
}
|
||||
}
|
||||
|
||||
/// LLM 调用周期 —— 管理一次或多次 LLM 请求的生命周期。
|
||||
pub struct LlmCycle {
|
||||
provider: Box<dyn LlmProvider>,
|
||||
config: CycleConfig,
|
||||
usage: CostTracker,
|
||||
messages: Vec<OpenaiChatMessage>,
|
||||
system_prompt: Option<String>,
|
||||
hook_executor: Option<Arc<HookExecutor>>,
|
||||
compact_config: Option<CompactConfig>,
|
||||
compact_state: CompactState,
|
||||
}
|
||||
|
||||
impl LlmCycle {
|
||||
/// 创建一个新的 LlmCycle。
|
||||
pub fn new(provider: Box<dyn LlmProvider>, config: CycleConfig) -> Self {
|
||||
Self {
|
||||
provider,
|
||||
@@ -47,30 +69,51 @@ impl LlmCycle {
|
||||
usage: CostTracker::default(),
|
||||
messages: Vec::new(),
|
||||
system_prompt: None,
|
||||
hook_executor: None,
|
||||
compact_config: None,
|
||||
compact_state: CompactState::new(),
|
||||
}
|
||||
}
|
||||
|
||||
/// 设置系统提示词。
|
||||
pub fn with_system_prompt(mut self, prompt: String) -> Self {
|
||||
self.system_prompt = Some(prompt);
|
||||
self
|
||||
}
|
||||
|
||||
/// 设置钩子执行器。
|
||||
pub fn with_hook_executor(mut self, executor: HookExecutor) -> Self {
|
||||
self.hook_executor = Some(Arc::new(executor));
|
||||
self
|
||||
}
|
||||
|
||||
/// 设置上下文压缩配置。
|
||||
pub fn with_compact_config(mut self, config: CompactConfig) -> Self {
|
||||
self.compact_config = Some(config);
|
||||
self
|
||||
}
|
||||
|
||||
/// 获取用量追踪器引用。
|
||||
pub fn usage(&self) -> &CostTracker {
|
||||
&self.usage
|
||||
}
|
||||
|
||||
/// 获取消息历史引用。
|
||||
pub fn messages(&self) -> &[OpenaiChatMessage] {
|
||||
&self.messages
|
||||
}
|
||||
|
||||
/// 清空消息历史。
|
||||
pub fn clear_messages(&mut self) {
|
||||
self.messages.clear();
|
||||
}
|
||||
|
||||
/// 重置用量统计。
|
||||
pub fn reset_usage(&mut self) {
|
||||
self.usage.reset();
|
||||
}
|
||||
|
||||
/// 提交用户消息并获取 LLM 响应。
|
||||
pub async fn submit(
|
||||
&mut self,
|
||||
prompt: String,
|
||||
@@ -78,31 +121,203 @@ impl LlmCycle {
|
||||
) -> Result<ChatResponse, LlmError> {
|
||||
self.messages.push(OpenaiChatMessage::user_text(prompt));
|
||||
|
||||
if let Some(ref config) = self.compact_config
|
||||
&& should_compact(&self.messages, config, &self.compact_state)
|
||||
{
|
||||
let freed = microcompact(&mut self.messages, config.keep_recent);
|
||||
if freed > 0 {
|
||||
self.compact_state.record_success();
|
||||
}
|
||||
}
|
||||
|
||||
let mut attempts = 0;
|
||||
|
||||
loop {
|
||||
let request = self.build_request(&tools);
|
||||
|
||||
if let Some(ref executor) = self.hook_executor {
|
||||
let ctx = HookContext::new(crate::llm::hooks::HookEvent::PreRequest)
|
||||
.with_request(&request);
|
||||
let results = executor
|
||||
.execute(crate::llm::hooks::HookEvent::PreRequest, &ctx)
|
||||
.await;
|
||||
if results.iter().any(|r| r.should_block) {
|
||||
let reason = results
|
||||
.iter()
|
||||
.find(|r| r.should_block)
|
||||
.and_then(|r| r.reason.clone())
|
||||
.unwrap_or_else(|| "Blocked by pre-request hook".to_string());
|
||||
return Err(LlmError::Other(reason));
|
||||
}
|
||||
}
|
||||
|
||||
match self.provider.chat(request).await {
|
||||
Ok(response) => {
|
||||
self.messages.push(response.message.clone());
|
||||
if let Some(ref executor) = self.hook_executor {
|
||||
let post_request = self.build_request(&tools);
|
||||
let ctx = HookContext::new(crate::llm::hooks::HookEvent::PostRequest)
|
||||
.with_request(&post_request);
|
||||
executor
|
||||
.execute(crate::llm::hooks::HookEvent::PostRequest, &ctx)
|
||||
.await;
|
||||
}
|
||||
|
||||
self.messages.push(response.message.clone());
|
||||
self.usage.add(&response.usage);
|
||||
|
||||
return Ok(response);
|
||||
}
|
||||
Err(e) if should_retry(&e) && attempts < self.config.retry.max_retries => {
|
||||
attempts += 1;
|
||||
|
||||
if let Some(ref executor) = self.hook_executor {
|
||||
let ctx = HookContext::new(crate::llm::hooks::HookEvent::OnRetry)
|
||||
.with_error(&e)
|
||||
.with_attempt(attempts);
|
||||
executor
|
||||
.execute(crate::llm::hooks::HookEvent::OnRetry, &ctx)
|
||||
.await;
|
||||
}
|
||||
|
||||
let delay = self.config.retry.compute_delay(attempts);
|
||||
tokio::time::sleep(delay).await;
|
||||
}
|
||||
Err(e) => {
|
||||
if let Some(ref executor) = self.hook_executor {
|
||||
let ctx =
|
||||
HookContext::new(crate::llm::hooks::HookEvent::OnError).with_error(&e);
|
||||
executor
|
||||
.execute(crate::llm::hooks::HookEvent::OnError, &ctx)
|
||||
.await;
|
||||
}
|
||||
|
||||
return Err(e);
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// 提交用户消息并返回语义事件流。
|
||||
///
|
||||
/// 与 `submit` 不同，该方法返回流式事件而非完整响应。
|
||||
/// 适用于需要实时处理 LLM 输出的场景。
|
||||
pub async fn submit_stream(
|
||||
&mut self,
|
||||
prompt: String,
|
||||
tools: Vec<ToolDefinition>,
|
||||
) -> Result<Pin<Box<dyn Stream<Item = StreamEvent> + Send>>, LlmError> {
|
||||
self.messages.push(OpenaiChatMessage::user_text(prompt));
|
||||
|
||||
if let Some(ref config) = self.compact_config
|
||||
&& should_compact(&self.messages, config, &self.compact_state)
|
||||
{
|
||||
let freed = microcompact(&mut self.messages, config.keep_recent);
|
||||
if freed > 0 {
|
||||
self.compact_state.record_success();
|
||||
}
|
||||
}
|
||||
|
||||
let request = self.build_request(&tools);
|
||||
|
||||
if let Some(ref executor) = self.hook_executor {
|
||||
let ctx = HookContext::new(crate::llm::hooks::HookEvent::PreRequest)
|
||||
.with_request(&request);
|
||||
let results = executor
|
||||
.execute(crate::llm::hooks::HookEvent::PreRequest, &ctx)
|
||||
.await;
|
||||
if results.iter().any(|r| r.should_block) {
|
||||
let reason = results
|
||||
.iter()
|
||||
.find(|r| r.should_block)
|
||||
.and_then(|r| r.reason.clone())
|
||||
.unwrap_or_else(|| "Blocked by pre-request hook".to_string());
|
||||
return Err(LlmError::Other(reason));
|
||||
}
|
||||
}
|
||||
|
||||
let chunk_stream = self.provider.chat_stream(request).await?;
|
||||
let hook_executor = self.hook_executor.clone();
|
||||
let post_request = self.build_request(&tools);
|
||||
|
||||
Ok(Box::pin(stream! {
|
||||
use futures_util::StreamExt;
|
||||
let mut chunk_stream = chunk_stream;
|
||||
|
||||
while let Some(result) = chunk_stream.next().await {
|
||||
match result {
|
||||
Ok(chunk) => {
|
||||
let mut assistant_text = String::new();
|
||||
let mut tool_started: Option<(String, String, String)> = None;
|
||||
|
||||
for choice in &chunk.choices {
|
||||
let delta = &choice.delta;
|
||||
|
||||
if let Some(content) = &delta.content {
|
||||
assistant_text.push_str(content);
|
||||
}
|
||||
|
||||
if let Some(tool_calls) = &delta.tool_calls
|
||||
&& let Some(tc) = tool_calls.first()
|
||||
{
|
||||
let crate::llm::types::OpenaiToolCall::Function { id, function } = tc;
|
||||
tool_started = Some((id.clone(), function.name.clone(), function.arguments.clone()));
|
||||
}
|
||||
}
|
||||
|
||||
if !assistant_text.is_empty() {
|
||||
yield StreamEvent::AssistantTextDelta { text: assistant_text };
|
||||
}
|
||||
|
||||
if let Some((tool_call_id, tool_name, arguments)) = tool_started {
|
||||
let args: serde_json::Value = serde_json::from_str(&arguments)
|
||||
.unwrap_or(serde_json::Value::Null);
|
||||
yield StreamEvent::ToolExecutionStarted {
|
||||
tool_name,
|
||||
input: args,
|
||||
tool_call_id,
|
||||
};
|
||||
}
|
||||
|
||||
for choice in &chunk.choices {
|
||||
if let Some(finish_reason) = &choice.finish_reason {
|
||||
yield StreamEvent::TurnComplete {
|
||||
reason: *finish_reason,
|
||||
};
|
||||
}
|
||||
}
|
||||
|
||||
if let Some(usage_info) = &chunk.usage {
|
||||
yield StreamEvent::CostUpdate { usage: *usage_info };
|
||||
}
|
||||
}
|
||||
Err(e) => {
|
||||
if let Some(ref executor) = hook_executor {
|
||||
let ctx = crate::llm::hooks::HookContext::new(
|
||||
crate::llm::hooks::HookEvent::OnError,
|
||||
)
|
||||
.with_error(&e);
|
||||
executor
|
||||
.execute(crate::llm::hooks::HookEvent::OnError, &ctx)
|
||||
.await;
|
||||
}
|
||||
yield StreamEvent::Error { message: e.to_string() };
|
||||
break;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
if let Some(ref executor) = hook_executor {
|
||||
let ctx = crate::llm::hooks::HookContext::new(
|
||||
crate::llm::hooks::HookEvent::PostRequest,
|
||||
)
|
||||
.with_request(&post_request);
|
||||
executor
|
||||
.execute(crate::llm::hooks::HookEvent::PostRequest, &ctx)
|
||||
.await;
|
||||
}
|
||||
}))
|
||||
}
|
||||
|
||||
fn build_request(&self, tools: &[ToolDefinition]) -> ChatRequest {
|
||||
let mut messages = self.messages.clone();
|
||||
|
||||
|
||||
Reference in New Issue
Block a user