452 lines
18 KiB
Rust
452 lines
18 KiB
Rust
//! [`StreamProvider`] trait——agent 循环与具体 LLM 后端之间的 I/O 边界。
|
||
//! The [`StreamProvider`] trait — the I/O boundary between the agent loop and
|
||
//! concrete LLM backends.
|
||
//!
|
||
//! provider 以 [`StreamEvent`] 流的形式输出。它们**绝不**通过返回 `Err` 来表示
|
||
//! API 失败——失败被编码为带相应 [`StopReason`](crate::model::StopReason) 的
|
||
//! [`StreamEvent::Error`]/[`StreamEvent::Done`]。返回 `Err` 仅留给真正意料之外
|
||
//! 的情况（例如 provider 无法被构造）。
|
||
//! Providers stream [`StreamEvent`]s. They must **never** return an `Err` to
|
||
//! signal an API failure — failures are encoded as
|
||
//! [`StreamEvent::Error`]/[`StreamEvent::Done`] with an appropriate
|
||
//! [`StopReason`](crate::model::StopReason). Returning `Err` is reserved for
|
||
//! truly unexpected conditions (e.g. the provider could not be constructed).
|
||
|
||
use crate::error::CoreResult;
|
||
use crate::model::*;
|
||
use focus_json::JsonValue;
|
||
use std::fmt;
|
||
|
||
/// agent 向 provider 发出的请求。
|
||
/// A request the agent makes to a provider.
|
||
#[derive(Debug, Clone)]
|
||
pub struct ProviderRequest {
|
||
/// 模型 id。
|
||
/// The model id.
|
||
pub model: String,
|
||
/// 系统提示词。
|
||
/// The system prompt.
|
||
pub system_prompt: String,
|
||
/// 对话消息列表。
|
||
/// The conversation messages.
|
||
pub messages: Vec<Message>,
|
||
/// 工具定义（原始 JSON）。
|
||
/// Tool definitions (raw JSON).
|
||
pub tools: JsonValue,
|
||
/// 最大输出 token 数。
|
||
/// Max output tokens.
|
||
pub max_tokens: Option<u64>,
|
||
/// 采样温度。
|
||
/// Sampling temperature.
|
||
pub temperature: Option<f64>,
|
||
}
|
||
|
||
/// provider 流式响应过程中发出的事件。
|
||
/// Events emitted by a provider while streaming a response.
|
||
///
|
||
/// 每个携带 delta 的变体都会在 `partial` 中携带当前完整的 [`AssistantMessage`]
|
||
/// 快照，对应 pi 的协议——这样循环本身无需累积 delta。
|
||
/// Every variant that carries a delta also carries the current full
|
||
/// [`AssistantMessage`] snapshot in `partial`, mirroring pi's protocol so the
|
||
/// loop never needs to accumulate deltas itself.
|
||
#[derive(Debug, Clone)]
|
||
pub enum StreamEvent {
|
||
/// 响应已开始；附带一个空/部分的 assistant 消息。
|
||
/// The response has started; here is an empty/partial assistant message.
|
||
Start { partial: AssistantMessage },
|
||
/// 在 `content_index` 处打开了一个文本块。
|
||
/// A text block opened at `content_index`.
|
||
TextStart {
|
||
content_index: usize,
|
||
partial: AssistantMessage,
|
||
},
|
||
/// 向 `content_index` 处的文本块追加了文本。
|
||
/// Text appended to the block at `content_index`.
|
||
TextDelta {
|
||
content_index: usize,
|
||
delta: String,
|
||
partial: AssistantMessage,
|
||
},
|
||
/// `content_index` 处的文本块结束。
|
||
/// The text block at `content_index` finished.
|
||
TextEnd {
|
||
content_index: usize,
|
||
partial: AssistantMessage,
|
||
},
|
||
/// 打开了一个思考块。
|
||
/// A thinking block opened.
|
||
ThinkingStart {
|
||
content_index: usize,
|
||
partial: AssistantMessage,
|
||
},
|
||
/// 追加了思考文本。
|
||
/// Thinking text appended.
|
||
ThinkingDelta {
|
||
content_index: usize,
|
||
delta: String,
|
||
partial: AssistantMessage,
|
||
},
|
||
/// 思考块结束。
|
||
/// The thinking block finished.
|
||
ThinkingEnd {
|
||
content_index: usize,
|
||
partial: AssistantMessage,
|
||
},
|
||
/// 打开了一个工具调用块。
|
||
/// A tool-call block opened.
|
||
ToolCallStart {
|
||
content_index: usize,
|
||
partial: AssistantMessage,
|
||
},
|
||
/// 向工具调用追加了部分 JSON 参数。
|
||
/// Partial JSON arguments appended to the tool call.
|
||
ToolCallDelta {
|
||
content_index: usize,
|
||
delta: String,
|
||
partial: AssistantMessage,
|
||
},
|
||
/// 工具调用结束，附带最终的 [`ToolCall`]。
|
||
/// The tool call finished with its final [`ToolCall`].
|
||
ToolCallEnd {
|
||
content_index: usize,
|
||
tool_call: ToolCall,
|
||
partial: AssistantMessage,
|
||
},
|
||
/// 响应正常完成。
|
||
/// The response completed normally.
|
||
Done { message: AssistantMessage },
|
||
/// 响应失败或被中止。
|
||
/// The response failed or was aborted.
|
||
Error { error: AssistantMessage },
|
||
}
|
||
|
||
impl fmt::Display for StreamEvent {
|
||
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
|
||
match self {
|
||
StreamEvent::Start { .. } => f.write_str("start"),
|
||
StreamEvent::TextStart { .. } => f.write_str("text_start"),
|
||
StreamEvent::TextDelta { .. } => f.write_str("text_delta"),
|
||
StreamEvent::TextEnd { .. } => f.write_str("text_end"),
|
||
StreamEvent::ThinkingStart { .. } => f.write_str("thinking_start"),
|
||
StreamEvent::ThinkingDelta { .. } => f.write_str("thinking_delta"),
|
||
StreamEvent::ThinkingEnd { .. } => f.write_str("thinking_end"),
|
||
StreamEvent::ToolCallStart { .. } => f.write_str("toolcall_start"),
|
||
StreamEvent::ToolCallDelta { .. } => f.write_str("toolcall_delta"),
|
||
StreamEvent::ToolCallEnd { .. } => f.write_str("toolcall_end"),
|
||
StreamEvent::Done { .. } => f.write_str("done"),
|
||
StreamEvent::Error { .. } => f.write_str("error"),
|
||
}
|
||
}
|
||
}
|
||
|
||
/// 一个辅助器：把 provider 的增量 delta 累积进 [`AssistantMessage`]，
|
||
/// 并发出带最新 partial 的 [`StreamEvent`]。
|
||
/// A helper that accumulates incremental provider deltas into an
|
||
/// [`AssistantMessage`] and emits [`StreamEvent`]s with up-to-date partials.
|
||
///
|
||
/// 像 OpenAI 这样的 provider 会以扁平序列流式输出内容/工具调用 delta，且没有显式
|
||
/// 的 block-start/stop 标记；本 reducer 填补这一空缺，使 provider 只需喂入文本/
|
||
/// 工具 delta 和一个 finish reason。
|
||
/// Providers like OpenAI stream content/tool-call deltas in a flat sequence
|
||
/// without explicit block-start/stop markers; this reducer fills that gap so a
|
||
/// provider only needs to feed text/tool deltas and a finish reason.
|
||
pub struct ProviderEventReducer {
|
||
model: String,
|
||
content: Vec<ContentBlock>,
|
||
/// 与 `content` 对齐的参数缓冲：非工具调用槽为 `None`。
|
||
/// Argument buffer aligned with `content`; non-tool-call slots are `None`.
|
||
args_by_index: Vec<Option<String>>,
|
||
/// provider 局部索引 → content 位置 的映射（支持交错的多工具调用）。
|
||
/// Mapping from provider-local index to content position (supports
|
||
/// interleaved multi-tool-call streams).
|
||
tool_call_pos: Vec<(usize, usize)>,
|
||
started: bool,
|
||
/// 当前打开的文本块在 `content` 中的位置。
|
||
/// Position of the currently open text block in `content`.
|
||
text_index: Option<usize>,
|
||
/// 当前打开的思考块在 `content` 中的位置。
|
||
/// Position of the currently open thinking block in `content`.
|
||
thinking_index: Option<usize>,
|
||
usage: Usage,
|
||
stop_reason: StopReason,
|
||
error_message: Option<String>,
|
||
}
|
||
|
||
impl ProviderEventReducer {
|
||
pub fn new(model: &str) -> Self {
|
||
Self {
|
||
model: model.to_string(),
|
||
content: Vec::new(),
|
||
args_by_index: Vec::new(),
|
||
tool_call_pos: Vec::new(),
|
||
started: false,
|
||
text_index: None,
|
||
thinking_index: None,
|
||
usage: Usage::default(),
|
||
stop_reason: StopReason::Stop,
|
||
error_message: None,
|
||
}
|
||
}
|
||
|
||
/// 用当前累积状态构造一条 partial [`AssistantMessage`] 快照。
|
||
/// Build a partial [`AssistantMessage`] snapshot from the current accumulator.
|
||
fn partial(&self) -> AssistantMessage {
|
||
AssistantMessage {
|
||
content: self.content.clone(),
|
||
model: self.model.clone(),
|
||
usage: self.usage.clone(),
|
||
stop_reason: self.stop_reason,
|
||
error_message: self.error_message.clone(),
|
||
timestamp: now_ms(),
|
||
}
|
||
}
|
||
|
||
/// 首次调用时发出 `Start` 事件；之后返回 `None`。
|
||
/// Emit the `Start` event on first call; `None` thereafter.
|
||
fn ensure_started(&mut self) -> Option<StreamEvent> {
|
||
if !self.started {
|
||
self.started = true;
|
||
return Some(StreamEvent::Start {
|
||
partial: self.partial(),
|
||
});
|
||
}
|
||
None
|
||
}
|
||
|
||
/// 喂入一个文本 delta；按需打开 Start + TextStart。
|
||
/// Feed a text delta. Opens a Start + TextStart as needed.
|
||
pub fn text_delta(&mut self, text: &str) -> Vec<StreamEvent> {
|
||
let mut out = Vec::new();
|
||
if let Some(ev) = self.ensure_started() {
|
||
out.push(ev);
|
||
}
|
||
if self.text_index.is_none() {
|
||
self.content.push(ContentBlock::Text(TextContent {
|
||
text: String::new(),
|
||
signature: None,
|
||
}));
|
||
self.args_by_index.push(None);
|
||
self.text_index = Some(self.content.len() - 1);
|
||
out.push(StreamEvent::TextStart {
|
||
content_index: self.content.len() - 1,
|
||
partial: self.partial(),
|
||
});
|
||
}
|
||
let idx = self.text_index.expect("text block open");
|
||
if let Some(ContentBlock::Text(t)) = self.content.get_mut(idx) {
|
||
t.text.push_str(text);
|
||
}
|
||
out.push(StreamEvent::TextDelta {
|
||
content_index: idx,
|
||
delta: text.to_string(),
|
||
partial: self.partial(),
|
||
});
|
||
out
|
||
}
|
||
|
||
/// 喂入一个思考 delta；按需打开 Start + ThinkingStart。
|
||
/// Feed a thinking delta. Opens a Start + ThinkingStart as needed.
|
||
pub fn thinking_delta(&mut self, text: &str) -> Vec<StreamEvent> {
|
||
let mut out = Vec::new();
|
||
if let Some(ev) = self.ensure_started() {
|
||
out.push(ev);
|
||
}
|
||
if self.thinking_index.is_none() {
|
||
self.content.push(ContentBlock::Thinking(ThinkingContent {
|
||
thinking: String::new(),
|
||
signature: None,
|
||
redacted: false,
|
||
}));
|
||
self.args_by_index.push(None);
|
||
self.thinking_index = Some(self.content.len() - 1);
|
||
out.push(StreamEvent::ThinkingStart {
|
||
content_index: self.content.len() - 1,
|
||
partial: self.partial(),
|
||
});
|
||
}
|
||
let idx = self.thinking_index.expect("thinking block open");
|
||
if let Some(ContentBlock::Thinking(t)) = self.content.get_mut(idx) {
|
||
t.thinking.push_str(text);
|
||
}
|
||
out.push(StreamEvent::ThinkingDelta {
|
||
content_index: idx,
|
||
delta: text.to_string(),
|
||
partial: self.partial(),
|
||
});
|
||
out
|
||
}
|
||
|
||
/// 在 `index`（provider 局部索引，映射到一个新内容槽）处打开工具调用块，
|
||
/// 使用指定的调用 id（API 提供的真实 id）。
|
||
/// Open a tool-call block at `index` (provider-local index, mapped to a
|
||
/// new content slot), using the given call id (the API's real id).
|
||
///
|
||
/// 使用 API 的真实 id 至关重要：跨多轮工具调用时，若每轮都从 `call_0`
|
||
/// 重新生成，回放历史时 id 会冲突，导致 API 无法配对 function_call 与
|
||
/// function_call_output。空 id 时退回自动生成。
|
||
/// Using the API's real id matters: across multiple tool-using turns,
|
||
/// re-generating `call_0` every turn collides on replay, breaking the
|
||
/// function_call ↔ function_call_output pairing. Empty ids fall back to
|
||
/// auto-generated ones.
|
||
pub fn tool_call_start_with_id(
|
||
&mut self,
|
||
index: usize,
|
||
name: &str,
|
||
id: &str,
|
||
) -> Vec<StreamEvent> {
|
||
let mut out = Vec::new();
|
||
if let Some(ev) = self.ensure_started() {
|
||
out.push(ev);
|
||
}
|
||
let call_id = if id.is_empty() {
|
||
format!("call_{}", self.tool_call_pos.len())
|
||
} else {
|
||
id.to_string()
|
||
};
|
||
self.content.push(ContentBlock::ToolCall(ToolCall {
|
||
id: call_id,
|
||
name: name.to_string(),
|
||
arguments: focus_json::JsonValue::obj(),
|
||
}));
|
||
let content_index = self.content.len() - 1;
|
||
self.args_by_index.push(Some(String::new()));
|
||
self.tool_call_pos.push((index, content_index));
|
||
out.push(StreamEvent::ToolCallStart {
|
||
content_index,
|
||
partial: self.partial(),
|
||
});
|
||
out
|
||
}
|
||
|
||
/// 在 `index`（provider 局部索引，映射到一个新内容槽）处打开工具调用块，
|
||
/// 自动生成调用 id。
|
||
/// Open a tool-call block at `index` (provider-local index, mapped to a
|
||
/// new content slot) with an auto-generated call id.
|
||
pub fn tool_call_start(&mut self, index: usize, name: &str) -> Vec<StreamEvent> {
|
||
self.tool_call_start_with_id(index, name, "")
|
||
}
|
||
|
||
/// 向 `index` 处的工具调用追加部分 JSON 参数（按 provider 局部索引寻址）。
|
||
/// Append partial JSON arguments to the tool call at `index` (addressed by
|
||
/// the provider-local index).
|
||
pub fn tool_call_delta(&mut self, index: usize, args: &str) -> Vec<StreamEvent> {
|
||
let mut out = Vec::new();
|
||
let pos = self
|
||
.tool_call_pos
|
||
.iter()
|
||
.find(|(i, _)| *i == index)
|
||
.map(|(_, c)| *c)
|
||
.unwrap_or_else(|| self.content.len().saturating_sub(1));
|
||
if let Some(slot) = self.args_by_index.get_mut(pos).and_then(|s| s.as_mut()) {
|
||
slot.push_str(args);
|
||
}
|
||
out.push(StreamEvent::ToolCallDelta {
|
||
content_index: pos,
|
||
delta: args.to_string(),
|
||
partial: self.partial(),
|
||
});
|
||
out
|
||
}
|
||
|
||
/// 设置停止原因。
|
||
/// Set the stop reason.
|
||
pub fn set_stop(&mut self, reason: StopReason) {
|
||
self.stop_reason = reason;
|
||
}
|
||
|
||
/// 收尾：关闭所有未闭合的工具调用，返回最终的 Done 消息。
|
||
/// Finalize: close any open tool calls, return the final message.
|
||
pub fn finish(&mut self) -> Option<AssistantMessage> {
|
||
// 通过解析累积的 args 来关闭工具调用。
|
||
// Close tool calls by parsing their accumulated args.
|
||
for (_, content_index) in &self.tool_call_pos {
|
||
let args = self
|
||
.args_by_index
|
||
.get(*content_index)
|
||
.and_then(|s| s.as_ref())
|
||
.cloned()
|
||
.unwrap_or_default();
|
||
if let Some(ContentBlock::ToolCall(tc)) = self.content.get_mut(*content_index) {
|
||
tc.arguments = focus_json::parse(&args).unwrap_or(focus_json::JsonValue::obj());
|
||
}
|
||
}
|
||
let mut msg = self.partial();
|
||
msg.usage = self.usage.clone();
|
||
Some(msg)
|
||
}
|
||
|
||
/// 为所有未闭合的内容块（文本/思考/工具调用）发出 end 事件；之后调用方应再调用
|
||
/// [`finish`](Self::finish) 取得最终 Done。按顺序返回这些 end 事件。
|
||
/// Emit end events for any open content blocks (text/thinking/tool calls),
|
||
/// then callers should call [`finish`](Self::finish) for the final message.
|
||
/// Returns the end events in order.
|
||
pub fn finalize_events(&mut self) -> Vec<StreamEvent> {
|
||
let mut out = Vec::new();
|
||
if let Some(idx) = self.text_index.take() {
|
||
out.push(StreamEvent::TextEnd {
|
||
content_index: idx,
|
||
partial: self.partial(),
|
||
});
|
||
}
|
||
if let Some(idx) = self.thinking_index.take() {
|
||
out.push(StreamEvent::ThinkingEnd {
|
||
content_index: idx,
|
||
partial: self.partial(),
|
||
});
|
||
}
|
||
// 关闭每个工具调用块。
|
||
// Close each tool call block.
|
||
for (_, content_index) in &self.tool_call_pos {
|
||
let args = self
|
||
.args_by_index
|
||
.get(*content_index)
|
||
.and_then(|s| s.as_ref())
|
||
.cloned()
|
||
.unwrap_or_default();
|
||
if let Some(ContentBlock::ToolCall(tc)) = self.content.get_mut(*content_index) {
|
||
tc.arguments = focus_json::parse(&args).unwrap_or(focus_json::JsonValue::obj());
|
||
let finalized = tc.clone();
|
||
out.push(StreamEvent::ToolCallEnd {
|
||
content_index: *content_index,
|
||
tool_call: finalized,
|
||
partial: self.partial(),
|
||
});
|
||
}
|
||
}
|
||
out
|
||
}
|
||
}
|
||
|
||
/// provider 产出的、装箱后的事件流。provider 持有底层连接/驱动循环；
|
||
/// agent 通过 [`StreamIterator::next_event`] 拉取事件。
|
||
/// The boxed stream of events a provider produces. The provider owns the
|
||
/// underlying connection/drive loop; the agent pulls events via
|
||
/// [`StreamIterator::next_event`].
|
||
pub type StreamResult = CoreResult<Box<dyn StreamIterator + Send>>;
|
||
|
||
/// 流事件的同步迭代器。provider 驱动其完成；agent 循环拉取事件直到 `None`。
|
||
/// A sync iterator over stream events. The provider drives this to completion;
|
||
/// the agent loop pulls events until `None`.
|
||
pub trait StreamIterator {
|
||
fn next_event(&mut self) -> Option<StreamEvent>;
|
||
}
|
||
|
||
/// 每个 LLM 后端都实现的 trait。
|
||
/// The trait every LLM backend implements.
|
||
pub trait StreamProvider: Send + Sync + fmt::Debug {
|
||
/// 开始流式响应。返回的迭代器会持续产出事件，直到终止性的 `Done`/`Error`。
|
||
/// Begin streaming a response. The returned iterator yields events until a
|
||
/// terminal `Done`/`Error`.
|
||
fn stream(&self, request: &ProviderRequest) -> StreamResult;
|
||
}
|
||
|
||
/// 装箱的 provider 本身也是 provider，这样 [`Agent`](crate::agent::Agent)
|
||
/// 可以用 `Box::new(some_provider)` 构造。
|
||
/// A boxed provider is itself a provider, so the [`Agent`](crate::agent::Agent)
|
||
/// can be constructed with `Box::new(some_provider)`.
|
||
impl StreamProvider for Box<dyn StreamProvider> {
|
||
fn stream(&self, request: &ProviderRequest) -> StreamResult {
|
||
(**self).stream(request)
|
||
}
|
||
}
|