feat(llm): 实现 Phase 0 剩余四个模块

实现 ProviderRegistry、HookExecutor、StreamEvents 和 Auto-compaction 模块,并集成到 LlmCycle 中
This commit is contained in:
徐涛
2026-06-02 08:51:42 +08:00
parent 69b6dd942b
commit 32f3edaf19
13 changed files with 1299 additions and 9 deletions
+122 -1
View File
@@ -1,12 +1,18 @@
use std::pin::Pin;
use std::time::Duration;
use async_trait::async_trait;
use bytes::Bytes;
use futures_core::stream::Stream;
use futures_util::StreamExt;
use reqwest::Client;
use tracing::{debug, error, info};
use super::LlmProvider;
use crate::llm::error::LlmError;
use crate::llm::types::{ChatRequest, ChatResponse, OpenaiChatResponse};
use crate::llm::types::{
ChatRequest, ChatResponse, OpenaiChatChunk, OpenaiChatResponse, StreamOptions,
};
pub struct OpenaiProvider {
http_client: Client,
@@ -111,4 +117,119 @@ impl LlmProvider for OpenaiProvider {
Ok(ChatResponse::from(chat_response))
}
async fn chat_stream(
&self,
mut request: ChatRequest,
) -> Result<
Pin<Box<dyn Stream<Item = Result<OpenaiChatChunk, LlmError>> + Send>>,
LlmError,
> {
let url = format!("{}/chat/completions", self.base_url.trim_end_matches('/'));
request.stream = Some(true);
request.stream_options = Some(StreamOptions {
include_usage: Some(true),
include_obfuscation: None,
});
info!(model = %request.model, "发送 LLM 流式请求");
let response = self
.http_client
.post(&url)
.header("Authorization", format!("Bearer {}", self.api_key))
.json(&request)
.send()
.await
.map_err(|e| {
error!(error = %e, "流式请求失败");
Self::map_reqwest_error(e)
})?;
let status = response.status();
let status_code: u16 = status.as_u16();
if !status.is_success() {
let body_text = response.text().await.unwrap_or_default();
error!(status = status_code, body = %body_text, "流式请求失败");
return Err(LlmError::Request {
status: status_code,
body: body_text,
});
}
let byte_stream = response.bytes_stream().map(|r| {
r.map_err(|e| LlmError::Other(format!("流式读取失败: {}", e)))
});
Ok(Box::pin(SseChunkStream::new(byte_stream)))
}
}
struct SseChunkStream<S> {
inner: S,
buffer: String,
}
impl<S: Stream<Item = Result<Bytes, LlmError>> + Unpin> SseChunkStream<S> {
fn new(stream: S) -> Self {
Self {
inner: stream,
buffer: String::new(),
}
}
}
impl<S: Stream<Item = Result<Bytes, LlmError>> + Unpin> Stream for SseChunkStream<S> {
type Item = Result<OpenaiChatChunk, LlmError>;
fn poll_next(mut self: Pin<&mut Self>, cx: &mut std::task::Context<'_>) -> std::task::Poll<Option<Self::Item>> {
loop {
if let Some(pos) = self.buffer.find("\n") {
let line = self.buffer.drain(..pos + 1).collect::<String>();
let trimmed = line.trim();
if trimmed.is_empty()
|| trimmed == "data: [DONE]"
|| trimmed == "[DONE]"
|| trimmed == "data:"
{
continue;
}
let data = if let Some(p) = trimmed.strip_prefix("data: ") {
p
} else {
trimmed
};
match serde_json::from_str::<OpenaiChatChunk>(data) {
Ok(chunk) => return std::task::Poll::Ready(Some(Ok(chunk))),
Err(e) => {
return std::task::Poll::Ready(Some(Err(LlmError::Other(format!(
"Chunk 解析失败: {} | raw: {}",
e, data
)))));
}
}
}
match Pin::new(&mut self.inner).poll_next(cx) {
std::task::Poll::Ready(Some(Ok(bytes))) => {
if let Ok(s) = std::str::from_utf8(&bytes) {
self.buffer.push_str(s);
}
}
std::task::Poll::Ready(Some(Err(e))) => {
return std::task::Poll::Ready(Some(Err(e)));
}
std::task::Poll::Ready(None) => {
if self.buffer.is_empty() {
return std::task::Poll::Ready(None);
}
self.buffer.clear();
return std::task::Poll::Ready(None);
}
std::task::Poll::Pending => return std::task::Poll::Pending,
}
}
}
}
+68
View File
@@ -0,0 +1,68 @@
//! Provider 注册表 —— 多 Provider 实例的注册与发现。
use std::collections::HashMap;
use crate::llm::error::LlmError;
use crate::llm::provider::{create_provider, LlmProvider, ProviderConfig, ProviderType};
/// Provider 注册表 —— 管理多个 LLM Provider 实例。
///
/// 支持注册命名 Provider、按名称查找、设置默认 Provider。
pub struct ProviderRegistry {
providers: HashMap<String, Box<dyn LlmProvider>>,
default_name: Option<String>,
}
impl Default for ProviderRegistry {
fn default() -> Self {
Self::new()
}
}
impl ProviderRegistry {
/// 创建一个空的注册表。
pub fn new() -> Self {
Self {
providers: HashMap::new(),
default_name: None,
}
}
/// 注册一个已初始化的 Provider 实例。
pub fn register(&mut self, name: impl Into<String>, provider: Box<dyn LlmProvider>) {
self.providers.insert(name.into(), provider);
}
/// 通过 ProviderType + ProviderConfig 创建并注册。
pub fn register_with_config(
&mut self,
name: impl Into<String>,
provider_type: ProviderType,
config: ProviderConfig,
) -> Result<(), LlmError> {
let provider = create_provider(provider_type, config)?;
self.register(name, provider);
Ok(())
}
/// 设置默认 Provider。
pub fn set_default(&mut self, name: &str) -> Result<(), LlmError> {
if !self.providers.contains_key(name) {
return Err(LlmError::Other(format!("Provider '{}' 不存在", name)));
}
self.default_name = Some(name.to_string());
Ok(())
}
/// 按名称查找 Provider。
pub fn get(&self, name: &str) -> Option<&dyn LlmProvider> {
self.providers.get(name).map(|p| p.as_ref())
}
/// 获取默认 Provider。
pub fn get_default(&self) -> Option<&dyn LlmProvider> {
self.default_name
.as_ref()
.and_then(|name| self.get(name))
}
}