feat(llm): 实现 Phase 0 剩余四个模块
实现 ProviderRegistry、HookExecutor、StreamEvents 和 Auto-compaction 模块,并集成到 LlmCycle 中
This commit is contained in:
+122
-1
@@ -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,
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
@@ -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))
|
||||
}
|
||||
}
|
||||
Reference in New Issue
Block a user