refactor: rename llm worker crate to engine
This commit is contained in:
@@ -0,0 +1,291 @@
|
||||
//! Closure-based event callback API
|
||||
//!
|
||||
//! Provides a closure-based alternative to implementing `Handler<K>` directly.
|
||||
//! Register callbacks on `Engine` via `on_text_block()`, `on_tool_use_block()`,
|
||||
//! `on_usage()`, etc.
|
||||
|
||||
use std::marker::PhantomData;
|
||||
|
||||
use crate::handler::{
|
||||
Handler, Kind, TextBlockEvent, TextBlockKind, ThinkingBlockEvent, ThinkingBlockKind,
|
||||
ToolUseBlockEvent, ToolUseBlockKind, ToolUseBlockStart,
|
||||
};
|
||||
use crate::tool::ToolCall;
|
||||
|
||||
// =============================================================================
|
||||
// TextBlock Closure Handler
|
||||
// =============================================================================
|
||||
|
||||
/// Callback scope for a text block.
|
||||
///
|
||||
/// Passed to the setup closure registered with `Engine::on_text_block()`.
|
||||
/// Register per-block callbacks via `on_delta()` and `on_stop()`.
|
||||
///
|
||||
/// # Examples
|
||||
///
|
||||
/// ```ignore
|
||||
/// engine.on_text_block(|block| {
|
||||
/// block.on_delta(|text| print!("{}", text));
|
||||
/// block.on_stop(|full_text| println!("\n--- {} chars ---", full_text.len()));
|
||||
/// });
|
||||
/// ```
|
||||
pub struct TextBlockScope {
|
||||
pub(crate) on_delta: Option<Box<dyn FnMut(&str) + Send + Sync>>,
|
||||
pub(crate) on_stop: Option<Box<dyn FnMut(&str) + Send + Sync>>,
|
||||
}
|
||||
|
||||
impl TextBlockScope {
|
||||
fn new() -> Self {
|
||||
Self {
|
||||
on_delta: None,
|
||||
on_stop: None,
|
||||
}
|
||||
}
|
||||
|
||||
/// Register a callback for each text delta (streaming fragment).
|
||||
pub fn on_delta(&mut self, f: impl FnMut(&str) + Send + Sync + 'static) {
|
||||
self.on_delta = Some(Box::new(f));
|
||||
}
|
||||
|
||||
/// Register a callback invoked when the block completes.
|
||||
///
|
||||
/// Receives the full accumulated text of the block.
|
||||
pub fn on_stop(&mut self, f: impl FnMut(&str) + Send + Sync + 'static) {
|
||||
self.on_stop = Some(Box::new(f));
|
||||
}
|
||||
}
|
||||
|
||||
/// Per-block state created by Timeline's scope lifecycle.
|
||||
#[derive(Default)]
|
||||
pub(crate) struct TextBlockClosureState {
|
||||
on_delta: Option<Box<dyn FnMut(&str) + Send + Sync>>,
|
||||
on_stop: Option<Box<dyn FnMut(&str) + Send + Sync>>,
|
||||
buffer: String,
|
||||
}
|
||||
|
||||
/// Closure-based `Handler<TextBlockKind>` adapter.
|
||||
pub(crate) struct ClosureTextBlockHandler {
|
||||
pub(crate) setup: Box<dyn FnMut(&mut TextBlockScope) + Send + Sync>,
|
||||
}
|
||||
|
||||
impl Handler<TextBlockKind> for ClosureTextBlockHandler {
|
||||
type Scope = TextBlockClosureState;
|
||||
|
||||
fn on_event(&mut self, scope: &mut Self::Scope, event: &TextBlockEvent) {
|
||||
match event {
|
||||
TextBlockEvent::Start(_) => {
|
||||
scope.buffer.clear();
|
||||
let mut builder = TextBlockScope::new();
|
||||
(self.setup)(&mut builder);
|
||||
scope.on_delta = builder.on_delta;
|
||||
scope.on_stop = builder.on_stop;
|
||||
}
|
||||
TextBlockEvent::Delta(text) => {
|
||||
scope.buffer.push_str(text);
|
||||
if let Some(f) = &mut scope.on_delta {
|
||||
f(text);
|
||||
}
|
||||
}
|
||||
TextBlockEvent::Stop(_) => {
|
||||
if let Some(f) = &mut scope.on_stop {
|
||||
f(&scope.buffer);
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// =============================================================================
|
||||
// ThinkingBlock Closure Handler
|
||||
// =============================================================================
|
||||
|
||||
/// Callback scope for a thinking block.
|
||||
///
|
||||
/// Mirrors `TextBlockScope`. Some providers (or some configurations)
|
||||
/// emit thinking metadata without plaintext deltas — in that case the
|
||||
/// block fires `Start` and `Stop` with no `Delta` in between, which is
|
||||
/// expected and not an error.
|
||||
pub struct ThinkingBlockScope {
|
||||
pub(crate) on_delta: Option<Box<dyn FnMut(&str) + Send + Sync>>,
|
||||
pub(crate) on_stop: Option<Box<dyn FnMut(&str) + Send + Sync>>,
|
||||
}
|
||||
|
||||
impl ThinkingBlockScope {
|
||||
fn new() -> Self {
|
||||
Self {
|
||||
on_delta: None,
|
||||
on_stop: None,
|
||||
}
|
||||
}
|
||||
|
||||
/// Register a callback for each thinking text delta (streaming fragment).
|
||||
pub fn on_delta(&mut self, f: impl FnMut(&str) + Send + Sync + 'static) {
|
||||
self.on_delta = Some(Box::new(f));
|
||||
}
|
||||
|
||||
/// Register a callback invoked when the block completes.
|
||||
///
|
||||
/// Receives the full accumulated thinking text. May be empty when
|
||||
/// the provider didn't emit any plaintext deltas.
|
||||
pub fn on_stop(&mut self, f: impl FnMut(&str) + Send + Sync + 'static) {
|
||||
self.on_stop = Some(Box::new(f));
|
||||
}
|
||||
}
|
||||
|
||||
#[derive(Default)]
|
||||
pub(crate) struct ThinkingBlockClosureState {
|
||||
on_delta: Option<Box<dyn FnMut(&str) + Send + Sync>>,
|
||||
on_stop: Option<Box<dyn FnMut(&str) + Send + Sync>>,
|
||||
buffer: String,
|
||||
}
|
||||
|
||||
pub(crate) struct ClosureThinkingBlockHandler {
|
||||
pub(crate) setup: Box<dyn FnMut(&mut ThinkingBlockScope) + Send + Sync>,
|
||||
}
|
||||
|
||||
impl Handler<ThinkingBlockKind> for ClosureThinkingBlockHandler {
|
||||
type Scope = ThinkingBlockClosureState;
|
||||
|
||||
fn on_event(&mut self, scope: &mut Self::Scope, event: &ThinkingBlockEvent) {
|
||||
match event {
|
||||
ThinkingBlockEvent::Start(_) => {
|
||||
scope.buffer.clear();
|
||||
let mut builder = ThinkingBlockScope::new();
|
||||
(self.setup)(&mut builder);
|
||||
scope.on_delta = builder.on_delta;
|
||||
scope.on_stop = builder.on_stop;
|
||||
}
|
||||
ThinkingBlockEvent::Delta(text) => {
|
||||
scope.buffer.push_str(text);
|
||||
if let Some(f) = &mut scope.on_delta {
|
||||
f(text);
|
||||
}
|
||||
}
|
||||
ThinkingBlockEvent::Stop(_) => {
|
||||
if let Some(f) = &mut scope.on_stop {
|
||||
f(&scope.buffer);
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// =============================================================================
|
||||
// ToolUseBlock Closure Handler
|
||||
// =============================================================================
|
||||
|
||||
/// Callback scope for a tool use block.
|
||||
///
|
||||
/// Passed to the setup closure registered with `Engine::on_tool_use_block()`.
|
||||
/// The setup closure also receives `&ToolUseBlockStart` with `id` and `name`.
|
||||
///
|
||||
/// # Examples
|
||||
///
|
||||
/// ```ignore
|
||||
/// engine.on_tool_use_block(|start, block| {
|
||||
/// println!("Tool: {} ({})", start.name, start.id);
|
||||
/// block.on_delta(|json| { /* streaming JSON fragment */ });
|
||||
/// block.on_stop(|call| println!("Done: {}", call.name));
|
||||
/// });
|
||||
/// ```
|
||||
pub struct ToolUseBlockScope {
|
||||
pub(crate) on_delta: Option<Box<dyn FnMut(&str) + Send + Sync>>,
|
||||
pub(crate) on_stop: Option<Box<dyn FnMut(&ToolCall) + Send + Sync>>,
|
||||
}
|
||||
|
||||
impl ToolUseBlockScope {
|
||||
fn new() -> Self {
|
||||
Self {
|
||||
on_delta: None,
|
||||
on_stop: None,
|
||||
}
|
||||
}
|
||||
|
||||
/// Register a callback for each JSON input delta (streaming fragment).
|
||||
pub fn on_delta(&mut self, f: impl FnMut(&str) + Send + Sync + 'static) {
|
||||
self.on_delta = Some(Box::new(f));
|
||||
}
|
||||
|
||||
/// Register a callback invoked when the block completes.
|
||||
///
|
||||
/// Receives the fully assembled `ToolCall` with parsed JSON input.
|
||||
pub fn on_stop(&mut self, f: impl FnMut(&ToolCall) + Send + Sync + 'static) {
|
||||
self.on_stop = Some(Box::new(f));
|
||||
}
|
||||
}
|
||||
|
||||
/// Per-block state for tool use closure handler.
|
||||
#[derive(Default)]
|
||||
pub(crate) struct ToolUseBlockClosureState {
|
||||
on_delta: Option<Box<dyn FnMut(&str) + Send + Sync>>,
|
||||
on_stop: Option<Box<dyn FnMut(&ToolCall) + Send + Sync>>,
|
||||
id: String,
|
||||
name: String,
|
||||
input_json: String,
|
||||
}
|
||||
|
||||
/// Closure-based `Handler<ToolUseBlockKind>` adapter.
|
||||
pub(crate) struct ClosureToolUseBlockHandler {
|
||||
pub(crate) setup: Box<dyn FnMut(&ToolUseBlockStart, &mut ToolUseBlockScope) + Send + Sync>,
|
||||
}
|
||||
|
||||
impl Handler<ToolUseBlockKind> for ClosureToolUseBlockHandler {
|
||||
type Scope = ToolUseBlockClosureState;
|
||||
|
||||
fn on_event(&mut self, scope: &mut Self::Scope, event: &ToolUseBlockEvent) {
|
||||
match event {
|
||||
ToolUseBlockEvent::Start(start) => {
|
||||
scope.id = start.id.clone();
|
||||
scope.name = start.name.clone();
|
||||
scope.input_json.clear();
|
||||
let mut builder = ToolUseBlockScope::new();
|
||||
(self.setup)(start, &mut builder);
|
||||
scope.on_delta = builder.on_delta;
|
||||
scope.on_stop = builder.on_stop;
|
||||
}
|
||||
ToolUseBlockEvent::InputJsonDelta(json) => {
|
||||
scope.input_json.push_str(json);
|
||||
if let Some(f) = &mut scope.on_delta {
|
||||
f(json);
|
||||
}
|
||||
}
|
||||
ToolUseBlockEvent::Stop(_) => {
|
||||
let input: serde_json::Value =
|
||||
serde_json::from_str(&scope.input_json).unwrap_or_default();
|
||||
let tool_call = ToolCall {
|
||||
id: std::mem::take(&mut scope.id),
|
||||
name: std::mem::take(&mut scope.name),
|
||||
input,
|
||||
};
|
||||
if let Some(f) = &mut scope.on_stop {
|
||||
f(&tool_call);
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// =============================================================================
|
||||
// Generic Meta Event Closure Handler
|
||||
// =============================================================================
|
||||
|
||||
/// Closure-based `Handler<K>` adapter for meta events (Usage, Status, Error).
|
||||
pub(crate) struct ClosureMetaHandler<F, K>
|
||||
where
|
||||
K: Kind,
|
||||
{
|
||||
pub(crate) callback: F,
|
||||
pub(crate) _kind: PhantomData<K>,
|
||||
}
|
||||
|
||||
impl<F, K> Handler<K> for ClosureMetaHandler<F, K>
|
||||
where
|
||||
F: FnMut(&K::Event) + Send + Sync,
|
||||
K: Kind,
|
||||
{
|
||||
type Scope = ();
|
||||
|
||||
fn on_event(&mut self, _scope: &mut (), event: &K::Event) {
|
||||
(self.callback)(event);
|
||||
}
|
||||
}
|
||||
File diff suppressed because it is too large
Load Diff
@@ -0,0 +1,5 @@
|
||||
//! Public event types for Engine layer
|
||||
//!
|
||||
//! Re-exports from the canonical event definitions in llm_client.
|
||||
|
||||
pub use crate::llm_client::event::*;
|
||||
@@ -0,0 +1,175 @@
|
||||
//! Handler/Kind Types
|
||||
//!
|
||||
//! Traits for processing events in the Timeline layer.
|
||||
//! By implementing custom handlers and registering them with Timeline,
|
||||
//! you can receive stream events.
|
||||
|
||||
use crate::timeline::event::*;
|
||||
|
||||
// =============================================================================
|
||||
// Kind Trait
|
||||
// =============================================================================
|
||||
|
||||
/// Marker trait defining event types
|
||||
///
|
||||
/// Each Kind specifies its corresponding event type.
|
||||
/// Handlers are implemented for this Kind, and multiple Handlers
|
||||
/// with different Scope types can be registered for the same Kind.
|
||||
pub trait Kind {
|
||||
/// Event type corresponding to this Kind
|
||||
type Event;
|
||||
}
|
||||
|
||||
// =============================================================================
|
||||
// Handler Trait
|
||||
// =============================================================================
|
||||
|
||||
/// Handler trait for processing events
|
||||
///
|
||||
/// Defines event processing for a specific `Kind`.
|
||||
/// `Scope` is state held during the block's lifecycle.
|
||||
///
|
||||
/// # Examples
|
||||
///
|
||||
/// ```ignore
|
||||
/// use llm_engine::timeline::{Handler, TextBlockEvent, TextBlockKind};
|
||||
///
|
||||
/// struct TextCollector {
|
||||
/// texts: Vec<String>,
|
||||
/// }
|
||||
///
|
||||
/// impl Handler<TextBlockKind> for TextCollector {
|
||||
/// type Scope = String; // Buffer per block
|
||||
///
|
||||
/// fn on_event(&mut self, buffer: &mut String, event: &TextBlockEvent) {
|
||||
/// match event {
|
||||
/// TextBlockEvent::Delta(text) => buffer.push_str(text),
|
||||
/// TextBlockEvent::Stop(_) => {
|
||||
/// self.texts.push(std::mem::take(buffer));
|
||||
/// }
|
||||
/// _ => {}
|
||||
/// }
|
||||
/// }
|
||||
/// }
|
||||
/// ```
|
||||
pub trait Handler<K: Kind> {
|
||||
/// Handler-specific scope type
|
||||
///
|
||||
/// Generated with `Default::default()` at block start,
|
||||
/// and destroyed at block end.
|
||||
type Scope: Default;
|
||||
|
||||
/// Process the event
|
||||
fn on_event(&mut self, scope: &mut Self::Scope, event: &K::Event);
|
||||
}
|
||||
|
||||
// =============================================================================
|
||||
// Meta Kind Definitions
|
||||
// =============================================================================
|
||||
|
||||
/// Usage Kind - for usage events
|
||||
pub struct UsageKind;
|
||||
impl Kind for UsageKind {
|
||||
type Event = UsageEvent;
|
||||
}
|
||||
|
||||
/// Ping Kind - for ping events
|
||||
pub struct PingKind;
|
||||
impl Kind for PingKind {
|
||||
type Event = PingEvent;
|
||||
}
|
||||
|
||||
/// Status Kind - for status events
|
||||
pub struct StatusKind;
|
||||
impl Kind for StatusKind {
|
||||
type Event = StatusEvent;
|
||||
}
|
||||
|
||||
/// Error Kind - for error events
|
||||
pub struct ErrorKind;
|
||||
impl Kind for ErrorKind {
|
||||
type Event = ErrorEvent;
|
||||
}
|
||||
|
||||
// =============================================================================
|
||||
// Block Kind Definitions
|
||||
// =============================================================================
|
||||
|
||||
/// TextBlock Kind - for text blocks
|
||||
pub struct TextBlockKind;
|
||||
impl Kind for TextBlockKind {
|
||||
type Event = TextBlockEvent;
|
||||
}
|
||||
|
||||
/// Text block events
|
||||
#[derive(Debug, Clone, PartialEq)]
|
||||
pub enum TextBlockEvent {
|
||||
Start(TextBlockStart),
|
||||
Delta(String),
|
||||
Stop(TextBlockStop),
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, PartialEq)]
|
||||
pub struct TextBlockStart {
|
||||
pub index: usize,
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, PartialEq)]
|
||||
pub struct TextBlockStop {
|
||||
pub index: usize,
|
||||
pub stop_reason: Option<StopReason>,
|
||||
}
|
||||
|
||||
/// ThinkingBlock Kind - for thinking blocks
|
||||
pub struct ThinkingBlockKind;
|
||||
impl Kind for ThinkingBlockKind {
|
||||
type Event = ThinkingBlockEvent;
|
||||
}
|
||||
|
||||
/// Thinking block events
|
||||
#[derive(Debug, Clone, PartialEq)]
|
||||
pub enum ThinkingBlockEvent {
|
||||
Start(ThinkingBlockStart),
|
||||
Delta(String),
|
||||
Stop(ThinkingBlockStop),
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, PartialEq)]
|
||||
pub struct ThinkingBlockStart {
|
||||
pub index: usize,
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, PartialEq)]
|
||||
pub struct ThinkingBlockStop {
|
||||
pub index: usize,
|
||||
pub reasoning: Option<ReasoningBlockData>,
|
||||
}
|
||||
|
||||
/// ToolUseBlock Kind - for tool use blocks
|
||||
pub struct ToolUseBlockKind;
|
||||
impl Kind for ToolUseBlockKind {
|
||||
type Event = ToolUseBlockEvent;
|
||||
}
|
||||
|
||||
/// Tool use block events
|
||||
#[derive(Debug, Clone, PartialEq)]
|
||||
pub enum ToolUseBlockEvent {
|
||||
Start(ToolUseBlockStart),
|
||||
/// JSON substring of tool arguments
|
||||
InputJsonDelta(String),
|
||||
Stop(ToolUseBlockStop),
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, PartialEq)]
|
||||
pub struct ToolUseBlockStart {
|
||||
pub index: usize,
|
||||
pub id: String,
|
||||
pub name: String,
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, PartialEq)]
|
||||
pub struct ToolUseBlockStop {
|
||||
pub index: usize,
|
||||
pub id: String,
|
||||
pub name: String,
|
||||
}
|
||||
@@ -0,0 +1,202 @@
|
||||
//! Interceptor - control flow delegation for the Engine execution loop
|
||||
//!
|
||||
//! Defines the [`Interceptor`] trait that upper layers (e.g. Pod) implement
|
||||
//! to inject orchestration decisions (approval, skip, pause, abort)
|
||||
//! into the Engine's turn loop without the Engine knowing about
|
||||
//! higher-level concepts.
|
||||
|
||||
use std::sync::Arc;
|
||||
|
||||
use async_trait::async_trait;
|
||||
|
||||
use crate::Item;
|
||||
use crate::tool::{Tool, ToolCall, ToolExecutionContext, ToolMeta, ToolResult};
|
||||
|
||||
// =============================================================================
|
||||
// Action Enums
|
||||
// =============================================================================
|
||||
|
||||
/// Action after prompt submission.
|
||||
#[derive(Debug, Clone, PartialEq)]
|
||||
pub enum PromptAction {
|
||||
/// Proceed normally.
|
||||
Continue,
|
||||
/// Cancel with a reason.
|
||||
Cancel(String),
|
||||
/// Proceed, and append these items to history right after the user
|
||||
/// message. Mirrors [`TurnEndAction::ContinueWithMessages`] for the
|
||||
/// submit edge: lets the upper layer attach resolver-produced
|
||||
/// system messages (e.g. `@<path>` file content) so they sit
|
||||
/// adjacent to the user message that referenced them.
|
||||
ContinueWith(Vec<Item>),
|
||||
}
|
||||
|
||||
/// Action before an LLM request.
|
||||
#[derive(Debug, Clone, PartialEq)]
|
||||
pub enum PreRequestAction {
|
||||
/// Proceed normally.
|
||||
Continue,
|
||||
/// Proceed after appending these items to durable engine history.
|
||||
///
|
||||
/// This is for upper-layer budget/status nudges that the model may react
|
||||
/// to: the items are committed before the request so later turns can see
|
||||
/// why the engine changed course.
|
||||
ContinueWith(Vec<Item>),
|
||||
/// Yield after appending these items to durable engine history.
|
||||
///
|
||||
/// This is for host-mediated pre-request appends that must be visible to
|
||||
/// usage accounting and compaction checks before the current LLM request is
|
||||
/// allowed to proceed.
|
||||
YieldWith(Vec<Item>),
|
||||
/// Cancel with a reason (treated as an error).
|
||||
Cancel(String),
|
||||
/// Yield control to the caller for external processing.
|
||||
///
|
||||
/// The Engine exits the turn loop cleanly with `EngineResult::Yielded`.
|
||||
/// The caller is expected to resume execution later.
|
||||
Yield,
|
||||
}
|
||||
|
||||
/// Action before a tool call.
|
||||
#[derive(Debug, Clone, PartialEq, Eq)]
|
||||
pub enum PreToolAction {
|
||||
/// Proceed with execution.
|
||||
Continue,
|
||||
/// Skip this tool call (do not execute).
|
||||
Skip,
|
||||
/// Do not execute the tool call; commit this synthetic result instead.
|
||||
///
|
||||
/// This preserves provider-visible `tool_use` / `tool_result` pairing
|
||||
/// without aborting the whole turn.
|
||||
SyntheticResult(ToolResult),
|
||||
/// Abort the entire run.
|
||||
Abort(String),
|
||||
/// Pause execution (can be resumed later).
|
||||
Pause,
|
||||
}
|
||||
|
||||
/// Action after a tool call.
|
||||
#[derive(Debug, Clone, PartialEq, Eq)]
|
||||
pub enum PostToolAction {
|
||||
/// Proceed normally.
|
||||
Continue,
|
||||
/// Abort the entire run.
|
||||
Abort(String),
|
||||
}
|
||||
|
||||
/// Action at the end of a turn (when LLM produces no tool calls).
|
||||
#[derive(Debug, Clone)]
|
||||
pub enum TurnEndAction {
|
||||
/// Turn is finished, return to caller.
|
||||
Finish,
|
||||
/// Continue with additional messages injected into history.
|
||||
ContinueWithMessages(Vec<Item>),
|
||||
/// Pause execution (can be resumed later).
|
||||
Pause,
|
||||
}
|
||||
|
||||
// =============================================================================
|
||||
// Context Types
|
||||
// =============================================================================
|
||||
|
||||
/// Context for pre-tool-call decisions.
|
||||
pub struct ToolCallInfo {
|
||||
/// Tool call information (modifiable).
|
||||
pub call: ToolCall,
|
||||
/// Tool meta information.
|
||||
pub meta: ToolMeta,
|
||||
/// Tool instance (for state access).
|
||||
pub tool: Arc<dyn Tool>,
|
||||
/// Response-local execution context for this call.
|
||||
pub context: ToolExecutionContext,
|
||||
}
|
||||
|
||||
/// Context for post-tool-call decisions.
|
||||
pub struct ToolResultInfo {
|
||||
/// Original tool call.
|
||||
pub call: ToolCall,
|
||||
/// Tool execution result (modifiable).
|
||||
pub result: ToolResult,
|
||||
/// Tool meta information.
|
||||
pub meta: ToolMeta,
|
||||
/// Tool instance (for state access).
|
||||
pub tool: Arc<dyn Tool>,
|
||||
/// Response-local execution context for this call.
|
||||
pub context: ToolExecutionContext,
|
||||
}
|
||||
|
||||
// =============================================================================
|
||||
// Interceptor Trait
|
||||
// =============================================================================
|
||||
|
||||
/// Intercepts the Engine execution loop at key decision points.
|
||||
///
|
||||
/// All methods have default implementations that let the Engine
|
||||
/// proceed without intervention. Upper layers (e.g. Pod) provide
|
||||
/// richer implementations for approval flows, permission checks, etc.
|
||||
#[async_trait]
|
||||
pub trait Interceptor: Send + Sync {
|
||||
/// Called after receiving user input, before adding to history.
|
||||
async fn on_prompt_submit(&self, _item: &mut Item) -> PromptAction {
|
||||
PromptAction::Continue
|
||||
}
|
||||
|
||||
/// Items that should be **committed to `engine.history`** just
|
||||
/// before the next LLM request. Returned items are `extend`ed into
|
||||
/// the persistent history (and therefore picked up by the per-turn
|
||||
/// clone that backs the LLM request, plus the usual
|
||||
/// history-persistence path).
|
||||
///
|
||||
/// Use this for inputs that arrive from outside the LLM and need
|
||||
/// to be reflected in the on-disk history — notifications,
|
||||
/// cross-Pod events, system reminders. Do **not** use
|
||||
/// [`Self::pre_llm_request`] for that purpose: it mutates a
|
||||
/// per-request clone, so any committed assistant response that
|
||||
/// reacts to the injection would have no visible trigger on the
|
||||
/// next turn (or after resume / compaction).
|
||||
///
|
||||
/// `pre_llm_request` remains the right place for purely
|
||||
/// reproducible per-request transformations (pruning, content
|
||||
/// trimming, cache anchors) that depend only on the existing
|
||||
/// history.
|
||||
async fn pending_history_appends(&self) -> Vec<Item> {
|
||||
Vec::new()
|
||||
}
|
||||
|
||||
/// Called before each LLM request. The context starts as a clone
|
||||
/// of `engine.history` (after `pending_history_appends` and the
|
||||
/// Engine's own prune projection have been applied).
|
||||
///
|
||||
/// Direct mutations to `context` remain request-local and are not persisted.
|
||||
/// If an interceptor derives a human/model-visible nudge from the current
|
||||
/// request context, return [`PreRequestAction::ContinueWith`] so the Engine
|
||||
/// commits it to history before the request is sent.
|
||||
async fn pre_llm_request(&self, _context: &mut Vec<Item>) -> PreRequestAction {
|
||||
PreRequestAction::Continue
|
||||
}
|
||||
|
||||
/// Called before each tool is executed.
|
||||
async fn pre_tool_call(&self, _info: &mut ToolCallInfo) -> PreToolAction {
|
||||
PreToolAction::Continue
|
||||
}
|
||||
|
||||
/// Called after each tool completes.
|
||||
async fn post_tool_call(&self, _info: &mut ToolResultInfo) -> PostToolAction {
|
||||
PostToolAction::Continue
|
||||
}
|
||||
|
||||
/// Called when a turn ends with no tool calls.
|
||||
async fn on_turn_end(&self, _history: &[Item]) -> TurnEndAction {
|
||||
TurnEndAction::Finish
|
||||
}
|
||||
|
||||
/// Called when execution is interrupted (abort or cancel).
|
||||
async fn on_abort(&self, _reason: &str) {}
|
||||
}
|
||||
|
||||
/// Default interceptor: no intervention. Engine proceeds through the loop
|
||||
/// without any external control flow decisions.
|
||||
pub(crate) struct DefaultInterceptor;
|
||||
|
||||
#[async_trait]
|
||||
impl Interceptor for DefaultInterceptor {}
|
||||
@@ -0,0 +1,65 @@
|
||||
//! llm-engine - LLM Engine Library
|
||||
//!
|
||||
//! Provides components for managing interactions with LLMs.
|
||||
//!
|
||||
//! # Main Components
|
||||
//!
|
||||
//! - [`Engine`] - Central component for managing LLM interactions
|
||||
//! - [`tool::Tool`] - Tools that can be invoked by the LLM
|
||||
//! - [`interceptor::Interceptor`] - Control-flow delegation for the execution loop
|
||||
//! - Closure-based event callbacks via `Engine::on_text_block()`, `on_tool_use_block()`, etc.
|
||||
//!
|
||||
//! # Quick Start
|
||||
//!
|
||||
//! ```ignore
|
||||
//! use llm_engine::{Engine, Item};
|
||||
//!
|
||||
//! // Create a Engine
|
||||
//! let mut engine = Engine::new(client)
|
||||
//! .system_prompt("You are a helpful assistant.");
|
||||
//!
|
||||
//! // Register tools (optional)
|
||||
//! // engine.register_tool(my_tool_definition)?;
|
||||
//!
|
||||
//! // Run the interaction
|
||||
//! let history = engine.run("Hello!").await?;
|
||||
//! ```
|
||||
//!
|
||||
//! # Cache Protection
|
||||
//!
|
||||
//! `run()` automatically locks the cache. To edit state between turns,
|
||||
//! call `unlock_cache()` first; the next `run()` re-locks automatically.
|
||||
//!
|
||||
//! ```ignore
|
||||
//! engine.run("user input").await?;
|
||||
//! engine.unlock_cache();
|
||||
//! engine.set_system_prompt("new prompt");
|
||||
//! engine.run("next input").await?;
|
||||
//! ```
|
||||
|
||||
mod engine;
|
||||
mod handler;
|
||||
mod message;
|
||||
|
||||
pub(crate) mod callback;
|
||||
pub mod event;
|
||||
pub mod interceptor;
|
||||
pub mod llm_client;
|
||||
pub mod prune;
|
||||
pub mod state;
|
||||
pub mod timeline;
|
||||
pub mod token_counter;
|
||||
pub mod tool;
|
||||
pub mod tool_server;
|
||||
pub mod usage_record;
|
||||
|
||||
pub use callback::{TextBlockScope, ThinkingBlockScope, ToolUseBlockScope};
|
||||
pub use engine::{
|
||||
Engine, EngineConfig, EngineError, EngineResult, EngineRunOutput, LlmRetryNotice,
|
||||
ToolRegistryError,
|
||||
};
|
||||
pub use handler::ToolUseBlockStart;
|
||||
pub use interceptor::Interceptor;
|
||||
pub use message::{ContentPart, Item, Message, Role};
|
||||
pub use tool::{ToolCall, ToolExecutionContext, ToolOutputLimits, ToolResult};
|
||||
pub use usage_record::UsageRecord;
|
||||
@@ -0,0 +1,57 @@
|
||||
//! `Scheme` 実装と通信層が要求する認証要件、および動的認証プロバイダ。
|
||||
//!
|
||||
//! マニフェスト側の型（`ModelConfig` / `SchemeKind` / `AuthRef`）は
|
||||
//! `crates/manifest` に置き、llm-engine はそれを知らずに済む。
|
||||
//! `AuthRequirement` は scheme が宣言する「この scheme はどんな認証を
|
||||
//! 期待するか」のランタイム記述で、manifest 側の `AuthRef` との
|
||||
//! 照合（`AuthRef → ResolvedAuth` 変換の適否）は `crates/provider`
|
||||
//! で行う。
|
||||
//!
|
||||
//! Codex OAuth のようにリクエスト毎にトークンが変わり得る認証は
|
||||
//! [`AuthProvider`] trait を `crates/provider` 側で実装し、
|
||||
//! [`super::transport::ResolvedAuth::Custom`] 経由で transport に渡す。
|
||||
|
||||
use async_trait::async_trait;
|
||||
use reqwest::header::{HeaderName, HeaderValue};
|
||||
|
||||
use super::error::ClientError;
|
||||
|
||||
/// `Scheme::required_auth()` が返す認証要件。
|
||||
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
|
||||
pub enum AuthRequirement {
|
||||
/// 認証を行わない（Ollama など）
|
||||
None,
|
||||
/// `Authorization: Bearer <token>` ヘッダ（token は API key 相当）
|
||||
Bearer,
|
||||
/// `x-api-key: <token>` ヘッダ（Anthropic 形式）
|
||||
XApiKey,
|
||||
/// クエリパラメータ `?<name>=<token>`（Gemini 形式）
|
||||
QueryParam { name: &'static str },
|
||||
/// 複合ヘッダ（Codex OAuth 等、`crates/provider` 側で解決）
|
||||
Custom,
|
||||
}
|
||||
|
||||
/// リクエスト毎に認証ヘッダを動的に組み立てるプロバイダ。
|
||||
///
|
||||
/// Codex OAuth のように access_token が refresh で更新されたり、
|
||||
/// `ChatGPT-Account-Id` / `X-OpenAI-Fedramp` のような複数ヘッダを
|
||||
/// 同時に注入する必要があるケースで使う。実体は `crates/provider`
|
||||
/// 側に置き、llm-engine は trait を知るだけ。
|
||||
///
|
||||
/// 返したヘッダはそのまま `HeaderMap` に挿入される。`Authorization`
|
||||
/// 含む scheme 既定の認証ヘッダは送出されないので、必要なら
|
||||
/// 実装側でセットすること。
|
||||
#[async_trait]
|
||||
pub trait AuthProvider: Send + Sync + std::fmt::Debug {
|
||||
/// 1 リクエスト分の認証ヘッダを返す。refresh が必要なら内部で行う。
|
||||
async fn headers(&self) -> Result<Vec<(HeaderName, HeaderValue)>, ClientError>;
|
||||
|
||||
/// ChatGPT Codex backend 向けの複合認証かどうか。
|
||||
///
|
||||
/// transport は provider crate の具象型を知らないため、この hook だけで
|
||||
/// Codex CLI 互換の wire behavior（conversation header / request compression 等）
|
||||
/// を切り替える。
|
||||
fn is_codex_backend(&self) -> bool {
|
||||
false
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,169 @@
|
||||
//! モデル能力メタデータ
|
||||
//!
|
||||
//! `ModelCapability` はモデルが持つ機能差を表現する。scheme は同じでも
|
||||
//! モデルごとに reasoning 可否や prompt caching 方式が違うため、scheme
|
||||
//! から分離して保持する。
|
||||
//!
|
||||
//! 値の供給経路は 2 通り:
|
||||
//! 1. scheme 実装側の `model_id → ModelCapability` 静的テーブル（既知モデル）
|
||||
//! 2. `ModelConfig::capability` での明示 override（未知モデル、または上書き）
|
||||
|
||||
use serde::{Deserialize, Deserializer, Serialize, Serializer};
|
||||
|
||||
/// モデル能力メタデータ
|
||||
#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
|
||||
pub struct ModelCapability {
|
||||
pub tool_calling: ToolCallingSupport,
|
||||
pub structured_output: StructuredOutput,
|
||||
#[serde(default)]
|
||||
pub reasoning: Option<ReasoningSupport>,
|
||||
#[serde(default)]
|
||||
pub vision: bool,
|
||||
pub prompt_caching: CacheStrategy,
|
||||
}
|
||||
|
||||
impl ModelCapability {
|
||||
/// 何もサポートしない安全側デフォルト。未知モデルのフォールバック用。
|
||||
pub const fn minimal() -> Self {
|
||||
Self {
|
||||
tool_calling: ToolCallingSupport::None,
|
||||
structured_output: StructuredOutput::None,
|
||||
reasoning: None,
|
||||
vision: false,
|
||||
prompt_caching: CacheStrategy::Auto,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// ツール呼び出しサポート
|
||||
#[derive(Debug, Clone, Copy, Serialize, Deserialize, PartialEq, Eq)]
|
||||
#[serde(rename_all = "snake_case")]
|
||||
pub enum ToolCallingSupport {
|
||||
/// 非サポート
|
||||
None,
|
||||
/// 1 回のレスポンスで 1 ツールのみ
|
||||
Sequential,
|
||||
/// 1 回のレスポンスで複数ツール並行
|
||||
Parallel,
|
||||
}
|
||||
|
||||
/// Structured output サポート
|
||||
#[derive(Debug, Clone, Copy, Serialize, Deserialize, PartialEq, Eq)]
|
||||
#[serde(rename_all = "snake_case")]
|
||||
pub enum StructuredOutput {
|
||||
None,
|
||||
/// `json_object` モード（スキーマなし JSON 強制）
|
||||
JsonObject,
|
||||
/// JSON Schema 指定で構造化出力
|
||||
JsonSchema,
|
||||
}
|
||||
|
||||
/// Reasoning（extended thinking）サポート
|
||||
#[derive(Debug, Clone, Copy, Serialize, Deserialize, PartialEq, Eq)]
|
||||
#[serde(rename_all = "snake_case")]
|
||||
pub enum ReasoningSupport {
|
||||
/// OpenAI 形式: `reasoning.effort` (low/medium/high)
|
||||
Effort,
|
||||
/// Anthropic 形式: `thinking.budget_tokens`
|
||||
BudgetTokens,
|
||||
/// 両対応（内部では共通 `ReasoningControl` として扱い、各 scheme で投影）
|
||||
Both,
|
||||
}
|
||||
|
||||
/// Prompt caching 戦略
|
||||
#[derive(Debug, Clone, Copy, Serialize, Deserialize, PartialEq, Eq)]
|
||||
#[serde(tag = "kind", rename_all = "snake_case")]
|
||||
pub enum CacheStrategy {
|
||||
/// Anthropic: `cache_control` マーカーを明示挿入
|
||||
Explicit { max_breakpoints: u8 },
|
||||
/// それ以外: サーバ側自動 prefix、または未サポート
|
||||
Auto,
|
||||
}
|
||||
|
||||
/// Reasoning 制御（共通型、scheme 側で各社形式に投影）。
|
||||
///
|
||||
/// 文字列は provider-native な effort label、数値は provider-native な
|
||||
/// thinking budget token として扱う。どちらか一方だけを型で表現する。
|
||||
#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)]
|
||||
#[serde(untagged)]
|
||||
pub enum ReasoningControl {
|
||||
Effort(ReasoningEffort),
|
||||
BudgetTokens(i32),
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, PartialEq, Eq)]
|
||||
pub enum ReasoningEffort {
|
||||
Minimal,
|
||||
Low,
|
||||
Medium,
|
||||
High,
|
||||
XHigh,
|
||||
Other(String),
|
||||
}
|
||||
|
||||
impl ReasoningEffort {
|
||||
pub fn as_str(&self) -> &str {
|
||||
match self {
|
||||
Self::Minimal => "minimal",
|
||||
Self::Low => "low",
|
||||
Self::Medium => "medium",
|
||||
Self::High => "high",
|
||||
Self::XHigh => "xhigh",
|
||||
Self::Other(label) => label.as_str(),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
impl From<String> for ReasoningEffort {
|
||||
fn from(value: String) -> Self {
|
||||
match value.as_str() {
|
||||
"minimal" => Self::Minimal,
|
||||
"low" => Self::Low,
|
||||
"medium" => Self::Medium,
|
||||
"high" => Self::High,
|
||||
"xhigh" => Self::XHigh,
|
||||
_ => Self::Other(value),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
impl Serialize for ReasoningEffort {
|
||||
fn serialize<S>(&self, serializer: S) -> Result<S::Ok, S::Error>
|
||||
where
|
||||
S: Serializer,
|
||||
{
|
||||
serializer.serialize_str(self.as_str())
|
||||
}
|
||||
}
|
||||
|
||||
impl<'de> Deserialize<'de> for ReasoningEffort {
|
||||
fn deserialize<D>(deserializer: D) -> Result<Self, D::Error>
|
||||
where
|
||||
D: Deserializer<'de>,
|
||||
{
|
||||
String::deserialize(deserializer).map(Self::from)
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::{ReasoningControl, ReasoningEffort};
|
||||
|
||||
#[test]
|
||||
fn reasoning_control_deserializes_effort_labels() {
|
||||
let known: ReasoningControl = serde_json::from_str(r#""xhigh""#).unwrap();
|
||||
assert_eq!(known, ReasoningControl::Effort(ReasoningEffort::XHigh));
|
||||
|
||||
let unknown: ReasoningControl = serde_json::from_str(r#""provider-native""#).unwrap();
|
||||
assert_eq!(
|
||||
unknown,
|
||||
ReasoningControl::Effort(ReasoningEffort::Other("provider-native".into()))
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn reasoning_control_deserializes_signed_budget() {
|
||||
let dynamic: ReasoningControl = serde_json::from_str("-1").unwrap();
|
||||
assert_eq!(dynamic, ReasoningControl::BudgetTokens(-1));
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,98 @@
|
||||
//! LLMクライアント共通trait定義
|
||||
|
||||
use std::pin::Pin;
|
||||
|
||||
use crate::llm_client::{ClientError, Request, RequestConfig, event::Event};
|
||||
use async_trait::async_trait;
|
||||
use futures::Stream;
|
||||
|
||||
/// 設定に関する警告
|
||||
///
|
||||
/// プロバイダがサポートしていない設定を使用した場合に返される。
|
||||
#[derive(Debug, Clone)]
|
||||
pub struct ConfigWarning {
|
||||
/// 設定オプション名
|
||||
pub option_name: &'static str,
|
||||
/// 警告メッセージ
|
||||
pub message: String,
|
||||
}
|
||||
|
||||
impl ConfigWarning {
|
||||
/// 新しい警告を作成
|
||||
pub fn unsupported(option_name: &'static str, provider_name: &str) -> Self {
|
||||
Self {
|
||||
option_name,
|
||||
message: format!(
|
||||
"'{}' is not supported by {} and will be ignored",
|
||||
option_name, provider_name
|
||||
),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
impl std::fmt::Display for ConfigWarning {
|
||||
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
|
||||
write!(f, "{}: {}", self.option_name, self.message)
|
||||
}
|
||||
}
|
||||
|
||||
pub type ResponseStream = Pin<Box<dyn Stream<Item = Result<Event, ClientError>> + Send>>;
|
||||
|
||||
/// LLMクライアントのtrait
|
||||
///
|
||||
/// 各プロバイダはこのtraitを実装し、統一されたインターフェースを提供する。
|
||||
#[async_trait]
|
||||
pub trait LlmClient: Send + Sync {
|
||||
/// ストリーミングリクエストを送信し、Eventストリームを返す
|
||||
///
|
||||
/// # Arguments
|
||||
/// * `request` - リクエスト情報
|
||||
///
|
||||
/// # Returns
|
||||
/// * `Ok(Stream)` - イベントストリーム
|
||||
/// * `Err(ClientError)` - エラー
|
||||
async fn stream(&self, request: Request) -> Result<ResponseStream, ClientError>;
|
||||
|
||||
/// Clone this client into a new `Box<dyn LlmClient>`.
|
||||
///
|
||||
/// Used when a second client instance is needed (e.g. for context
|
||||
/// compaction) without access to the original construction parameters.
|
||||
fn clone_boxed(&self) -> Box<dyn LlmClient>;
|
||||
|
||||
/// 設定をバリデーションし、未サポートの設定があれば警告を返す
|
||||
///
|
||||
/// # Arguments
|
||||
/// * `config` - バリデーション対象の設定
|
||||
///
|
||||
/// # Returns
|
||||
/// サポートされていない設定に対する警告のリスト
|
||||
fn validate_config(&self, config: &RequestConfig) -> Vec<ConfigWarning> {
|
||||
// デフォルト実装: 全ての設定をサポート
|
||||
let _ = config;
|
||||
Vec::new()
|
||||
}
|
||||
}
|
||||
|
||||
impl Clone for Box<dyn LlmClient> {
|
||||
fn clone(&self) -> Self {
|
||||
self.clone_boxed()
|
||||
}
|
||||
}
|
||||
|
||||
/// `Box<dyn LlmClient>` に対する `LlmClient` の実装
|
||||
///
|
||||
/// これにより、動的ディスパッチを使用するクライアントも `Engine` で利用可能になる。
|
||||
#[async_trait]
|
||||
impl LlmClient for Box<dyn LlmClient> {
|
||||
async fn stream(&self, request: Request) -> Result<ResponseStream, ClientError> {
|
||||
(**self).stream(request).await
|
||||
}
|
||||
|
||||
fn clone_boxed(&self) -> Box<dyn LlmClient> {
|
||||
(**self).clone_boxed()
|
||||
}
|
||||
|
||||
fn validate_config(&self, config: &RequestConfig) -> Vec<ConfigWarning> {
|
||||
(**self).validate_config(config)
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,172 @@
|
||||
//! LLMクライアントエラー型
|
||||
|
||||
use std::{fmt, time::Duration};
|
||||
|
||||
/// LLMクライアントのエラー
|
||||
#[derive(Debug)]
|
||||
pub enum ClientError {
|
||||
/// HTTPリクエストエラー
|
||||
Http(reqwest::Error),
|
||||
/// JSONパースエラー
|
||||
Json(serde_json::Error),
|
||||
/// SSEパースエラー
|
||||
Sse(String),
|
||||
/// APIエラー (プロバイダからのエラーレスポンス)
|
||||
Api {
|
||||
status: Option<u16>,
|
||||
code: Option<String>,
|
||||
message: String,
|
||||
retry_after: Option<Duration>,
|
||||
},
|
||||
/// A request lifecycle phase exceeded its hard timeout.
|
||||
Timeout {
|
||||
phase: &'static str,
|
||||
timeout: Duration,
|
||||
},
|
||||
/// 設定エラー
|
||||
Config(String),
|
||||
}
|
||||
|
||||
impl fmt::Display for ClientError {
|
||||
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
|
||||
match self {
|
||||
ClientError::Http(e) => write!(f, "HTTP error: {}", e),
|
||||
ClientError::Json(e) => write!(f, "JSON parse error: {}", e),
|
||||
ClientError::Sse(msg) => write!(f, "SSE parse error: {}", msg),
|
||||
ClientError::Api {
|
||||
status,
|
||||
code,
|
||||
message,
|
||||
..
|
||||
} => {
|
||||
write!(f, "API error")?;
|
||||
if let Some(s) = status {
|
||||
write!(f, " (status: {})", s)?;
|
||||
}
|
||||
if let Some(c) = code {
|
||||
write!(f, " [{}]", c)?;
|
||||
}
|
||||
write!(f, ": {}", message)
|
||||
}
|
||||
ClientError::Timeout { phase, timeout } => {
|
||||
write!(f, "{phase} timed out after {}s", timeout.as_secs())
|
||||
}
|
||||
ClientError::Config(msg) => write!(f, "Config error: {}", msg),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
impl std::error::Error for ClientError {
|
||||
fn source(&self) -> Option<&(dyn std::error::Error + 'static)> {
|
||||
match self {
|
||||
ClientError::Http(e) => Some(e),
|
||||
ClientError::Json(e) => Some(e),
|
||||
_ => None,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
impl From<reqwest::Error> for ClientError {
|
||||
fn from(err: reqwest::Error) -> Self {
|
||||
ClientError::Http(err)
|
||||
}
|
||||
}
|
||||
|
||||
impl From<serde_json::Error> for ClientError {
|
||||
fn from(err: serde_json::Error) -> Self {
|
||||
ClientError::Json(err)
|
||||
}
|
||||
}
|
||||
|
||||
impl ClientError {
|
||||
pub fn status(&self) -> Option<u16> {
|
||||
match self {
|
||||
ClientError::Api { status, .. } => *status,
|
||||
_ => None,
|
||||
}
|
||||
}
|
||||
|
||||
pub fn retry_after(&self) -> Option<Duration> {
|
||||
match self {
|
||||
ClientError::Api { retry_after, .. } => *retry_after,
|
||||
_ => None,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// transient な失敗としてリトライ対象になるかを判定する。
|
||||
///
|
||||
/// 対象:
|
||||
/// - `Api { status }` のうち 408 / 425 / 429 / 500 / 502 / 503 / 504 / 529
|
||||
/// - `Http(reqwest::Error)` のうち `is_connect()` または `is_timeout()`
|
||||
/// - `Timeout { .. }` の lifecycle hard timeout
|
||||
///
|
||||
/// それ以外（Json、Sse、Config、上記以外の Api ステータス）は false。
|
||||
/// SSE 読み出し開始後の失敗は呼び出し側で `Sse` として上に流すため、
|
||||
/// ここで対象外にしておけば自動的に弾かれる。
|
||||
pub fn is_retryable(error: &ClientError) -> bool {
|
||||
match error {
|
||||
ClientError::Api {
|
||||
status: Some(code), ..
|
||||
} => matches!(*code, 408 | 425 | 429 | 500 | 502 | 503 | 504 | 529),
|
||||
ClientError::Api { status: None, .. } => false,
|
||||
ClientError::Timeout { .. } => true,
|
||||
ClientError::Http(e) => e.is_connect() || e.is_timeout(),
|
||||
ClientError::Json(_) | ClientError::Sse(_) | ClientError::Config(_) => false,
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
|
||||
fn api_err(status: Option<u16>) -> ClientError {
|
||||
ClientError::Api {
|
||||
status,
|
||||
code: None,
|
||||
message: String::new(),
|
||||
retry_after: None,
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn retryable_status_codes() {
|
||||
for code in [408u16, 425, 429, 500, 502, 503, 504, 529] {
|
||||
assert!(
|
||||
is_retryable(&api_err(Some(code))),
|
||||
"status {code} should be retryable",
|
||||
);
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn non_retryable_status_codes() {
|
||||
for code in [400u16, 401, 403, 404, 409, 410, 422, 501] {
|
||||
assert!(
|
||||
!is_retryable(&api_err(Some(code))),
|
||||
"status {code} should not be retryable",
|
||||
);
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn api_without_status_not_retryable() {
|
||||
assert!(!is_retryable(&api_err(None)));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn lifecycle_timeout_is_retryable() {
|
||||
assert!(is_retryable(&ClientError::Timeout {
|
||||
phase: "stream_open",
|
||||
timeout: Duration::from_secs(30),
|
||||
}));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn json_sse_config_not_retryable() {
|
||||
let json_err = serde_json::from_str::<serde_json::Value>("not json").unwrap_err();
|
||||
assert!(!is_retryable(&ClientError::Json(json_err)));
|
||||
assert!(!is_retryable(&ClientError::Sse("boom".into())));
|
||||
assert!(!is_retryable(&ClientError::Config("boom".into())));
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,348 @@
|
||||
//! LLMクライアント層のイベント型
|
||||
//!
|
||||
//! 各LLMプロバイダからのストリーミングレスポンスを表現するイベント型。
|
||||
|
||||
use serde::{Deserialize, Serialize};
|
||||
|
||||
// =============================================================================
|
||||
// Core Event Types (from llm_client layer)
|
||||
// =============================================================================
|
||||
|
||||
/// LLMからのストリーミングイベント
|
||||
///
|
||||
/// 各LLMプロバイダからのレスポンスは、この`Event`のストリームとして
|
||||
/// 統一的に処理されます。
|
||||
///
|
||||
/// # イベントの種類
|
||||
///
|
||||
/// - **メタイベント**: `Ping`, `Usage`, `Status`, `Error`, `UnhandledSse`
|
||||
/// - **ブロックイベント**: `BlockStart`, `BlockDelta`, `BlockStop`, `BlockAbort`
|
||||
///
|
||||
/// # ブロックのライフサイクル
|
||||
///
|
||||
/// テキスト、thinking、ツール呼び出しは、`BlockStart` → `BlockDelta`(複数) → `BlockStop`
|
||||
/// の順序でイベントが発生します。thinking の round-trip metadata は
|
||||
/// `BlockStop.reasoning` に載ります。
|
||||
#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
|
||||
pub enum Event {
|
||||
/// ハートビート
|
||||
Ping(PingEvent),
|
||||
/// トークン使用量
|
||||
Usage(UsageEvent),
|
||||
/// ストリームのステータス変化
|
||||
Status(StatusEvent),
|
||||
/// エラー発生
|
||||
Error(ErrorEvent),
|
||||
/// Scheme が生成内容として解釈しない未対応 SSE イベント。
|
||||
///
|
||||
/// stream trace 用の観測イベントであり、timeline / history には反映しない。
|
||||
UnhandledSse(UnhandledSseEvent),
|
||||
|
||||
/// ブロック開始（テキスト、ツール使用等）
|
||||
BlockStart(BlockStart),
|
||||
/// ブロックの差分データ
|
||||
BlockDelta(BlockDelta),
|
||||
/// ブロック正常終了
|
||||
BlockStop(BlockStop),
|
||||
/// ブロック中断
|
||||
BlockAbort(BlockAbort),
|
||||
}
|
||||
|
||||
// =============================================================================
|
||||
// Meta Events
|
||||
// =============================================================================
|
||||
|
||||
/// Pingイベント（ハートビート）
|
||||
#[derive(Debug, Clone, PartialEq, Default, Serialize, Deserialize)]
|
||||
pub struct PingEvent {
|
||||
pub timestamp: Option<u64>,
|
||||
}
|
||||
|
||||
/// 使用量イベント
|
||||
///
|
||||
/// プロバイダから受信した 1 LLM リクエスト分のトークン会計。
|
||||
/// 各 scheme で正規化され、フィールドの意味は全プロバイダ共通:
|
||||
///
|
||||
/// - `input_tokens` は **送信した prompt prefix 全体の占有量**（プロンプト全長）。
|
||||
/// キャッシュヒット分も含まれる。Anthropic は raw API では非キャッシュ分のみを
|
||||
/// `input_tokens` として返すため、`AnthropicScheme::convert_usage` で
|
||||
/// `cache_read + cache_creation` を加算してこの規約に揃えている。
|
||||
/// - `cache_read_input_tokens` / `cache_creation_input_tokens` は上記の内訳で、
|
||||
/// 料金会計用。占有量からは差し引かない。
|
||||
#[derive(Debug, Clone, PartialEq, Default, Serialize, Deserialize)]
|
||||
pub struct UsageEvent {
|
||||
/// 送信した prompt prefix の総トークン数（占有量、キャッシュ込み）
|
||||
pub input_tokens: Option<u64>,
|
||||
/// このリクエストで生成された出力トークン数
|
||||
pub output_tokens: Option<u64>,
|
||||
/// `input_tokens + output_tokens`
|
||||
pub total_tokens: Option<u64>,
|
||||
/// `input_tokens` のうちキャッシュから読まれた分（割引料金）
|
||||
pub cache_read_input_tokens: Option<u64>,
|
||||
/// `input_tokens` のうちこのリクエストでキャッシュに書かれた分（割増料金、Anthropic）
|
||||
pub cache_creation_input_tokens: Option<u64>,
|
||||
}
|
||||
|
||||
/// ステータスイベント
|
||||
#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
|
||||
pub struct StatusEvent {
|
||||
pub status: ResponseStatus,
|
||||
}
|
||||
|
||||
/// レスポンスステータス
|
||||
#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
|
||||
pub enum ResponseStatus {
|
||||
/// ストリーム開始
|
||||
Started,
|
||||
/// 正常完了
|
||||
Completed,
|
||||
/// キャンセルされた
|
||||
Cancelled,
|
||||
/// エラー発生
|
||||
Failed,
|
||||
}
|
||||
|
||||
/// エラーイベント
|
||||
#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
|
||||
pub struct ErrorEvent {
|
||||
pub code: Option<String>,
|
||||
pub message: String,
|
||||
}
|
||||
|
||||
/// 未対応 SSE イベントの観測用メタイベント。
|
||||
///
|
||||
/// `data_preview` は provider から受け取った raw SSE data の bounded preview、
|
||||
/// `data_len` は preview 前の raw data byte length。
|
||||
#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
|
||||
pub struct UnhandledSseEvent {
|
||||
pub provider: String,
|
||||
pub event_type: String,
|
||||
pub data_preview: String,
|
||||
pub data_len: usize,
|
||||
}
|
||||
|
||||
// =============================================================================
|
||||
// Block Types
|
||||
// =============================================================================
|
||||
|
||||
/// ブロックの種別
|
||||
#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, Serialize, Deserialize)]
|
||||
pub enum BlockType {
|
||||
/// テキスト生成
|
||||
Text,
|
||||
/// 思考 (Claude Extended Thinking等)
|
||||
Thinking,
|
||||
/// ツール呼び出し
|
||||
ToolUse,
|
||||
/// ツール結果
|
||||
ToolResult,
|
||||
}
|
||||
|
||||
/// ブロック開始イベント
|
||||
#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
|
||||
pub struct BlockStart {
|
||||
/// ブロックのインデックス
|
||||
pub index: usize,
|
||||
/// ブロックの種別
|
||||
pub block_type: BlockType,
|
||||
/// ブロック固有のメタデータ
|
||||
pub metadata: BlockMetadata,
|
||||
}
|
||||
|
||||
impl BlockStart {
|
||||
pub fn block_type(&self) -> BlockType {
|
||||
self.block_type
|
||||
}
|
||||
}
|
||||
|
||||
/// ブロックのメタデータ
|
||||
#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
|
||||
pub enum BlockMetadata {
|
||||
Text,
|
||||
Thinking,
|
||||
ToolUse { id: String, name: String },
|
||||
ToolResult { tool_use_id: String },
|
||||
}
|
||||
|
||||
/// ブロックデルタイベント
|
||||
#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
|
||||
pub struct BlockDelta {
|
||||
/// ブロックのインデックス
|
||||
pub index: usize,
|
||||
/// デルタの内容
|
||||
pub delta: DeltaContent,
|
||||
}
|
||||
|
||||
/// デルタの内容
|
||||
#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
|
||||
pub enum DeltaContent {
|
||||
/// テキストデルタ
|
||||
Text(String),
|
||||
/// 思考デルタ
|
||||
Thinking(String),
|
||||
/// ツール引数のJSON部分文字列
|
||||
InputJson(String),
|
||||
}
|
||||
|
||||
impl DeltaContent {
|
||||
/// デルタのブロック種別を取得
|
||||
pub fn block_type(&self) -> BlockType {
|
||||
match self {
|
||||
DeltaContent::Text(_) => BlockType::Text,
|
||||
DeltaContent::Thinking(_) => BlockType::Thinking,
|
||||
DeltaContent::InputJson(_) => BlockType::ToolUse,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// ブロック停止イベント
|
||||
#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
|
||||
pub struct BlockStop {
|
||||
/// ブロックのインデックス
|
||||
pub index: usize,
|
||||
/// ブロックの種別
|
||||
pub block_type: BlockType,
|
||||
/// 停止理由
|
||||
pub stop_reason: Option<StopReason>,
|
||||
/// Thinking block の停止時に確定した reasoning round-trip metadata。
|
||||
///
|
||||
/// `None` の Thinking block は live streaming / trace 用で、history に
|
||||
/// `Item::Reasoning` として永続化しない。`Some` の場合は block lifecycle
|
||||
/// が永続化の authoritative source になる。
|
||||
pub reasoning: Option<ReasoningBlockData>,
|
||||
}
|
||||
|
||||
impl BlockStop {
|
||||
pub fn block_type(&self) -> BlockType {
|
||||
self.block_type
|
||||
}
|
||||
}
|
||||
|
||||
/// ブロック中断イベント
|
||||
#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
|
||||
pub struct BlockAbort {
|
||||
/// ブロックのインデックス
|
||||
pub index: usize,
|
||||
/// ブロックの種別
|
||||
pub block_type: BlockType,
|
||||
/// 中断理由
|
||||
pub reason: String,
|
||||
}
|
||||
|
||||
impl BlockAbort {
|
||||
pub fn block_type(&self) -> BlockType {
|
||||
self.block_type
|
||||
}
|
||||
}
|
||||
|
||||
/// Thinking block stop で確定した reasoning material。
|
||||
///
|
||||
/// `Item::Reasoning` の round-trip に必要な provider material を保持する。
|
||||
/// `text` は deltas から収集した本文を上書きするために使う（metadata-only
|
||||
/// reasoning block や provider completion event で全文が届くケース）。
|
||||
#[derive(Debug, Clone, PartialEq, Default, Serialize, Deserialize)]
|
||||
pub struct ReasoningBlockData {
|
||||
/// scheme 側で観測した item id（OpenAI Responses の `id`）。
|
||||
pub id: Option<String>,
|
||||
/// reasoning 本体テキスト。`None` の場合は block delta 収集結果を使う。
|
||||
pub text: Option<String>,
|
||||
/// summary (OpenAI Responses の `summary_text[]`)。他 scheme は空。
|
||||
pub summary: Vec<String>,
|
||||
/// 暗号化された opaque blob（Anthropic `redacted_thinking.data` /
|
||||
/// OpenAI Responses `encrypted_content`）。
|
||||
pub encrypted_content: Option<String>,
|
||||
/// Anthropic extended thinking signature。round-trip 必須。
|
||||
pub signature: Option<String>,
|
||||
}
|
||||
|
||||
/// 停止理由
|
||||
#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
|
||||
pub enum StopReason {
|
||||
/// 自然終了
|
||||
EndTurn,
|
||||
/// 最大トークン数到達
|
||||
MaxTokens,
|
||||
/// ストップシーケンス到達
|
||||
StopSequence,
|
||||
/// ツール使用
|
||||
ToolUse,
|
||||
}
|
||||
|
||||
// =============================================================================
|
||||
// Builder / Factory helpers
|
||||
// =============================================================================
|
||||
|
||||
impl Event {
|
||||
/// テキストブロック開始イベントを作成
|
||||
pub fn text_block_start(index: usize) -> Self {
|
||||
Event::BlockStart(BlockStart {
|
||||
index,
|
||||
block_type: BlockType::Text,
|
||||
metadata: BlockMetadata::Text,
|
||||
})
|
||||
}
|
||||
|
||||
/// テキストデルタイベントを作成
|
||||
pub fn text_delta(index: usize, text: impl Into<String>) -> Self {
|
||||
Event::BlockDelta(BlockDelta {
|
||||
index,
|
||||
delta: DeltaContent::Text(text.into()),
|
||||
})
|
||||
}
|
||||
|
||||
/// テキストブロック停止イベントを作成
|
||||
pub fn text_block_stop(index: usize, stop_reason: Option<StopReason>) -> Self {
|
||||
Event::BlockStop(BlockStop {
|
||||
index,
|
||||
block_type: BlockType::Text,
|
||||
stop_reason,
|
||||
reasoning: None,
|
||||
})
|
||||
}
|
||||
|
||||
/// ツール使用ブロック開始イベントを作成
|
||||
pub fn tool_use_start(index: usize, id: impl Into<String>, name: impl Into<String>) -> Self {
|
||||
Event::BlockStart(BlockStart {
|
||||
index,
|
||||
block_type: BlockType::ToolUse,
|
||||
metadata: BlockMetadata::ToolUse {
|
||||
id: id.into(),
|
||||
name: name.into(),
|
||||
},
|
||||
})
|
||||
}
|
||||
|
||||
/// ツール引数デルタイベントを作成
|
||||
pub fn tool_input_delta(index: usize, json: impl Into<String>) -> Self {
|
||||
Event::BlockDelta(BlockDelta {
|
||||
index,
|
||||
delta: DeltaContent::InputJson(json.into()),
|
||||
})
|
||||
}
|
||||
|
||||
/// ツール使用ブロック停止イベントを作成
|
||||
pub fn tool_use_stop(index: usize) -> Self {
|
||||
Event::BlockStop(BlockStop {
|
||||
index,
|
||||
block_type: BlockType::ToolUse,
|
||||
stop_reason: Some(StopReason::ToolUse),
|
||||
reasoning: None,
|
||||
})
|
||||
}
|
||||
|
||||
/// 使用量イベントを作成
|
||||
pub fn usage(input_tokens: u64, output_tokens: u64) -> Self {
|
||||
Event::Usage(UsageEvent {
|
||||
input_tokens: Some(input_tokens),
|
||||
output_tokens: Some(output_tokens),
|
||||
total_tokens: Some(input_tokens + output_tokens),
|
||||
cache_read_input_tokens: None,
|
||||
cache_creation_input_tokens: None,
|
||||
})
|
||||
}
|
||||
|
||||
/// Pingイベントを作成
|
||||
pub fn ping() -> Self {
|
||||
Event::Ping(PingEvent { timestamp: None })
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,35 @@
|
||||
//! LLMクライアント層
|
||||
//!
|
||||
//! 各LLMプロバイダと通信し、統一された[`Event`]
|
||||
//! ストリームを出力します。
|
||||
//!
|
||||
//! # サポートするプロバイダ
|
||||
//!
|
||||
//! - Anthropic (Claude)
|
||||
//! - OpenAI (GPT-4, etc.)
|
||||
//! - Google (Gemini)
|
||||
//! - Ollama (ローカルLLM)
|
||||
//!
|
||||
//! # アーキテクチャ
|
||||
//!
|
||||
//! - [`LlmClient`] - プロバイダ共通のtrait
|
||||
//! - `providers`: プロバイダ固有のクライアント実装
|
||||
//! - `scheme`: APIスキーマ（リクエスト/レスポンス変換）
|
||||
|
||||
pub mod auth;
|
||||
pub mod capability;
|
||||
pub mod client;
|
||||
pub mod error;
|
||||
pub mod event;
|
||||
pub mod types;
|
||||
|
||||
pub mod retry;
|
||||
pub mod scheme;
|
||||
pub mod transport;
|
||||
|
||||
pub use auth::*;
|
||||
pub use capability::*;
|
||||
pub use client::*;
|
||||
pub use error::*;
|
||||
pub use event::*;
|
||||
pub use types::*;
|
||||
@@ -0,0 +1,104 @@
|
||||
//! LLM response stream を開く前の transient error 向けリトライポリシー。
|
||||
//!
|
||||
//! Engine が `LlmClient::stream` の open error に対して `is_retryable` を見て
|
||||
//! retry / backoff / TUI event / cancellation をまとめて管理する。
|
||||
//! SSE 読み出し開始後の失敗は対象外。
|
||||
|
||||
use std::time::Duration;
|
||||
|
||||
/// 指数バックオフ + ジッター + 累積タイムアウトを表すポリシー。
|
||||
///
|
||||
/// `Default` は llm-engine 全体の固定値を返す。manifest 経由の上書きが
|
||||
/// 必要になったら拡張する（現状は不要 → `tickets/llm-engine-transient-retry.md`）。
|
||||
#[derive(Debug, Clone)]
|
||||
pub struct RetryPolicy {
|
||||
/// 指数の基準値。`base * 2^attempt` を `cap` で頭打ちにした上限から
|
||||
/// フルジッターで実際の wait を抽選する。
|
||||
pub base: Duration,
|
||||
/// 1 回あたりの wait の上限。
|
||||
pub cap: Duration,
|
||||
/// 試行の合計回数（初回 + リトライ）。`1` ならリトライしない。
|
||||
pub max_attempts: u32,
|
||||
/// 初回送信開始からの累積タイムアウト。これを超える wait は打ち切る。
|
||||
pub total_timeout: Duration,
|
||||
}
|
||||
|
||||
impl Default for RetryPolicy {
|
||||
fn default() -> Self {
|
||||
Self {
|
||||
base: Duration::from_millis(500),
|
||||
cap: Duration::from_secs(10),
|
||||
max_attempts: 4,
|
||||
total_timeout: Duration::from_secs(40),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
impl RetryPolicy {
|
||||
/// `attempt` 回目の失敗（0-indexed）後に待つ時間を返す。
|
||||
/// `Retry-After` で上書きしたい場合は呼び出さず、その値をそのまま使う。
|
||||
pub fn backoff(&self, attempt: u32) -> Duration {
|
||||
let shift = attempt.min(20);
|
||||
let base_nanos = self.base.as_nanos() as u64;
|
||||
let exp_nanos = base_nanos.saturating_mul(1u64 << shift);
|
||||
let cap_nanos = self.cap.as_nanos() as u64;
|
||||
let upper = exp_nanos.min(cap_nanos);
|
||||
Duration::from_nanos(jitter_nanos(upper))
|
||||
}
|
||||
}
|
||||
|
||||
/// `[0, max_nanos]` から擬似乱数的に 1 つ取り出す。`SystemTime` の
|
||||
/// 下位ビットを splitmix64 で攪拌するだけの軽量実装で、暗号的乱数性は
|
||||
/// 持たないがフルジッターのぶつかり回避には十分。
|
||||
fn jitter_nanos(max_nanos: u64) -> u64 {
|
||||
if max_nanos == 0 {
|
||||
return 0;
|
||||
}
|
||||
let seed = std::time::SystemTime::now()
|
||||
.duration_since(std::time::UNIX_EPOCH)
|
||||
.map(|d| d.as_nanos() as u64)
|
||||
.unwrap_or(0);
|
||||
let mut x = seed.wrapping_add(0x9E37_79B9_7F4A_7C15);
|
||||
x = (x ^ (x >> 30)).wrapping_mul(0xBF58_476D_1CE4_E5B9);
|
||||
x = (x ^ (x >> 27)).wrapping_mul(0x94D0_49BB_1331_11EB);
|
||||
x ^= x >> 31;
|
||||
x % (max_nanos + 1)
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
|
||||
#[test]
|
||||
fn default_policy_values() {
|
||||
let p = RetryPolicy::default();
|
||||
assert_eq!(p.base, Duration::from_millis(500));
|
||||
assert_eq!(p.cap, Duration::from_secs(10));
|
||||
assert_eq!(p.max_attempts, 4);
|
||||
assert_eq!(p.total_timeout, Duration::from_secs(40));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn backoff_respects_cap() {
|
||||
let p = RetryPolicy::default();
|
||||
for attempt in 0..30u32 {
|
||||
assert!(
|
||||
p.backoff(attempt) <= p.cap,
|
||||
"attempt {attempt} exceeded cap",
|
||||
);
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn backoff_zero_when_base_zero() {
|
||||
let p = RetryPolicy {
|
||||
base: Duration::ZERO,
|
||||
cap: Duration::from_secs(10),
|
||||
max_attempts: 4,
|
||||
total_timeout: Duration::from_secs(30),
|
||||
};
|
||||
for attempt in 0..5 {
|
||||
assert_eq!(p.backoff(attempt), Duration::ZERO);
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,23 @@
|
||||
//! Anthropic scheme の wire-level 既定 capability。
|
||||
//!
|
||||
//! モデル ID 固有のテーブル(`claude-*` など)は高レベル構築層
|
||||
//! (`provider::capability`)の責務。ここでは未知モデルでも「この wire で
|
||||
//! 安全に送れる最小共通項」を返すだけに留める。
|
||||
|
||||
use crate::llm_client::capability::{
|
||||
CacheStrategy, ModelCapability, StructuredOutput, ToolCallingSupport,
|
||||
};
|
||||
|
||||
/// Scheme 既定の capability。
|
||||
///
|
||||
/// Ollama の `/v1/messages` 流用を想定して `cache_control` を送らない
|
||||
/// `CacheStrategy::Auto` にする。
|
||||
pub(crate) fn default_capability() -> ModelCapability {
|
||||
ModelCapability {
|
||||
tool_calling: ToolCallingSupport::Parallel,
|
||||
structured_output: StructuredOutput::JsonSchema,
|
||||
reasoning: None,
|
||||
vision: false,
|
||||
prompt_caching: CacheStrategy::Auto,
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,659 @@
|
||||
//! Anthropic SSEイベントパース
|
||||
//!
|
||||
//! Anthropic Messages APIのSSEイベントをパースし、統一Event型に変換
|
||||
|
||||
use crate::llm_client::{
|
||||
ClientError,
|
||||
event::{
|
||||
BlockDelta, BlockMetadata, BlockStart, BlockStop, BlockType, DeltaContent, ErrorEvent,
|
||||
Event, PingEvent, ResponseStatus, StatusEvent, UsageEvent,
|
||||
},
|
||||
};
|
||||
use serde::Deserialize;
|
||||
|
||||
use super::AnthropicScheme;
|
||||
use super::scheme_impl::{AnthropicState, PendingThinking};
|
||||
|
||||
/// Anthropic SSEイベントタイプ
|
||||
#[derive(Debug, Clone, PartialEq, Eq)]
|
||||
pub(crate) enum AnthropicEventType {
|
||||
MessageStart,
|
||||
ContentBlockStart,
|
||||
ContentBlockDelta,
|
||||
ContentBlockStop,
|
||||
MessageDelta,
|
||||
MessageStop,
|
||||
Ping,
|
||||
Error,
|
||||
}
|
||||
|
||||
impl AnthropicEventType {
|
||||
/// イベントタイプ文字列からパース
|
||||
pub(crate) fn parse(s: &str) -> Option<Self> {
|
||||
match s {
|
||||
"message_start" => Some(Self::MessageStart),
|
||||
"content_block_start" => Some(Self::ContentBlockStart),
|
||||
"content_block_delta" => Some(Self::ContentBlockDelta),
|
||||
"content_block_stop" => Some(Self::ContentBlockStop),
|
||||
"message_delta" => Some(Self::MessageDelta),
|
||||
"message_stop" => Some(Self::MessageStop),
|
||||
"ping" => Some(Self::Ping),
|
||||
"error" => Some(Self::Error),
|
||||
_ => None,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// ============================================================================
|
||||
// SSEイベントのJSON構造
|
||||
// ============================================================================
|
||||
|
||||
/// message_start イベント
|
||||
#[derive(Debug, Deserialize)]
|
||||
pub(crate) struct MessageStartEvent {
|
||||
pub message: MessageStartMessage,
|
||||
}
|
||||
|
||||
#[allow(dead_code)]
|
||||
#[derive(Debug, Deserialize)]
|
||||
pub(crate) struct MessageStartMessage {
|
||||
pub id: String,
|
||||
pub model: String,
|
||||
pub usage: Option<UsageData>,
|
||||
}
|
||||
|
||||
/// content_block_start イベント
|
||||
#[derive(Debug, Deserialize)]
|
||||
pub(crate) struct ContentBlockStartEvent {
|
||||
pub index: usize,
|
||||
pub content_block: ContentBlock,
|
||||
}
|
||||
|
||||
#[allow(dead_code)]
|
||||
#[derive(Debug, Deserialize)]
|
||||
#[serde(tag = "type")]
|
||||
pub(crate) enum ContentBlock {
|
||||
#[serde(rename = "text")]
|
||||
Text { text: String },
|
||||
#[serde(rename = "thinking")]
|
||||
Thinking {
|
||||
#[serde(default)]
|
||||
thinking: String,
|
||||
/// 非ストリーミングレスポンス由来の初期 signature（通常はストリームでは
|
||||
/// 空 → `signature_delta` で埋まる）。
|
||||
#[serde(default)]
|
||||
signature: Option<String>,
|
||||
},
|
||||
#[serde(rename = "redacted_thinking")]
|
||||
RedactedThinking {
|
||||
/// 暗号化された opaque blob。signature ではなく、まるごと
|
||||
/// `redacted_thinking.data` として送り返す必要がある。
|
||||
#[serde(default)]
|
||||
data: String,
|
||||
},
|
||||
#[serde(rename = "tool_use")]
|
||||
ToolUse {
|
||||
id: String,
|
||||
name: String,
|
||||
input: serde_json::Value,
|
||||
},
|
||||
}
|
||||
|
||||
/// content_block_delta イベント
|
||||
#[derive(Debug, Deserialize)]
|
||||
pub(crate) struct ContentBlockDeltaEvent {
|
||||
pub index: usize,
|
||||
pub delta: DeltaBlock,
|
||||
}
|
||||
|
||||
#[allow(dead_code)]
|
||||
#[derive(Debug, Deserialize)]
|
||||
#[serde(tag = "type")]
|
||||
pub(crate) enum DeltaBlock {
|
||||
#[serde(rename = "text_delta")]
|
||||
TextDelta { text: String },
|
||||
#[serde(rename = "thinking_delta")]
|
||||
ThinkingDelta { thinking: String },
|
||||
#[serde(rename = "input_json_delta")]
|
||||
InputJsonDelta { partial_json: String },
|
||||
#[serde(rename = "signature_delta")]
|
||||
SignatureDelta { signature: String },
|
||||
}
|
||||
|
||||
/// content_block_stop イベント
|
||||
#[derive(Debug, Deserialize)]
|
||||
pub(crate) struct ContentBlockStopEvent {
|
||||
pub index: usize,
|
||||
}
|
||||
|
||||
/// message_delta イベント
|
||||
#[allow(dead_code)]
|
||||
#[derive(Debug, Deserialize)]
|
||||
pub(crate) struct MessageDeltaEvent {
|
||||
pub delta: MessageDeltaData,
|
||||
pub usage: Option<UsageData>,
|
||||
}
|
||||
|
||||
#[allow(dead_code)]
|
||||
#[derive(Debug, Deserialize)]
|
||||
pub(crate) struct MessageDeltaData {
|
||||
pub stop_reason: Option<String>,
|
||||
pub stop_sequence: Option<String>,
|
||||
}
|
||||
|
||||
/// 使用量データ
|
||||
#[derive(Debug, Deserialize)]
|
||||
pub(crate) struct UsageData {
|
||||
pub input_tokens: Option<u64>,
|
||||
pub output_tokens: Option<u64>,
|
||||
pub cache_read_input_tokens: Option<u64>,
|
||||
pub cache_creation_input_tokens: Option<u64>,
|
||||
}
|
||||
|
||||
/// エラーイベント
|
||||
#[derive(Debug, Deserialize)]
|
||||
pub(crate) struct ErrorEventData {
|
||||
pub error: ErrorDetail,
|
||||
}
|
||||
|
||||
#[derive(Debug, Deserialize)]
|
||||
pub(crate) struct ErrorDetail {
|
||||
#[serde(rename = "type")]
|
||||
pub error_type: String,
|
||||
pub message: String,
|
||||
}
|
||||
|
||||
// ============================================================================
|
||||
// イベント変換
|
||||
// ============================================================================
|
||||
|
||||
impl AnthropicScheme {
|
||||
/// SSEイベントをEvent型に変換
|
||||
///
|
||||
/// # Arguments
|
||||
/// * `event_type` - SSEイベントタイプ
|
||||
/// * `data` - イベントデータJSON文字列
|
||||
///
|
||||
/// # Returns
|
||||
/// * `Ok(Some(Event))` - 変換成功
|
||||
/// * `Ok(None)` - イベントを無視（unknown event等）
|
||||
/// * `Err(ClientError)` - パースエラー
|
||||
pub(crate) fn parse_event(
|
||||
&self,
|
||||
event_type: &str,
|
||||
data: &str,
|
||||
) -> Result<Option<Event>, ClientError> {
|
||||
let Some(event_type) = AnthropicEventType::parse(event_type) else {
|
||||
// Unknown event type, ignore
|
||||
return Ok(None);
|
||||
};
|
||||
|
||||
match event_type {
|
||||
AnthropicEventType::MessageStart => {
|
||||
let event: MessageStartEvent = serde_json::from_str(data)?;
|
||||
// message_start時にUsageイベントがあれば出力
|
||||
if let Some(usage) = event.message.usage {
|
||||
return Ok(Some(Event::Usage(self.convert_usage(&usage))));
|
||||
}
|
||||
// Statusイベントとして開始を通知
|
||||
Ok(Some(Event::Status(StatusEvent {
|
||||
status: ResponseStatus::Started,
|
||||
})))
|
||||
}
|
||||
AnthropicEventType::ContentBlockStart => {
|
||||
let event: ContentBlockStartEvent = serde_json::from_str(data)?;
|
||||
Ok(Some(self.convert_block_start(&event)))
|
||||
}
|
||||
AnthropicEventType::ContentBlockDelta => {
|
||||
let event: ContentBlockDeltaEvent = serde_json::from_str(data)?;
|
||||
Ok(self.convert_block_delta(&event))
|
||||
}
|
||||
AnthropicEventType::ContentBlockStop => {
|
||||
let event: ContentBlockStopEvent = serde_json::from_str(data)?;
|
||||
// Note: BlockStopにはblock_typeが必要だが、AnthropicはStopイベントに含めない
|
||||
// Timeline層がBlockStartを追跡して正しいblock_typeを知る
|
||||
Ok(Some(Event::BlockStop(BlockStop {
|
||||
index: event.index,
|
||||
block_type: BlockType::Text, // Timeline層で上書きされる
|
||||
stop_reason: None,
|
||||
reasoning: None,
|
||||
})))
|
||||
}
|
||||
AnthropicEventType::MessageDelta => {
|
||||
let event: MessageDeltaEvent = serde_json::from_str(data)?;
|
||||
// Usage情報があれば出力
|
||||
if let Some(usage) = event.usage {
|
||||
return Ok(Some(Event::Usage(self.convert_usage(&usage))));
|
||||
}
|
||||
Ok(None)
|
||||
}
|
||||
AnthropicEventType::MessageStop => Ok(Some(Event::Status(StatusEvent {
|
||||
status: ResponseStatus::Completed,
|
||||
}))),
|
||||
AnthropicEventType::Ping => Ok(Some(Event::Ping(PingEvent { timestamp: None }))),
|
||||
AnthropicEventType::Error => {
|
||||
let event: ErrorEventData = serde_json::from_str(data)?;
|
||||
Ok(Some(Event::Error(ErrorEvent {
|
||||
code: Some(event.error.error_type),
|
||||
message: event.error.message,
|
||||
})))
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
fn convert_block_start(&self, event: &ContentBlockStartEvent) -> Event {
|
||||
let (block_type, metadata) = match &event.content_block {
|
||||
ContentBlock::Text { .. } => (BlockType::Text, BlockMetadata::Text),
|
||||
ContentBlock::Thinking { .. } | ContentBlock::RedactedThinking { .. } => {
|
||||
(BlockType::Thinking, BlockMetadata::Thinking)
|
||||
}
|
||||
ContentBlock::ToolUse { id, name, .. } => (
|
||||
BlockType::ToolUse,
|
||||
BlockMetadata::ToolUse {
|
||||
id: id.clone(),
|
||||
name: name.clone(),
|
||||
},
|
||||
),
|
||||
};
|
||||
|
||||
Event::BlockStart(BlockStart {
|
||||
index: event.index,
|
||||
block_type,
|
||||
metadata,
|
||||
})
|
||||
}
|
||||
|
||||
fn convert_block_delta(&self, event: &ContentBlockDeltaEvent) -> Option<Event> {
|
||||
let delta = match &event.delta {
|
||||
DeltaBlock::TextDelta { text } => DeltaContent::Text(text.clone()),
|
||||
DeltaBlock::ThinkingDelta { thinking } => DeltaContent::Thinking(thinking.clone()),
|
||||
DeltaBlock::InputJsonDelta { partial_json } => {
|
||||
DeltaContent::InputJson(partial_json.clone())
|
||||
}
|
||||
DeltaBlock::SignatureDelta { .. } => {
|
||||
// signature_delta は無視
|
||||
return None;
|
||||
}
|
||||
};
|
||||
|
||||
Some(Event::BlockDelta(BlockDelta {
|
||||
index: event.index,
|
||||
delta,
|
||||
}))
|
||||
}
|
||||
|
||||
/// state を持ち回す上位パース。
|
||||
///
|
||||
/// `parse_event` の単発 Event に加えて、以下を行う:
|
||||
/// - `content_block_stop` の `block_type` を直前の Start 値で書き戻す
|
||||
/// - `thinking` / `redacted_thinking` ブロックの本体・signature・data を
|
||||
/// `state.pending_thinking` に蓄積し、`content_block_stop` の Thinking
|
||||
/// BlockStop metadata に載せる
|
||||
/// - `signature_delta` を蓄積（Stream channel には流さず、reasoning metadata
|
||||
/// にだけ反映する）
|
||||
pub(crate) fn parse_with_state(
|
||||
&self,
|
||||
event_type: &str,
|
||||
data: &str,
|
||||
state: &mut AnthropicState,
|
||||
) -> Result<Vec<Event>, ClientError> {
|
||||
let Some(parsed_event_type) = AnthropicEventType::parse(event_type) else {
|
||||
return Ok(Vec::new());
|
||||
};
|
||||
|
||||
// signature_delta はストリーム表示には流さず、state にだけ蓄積。
|
||||
// それ以外は parse_event で標準 Event 化する。
|
||||
let mut emitted: Vec<Event> = Vec::new();
|
||||
|
||||
match parsed_event_type {
|
||||
AnthropicEventType::ContentBlockStart => {
|
||||
let raw: ContentBlockStartEvent = serde_json::from_str(data)?;
|
||||
state.current_block_type = Some(match &raw.content_block {
|
||||
ContentBlock::Text { .. } => BlockType::Text,
|
||||
ContentBlock::Thinking { .. } | ContentBlock::RedactedThinking { .. } => {
|
||||
BlockType::Thinking
|
||||
}
|
||||
ContentBlock::ToolUse { .. } => BlockType::ToolUse,
|
||||
});
|
||||
match &raw.content_block {
|
||||
ContentBlock::Thinking {
|
||||
thinking,
|
||||
signature,
|
||||
} => {
|
||||
state.pending_thinking = Some(PendingThinking {
|
||||
text: thinking.clone(),
|
||||
signature: signature.clone(),
|
||||
redacted_data: None,
|
||||
});
|
||||
}
|
||||
ContentBlock::RedactedThinking { data: blob } => {
|
||||
state.pending_thinking = Some(PendingThinking {
|
||||
text: String::new(),
|
||||
signature: None,
|
||||
redacted_data: Some(blob.clone()),
|
||||
});
|
||||
}
|
||||
_ => {}
|
||||
}
|
||||
emitted.push(self.convert_block_start(&raw));
|
||||
}
|
||||
AnthropicEventType::ContentBlockDelta => {
|
||||
let raw: ContentBlockDeltaEvent = serde_json::from_str(data)?;
|
||||
match &raw.delta {
|
||||
DeltaBlock::ThinkingDelta { thinking } => {
|
||||
if let Some(pending) = state.pending_thinking.as_mut() {
|
||||
pending.text.push_str(thinking);
|
||||
}
|
||||
emitted.push(Event::BlockDelta(BlockDelta {
|
||||
index: raw.index,
|
||||
delta: DeltaContent::Thinking(thinking.clone()),
|
||||
}));
|
||||
}
|
||||
DeltaBlock::SignatureDelta { signature } => {
|
||||
if let Some(pending) = state.pending_thinking.as_mut() {
|
||||
// 通常 1 回しか来ないが、複数 fragment 来ても連結しておく
|
||||
match &mut pending.signature {
|
||||
Some(acc) => acc.push_str(signature),
|
||||
None => pending.signature = Some(signature.clone()),
|
||||
}
|
||||
}
|
||||
}
|
||||
DeltaBlock::TextDelta { text } => {
|
||||
emitted.push(Event::BlockDelta(BlockDelta {
|
||||
index: raw.index,
|
||||
delta: DeltaContent::Text(text.clone()),
|
||||
}));
|
||||
}
|
||||
DeltaBlock::InputJsonDelta { partial_json } => {
|
||||
emitted.push(Event::BlockDelta(BlockDelta {
|
||||
index: raw.index,
|
||||
delta: DeltaContent::InputJson(partial_json.clone()),
|
||||
}));
|
||||
}
|
||||
}
|
||||
}
|
||||
AnthropicEventType::ContentBlockStop => {
|
||||
let raw: ContentBlockStopEvent = serde_json::from_str(data)?;
|
||||
let block_type = state.current_block_type.take().unwrap_or(BlockType::Text);
|
||||
let reasoning = if matches!(block_type, BlockType::Thinking) {
|
||||
state
|
||||
.pending_thinking
|
||||
.take()
|
||||
.map(PendingThinking::into_reasoning)
|
||||
} else {
|
||||
state.pending_thinking.take();
|
||||
None
|
||||
};
|
||||
emitted.push(Event::BlockStop(BlockStop {
|
||||
index: raw.index,
|
||||
block_type,
|
||||
stop_reason: None,
|
||||
reasoning,
|
||||
}));
|
||||
}
|
||||
// 残りは state を必要としない。既存 parse_event に委譲。
|
||||
_ => {
|
||||
if let Some(event) = self.parse_event(event_type, data)? {
|
||||
emitted.push(event);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
Ok(emitted)
|
||||
}
|
||||
|
||||
fn convert_usage(&self, usage: &UsageData) -> UsageEvent {
|
||||
// Anthropic の `input_tokens` は **キャッシュ外** の入力トークンのみで、
|
||||
// プロンプト全長は input_tokens + cache_read + cache_creation。
|
||||
// UsageEvent の `input_tokens` には「占有量（プロンプト全長）」を載せる
|
||||
// 規約に合わせて、ここでキャッシュ分を足し込む。
|
||||
// cache_read_input_tokens / cache_creation_input_tokens は内訳として
|
||||
// 別フィールドに残るので、料金計算側で `input - cache_read - cache_creation`
|
||||
// により非キャッシュ入力分は逆算可能。
|
||||
let raw_input = usage.input_tokens.unwrap_or(0);
|
||||
let cache_read = usage.cache_read_input_tokens.unwrap_or(0);
|
||||
let cache_creation = usage.cache_creation_input_tokens.unwrap_or(0);
|
||||
let input_total = raw_input + cache_read + cache_creation;
|
||||
let output = usage.output_tokens.unwrap_or(0);
|
||||
|
||||
UsageEvent {
|
||||
input_tokens: usage.input_tokens.map(|_| input_total),
|
||||
output_tokens: usage.output_tokens,
|
||||
total_tokens: Some(input_total + output),
|
||||
cache_read_input_tokens: usage.cache_read_input_tokens,
|
||||
cache_creation_input_tokens: usage.cache_creation_input_tokens,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
|
||||
#[test]
|
||||
fn test_parse_message_start() {
|
||||
let scheme = AnthropicScheme::new();
|
||||
let data = r#"{"type":"message_start","message":{"id":"msg_123","type":"message","role":"assistant","content":[],"model":"claude-sonnet-4-20250514","stop_reason":null,"stop_sequence":null,"usage":{"input_tokens":10,"output_tokens":0}}}"#;
|
||||
|
||||
let event = scheme.parse_event("message_start", data).unwrap().unwrap();
|
||||
match event {
|
||||
Event::Usage(u) => {
|
||||
// キャッシュなしなので input_total = raw_input = 10
|
||||
assert_eq!(u.input_tokens, Some(10));
|
||||
}
|
||||
_ => panic!("Expected Usage event"),
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_convert_usage_includes_cache_in_input_total() {
|
||||
// Anthropic の input_tokens はキャッシュ外のみで、占有量は
|
||||
// input + cache_read + cache_creation。
|
||||
// UsageEvent.input_tokens は占有量に正規化される。
|
||||
let scheme = AnthropicScheme::new();
|
||||
let usage = UsageData {
|
||||
input_tokens: Some(100),
|
||||
output_tokens: Some(50),
|
||||
cache_read_input_tokens: Some(800),
|
||||
cache_creation_input_tokens: Some(200),
|
||||
};
|
||||
let event = scheme.convert_usage(&usage);
|
||||
// 100 + 800 + 200 = 1100
|
||||
assert_eq!(event.input_tokens, Some(1100));
|
||||
assert_eq!(event.cache_read_input_tokens, Some(800));
|
||||
assert_eq!(event.cache_creation_input_tokens, Some(200));
|
||||
assert_eq!(event.total_tokens, Some(1150));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_parse_content_block_start_text() {
|
||||
let scheme = AnthropicScheme::new();
|
||||
let data =
|
||||
r#"{"type":"content_block_start","index":0,"content_block":{"type":"text","text":""}}"#;
|
||||
|
||||
let event = scheme
|
||||
.parse_event("content_block_start", data)
|
||||
.unwrap()
|
||||
.unwrap();
|
||||
match event {
|
||||
Event::BlockStart(s) => {
|
||||
assert_eq!(s.index, 0);
|
||||
assert_eq!(s.block_type, BlockType::Text);
|
||||
}
|
||||
_ => panic!("Expected BlockStart event"),
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_parse_content_block_delta_text() {
|
||||
let scheme = AnthropicScheme::new();
|
||||
let data = r#"{"type":"content_block_delta","index":0,"delta":{"type":"text_delta","text":"Hello"}}"#;
|
||||
|
||||
let event = scheme
|
||||
.parse_event("content_block_delta", data)
|
||||
.unwrap()
|
||||
.unwrap();
|
||||
match event {
|
||||
Event::BlockDelta(d) => {
|
||||
assert_eq!(d.index, 0);
|
||||
match d.delta {
|
||||
DeltaContent::Text(t) => assert_eq!(t, "Hello"),
|
||||
_ => panic!("Expected Text delta"),
|
||||
}
|
||||
}
|
||||
_ => panic!("Expected BlockDelta event"),
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_parse_tool_use_start() {
|
||||
let scheme = AnthropicScheme::new();
|
||||
let data = r#"{"type":"content_block_start","index":1,"content_block":{"type":"tool_use","id":"toolu_123","name":"get_weather","input":{}}}"#;
|
||||
|
||||
let event = scheme
|
||||
.parse_event("content_block_start", data)
|
||||
.unwrap()
|
||||
.unwrap();
|
||||
match event {
|
||||
Event::BlockStart(s) => {
|
||||
assert_eq!(s.block_type, BlockType::ToolUse);
|
||||
match s.metadata {
|
||||
BlockMetadata::ToolUse { id, name } => {
|
||||
assert_eq!(id, "toolu_123");
|
||||
assert_eq!(name, "get_weather");
|
||||
}
|
||||
_ => panic!("Expected ToolUse metadata"),
|
||||
}
|
||||
}
|
||||
_ => panic!("Expected BlockStart event"),
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn thinking_block_stop_carries_reasoning_with_signature() {
|
||||
// thinking ブロックが完了したら reasoning metadata に text+signature が乗ること
|
||||
let scheme = AnthropicScheme::new();
|
||||
let mut state = AnthropicState::default();
|
||||
|
||||
let evs = scheme
|
||||
.parse_with_state(
|
||||
"content_block_start",
|
||||
r#"{"type":"content_block_start","index":0,"content_block":{"type":"thinking","thinking":""}}"#,
|
||||
&mut state,
|
||||
)
|
||||
.unwrap();
|
||||
assert!(matches!(evs[0], Event::BlockStart(_)));
|
||||
|
||||
scheme
|
||||
.parse_with_state(
|
||||
"content_block_delta",
|
||||
r#"{"type":"content_block_delta","index":0,"delta":{"type":"thinking_delta","thinking":"hello "}}"#,
|
||||
&mut state,
|
||||
)
|
||||
.unwrap();
|
||||
scheme
|
||||
.parse_with_state(
|
||||
"content_block_delta",
|
||||
r#"{"type":"content_block_delta","index":0,"delta":{"type":"thinking_delta","thinking":"world"}}"#,
|
||||
&mut state,
|
||||
)
|
||||
.unwrap();
|
||||
scheme
|
||||
.parse_with_state(
|
||||
"content_block_delta",
|
||||
r#"{"type":"content_block_delta","index":0,"delta":{"type":"signature_delta","signature":"SIG-XYZ"}}"#,
|
||||
&mut state,
|
||||
)
|
||||
.unwrap();
|
||||
|
||||
let stop_evs = scheme
|
||||
.parse_with_state(
|
||||
"content_block_stop",
|
||||
r#"{"type":"content_block_stop","index":0}"#,
|
||||
&mut state,
|
||||
)
|
||||
.unwrap();
|
||||
assert_eq!(stop_evs.len(), 1);
|
||||
let Event::BlockStop(stop) = &stop_evs[0] else {
|
||||
panic!("expected BlockStop, got {:?}", stop_evs[0]);
|
||||
};
|
||||
let reasoning = stop.reasoning.as_ref().expect("reasoning metadata");
|
||||
assert_eq!(reasoning.text.as_deref(), Some("hello world"));
|
||||
assert_eq!(reasoning.signature.as_deref(), Some("SIG-XYZ"));
|
||||
assert!(reasoning.encrypted_content.is_none());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn redacted_thinking_stop_carries_reasoning_with_data() {
|
||||
let scheme = AnthropicScheme::new();
|
||||
let mut state = AnthropicState::default();
|
||||
|
||||
scheme
|
||||
.parse_with_state(
|
||||
"content_block_start",
|
||||
r#"{"type":"content_block_start","index":0,"content_block":{"type":"redacted_thinking","data":"opaque-blob"}}"#,
|
||||
&mut state,
|
||||
)
|
||||
.unwrap();
|
||||
let stop_evs = scheme
|
||||
.parse_with_state(
|
||||
"content_block_stop",
|
||||
r#"{"type":"content_block_stop","index":0}"#,
|
||||
&mut state,
|
||||
)
|
||||
.unwrap();
|
||||
assert_eq!(stop_evs.len(), 1);
|
||||
let Event::BlockStop(stop) = &stop_evs[0] else {
|
||||
panic!("expected BlockStop");
|
||||
};
|
||||
let reasoning = stop.reasoning.as_ref().expect("reasoning metadata");
|
||||
assert_eq!(reasoning.text.as_deref(), Some(""));
|
||||
assert!(reasoning.signature.is_none());
|
||||
assert_eq!(reasoning.encrypted_content.as_deref(), Some("opaque-blob"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn text_block_stop_has_no_reasoning_metadata() {
|
||||
let scheme = AnthropicScheme::new();
|
||||
let mut state = AnthropicState::default();
|
||||
|
||||
scheme
|
||||
.parse_with_state(
|
||||
"content_block_start",
|
||||
r#"{"type":"content_block_start","index":0,"content_block":{"type":"text","text":""}}"#,
|
||||
&mut state,
|
||||
)
|
||||
.unwrap();
|
||||
scheme
|
||||
.parse_with_state(
|
||||
"content_block_delta",
|
||||
r#"{"type":"content_block_delta","index":0,"delta":{"type":"text_delta","text":"hi"}}"#,
|
||||
&mut state,
|
||||
)
|
||||
.unwrap();
|
||||
let stop_evs = scheme
|
||||
.parse_with_state(
|
||||
"content_block_stop",
|
||||
r#"{"type":"content_block_stop","index":0}"#,
|
||||
&mut state,
|
||||
)
|
||||
.unwrap();
|
||||
assert_eq!(stop_evs.len(), 1);
|
||||
let Event::BlockStop(stop) = &stop_evs[0] else {
|
||||
panic!("expected BlockStop");
|
||||
};
|
||||
assert!(stop.reasoning.is_none());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_parse_ping() {
|
||||
let scheme = AnthropicScheme::new();
|
||||
let data = r#"{"type":"ping"}"#;
|
||||
|
||||
let event = scheme.parse_event("ping", data).unwrap().unwrap();
|
||||
match event {
|
||||
Event::Ping(_) => {}
|
||||
_ => panic!("Expected Ping event"),
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,44 @@
|
||||
//! Anthropic Messages API スキーマ
|
||||
//!
|
||||
//! - リクエストJSON生成
|
||||
//! - SSEイベントパース → Event変換
|
||||
|
||||
mod capability;
|
||||
mod events;
|
||||
mod request;
|
||||
mod scheme_impl;
|
||||
|
||||
pub use scheme_impl::AnthropicState;
|
||||
|
||||
/// Anthropicスキーマ
|
||||
///
|
||||
/// Anthropic Messages APIのリクエスト/レスポンス変換を担当
|
||||
#[derive(Debug, Clone)]
|
||||
pub struct AnthropicScheme {
|
||||
/// APIバージョン
|
||||
pub api_version: String,
|
||||
/// 細粒度ツールストリーミングを有効にするか
|
||||
pub fine_grained_tool_streaming: bool,
|
||||
}
|
||||
|
||||
impl Default for AnthropicScheme {
|
||||
fn default() -> Self {
|
||||
Self {
|
||||
api_version: "2023-06-01".to_string(),
|
||||
fine_grained_tool_streaming: true,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
impl AnthropicScheme {
|
||||
/// 新しいスキーマを作成
|
||||
pub fn new() -> Self {
|
||||
Self::default()
|
||||
}
|
||||
|
||||
/// 細粒度ツールストリーミングを有効/無効にする
|
||||
pub fn with_fine_grained_tool_streaming(mut self, enabled: bool) -> Self {
|
||||
self.fine_grained_tool_streaming = enabled;
|
||||
self
|
||||
}
|
||||
}
|
||||
File diff suppressed because it is too large
Load Diff
@@ -0,0 +1,107 @@
|
||||
//! `impl Scheme for AnthropicScheme`
|
||||
//!
|
||||
//! Anthropic Messages API の wire 表現に必要な URL・ヘッダ・SSE パース・
|
||||
//! リクエスト body 生成を共通 `Scheme` trait にぶら下げる。
|
||||
|
||||
use serde_json::Value;
|
||||
|
||||
use crate::llm_client::{
|
||||
ClientError,
|
||||
auth::AuthRequirement,
|
||||
capability::ModelCapability,
|
||||
event::{BlockType, Event, ReasoningBlockData},
|
||||
scheme::Scheme,
|
||||
types::Request,
|
||||
};
|
||||
|
||||
use super::AnthropicScheme;
|
||||
|
||||
/// Anthropic の SSE パースで必要な状態。
|
||||
///
|
||||
/// 1. `content_block_stop` イベントは `block_type` を持たない仕様なので、
|
||||
/// 直前の `content_block_start` で観測した `block_type` を保持して
|
||||
/// `BlockStop` に書き戻す。
|
||||
/// 2. `thinking` ブロック中の `thinking_delta` テキストと `signature_delta`
|
||||
/// 署名、および `redacted_thinking` ブロックの `data` を蓄積し、
|
||||
/// `content_block_stop` の Thinking block metadata として返す
|
||||
/// （round-trip 永続化のため）。
|
||||
#[derive(Debug, Default)]
|
||||
pub struct AnthropicState {
|
||||
pub(crate) current_block_type: Option<BlockType>,
|
||||
pub(crate) pending_thinking: Option<PendingThinking>,
|
||||
}
|
||||
|
||||
/// 1 つの `thinking` または `redacted_thinking` content_block の蓄積バッファ。
|
||||
#[derive(Debug, Default)]
|
||||
pub(crate) struct PendingThinking {
|
||||
pub(crate) text: String,
|
||||
pub(crate) signature: Option<String>,
|
||||
pub(crate) redacted_data: Option<String>,
|
||||
}
|
||||
|
||||
impl PendingThinking {
|
||||
pub(crate) fn into_reasoning(self) -> ReasoningBlockData {
|
||||
ReasoningBlockData {
|
||||
id: None,
|
||||
text: Some(self.text),
|
||||
summary: Vec::new(),
|
||||
encrypted_content: self.redacted_data,
|
||||
signature: self.signature,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
impl Scheme for AnthropicScheme {
|
||||
type State = AnthropicState;
|
||||
|
||||
fn default_base_url(&self) -> &'static str {
|
||||
"https://api.anthropic.com"
|
||||
}
|
||||
|
||||
fn path(&self, _model_id: &str) -> String {
|
||||
"/v1/messages".to_string()
|
||||
}
|
||||
|
||||
fn required_auth(&self) -> AuthRequirement {
|
||||
// Ollama の `/v1/messages` 互換では認証が要らないが、それは
|
||||
// `AuthRef::None` + `build_headers` 側の「ResolvedAuth::None
|
||||
// なら何もしない」分岐で吸収する（`accepts` 判定で弾かれない
|
||||
// よう、現状は XApiKey を要求しつつ、None 側でもパスするよう
|
||||
// にする戦略）。
|
||||
AuthRequirement::XApiKey
|
||||
}
|
||||
|
||||
fn additional_headers(&self) -> Vec<(&'static str, String)> {
|
||||
let mut headers = vec![("anthropic-version", self.api_version.clone())];
|
||||
if self.fine_grained_tool_streaming {
|
||||
headers.push((
|
||||
"anthropic-beta",
|
||||
"fine-grained-tool-streaming-2025-05-14".to_string(),
|
||||
));
|
||||
}
|
||||
headers
|
||||
}
|
||||
|
||||
fn build_request_body(
|
||||
&self,
|
||||
model_id: &str,
|
||||
request: &Request,
|
||||
capability: &ModelCapability,
|
||||
) -> Value {
|
||||
let req = self.build_request(model_id, request, capability);
|
||||
serde_json::to_value(&req).expect("AnthropicRequest is always serialisable")
|
||||
}
|
||||
|
||||
fn parse_sse(
|
||||
&self,
|
||||
event_type: &str,
|
||||
data: &str,
|
||||
state: &mut Self::State,
|
||||
) -> Result<Vec<Event>, ClientError> {
|
||||
self.parse_with_state(event_type, data, state)
|
||||
}
|
||||
|
||||
fn default_capability(&self) -> ModelCapability {
|
||||
super::capability::default_capability()
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,20 @@
|
||||
//! Gemini scheme の wire-level 既定 capability。
|
||||
//!
|
||||
//! モデル ID 固有のテーブル(`gemini-*` バージョン別の reasoning 有無)は
|
||||
//! 高レベル構築層(`provider::capability`)の責務。ここでは wire の
|
||||
//! 保守的 default のみ。
|
||||
|
||||
use crate::llm_client::capability::{
|
||||
CacheStrategy, ModelCapability, StructuredOutput, ToolCallingSupport,
|
||||
};
|
||||
|
||||
/// Scheme 既定の capability(未知モデル / 未明示モデル用)。
|
||||
pub(crate) fn default_capability() -> ModelCapability {
|
||||
ModelCapability {
|
||||
tool_calling: ToolCallingSupport::Parallel,
|
||||
structured_output: StructuredOutput::JsonSchema,
|
||||
reasoning: None,
|
||||
vision: true,
|
||||
prompt_caching: CacheStrategy::Auto,
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,329 @@
|
||||
//! Gemini SSEイベントパース
|
||||
//!
|
||||
//! Google Gemini APIのSSEイベントをパースし、統一Event型に変換
|
||||
|
||||
use crate::llm_client::{
|
||||
ClientError,
|
||||
event::{BlockMetadata, BlockStart, BlockStop, BlockType, Event, StopReason, UsageEvent},
|
||||
};
|
||||
use serde::Deserialize;
|
||||
|
||||
use super::GeminiScheme;
|
||||
|
||||
// ============================================================================
|
||||
// SSEイベントのJSON構造
|
||||
// ============================================================================
|
||||
|
||||
/// Gemini GenerateContentResponse (ストリーミングチャンク)
|
||||
#[allow(dead_code)]
|
||||
#[derive(Debug, Deserialize)]
|
||||
#[serde(rename_all = "camelCase")]
|
||||
pub(crate) struct GenerateContentResponse {
|
||||
/// 候補
|
||||
pub candidates: Option<Vec<Candidate>>,
|
||||
/// 使用量メタデータ
|
||||
pub usage_metadata: Option<UsageMetadata>,
|
||||
/// プロンプトフィードバック
|
||||
pub prompt_feedback: Option<PromptFeedback>,
|
||||
/// モデルバージョン
|
||||
pub model_version: Option<String>,
|
||||
}
|
||||
|
||||
/// 候補
|
||||
#[allow(dead_code)]
|
||||
#[derive(Debug, Deserialize)]
|
||||
#[serde(rename_all = "camelCase")]
|
||||
pub(crate) struct Candidate {
|
||||
/// コンテンツ
|
||||
pub content: Option<CandidateContent>,
|
||||
/// 完了理由
|
||||
pub finish_reason: Option<String>,
|
||||
/// インデックス
|
||||
pub index: Option<usize>,
|
||||
/// 安全性評価
|
||||
pub safety_ratings: Option<Vec<SafetyRating>>,
|
||||
}
|
||||
|
||||
/// 候補コンテンツ
|
||||
#[allow(dead_code)]
|
||||
#[derive(Debug, Deserialize)]
|
||||
pub(crate) struct CandidateContent {
|
||||
/// パーツ
|
||||
pub parts: Option<Vec<CandidatePart>>,
|
||||
/// ロール
|
||||
pub role: Option<String>,
|
||||
}
|
||||
|
||||
/// 候補パーツ
|
||||
#[derive(Debug, Deserialize)]
|
||||
#[serde(rename_all = "camelCase")]
|
||||
pub(crate) struct CandidatePart {
|
||||
/// テキスト
|
||||
pub text: Option<String>,
|
||||
/// 関数呼び出し
|
||||
pub function_call: Option<FunctionCall>,
|
||||
}
|
||||
|
||||
/// 関数呼び出し
|
||||
#[derive(Debug, Deserialize)]
|
||||
pub(crate) struct FunctionCall {
|
||||
/// 関数名
|
||||
pub name: String,
|
||||
/// 引数
|
||||
pub args: Option<serde_json::Value>,
|
||||
}
|
||||
|
||||
/// 使用量メタデータ
|
||||
#[derive(Debug, Deserialize)]
|
||||
#[serde(rename_all = "camelCase")]
|
||||
pub(crate) struct UsageMetadata {
|
||||
/// プロンプトトークン数
|
||||
pub prompt_token_count: Option<u64>,
|
||||
/// 候補トークン数
|
||||
pub candidates_token_count: Option<u64>,
|
||||
/// 合計トークン数
|
||||
pub total_token_count: Option<u64>,
|
||||
}
|
||||
|
||||
/// プロンプトフィードバック
|
||||
#[allow(dead_code)]
|
||||
#[derive(Debug, Deserialize)]
|
||||
#[serde(rename_all = "camelCase")]
|
||||
pub(crate) struct PromptFeedback {
|
||||
/// ブロック理由
|
||||
pub block_reason: Option<String>,
|
||||
/// 安全性評価
|
||||
pub safety_ratings: Option<Vec<SafetyRating>>,
|
||||
}
|
||||
|
||||
/// 安全性評価
|
||||
#[allow(dead_code)]
|
||||
#[derive(Debug, Deserialize)]
|
||||
pub(crate) struct SafetyRating {
|
||||
/// カテゴリ
|
||||
pub category: Option<String>,
|
||||
/// 確率
|
||||
pub probability: Option<String>,
|
||||
}
|
||||
|
||||
// ============================================================================
|
||||
// イベント変換
|
||||
// ============================================================================
|
||||
|
||||
impl GeminiScheme {
|
||||
/// SSEデータをEvent型に変換
|
||||
///
|
||||
/// # Arguments
|
||||
/// * `data` - SSEイベントデータJSON文字列
|
||||
///
|
||||
/// # Returns
|
||||
/// * `Ok(Some(Vec<Event>))` - 変換成功
|
||||
/// * `Ok(None)` - イベントを無視
|
||||
/// * `Err(ClientError)` - パースエラー
|
||||
pub(crate) fn parse_event(&self, data: &str) -> Result<Option<Vec<Event>>, ClientError> {
|
||||
// データが空または無効な場合はスキップ
|
||||
if data.is_empty() || data == "[DONE]" {
|
||||
return Ok(None);
|
||||
}
|
||||
|
||||
let response: GenerateContentResponse =
|
||||
serde_json::from_str(data).map_err(|e| ClientError::Api {
|
||||
status: None,
|
||||
code: Some("parse_error".to_string()),
|
||||
message: format!("Failed to parse Gemini SSE data: {} -> {}", e, data),
|
||||
retry_after: None,
|
||||
})?;
|
||||
|
||||
let mut events = Vec::new();
|
||||
|
||||
// 使用量メタデータ
|
||||
if let Some(usage) = response.usage_metadata {
|
||||
events.push(self.convert_usage(&usage));
|
||||
}
|
||||
|
||||
// 候補を処理
|
||||
if let Some(candidates) = response.candidates {
|
||||
for candidate in candidates {
|
||||
let candidate_index = candidate.index.unwrap_or(0);
|
||||
|
||||
if let Some(content) = candidate.content {
|
||||
if let Some(parts) = content.parts {
|
||||
for (part_index, part) in parts.iter().enumerate() {
|
||||
// テキストデルタ
|
||||
if let Some(text) = &part.text {
|
||||
if !text.is_empty() {
|
||||
// Geminiは明示的なBlockStartを送らないため、
|
||||
// TextDeltaを直接送る（Timelineが暗黙的に開始を処理）
|
||||
events.push(Event::text_delta(part_index, text.clone()));
|
||||
}
|
||||
}
|
||||
|
||||
// 関数呼び出し
|
||||
if let Some(function_call) = &part.function_call {
|
||||
// 関数呼び出しの開始
|
||||
// Geminiでは関数呼び出しは一度に送られることが多い
|
||||
// ストリーミング引数が有効な場合は部分的に送られる可能性がある
|
||||
|
||||
// 関数呼び出しIDはGeminiにはないので、名前をIDとして使用
|
||||
let function_id = format!("call_{}", function_call.name);
|
||||
|
||||
events.push(Event::BlockStart(BlockStart {
|
||||
index: candidate_index * 10 + part_index, // 複合インデックス
|
||||
block_type: BlockType::ToolUse,
|
||||
metadata: BlockMetadata::ToolUse {
|
||||
id: function_id,
|
||||
name: function_call.name.clone(),
|
||||
},
|
||||
}));
|
||||
|
||||
// 引数がある場合はデルタとして送る
|
||||
if let Some(args) = &function_call.args {
|
||||
let args_str = serde_json::to_string(args).unwrap_or_default();
|
||||
if !args_str.is_empty() && args_str != "null" {
|
||||
events.push(Event::tool_input_delta(
|
||||
candidate_index * 10 + part_index,
|
||||
args_str,
|
||||
));
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// 完了理由
|
||||
if let Some(finish_reason) = candidate.finish_reason {
|
||||
let stop_reason = match finish_reason.as_str() {
|
||||
"STOP" => Some(StopReason::EndTurn),
|
||||
"MAX_TOKENS" => Some(StopReason::MaxTokens),
|
||||
"SAFETY" | "RECITATION" | "OTHER" => Some(StopReason::EndTurn),
|
||||
_ => None,
|
||||
};
|
||||
|
||||
// テキストブロックの停止
|
||||
events.push(Event::BlockStop(BlockStop {
|
||||
index: candidate_index,
|
||||
block_type: BlockType::Text,
|
||||
stop_reason,
|
||||
reasoning: None,
|
||||
}));
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
if events.is_empty() {
|
||||
Ok(None)
|
||||
} else {
|
||||
Ok(Some(events))
|
||||
}
|
||||
}
|
||||
|
||||
fn convert_usage(&self, usage: &UsageMetadata) -> Event {
|
||||
Event::Usage(UsageEvent {
|
||||
input_tokens: usage.prompt_token_count,
|
||||
output_tokens: usage.candidates_token_count,
|
||||
total_tokens: usage.total_token_count,
|
||||
cache_read_input_tokens: None,
|
||||
cache_creation_input_tokens: None,
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
use crate::llm_client::event::DeltaContent;
|
||||
|
||||
#[test]
|
||||
fn test_parse_text_response() {
|
||||
let scheme = GeminiScheme::new();
|
||||
let data =
|
||||
r#"{"candidates":[{"content":{"parts":[{"text":"Hello"}],"role":"model"},"index":0}]}"#;
|
||||
|
||||
let events = scheme.parse_event(data).unwrap().unwrap();
|
||||
assert_eq!(events.len(), 1);
|
||||
|
||||
if let Event::BlockDelta(delta) = &events[0] {
|
||||
assert_eq!(delta.index, 0);
|
||||
if let DeltaContent::Text(text) = &delta.delta {
|
||||
assert_eq!(text, "Hello");
|
||||
} else {
|
||||
panic!("Expected text delta");
|
||||
}
|
||||
} else {
|
||||
panic!("Expected BlockDelta");
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_parse_usage_metadata() {
|
||||
let scheme = GeminiScheme::new();
|
||||
let data = r#"{"candidates":[{"content":{"parts":[{"text":"Hi"}],"role":"model"},"index":0}],"usageMetadata":{"promptTokenCount":10,"candidatesTokenCount":5,"totalTokenCount":15}}"#;
|
||||
|
||||
let events = scheme.parse_event(data).unwrap().unwrap();
|
||||
|
||||
// Usageイベントが含まれるはず
|
||||
let usage_event = events.iter().find(|e| matches!(e, Event::Usage(_)));
|
||||
assert!(usage_event.is_some());
|
||||
|
||||
if let Event::Usage(usage) = usage_event.unwrap() {
|
||||
assert_eq!(usage.input_tokens, Some(10));
|
||||
assert_eq!(usage.output_tokens, Some(5));
|
||||
assert_eq!(usage.total_tokens, Some(15));
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_parse_function_call() {
|
||||
let scheme = GeminiScheme::new();
|
||||
let data = r#"{"candidates":[{"content":{"parts":[{"functionCall":{"name":"get_weather","args":{"location":"Tokyo"}}}],"role":"model"},"index":0}]}"#;
|
||||
|
||||
let events = scheme.parse_event(data).unwrap().unwrap();
|
||||
|
||||
// BlockStartイベントがあるはず
|
||||
let start_event = events.iter().find(|e| matches!(e, Event::BlockStart(_)));
|
||||
assert!(start_event.is_some());
|
||||
|
||||
if let Event::BlockStart(start) = start_event.unwrap() {
|
||||
assert_eq!(start.block_type, BlockType::ToolUse);
|
||||
if let BlockMetadata::ToolUse { id: _, name } = &start.metadata {
|
||||
assert_eq!(name, "get_weather");
|
||||
} else {
|
||||
panic!("Expected ToolUse metadata");
|
||||
}
|
||||
}
|
||||
|
||||
// 引数デルタもあるはず
|
||||
let delta_event = events.iter().find(|e| {
|
||||
if let Event::BlockDelta(d) = e {
|
||||
matches!(d.delta, DeltaContent::InputJson(_))
|
||||
} else {
|
||||
false
|
||||
}
|
||||
});
|
||||
assert!(delta_event.is_some());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_parse_finish_reason() {
|
||||
let scheme = GeminiScheme::new();
|
||||
let data = r#"{"candidates":[{"content":{"parts":[{"text":"Done"}],"role":"model"},"finishReason":"STOP","index":0}]}"#;
|
||||
|
||||
let events = scheme.parse_event(data).unwrap().unwrap();
|
||||
|
||||
// BlockStopイベントがあるはず
|
||||
let stop_event = events.iter().find(|e| matches!(e, Event::BlockStop(_)));
|
||||
assert!(stop_event.is_some());
|
||||
|
||||
if let Event::BlockStop(stop) = stop_event.unwrap() {
|
||||
assert_eq!(stop.stop_reason, Some(StopReason::EndTurn));
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_parse_empty_data() {
|
||||
let scheme = GeminiScheme::new();
|
||||
assert!(scheme.parse_event("").unwrap().is_none());
|
||||
assert!(scheme.parse_event("[DONE]").unwrap().is_none());
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,31 @@
|
||||
//! Google Gemini API スキーマ
|
||||
//!
|
||||
//! - リクエストJSON生成
|
||||
//! - SSEイベントパース → Event変換
|
||||
|
||||
mod capability;
|
||||
mod events;
|
||||
mod request;
|
||||
mod scheme_impl;
|
||||
|
||||
/// Geminiスキーマ
|
||||
///
|
||||
/// Google Gemini APIのリクエスト/レスポンス変換を担当
|
||||
#[derive(Debug, Clone, Default)]
|
||||
pub struct GeminiScheme {
|
||||
/// ストリーミング関数呼び出し引数を有効にするか
|
||||
pub stream_function_call_arguments: bool,
|
||||
}
|
||||
|
||||
impl GeminiScheme {
|
||||
/// 新しいスキーマを作成
|
||||
pub fn new() -> Self {
|
||||
Self::default()
|
||||
}
|
||||
|
||||
/// ストリーミング関数呼び出し引数を有効/無効にする
|
||||
pub fn with_stream_function_call_arguments(mut self, enabled: bool) -> Self {
|
||||
self.stream_function_call_arguments = enabled;
|
||||
self
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,496 @@
|
||||
//! Gemini Request Builder
|
||||
//!
|
||||
//! Converts Open Responses native Item model to Google Gemini API format.
|
||||
|
||||
use serde::Serialize;
|
||||
use serde_json::Value;
|
||||
|
||||
use crate::llm_client::{
|
||||
Request,
|
||||
capability::{ModelCapability, ReasoningControl, ReasoningSupport},
|
||||
types::{Item, Role, ToolDefinition, parse_tool_arguments},
|
||||
};
|
||||
|
||||
use super::GeminiScheme;
|
||||
|
||||
/// Gemini API request body
|
||||
#[derive(Debug, Serialize)]
|
||||
#[serde(rename_all = "camelCase")]
|
||||
pub(crate) struct GeminiRequest {
|
||||
/// Contents (conversation history)
|
||||
pub contents: Vec<GeminiContent>,
|
||||
/// System instruction
|
||||
#[serde(skip_serializing_if = "Option::is_none")]
|
||||
pub system_instruction: Option<GeminiContent>,
|
||||
/// Tool definitions
|
||||
#[serde(skip_serializing_if = "Vec::is_empty")]
|
||||
pub tools: Vec<GeminiTool>,
|
||||
/// Tool config
|
||||
#[serde(skip_serializing_if = "Option::is_none")]
|
||||
pub tool_config: Option<GeminiToolConfig>,
|
||||
/// Generation config
|
||||
#[serde(skip_serializing_if = "Option::is_none")]
|
||||
pub generation_config: Option<GeminiGenerationConfig>,
|
||||
}
|
||||
|
||||
/// Gemini content
|
||||
#[derive(Debug, Serialize)]
|
||||
pub(crate) struct GeminiContent {
|
||||
/// Role
|
||||
pub role: String,
|
||||
/// Parts
|
||||
pub parts: Vec<GeminiPart>,
|
||||
}
|
||||
|
||||
/// Gemini part
|
||||
#[derive(Debug, Serialize)]
|
||||
#[serde(untagged)]
|
||||
pub(crate) enum GeminiPart {
|
||||
/// Text part
|
||||
Text { text: String },
|
||||
/// Function call part
|
||||
FunctionCall {
|
||||
#[serde(rename = "functionCall")]
|
||||
function_call: GeminiFunctionCall,
|
||||
},
|
||||
/// Function response part
|
||||
FunctionResponse {
|
||||
#[serde(rename = "functionResponse")]
|
||||
function_response: GeminiFunctionResponse,
|
||||
},
|
||||
}
|
||||
|
||||
/// Gemini function call
|
||||
#[derive(Debug, Serialize)]
|
||||
pub(crate) struct GeminiFunctionCall {
|
||||
pub name: String,
|
||||
pub args: Value,
|
||||
}
|
||||
|
||||
/// Gemini function response
|
||||
#[derive(Debug, Serialize)]
|
||||
pub(crate) struct GeminiFunctionResponse {
|
||||
pub name: String,
|
||||
pub response: GeminiFunctionResponseContent,
|
||||
}
|
||||
|
||||
/// Gemini function response content
|
||||
#[derive(Debug, Serialize)]
|
||||
pub(crate) struct GeminiFunctionResponseContent {
|
||||
pub name: String,
|
||||
pub content: Value,
|
||||
}
|
||||
|
||||
/// Gemini tool definition
|
||||
#[derive(Debug, Serialize)]
|
||||
#[serde(rename_all = "camelCase")]
|
||||
pub(crate) struct GeminiTool {
|
||||
/// Function declarations
|
||||
pub function_declarations: Vec<GeminiFunctionDeclaration>,
|
||||
}
|
||||
|
||||
/// Gemini function declaration
|
||||
#[derive(Debug, Serialize)]
|
||||
pub(crate) struct GeminiFunctionDeclaration {
|
||||
/// Function name
|
||||
pub name: String,
|
||||
/// Description
|
||||
#[serde(skip_serializing_if = "Option::is_none")]
|
||||
pub description: Option<String>,
|
||||
/// Parameter schema
|
||||
pub parameters: Value,
|
||||
}
|
||||
|
||||
/// Gemini tool config
|
||||
#[derive(Debug, Serialize)]
|
||||
#[serde(rename_all = "camelCase")]
|
||||
pub(crate) struct GeminiToolConfig {
|
||||
/// Function calling config
|
||||
pub function_calling_config: GeminiFunctionCallingConfig,
|
||||
}
|
||||
|
||||
/// Gemini function calling config
|
||||
#[derive(Debug, Serialize)]
|
||||
#[serde(rename_all = "camelCase")]
|
||||
pub(crate) struct GeminiFunctionCallingConfig {
|
||||
/// Mode: AUTO, ANY, NONE
|
||||
#[serde(skip_serializing_if = "Option::is_none")]
|
||||
pub mode: Option<String>,
|
||||
/// Enable streaming function call arguments
|
||||
#[serde(skip_serializing_if = "Option::is_none")]
|
||||
pub stream_function_call_arguments: Option<bool>,
|
||||
}
|
||||
|
||||
/// Gemini generation config
|
||||
#[derive(Debug, Serialize)]
|
||||
#[serde(rename_all = "camelCase")]
|
||||
pub(crate) struct GeminiGenerationConfig {
|
||||
/// Max output tokens
|
||||
#[serde(skip_serializing_if = "Option::is_none")]
|
||||
pub max_output_tokens: Option<u32>,
|
||||
/// Temperature
|
||||
#[serde(skip_serializing_if = "Option::is_none")]
|
||||
pub temperature: Option<f32>,
|
||||
/// Top P
|
||||
#[serde(skip_serializing_if = "Option::is_none")]
|
||||
pub top_p: Option<f32>,
|
||||
/// Top K
|
||||
#[serde(skip_serializing_if = "Option::is_none")]
|
||||
pub top_k: Option<u32>,
|
||||
/// Stop sequences
|
||||
#[serde(skip_serializing_if = "Vec::is_empty")]
|
||||
pub stop_sequences: Vec<String>,
|
||||
/// Thinking / reasoning 設定（Gemini 2.5 以降）。
|
||||
#[serde(skip_serializing_if = "Option::is_none")]
|
||||
pub thinking_config: Option<GeminiThinkingConfig>,
|
||||
}
|
||||
|
||||
/// Gemini thinking config (gemini-2.5 以降)
|
||||
#[derive(Debug, Serialize)]
|
||||
#[serde(rename_all = "camelCase")]
|
||||
pub(crate) struct GeminiThinkingConfig {
|
||||
/// Token budget for thinking. `-1` means dynamic.
|
||||
pub thinking_budget: i32,
|
||||
}
|
||||
|
||||
impl GeminiScheme {
|
||||
/// Build Gemini request from Request
|
||||
pub(crate) fn build_request(
|
||||
&self,
|
||||
request: &Request,
|
||||
capability: &ModelCapability,
|
||||
) -> GeminiRequest {
|
||||
let contents = self.convert_items_to_contents(&request.items);
|
||||
|
||||
// System prompt
|
||||
let system_instruction = request.system_prompt.as_ref().map(|s| GeminiContent {
|
||||
role: "user".to_string(),
|
||||
parts: vec![GeminiPart::Text { text: s.clone() }],
|
||||
});
|
||||
|
||||
// Tools
|
||||
let tools = if request.tools.is_empty() {
|
||||
vec![]
|
||||
} else {
|
||||
vec![GeminiTool {
|
||||
function_declarations: request.tools.iter().map(|t| self.convert_tool(t)).collect(),
|
||||
}]
|
||||
};
|
||||
|
||||
// Tool config
|
||||
let tool_config = if !request.tools.is_empty() {
|
||||
Some(GeminiToolConfig {
|
||||
function_calling_config: GeminiFunctionCallingConfig {
|
||||
mode: Some("AUTO".to_string()),
|
||||
stream_function_call_arguments: if self.stream_function_call_arguments {
|
||||
Some(true)
|
||||
} else {
|
||||
None
|
||||
},
|
||||
},
|
||||
})
|
||||
} else {
|
||||
None
|
||||
};
|
||||
|
||||
// Reasoning の投影: capability が BudgetTokens / Both をサポートし、
|
||||
// request 側で budget_tokens が指定されているときだけ thinking_config を付ける。
|
||||
let supports_budget = matches!(
|
||||
capability.reasoning,
|
||||
Some(ReasoningSupport::BudgetTokens | ReasoningSupport::Both),
|
||||
);
|
||||
let thinking_config = request
|
||||
.config
|
||||
.reasoning
|
||||
.as_ref()
|
||||
.filter(|_| supports_budget)
|
||||
.and_then(|rc| match rc {
|
||||
ReasoningControl::BudgetTokens(budget) => Some(GeminiThinkingConfig {
|
||||
thinking_budget: *budget,
|
||||
}),
|
||||
ReasoningControl::Effort(_) => None,
|
||||
});
|
||||
|
||||
// Generation config
|
||||
let generation_config = Some(GeminiGenerationConfig {
|
||||
max_output_tokens: request.config.max_tokens,
|
||||
temperature: request.config.temperature,
|
||||
top_p: request.config.top_p,
|
||||
top_k: request.config.top_k,
|
||||
stop_sequences: request.config.stop_sequences.clone(),
|
||||
thinking_config,
|
||||
});
|
||||
|
||||
GeminiRequest {
|
||||
contents,
|
||||
system_instruction,
|
||||
tools,
|
||||
tool_config,
|
||||
generation_config,
|
||||
}
|
||||
}
|
||||
|
||||
/// Convert Open Responses Items to Gemini Contents
|
||||
///
|
||||
/// Gemini uses:
|
||||
/// - role "user" for user messages and function responses
|
||||
/// - role "model" for assistant messages and function calls
|
||||
fn convert_items_to_contents(&self, items: &[Item]) -> Vec<GeminiContent> {
|
||||
let mut contents = Vec::new();
|
||||
let mut pending_model_parts: Vec<GeminiPart> = Vec::new();
|
||||
let mut pending_user_parts: Vec<GeminiPart> = Vec::new();
|
||||
|
||||
for item in items {
|
||||
match item {
|
||||
Item::Message { role, content, .. } => {
|
||||
// Flush pending parts
|
||||
self.flush_pending_parts(
|
||||
&mut contents,
|
||||
&mut pending_model_parts,
|
||||
&mut pending_user_parts,
|
||||
);
|
||||
|
||||
let gemini_role = match role {
|
||||
Role::User | Role::System => "user",
|
||||
Role::Assistant => "model",
|
||||
};
|
||||
|
||||
let parts: Vec<GeminiPart> = content
|
||||
.iter()
|
||||
.map(|p| GeminiPart::Text {
|
||||
text: p.as_text().to_string(),
|
||||
})
|
||||
.collect();
|
||||
|
||||
contents.push(GeminiContent {
|
||||
role: gemini_role.to_string(),
|
||||
parts,
|
||||
});
|
||||
}
|
||||
|
||||
Item::ToolCall {
|
||||
name, arguments, ..
|
||||
} => {
|
||||
// Flush pending user parts first
|
||||
if !pending_user_parts.is_empty() {
|
||||
contents.push(GeminiContent {
|
||||
role: "user".to_string(),
|
||||
parts: std::mem::take(&mut pending_user_parts),
|
||||
});
|
||||
}
|
||||
|
||||
// Parse arguments (normalize non-object / legacy "null" payloads to {})
|
||||
let args = parse_tool_arguments(arguments);
|
||||
|
||||
pending_model_parts.push(GeminiPart::FunctionCall {
|
||||
function_call: GeminiFunctionCall {
|
||||
name: name.clone(),
|
||||
args,
|
||||
},
|
||||
});
|
||||
}
|
||||
|
||||
Item::ToolResult {
|
||||
call_id,
|
||||
summary,
|
||||
content,
|
||||
..
|
||||
} => {
|
||||
// Flush pending model parts first
|
||||
if !pending_model_parts.is_empty() {
|
||||
contents.push(GeminiContent {
|
||||
role: "model".to_string(),
|
||||
parts: std::mem::take(&mut pending_model_parts),
|
||||
});
|
||||
}
|
||||
|
||||
let text = match content {
|
||||
Some(c) => format!("{summary}\n{c}"),
|
||||
None => summary.clone(),
|
||||
};
|
||||
pending_user_parts.push(GeminiPart::FunctionResponse {
|
||||
function_response: GeminiFunctionResponse {
|
||||
name: call_id.clone(),
|
||||
response: GeminiFunctionResponseContent {
|
||||
name: call_id.clone(),
|
||||
content: Value::String(text),
|
||||
},
|
||||
},
|
||||
});
|
||||
}
|
||||
|
||||
Item::Reasoning { text, .. } => {
|
||||
// Flush pending user parts first
|
||||
if !pending_user_parts.is_empty() {
|
||||
contents.push(GeminiContent {
|
||||
role: "user".to_string(),
|
||||
parts: std::mem::take(&mut pending_user_parts),
|
||||
});
|
||||
}
|
||||
|
||||
// Reasoning is treated as model text in Gemini
|
||||
pending_model_parts.push(GeminiPart::Text { text: text.clone() });
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// Flush remaining pending parts
|
||||
self.flush_pending_parts(
|
||||
&mut contents,
|
||||
&mut pending_model_parts,
|
||||
&mut pending_user_parts,
|
||||
);
|
||||
|
||||
contents
|
||||
}
|
||||
|
||||
fn flush_pending_parts(
|
||||
&self,
|
||||
contents: &mut Vec<GeminiContent>,
|
||||
pending_model_parts: &mut Vec<GeminiPart>,
|
||||
pending_user_parts: &mut Vec<GeminiPart>,
|
||||
) {
|
||||
if !pending_model_parts.is_empty() {
|
||||
contents.push(GeminiContent {
|
||||
role: "model".to_string(),
|
||||
parts: std::mem::take(pending_model_parts),
|
||||
});
|
||||
}
|
||||
if !pending_user_parts.is_empty() {
|
||||
contents.push(GeminiContent {
|
||||
role: "user".to_string(),
|
||||
parts: std::mem::take(pending_user_parts),
|
||||
});
|
||||
}
|
||||
}
|
||||
|
||||
fn convert_tool(&self, tool: &ToolDefinition) -> GeminiFunctionDeclaration {
|
||||
GeminiFunctionDeclaration {
|
||||
name: tool.name.clone(),
|
||||
description: tool.description.clone(),
|
||||
parameters: tool.input_schema.clone(),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
use crate::llm_client::capability::{
|
||||
CacheStrategy, ReasoningEffort, StructuredOutput, ToolCallingSupport,
|
||||
};
|
||||
|
||||
fn cap() -> ModelCapability {
|
||||
ModelCapability {
|
||||
tool_calling: ToolCallingSupport::Parallel,
|
||||
structured_output: StructuredOutput::JsonSchema,
|
||||
reasoning: None,
|
||||
vision: true,
|
||||
prompt_caching: CacheStrategy::Auto,
|
||||
}
|
||||
}
|
||||
|
||||
fn cap_budget_reasoning() -> ModelCapability {
|
||||
ModelCapability {
|
||||
reasoning: Some(ReasoningSupport::BudgetTokens),
|
||||
..cap()
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_build_simple_request() {
|
||||
let scheme = GeminiScheme::new();
|
||||
let request = Request::new()
|
||||
.system("You are a helpful assistant.")
|
||||
.user("Hello!");
|
||||
|
||||
let gemini_req = scheme.build_request(&request, &cap());
|
||||
|
||||
assert!(gemini_req.system_instruction.is_some());
|
||||
assert_eq!(gemini_req.contents.len(), 1);
|
||||
assert_eq!(gemini_req.contents[0].role, "user");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_build_request_with_tool() {
|
||||
let scheme = GeminiScheme::new();
|
||||
let request = Request::new().user("What's the weather?").tool(
|
||||
ToolDefinition::new("get_weather")
|
||||
.description("Get current weather")
|
||||
.input_schema(serde_json::json!({
|
||||
"type": "object",
|
||||
"properties": {
|
||||
"location": { "type": "string" }
|
||||
},
|
||||
"required": ["location"]
|
||||
})),
|
||||
);
|
||||
|
||||
let gemini_req = scheme.build_request(&request, &cap());
|
||||
|
||||
assert_eq!(gemini_req.tools.len(), 1);
|
||||
assert_eq!(gemini_req.tools[0].function_declarations.len(), 1);
|
||||
assert_eq!(
|
||||
gemini_req.tools[0].function_declarations[0].name,
|
||||
"get_weather"
|
||||
);
|
||||
assert!(gemini_req.tool_config.is_some());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_assistant_role_is_model() {
|
||||
let scheme = GeminiScheme::new();
|
||||
let request = Request::new().user("Hello").assistant("Hi there!");
|
||||
|
||||
let gemini_req = scheme.build_request(&request, &cap());
|
||||
|
||||
assert_eq!(gemini_req.contents.len(), 2);
|
||||
assert_eq!(gemini_req.contents[0].role, "user");
|
||||
assert_eq!(gemini_req.contents[1].role, "model");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_tool_call_and_result() {
|
||||
let scheme = GeminiScheme::new();
|
||||
let request = Request::new()
|
||||
.user("What's the weather?")
|
||||
.item(Item::tool_call(
|
||||
"call_123",
|
||||
"get_weather",
|
||||
r#"{"city":"Tokyo"}"#,
|
||||
))
|
||||
.item(Item::tool_result("call_123", "Sunny, 25°C"));
|
||||
|
||||
let gemini_req = scheme.build_request(&request, &cap());
|
||||
|
||||
assert_eq!(gemini_req.contents.len(), 3);
|
||||
assert_eq!(gemini_req.contents[0].role, "user");
|
||||
assert_eq!(gemini_req.contents[1].role, "model");
|
||||
assert_eq!(gemini_req.contents[2].role, "user");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn thinking_budget_projected_when_supported() {
|
||||
let scheme = GeminiScheme::new();
|
||||
let mut request = Request::new().user("think");
|
||||
request.config.reasoning = Some(ReasoningControl::BudgetTokens(-1));
|
||||
|
||||
let gemini_req = scheme.build_request(&request, &cap_budget_reasoning());
|
||||
let config = gemini_req.generation_config.expect("generation config");
|
||||
let thinking = config.thinking_config.expect("thinking config");
|
||||
|
||||
assert_eq!(thinking.thinking_budget, -1);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn effort_reasoning_not_projected_to_gemini() {
|
||||
let scheme = GeminiScheme::new();
|
||||
let mut request = Request::new().user("think");
|
||||
request.config.reasoning = Some(ReasoningControl::Effort(ReasoningEffort::Medium));
|
||||
|
||||
let gemini_req = scheme.build_request(&request, &cap_budget_reasoning());
|
||||
let config = gemini_req.generation_config.expect("generation config");
|
||||
|
||||
assert!(config.thinking_config.is_none());
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,49 @@
|
||||
//! `impl Scheme for GeminiScheme`
|
||||
|
||||
use serde_json::Value;
|
||||
|
||||
use crate::llm_client::{
|
||||
ClientError, auth::AuthRequirement, capability::ModelCapability, event::Event, scheme::Scheme,
|
||||
types::Request,
|
||||
};
|
||||
|
||||
use super::GeminiScheme;
|
||||
|
||||
impl Scheme for GeminiScheme {
|
||||
type State = ();
|
||||
|
||||
fn default_base_url(&self) -> &'static str {
|
||||
"https://generativelanguage.googleapis.com"
|
||||
}
|
||||
|
||||
fn path(&self, model_id: &str) -> String {
|
||||
format!("/v1beta/models/{model_id}:streamGenerateContent?alt=sse")
|
||||
}
|
||||
|
||||
fn required_auth(&self) -> AuthRequirement {
|
||||
AuthRequirement::QueryParam { name: "key" }
|
||||
}
|
||||
|
||||
fn build_request_body(
|
||||
&self,
|
||||
_model_id: &str,
|
||||
request: &Request,
|
||||
capability: &ModelCapability,
|
||||
) -> Value {
|
||||
let req = self.build_request(request, capability);
|
||||
serde_json::to_value(&req).expect("GeminiRequest is always serialisable")
|
||||
}
|
||||
|
||||
fn parse_sse(
|
||||
&self,
|
||||
_event_type: &str,
|
||||
data: &str,
|
||||
_state: &mut Self::State,
|
||||
) -> Result<Vec<Event>, ClientError> {
|
||||
Ok(self.parse_event(data)?.unwrap_or_default())
|
||||
}
|
||||
|
||||
fn default_capability(&self) -> ModelCapability {
|
||||
super::capability::default_capability()
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,92 @@
|
||||
//! APIスキーマ定義
|
||||
//!
|
||||
//! 各APIスキーマごとの変換ロジック
|
||||
//! - リクエスト変換: Request → プロバイダ固有JSON
|
||||
//! - レスポンス変換: SSEイベント → Event
|
||||
//!
|
||||
//! [`Scheme`] trait により `HttpTransport<S>` から scheme 固有の差分
|
||||
//! （パス、ヘッダ、認証要件、body 生成、SSE パース）をすべて委譲する。
|
||||
|
||||
pub mod anthropic;
|
||||
pub mod gemini;
|
||||
pub mod openai_chat;
|
||||
pub mod openai_responses;
|
||||
|
||||
use serde_json::Value;
|
||||
|
||||
use super::auth::AuthRequirement;
|
||||
use super::capability::ModelCapability;
|
||||
use super::client::ConfigWarning;
|
||||
use super::error::ClientError;
|
||||
use super::event::Event;
|
||||
use super::types::{Request, RequestConfig};
|
||||
|
||||
/// wire scheme の抽象。各プロバイダの API 仕様ごとに 1 つ実装する。
|
||||
///
|
||||
/// `HttpTransport<S: Scheme>` が URL 組立・認証ヘッダ挿入・SSE パース
|
||||
/// のループを担い、`Scheme` 実装は各仕様固有の差分のみ提供する。
|
||||
///
|
||||
/// # 状態
|
||||
///
|
||||
/// SSE パースでフレーム間に状態を保つ必要がある scheme（Anthropic の
|
||||
/// `BlockStop` に `block_type` が載らない仕様の補完など）は
|
||||
/// [`Scheme::State`] に中間状態を表す型を置く。
|
||||
/// 状態を持たない scheme は `type State = ()` とする。
|
||||
pub trait Scheme: Clone + Send + Sync + 'static {
|
||||
/// SSE パースのフレーム間で共有する状態。`HttpTransport` が
|
||||
/// ストリーム開始時に `Default::default()` を一度だけ作り、
|
||||
/// フレームごとに `&mut` で渡す。
|
||||
type State: Default + Send + 'static;
|
||||
|
||||
/// scheme のベース URL（`ModelConfig::base_url` 未指定時のデフォルト）
|
||||
fn default_base_url(&self) -> &'static str;
|
||||
|
||||
/// リクエスト先の相対パス。Gemini のようにモデル名をパスに埋め込む
|
||||
/// プロバイダもあるため、モデル ID を受け取る。
|
||||
fn path(&self, model_id: &str) -> String;
|
||||
|
||||
/// この scheme が要求する認証形式。`build_client` 時に
|
||||
/// `manifest::AuthRef` と照合する。
|
||||
fn required_auth(&self) -> AuthRequirement;
|
||||
|
||||
/// `Content-Type` 以外の追加ヘッダ。`anthropic-version` / `anthropic-beta` 等。
|
||||
fn additional_headers(&self) -> Vec<(&'static str, String)> {
|
||||
Vec::new()
|
||||
}
|
||||
|
||||
/// リクエスト body を生成する。`capability` は `CacheStrategy` や
|
||||
/// `ReasoningSupport` を参照して scheme 側の挙動を分岐させるため
|
||||
/// に渡される。
|
||||
fn build_request_body(
|
||||
&self,
|
||||
model_id: &str,
|
||||
request: &Request,
|
||||
capability: &ModelCapability,
|
||||
) -> Value;
|
||||
|
||||
/// SSE イベント 1 件を 0 個以上の [`Event`] に変換する。
|
||||
///
|
||||
/// `event_type` は SSE フレームの `event:` フィールド、`data` は
|
||||
/// `data:` フィールド。`[DONE]` 等の終端マーカーは実装側で判定する。
|
||||
/// `state` はストリーム単位で共有される可変状態。
|
||||
fn parse_sse(
|
||||
&self,
|
||||
event_type: &str,
|
||||
data: &str,
|
||||
state: &mut Self::State,
|
||||
) -> Result<Vec<Event>, ClientError>;
|
||||
|
||||
/// scheme 既定の capability。モデル ID に関係なく、この wire で
|
||||
/// 安全に送れる最小共通項を返す。既知モデル ID の能力テーブルは
|
||||
/// `provider::capability::lookup` 側(高レベル構築層)の責務で、
|
||||
/// scheme はここには関与しない。
|
||||
fn default_capability(&self) -> ModelCapability;
|
||||
|
||||
/// scheme 側でサポートしていない `RequestConfig` フィールドを
|
||||
/// 警告として返す（例: OpenAI Chat は `top_k` 非対応）。
|
||||
/// デフォルトは空 Vec。
|
||||
fn validate_config(&self, config: &RequestConfig) -> Vec<ConfigWarning> {
|
||||
let _ = config;
|
||||
Vec::new()
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,20 @@
|
||||
//! OpenAI Chat Completions scheme の wire-level 既定 capability。
|
||||
//!
|
||||
//! モデル ID 固有のテーブル(`gpt-5` 系など)は高レベル構築層
|
||||
//! (`provider::capability`)の責務。ここでは wire の保守的 default のみ。
|
||||
|
||||
use crate::llm_client::capability::{
|
||||
CacheStrategy, ModelCapability, StructuredOutput, ToolCallingSupport,
|
||||
};
|
||||
|
||||
/// Scheme 既定の capability。OpenAI 互換ルーター系(xAI / Groq / OpenRouter 等)
|
||||
/// で未知モデル ID を受けたときのフォールバックに使う。
|
||||
pub(crate) fn default_capability() -> ModelCapability {
|
||||
ModelCapability {
|
||||
tool_calling: ToolCallingSupport::Parallel,
|
||||
structured_output: StructuredOutput::JsonSchema,
|
||||
reasoning: None,
|
||||
vision: false,
|
||||
prompt_caching: CacheStrategy::Auto,
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,213 @@
|
||||
//! OpenAI SSEイベントパース
|
||||
|
||||
use crate::llm_client::{
|
||||
ClientError,
|
||||
event::{Event, StopReason, UsageEvent},
|
||||
};
|
||||
use serde::Deserialize;
|
||||
|
||||
use super::OpenAIScheme;
|
||||
|
||||
/// OpenAI Streaming Chat Response Chunk
|
||||
#[allow(dead_code)]
|
||||
#[derive(Debug, Deserialize)]
|
||||
pub(crate) struct ChatCompletionChunk {
|
||||
pub id: String,
|
||||
pub object: String,
|
||||
pub created: u64,
|
||||
pub model: String,
|
||||
pub choices: Vec<ChunkChoice>,
|
||||
pub usage: Option<ChunkUsage>,
|
||||
}
|
||||
|
||||
#[allow(dead_code)]
|
||||
#[derive(Debug, Deserialize)]
|
||||
pub(crate) struct ChunkChoice {
|
||||
pub index: usize,
|
||||
pub delta: ChunkDelta,
|
||||
pub finish_reason: Option<String>,
|
||||
}
|
||||
|
||||
#[allow(dead_code)]
|
||||
#[derive(Debug, Deserialize)]
|
||||
pub(crate) struct ChunkDelta {
|
||||
pub role: Option<String>,
|
||||
pub content: Option<String>,
|
||||
pub tool_calls: Option<Vec<ChunkToolCall>>,
|
||||
}
|
||||
|
||||
#[allow(dead_code)]
|
||||
#[derive(Debug, Deserialize)]
|
||||
pub(crate) struct ChunkToolCall {
|
||||
pub index: usize,
|
||||
pub id: Option<String>,
|
||||
#[serde(rename = "type")]
|
||||
pub call_type: Option<String>,
|
||||
pub function: Option<ChunkFunction>,
|
||||
}
|
||||
|
||||
#[allow(dead_code)]
|
||||
#[derive(Debug, Deserialize)]
|
||||
pub(crate) struct ChunkFunction {
|
||||
pub name: Option<String>,
|
||||
pub arguments: Option<String>,
|
||||
}
|
||||
|
||||
#[derive(Debug, Deserialize)]
|
||||
pub(crate) struct ChunkUsage {
|
||||
pub prompt_tokens: u64,
|
||||
pub completion_tokens: u64,
|
||||
pub total_tokens: u64,
|
||||
}
|
||||
|
||||
impl OpenAIScheme {
|
||||
/// SSEデータのパースとEventへの変換
|
||||
///
|
||||
/// OpenAI APIはBlockStartイベントを明示的に送信しない。
|
||||
/// Timeline層が暗黙的なBlockStartを処理する。
|
||||
pub fn parse_event(&self, data: &str) -> Result<Option<Vec<Event>>, ClientError> {
|
||||
if data == "[DONE]" {
|
||||
return Ok(None);
|
||||
}
|
||||
|
||||
let chunk: ChatCompletionChunk =
|
||||
serde_json::from_str(data).map_err(|e| ClientError::Api {
|
||||
status: None,
|
||||
code: Some("parse_error".to_string()),
|
||||
message: format!("Failed to parse SSE data: {} -> {}", e, data),
|
||||
retry_after: None,
|
||||
})?;
|
||||
|
||||
let mut events = Vec::new();
|
||||
|
||||
// Usage handling
|
||||
if let Some(usage) = chunk.usage {
|
||||
events.push(Event::Usage(UsageEvent {
|
||||
input_tokens: Some(usage.prompt_tokens),
|
||||
output_tokens: Some(usage.completion_tokens),
|
||||
total_tokens: Some(usage.total_tokens),
|
||||
cache_read_input_tokens: None,
|
||||
cache_creation_input_tokens: None,
|
||||
}));
|
||||
}
|
||||
|
||||
for choice in chunk.choices {
|
||||
// Text Content Delta
|
||||
if let Some(content) = choice.delta.content {
|
||||
// OpenAI APIはBlockStartを送らないため、デルタのみを発行
|
||||
// Timeline層が暗黙的なBlockStartを処理する
|
||||
events.push(Event::text_delta(choice.index, content));
|
||||
}
|
||||
|
||||
// Tool Call Delta
|
||||
if let Some(tool_calls) = choice.delta.tool_calls {
|
||||
for tool_call in tool_calls {
|
||||
// Start of tool call (has ID)
|
||||
if let Some(id) = tool_call.id {
|
||||
let name = tool_call
|
||||
.function
|
||||
.as_ref()
|
||||
.and_then(|f| f.name.clone())
|
||||
.unwrap_or_default();
|
||||
events.push(Event::tool_use_start(tool_call.index, id, name));
|
||||
}
|
||||
|
||||
// Arguments delta
|
||||
if let Some(function) = tool_call.function {
|
||||
if let Some(args) = function.arguments {
|
||||
if !args.is_empty() {
|
||||
events.push(Event::tool_input_delta(tool_call.index, args));
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// Finish Reason
|
||||
if let Some(finish_reason) = choice.finish_reason {
|
||||
let stop_reason = match finish_reason.as_str() {
|
||||
"stop" => Some(StopReason::EndTurn),
|
||||
"length" => Some(StopReason::MaxTokens),
|
||||
"tool_calls" | "function_call" => Some(StopReason::ToolUse),
|
||||
_ => Some(StopReason::EndTurn),
|
||||
};
|
||||
|
||||
let is_tool_finish =
|
||||
finish_reason == "tool_calls" || finish_reason == "function_call";
|
||||
|
||||
if is_tool_finish {
|
||||
// ツール呼び出し終了
|
||||
// Note: OpenAIはどのツールが終了したか明示しないため、
|
||||
// Timeline層で適切に処理する必要がある
|
||||
} else {
|
||||
// テキスト終了
|
||||
events.push(Event::text_block_stop(choice.index, stop_reason));
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
if events.is_empty() {
|
||||
Ok(None)
|
||||
} else {
|
||||
Ok(Some(events))
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
use crate::llm_client::event::DeltaContent;
|
||||
|
||||
#[test]
|
||||
fn test_parse_text_delta() {
|
||||
let scheme = OpenAIScheme::new();
|
||||
let data = r#"{"id":"chatcmpl-123","object":"chat.completion.chunk","created":1694268190,"model":"gpt-4o","choices":[{"index":0,"delta":{"content":"Hello"},"finish_reason":null}]}"#;
|
||||
|
||||
let events = scheme.parse_event(data).unwrap().unwrap();
|
||||
// OpenAIはBlockStartを発行しないため、デルタのみ
|
||||
assert_eq!(events.len(), 1);
|
||||
|
||||
if let Event::BlockDelta(delta) = &events[0] {
|
||||
assert_eq!(delta.index, 0);
|
||||
if let DeltaContent::Text(text) = &delta.delta {
|
||||
assert_eq!(text, "Hello");
|
||||
} else {
|
||||
panic!("Expected text delta");
|
||||
}
|
||||
} else {
|
||||
panic!("Expected BlockDelta");
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_parse_tool_call() {
|
||||
let scheme = OpenAIScheme::new();
|
||||
// Start of tool call
|
||||
let data_start = r#"{"id":"chatcmpl-123","object":"chat.completion.chunk","created":1694268190,"model":"gpt-4o","choices":[{"index":0,"delta":{"tool_calls":[{"index":0,"id":"call_abc","type":"function","function":{"name":"get_weather","arguments":""}}]},"finish_reason":null}]}"#;
|
||||
|
||||
let events = scheme.parse_event(data_start).unwrap().unwrap();
|
||||
assert_eq!(events.len(), 1);
|
||||
if let Event::BlockStart(start) = &events[0] {
|
||||
assert_eq!(start.index, 0);
|
||||
if let crate::llm_client::event::BlockMetadata::ToolUse { id, name } = &start.metadata {
|
||||
assert_eq!(id, "call_abc");
|
||||
assert_eq!(name, "get_weather");
|
||||
} else {
|
||||
panic!("Expected ToolUse metadata");
|
||||
}
|
||||
}
|
||||
|
||||
// Tool arguments delta
|
||||
let data_arg = r#"{"id":"chatcmpl-123","object":"chat.completion.chunk","created":1694268190,"model":"gpt-4o","choices":[{"index":0,"delta":{"tool_calls":[{"index":0,"function":{"arguments":"{}}"}}]},"finish_reason":null}]}"#;
|
||||
let events = scheme.parse_event(data_arg).unwrap().unwrap();
|
||||
assert_eq!(events.len(), 1);
|
||||
if let Event::BlockDelta(delta) = &events[0] {
|
||||
if let DeltaContent::InputJson(json) = &delta.delta {
|
||||
assert_eq!(json, "{}}");
|
||||
} else {
|
||||
panic!("Expected input json delta");
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,33 @@
|
||||
//! OpenAI Chat Completions API スキーマ
|
||||
//!
|
||||
//! - リクエストJSON生成
|
||||
//! - SSEイベントパース → Event変換
|
||||
|
||||
pub(crate) mod capability;
|
||||
mod events;
|
||||
mod request;
|
||||
mod scheme_impl;
|
||||
|
||||
/// OpenAIスキーマ
|
||||
///
|
||||
/// OpenAI Chat Completions API (および互換API) のリクエスト/レスポンス変換を担当
|
||||
#[derive(Debug, Clone, Default)]
|
||||
pub struct OpenAIScheme {
|
||||
/// モデル名 (リクエスト時に指定されるが、デフォルト値として保持も可能)
|
||||
pub model: Option<String>,
|
||||
/// レガシーなmax_tokensを使用するか (Ollama互換用)
|
||||
pub use_legacy_max_tokens: bool,
|
||||
}
|
||||
|
||||
impl OpenAIScheme {
|
||||
/// 新しいスキーマを作成
|
||||
pub fn new() -> Self {
|
||||
Self::default()
|
||||
}
|
||||
|
||||
/// レガシーなmax_tokensを使用するか設定
|
||||
pub fn with_legacy_max_tokens(mut self, use_legacy: bool) -> Self {
|
||||
self.use_legacy_max_tokens = use_legacy;
|
||||
self
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,442 @@
|
||||
//! OpenAI Request Builder
|
||||
//!
|
||||
//! Converts Open Responses native Item model to OpenAI Chat Completions API format.
|
||||
|
||||
use serde::Serialize;
|
||||
use serde_json::Value;
|
||||
|
||||
use crate::llm_client::{
|
||||
Request,
|
||||
capability::{ModelCapability, ReasoningControl, ReasoningSupport},
|
||||
types::{Item, Role, ToolDefinition, parse_tool_arguments},
|
||||
};
|
||||
|
||||
use super::OpenAIScheme;
|
||||
|
||||
/// OpenAI API request body
|
||||
#[derive(Debug, Serialize)]
|
||||
pub(crate) struct OpenAIRequest {
|
||||
pub model: String,
|
||||
#[serde(skip_serializing_if = "Option::is_none")]
|
||||
pub max_completion_tokens: Option<u32>,
|
||||
#[serde(skip_serializing_if = "Option::is_none")]
|
||||
pub max_tokens: Option<u32>, // Legacy field for compatibility (e.g. Ollama)
|
||||
#[serde(skip_serializing_if = "Option::is_none")]
|
||||
pub temperature: Option<f32>,
|
||||
#[serde(skip_serializing_if = "Option::is_none")]
|
||||
pub top_p: Option<f32>,
|
||||
#[serde(skip_serializing_if = "Vec::is_empty")]
|
||||
pub stop: Vec<String>,
|
||||
pub stream: bool,
|
||||
#[serde(skip_serializing_if = "Option::is_none")]
|
||||
pub stream_options: Option<StreamOptions>,
|
||||
pub messages: Vec<OpenAIMessage>,
|
||||
#[serde(skip_serializing_if = "Vec::is_empty")]
|
||||
pub tools: Vec<OpenAITool>,
|
||||
#[serde(skip_serializing_if = "Option::is_none")]
|
||||
pub tool_choice: Option<String>,
|
||||
/// Reasoning effort（o1 / o3 / o4 / gpt-5 系で有効）。
|
||||
#[serde(skip_serializing_if = "Option::is_none")]
|
||||
pub reasoning_effort: Option<String>,
|
||||
}
|
||||
|
||||
#[derive(Debug, Serialize)]
|
||||
pub(crate) struct StreamOptions {
|
||||
pub include_usage: bool,
|
||||
}
|
||||
|
||||
/// OpenAI message
|
||||
#[derive(Debug, Serialize)]
|
||||
pub(crate) struct OpenAIMessage {
|
||||
pub role: String,
|
||||
pub content: Option<OpenAIContent>,
|
||||
#[serde(skip_serializing_if = "Vec::is_empty")]
|
||||
pub tool_calls: Vec<OpenAIToolCall>,
|
||||
#[serde(skip_serializing_if = "Option::is_none")]
|
||||
pub tool_call_id: Option<String>,
|
||||
#[serde(skip_serializing_if = "Option::is_none")]
|
||||
pub name: Option<String>,
|
||||
}
|
||||
|
||||
/// OpenAI content
|
||||
#[allow(dead_code)]
|
||||
#[derive(Debug, Serialize)]
|
||||
#[serde(untagged)]
|
||||
pub(crate) enum OpenAIContent {
|
||||
Text(String),
|
||||
Parts(Vec<OpenAIContentPart>),
|
||||
}
|
||||
|
||||
/// OpenAI content part
|
||||
#[allow(dead_code)]
|
||||
#[derive(Debug, Serialize)]
|
||||
#[serde(tag = "type")]
|
||||
pub(crate) enum OpenAIContentPart {
|
||||
#[serde(rename = "text")]
|
||||
Text { text: String },
|
||||
#[serde(rename = "image_url")]
|
||||
ImageUrl { image_url: ImageUrl },
|
||||
}
|
||||
|
||||
#[derive(Debug, Serialize)]
|
||||
pub(crate) struct ImageUrl {
|
||||
pub url: String,
|
||||
}
|
||||
|
||||
/// OpenAI tool definition
|
||||
#[derive(Debug, Serialize)]
|
||||
pub(crate) struct OpenAITool {
|
||||
pub r#type: String,
|
||||
pub function: OpenAIToolFunction,
|
||||
}
|
||||
|
||||
#[derive(Debug, Serialize)]
|
||||
pub(crate) struct OpenAIToolFunction {
|
||||
pub name: String,
|
||||
#[serde(skip_serializing_if = "Option::is_none")]
|
||||
pub description: Option<String>,
|
||||
pub parameters: Value,
|
||||
}
|
||||
|
||||
/// OpenAI tool call in message
|
||||
#[derive(Debug, Serialize)]
|
||||
pub(crate) struct OpenAIToolCall {
|
||||
pub id: String,
|
||||
pub r#type: String,
|
||||
pub function: OpenAIToolCallFunction,
|
||||
}
|
||||
|
||||
#[derive(Debug, Serialize)]
|
||||
pub(crate) struct OpenAIToolCallFunction {
|
||||
pub name: String,
|
||||
pub arguments: String,
|
||||
}
|
||||
|
||||
impl OpenAIScheme {
|
||||
/// Build OpenAI request from Request
|
||||
pub(crate) fn build_request(
|
||||
&self,
|
||||
model: &str,
|
||||
request: &Request,
|
||||
capability: &ModelCapability,
|
||||
) -> OpenAIRequest {
|
||||
let mut messages = Vec::new();
|
||||
|
||||
// Add system message if present
|
||||
if let Some(system) = &request.system_prompt {
|
||||
messages.push(OpenAIMessage {
|
||||
role: "system".to_string(),
|
||||
content: Some(OpenAIContent::Text(system.clone())),
|
||||
tool_calls: vec![],
|
||||
tool_call_id: None,
|
||||
name: None,
|
||||
});
|
||||
}
|
||||
|
||||
// Convert items to messages
|
||||
messages.extend(self.convert_items_to_messages(&request.items));
|
||||
|
||||
let tools = request.tools.iter().map(|t| self.convert_tool(t)).collect();
|
||||
|
||||
let (max_tokens, max_completion_tokens) = if self.use_legacy_max_tokens {
|
||||
(request.config.max_tokens, None)
|
||||
} else {
|
||||
(None, request.config.max_tokens)
|
||||
};
|
||||
|
||||
// Reasoning の投影: capability が Effort / Both をサポートし、
|
||||
// request 側で effort が指定されているときだけ reasoning_effort を付ける。
|
||||
let supports_effort = matches!(
|
||||
capability.reasoning,
|
||||
Some(ReasoningSupport::Effort | ReasoningSupport::Both),
|
||||
);
|
||||
let reasoning_effort = request
|
||||
.config
|
||||
.reasoning
|
||||
.as_ref()
|
||||
.filter(|_| supports_effort)
|
||||
.and_then(|rc| match rc {
|
||||
ReasoningControl::Effort(effort) => Some(effort.as_str().to_string()),
|
||||
ReasoningControl::BudgetTokens(_) => None,
|
||||
});
|
||||
|
||||
OpenAIRequest {
|
||||
model: model.to_string(),
|
||||
max_completion_tokens,
|
||||
max_tokens,
|
||||
temperature: request.config.temperature,
|
||||
top_p: request.config.top_p,
|
||||
stop: request.config.stop_sequences.clone(),
|
||||
stream: true,
|
||||
stream_options: Some(StreamOptions {
|
||||
include_usage: true,
|
||||
}),
|
||||
messages,
|
||||
tools,
|
||||
tool_choice: None,
|
||||
reasoning_effort,
|
||||
}
|
||||
}
|
||||
|
||||
/// Convert Open Responses Items to OpenAI Messages
|
||||
///
|
||||
/// OpenAI uses a message-based model where:
|
||||
/// - User messages have role "user"
|
||||
/// - Assistant messages have role "assistant"
|
||||
/// - Tool calls are within assistant messages as tool_calls array
|
||||
/// - Tool results have role "tool" with tool_call_id
|
||||
fn convert_items_to_messages(&self, items: &[Item]) -> Vec<OpenAIMessage> {
|
||||
let mut messages = Vec::new();
|
||||
let mut pending_tool_calls: Vec<OpenAIToolCall> = Vec::new();
|
||||
let mut pending_assistant_text: Option<String> = None;
|
||||
|
||||
for item in items {
|
||||
match item {
|
||||
Item::Message { role, content, .. } => {
|
||||
// Flush pending tool calls
|
||||
self.flush_pending_assistant(
|
||||
&mut messages,
|
||||
&mut pending_tool_calls,
|
||||
&mut pending_assistant_text,
|
||||
);
|
||||
|
||||
let openai_role = match role {
|
||||
Role::User => "user",
|
||||
Role::Assistant => "assistant",
|
||||
Role::System => "system",
|
||||
};
|
||||
|
||||
let text_content: String = content
|
||||
.iter()
|
||||
.map(|p| p.as_text())
|
||||
.collect::<Vec<_>>()
|
||||
.join("");
|
||||
|
||||
messages.push(OpenAIMessage {
|
||||
role: openai_role.to_string(),
|
||||
content: Some(OpenAIContent::Text(text_content)),
|
||||
tool_calls: vec![],
|
||||
tool_call_id: None,
|
||||
name: None,
|
||||
});
|
||||
}
|
||||
|
||||
Item::ToolCall {
|
||||
call_id,
|
||||
name,
|
||||
arguments,
|
||||
..
|
||||
} => {
|
||||
// Normalize non-object / legacy "null" payloads to "{}" so
|
||||
// OpenAI gets a valid JSON object string.
|
||||
let normalized_args = parse_tool_arguments(arguments).to_string();
|
||||
pending_tool_calls.push(OpenAIToolCall {
|
||||
id: call_id.clone(),
|
||||
r#type: "function".to_string(),
|
||||
function: OpenAIToolCallFunction {
|
||||
name: name.clone(),
|
||||
arguments: normalized_args,
|
||||
},
|
||||
});
|
||||
}
|
||||
|
||||
Item::ToolResult {
|
||||
call_id,
|
||||
summary,
|
||||
content,
|
||||
..
|
||||
} => {
|
||||
// Flush pending tool calls before tool result
|
||||
self.flush_pending_assistant(
|
||||
&mut messages,
|
||||
&mut pending_tool_calls,
|
||||
&mut pending_assistant_text,
|
||||
);
|
||||
|
||||
let text = match content {
|
||||
Some(c) => format!("{summary}\n{c}"),
|
||||
None => summary.clone(),
|
||||
};
|
||||
messages.push(OpenAIMessage {
|
||||
role: "tool".to_string(),
|
||||
content: Some(OpenAIContent::Text(text)),
|
||||
tool_calls: vec![],
|
||||
tool_call_id: Some(call_id.clone()),
|
||||
name: None,
|
||||
});
|
||||
}
|
||||
|
||||
Item::Reasoning { text, .. } => {
|
||||
// Reasoning is treated as assistant text in OpenAI
|
||||
// (OpenAI doesn't have native reasoning support like Claude)
|
||||
if let Some(ref mut existing) = pending_assistant_text {
|
||||
existing.push_str(text);
|
||||
} else {
|
||||
pending_assistant_text = Some(text.clone());
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// Flush remaining pending items
|
||||
self.flush_pending_assistant(
|
||||
&mut messages,
|
||||
&mut pending_tool_calls,
|
||||
&mut pending_assistant_text,
|
||||
);
|
||||
|
||||
messages
|
||||
}
|
||||
|
||||
fn flush_pending_assistant(
|
||||
&self,
|
||||
messages: &mut Vec<OpenAIMessage>,
|
||||
pending_tool_calls: &mut Vec<OpenAIToolCall>,
|
||||
pending_assistant_text: &mut Option<String>,
|
||||
) {
|
||||
if !pending_tool_calls.is_empty() || pending_assistant_text.is_some() {
|
||||
messages.push(OpenAIMessage {
|
||||
role: "assistant".to_string(),
|
||||
content: pending_assistant_text.take().map(OpenAIContent::Text),
|
||||
tool_calls: std::mem::take(pending_tool_calls),
|
||||
tool_call_id: None,
|
||||
name: None,
|
||||
});
|
||||
}
|
||||
}
|
||||
|
||||
fn convert_tool(&self, tool: &ToolDefinition) -> OpenAITool {
|
||||
OpenAITool {
|
||||
r#type: "function".to_string(),
|
||||
function: OpenAIToolFunction {
|
||||
name: tool.name.clone(),
|
||||
description: tool.description.clone(),
|
||||
parameters: tool.input_schema.clone(),
|
||||
},
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
use crate::llm_client::capability::{
|
||||
CacheStrategy, ReasoningEffort, StructuredOutput, ToolCallingSupport,
|
||||
};
|
||||
|
||||
fn cap() -> ModelCapability {
|
||||
ModelCapability {
|
||||
tool_calling: ToolCallingSupport::Parallel,
|
||||
structured_output: StructuredOutput::JsonSchema,
|
||||
reasoning: None,
|
||||
vision: false,
|
||||
prompt_caching: CacheStrategy::Auto,
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_build_simple_request() {
|
||||
let scheme = OpenAIScheme::new();
|
||||
let request = Request::new().system("System prompt").user("Hello");
|
||||
|
||||
let body = scheme.build_request("gpt-4o", &request, &cap());
|
||||
|
||||
assert_eq!(body.model, "gpt-4o");
|
||||
assert_eq!(body.messages.len(), 2);
|
||||
assert_eq!(body.messages[0].role, "system");
|
||||
assert_eq!(body.messages[1].role, "user");
|
||||
|
||||
if let Some(OpenAIContent::Text(text)) = &body.messages[0].content {
|
||||
assert_eq!(text, "System prompt");
|
||||
} else {
|
||||
panic!("Expected text content");
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_build_request_with_tool() {
|
||||
let scheme = OpenAIScheme::new();
|
||||
let request = Request::new()
|
||||
.user("Check weather")
|
||||
.tool(ToolDefinition::new("weather").description("Get weather"));
|
||||
|
||||
let body = scheme.build_request("gpt-4o", &request, &cap());
|
||||
assert_eq!(body.tools.len(), 1);
|
||||
assert_eq!(body.tools[0].function.name, "weather");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_build_request_legacy_max_tokens() {
|
||||
let scheme = OpenAIScheme::new().with_legacy_max_tokens(true);
|
||||
let request = Request::new().user("Hello").max_tokens(100);
|
||||
|
||||
let body = scheme.build_request("llama3", &request, &cap());
|
||||
|
||||
assert_eq!(body.max_tokens, Some(100));
|
||||
assert!(body.max_completion_tokens.is_none());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_build_request_modern_max_tokens() {
|
||||
let scheme = OpenAIScheme::new();
|
||||
let request = Request::new().user("Hello").max_tokens(100);
|
||||
|
||||
let body = scheme.build_request("gpt-4o", &request, &cap());
|
||||
|
||||
assert_eq!(body.max_completion_tokens, Some(100));
|
||||
assert!(body.max_tokens.is_none());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn reasoning_effort_projected_when_supported() {
|
||||
let scheme = OpenAIScheme::new();
|
||||
let mut request = Request::new().user("Hello");
|
||||
request.config.reasoning = Some(ReasoningControl::Effort(ReasoningEffort::Other(
|
||||
"provider-native".into(),
|
||||
)));
|
||||
let capability = ModelCapability {
|
||||
reasoning: Some(ReasoningSupport::Effort),
|
||||
..cap()
|
||||
};
|
||||
|
||||
let body = scheme.build_request("gpt-5", &request, &capability);
|
||||
|
||||
assert_eq!(body.reasoning_effort.as_deref(), Some("provider-native"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn budget_reasoning_not_projected_to_openai_chat() {
|
||||
let scheme = OpenAIScheme::new();
|
||||
let mut request = Request::new().user("Hello");
|
||||
request.config.reasoning = Some(ReasoningControl::BudgetTokens(4096));
|
||||
let capability = ModelCapability {
|
||||
reasoning: Some(ReasoningSupport::Both),
|
||||
..cap()
|
||||
};
|
||||
|
||||
let body = scheme.build_request("gpt-5", &request, &capability);
|
||||
|
||||
assert!(body.reasoning_effort.is_none());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_tool_call_and_result() {
|
||||
let scheme = OpenAIScheme::new();
|
||||
let request = Request::new()
|
||||
.user("Check weather")
|
||||
.item(Item::tool_call(
|
||||
"call_123",
|
||||
"get_weather",
|
||||
r#"{"city":"Tokyo"}"#,
|
||||
))
|
||||
.item(Item::tool_result("call_123", "Sunny, 25°C"));
|
||||
|
||||
let body = scheme.build_request("gpt-4o", &request, &cap());
|
||||
|
||||
assert_eq!(body.messages.len(), 3);
|
||||
assert_eq!(body.messages[0].role, "user");
|
||||
assert_eq!(body.messages[1].role, "assistant");
|
||||
assert_eq!(body.messages[1].tool_calls.len(), 1);
|
||||
assert_eq!(body.messages[2].role, "tool");
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,67 @@
|
||||
//! `impl Scheme for OpenAIScheme`
|
||||
|
||||
use serde_json::Value;
|
||||
|
||||
use crate::llm_client::{
|
||||
ClientError,
|
||||
auth::AuthRequirement,
|
||||
capability::ModelCapability,
|
||||
client::ConfigWarning,
|
||||
event::Event,
|
||||
scheme::Scheme,
|
||||
types::{Request, RequestConfig},
|
||||
};
|
||||
|
||||
use super::OpenAIScheme;
|
||||
|
||||
impl Scheme for OpenAIScheme {
|
||||
type State = ();
|
||||
|
||||
fn default_base_url(&self) -> &'static str {
|
||||
"https://api.openai.com"
|
||||
}
|
||||
|
||||
fn path(&self, _model_id: &str) -> String {
|
||||
"/v1/chat/completions".to_string()
|
||||
}
|
||||
|
||||
fn required_auth(&self) -> AuthRequirement {
|
||||
AuthRequirement::Bearer
|
||||
}
|
||||
|
||||
fn build_request_body(
|
||||
&self,
|
||||
model_id: &str,
|
||||
request: &Request,
|
||||
capability: &ModelCapability,
|
||||
) -> Value {
|
||||
let req = self.build_request(model_id, request, capability);
|
||||
serde_json::to_value(&req).expect("OpenAIRequest is always serialisable")
|
||||
}
|
||||
|
||||
fn parse_sse(
|
||||
&self,
|
||||
_event_type: &str,
|
||||
data: &str,
|
||||
_state: &mut Self::State,
|
||||
) -> Result<Vec<Event>, ClientError> {
|
||||
// `data: [DONE]` は終端マーカー
|
||||
if data.trim() == "[DONE]" {
|
||||
return Ok(Vec::new());
|
||||
}
|
||||
Ok(self.parse_event(data)?.unwrap_or_default())
|
||||
}
|
||||
|
||||
fn default_capability(&self) -> ModelCapability {
|
||||
super::capability::default_capability()
|
||||
}
|
||||
|
||||
fn validate_config(&self, config: &RequestConfig) -> Vec<ConfigWarning> {
|
||||
let mut warnings = Vec::new();
|
||||
// OpenAI Chat Completions API は top_k を受け付けない
|
||||
if config.top_k.is_some() {
|
||||
warnings.push(ConfigWarning::unsupported("top_k", "OpenAI Chat"));
|
||||
}
|
||||
warnings
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,18 @@
|
||||
//! OpenAI Responses scheme の wire-level 既定 capability。
|
||||
//!
|
||||
//! モデル ID 固有のテーブル(`gpt-5` / `codex-` 系など)は高レベル構築層
|
||||
//! (`provider::capability`)の責務。ここでは wire の保守的 default のみ。
|
||||
|
||||
use crate::llm_client::capability::{
|
||||
CacheStrategy, ModelCapability, StructuredOutput, ToolCallingSupport,
|
||||
};
|
||||
|
||||
pub(crate) fn default_capability() -> ModelCapability {
|
||||
ModelCapability {
|
||||
tool_calling: ToolCallingSupport::Parallel,
|
||||
structured_output: StructuredOutput::JsonSchema,
|
||||
reasoning: None,
|
||||
vision: false,
|
||||
prompt_caching: CacheStrategy::Auto,
|
||||
}
|
||||
}
|
||||
File diff suppressed because it is too large
Load Diff
@@ -0,0 +1,84 @@
|
||||
//! OpenAI Responses API スキーマ (`/v1/responses`)
|
||||
//!
|
||||
//! Chat Completions とは別物の item-based wire format。reasoning item と
|
||||
//! function_call item が first-class で、SSE イベントも `response.*` 名前空間で
|
||||
//! 流れる。ChatGPT OAuth 経路 (codex) は本 scheme 必須。
|
||||
//!
|
||||
//! - リクエスト JSON 生成: [`request`]
|
||||
//! - SSE イベントパース → [`Event`](crate::llm_client::event::Event) 変換: [`events`]
|
||||
|
||||
mod capability;
|
||||
mod events;
|
||||
mod request;
|
||||
mod scheme_impl;
|
||||
|
||||
pub use scheme_impl::OpenAIResponsesState;
|
||||
|
||||
/// OpenAI Responses scheme 本体。
|
||||
///
|
||||
/// `store` / `include_encrypted_content` / `send_max_output_tokens` /
|
||||
/// `send_sampling_params` は scheme 固定の wire 設定で、デフォルトは
|
||||
/// 公式 OpenAI Responses API 向け (stateless + ZDR + `max_output_tokens`
|
||||
/// / `temperature` / `top_p` 送出可)。ChatGPT backend (codex-oauth) の
|
||||
/// ように受理パラメータが subset の経路では provider 層で
|
||||
/// `send_max_output_tokens=false` / `send_sampling_params=false` に
|
||||
/// 上書きする。`ModelCapability` には入れない（モデル能力ではなく wire policy）。
|
||||
#[derive(Debug, Clone)]
|
||||
pub struct OpenAIResponsesScheme {
|
||||
/// サーバ側に response を保存するか。ZDR/stateless 運用では `false`。
|
||||
pub store: bool,
|
||||
/// `include: ["reasoning.encrypted_content"]` を付けるか。
|
||||
/// `store=false` で reasoning を使うなら必須。
|
||||
pub include_encrypted_content: bool,
|
||||
/// `max_output_tokens` を body に載せるか。公式 OpenAI Responses API は
|
||||
/// 受理するが、ChatGPT backend (codex-oauth) は `Unsupported parameter`
|
||||
/// で 400 を返すため、その経路では `false` にする。
|
||||
pub send_max_output_tokens: bool,
|
||||
/// `temperature` / `top_p` を body に載せるか。公式 OpenAI Responses API
|
||||
/// は受理するが、ChatGPT backend (codex-oauth) は `Unsupported parameter`
|
||||
/// で 400 を返すため、その経路では `false` にする。
|
||||
pub send_sampling_params: bool,
|
||||
}
|
||||
|
||||
impl Default for OpenAIResponsesScheme {
|
||||
fn default() -> Self {
|
||||
Self {
|
||||
store: false,
|
||||
include_encrypted_content: true,
|
||||
send_max_output_tokens: true,
|
||||
send_sampling_params: true,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
impl OpenAIResponsesScheme {
|
||||
/// デフォルト設定 (`store=false`, `include=["reasoning.encrypted_content"]`,
|
||||
/// `send_max_output_tokens=true`, `send_sampling_params=true`)。
|
||||
pub fn new() -> Self {
|
||||
Self::default()
|
||||
}
|
||||
|
||||
/// `store` を上書き。
|
||||
pub fn with_store(mut self, store: bool) -> Self {
|
||||
self.store = store;
|
||||
self
|
||||
}
|
||||
|
||||
/// `include: ["reasoning.encrypted_content"]` の有無を上書き。
|
||||
pub fn with_include_encrypted_content(mut self, include: bool) -> Self {
|
||||
self.include_encrypted_content = include;
|
||||
self
|
||||
}
|
||||
|
||||
/// `max_output_tokens` を body に載せるかを上書き。
|
||||
pub fn with_send_max_output_tokens(mut self, send: bool) -> Self {
|
||||
self.send_max_output_tokens = send;
|
||||
self
|
||||
}
|
||||
|
||||
/// `temperature` / `top_p` を body に載せるかを上書き。
|
||||
pub fn with_send_sampling_params(mut self, send: bool) -> Self {
|
||||
self.send_sampling_params = send;
|
||||
self
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,697 @@
|
||||
//! OpenAI Responses API リクエスト body 生成
|
||||
//!
|
||||
//! Chat Completions の `messages` と違い、Responses は `input[]` の
|
||||
//! item 配列で reasoning / function_call / function_call_output が
|
||||
//! first-class。`Item` を素に近い形で `input[]` に投影できる。
|
||||
|
||||
use serde::{Serialize, Serializer};
|
||||
use serde_json::Value;
|
||||
|
||||
use crate::llm_client::{
|
||||
Request,
|
||||
capability::{ModelCapability, ReasoningControl, ReasoningSupport},
|
||||
types::{ContentPart, Item, Role, ToolDefinition, parse_tool_arguments},
|
||||
};
|
||||
|
||||
use super::OpenAIResponsesScheme;
|
||||
|
||||
/// `/v1/responses` のリクエスト body。
|
||||
#[derive(Debug, Serialize)]
|
||||
pub(crate) struct ResponsesRequest {
|
||||
pub model: String,
|
||||
/// システムプロンプト相当。`input[]` とは別フィールド。
|
||||
#[serde(skip_serializing_if = "Option::is_none")]
|
||||
pub instructions: Option<String>,
|
||||
pub input: Vec<InputItem>,
|
||||
#[serde(skip_serializing_if = "Vec::is_empty")]
|
||||
pub tools: Vec<ResponseTool>,
|
||||
/// 常時 `"auto"` を送る。scheme 固定値。
|
||||
pub tool_choice: &'static str,
|
||||
/// 常時 `true` を送る。scheme 固定値。
|
||||
pub parallel_tool_calls: bool,
|
||||
#[serde(skip_serializing_if = "Option::is_none")]
|
||||
pub reasoning: Option<ReasoningConfig>,
|
||||
/// ZDR / stateless 運用では `false`。
|
||||
pub store: bool,
|
||||
/// 常時 `true`。
|
||||
pub stream: bool,
|
||||
/// `["reasoning.encrypted_content"]` 等。
|
||||
#[serde(skip_serializing_if = "Vec::is_empty")]
|
||||
pub include: Vec<&'static str>,
|
||||
/// 公式 OpenAI Responses API では受理されるが、ChatGPT backend
|
||||
/// (codex-oauth) は 400 で弾く。scheme の `send_max_output_tokens`
|
||||
/// が `false` のときは `None` のまま送る (skip_serializing_if で除外)。
|
||||
#[serde(skip_serializing_if = "Option::is_none")]
|
||||
pub max_output_tokens: Option<u32>,
|
||||
/// 公式 OpenAI Responses API では受理されるが、ChatGPT backend
|
||||
/// (codex-oauth) は `temperature` / `top_p` を 400 で弾く。scheme の
|
||||
/// `send_sampling_params` が `false` のときは `None` のまま送る。
|
||||
#[serde(skip_serializing_if = "Option::is_none")]
|
||||
pub temperature: Option<f32>,
|
||||
#[serde(skip_serializing_if = "Option::is_none")]
|
||||
pub top_p: Option<f32>,
|
||||
/// 会話単位の安定キー。ChatGPT backend (codex-oauth) は明示キーが
|
||||
/// 無いとプロンプトキャッシュがほぼ効かない。pod 側は `SegmentId`
|
||||
/// を渡す。`Request::cache_key` が `None` のときはキー自体を送らない。
|
||||
#[serde(skip_serializing_if = "Option::is_none")]
|
||||
pub prompt_cache_key: Option<String>,
|
||||
}
|
||||
|
||||
/// reasoning 制御。
|
||||
#[derive(Debug, Serialize)]
|
||||
pub(crate) struct ReasoningConfig {
|
||||
#[serde(skip_serializing_if = "Option::is_none")]
|
||||
pub effort: Option<String>,
|
||||
/// summary の出力制御。`"auto"` 固定で summary_text を受け取る。
|
||||
pub summary: &'static str,
|
||||
}
|
||||
|
||||
/// `input[]` の 1 要素。
|
||||
///
|
||||
/// Responses API の item 型を素に近い形で投影する。未対応 type は
|
||||
/// 無視（reasoning 送信時に `content: []` の場合は `None` として弾く）。
|
||||
#[derive(Debug, Serialize)]
|
||||
#[serde(tag = "type", rename_all = "snake_case")]
|
||||
pub(crate) enum InputItem {
|
||||
/// 会話メッセージ。user / assistant / developer のいずれか。
|
||||
/// `Role::System` items は `developer` として投影する（ChatGPT
|
||||
/// backend が `role: "system"` を拒否するため。Codex CLI も
|
||||
/// system 相当の挿入には DeveloperInstructions = `role: "developer"`
|
||||
/// を使う）。
|
||||
Message {
|
||||
role: &'static str,
|
||||
content: Vec<InputContent>,
|
||||
},
|
||||
/// 過去の function tool 呼び出し（assistant 側）。
|
||||
FunctionCall {
|
||||
call_id: String,
|
||||
name: String,
|
||||
/// JSON 文字列（object でなくても正規化済み）。
|
||||
arguments: String,
|
||||
},
|
||||
/// function tool の結果（user 側）。
|
||||
FunctionCallOutput {
|
||||
call_id: String,
|
||||
/// Responses は文字列 or 構造化 output を許すが、ここでは
|
||||
/// `summary` + `content` を改行連結した文字列で送る。
|
||||
output: String,
|
||||
},
|
||||
/// reasoning item。`encrypted_content` があれば必ず添える。
|
||||
Reasoning {
|
||||
#[serde(skip_serializing_if = "Option::is_none")]
|
||||
id: Option<String>,
|
||||
/// Responses API は reasoning item に `summary` フィールドを必須で
|
||||
/// 要求する（中身が空でも `[]` として送る必要がある）。GPT-5 など
|
||||
/// summary を返さないモデル + reasoning effort 指定なしのターンでは
|
||||
/// summary text が一切付かないので、ここを skip すると 400
|
||||
/// "Missing required parameter: 'input[N].summary'" で弾かれる。
|
||||
summary: Vec<ReasoningSummaryPart>,
|
||||
#[serde(skip_serializing_if = "Vec::is_empty")]
|
||||
content: Vec<ReasoningContentPart>,
|
||||
#[serde(skip_serializing_if = "Option::is_none")]
|
||||
encrypted_content: Option<String>,
|
||||
},
|
||||
}
|
||||
|
||||
/// メッセージ content_part。role で input/output を使い分ける。
|
||||
#[derive(Debug, Serialize)]
|
||||
#[serde(tag = "type", rename_all = "snake_case")]
|
||||
pub(crate) enum InputContent {
|
||||
/// user / developer 側のテキスト
|
||||
InputText { text: String },
|
||||
/// assistant 側のテキスト
|
||||
OutputText { text: String },
|
||||
}
|
||||
|
||||
#[derive(Debug, Serialize)]
|
||||
#[serde(tag = "type", rename_all = "snake_case")]
|
||||
pub(crate) enum ReasoningSummaryPart {
|
||||
SummaryText { text: String },
|
||||
}
|
||||
|
||||
#[derive(Debug, Serialize)]
|
||||
#[serde(tag = "type", rename_all = "snake_case")]
|
||||
pub(crate) enum ReasoningContentPart {
|
||||
ReasoningText { text: String },
|
||||
}
|
||||
|
||||
/// Responses 用 tool 定義。Chat と違い function キーでネストせず
|
||||
/// トップレベルに `name` / `parameters` が載る。
|
||||
#[derive(Debug, Serialize)]
|
||||
pub(crate) struct ResponseTool {
|
||||
#[serde(rename = "type")]
|
||||
pub r#type: &'static str,
|
||||
pub name: String,
|
||||
#[serde(skip_serializing_if = "Option::is_none")]
|
||||
pub description: Option<String>,
|
||||
/// OpenAI Responses API は `type:"object"` のパラメータスキーマに
|
||||
/// `properties` が存在することを要求する。schemars は引数なし struct
|
||||
/// から `properties` を含まない最小スキーマを出すので、serialize
|
||||
/// 時に空オブジェクトを補う。
|
||||
#[serde(serialize_with = "serialize_parameters")]
|
||||
pub parameters: Value,
|
||||
/// Structured output モード制御。デフォルト false。
|
||||
pub strict: bool,
|
||||
}
|
||||
|
||||
fn serialize_parameters<S: Serializer>(value: &Value, s: S) -> Result<S::Ok, S::Error> {
|
||||
if let Some(obj) = value.as_object()
|
||||
&& obj.get("type").and_then(Value::as_str) == Some("object")
|
||||
&& !obj.contains_key("properties")
|
||||
{
|
||||
let mut patched = obj.clone();
|
||||
patched.insert("properties".to_string(), Value::Object(Default::default()));
|
||||
return Value::Object(patched).serialize(s);
|
||||
}
|
||||
value.serialize(s)
|
||||
}
|
||||
|
||||
impl OpenAIResponsesScheme {
|
||||
/// `Request` から wire 形式の body を組み立てる。
|
||||
pub(crate) fn build_request(
|
||||
&self,
|
||||
model: &str,
|
||||
request: &Request,
|
||||
capability: &ModelCapability,
|
||||
) -> ResponsesRequest {
|
||||
let input = convert_items_to_input(&request.items);
|
||||
let tools = request.tools.iter().map(convert_tool).collect();
|
||||
|
||||
// Reasoning 投影: capability が Effort / Both をサポートし、かつ
|
||||
// request 側で effort が指定されているときだけ reasoning を付ける。
|
||||
let supports_effort = matches!(
|
||||
capability.reasoning,
|
||||
Some(ReasoningSupport::Effort | ReasoningSupport::Both),
|
||||
);
|
||||
let reasoning = request
|
||||
.config
|
||||
.reasoning
|
||||
.as_ref()
|
||||
.filter(|_| supports_effort)
|
||||
.map(|effort| ReasoningConfig {
|
||||
effort: match effort {
|
||||
ReasoningControl::Effort(effort) => Some(effort.as_str().to_string()),
|
||||
ReasoningControl::BudgetTokens(_) => None,
|
||||
},
|
||||
summary: "auto",
|
||||
})
|
||||
.filter(|reasoning| reasoning.effort.is_some());
|
||||
|
||||
let include: Vec<&'static str> = if self.include_encrypted_content {
|
||||
vec!["reasoning.encrypted_content"]
|
||||
} else {
|
||||
Vec::new()
|
||||
};
|
||||
|
||||
ResponsesRequest {
|
||||
model: model.to_string(),
|
||||
instructions: request.system_prompt.clone(),
|
||||
input,
|
||||
tools,
|
||||
tool_choice: "auto",
|
||||
parallel_tool_calls: true,
|
||||
reasoning,
|
||||
store: self.store,
|
||||
stream: true,
|
||||
include,
|
||||
max_output_tokens: if self.send_max_output_tokens {
|
||||
request.config.max_tokens
|
||||
} else {
|
||||
None
|
||||
},
|
||||
temperature: if self.send_sampling_params {
|
||||
request.config.temperature
|
||||
} else {
|
||||
None
|
||||
},
|
||||
top_p: if self.send_sampling_params {
|
||||
request.config.top_p
|
||||
} else {
|
||||
None
|
||||
},
|
||||
prompt_cache_key: request.cache_key.clone(),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// `Item` 列を `input[]` に変換する。
|
||||
fn convert_items_to_input(items: &[Item]) -> Vec<InputItem> {
|
||||
let mut out = Vec::with_capacity(items.len());
|
||||
for item in items {
|
||||
match item {
|
||||
Item::Message { role, content, .. } => {
|
||||
let (role_str, text_variant): (&'static str, fn(String) -> InputContent) =
|
||||
match role {
|
||||
Role::User => ("user", |t| InputContent::InputText { text: t }),
|
||||
Role::Assistant => ("assistant", |t| InputContent::OutputText { text: t }),
|
||||
Role::System => ("developer", |t| InputContent::InputText { text: t }),
|
||||
};
|
||||
let parts: Vec<InputContent> = content
|
||||
.iter()
|
||||
.map(|p| match p {
|
||||
ContentPart::Text { text } => text_variant(text.clone()),
|
||||
ContentPart::Refusal { refusal } => text_variant(refusal.clone()),
|
||||
})
|
||||
.collect();
|
||||
out.push(InputItem::Message {
|
||||
role: role_str,
|
||||
content: parts,
|
||||
});
|
||||
}
|
||||
Item::ToolCall {
|
||||
call_id,
|
||||
name,
|
||||
arguments,
|
||||
..
|
||||
} => {
|
||||
// 非 object / 旧形式の "null" を "{}" に正規化。
|
||||
let normalized = parse_tool_arguments(arguments).to_string();
|
||||
out.push(InputItem::FunctionCall {
|
||||
call_id: call_id.clone(),
|
||||
name: name.clone(),
|
||||
arguments: normalized,
|
||||
});
|
||||
}
|
||||
Item::ToolResult {
|
||||
call_id,
|
||||
summary,
|
||||
content,
|
||||
..
|
||||
} => {
|
||||
let text = match content {
|
||||
Some(c) => format!("{summary}\n{c}"),
|
||||
None => summary.clone(),
|
||||
};
|
||||
out.push(InputItem::FunctionCallOutput {
|
||||
call_id: call_id.clone(),
|
||||
output: text,
|
||||
});
|
||||
}
|
||||
Item::Reasoning {
|
||||
id,
|
||||
text,
|
||||
summary,
|
||||
encrypted_content,
|
||||
..
|
||||
} => {
|
||||
let summary_parts = summary
|
||||
.iter()
|
||||
.filter(|s| !s.is_empty())
|
||||
.map(|s| ReasoningSummaryPart::SummaryText { text: s.clone() })
|
||||
.collect();
|
||||
let content_parts = if text.is_empty() {
|
||||
Vec::new()
|
||||
} else {
|
||||
vec![ReasoningContentPart::ReasoningText { text: text.clone() }]
|
||||
};
|
||||
out.push(InputItem::Reasoning {
|
||||
id: id.clone(),
|
||||
summary: summary_parts,
|
||||
content: content_parts,
|
||||
encrypted_content: encrypted_content.clone(),
|
||||
});
|
||||
}
|
||||
}
|
||||
}
|
||||
out
|
||||
}
|
||||
|
||||
fn convert_tool(tool: &ToolDefinition) -> ResponseTool {
|
||||
ResponseTool {
|
||||
r#type: "function",
|
||||
name: tool.name.clone(),
|
||||
description: tool.description.clone(),
|
||||
parameters: tool.input_schema.clone(),
|
||||
strict: false,
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
use crate::llm_client::capability::{
|
||||
CacheStrategy, ModelCapability, ReasoningControl, ReasoningEffort, ReasoningSupport,
|
||||
StructuredOutput, ToolCallingSupport,
|
||||
};
|
||||
|
||||
fn cap_with_reasoning() -> ModelCapability {
|
||||
ModelCapability {
|
||||
tool_calling: ToolCallingSupport::Parallel,
|
||||
structured_output: StructuredOutput::JsonSchema,
|
||||
reasoning: Some(ReasoningSupport::Effort),
|
||||
vision: true,
|
||||
prompt_caching: CacheStrategy::Auto,
|
||||
}
|
||||
}
|
||||
|
||||
fn cap_no_reasoning() -> ModelCapability {
|
||||
ModelCapability {
|
||||
reasoning: None,
|
||||
..cap_with_reasoning()
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn scheme_defaults_to_stateless_zdr() {
|
||||
let s = OpenAIResponsesScheme::new();
|
||||
assert!(!s.store);
|
||||
assert!(s.include_encrypted_content);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn includes_encrypted_content_when_enabled() {
|
||||
let scheme = OpenAIResponsesScheme::new();
|
||||
let req = Request::new().user("hi");
|
||||
let body = scheme.build_request("gpt-5", &req, &cap_with_reasoning());
|
||||
assert_eq!(body.include, vec!["reasoning.encrypted_content"]);
|
||||
assert!(!body.store);
|
||||
assert!(body.stream);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn instructions_from_system_prompt() {
|
||||
let scheme = OpenAIResponsesScheme::new();
|
||||
let req = Request::new().system("be terse").user("hi");
|
||||
let body = scheme.build_request("gpt-5", &req, &cap_with_reasoning());
|
||||
assert_eq!(body.instructions.as_deref(), Some("be terse"));
|
||||
assert_eq!(body.input.len(), 1);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn tool_choice_and_parallel_are_fixed() {
|
||||
let scheme = OpenAIResponsesScheme::new();
|
||||
let req = Request::new().user("hi");
|
||||
let body = scheme.build_request("gpt-5", &req, &cap_with_reasoning());
|
||||
assert_eq!(body.tool_choice, "auto");
|
||||
assert!(body.parallel_tool_calls);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn user_message_uses_input_text() {
|
||||
let scheme = OpenAIResponsesScheme::new();
|
||||
let req = Request::new().user("hi");
|
||||
let body = scheme.build_request("gpt-5", &req, &cap_with_reasoning());
|
||||
match &body.input[0] {
|
||||
InputItem::Message { role, content } => {
|
||||
assert_eq!(*role, "user");
|
||||
assert_eq!(content.len(), 1);
|
||||
assert!(matches!(&content[0], InputContent::InputText { text } if text == "hi"));
|
||||
}
|
||||
_ => panic!("expected message"),
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn system_role_item_is_projected_as_developer() {
|
||||
// ChatGPT backend (codex-oauth) は input[] の `role: "system"` を
|
||||
// "System messages are not allowed" で 400 拒否する。in-conversation
|
||||
// な system note (notify / fs_view auto-read / compaction summary) は
|
||||
// `role: "developer"` として投影し、両 backend で受理されるようにする。
|
||||
let scheme = OpenAIResponsesScheme::new();
|
||||
let req = Request::new()
|
||||
.user("hi")
|
||||
.item(Item::system_message("[notify] hello"));
|
||||
let body = scheme.build_request("gpt-5", &req, &cap_with_reasoning());
|
||||
match &body.input[1] {
|
||||
InputItem::Message { role, content } => {
|
||||
assert_eq!(*role, "developer");
|
||||
assert!(
|
||||
matches!(&content[0], InputContent::InputText { text } if text == "[notify] hello"),
|
||||
);
|
||||
}
|
||||
_ => panic!("expected message"),
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn assistant_message_uses_output_text() {
|
||||
let scheme = OpenAIResponsesScheme::new();
|
||||
let req = Request::new().user("hi").assistant("hello");
|
||||
let body = scheme.build_request("gpt-5", &req, &cap_with_reasoning());
|
||||
match &body.input[1] {
|
||||
InputItem::Message { role, content } => {
|
||||
assert_eq!(*role, "assistant");
|
||||
assert!(
|
||||
matches!(&content[0], InputContent::OutputText { text } if text == "hello")
|
||||
);
|
||||
}
|
||||
_ => panic!("expected message"),
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn tool_call_and_result_become_function_items() {
|
||||
let scheme = OpenAIResponsesScheme::new();
|
||||
let req = Request::new()
|
||||
.user("run")
|
||||
.item(Item::tool_call("c1", "t", r#"{"a":1}"#))
|
||||
.item(Item::tool_result("c1", "ok"));
|
||||
let body = scheme.build_request("gpt-5", &req, &cap_with_reasoning());
|
||||
assert!(matches!(body.input[1], InputItem::FunctionCall { .. }));
|
||||
assert!(matches!(
|
||||
body.input[2],
|
||||
InputItem::FunctionCallOutput { .. }
|
||||
));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn reasoning_item_round_trips_encrypted_content() {
|
||||
let scheme = OpenAIResponsesScheme::new();
|
||||
let item = Item::reasoning("inner")
|
||||
.with_reasoning_summary(vec!["s1".into()])
|
||||
.with_encrypted_content("ENC");
|
||||
let req = Request::new().user("hi").item(item);
|
||||
let body = scheme.build_request("gpt-5", &req, &cap_with_reasoning());
|
||||
match &body.input[1] {
|
||||
InputItem::Reasoning {
|
||||
summary,
|
||||
content,
|
||||
encrypted_content,
|
||||
..
|
||||
} => {
|
||||
assert_eq!(summary.len(), 1);
|
||||
assert_eq!(content.len(), 1);
|
||||
assert_eq!(encrypted_content.as_deref(), Some("ENC"));
|
||||
}
|
||||
_ => panic!("expected reasoning"),
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn persisted_reasoning_items_are_preserved_across_user_turns() {
|
||||
let scheme = OpenAIResponsesScheme::new();
|
||||
let old_reasoning = Item::reasoning("old").with_encrypted_content("OLD_ENC");
|
||||
let current_reasoning = Item::reasoning("current").with_encrypted_content("CURRENT_ENC");
|
||||
let req = Request::new()
|
||||
.user("old prompt")
|
||||
.item(old_reasoning)
|
||||
.assistant("old answer")
|
||||
.user("new prompt")
|
||||
.item(current_reasoning);
|
||||
let body = scheme.build_request("gpt-5", &req, &cap_with_reasoning());
|
||||
let encrypted: Vec<_> = body
|
||||
.input
|
||||
.iter()
|
||||
.filter_map(|item| match item {
|
||||
InputItem::Reasoning {
|
||||
encrypted_content, ..
|
||||
} => encrypted_content.as_deref(),
|
||||
_ => None,
|
||||
})
|
||||
.collect();
|
||||
assert_eq!(encrypted, vec!["OLD_ENC", "CURRENT_ENC"]);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn reasoning_is_kept_across_function_call_loop() {
|
||||
let scheme = OpenAIResponsesScheme::new();
|
||||
let req = Request::new()
|
||||
.user("run tool")
|
||||
.item(Item::reasoning("plan").with_encrypted_content("ENC"))
|
||||
.item(Item::tool_call("c1", "tool", "{}"))
|
||||
.item(Item::tool_result("c1", "ok"));
|
||||
let body = scheme.build_request("gpt-5", &req, &cap_with_reasoning());
|
||||
assert!(matches!(body.input[1], InputItem::Reasoning { .. }));
|
||||
assert!(matches!(body.input[2], InputItem::FunctionCall { .. }));
|
||||
assert!(matches!(
|
||||
body.input[3],
|
||||
InputItem::FunctionCallOutput { .. }
|
||||
));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn reasoning_summary_field_is_always_serialized() {
|
||||
// Responses API は reasoning item に `summary` を必須で要求する。
|
||||
// summary が空でも wire 上に `summary: []` として残らないと、
|
||||
// ChatGPT backend (codex-oauth) が
|
||||
// 400 invalid_request_error: Missing required parameter:
|
||||
// 'input[N].summary'.
|
||||
// で弾く。GPT-5 + reasoning effort 未指定のターンでは summary text
|
||||
// が付かないことがあるため、空のままでも skip しないこと。
|
||||
let scheme = OpenAIResponsesScheme::new();
|
||||
let item = Item::reasoning("").with_encrypted_content("ENC");
|
||||
let req = Request::new().user("hi").item(item);
|
||||
let body = scheme.build_request("gpt-5", &req, &cap_with_reasoning());
|
||||
let json = serde_json::to_value(&body).unwrap();
|
||||
let reasoning_item = &json["input"][1];
|
||||
assert_eq!(reasoning_item["type"], "reasoning");
|
||||
assert!(
|
||||
reasoning_item.get("summary").is_some(),
|
||||
"summary key must be present even when empty, got: {reasoning_item}"
|
||||
);
|
||||
assert_eq!(reasoning_item["summary"], serde_json::json!([]));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn reasoning_effort_projected_when_supported() {
|
||||
let scheme = OpenAIResponsesScheme::new();
|
||||
let mut req = Request::new().user("hi");
|
||||
req.config.reasoning = Some(ReasoningControl::Effort(ReasoningEffort::High));
|
||||
let body = scheme.build_request("gpt-5", &req, &cap_with_reasoning());
|
||||
let reasoning = body.reasoning.expect("reasoning should be set");
|
||||
assert_eq!(reasoning.effort.as_deref(), Some("high"));
|
||||
assert_eq!(reasoning.summary, "auto");
|
||||
let json = serde_json::to_value(reasoning).unwrap();
|
||||
assert!(
|
||||
json.get("context").is_none(),
|
||||
"reasoning.context must not be serialized, got: {json}"
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn reasoning_omitted_when_unsupported() {
|
||||
let scheme = OpenAIResponsesScheme::new();
|
||||
let mut req = Request::new().user("hi");
|
||||
req.config.reasoning = Some(ReasoningControl::Effort(ReasoningEffort::High));
|
||||
let body = scheme.build_request("gpt-4o", &req, &cap_no_reasoning());
|
||||
assert!(body.reasoning.is_none());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn max_output_tokens_passed_through_by_default() {
|
||||
let scheme = OpenAIResponsesScheme::new();
|
||||
let req = Request::new().user("hi").max_tokens(100);
|
||||
let body = scheme.build_request("gpt-5", &req, &cap_with_reasoning());
|
||||
assert_eq!(body.max_output_tokens, Some(100));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn max_output_tokens_dropped_when_send_disabled() {
|
||||
let scheme = OpenAIResponsesScheme::new().with_send_max_output_tokens(false);
|
||||
let req = Request::new().user("hi").max_tokens(100);
|
||||
let body = scheme.build_request("gpt-5", &req, &cap_with_reasoning());
|
||||
assert_eq!(body.max_output_tokens, None);
|
||||
let json = serde_json::to_value(&body).unwrap();
|
||||
assert!(
|
||||
json.get("max_output_tokens").is_none(),
|
||||
"max_output_tokens key must not appear in serialised body, got: {json}"
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn sampling_params_passed_through_by_default() {
|
||||
let scheme = OpenAIResponsesScheme::new();
|
||||
let req = Request::new().user("hi").temperature(0.4).top_p(0.9);
|
||||
let body = scheme.build_request("gpt-5", &req, &cap_with_reasoning());
|
||||
assert_eq!(body.temperature, Some(0.4));
|
||||
assert_eq!(body.top_p, Some(0.9));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn sampling_params_dropped_when_send_disabled() {
|
||||
let scheme = OpenAIResponsesScheme::new().with_send_sampling_params(false);
|
||||
let req = Request::new().user("hi").temperature(0.4).top_p(0.9);
|
||||
let body = scheme.build_request("gpt-5", &req, &cap_with_reasoning());
|
||||
assert_eq!(body.temperature, None);
|
||||
assert_eq!(body.top_p, None);
|
||||
let json = serde_json::to_value(&body).unwrap();
|
||||
assert!(
|
||||
json.get("temperature").is_none() && json.get("top_p").is_none(),
|
||||
"temperature/top_p keys must not appear in serialised body, got: {json}"
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn prompt_cache_key_passed_through_when_set() {
|
||||
let scheme = OpenAIResponsesScheme::new();
|
||||
let req = Request::new().user("hi").cache_key("session-abc");
|
||||
let body = scheme.build_request("gpt-5", &req, &cap_with_reasoning());
|
||||
assert_eq!(body.prompt_cache_key.as_deref(), Some("session-abc"));
|
||||
let json = serde_json::to_value(&body).unwrap();
|
||||
assert_eq!(json["prompt_cache_key"], "session-abc");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn prompt_cache_key_omitted_when_none() {
|
||||
let scheme = OpenAIResponsesScheme::new();
|
||||
let req = Request::new().user("hi");
|
||||
let body = scheme.build_request("gpt-5", &req, &cap_with_reasoning());
|
||||
assert!(body.prompt_cache_key.is_none());
|
||||
let json = serde_json::to_value(&body).unwrap();
|
||||
assert!(
|
||||
json.get("prompt_cache_key").is_none(),
|
||||
"prompt_cache_key key must not appear in serialised body, got: {json}"
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn tool_schema_without_properties_is_normalized() {
|
||||
// schemars は引数なし struct から `type:"object"` だけのスキーマを
|
||||
// 吐く。OpenAI Responses は `properties` 欠落を 400 で拒否するので
|
||||
// 送る直前に空オブジェクトを補うのを確認。
|
||||
let scheme = OpenAIResponsesScheme::new();
|
||||
let raw_schema = serde_json::json!({ "type": "object" });
|
||||
let req = Request::new().tool(
|
||||
ToolDefinition::new("empty")
|
||||
.description("no args")
|
||||
.input_schema(raw_schema),
|
||||
);
|
||||
let body = scheme.build_request("gpt-5", &req, &cap_with_reasoning());
|
||||
let json = serde_json::to_value(&body).unwrap();
|
||||
assert_eq!(json["tools"][0]["parameters"]["type"], "object");
|
||||
assert!(
|
||||
json["tools"][0]["parameters"]["properties"].is_object(),
|
||||
"properties must be present as an object, got: {}",
|
||||
json["tools"][0]["parameters"]
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn tool_schema_with_properties_is_untouched() {
|
||||
let scheme = OpenAIResponsesScheme::new();
|
||||
let raw_schema = serde_json::json!({
|
||||
"type": "object",
|
||||
"properties": { "path": { "type": "string" } },
|
||||
"required": ["path"]
|
||||
});
|
||||
let req = Request::new().tool(
|
||||
ToolDefinition::new("t")
|
||||
.description("d")
|
||||
.input_schema(raw_schema.clone()),
|
||||
);
|
||||
let body = scheme.build_request("gpt-5", &req, &cap_with_reasoning());
|
||||
let json = serde_json::to_value(&body).unwrap();
|
||||
assert_eq!(json["tools"][0]["parameters"], raw_schema);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn serialized_body_has_expected_shape() {
|
||||
// wire 形式が崩れていないかのスモークテスト
|
||||
let scheme = OpenAIResponsesScheme::new();
|
||||
let req = Request::new()
|
||||
.system("sys")
|
||||
.user("hi")
|
||||
.tool(ToolDefinition::new("t").description("d"));
|
||||
let body = scheme.build_request("gpt-5", &req, &cap_with_reasoning());
|
||||
let json = serde_json::to_value(&body).unwrap();
|
||||
assert_eq!(json["model"], "gpt-5");
|
||||
assert_eq!(json["instructions"], "sys");
|
||||
assert_eq!(json["tool_choice"], "auto");
|
||||
assert_eq!(json["parallel_tool_calls"], true);
|
||||
assert_eq!(json["store"], false);
|
||||
assert_eq!(json["stream"], true);
|
||||
assert_eq!(json["include"][0], "reasoning.encrypted_content");
|
||||
assert_eq!(json["tools"][0]["type"], "function");
|
||||
assert_eq!(json["tools"][0]["name"], "t");
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,88 @@
|
||||
//! `impl Scheme for OpenAIResponsesScheme`
|
||||
|
||||
use serde_json::Value;
|
||||
|
||||
use crate::llm_client::{
|
||||
ClientError,
|
||||
auth::AuthRequirement,
|
||||
capability::ModelCapability,
|
||||
client::ConfigWarning,
|
||||
event::Event,
|
||||
scheme::Scheme,
|
||||
types::{Request, RequestConfig},
|
||||
};
|
||||
|
||||
use super::OpenAIResponsesScheme;
|
||||
|
||||
pub use super::events::OpenAIResponsesState;
|
||||
|
||||
impl Scheme for OpenAIResponsesScheme {
|
||||
type State = OpenAIResponsesState;
|
||||
|
||||
fn default_base_url(&self) -> &'static str {
|
||||
// `/v1` は base_url 側に寄せる。ChatGPT OAuth 経由のときは
|
||||
// `https://chatgpt.com/backend-api/codex` を base にすれば同じ
|
||||
// `/responses` path で両系統を吸収できる（Codex CLI 準拠）。
|
||||
"https://api.openai.com/v1"
|
||||
}
|
||||
|
||||
fn path(&self, _model_id: &str) -> String {
|
||||
"/responses".to_string()
|
||||
}
|
||||
|
||||
fn required_auth(&self) -> AuthRequirement {
|
||||
AuthRequirement::Bearer
|
||||
}
|
||||
|
||||
fn build_request_body(
|
||||
&self,
|
||||
model_id: &str,
|
||||
request: &Request,
|
||||
capability: &ModelCapability,
|
||||
) -> Value {
|
||||
let body = self.build_request(model_id, request, capability);
|
||||
serde_json::to_value(&body).expect("ResponsesRequest is always serialisable")
|
||||
}
|
||||
|
||||
fn parse_sse(
|
||||
&self,
|
||||
event_type: &str,
|
||||
data: &str,
|
||||
state: &mut Self::State,
|
||||
) -> Result<Vec<Event>, ClientError> {
|
||||
super::events::parse_sse(event_type, data, state)
|
||||
}
|
||||
|
||||
fn default_capability(&self) -> ModelCapability {
|
||||
super::capability::default_capability()
|
||||
}
|
||||
|
||||
fn validate_config(&self, config: &RequestConfig) -> Vec<ConfigWarning> {
|
||||
let mut warnings = Vec::new();
|
||||
// ChatGPT backend (codex-oauth) は `max_output_tokens` を 400 で弾く。
|
||||
// scheme 構築時に `send_max_output_tokens=false` で組まれていれば
|
||||
// body 投影は止まっているので、ユーザの意図が落ちることだけを通知する。
|
||||
if !self.send_max_output_tokens && config.max_tokens.is_some() {
|
||||
warnings.push(ConfigWarning::unsupported(
|
||||
"max_tokens",
|
||||
"OpenAI Responses (ChatGPT backend)",
|
||||
));
|
||||
}
|
||||
// 同上、`temperature` / `top_p` も ChatGPT backend では 400 で弾かれる。
|
||||
if !self.send_sampling_params {
|
||||
if config.temperature.is_some() {
|
||||
warnings.push(ConfigWarning::unsupported(
|
||||
"temperature",
|
||||
"OpenAI Responses (ChatGPT backend)",
|
||||
));
|
||||
}
|
||||
if config.top_p.is_some() {
|
||||
warnings.push(ConfigWarning::unsupported(
|
||||
"top_p",
|
||||
"OpenAI Responses (ChatGPT backend)",
|
||||
));
|
||||
}
|
||||
}
|
||||
warnings
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,908 @@
|
||||
//! `HttpTransport<S: Scheme>`: すべての LLM wire scheme を共通の 1 本の
|
||||
//! HTTP クライアントで扱う。
|
||||
//!
|
||||
//! 旧 `providers/{anthropic,openai,gemini,ollama}.rs` を置き換える。
|
||||
//! scheme 固有の差分は [`Scheme`] trait 実装に委譲する。
|
||||
|
||||
use std::pin::Pin;
|
||||
use std::sync::Arc;
|
||||
use std::time::{Duration, Instant};
|
||||
|
||||
use async_trait::async_trait;
|
||||
use eventsource_stream::Eventsource;
|
||||
use futures::{Stream, StreamExt, TryStreamExt};
|
||||
use reqwest::header::{
|
||||
ACCEPT, CONTENT_ENCODING, CONTENT_LENGTH, CONTENT_TYPE, HeaderMap, HeaderName, HeaderValue,
|
||||
RETRY_AFTER, TRANSFER_ENCODING,
|
||||
};
|
||||
use serde_json::{Map, Value, json};
|
||||
|
||||
use super::auth::{AuthProvider, AuthRequirement};
|
||||
use super::capability::ModelCapability;
|
||||
use super::client::{ConfigWarning, LlmClient, ResponseStream};
|
||||
use super::error::ClientError;
|
||||
use super::event::Event;
|
||||
use super::scheme::Scheme;
|
||||
use super::types::{Request, RequestConfig};
|
||||
|
||||
pub const DEFAULT_STREAM_OPEN_TIMEOUT: Duration = Duration::from_secs(20);
|
||||
pub const DEFAULT_FIRST_STREAM_EVENT_TIMEOUT: Duration = Duration::from_secs(30);
|
||||
|
||||
/// `AuthRef` を解決したランタイム表現。`crates/provider` が構築する。
|
||||
///
|
||||
/// - `None`: 認証ヘッダを送らない（Ollama 等の opt-out）
|
||||
/// - `ApiKey`: 静的な API key 文字列
|
||||
/// - `Custom`: リクエスト毎に動的にヘッダを組み立てる（Codex OAuth 等）
|
||||
#[derive(Debug, Clone)]
|
||||
pub enum ResolvedAuth {
|
||||
None,
|
||||
ApiKey(String),
|
||||
Custom(Arc<dyn AuthProvider>),
|
||||
}
|
||||
|
||||
impl ResolvedAuth {
|
||||
/// 認証要件と実際の解決値が噛み合うか検査する。構築時検証用。
|
||||
///
|
||||
/// - `ResolvedAuth::None` は認証を付けない宣言なので、どの
|
||||
/// `AuthRequirement` でも受け入れる（Ollama の Anthropic scheme
|
||||
/// 流用は `required_auth = XApiKey` だが認証ヘッダなしで動く）
|
||||
/// - `ResolvedAuth::Custom` は「ヘッダ組立を全部こちらで行う」
|
||||
/// 宣言なので、scheme が要求する形式によらず受け入れる
|
||||
pub fn matches(&self, req: AuthRequirement) -> bool {
|
||||
match (self, req) {
|
||||
(Self::None, _) => true,
|
||||
(Self::Custom(_), _) => true,
|
||||
(
|
||||
Self::ApiKey(_),
|
||||
AuthRequirement::Bearer
|
||||
| AuthRequirement::XApiKey
|
||||
| AuthRequirement::QueryParam { .. },
|
||||
) => true,
|
||||
_ => false,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
fn header_value_for_diagnostics(headers: &HeaderMap, name: &str) -> Option<String> {
|
||||
headers
|
||||
.get(name)
|
||||
.and_then(|value| value.to_str().ok())
|
||||
.map(str::trim)
|
||||
.filter(|value| !value.is_empty())
|
||||
.map(ToOwned::to_owned)
|
||||
}
|
||||
|
||||
fn response_header_diagnostics(headers: &HeaderMap) -> serde_json::Value {
|
||||
serde_json::json!({
|
||||
"content_type": header_value_for_diagnostics(headers, CONTENT_TYPE.as_str()),
|
||||
"content_encoding": header_value_for_diagnostics(headers, CONTENT_ENCODING.as_str()),
|
||||
"transfer_encoding": header_value_for_diagnostics(headers, TRANSFER_ENCODING.as_str()),
|
||||
"content_length": header_value_for_diagnostics(headers, CONTENT_LENGTH.as_str()),
|
||||
})
|
||||
}
|
||||
|
||||
fn request_header_diagnostics(headers: &HeaderMap) -> serde_json::Value {
|
||||
serde_json::json!({
|
||||
"content_type": header_value_for_diagnostics(headers, CONTENT_TYPE.as_str()),
|
||||
"content_encoding": header_value_for_diagnostics(headers, CONTENT_ENCODING.as_str()),
|
||||
"accept": header_value_for_diagnostics(headers, ACCEPT.as_str()),
|
||||
"openai_beta": header_value_for_diagnostics(headers, "openai-beta"),
|
||||
"session_id_present": headers.contains_key("session-id"),
|
||||
"thread_id_present": headers.contains_key("thread-id"),
|
||||
"legacy_session_id_present": headers.contains_key("session_id"),
|
||||
"legacy_thread_id_present": headers.contains_key("thread_id"),
|
||||
"x_client_request_id_present": headers.contains_key("x-client-request-id"),
|
||||
"chatgpt_account_id_present": headers.contains_key("chatgpt-account-id"),
|
||||
})
|
||||
}
|
||||
|
||||
fn sse_error_context(status: u16, headers: &serde_json::Value, source: &str) -> String {
|
||||
let field = |name: &str| {
|
||||
headers
|
||||
.get(name)
|
||||
.and_then(serde_json::Value::as_str)
|
||||
.unwrap_or("<none>")
|
||||
};
|
||||
format!(
|
||||
"SSE stream parse failed after HTTP {status}: {source}; content-type={}, content-encoding={}, transfer-encoding={}, content-length={}",
|
||||
field("content_type"),
|
||||
field("content_encoding"),
|
||||
field("transfer_encoding"),
|
||||
field("content_length")
|
||||
)
|
||||
}
|
||||
|
||||
/// scheme 共通の HTTP 通信層。
|
||||
pub struct HttpTransport<S: Scheme> {
|
||||
http_client: reqwest::Client,
|
||||
scheme: S,
|
||||
model_id: String,
|
||||
base_url: String,
|
||||
auth: ResolvedAuth,
|
||||
capability: ModelCapability,
|
||||
}
|
||||
|
||||
impl<S: Scheme> HttpTransport<S> {
|
||||
/// 新しい transport を作る。`base_url` は末尾スラッシュの有無を
|
||||
/// どちらでも受け付ける（内部で正規化）。
|
||||
pub fn new(
|
||||
scheme: S,
|
||||
model_id: impl Into<String>,
|
||||
base_url: impl Into<String>,
|
||||
auth: ResolvedAuth,
|
||||
capability: ModelCapability,
|
||||
) -> Self {
|
||||
let base_url = base_url.into();
|
||||
let base_url = base_url.trim_end_matches('/').to_string();
|
||||
Self {
|
||||
http_client: reqwest::Client::new(),
|
||||
scheme,
|
||||
model_id: model_id.into(),
|
||||
base_url,
|
||||
auth,
|
||||
capability,
|
||||
}
|
||||
}
|
||||
|
||||
/// カスタム HTTP クライアントを差し込む（テスト等）。
|
||||
pub fn with_http_client(mut self, client: reqwest::Client) -> Self {
|
||||
self.http_client = client;
|
||||
self
|
||||
}
|
||||
|
||||
fn build_url(&self) -> String {
|
||||
let path = self.scheme.path(&self.model_id);
|
||||
let url = format!("{}{}", self.base_url, path);
|
||||
// Gemini のようにクエリパラメータで認証する場合は URL にキーを追記する
|
||||
if let (AuthRequirement::QueryParam { name }, ResolvedAuth::ApiKey(key)) =
|
||||
(self.scheme.required_auth(), &self.auth)
|
||||
{
|
||||
let sep = if url.contains('?') { '&' } else { '?' };
|
||||
format!("{url}{sep}{name}={key}")
|
||||
} else {
|
||||
url
|
||||
}
|
||||
}
|
||||
|
||||
async fn build_headers(&self) -> Result<HeaderMap, ClientError> {
|
||||
let mut headers = HeaderMap::new();
|
||||
headers.insert(CONTENT_TYPE, HeaderValue::from_static("application/json"));
|
||||
|
||||
match (&self.auth, self.scheme.required_auth()) {
|
||||
(ResolvedAuth::None, _) | (_, AuthRequirement::None) => {}
|
||||
(ResolvedAuth::Custom(provider), _) => {
|
||||
for (name, mut value) in provider.headers().await? {
|
||||
value.set_sensitive(true);
|
||||
headers.insert(name, value);
|
||||
}
|
||||
}
|
||||
(ResolvedAuth::ApiKey(key), AuthRequirement::Bearer) => {
|
||||
let mut val = HeaderValue::from_str(&format!("Bearer {key}"))
|
||||
.map_err(|e| ClientError::Config(format!("invalid api key: {e}")))?;
|
||||
val.set_sensitive(true);
|
||||
headers.insert("Authorization", val);
|
||||
}
|
||||
(ResolvedAuth::ApiKey(key), AuthRequirement::XApiKey) => {
|
||||
let mut val = HeaderValue::from_str(key.as_str())
|
||||
.map_err(|e| ClientError::Config(format!("invalid api key: {e}")))?;
|
||||
val.set_sensitive(true);
|
||||
headers.insert("x-api-key", val);
|
||||
}
|
||||
(_, AuthRequirement::QueryParam { .. }) => {
|
||||
// クエリパラメータは `build_url` で付与済み
|
||||
}
|
||||
(ResolvedAuth::ApiKey(_), AuthRequirement::Custom) => {
|
||||
// scheme が Custom を要求する組合せに ApiKey は流れてこない想定
|
||||
// （`matches()` で弾かれる）。安全側で何もしない
|
||||
}
|
||||
}
|
||||
|
||||
for (name, value) in self.scheme.additional_headers() {
|
||||
let hv = HeaderValue::from_str(&value)
|
||||
.map_err(|e| ClientError::Config(format!("invalid header {name}: {e}")))?;
|
||||
headers.insert(name, hv);
|
||||
}
|
||||
|
||||
Ok(headers)
|
||||
}
|
||||
|
||||
fn is_codex_backend(&self) -> bool {
|
||||
match &self.auth {
|
||||
ResolvedAuth::Custom(provider) => provider.is_codex_backend(),
|
||||
_ => false,
|
||||
}
|
||||
}
|
||||
|
||||
fn apply_stream_headers(
|
||||
&self,
|
||||
headers: &mut HeaderMap,
|
||||
request: &Request,
|
||||
) -> Result<(), ClientError> {
|
||||
headers.insert(ACCEPT, HeaderValue::from_static("text/event-stream"));
|
||||
|
||||
if self.is_codex_backend()
|
||||
&& let Some(cache_key) = request.cache_key.as_deref()
|
||||
{
|
||||
let value = HeaderValue::from_str(cache_key).map_err(|e| {
|
||||
ClientError::Config(format!("invalid Codex conversation header: {e}"))
|
||||
})?;
|
||||
// Codex CLI sends hyphenated session/thread headers to the
|
||||
// ChatGPT Codex backend. Keep the legacy underscore header for
|
||||
// existing traces/backends while exposing the current Codex shape.
|
||||
headers.insert(HeaderName::from_static("session-id"), value.clone());
|
||||
headers.insert(HeaderName::from_static("thread-id"), value.clone());
|
||||
headers.insert(HeaderName::from_static("session_id"), value.clone());
|
||||
headers.insert(HeaderName::from_static("x-client-request-id"), value);
|
||||
}
|
||||
|
||||
Ok(())
|
||||
}
|
||||
|
||||
fn encode_request_body(
|
||||
&self,
|
||||
body: &serde_json::Value,
|
||||
headers: &mut HeaderMap,
|
||||
) -> Result<RequestBody, ClientError> {
|
||||
if !self.is_codex_backend() {
|
||||
return Ok(RequestBody::Json(body.clone()));
|
||||
}
|
||||
|
||||
let raw = serde_json::to_vec(body)?;
|
||||
let raw_json_bytes = raw.len();
|
||||
let compressed = zstd::stream::encode_all(std::io::Cursor::new(raw), 3)
|
||||
.map_err(|e| ClientError::Config(format!("failed to zstd-compress request: {e}")))?;
|
||||
headers.insert(CONTENT_ENCODING, HeaderValue::from_static("zstd"));
|
||||
Ok(RequestBody::CompressedJson {
|
||||
bytes: compressed,
|
||||
raw_json_bytes,
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
enum RequestBody {
|
||||
Json(serde_json::Value),
|
||||
CompressedJson {
|
||||
bytes: Vec<u8>,
|
||||
raw_json_bytes: usize,
|
||||
},
|
||||
}
|
||||
|
||||
impl RequestBody {
|
||||
fn encoding(&self) -> &'static str {
|
||||
match self {
|
||||
Self::Json(_) => "json",
|
||||
Self::CompressedJson { .. } => "zstd",
|
||||
}
|
||||
}
|
||||
|
||||
fn raw_json_bytes(&self) -> Option<usize> {
|
||||
match self {
|
||||
Self::Json(body) => serde_json::to_vec(body).ok().map(|bytes| bytes.len()),
|
||||
Self::CompressedJson { raw_json_bytes, .. } => Some(*raw_json_bytes),
|
||||
}
|
||||
}
|
||||
|
||||
fn wire_bytes(&self) -> Option<usize> {
|
||||
match self {
|
||||
Self::Json(body) => serde_json::to_vec(body).ok().map(|bytes| bytes.len()),
|
||||
Self::CompressedJson { bytes, .. } => Some(bytes.len()),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
fn auth_kind(auth: &ResolvedAuth) -> &'static str {
|
||||
match auth {
|
||||
ResolvedAuth::None => "none",
|
||||
ResolvedAuth::ApiKey(_) => "api_key",
|
||||
ResolvedAuth::Custom(_) => "custom",
|
||||
}
|
||||
}
|
||||
|
||||
fn emit_transport_trace(request: &Request, label: &str, data: Value) {
|
||||
if let Some(trace) = &request.transport_trace {
|
||||
trace.emit(label, data);
|
||||
}
|
||||
}
|
||||
|
||||
fn json_value_kind(value: &Value) -> &'static str {
|
||||
match value {
|
||||
Value::Null => "null",
|
||||
Value::Bool(_) => "bool",
|
||||
Value::Number(_) => "number",
|
||||
Value::String(_) => "string",
|
||||
Value::Array(_) => "array",
|
||||
Value::Object(_) => "object",
|
||||
}
|
||||
}
|
||||
|
||||
fn request_body_shape_payload(body: &Value) -> Value {
|
||||
let mut map = Map::new();
|
||||
if let Some(input) = body.get("input").and_then(Value::as_array) {
|
||||
let items_json_bytes = serde_json::to_vec(input).map(|bytes| bytes.len()).ok();
|
||||
let mut reasoning_items = 0usize;
|
||||
let mut reasoning_encrypted_content_count = 0usize;
|
||||
let mut reasoning_encrypted_content_bytes = 0usize;
|
||||
for item in input {
|
||||
if item.get("type").and_then(Value::as_str) != Some("reasoning") {
|
||||
continue;
|
||||
}
|
||||
reasoning_items += 1;
|
||||
if let Some(encrypted) = item.get("encrypted_content").and_then(Value::as_str) {
|
||||
reasoning_encrypted_content_count += 1;
|
||||
reasoning_encrypted_content_bytes += encrypted.len();
|
||||
}
|
||||
}
|
||||
map.insert("items_len".to_string(), json!(input.len()));
|
||||
map.insert("items_json_bytes".to_string(), json!(items_json_bytes));
|
||||
map.insert("reasoning_items".to_string(), json!(reasoning_items));
|
||||
map.insert(
|
||||
"reasoning_encrypted_content_count".to_string(),
|
||||
json!(reasoning_encrypted_content_count),
|
||||
);
|
||||
map.insert(
|
||||
"reasoning_encrypted_content_bytes".to_string(),
|
||||
json!(reasoning_encrypted_content_bytes),
|
||||
);
|
||||
}
|
||||
Value::Object(map)
|
||||
}
|
||||
|
||||
fn api_error_code(error: &ClientError) -> Option<&str> {
|
||||
match error {
|
||||
ClientError::Api { code, .. } => code.as_deref(),
|
||||
_ => None,
|
||||
}
|
||||
}
|
||||
|
||||
fn is_context_length_exceeded(error: &ClientError) -> bool {
|
||||
match error {
|
||||
ClientError::Api { code, message, .. } => {
|
||||
code.as_deref() == Some("context_length_exceeded")
|
||||
|| message.contains("context_length_exceeded")
|
||||
}
|
||||
_ => false,
|
||||
}
|
||||
}
|
||||
|
||||
async fn response_with_timeout(
|
||||
future: impl std::future::Future<Output = Result<reqwest::Response, reqwest::Error>>,
|
||||
timeout: Duration,
|
||||
phase: &'static str,
|
||||
) -> Result<reqwest::Response, ClientError> {
|
||||
tokio::time::timeout(timeout, future)
|
||||
.await
|
||||
.map_err(|_| ClientError::Timeout { phase, timeout })?
|
||||
.map_err(ClientError::Http)
|
||||
}
|
||||
|
||||
impl<S: Scheme + Clone> Clone for HttpTransport<S> {
|
||||
fn clone(&self) -> Self {
|
||||
Self {
|
||||
http_client: self.http_client.clone(),
|
||||
scheme: self.scheme.clone(),
|
||||
model_id: self.model_id.clone(),
|
||||
base_url: self.base_url.clone(),
|
||||
auth: self.auth.clone(),
|
||||
capability: self.capability.clone(),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// エラーレスポンスを `ClientError::Api` に変換する。
|
||||
async fn classify_error_response(resp: reqwest::Response) -> ClientError {
|
||||
let status = resp.status().as_u16();
|
||||
let retry_after = resp
|
||||
.headers()
|
||||
.get(RETRY_AFTER)
|
||||
.and_then(|v| v.to_str().ok())
|
||||
.and_then(|s| s.trim().parse::<u64>().ok())
|
||||
.map(Duration::from_secs);
|
||||
let text = resp.text().await.unwrap_or_default();
|
||||
if let Ok(json) = serde_json::from_str::<serde_json::Value>(&text) {
|
||||
let error = json.get("error").unwrap_or(&json);
|
||||
let code = error
|
||||
.get("code")
|
||||
.and_then(|v| v.as_str())
|
||||
.or_else(|| error.get("type").and_then(|v| v.as_str()))
|
||||
.map(String::from);
|
||||
let message = error
|
||||
.get("message")
|
||||
.and_then(|v| v.as_str())
|
||||
.unwrap_or(&text)
|
||||
.to_string();
|
||||
ClientError::Api {
|
||||
status: Some(status),
|
||||
code,
|
||||
message,
|
||||
retry_after,
|
||||
}
|
||||
} else {
|
||||
ClientError::Api {
|
||||
status: Some(status),
|
||||
code: None,
|
||||
message: text,
|
||||
retry_after,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[async_trait]
|
||||
impl<S: Scheme + Clone + 'static> LlmClient for HttpTransport<S> {
|
||||
fn clone_boxed(&self) -> Box<dyn LlmClient> {
|
||||
Box::new(self.clone())
|
||||
}
|
||||
|
||||
fn validate_config(&self, config: &RequestConfig) -> Vec<ConfigWarning> {
|
||||
self.scheme.validate_config(config)
|
||||
}
|
||||
|
||||
async fn stream(&self, request: Request) -> Result<ResponseStream, ClientError> {
|
||||
let total_started = Instant::now();
|
||||
let path = self.scheme.path(&self.model_id);
|
||||
emit_transport_trace(
|
||||
&request,
|
||||
"transport_start",
|
||||
json!({
|
||||
"model": &self.model_id,
|
||||
"path": path,
|
||||
"auth_kind": auth_kind(&self.auth),
|
||||
"required_auth": format!("{:?}", self.scheme.required_auth()),
|
||||
"codex_backend": self.is_codex_backend(),
|
||||
"cache_key_present": request.cache_key.is_some(),
|
||||
"stream_open_timeout_ms": DEFAULT_STREAM_OPEN_TIMEOUT.as_millis() as u64,
|
||||
}),
|
||||
);
|
||||
|
||||
let url = self.build_url();
|
||||
let headers_started = Instant::now();
|
||||
emit_transport_trace(
|
||||
&request,
|
||||
"transport_headers_start",
|
||||
json!({
|
||||
"auth_kind": auth_kind(&self.auth),
|
||||
"required_auth": format!("{:?}", self.scheme.required_auth()),
|
||||
}),
|
||||
);
|
||||
let mut headers = match self.build_headers().await {
|
||||
Ok(headers) => {
|
||||
emit_transport_trace(
|
||||
&request,
|
||||
"transport_headers_done",
|
||||
json!({
|
||||
"elapsed_ms": headers_started.elapsed().as_millis() as u64,
|
||||
"headers_len": headers.len(),
|
||||
"headers": request_header_diagnostics(&headers),
|
||||
}),
|
||||
);
|
||||
headers
|
||||
}
|
||||
Err(error) => {
|
||||
emit_transport_trace(
|
||||
&request,
|
||||
"transport_headers_error",
|
||||
json!({
|
||||
"elapsed_ms": headers_started.elapsed().as_millis() as u64,
|
||||
"error": error.to_string(),
|
||||
}),
|
||||
);
|
||||
return Err(error);
|
||||
}
|
||||
};
|
||||
|
||||
let stream_headers_started = Instant::now();
|
||||
if let Err(error) = self.apply_stream_headers(&mut headers, &request) {
|
||||
emit_transport_trace(
|
||||
&request,
|
||||
"transport_stream_headers_error",
|
||||
json!({
|
||||
"elapsed_ms": stream_headers_started.elapsed().as_millis() as u64,
|
||||
"error": error.to_string(),
|
||||
}),
|
||||
);
|
||||
return Err(error);
|
||||
}
|
||||
emit_transport_trace(
|
||||
&request,
|
||||
"transport_stream_headers_done",
|
||||
json!({
|
||||
"elapsed_ms": stream_headers_started.elapsed().as_millis() as u64,
|
||||
"headers_len": headers.len(),
|
||||
"headers": request_header_diagnostics(&headers),
|
||||
}),
|
||||
);
|
||||
|
||||
let body_started = Instant::now();
|
||||
emit_transport_trace(&request, "transport_body_build_start", json!({}));
|
||||
let body = self
|
||||
.scheme
|
||||
.build_request_body(&self.model_id, &request, &self.capability);
|
||||
let body_shape = request_body_shape_payload(&body);
|
||||
emit_transport_trace(
|
||||
&request,
|
||||
"transport_body_build_done",
|
||||
json!({
|
||||
"elapsed_ms": body_started.elapsed().as_millis() as u64,
|
||||
"body_kind": json_value_kind(&body),
|
||||
"request_shape": body_shape.clone(),
|
||||
}),
|
||||
);
|
||||
|
||||
let encode_started = Instant::now();
|
||||
let request_body = match self.encode_request_body(&body, &mut headers) {
|
||||
Ok(body) => body,
|
||||
Err(error) => {
|
||||
emit_transport_trace(
|
||||
&request,
|
||||
"transport_body_encode_error",
|
||||
json!({
|
||||
"elapsed_ms": encode_started.elapsed().as_millis() as u64,
|
||||
"error": error.to_string(),
|
||||
}),
|
||||
);
|
||||
return Err(error);
|
||||
}
|
||||
};
|
||||
let final_request_headers = request_header_diagnostics(&headers);
|
||||
let body_compression = request_body.encoding().to_string();
|
||||
emit_transport_trace(
|
||||
&request,
|
||||
"transport_body_encode_done",
|
||||
json!({
|
||||
"elapsed_ms": encode_started.elapsed().as_millis() as u64,
|
||||
"encoding": body_compression.as_str(),
|
||||
"body_compression": body_compression.as_str(),
|
||||
"raw_json_bytes": request_body.raw_json_bytes(),
|
||||
"wire_bytes": request_body.wire_bytes(),
|
||||
"request_shape": body_shape.clone(),
|
||||
"headers": final_request_headers.clone(),
|
||||
}),
|
||||
);
|
||||
|
||||
let builder = self.http_client.post(&url).headers(headers);
|
||||
let builder = match request_body {
|
||||
RequestBody::Json(body) => builder.json(&body),
|
||||
RequestBody::CompressedJson { bytes, .. } => builder.body(bytes),
|
||||
};
|
||||
|
||||
let send_started = Instant::now();
|
||||
emit_transport_trace(&request, "transport_http_send_start", json!({}));
|
||||
let response =
|
||||
match response_with_timeout(builder.send(), DEFAULT_STREAM_OPEN_TIMEOUT, "stream_open")
|
||||
.await
|
||||
{
|
||||
Ok(response) => {
|
||||
let response_headers = response_header_diagnostics(response.headers());
|
||||
emit_transport_trace(
|
||||
&request,
|
||||
"transport_http_headers_received",
|
||||
json!({
|
||||
"elapsed_ms": send_started.elapsed().as_millis() as u64,
|
||||
"status": response.status().as_u16(),
|
||||
"success": response.status().is_success(),
|
||||
"headers": response_headers,
|
||||
}),
|
||||
);
|
||||
response
|
||||
}
|
||||
Err(error) => {
|
||||
emit_transport_trace(
|
||||
&request,
|
||||
"transport_http_send_error",
|
||||
json!({
|
||||
"elapsed_ms": send_started.elapsed().as_millis() as u64,
|
||||
"error": error.to_string(),
|
||||
}),
|
||||
);
|
||||
return Err(error);
|
||||
}
|
||||
};
|
||||
|
||||
if !response.status().is_success() {
|
||||
let status = response.status().as_u16();
|
||||
let retry_after_present = response.headers().get(RETRY_AFTER).is_some();
|
||||
let error = classify_error_response(response).await;
|
||||
let context_length_exceeded = is_context_length_exceeded(&error);
|
||||
emit_transport_trace(
|
||||
&request,
|
||||
"transport_http_status_error",
|
||||
json!({
|
||||
"status": status,
|
||||
"retry_after_present": retry_after_present,
|
||||
"api_error_code": api_error_code(&error),
|
||||
"context_length_exceeded": context_length_exceeded,
|
||||
"provider_usage_absent": context_length_exceeded,
|
||||
"request_shape": body_shape.clone(),
|
||||
"request_headers": final_request_headers.clone(),
|
||||
"body_compression": body_compression.as_str(),
|
||||
}),
|
||||
);
|
||||
return Err(error);
|
||||
}
|
||||
|
||||
emit_transport_trace(
|
||||
&request,
|
||||
"transport_stream_ready",
|
||||
json!({
|
||||
"elapsed_ms": total_started.elapsed().as_millis() as u64,
|
||||
}),
|
||||
);
|
||||
|
||||
let scheme = self.scheme.clone();
|
||||
let status = response.status().as_u16();
|
||||
let response_headers = response_header_diagnostics(response.headers());
|
||||
let transport_trace = request.transport_trace.clone();
|
||||
let byte_stream = response.bytes_stream().map_err(std::io::Error::other);
|
||||
let event_stream = byte_stream.eventsource();
|
||||
|
||||
// scheme 固有のパース状態をストリーム単位で保持する
|
||||
let mut state = <S::State as Default>::default();
|
||||
|
||||
let stream = event_stream
|
||||
.map(move |result| match result {
|
||||
Ok(frame) => match scheme.parse_sse(&frame.event, &frame.data, &mut state) {
|
||||
Ok(events) => Ok(events),
|
||||
Err(e) => Err(e),
|
||||
},
|
||||
Err(e) => {
|
||||
let source = e.to_string();
|
||||
let message = sse_error_context(status, &response_headers, &source);
|
||||
if let Some(trace) = &transport_trace {
|
||||
trace.emit(
|
||||
"transport_sse_parse_error",
|
||||
json!({
|
||||
"status": status,
|
||||
"headers": response_headers.clone(),
|
||||
"error": source,
|
||||
}),
|
||||
);
|
||||
}
|
||||
Err(ClientError::Sse(message))
|
||||
}
|
||||
})
|
||||
.map(|res| {
|
||||
let s: Pin<Box<dyn Stream<Item = Result<Event, ClientError>> + Send>> = match res {
|
||||
Ok(events) => Box::pin(futures::stream::iter(events.into_iter().map(Ok))),
|
||||
Err(e) => Box::pin(futures::stream::once(async move { Err(e) })),
|
||||
};
|
||||
s
|
||||
})
|
||||
.flatten();
|
||||
|
||||
Ok(Box::pin(stream))
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
use serde_json::json;
|
||||
|
||||
#[derive(Debug)]
|
||||
struct TestAuthProvider {
|
||||
codex: bool,
|
||||
}
|
||||
|
||||
#[async_trait]
|
||||
impl AuthProvider for TestAuthProvider {
|
||||
async fn headers(&self) -> Result<Vec<(HeaderName, HeaderValue)>, ClientError> {
|
||||
Ok(vec![
|
||||
(
|
||||
HeaderName::from_static("authorization"),
|
||||
HeaderValue::from_static("Bearer test-token"),
|
||||
),
|
||||
(
|
||||
HeaderName::from_static("chatgpt-account-id"),
|
||||
HeaderValue::from_static("account-1"),
|
||||
),
|
||||
])
|
||||
}
|
||||
|
||||
fn is_codex_backend(&self) -> bool {
|
||||
self.codex
|
||||
}
|
||||
}
|
||||
|
||||
#[derive(Clone)]
|
||||
struct TestScheme;
|
||||
|
||||
impl Scheme for TestScheme {
|
||||
type State = ();
|
||||
|
||||
fn default_base_url(&self) -> &'static str {
|
||||
"https://example.test"
|
||||
}
|
||||
|
||||
fn path(&self, _model_id: &str) -> String {
|
||||
"/responses".to_string()
|
||||
}
|
||||
|
||||
fn required_auth(&self) -> AuthRequirement {
|
||||
AuthRequirement::Bearer
|
||||
}
|
||||
|
||||
fn build_request_body(
|
||||
&self,
|
||||
model_id: &str,
|
||||
request: &Request,
|
||||
_capability: &ModelCapability,
|
||||
) -> serde_json::Value {
|
||||
json!({
|
||||
"model": model_id,
|
||||
"input_len": request.items.len(),
|
||||
"prompt_cache_key": request.cache_key,
|
||||
})
|
||||
}
|
||||
|
||||
fn parse_sse(
|
||||
&self,
|
||||
_event_type: &str,
|
||||
_data: &str,
|
||||
_state: &mut Self::State,
|
||||
) -> Result<Vec<Event>, ClientError> {
|
||||
Ok(Vec::new())
|
||||
}
|
||||
|
||||
fn default_capability(&self) -> ModelCapability {
|
||||
ModelCapability::minimal()
|
||||
}
|
||||
}
|
||||
|
||||
fn transport(auth: ResolvedAuth) -> HttpTransport<TestScheme> {
|
||||
HttpTransport::new(
|
||||
TestScheme,
|
||||
"gpt-test",
|
||||
"https://example.test",
|
||||
auth,
|
||||
ModelCapability::minimal(),
|
||||
)
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn sse_error_context_includes_response_headers() {
|
||||
let headers = json!({
|
||||
"content_type": "application/octet-stream",
|
||||
"content_encoding": "gzip",
|
||||
"transfer_encoding": "chunked",
|
||||
"content_length": "123",
|
||||
});
|
||||
|
||||
let message = sse_error_context(200, &headers, "stream did not contain valid UTF-8");
|
||||
|
||||
assert!(message.contains("HTTP 200"));
|
||||
assert!(message.contains("stream did not contain valid UTF-8"));
|
||||
assert!(message.contains("content-type=application/octet-stream"));
|
||||
assert!(message.contains("content-encoding=gzip"));
|
||||
assert!(message.contains("transfer-encoding=chunked"));
|
||||
assert!(message.contains("content-length=123"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn response_header_diagnostics_redacts_to_safe_header_subset() {
|
||||
let mut headers = HeaderMap::new();
|
||||
headers.insert(CONTENT_TYPE, HeaderValue::from_static("text/event-stream"));
|
||||
headers.insert(CONTENT_ENCODING, HeaderValue::from_static("identity"));
|
||||
headers.insert("authorization", HeaderValue::from_static("Bearer secret"));
|
||||
|
||||
let diagnostics = response_header_diagnostics(&headers);
|
||||
|
||||
assert_eq!(diagnostics["content_type"], "text/event-stream");
|
||||
assert_eq!(diagnostics["content_encoding"], "identity");
|
||||
assert!(diagnostics.get("authorization").is_none());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn request_body_shape_counts_reasoning_encrypted_content() {
|
||||
let payload = request_body_shape_payload(&json!({
|
||||
"reasoning": { "summary": "auto" },
|
||||
"input": [
|
||||
{ "type": "message", "role": "user", "content": [] },
|
||||
{ "type": "reasoning", "encrypted_content": "abc", "summary": [] },
|
||||
{ "type": "reasoning", "encrypted_content": "defgh", "summary": [] }
|
||||
]
|
||||
}));
|
||||
assert_eq!(payload["items_len"], 3);
|
||||
assert_eq!(payload["reasoning_items"], 2);
|
||||
assert_eq!(payload["reasoning_encrypted_content_count"], 2);
|
||||
assert_eq!(payload["reasoning_encrypted_content_bytes"], 8);
|
||||
assert!(payload["items_json_bytes"].as_u64().unwrap() > 0);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn response_timeout_returns_retryable_lifecycle_timeout() {
|
||||
let err = response_with_timeout(
|
||||
std::future::pending::<Result<reqwest::Response, reqwest::Error>>(),
|
||||
Duration::from_millis(5),
|
||||
"stream_open",
|
||||
)
|
||||
.await
|
||||
.unwrap_err();
|
||||
|
||||
assert!(crate::llm_client::error::is_retryable(&err));
|
||||
assert!(matches!(
|
||||
err,
|
||||
ClientError::Timeout {
|
||||
phase: "stream_open",
|
||||
..
|
||||
}
|
||||
));
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn codex_backend_adds_conversation_headers_and_zstd_body() {
|
||||
let transport = transport(ResolvedAuth::Custom(Arc::new(TestAuthProvider {
|
||||
codex: true,
|
||||
})));
|
||||
let request = Request::new().user("hello").cache_key("segment-123");
|
||||
let mut headers = transport.build_headers().await.unwrap();
|
||||
transport
|
||||
.apply_stream_headers(&mut headers, &request)
|
||||
.unwrap();
|
||||
let body = transport.scheme.build_request_body(
|
||||
&transport.model_id,
|
||||
&request,
|
||||
&transport.capability,
|
||||
);
|
||||
let encoded = transport.encode_request_body(&body, &mut headers).unwrap();
|
||||
|
||||
assert_eq!(headers.get(ACCEPT).unwrap(), "text/event-stream");
|
||||
assert_eq!(headers.get("session-id").unwrap(), "segment-123");
|
||||
assert_eq!(headers.get("thread-id").unwrap(), "segment-123");
|
||||
assert_eq!(headers.get("session_id").unwrap(), "segment-123");
|
||||
assert_eq!(headers.get("x-client-request-id").unwrap(), "segment-123");
|
||||
assert_eq!(headers.get(CONTENT_ENCODING).unwrap(), "zstd");
|
||||
|
||||
let diagnostics = request_header_diagnostics(&headers);
|
||||
assert_eq!(diagnostics["content_type"], "application/json");
|
||||
assert_eq!(diagnostics["content_encoding"], "zstd");
|
||||
assert_eq!(diagnostics["accept"], "text/event-stream");
|
||||
assert!(diagnostics["session_id_present"].as_bool().unwrap());
|
||||
assert!(diagnostics["thread_id_present"].as_bool().unwrap());
|
||||
assert!(diagnostics["legacy_session_id_present"].as_bool().unwrap());
|
||||
assert!(
|
||||
diagnostics["x_client_request_id_present"]
|
||||
.as_bool()
|
||||
.unwrap()
|
||||
);
|
||||
assert!(diagnostics["chatgpt_account_id_present"].as_bool().unwrap());
|
||||
|
||||
let RequestBody::CompressedJson {
|
||||
bytes: compressed,
|
||||
raw_json_bytes,
|
||||
} = encoded
|
||||
else {
|
||||
panic!("Codex backend request body must be zstd-compressed");
|
||||
};
|
||||
assert!(raw_json_bytes > 0);
|
||||
let decoded = zstd::stream::decode_all(std::io::Cursor::new(compressed)).unwrap();
|
||||
let decoded: serde_json::Value = serde_json::from_slice(&decoded).unwrap();
|
||||
assert_eq!(decoded["prompt_cache_key"], "segment-123");
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn non_codex_request_does_not_get_codex_only_headers_or_compression() {
|
||||
let transport = transport(ResolvedAuth::ApiKey("api-key".to_string()));
|
||||
let request = Request::new().user("hello").cache_key("segment-123");
|
||||
let mut headers = transport.build_headers().await.unwrap();
|
||||
transport
|
||||
.apply_stream_headers(&mut headers, &request)
|
||||
.unwrap();
|
||||
let body = transport.scheme.build_request_body(
|
||||
&transport.model_id,
|
||||
&request,
|
||||
&transport.capability,
|
||||
);
|
||||
let encoded = transport.encode_request_body(&body, &mut headers).unwrap();
|
||||
|
||||
assert_eq!(headers.get(ACCEPT).unwrap(), "text/event-stream");
|
||||
assert!(headers.get("session-id").is_none());
|
||||
assert!(headers.get("thread-id").is_none());
|
||||
assert!(headers.get("session_id").is_none());
|
||||
assert!(headers.get("x-client-request-id").is_none());
|
||||
assert!(headers.get(CONTENT_ENCODING).is_none());
|
||||
|
||||
let RequestBody::Json(decoded) = encoded else {
|
||||
panic!("non-Codex request body must remain normal JSON");
|
||||
};
|
||||
assert_eq!(decoded["prompt_cache_key"], "segment-123");
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,784 @@
|
||||
//! LLM Client Common Types
|
||||
//!
|
||||
//! Core conversation types for yoi's LLM interaction model.
|
||||
//! The core abstraction is `Item` which represents different types of conversation elements:
|
||||
//! - Message items (user/assistant messages with content parts)
|
||||
//! - ToolCall items (tool invocations)
|
||||
//! - ToolResult items (tool results)
|
||||
//! - Reasoning items (extended thinking)
|
||||
|
||||
use std::{fmt, sync::Arc};
|
||||
|
||||
use serde::{Deserialize, Serialize};
|
||||
|
||||
fn is_false(value: &bool) -> bool {
|
||||
!*value
|
||||
}
|
||||
|
||||
// ============================================================================
|
||||
// Item - The core unit of conversation
|
||||
// ============================================================================
|
||||
|
||||
/// Item ID type for tracking items in a conversation
|
||||
pub type ItemId = String;
|
||||
|
||||
/// Call ID type for linking function calls to their outputs
|
||||
pub type CallId = String;
|
||||
|
||||
/// Callback sink for request-local transport lifecycle diagnostics.
|
||||
///
|
||||
/// This is carried on [`Request`] so generic [`crate::llm_client::LlmClient`]
|
||||
/// implementations can emit fine-grained transport milestones without widening
|
||||
/// the trait method signature. The callback must never receive request body
|
||||
/// contents or secret header values.
|
||||
#[derive(Clone)]
|
||||
pub struct RequestTrace {
|
||||
callback: Arc<dyn Fn(&str, serde_json::Value) + Send + Sync>,
|
||||
}
|
||||
|
||||
impl RequestTrace {
|
||||
pub fn new(callback: impl Fn(&str, serde_json::Value) + Send + Sync + 'static) -> Self {
|
||||
Self {
|
||||
callback: Arc::new(callback),
|
||||
}
|
||||
}
|
||||
|
||||
pub fn emit(&self, label: &str, data: serde_json::Value) {
|
||||
(self.callback)(label, data);
|
||||
}
|
||||
}
|
||||
|
||||
impl fmt::Debug for RequestTrace {
|
||||
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
|
||||
f.debug_struct("RequestTrace").finish_non_exhaustive()
|
||||
}
|
||||
}
|
||||
|
||||
/// Conversation item - the primary unit of conversation history
|
||||
///
|
||||
/// Items represent discrete elements in a conversation. Tool calls and reasoning
|
||||
/// are first-class items rather than parts of messages.
|
||||
///
|
||||
/// # Examples
|
||||
///
|
||||
/// ```ignore
|
||||
/// use llm_engine::Item;
|
||||
///
|
||||
/// let user = Item::user_message("Hello!");
|
||||
/// let assistant = Item::assistant_message("Hi there!");
|
||||
/// let call = Item::tool_call("call_123", "get_weather", json!({"city": "Tokyo"}));
|
||||
/// let result = Item::tool_result("call_123", "Sunny, 25°C");
|
||||
/// ```
|
||||
#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
|
||||
#[serde(tag = "type", rename_all = "snake_case")]
|
||||
pub enum Item {
|
||||
/// User or assistant message with content parts
|
||||
Message {
|
||||
/// Optional item ID
|
||||
#[serde(skip_serializing_if = "Option::is_none")]
|
||||
id: Option<ItemId>,
|
||||
/// Message role
|
||||
role: Role,
|
||||
/// Content parts
|
||||
content: Vec<ContentPart>,
|
||||
/// Item status
|
||||
#[serde(skip_serializing_if = "Option::is_none")]
|
||||
status: Option<ItemStatus>,
|
||||
},
|
||||
|
||||
/// Tool call from the assistant
|
||||
ToolCall {
|
||||
/// Optional item ID
|
||||
#[serde(skip_serializing_if = "Option::is_none")]
|
||||
id: Option<ItemId>,
|
||||
/// Call ID for linking to result
|
||||
call_id: CallId,
|
||||
/// Tool name
|
||||
name: String,
|
||||
/// Tool arguments as JSON string
|
||||
arguments: String,
|
||||
/// Item status
|
||||
#[serde(skip_serializing_if = "Option::is_none")]
|
||||
status: Option<ItemStatus>,
|
||||
},
|
||||
|
||||
/// Tool call result
|
||||
ToolResult {
|
||||
/// Optional item ID
|
||||
#[serde(skip_serializing_if = "Option::is_none")]
|
||||
id: Option<ItemId>,
|
||||
/// Call ID linking to the tool call
|
||||
call_id: CallId,
|
||||
/// Short summary (always kept in history, survives pruning)
|
||||
summary: String,
|
||||
/// Detailed output (removed by pruning when old enough)
|
||||
#[serde(default, skip_serializing_if = "Option::is_none")]
|
||||
content: Option<String>,
|
||||
/// Whether the tool result represents an execution error.
|
||||
#[serde(default, skip_serializing_if = "is_false")]
|
||||
is_error: bool,
|
||||
},
|
||||
|
||||
/// Reasoning/thinking item
|
||||
Reasoning {
|
||||
/// Optional item ID
|
||||
#[serde(skip_serializing_if = "Option::is_none")]
|
||||
id: Option<ItemId>,
|
||||
/// Reasoning text（reasoning body, `reasoning_text.delta` の累積）
|
||||
text: String,
|
||||
/// Reasoning summary（OpenAI Responses の `summary_text[]` を格納。
|
||||
/// 他 scheme は空）
|
||||
#[serde(default, skip_serializing_if = "Vec::is_empty")]
|
||||
summary: Vec<String>,
|
||||
/// サーバから返された暗号化済み reasoning blob。ZDR / `store=false`
|
||||
/// 運用で stateless に再送するときそのまま添える必要がある。
|
||||
/// Anthropic の `redacted_thinking.data` もここに格納する。
|
||||
#[serde(default, skip_serializing_if = "Option::is_none")]
|
||||
encrypted_content: Option<String>,
|
||||
/// Anthropic extended thinking の `signature`。新世代 Claude
|
||||
/// (Opus 4.5+/Sonnet 4.6+) では同一論理ターン内の `thinking`
|
||||
/// ブロックを送り返す際に必須。改ざん検知に使われる。他 scheme
|
||||
/// では `None`。
|
||||
#[serde(default, skip_serializing_if = "Option::is_none")]
|
||||
signature: Option<String>,
|
||||
/// Item status
|
||||
#[serde(skip_serializing_if = "Option::is_none")]
|
||||
status: Option<ItemStatus>,
|
||||
},
|
||||
}
|
||||
|
||||
impl Item {
|
||||
// ========================================================================
|
||||
// Message constructors
|
||||
// ========================================================================
|
||||
|
||||
/// Create a system message item with text content.
|
||||
///
|
||||
/// System items in history are sent as `role: "system"` on OpenAI,
|
||||
/// and as `role: "user"` on Anthropic/Gemini (which lack a system
|
||||
/// role in conversation items).
|
||||
pub fn system_message(text: impl Into<String>) -> Self {
|
||||
Self::Message {
|
||||
id: None,
|
||||
role: Role::System,
|
||||
content: vec![ContentPart::Text { text: text.into() }],
|
||||
status: None,
|
||||
}
|
||||
}
|
||||
|
||||
/// Create a user message item with text content
|
||||
pub fn user_message(text: impl Into<String>) -> Self {
|
||||
Self::Message {
|
||||
id: None,
|
||||
role: Role::User,
|
||||
content: vec![ContentPart::Text { text: text.into() }],
|
||||
status: None,
|
||||
}
|
||||
}
|
||||
|
||||
/// Create a user message item with multiple content parts
|
||||
pub fn user_message_parts(parts: Vec<ContentPart>) -> Self {
|
||||
Self::Message {
|
||||
id: None,
|
||||
role: Role::User,
|
||||
content: parts,
|
||||
status: None,
|
||||
}
|
||||
}
|
||||
|
||||
/// Create an assistant message item with text content
|
||||
pub fn assistant_message(text: impl Into<String>) -> Self {
|
||||
Self::Message {
|
||||
id: None,
|
||||
role: Role::Assistant,
|
||||
content: vec![ContentPart::Text { text: text.into() }],
|
||||
status: None,
|
||||
}
|
||||
}
|
||||
|
||||
/// Create an assistant message item with multiple content parts
|
||||
pub fn assistant_message_parts(parts: Vec<ContentPart>) -> Self {
|
||||
Self::Message {
|
||||
id: None,
|
||||
role: Role::Assistant,
|
||||
content: parts,
|
||||
status: None,
|
||||
}
|
||||
}
|
||||
|
||||
// ========================================================================
|
||||
// Tool call constructors
|
||||
// ========================================================================
|
||||
|
||||
/// Create a tool call item
|
||||
pub fn tool_call(
|
||||
call_id: impl Into<String>,
|
||||
name: impl Into<String>,
|
||||
arguments: impl Into<String>,
|
||||
) -> Self {
|
||||
Self::ToolCall {
|
||||
id: None,
|
||||
call_id: call_id.into(),
|
||||
name: name.into(),
|
||||
arguments: arguments.into(),
|
||||
status: None,
|
||||
}
|
||||
}
|
||||
|
||||
/// Create a tool call item from a JSON value
|
||||
pub fn tool_call_json(
|
||||
call_id: impl Into<String>,
|
||||
name: impl Into<String>,
|
||||
arguments: serde_json::Value,
|
||||
) -> Self {
|
||||
Self::tool_call(call_id, name, arguments.to_string())
|
||||
}
|
||||
|
||||
/// Create a tool result item with summary only (no content).
|
||||
pub fn tool_result(call_id: impl Into<String>, summary: impl Into<String>) -> Self {
|
||||
Self::tool_result_item(call_id, summary, None, false)
|
||||
}
|
||||
|
||||
/// Create an error tool result item with summary only (no content).
|
||||
pub fn tool_result_error(call_id: impl Into<String>, summary: impl Into<String>) -> Self {
|
||||
Self::tool_result_item(call_id, summary, None, true)
|
||||
}
|
||||
|
||||
/// Create a tool result item with summary, optional content, and error flag.
|
||||
pub fn tool_result_item(
|
||||
call_id: impl Into<String>,
|
||||
summary: impl Into<String>,
|
||||
content: Option<String>,
|
||||
is_error: bool,
|
||||
) -> Self {
|
||||
Self::ToolResult {
|
||||
id: None,
|
||||
call_id: call_id.into(),
|
||||
summary: summary.into(),
|
||||
content,
|
||||
is_error,
|
||||
}
|
||||
}
|
||||
|
||||
/// Create a tool result item with summary and content.
|
||||
pub fn tool_result_with_content(
|
||||
call_id: impl Into<String>,
|
||||
summary: impl Into<String>,
|
||||
content: impl Into<String>,
|
||||
) -> Self {
|
||||
Self::tool_result_item(call_id, summary, Some(content.into()), false)
|
||||
}
|
||||
|
||||
// ========================================================================
|
||||
// Reasoning constructors
|
||||
// ========================================================================
|
||||
|
||||
/// Create a reasoning item
|
||||
pub fn reasoning(text: impl Into<String>) -> Self {
|
||||
Self::Reasoning {
|
||||
id: None,
|
||||
text: text.into(),
|
||||
summary: Vec::new(),
|
||||
encrypted_content: None,
|
||||
signature: None,
|
||||
status: None,
|
||||
}
|
||||
}
|
||||
|
||||
/// Set reasoning summary on a `Reasoning` item. No-op on other variants.
|
||||
pub fn with_reasoning_summary(mut self, new_summary: Vec<String>) -> Self {
|
||||
if let Self::Reasoning { summary, .. } = &mut self {
|
||||
*summary = new_summary;
|
||||
}
|
||||
self
|
||||
}
|
||||
|
||||
/// Set `encrypted_content` on a `Reasoning` item. No-op on other variants.
|
||||
pub fn with_encrypted_content(mut self, content: impl Into<String>) -> Self {
|
||||
if let Self::Reasoning {
|
||||
encrypted_content, ..
|
||||
} = &mut self
|
||||
{
|
||||
*encrypted_content = Some(content.into());
|
||||
}
|
||||
self
|
||||
}
|
||||
|
||||
/// Set Anthropic `signature` on a `Reasoning` item. No-op on other variants.
|
||||
pub fn with_signature(mut self, sig: impl Into<String>) -> Self {
|
||||
if let Self::Reasoning { signature, .. } = &mut self {
|
||||
*signature = Some(sig.into());
|
||||
}
|
||||
self
|
||||
}
|
||||
|
||||
// ========================================================================
|
||||
// Builder methods
|
||||
// ========================================================================
|
||||
|
||||
/// Set the item ID
|
||||
pub fn with_id(mut self, id: impl Into<String>) -> Self {
|
||||
match &mut self {
|
||||
Self::Message { id: item_id, .. } => *item_id = Some(id.into()),
|
||||
Self::ToolCall { id: item_id, .. } => *item_id = Some(id.into()),
|
||||
Self::ToolResult { id: item_id, .. } => *item_id = Some(id.into()),
|
||||
Self::Reasoning { id: item_id, .. } => *item_id = Some(id.into()),
|
||||
}
|
||||
self
|
||||
}
|
||||
|
||||
/// Set the item status
|
||||
pub fn with_status(mut self, new_status: ItemStatus) -> Self {
|
||||
match &mut self {
|
||||
Self::Message { status, .. } => *status = Some(new_status),
|
||||
Self::ToolCall { status, .. } => *status = Some(new_status),
|
||||
Self::ToolResult { .. } => {} // Result items don't have status
|
||||
Self::Reasoning { status, .. } => *status = Some(new_status),
|
||||
}
|
||||
self
|
||||
}
|
||||
|
||||
// ========================================================================
|
||||
// Accessors
|
||||
// ========================================================================
|
||||
|
||||
/// Get the item ID if set
|
||||
pub fn id(&self) -> Option<&str> {
|
||||
match self {
|
||||
Self::Message { id, .. } => id.as_deref(),
|
||||
Self::ToolCall { id, .. } => id.as_deref(),
|
||||
Self::ToolResult { id, .. } => id.as_deref(),
|
||||
Self::Reasoning { id, .. } => id.as_deref(),
|
||||
}
|
||||
}
|
||||
|
||||
/// Get the item type as a string
|
||||
pub fn item_type(&self) -> &'static str {
|
||||
match self {
|
||||
Self::Message { .. } => "message",
|
||||
Self::ToolCall { .. } => "tool_call",
|
||||
Self::ToolResult { .. } => "tool_result",
|
||||
Self::Reasoning { .. } => "reasoning",
|
||||
}
|
||||
}
|
||||
|
||||
/// Check if this is a user message
|
||||
pub fn is_user_message(&self) -> bool {
|
||||
matches!(
|
||||
self,
|
||||
Self::Message {
|
||||
role: Role::User,
|
||||
..
|
||||
}
|
||||
)
|
||||
}
|
||||
|
||||
/// Check if this is an assistant message
|
||||
pub fn is_assistant_message(&self) -> bool {
|
||||
matches!(
|
||||
self,
|
||||
Self::Message {
|
||||
role: Role::Assistant,
|
||||
..
|
||||
}
|
||||
)
|
||||
}
|
||||
|
||||
/// Check if this is a tool call
|
||||
pub fn is_tool_call(&self) -> bool {
|
||||
matches!(self, Self::ToolCall { .. })
|
||||
}
|
||||
|
||||
/// Check if this is a tool result
|
||||
pub fn is_tool_result(&self) -> bool {
|
||||
matches!(self, Self::ToolResult { .. })
|
||||
}
|
||||
|
||||
/// Check if this is a reasoning item
|
||||
pub fn is_reasoning(&self) -> bool {
|
||||
matches!(self, Self::Reasoning { .. })
|
||||
}
|
||||
|
||||
/// Get text content if this is a simple text message
|
||||
pub fn as_text(&self) -> Option<&str> {
|
||||
match self {
|
||||
Self::Message { content, .. } if content.len() == 1 => match &content[0] {
|
||||
ContentPart::Text { text } => Some(text),
|
||||
_ => None,
|
||||
},
|
||||
_ => None,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// Parse a ToolCall `arguments` string into a JSON object.
|
||||
///
|
||||
/// Tool call arguments must be a JSON object at the provider API level
|
||||
/// (Anthropic rejects non-object `tool_use.input`). This helper normalizes
|
||||
/// anything that is not a JSON object — empty string, the literal `"null"`,
|
||||
/// arrays, scalars, or parse failures — to an empty object `{}`.
|
||||
pub fn parse_tool_arguments(arguments: &str) -> serde_json::Value {
|
||||
match serde_json::from_str::<serde_json::Value>(arguments) {
|
||||
Ok(value) if value.is_object() => value,
|
||||
_ => serde_json::Value::Object(serde_json::Map::new()),
|
||||
}
|
||||
}
|
||||
|
||||
// ============================================================================
|
||||
// Content Parts - Components within message items
|
||||
// ============================================================================
|
||||
|
||||
/// Content part within a message item
|
||||
///
|
||||
/// Text content is role-agnostic; the containing Item's Role determines
|
||||
/// whether it's user input or assistant output.
|
||||
#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
|
||||
#[serde(tag = "type", rename_all = "snake_case")]
|
||||
pub enum ContentPart {
|
||||
/// Text content
|
||||
Text {
|
||||
/// The text content
|
||||
text: String,
|
||||
},
|
||||
|
||||
/// Refusal content (for assistant messages)
|
||||
Refusal {
|
||||
/// The refusal message
|
||||
refusal: String,
|
||||
},
|
||||
}
|
||||
|
||||
impl ContentPart {
|
||||
/// Create a text part
|
||||
pub fn text(text: impl Into<String>) -> Self {
|
||||
Self::Text { text: text.into() }
|
||||
}
|
||||
|
||||
/// Create a refusal part
|
||||
pub fn refusal(refusal: impl Into<String>) -> Self {
|
||||
Self::Refusal {
|
||||
refusal: refusal.into(),
|
||||
}
|
||||
}
|
||||
|
||||
/// Get the text content regardless of type
|
||||
pub fn as_text(&self) -> &str {
|
||||
match self {
|
||||
Self::Text { text } => text,
|
||||
Self::Refusal { refusal } => refusal,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// ============================================================================
|
||||
// Role and Status
|
||||
// ============================================================================
|
||||
|
||||
/// Message role
|
||||
#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
|
||||
#[serde(rename_all = "lowercase")]
|
||||
pub enum Role {
|
||||
/// User
|
||||
User,
|
||||
/// Assistant
|
||||
Assistant,
|
||||
/// System (for system prompts, not typically used in items)
|
||||
System,
|
||||
}
|
||||
|
||||
/// Item status
|
||||
#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
|
||||
#[serde(rename_all = "lowercase")]
|
||||
pub enum ItemStatus {
|
||||
/// Item is being generated
|
||||
InProgress,
|
||||
/// Item completed successfully
|
||||
Completed,
|
||||
/// Item was truncated (e.g., max tokens)
|
||||
Incomplete,
|
||||
}
|
||||
|
||||
// ============================================================================
|
||||
// Request Types
|
||||
// ============================================================================
|
||||
|
||||
/// LLM Request
|
||||
#[derive(Debug, Clone, Default)]
|
||||
pub struct Request {
|
||||
/// System prompt (instructions)
|
||||
pub system_prompt: Option<String>,
|
||||
/// Input items (conversation history)
|
||||
pub items: Vec<Item>,
|
||||
/// Tool definitions
|
||||
pub tools: Vec<ToolDefinition>,
|
||||
/// Request configuration
|
||||
pub config: RequestConfig,
|
||||
/// Index into `items` marking the end of a stable, cacheable prefix.
|
||||
///
|
||||
/// Higher layers that know about durable prefix boundaries (e.g. a
|
||||
/// post-compaction summary) set this so that caching-aware providers
|
||||
/// (Anthropic today) can place a long-lived cache breakpoint there.
|
||||
/// Providers without prompt caching ignore the field.
|
||||
pub cache_anchor: Option<usize>,
|
||||
/// 会話単位の安定キー。`prompt_cache_key` として送られる
|
||||
/// (OpenAI Responses)。ChatGPT backend (codex-oauth) は明示キーが
|
||||
/// 無いと org/project ハッシュ衝突でプロンプトキャッシュが
|
||||
/// ほぼヒットしないため、pod 側で `SegmentId` を渡す運用を想定。
|
||||
/// `cache_anchor` と違い名前空間キーであり、`prefix anchor` とは
|
||||
/// 別の概念。`cache_anchor` を読まない provider と同じく、
|
||||
/// `prompt_cache_key` を持たない provider は無視する。
|
||||
pub cache_key: Option<String>,
|
||||
/// Request-local diagnostics sink for transport lifecycle tracing.
|
||||
#[doc(hidden)]
|
||||
pub transport_trace: Option<RequestTrace>,
|
||||
}
|
||||
|
||||
impl Request {
|
||||
/// Create a new empty request
|
||||
pub fn new() -> Self {
|
||||
Self::default()
|
||||
}
|
||||
|
||||
/// Set the system prompt
|
||||
pub fn system(mut self, prompt: impl Into<String>) -> Self {
|
||||
self.system_prompt = Some(prompt.into());
|
||||
self
|
||||
}
|
||||
|
||||
/// Add a user message
|
||||
pub fn user(mut self, content: impl Into<String>) -> Self {
|
||||
self.items.push(Item::user_message(content));
|
||||
self
|
||||
}
|
||||
|
||||
/// Add an assistant message
|
||||
pub fn assistant(mut self, content: impl Into<String>) -> Self {
|
||||
self.items.push(Item::assistant_message(content));
|
||||
self
|
||||
}
|
||||
|
||||
/// Add an item
|
||||
pub fn item(mut self, item: Item) -> Self {
|
||||
self.items.push(item);
|
||||
self
|
||||
}
|
||||
|
||||
/// Add multiple items
|
||||
pub fn items(mut self, items: impl IntoIterator<Item = Item>) -> Self {
|
||||
self.items.extend(items);
|
||||
self
|
||||
}
|
||||
|
||||
/// Add a tool definition
|
||||
pub fn tool(mut self, tool: ToolDefinition) -> Self {
|
||||
self.tools.push(tool);
|
||||
self
|
||||
}
|
||||
|
||||
/// Set the request config
|
||||
pub fn config(mut self, config: RequestConfig) -> Self {
|
||||
self.config = config;
|
||||
self
|
||||
}
|
||||
|
||||
/// Attach a request-local transport trace callback.
|
||||
pub fn transport_trace(
|
||||
mut self,
|
||||
callback: impl Fn(&str, serde_json::Value) + Send + Sync + 'static,
|
||||
) -> Self {
|
||||
self.transport_trace = Some(RequestTrace::new(callback));
|
||||
self
|
||||
}
|
||||
|
||||
/// Set max tokens
|
||||
pub fn max_tokens(mut self, max_tokens: u32) -> Self {
|
||||
self.config.max_tokens = Some(max_tokens);
|
||||
self
|
||||
}
|
||||
|
||||
/// Set temperature
|
||||
pub fn temperature(mut self, temperature: f32) -> Self {
|
||||
self.config.temperature = Some(temperature);
|
||||
self
|
||||
}
|
||||
|
||||
/// Set top_p
|
||||
pub fn top_p(mut self, top_p: f32) -> Self {
|
||||
self.config.top_p = Some(top_p);
|
||||
self
|
||||
}
|
||||
|
||||
/// Set top_k
|
||||
pub fn top_k(mut self, top_k: u32) -> Self {
|
||||
self.config.top_k = Some(top_k);
|
||||
self
|
||||
}
|
||||
|
||||
/// Add a stop sequence
|
||||
pub fn stop_sequence(mut self, sequence: impl Into<String>) -> Self {
|
||||
self.config.stop_sequences.push(sequence.into());
|
||||
self
|
||||
}
|
||||
|
||||
/// Set the conversation cache key.
|
||||
///
|
||||
/// 詳細は [`Request::cache_key`] のフィールドコメント参照。
|
||||
pub fn cache_key(mut self, key: impl Into<String>) -> Self {
|
||||
self.cache_key = Some(key.into());
|
||||
self
|
||||
}
|
||||
}
|
||||
|
||||
// ============================================================================
|
||||
// Tool Definition
|
||||
// ============================================================================
|
||||
|
||||
/// Tool (function) definition
|
||||
#[derive(Debug, Clone, Serialize, Deserialize)]
|
||||
pub struct ToolDefinition {
|
||||
/// Tool name
|
||||
pub name: String,
|
||||
/// Tool description
|
||||
pub description: Option<String>,
|
||||
/// Input schema (JSON Schema)
|
||||
pub input_schema: serde_json::Value,
|
||||
}
|
||||
|
||||
impl ToolDefinition {
|
||||
/// Create a new tool definition
|
||||
pub fn new(name: impl Into<String>) -> Self {
|
||||
Self {
|
||||
name: name.into(),
|
||||
description: None,
|
||||
input_schema: serde_json::json!({
|
||||
"type": "object",
|
||||
"properties": {}
|
||||
}),
|
||||
}
|
||||
}
|
||||
|
||||
/// Set the description
|
||||
pub fn description(mut self, desc: impl Into<String>) -> Self {
|
||||
self.description = Some(desc.into());
|
||||
self
|
||||
}
|
||||
|
||||
/// Set the input schema
|
||||
pub fn input_schema(mut self, schema: serde_json::Value) -> Self {
|
||||
self.input_schema = schema;
|
||||
self
|
||||
}
|
||||
}
|
||||
|
||||
// ============================================================================
|
||||
// Request Config
|
||||
// ============================================================================
|
||||
|
||||
/// Request configuration
|
||||
#[derive(Debug, Clone, Default, Serialize, Deserialize)]
|
||||
pub struct RequestConfig {
|
||||
/// Maximum tokens to generate
|
||||
pub max_tokens: Option<u32>,
|
||||
/// Temperature (randomness)
|
||||
pub temperature: Option<f32>,
|
||||
/// Top P (nucleus sampling)
|
||||
pub top_p: Option<f32>,
|
||||
/// Top K
|
||||
pub top_k: Option<u32>,
|
||||
/// Stop sequences
|
||||
pub stop_sequences: Vec<String>,
|
||||
/// Reasoning / extended-thinking 制御（共通型、scheme 側で各社形式に投影）。
|
||||
///
|
||||
/// `None` のときは何も送らない。`Some` でも scheme の
|
||||
/// `ModelCapability::reasoning` が `None` なら無視される。
|
||||
#[serde(default)]
|
||||
pub reasoning: Option<crate::llm_client::capability::ReasoningControl>,
|
||||
}
|
||||
|
||||
impl RequestConfig {
|
||||
/// Create a new default config
|
||||
pub fn new() -> Self {
|
||||
Self::default()
|
||||
}
|
||||
|
||||
/// Set max tokens
|
||||
pub fn with_max_tokens(mut self, max_tokens: u32) -> Self {
|
||||
self.max_tokens = Some(max_tokens);
|
||||
self
|
||||
}
|
||||
|
||||
/// Set temperature
|
||||
pub fn with_temperature(mut self, temperature: f32) -> Self {
|
||||
self.temperature = Some(temperature);
|
||||
self
|
||||
}
|
||||
|
||||
/// Set top_p
|
||||
pub fn with_top_p(mut self, top_p: f32) -> Self {
|
||||
self.top_p = Some(top_p);
|
||||
self
|
||||
}
|
||||
|
||||
/// Set top_k
|
||||
pub fn with_top_k(mut self, top_k: u32) -> Self {
|
||||
self.top_k = Some(top_k);
|
||||
self
|
||||
}
|
||||
|
||||
/// Add a stop sequence
|
||||
pub fn with_stop_sequence(mut self, sequence: impl Into<String>) -> Self {
|
||||
self.stop_sequences.push(sequence.into());
|
||||
self
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod parse_tool_arguments_tests {
|
||||
use super::parse_tool_arguments;
|
||||
use serde_json::{Value, json};
|
||||
|
||||
fn empty_object() -> Value {
|
||||
Value::Object(serde_json::Map::new())
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn empty_string_normalizes_to_object() {
|
||||
assert_eq!(parse_tool_arguments(""), empty_object());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn literal_null_normalizes_to_object() {
|
||||
// 既存セッションに残っている "null" が resume 時に復旧できること
|
||||
assert_eq!(parse_tool_arguments("null"), empty_object());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn array_normalizes_to_object() {
|
||||
assert_eq!(parse_tool_arguments("[1, 2, 3]"), empty_object());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn scalar_normalizes_to_object() {
|
||||
assert_eq!(parse_tool_arguments("42"), empty_object());
|
||||
assert_eq!(parse_tool_arguments("\"str\""), empty_object());
|
||||
assert_eq!(parse_tool_arguments("true"), empty_object());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn invalid_json_normalizes_to_object() {
|
||||
assert_eq!(parse_tool_arguments("{not json"), empty_object());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn valid_object_passes_through() {
|
||||
assert_eq!(
|
||||
parse_tool_arguments(r#"{"city":"Tokyo","days":3}"#),
|
||||
json!({"city": "Tokyo", "days": 3}),
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn empty_object_passes_through() {
|
||||
assert_eq!(parse_tool_arguments("{}"), empty_object());
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,14 @@
|
||||
//! Message and Item Types
|
||||
//!
|
||||
//! Core types for representing conversation items.
|
||||
//!
|
||||
//! The primary type is [`Item`], which represents different kinds of conversation
|
||||
//! elements: messages, tool calls, tool results, and reasoning.
|
||||
|
||||
// Re-export all types from llm_client::types
|
||||
pub use crate::llm_client::types::{ContentPart, Item, Role};
|
||||
|
||||
/// Convenience alias
|
||||
///
|
||||
/// Messages are just one type of Item.
|
||||
pub type Message = Item;
|
||||
@@ -0,0 +1,451 @@
|
||||
//! Prune — context projection for old tool-result content.
|
||||
//!
|
||||
//! LLM 送信時のコンテキストから古い [`Item::ToolResult`] の `content` を
|
||||
//! 省略して、コンテキスト窓のトークンを回収する。`summary` は残すので
|
||||
//! 「何が起きたか」の痕跡は保たれる。
|
||||
//!
|
||||
//! # 設計方針
|
||||
//!
|
||||
//! Prune は **コンテキスト射影** であり、history の変換ではない。
|
||||
//! この crate が提供するのは pure な候補抽出 [`prunable_indices`] のみで、
|
||||
//! 射影の適用は上位層（`pod::prune_hook` 等）が LLM に送る一時コンテキスト
|
||||
//! に対してだけ行う。Engine の永続履歴は決して変更されない。
|
||||
//!
|
||||
//! 保護境界は末尾 token budget で決めるが、この crate は usage 履歴を
|
||||
//! 所有しない。prefix ごとの token 推定値と savings 推定は上位層から
|
||||
//! callback で注入される。
|
||||
|
||||
use serde::{Deserialize, Serialize};
|
||||
|
||||
use crate::llm_client::types::Item;
|
||||
use crate::token_counter::{EstimateSource, TokenEstimate};
|
||||
|
||||
/// Callback that returns token estimates for every prefix boundary of the
|
||||
/// supplied request history.
|
||||
///
|
||||
/// The returned slice must have `history.len() + 1` entries where entry `i`
|
||||
/// estimates the token count of `history[..i]`. Returning a malformed vector,
|
||||
/// or estimates whose source is [`EstimateSource::NoData`], makes prune treat
|
||||
/// the request as having no candidates.
|
||||
pub type TokenEstimator = Box<dyn Fn(&[Item]) -> Vec<TokenEstimate> + Send + Sync>;
|
||||
|
||||
/// Callback that estimates the token savings for projecting the
|
||||
/// `ToolResult.content` out of `history[i]` for each `i` in `indices`.
|
||||
///
|
||||
/// Injected into [`crate::Engine`] via `set_savings_estimator` so the
|
||||
/// Engine can make `min_savings` decisions without knowing about usage
|
||||
/// measurement sources. Return `0` to signal "no data / refuse to prune".
|
||||
///
|
||||
/// 推定対象は「drop する範囲全体」ではなく「content を None にする差分」
|
||||
/// であることに注意。item 自体（summary 等）は残るので、この callback は
|
||||
/// 実際の projection と一致する savings を返す必要がある。
|
||||
pub type SavingsEstimator = Box<dyn Fn(&[Item], &[usize]) -> u64 + Send + Sync>;
|
||||
|
||||
/// Result of one prune evaluation pass, surfaced to the optional
|
||||
/// [`PruneObserver`] for instrumentation.
|
||||
///
|
||||
/// Engine は LLM リクエストごとに 1 回 prune の評価をし、その結果を
|
||||
/// （observer が登録されていれば）この値で通知する。fire/skip の判定
|
||||
/// 結果と、判定材料になった候補数 / 推定 savings / 保護領域の先頭 index を持つ。
|
||||
#[derive(Debug, Clone)]
|
||||
pub struct PruneEvaluation {
|
||||
/// `prunable_indices` の長さ。`Skipped::NoCandidates` の時は 0。
|
||||
pub candidate_count: usize,
|
||||
/// 推定された savings (tokens)。`NoCandidates` の時は 0。
|
||||
pub estimated_savings: u64,
|
||||
/// Token budget で保護される suffix の先頭 item index。
|
||||
/// usage 推定が `NoData` で境界が決まらない場合は `None`。
|
||||
pub protected_start_index: Option<usize>,
|
||||
/// 判定結果。
|
||||
pub decision: PruneDecision,
|
||||
}
|
||||
|
||||
/// Outcome of one prune evaluation. Each variant is one branch of the
|
||||
/// "fire vs skip" decision tree the Engine walks before each LLM request.
|
||||
#[derive(Debug, Clone, PartialEq, Eq)]
|
||||
pub enum PruneDecision {
|
||||
/// `prunable_indices` が空 → 何もしない。
|
||||
SkippedNoCandidates,
|
||||
/// 候補はあったが推定 savings が `min_savings` 未満 → 何もしない。
|
||||
SkippedBelowMinSavings,
|
||||
/// 候補があり savings >= min_savings → projection を適用した。
|
||||
/// `pruned_count` は `project()` が実際に書き換えた item 数
|
||||
/// （既に content=None だった候補は 0 計上）。
|
||||
Fired { pruned_count: usize },
|
||||
}
|
||||
|
||||
/// Optional observer invoked after each prune evaluation, regardless of
|
||||
/// branch. Pod 等の上位層が install して metrics を発行する。
|
||||
pub type PruneObserver = Box<dyn Fn(&PruneEvaluation) + Send + Sync>;
|
||||
|
||||
/// Configuration for the Prune algorithm.
|
||||
#[derive(Debug, Clone, Serialize, Deserialize)]
|
||||
pub struct PruneConfig {
|
||||
/// Token budget at the history tail protected from pruning.
|
||||
#[serde(default = "default_protected_tokens")]
|
||||
pub protected_tokens: u64,
|
||||
|
||||
/// Minimum token savings required to actually prune. If the prunable
|
||||
/// content is smaller than this, the caller should skip to avoid
|
||||
/// pointless KV-cache invalidation. The unit is tokens; the caller
|
||||
/// is responsible for measuring savings via a usage-history-aware
|
||||
/// estimator and comparing against this threshold.
|
||||
#[serde(default = "default_min_savings")]
|
||||
pub min_savings: u64,
|
||||
}
|
||||
|
||||
fn default_protected_tokens() -> u64 {
|
||||
8000
|
||||
}
|
||||
fn default_min_savings() -> u64 {
|
||||
4096
|
||||
}
|
||||
|
||||
impl Default for PruneConfig {
|
||||
fn default() -> Self {
|
||||
Self {
|
||||
protected_tokens: default_protected_tokens(),
|
||||
min_savings: default_min_savings(),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// Set `content = None` on each `Item::ToolResult` at the given indices.
|
||||
///
|
||||
/// Returns the number of items that were actually modified — items that
|
||||
/// are already content-less are counted as 0. Intended for use on a
|
||||
/// request-context clone (never on a persistent history).
|
||||
pub fn project(items: &mut [Item], indices: &[usize]) -> usize {
|
||||
let mut count = 0;
|
||||
for &i in indices {
|
||||
if let Item::ToolResult { content, .. } = &mut items[i]
|
||||
&& content.is_some()
|
||||
{
|
||||
*content = None;
|
||||
count += 1;
|
||||
}
|
||||
}
|
||||
count
|
||||
}
|
||||
|
||||
/// Indices of `Item::ToolResult { content: Some(_), .. }` that lie before
|
||||
/// the suffix protected by `protected_tokens`. Pure: does not mutate `items`.
|
||||
///
|
||||
/// Returns an empty vector when token estimates are unavailable (`NoData`) or
|
||||
/// no prunable candidates exist.
|
||||
pub fn prunable_indices(
|
||||
items: &[Item],
|
||||
protected_tokens: u64,
|
||||
token_estimates: &[TokenEstimate],
|
||||
) -> Vec<usize> {
|
||||
evaluate_candidates(items, protected_tokens, token_estimates).0
|
||||
}
|
||||
|
||||
/// Same as [`prunable_indices`] but also returns the start index of the
|
||||
/// protected suffix. `None` means the token boundary could not be determined
|
||||
/// (currently because usage estimates were `NoData` or malformed).
|
||||
pub fn evaluate_candidates(
|
||||
items: &[Item],
|
||||
protected_tokens: u64,
|
||||
token_estimates: &[TokenEstimate],
|
||||
) -> (Vec<usize>, Option<usize>) {
|
||||
let Some(protected_start) = protected_start_index(items, protected_tokens, token_estimates)
|
||||
else {
|
||||
return (Vec::new(), None);
|
||||
};
|
||||
|
||||
let candidates = items[..protected_start]
|
||||
.iter()
|
||||
.enumerate()
|
||||
.filter_map(|(i, item)| match item {
|
||||
Item::ToolResult {
|
||||
content: Some(_), ..
|
||||
} => Some(i),
|
||||
_ => None,
|
||||
})
|
||||
.collect();
|
||||
(candidates, Some(protected_start))
|
||||
}
|
||||
|
||||
fn protected_start_index(
|
||||
items: &[Item],
|
||||
protected_tokens: u64,
|
||||
token_estimates: &[TokenEstimate],
|
||||
) -> Option<usize> {
|
||||
if token_estimates.len() != items.len() + 1 {
|
||||
return None;
|
||||
}
|
||||
let total = token_estimates[items.len()];
|
||||
if total.source == EstimateSource::NoData {
|
||||
return None;
|
||||
}
|
||||
if protected_tokens == 0 {
|
||||
return Some(items.len());
|
||||
}
|
||||
|
||||
let mut protected_start = items.len();
|
||||
for idx in (0..items.len()).rev() {
|
||||
let prefix = token_estimates[idx];
|
||||
if prefix.source == EstimateSource::NoData {
|
||||
return None;
|
||||
}
|
||||
protected_start = idx;
|
||||
let tail_tokens = total.tokens.saturating_sub(prefix.tokens);
|
||||
if tail_tokens >= protected_tokens {
|
||||
break;
|
||||
}
|
||||
}
|
||||
Some(protected_start)
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
|
||||
/// Helper: build a history with interleaved user messages and tool results.
|
||||
fn make_history(turns: &[(&str, Vec<(&str, Option<&str>)>)]) -> Vec<Item> {
|
||||
let mut items = Vec::new();
|
||||
for (user_msg, tool_results) in turns {
|
||||
items.push(Item::user_message(*user_msg));
|
||||
items.push(Item::assistant_message("ok"));
|
||||
for (i, (summary, content)) in tool_results.iter().enumerate() {
|
||||
let call_id = format!("call_{}", items.len() + i);
|
||||
items.push(Item::tool_call(&call_id, "some_tool", "{}"));
|
||||
match content {
|
||||
Some(c) => items.push(Item::tool_result_with_content(&call_id, *summary, *c)),
|
||||
None => items.push(Item::tool_result(&call_id, *summary)),
|
||||
}
|
||||
}
|
||||
}
|
||||
items
|
||||
}
|
||||
|
||||
fn measured_prefix(tokens: &[u64]) -> Vec<TokenEstimate> {
|
||||
tokens
|
||||
.iter()
|
||||
.copied()
|
||||
.map(|tokens| TokenEstimate {
|
||||
tokens,
|
||||
source: EstimateSource::Measured,
|
||||
})
|
||||
.collect()
|
||||
}
|
||||
|
||||
fn uniform_estimates(items: &[Item], item_tokens: u64) -> Vec<TokenEstimate> {
|
||||
let mut tokens = Vec::with_capacity(items.len() + 1);
|
||||
for i in 0..=items.len() {
|
||||
tokens.push(i as u64 * item_tokens);
|
||||
}
|
||||
measured_prefix(&tokens)
|
||||
}
|
||||
|
||||
fn estimates_from_item_tokens(item_tokens: &[u64]) -> Vec<TokenEstimate> {
|
||||
let mut prefix = Vec::with_capacity(item_tokens.len() + 1);
|
||||
let mut acc = 0;
|
||||
prefix.push(acc);
|
||||
for tokens in item_tokens {
|
||||
acc += tokens;
|
||||
prefix.push(acc);
|
||||
}
|
||||
measured_prefix(&prefix)
|
||||
}
|
||||
|
||||
fn no_data_estimates(items: &[Item]) -> Vec<TokenEstimate> {
|
||||
(0..=items.len())
|
||||
.map(|i| TokenEstimate {
|
||||
tokens: i as u64,
|
||||
source: if i == 0 {
|
||||
EstimateSource::Measured
|
||||
} else {
|
||||
EstimateSource::NoData
|
||||
},
|
||||
})
|
||||
.collect()
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn no_candidates_when_estimate_has_no_data() {
|
||||
let items = make_history(&[("turn1", vec![("summary1", Some("big content here"))])]);
|
||||
let estimates = no_data_estimates(&items);
|
||||
let (candidates, protected_start) = evaluate_candidates(&items, 10, &estimates);
|
||||
assert!(candidates.is_empty());
|
||||
assert_eq!(protected_start, None);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn no_candidates_when_history_fits_in_protected_tokens() {
|
||||
let items = make_history(&[
|
||||
("turn1", vec![("summary1", Some("big content here"))]),
|
||||
("turn2", vec![("summary2", Some("more content"))]),
|
||||
]);
|
||||
let estimates = uniform_estimates(&items, 10);
|
||||
assert!(prunable_indices(&items, 10_000, &estimates).is_empty());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn candidates_before_token_protected_suffix() {
|
||||
let big = "x".repeat(4096 * 4);
|
||||
let items = make_history(&[
|
||||
("turn1", vec![("s1", Some(&big))]),
|
||||
("turn2", vec![("s2", Some(&big))]),
|
||||
("turn3", vec![("s3", Some("keep me"))]),
|
||||
("turn4", vec![("s4", Some("keep me too"))]),
|
||||
]);
|
||||
let estimates = uniform_estimates(&items, 10);
|
||||
let candidates = prunable_indices(&items, 80, &estimates);
|
||||
assert_eq!(candidates.len(), 2);
|
||||
// suffix budget 80 tokens protects turn3+turn4 (8 items), so only s1/s2 are candidates.
|
||||
for &i in &candidates {
|
||||
if let Item::ToolResult { summary, .. } = &items[i] {
|
||||
assert!(summary == "s1" || summary == "s2");
|
||||
} else {
|
||||
panic!("non tool-result selected");
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn single_long_task_gets_candidates_without_multiple_user_turns() {
|
||||
let big = "x".repeat(4096 * 8);
|
||||
let items = make_history(&[(
|
||||
"one long task",
|
||||
vec![
|
||||
("s1", Some(&big)),
|
||||
("s2", Some(&big)),
|
||||
("s3", Some(&big)),
|
||||
("s4", Some(&big)),
|
||||
],
|
||||
)]);
|
||||
// user + assistant are cheap; every ToolCall is cheap; every ToolResult is heavy.
|
||||
let item_tokens = vec![1, 1, 1, 5_000, 1, 5_000, 1, 5_000, 1, 5_000];
|
||||
let estimates = estimates_from_item_tokens(&item_tokens);
|
||||
|
||||
let (candidates, protected_start) = evaluate_candidates(&items, 8_000, &estimates);
|
||||
|
||||
assert_eq!(protected_start, Some(7));
|
||||
assert_eq!(candidates.len(), 2);
|
||||
for &i in &candidates {
|
||||
if let Item::ToolResult { summary, .. } = &items[i] {
|
||||
assert!(summary == "s1" || summary == "s2");
|
||||
} else {
|
||||
panic!("non tool-result selected");
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn already_pruned_items_excluded_from_candidates() {
|
||||
let items = make_history(&[
|
||||
("turn1", vec![("s1", None)]), // already pruned (content=None)
|
||||
("turn2", vec![]),
|
||||
("turn3", vec![]),
|
||||
("turn4", vec![]),
|
||||
]);
|
||||
let estimates = uniform_estimates(&items, 10);
|
||||
assert!(prunable_indices(&items, 20, &estimates).is_empty());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn project_drops_content_and_counts_modifications() {
|
||||
let big = "x".repeat(64);
|
||||
let mut items = make_history(&[
|
||||
("turn1", vec![("s1", Some(&big))]),
|
||||
("turn2", vec![("s2", Some(&big))]),
|
||||
("turn3", vec![("s3", Some("keep me"))]),
|
||||
("turn4", vec![("s4", Some("keep me too"))]),
|
||||
]);
|
||||
let estimates = uniform_estimates(&items, 10);
|
||||
let candidates = prunable_indices(&items, 80, &estimates);
|
||||
let count = project(&mut items, &candidates);
|
||||
assert_eq!(count, 2);
|
||||
|
||||
for item in &items {
|
||||
if let Item::ToolResult {
|
||||
summary, content, ..
|
||||
} = item
|
||||
{
|
||||
if summary == "s1" || summary == "s2" {
|
||||
assert!(content.is_none(), "old content should be projected out");
|
||||
} else {
|
||||
assert!(content.is_some(), "protected content should remain");
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn project_skips_already_pruned_items() {
|
||||
// indices points at an item whose content is already None.
|
||||
// project() should count it as 0 modifications.
|
||||
let mut items = make_history(&[
|
||||
("turn1", vec![("s1", None)]),
|
||||
("turn2", vec![("s2", Some("hello"))]),
|
||||
]);
|
||||
// Manually target s1 even though it's already None.
|
||||
let target = items
|
||||
.iter()
|
||||
.position(|it| matches!(it, Item::ToolResult { summary, .. } if summary == "s1"))
|
||||
.unwrap();
|
||||
let count = project(&mut items, &[target]);
|
||||
assert_eq!(count, 0);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn project_is_idempotent() {
|
||||
let big = "x".repeat(64);
|
||||
let mut items = make_history(&[
|
||||
("turn1", vec![("s1", Some(&big))]),
|
||||
("turn2", vec![]),
|
||||
("turn3", vec![]),
|
||||
("turn4", vec![]),
|
||||
]);
|
||||
let estimates = uniform_estimates(&items, 10);
|
||||
let candidates = prunable_indices(&items, 20, &estimates);
|
||||
assert_eq!(project(&mut items, &candidates), 1);
|
||||
// 2 周目: 候補は一度の prunable_indices 結果を使い回しても 0 件。
|
||||
assert_eq!(project(&mut items, &candidates), 0);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn evaluate_candidates_returns_protected_start_index() {
|
||||
let big = "x".repeat(64);
|
||||
let items = make_history(&[
|
||||
("turn1", vec![("s1", Some(&big))]),
|
||||
("turn2", vec![("s2", Some(&big))]),
|
||||
("turn3", vec![("s3", Some("keep"))]),
|
||||
("turn4", vec![("s4", Some("keep too"))]),
|
||||
]);
|
||||
let estimates = uniform_estimates(&items, 10);
|
||||
let (candidates, protected_start) = evaluate_candidates(&items, 80, &estimates);
|
||||
assert_eq!(candidates.len(), 2);
|
||||
// protected_tokens=80 → protected suffix is turn3+turn4, starting at index 8.
|
||||
assert_eq!(protected_start, Some(8));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn evaluate_candidates_reports_zero_start_when_everything_is_protected() {
|
||||
let items = make_history(&[("only", vec![("s", Some("x"))])]);
|
||||
let estimates = uniform_estimates(&items, 10);
|
||||
let (candidates, protected_start) = evaluate_candidates(&items, 10_000, &estimates);
|
||||
assert!(candidates.is_empty());
|
||||
assert_eq!(protected_start, Some(0));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn zero_protected_tokens_allows_all_tool_results_as_candidates() {
|
||||
let big = "x".repeat(64);
|
||||
let items = make_history(&[("turn1", vec![("s1", Some(&big)), ("s2", Some(&big))])]);
|
||||
let estimates = uniform_estimates(&items, 10);
|
||||
let (candidates, protected_start) = evaluate_candidates(&items, 0, &estimates);
|
||||
assert_eq!(protected_start, Some(items.len()));
|
||||
assert_eq!(candidates.len(), 2);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn malformed_estimate_vector_is_treated_as_no_boundary() {
|
||||
let items = make_history(&[("turn1", vec![("s1", Some("x"))])]);
|
||||
let (candidates, protected_start) = evaluate_candidates(&items, 10, &[]);
|
||||
assert!(candidates.is_empty());
|
||||
assert_eq!(protected_start, None);
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,60 @@
|
||||
//! Engine State
|
||||
//!
|
||||
//! State marker types for cache protection using the Type-state pattern.
|
||||
//! Engine has state transitions from `Mutable` → `Locked`.
|
||||
|
||||
/// Marker trait representing Engine state
|
||||
///
|
||||
/// This trait is sealed and cannot be implemented externally.
|
||||
pub trait EngineState: private::Sealed + Send + Sync + 'static {}
|
||||
|
||||
mod private {
|
||||
pub trait Sealed {}
|
||||
}
|
||||
|
||||
/// Mutable state (editable)
|
||||
///
|
||||
/// In this state, the following operations are available:
|
||||
/// - Setting/changing system prompt
|
||||
/// - Editing message history (add, delete, clear)
|
||||
/// - Registering tools and hooks
|
||||
///
|
||||
/// Can transition to [`Locked`] state via `Engine::lock()`.
|
||||
///
|
||||
/// # Examples
|
||||
///
|
||||
/// ```ignore
|
||||
/// use llm_engine::Engine;
|
||||
///
|
||||
/// let mut engine = Engine::new(client)
|
||||
/// .system_prompt("You are helpful.");
|
||||
///
|
||||
/// // History can be edited
|
||||
/// engine.push_message(Message::user("Hello"));
|
||||
/// engine.clear_history();
|
||||
///
|
||||
/// // Lock to protected state
|
||||
/// let locked = engine.lock();
|
||||
/// ```
|
||||
#[derive(Debug, Clone, Copy, Default)]
|
||||
pub struct Mutable;
|
||||
|
||||
impl private::Sealed for Mutable {}
|
||||
impl EngineState for Mutable {}
|
||||
|
||||
/// Cache locked state (cache protected)
|
||||
///
|
||||
/// In this state, the following restrictions apply:
|
||||
/// - System prompt cannot be changed
|
||||
/// - Existing message history cannot be modified (only appending to the end)
|
||||
///
|
||||
/// To ensure LLM API KV cache hits,
|
||||
/// using this state during execution is recommended.
|
||||
///
|
||||
/// Can return to [`Mutable`] state via `Engine::unlock()`,
|
||||
/// but note that cache protection will be released.
|
||||
#[derive(Debug, Clone, Copy, Default)]
|
||||
pub struct Locked;
|
||||
|
||||
impl private::Sealed for Locked {}
|
||||
impl EngineState for Locked {}
|
||||
@@ -0,0 +1,5 @@
|
||||
//! Timeline層のイベント型
|
||||
//!
|
||||
//! llm_client層のイベント型をそのまま使用する。
|
||||
|
||||
pub use crate::llm_client::event::*;
|
||||
@@ -0,0 +1,50 @@
|
||||
//! Timeline層
|
||||
//!
|
||||
//! LLMからのイベントストリームを受信し、登録されたHandlerにディスパッチします。
|
||||
//!
|
||||
//! # 主要コンポーネント
|
||||
//!
|
||||
//! - [`Timeline`] - イベントストリームの管理とディスパッチ
|
||||
//! - [`Handler`] - イベントを処理するトレイト
|
||||
//! - [`TextBlockCollector`] - テキストブロックを収集するHandler
|
||||
//! - [`ThinkingBlockCollector`] - reasoning material を Thinking block から収集するHandler
|
||||
//! - [`ToolCallCollector`] - ツール呼び出しを収集するHandler
|
||||
|
||||
pub mod event;
|
||||
mod text_block_collector;
|
||||
mod thinking_block_collector;
|
||||
mod timeline;
|
||||
mod tool_call_collector;
|
||||
|
||||
// 公開API
|
||||
pub use event::*;
|
||||
pub use text_block_collector::TextBlockCollector;
|
||||
pub use thinking_block_collector::ThinkingBlockCollector;
|
||||
pub use timeline::Timeline;
|
||||
pub use tool_call_collector::ToolCallCollector;
|
||||
|
||||
// 型定義からのre-export
|
||||
pub use crate::handler::{
|
||||
// Meta Kinds
|
||||
ErrorKind,
|
||||
// Core traits
|
||||
Handler,
|
||||
Kind,
|
||||
PingKind,
|
||||
StatusKind,
|
||||
// Block Events
|
||||
TextBlockEvent,
|
||||
// Block Kinds
|
||||
TextBlockKind,
|
||||
TextBlockStart,
|
||||
TextBlockStop,
|
||||
ThinkingBlockEvent,
|
||||
ThinkingBlockKind,
|
||||
ThinkingBlockStart,
|
||||
ThinkingBlockStop,
|
||||
ToolUseBlockEvent,
|
||||
ToolUseBlockKind,
|
||||
ToolUseBlockStart,
|
||||
ToolUseBlockStop,
|
||||
UsageKind,
|
||||
};
|
||||
@@ -0,0 +1,131 @@
|
||||
//! TextBlockCollector - テキストブロック収集用ハンドラ
|
||||
//!
|
||||
//! TimelineのTextBlockHandler として登録され、
|
||||
//! ストリーム中のテキストブロックを収集する。
|
||||
|
||||
use crate::handler::{Handler, TextBlockEvent, TextBlockKind};
|
||||
use std::sync::{Arc, Mutex};
|
||||
|
||||
/// TextBlockから収集したテキスト情報を保持
|
||||
#[derive(Debug, Default)]
|
||||
pub struct TextCollectorState {
|
||||
/// 蓄積中のテキスト
|
||||
buffer: String,
|
||||
}
|
||||
|
||||
/// TextBlockCollector - テキストブロックハンドラ
|
||||
///
|
||||
/// Timelineに登録してTextBlockイベントを受信し、
|
||||
/// 完了したテキストブロックを収集する。
|
||||
#[derive(Clone)]
|
||||
pub struct TextBlockCollector {
|
||||
/// 収集されたテキストブロック
|
||||
collected: Arc<Mutex<Vec<String>>>,
|
||||
}
|
||||
|
||||
impl TextBlockCollector {
|
||||
/// 新しいTextBlockCollectorを作成
|
||||
pub fn new() -> Self {
|
||||
Self {
|
||||
collected: Arc::new(Mutex::new(Vec::new())),
|
||||
}
|
||||
}
|
||||
|
||||
/// 収集されたテキストを取得してクリア
|
||||
pub fn take_collected(&self) -> Vec<String> {
|
||||
let mut guard = self.collected.lock().unwrap();
|
||||
std::mem::take(&mut *guard)
|
||||
}
|
||||
|
||||
/// 収集されたテキストの参照を取得
|
||||
pub fn collected(&self) -> Vec<String> {
|
||||
self.collected.lock().unwrap().clone()
|
||||
}
|
||||
|
||||
/// 収集されたテキストがあるかどうか
|
||||
pub fn has_content(&self) -> bool {
|
||||
!self.collected.lock().unwrap().is_empty()
|
||||
}
|
||||
|
||||
/// 収集をクリア
|
||||
pub fn clear(&self) {
|
||||
self.collected.lock().unwrap().clear();
|
||||
}
|
||||
}
|
||||
|
||||
impl Default for TextBlockCollector {
|
||||
fn default() -> Self {
|
||||
Self::new()
|
||||
}
|
||||
}
|
||||
|
||||
impl Handler<TextBlockKind> for TextBlockCollector {
|
||||
type Scope = TextCollectorState;
|
||||
|
||||
fn on_event(&mut self, scope: &mut Self::Scope, event: &TextBlockEvent) {
|
||||
match event {
|
||||
TextBlockEvent::Start(_) => {
|
||||
scope.buffer.clear();
|
||||
}
|
||||
TextBlockEvent::Delta(text) => {
|
||||
scope.buffer.push_str(text);
|
||||
}
|
||||
TextBlockEvent::Stop(_) => {
|
||||
// ブロック完了時にテキストを確定
|
||||
if !scope.buffer.is_empty() {
|
||||
let text = std::mem::take(&mut scope.buffer);
|
||||
self.collected.lock().unwrap().push(text);
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
use crate::timeline::Timeline;
|
||||
use crate::timeline::event::Event;
|
||||
|
||||
/// TextBlockCollectorが単一のテキストブロックを正しく収集することを確認
|
||||
#[test]
|
||||
fn test_collect_single_text_block() {
|
||||
let collector = TextBlockCollector::new();
|
||||
let mut timeline = Timeline::new();
|
||||
timeline.on_text_block(collector.clone());
|
||||
|
||||
// テキストブロックのイベントシーケンスをディスパッチ
|
||||
timeline.dispatch(&Event::text_block_start(0));
|
||||
timeline.dispatch(&Event::text_delta(0, "Hello, "));
|
||||
timeline.dispatch(&Event::text_delta(0, "World!"));
|
||||
timeline.dispatch(&Event::text_block_stop(0, None));
|
||||
|
||||
// 収集されたテキストを確認
|
||||
let texts = collector.take_collected();
|
||||
assert_eq!(texts.len(), 1);
|
||||
assert_eq!(texts[0], "Hello, World!");
|
||||
}
|
||||
|
||||
/// TextBlockCollectorが複数のテキストブロックを正しく収集することを確認
|
||||
#[test]
|
||||
fn test_collect_multiple_text_blocks() {
|
||||
let collector = TextBlockCollector::new();
|
||||
let mut timeline = Timeline::new();
|
||||
timeline.on_text_block(collector.clone());
|
||||
|
||||
// 1つ目のテキストブロック
|
||||
timeline.dispatch(&Event::text_block_start(0));
|
||||
timeline.dispatch(&Event::text_delta(0, "First"));
|
||||
timeline.dispatch(&Event::text_block_stop(0, None));
|
||||
|
||||
// 2つ目のテキストブロック
|
||||
timeline.dispatch(&Event::text_block_start(1));
|
||||
timeline.dispatch(&Event::text_delta(1, "Second"));
|
||||
timeline.dispatch(&Event::text_block_stop(1, None));
|
||||
|
||||
let texts = collector.take_collected();
|
||||
assert_eq!(texts.len(), 2);
|
||||
assert_eq!(texts[0], "First");
|
||||
assert_eq!(texts[1], "Second");
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,138 @@
|
||||
//! `ThinkingBlockCollector` - reasoning material from Thinking block lifecycle.
|
||||
//!
|
||||
//! Scheme implementations emit Thinking BlockStart/Delta/Stop events for live
|
||||
//! streaming. A Thinking block stop with `ReasoningBlockData` is also the single
|
||||
//! authoritative persistence signal for `Item::Reasoning` round-trip material.
|
||||
|
||||
use std::sync::{Arc, Mutex};
|
||||
|
||||
use crate::handler::{Handler, ThinkingBlockEvent, ThinkingBlockKind};
|
||||
use crate::llm_client::event::ReasoningBlockData;
|
||||
|
||||
/// Reasoning material collected from completed Thinking blocks.
|
||||
#[derive(Clone, Default)]
|
||||
pub struct ThinkingBlockCollector {
|
||||
collected: Arc<Mutex<Vec<ReasoningBlockData>>>,
|
||||
}
|
||||
|
||||
impl ThinkingBlockCollector {
|
||||
pub fn new() -> Self {
|
||||
Self::default()
|
||||
}
|
||||
|
||||
/// 収集済み item を取り出してクリア
|
||||
pub fn take_collected(&self) -> Vec<ReasoningBlockData> {
|
||||
let mut guard = self.collected.lock().unwrap();
|
||||
std::mem::take(&mut *guard)
|
||||
}
|
||||
|
||||
/// 収集をクリア
|
||||
pub fn clear(&self) {
|
||||
self.collected.lock().unwrap().clear();
|
||||
}
|
||||
}
|
||||
|
||||
impl Handler<ThinkingBlockKind> for ThinkingBlockCollector {
|
||||
type Scope = String;
|
||||
|
||||
fn on_event(&mut self, scope: &mut Self::Scope, event: &ThinkingBlockEvent) {
|
||||
match event {
|
||||
ThinkingBlockEvent::Start(_) => scope.clear(),
|
||||
ThinkingBlockEvent::Delta(text) => scope.push_str(text),
|
||||
ThinkingBlockEvent::Stop(stop) => {
|
||||
if let Some(mut reasoning) = stop.reasoning.clone() {
|
||||
if reasoning.text.is_none() {
|
||||
reasoning.text = Some(scope.clone());
|
||||
}
|
||||
self.collected.lock().unwrap().push(reasoning);
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
use crate::llm_client::event::{
|
||||
BlockMetadata, BlockStart, BlockStop, BlockType, DeltaContent, Event, ReasoningBlockData,
|
||||
};
|
||||
use crate::timeline::Timeline;
|
||||
|
||||
#[test]
|
||||
fn collects_in_order_from_thinking_block_stops() {
|
||||
let collector = ThinkingBlockCollector::new();
|
||||
let mut timeline = Timeline::new();
|
||||
timeline.on_thinking_block(collector.clone());
|
||||
|
||||
timeline.dispatch(&Event::BlockStart(BlockStart {
|
||||
index: 0,
|
||||
block_type: BlockType::Thinking,
|
||||
metadata: BlockMetadata::Thinking,
|
||||
}));
|
||||
timeline.dispatch(&Event::BlockDelta(crate::llm_client::event::BlockDelta {
|
||||
index: 0,
|
||||
delta: DeltaContent::Thinking("first".into()),
|
||||
}));
|
||||
timeline.dispatch(&Event::BlockStop(BlockStop {
|
||||
index: 0,
|
||||
block_type: BlockType::Thinking,
|
||||
stop_reason: None,
|
||||
reasoning: Some(ReasoningBlockData {
|
||||
id: Some("r1".into()),
|
||||
signature: Some("sig1".into()),
|
||||
..Default::default()
|
||||
}),
|
||||
}));
|
||||
|
||||
timeline.dispatch(&Event::BlockStart(BlockStart {
|
||||
index: 1,
|
||||
block_type: BlockType::Thinking,
|
||||
metadata: BlockMetadata::Thinking,
|
||||
}));
|
||||
timeline.dispatch(&Event::BlockStop(BlockStop {
|
||||
index: 1,
|
||||
block_type: BlockType::Thinking,
|
||||
stop_reason: None,
|
||||
reasoning: Some(ReasoningBlockData {
|
||||
id: Some("r2".into()),
|
||||
text: Some("second".into()),
|
||||
..Default::default()
|
||||
}),
|
||||
}));
|
||||
|
||||
let items = collector.take_collected();
|
||||
assert_eq!(items.len(), 2);
|
||||
assert_eq!(items[0].text.as_deref(), Some("first"));
|
||||
assert_eq!(items[0].signature.as_deref(), Some("sig1"));
|
||||
assert_eq!(items[1].text.as_deref(), Some("second"));
|
||||
|
||||
// take は drain なので 2 度目は空
|
||||
assert!(collector.take_collected().is_empty());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn ignores_streaming_only_thinking_blocks() {
|
||||
let collector = ThinkingBlockCollector::new();
|
||||
let mut timeline = Timeline::new();
|
||||
timeline.on_thinking_block(collector.clone());
|
||||
|
||||
timeline.dispatch(&Event::BlockStart(BlockStart {
|
||||
index: 0,
|
||||
block_type: BlockType::Thinking,
|
||||
metadata: BlockMetadata::Thinking,
|
||||
}));
|
||||
timeline.dispatch(&Event::BlockDelta(crate::llm_client::event::BlockDelta {
|
||||
index: 0,
|
||||
delta: DeltaContent::Thinking("live only".into()),
|
||||
}));
|
||||
timeline.dispatch(&Event::BlockStop(BlockStop {
|
||||
index: 0,
|
||||
block_type: BlockType::Thinking,
|
||||
stop_reason: None,
|
||||
reasoning: None,
|
||||
}));
|
||||
|
||||
assert!(collector.take_collected().is_empty());
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,764 @@
|
||||
//! Timeline層
|
||||
//!
|
||||
//! LLMからのイベントストリームを受信し、登録されたHandlerにディスパッチします。
|
||||
//! 通常はEngine経由で使用しますが、直接使用することも可能です。
|
||||
|
||||
use std::{collections::HashMap, marker::PhantomData};
|
||||
|
||||
use super::event::*;
|
||||
use crate::handler::*;
|
||||
|
||||
// =============================================================================
|
||||
// Helpers
|
||||
// =============================================================================
|
||||
|
||||
/// 1リクエスト内で受信した複数 UsageEvent をマージする。
|
||||
/// 各フィールドについて新しい値が `Some` ならそれで上書き。
|
||||
/// プロバイダによっては input/cache 系を最初の event だけに載せ、
|
||||
/// output_tokens を後続 event で更新するため、最後の値だけを取るのではなく
|
||||
/// フィールド単位で latest-non-None を取る。
|
||||
fn merge_usage(acc: &mut UsageEvent, new: &UsageEvent) {
|
||||
if new.input_tokens.is_some() {
|
||||
acc.input_tokens = new.input_tokens;
|
||||
}
|
||||
if new.output_tokens.is_some() {
|
||||
acc.output_tokens = new.output_tokens;
|
||||
}
|
||||
if new.total_tokens.is_some() {
|
||||
acc.total_tokens = new.total_tokens;
|
||||
}
|
||||
if new.cache_read_input_tokens.is_some() {
|
||||
acc.cache_read_input_tokens = new.cache_read_input_tokens;
|
||||
}
|
||||
if new.cache_creation_input_tokens.is_some() {
|
||||
acc.cache_creation_input_tokens = new.cache_creation_input_tokens;
|
||||
}
|
||||
}
|
||||
|
||||
// =============================================================================
|
||||
// Type-erased Handler
|
||||
// =============================================================================
|
||||
|
||||
/// 型消去された`Handler` trait
|
||||
///
|
||||
/// 各Handlerは独自のScope型を持つため、Timelineで保持するには型消去が必要です。
|
||||
/// 通常は直接使用せず、`Timeline::on_text_block()`などのメソッド経由で
|
||||
/// 自動的にラップされます。
|
||||
pub trait ErasedHandler<K: Kind>: Send + Sync {
|
||||
/// イベントをディスパッチ
|
||||
fn dispatch(&mut self, event: &K::Event);
|
||||
/// スコープを開始
|
||||
fn start_scope(&mut self);
|
||||
}
|
||||
|
||||
/// `Handler<K>`を`ErasedHandler<K>`として扱うためのラッパー
|
||||
pub struct HandlerWrapper<H, K>
|
||||
where
|
||||
H: Handler<K>,
|
||||
K: Kind,
|
||||
{
|
||||
handler: H,
|
||||
scope: Option<H::Scope>,
|
||||
// fn() -> K は常にSend+Syncなので、Kの制約に関係なくSendを満たせる
|
||||
_kind: PhantomData<fn() -> K>,
|
||||
}
|
||||
|
||||
impl<H, K> HandlerWrapper<H, K>
|
||||
where
|
||||
H: Handler<K>,
|
||||
K: Kind,
|
||||
{
|
||||
pub fn new(handler: H) -> Self {
|
||||
Self {
|
||||
handler,
|
||||
scope: None,
|
||||
_kind: PhantomData,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
impl<H, K> ErasedHandler<K> for HandlerWrapper<H, K>
|
||||
where
|
||||
H: Handler<K> + Send + Sync,
|
||||
K: Kind,
|
||||
H::Scope: Send + Sync,
|
||||
{
|
||||
fn dispatch(&mut self, event: &K::Event) {
|
||||
if let Some(scope) = &mut self.scope {
|
||||
self.handler.on_event(scope, event);
|
||||
}
|
||||
}
|
||||
|
||||
fn start_scope(&mut self) {
|
||||
self.scope = Some(H::Scope::default());
|
||||
}
|
||||
}
|
||||
|
||||
// =============================================================================
|
||||
// Block Handler Registry
|
||||
// =============================================================================
|
||||
|
||||
/// ブロックハンドラーの型消去trait
|
||||
trait ErasedBlockHandler: Send + Sync {
|
||||
fn dispatch_start(&mut self, start: &BlockStart);
|
||||
fn dispatch_delta(&mut self, delta: &BlockDelta);
|
||||
fn dispatch_stop(&mut self, stop: &BlockStop);
|
||||
fn dispatch_abort(&mut self, abort: &BlockAbort);
|
||||
fn start_scope(&mut self);
|
||||
fn end_scope(&mut self);
|
||||
/// スコープがアクティブかどうか
|
||||
fn has_scope(&self) -> bool;
|
||||
}
|
||||
|
||||
/// TextBlockKind用のラッパー
|
||||
struct TextBlockHandlerWrapper<H>
|
||||
where
|
||||
H: Handler<TextBlockKind>,
|
||||
{
|
||||
handler: H,
|
||||
scope: Option<H::Scope>,
|
||||
}
|
||||
|
||||
impl<H> TextBlockHandlerWrapper<H>
|
||||
where
|
||||
H: Handler<TextBlockKind>,
|
||||
{
|
||||
fn new(handler: H) -> Self {
|
||||
Self {
|
||||
handler,
|
||||
scope: None,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
impl<H> ErasedBlockHandler for TextBlockHandlerWrapper<H>
|
||||
where
|
||||
H: Handler<TextBlockKind> + Send + Sync,
|
||||
H::Scope: Send + Sync,
|
||||
{
|
||||
fn dispatch_start(&mut self, start: &BlockStart) {
|
||||
if let Some(scope) = &mut self.scope {
|
||||
self.handler.on_event(
|
||||
scope,
|
||||
&TextBlockEvent::Start(TextBlockStart { index: start.index }),
|
||||
);
|
||||
}
|
||||
}
|
||||
|
||||
fn dispatch_delta(&mut self, delta: &BlockDelta) {
|
||||
if let Some(scope) = &mut self.scope {
|
||||
if let DeltaContent::Text(text) = &delta.delta {
|
||||
self.handler
|
||||
.on_event(scope, &TextBlockEvent::Delta(text.clone()));
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
fn dispatch_stop(&mut self, stop: &BlockStop) {
|
||||
if let Some(scope) = &mut self.scope {
|
||||
self.handler.on_event(
|
||||
scope,
|
||||
&TextBlockEvent::Stop(TextBlockStop {
|
||||
index: stop.index,
|
||||
stop_reason: stop.stop_reason.clone(),
|
||||
}),
|
||||
);
|
||||
}
|
||||
}
|
||||
|
||||
fn dispatch_abort(&mut self, _abort: &BlockAbort) {
|
||||
// TextBlockはabortを特別扱いしない（スコープ終了のみ）
|
||||
}
|
||||
|
||||
fn start_scope(&mut self) {
|
||||
self.scope = Some(H::Scope::default());
|
||||
}
|
||||
|
||||
fn end_scope(&mut self) {
|
||||
self.scope = None;
|
||||
}
|
||||
|
||||
fn has_scope(&self) -> bool {
|
||||
self.scope.is_some()
|
||||
}
|
||||
}
|
||||
|
||||
/// ThinkingBlockKind用のラッパー
|
||||
struct ThinkingBlockHandlerWrapper<H>
|
||||
where
|
||||
H: Handler<ThinkingBlockKind>,
|
||||
{
|
||||
handler: H,
|
||||
scopes: HashMap<usize, H::Scope>,
|
||||
}
|
||||
|
||||
impl<H> ThinkingBlockHandlerWrapper<H>
|
||||
where
|
||||
H: Handler<ThinkingBlockKind>,
|
||||
{
|
||||
fn new(handler: H) -> Self {
|
||||
Self {
|
||||
handler,
|
||||
scopes: HashMap::new(),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
impl<H> ErasedBlockHandler for ThinkingBlockHandlerWrapper<H>
|
||||
where
|
||||
H: Handler<ThinkingBlockKind> + Send + Sync,
|
||||
H::Scope: Send + Sync,
|
||||
{
|
||||
fn dispatch_start(&mut self, start: &BlockStart) {
|
||||
let scope = self.scopes.entry(start.index).or_default();
|
||||
self.handler.on_event(
|
||||
scope,
|
||||
&ThinkingBlockEvent::Start(ThinkingBlockStart { index: start.index }),
|
||||
);
|
||||
}
|
||||
|
||||
fn dispatch_delta(&mut self, delta: &BlockDelta) {
|
||||
if let DeltaContent::Thinking(text) = &delta.delta {
|
||||
let scope = self.scopes.entry(delta.index).or_default();
|
||||
self.handler
|
||||
.on_event(scope, &ThinkingBlockEvent::Delta(text.clone()));
|
||||
}
|
||||
}
|
||||
|
||||
fn dispatch_stop(&mut self, stop: &BlockStop) {
|
||||
if let Some(mut scope) = self.scopes.remove(&stop.index) {
|
||||
self.handler.on_event(
|
||||
&mut scope,
|
||||
&ThinkingBlockEvent::Stop(ThinkingBlockStop {
|
||||
index: stop.index,
|
||||
reasoning: stop.reasoning.clone(),
|
||||
}),
|
||||
);
|
||||
}
|
||||
}
|
||||
|
||||
fn dispatch_abort(&mut self, _abort: &BlockAbort) {
|
||||
self.scopes.clear();
|
||||
}
|
||||
|
||||
fn start_scope(&mut self) {}
|
||||
|
||||
fn end_scope(&mut self) {}
|
||||
|
||||
fn has_scope(&self) -> bool {
|
||||
!self.scopes.is_empty()
|
||||
}
|
||||
}
|
||||
|
||||
/// ToolUseBlockKind用のラッパー
|
||||
struct ToolUseBlockHandlerWrapper<H>
|
||||
where
|
||||
H: Handler<ToolUseBlockKind>,
|
||||
{
|
||||
handler: H,
|
||||
scope: Option<H::Scope>,
|
||||
current_tool: Option<(String, String)>, // (id, name)
|
||||
}
|
||||
|
||||
impl<H> ToolUseBlockHandlerWrapper<H>
|
||||
where
|
||||
H: Handler<ToolUseBlockKind>,
|
||||
{
|
||||
fn new(handler: H) -> Self {
|
||||
Self {
|
||||
handler,
|
||||
scope: None,
|
||||
current_tool: None,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
impl<H> ErasedBlockHandler for ToolUseBlockHandlerWrapper<H>
|
||||
where
|
||||
H: Handler<ToolUseBlockKind> + Send + Sync,
|
||||
H::Scope: Send + Sync,
|
||||
{
|
||||
fn dispatch_start(&mut self, start: &BlockStart) {
|
||||
if let Some(scope) = &mut self.scope {
|
||||
if let BlockMetadata::ToolUse { id, name } = &start.metadata {
|
||||
self.current_tool = Some((id.clone(), name.clone()));
|
||||
self.handler.on_event(
|
||||
scope,
|
||||
&ToolUseBlockEvent::Start(ToolUseBlockStart {
|
||||
index: start.index,
|
||||
id: id.clone(),
|
||||
name: name.clone(),
|
||||
}),
|
||||
);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
fn dispatch_delta(&mut self, delta: &BlockDelta) {
|
||||
if let Some(scope) = &mut self.scope {
|
||||
if let DeltaContent::InputJson(json) = &delta.delta {
|
||||
self.handler
|
||||
.on_event(scope, &ToolUseBlockEvent::InputJsonDelta(json.clone()));
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
fn dispatch_stop(&mut self, stop: &BlockStop) {
|
||||
if let Some(scope) = &mut self.scope {
|
||||
if let Some((id, name)) = self.current_tool.take() {
|
||||
self.handler.on_event(
|
||||
scope,
|
||||
&ToolUseBlockEvent::Stop(ToolUseBlockStop {
|
||||
index: stop.index,
|
||||
id,
|
||||
name,
|
||||
}),
|
||||
);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
fn dispatch_abort(&mut self, _abort: &BlockAbort) {
|
||||
self.current_tool = None;
|
||||
}
|
||||
|
||||
fn start_scope(&mut self) {
|
||||
self.scope = Some(H::Scope::default());
|
||||
}
|
||||
|
||||
fn end_scope(&mut self) {
|
||||
self.scope = None;
|
||||
self.current_tool = None;
|
||||
}
|
||||
|
||||
fn has_scope(&self) -> bool {
|
||||
self.scope.is_some()
|
||||
}
|
||||
}
|
||||
|
||||
// =============================================================================
|
||||
// Timeline
|
||||
// =============================================================================
|
||||
|
||||
/// イベントストリームの管理とハンドラへのディスパッチ
|
||||
///
|
||||
/// LLMからのイベントを受信し、登録されたハンドラに振り分けます。
|
||||
/// ブロック系イベントはスコープ管理付きで処理されます。
|
||||
///
|
||||
/// # Examples
|
||||
///
|
||||
/// ```ignore
|
||||
/// use llm_engine::{Timeline, Handler, TextBlockKind, TextBlockEvent};
|
||||
///
|
||||
/// struct MyHandler;
|
||||
/// impl Handler<TextBlockKind> for MyHandler {
|
||||
/// type Scope = String;
|
||||
/// fn on_event(&mut self, buffer: &mut String, event: &TextBlockEvent) {
|
||||
/// if let TextBlockEvent::Delta(text) = event {
|
||||
/// buffer.push_str(text);
|
||||
/// }
|
||||
/// }
|
||||
/// }
|
||||
///
|
||||
/// let mut timeline = Timeline::new();
|
||||
/// timeline.on_text_block(MyHandler);
|
||||
/// ```
|
||||
///
|
||||
/// # サポートするイベント種別
|
||||
///
|
||||
/// - **メタ系**: Usage, Ping, Status, Error
|
||||
/// - **ブロック系**: TextBlock, ThinkingBlock, ToolUseBlock
|
||||
pub struct Timeline {
|
||||
// Meta系ハンドラー
|
||||
usage_handlers: Vec<Box<dyn ErasedHandler<UsageKind>>>,
|
||||
ping_handlers: Vec<Box<dyn ErasedHandler<PingKind>>>,
|
||||
status_handlers: Vec<Box<dyn ErasedHandler<StatusKind>>>,
|
||||
error_handlers: Vec<Box<dyn ErasedHandler<ErrorKind>>>,
|
||||
// Block系ハンドラー（BlockTypeごとにグループ化）
|
||||
text_block_handlers: Vec<Box<dyn ErasedBlockHandler>>,
|
||||
thinking_block_handlers: Vec<Box<dyn ErasedBlockHandler>>,
|
||||
tool_use_block_handlers: Vec<Box<dyn ErasedBlockHandler>>,
|
||||
|
||||
// 現在アクティブなブロック
|
||||
current_block: Option<BlockType>,
|
||||
|
||||
// 1リクエスト内で受信した Usage event の集約バッファ。
|
||||
// Anthropic は message_start と message_delta、Gemini は各チャンクと、
|
||||
// 多くのプロバイダが複数 Usage を発行するため、リクエスト境界で
|
||||
// 1度だけ発火するためにここでマージする。flush_usage() で発火する。
|
||||
pending_usage: Option<UsageEvent>,
|
||||
}
|
||||
|
||||
impl Default for Timeline {
|
||||
fn default() -> Self {
|
||||
Self::new()
|
||||
}
|
||||
}
|
||||
|
||||
impl Timeline {
|
||||
pub fn new() -> Self {
|
||||
Self {
|
||||
usage_handlers: Vec::new(),
|
||||
ping_handlers: Vec::new(),
|
||||
status_handlers: Vec::new(),
|
||||
error_handlers: Vec::new(),
|
||||
text_block_handlers: Vec::new(),
|
||||
thinking_block_handlers: Vec::new(),
|
||||
tool_use_block_handlers: Vec::new(),
|
||||
current_block: None,
|
||||
pending_usage: None,
|
||||
}
|
||||
}
|
||||
|
||||
// =========================================================================
|
||||
// Handler Registration
|
||||
// =========================================================================
|
||||
|
||||
/// UsageKind用のHandlerを登録
|
||||
pub fn on_usage<H>(&mut self, handler: H) -> &mut Self
|
||||
where
|
||||
H: Handler<UsageKind> + Send + Sync + 'static,
|
||||
H::Scope: Send + Sync,
|
||||
{
|
||||
// Meta系はデフォルトでスコープを開始しておく
|
||||
let mut wrapper = HandlerWrapper::new(handler);
|
||||
wrapper.start_scope();
|
||||
self.usage_handlers.push(Box::new(wrapper));
|
||||
self
|
||||
}
|
||||
|
||||
/// PingKind用のHandlerを登録
|
||||
pub fn on_ping<H>(&mut self, handler: H) -> &mut Self
|
||||
where
|
||||
H: Handler<PingKind> + Send + Sync + 'static,
|
||||
H::Scope: Send + Sync,
|
||||
{
|
||||
let mut wrapper = HandlerWrapper::new(handler);
|
||||
wrapper.start_scope();
|
||||
self.ping_handlers.push(Box::new(wrapper));
|
||||
self
|
||||
}
|
||||
|
||||
/// StatusKind用のHandlerを登録
|
||||
pub fn on_status<H>(&mut self, handler: H) -> &mut Self
|
||||
where
|
||||
H: Handler<StatusKind> + Send + Sync + 'static,
|
||||
H::Scope: Send + Sync,
|
||||
{
|
||||
let mut wrapper = HandlerWrapper::new(handler);
|
||||
wrapper.start_scope();
|
||||
self.status_handlers.push(Box::new(wrapper));
|
||||
self
|
||||
}
|
||||
|
||||
/// ErrorKind用のHandlerを登録
|
||||
pub fn on_error<H>(&mut self, handler: H) -> &mut Self
|
||||
where
|
||||
H: Handler<ErrorKind> + Send + Sync + 'static,
|
||||
H::Scope: Send + Sync,
|
||||
{
|
||||
let mut wrapper = HandlerWrapper::new(handler);
|
||||
wrapper.start_scope();
|
||||
self.error_handlers.push(Box::new(wrapper));
|
||||
self
|
||||
}
|
||||
|
||||
/// TextBlockKind用のHandlerを登録
|
||||
pub fn on_text_block<H>(&mut self, handler: H) -> &mut Self
|
||||
where
|
||||
H: Handler<TextBlockKind> + Send + Sync + 'static,
|
||||
H::Scope: Send + Sync,
|
||||
{
|
||||
self.text_block_handlers
|
||||
.push(Box::new(TextBlockHandlerWrapper::new(handler)));
|
||||
self
|
||||
}
|
||||
|
||||
/// ThinkingBlockKind用のHandlerを登録
|
||||
pub fn on_thinking_block<H>(&mut self, handler: H) -> &mut Self
|
||||
where
|
||||
H: Handler<ThinkingBlockKind> + Send + Sync + 'static,
|
||||
H::Scope: Send + Sync,
|
||||
{
|
||||
self.thinking_block_handlers
|
||||
.push(Box::new(ThinkingBlockHandlerWrapper::new(handler)));
|
||||
self
|
||||
}
|
||||
|
||||
/// ToolUseBlockKind用のHandlerを登録
|
||||
pub fn on_tool_use_block<H>(&mut self, handler: H) -> &mut Self
|
||||
where
|
||||
H: Handler<ToolUseBlockKind> + Send + Sync + 'static,
|
||||
H::Scope: Send + Sync,
|
||||
{
|
||||
self.tool_use_block_handlers
|
||||
.push(Box::new(ToolUseBlockHandlerWrapper::new(handler)));
|
||||
self
|
||||
}
|
||||
|
||||
// =========================================================================
|
||||
// Event Dispatch
|
||||
// =========================================================================
|
||||
|
||||
/// メインのディスパッチエントリポイント
|
||||
pub fn dispatch(&mut self, event: &Event) {
|
||||
match event {
|
||||
// Meta系: 即時ディスパッチ（登録順）
|
||||
Event::Usage(u) => self.dispatch_usage(u),
|
||||
Event::Ping(p) => self.dispatch_ping(p),
|
||||
Event::Status(s) => self.dispatch_status(s),
|
||||
Event::Error(e) => self.dispatch_error(e),
|
||||
// Observability-only event: stream trace records it before timeline dispatch.
|
||||
Event::UnhandledSse(_) => {}
|
||||
|
||||
// Block系: スコープ管理しながらディスパッチ
|
||||
Event::BlockStart(s) => self.handle_block_start(s),
|
||||
Event::BlockDelta(d) => self.handle_block_delta(d),
|
||||
Event::BlockStop(s) => self.handle_block_stop(s),
|
||||
Event::BlockAbort(a) => self.handle_block_abort(a),
|
||||
}
|
||||
}
|
||||
|
||||
/// Usage event を即時には dispatch せず、pending_usage にマージする。
|
||||
/// 1リクエスト内で複数の Usage event が来ても、ハンドラには 1 度だけ
|
||||
/// 最終値を渡したいため。flush_usage() で発火する。
|
||||
fn dispatch_usage(&mut self, event: &UsageEvent) {
|
||||
match &mut self.pending_usage {
|
||||
Some(acc) => merge_usage(acc, event),
|
||||
None => self.pending_usage = Some(event.clone()),
|
||||
}
|
||||
}
|
||||
|
||||
/// pending_usage を usage_handlers に発火し、バッファをクリアする。
|
||||
/// 1リクエスト分のストリーム終了時に1回だけ呼ぶ想定。
|
||||
/// pending_usage が空ならば何もしない。
|
||||
pub fn flush_usage(&mut self) {
|
||||
if let Some(event) = self.pending_usage.take() {
|
||||
for handler in &mut self.usage_handlers {
|
||||
handler.dispatch(&event);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
fn dispatch_ping(&mut self, event: &PingEvent) {
|
||||
for handler in &mut self.ping_handlers {
|
||||
handler.dispatch(event);
|
||||
}
|
||||
}
|
||||
|
||||
fn dispatch_status(&mut self, event: &StatusEvent) {
|
||||
for handler in &mut self.status_handlers {
|
||||
handler.dispatch(event);
|
||||
}
|
||||
}
|
||||
|
||||
fn dispatch_error(&mut self, event: &ErrorEvent) {
|
||||
for handler in &mut self.error_handlers {
|
||||
handler.dispatch(event);
|
||||
}
|
||||
}
|
||||
|
||||
fn handle_block_start(&mut self, start: &BlockStart) {
|
||||
self.current_block = Some(start.block_type);
|
||||
|
||||
let handlers = self.get_block_handlers_mut(start.block_type);
|
||||
for handler in handlers {
|
||||
handler.start_scope();
|
||||
handler.dispatch_start(start);
|
||||
}
|
||||
}
|
||||
|
||||
fn handle_block_delta(&mut self, delta: &BlockDelta) {
|
||||
let block_type = delta.delta.block_type();
|
||||
|
||||
// OpenAIなどのプロバイダはBlockStartを送らない場合があるため、
|
||||
// Deltaが来たときにスコープがなければ暗黙的に開始する
|
||||
if self.current_block.is_none() {
|
||||
self.current_block = Some(block_type);
|
||||
}
|
||||
|
||||
let handlers = self.get_block_handlers_mut(block_type);
|
||||
for handler in handlers {
|
||||
// スコープがなければ暗黙的に開始
|
||||
if !handler.has_scope() {
|
||||
handler.start_scope();
|
||||
}
|
||||
handler.dispatch_delta(delta);
|
||||
}
|
||||
}
|
||||
|
||||
fn handle_block_stop(&mut self, stop: &BlockStop) {
|
||||
let handlers = self.get_block_handlers_mut(stop.block_type);
|
||||
for handler in handlers {
|
||||
handler.dispatch_stop(stop);
|
||||
handler.end_scope();
|
||||
}
|
||||
self.current_block = None;
|
||||
}
|
||||
|
||||
fn handle_block_abort(&mut self, abort: &BlockAbort) {
|
||||
let handlers = self.get_block_handlers_mut(abort.block_type);
|
||||
for handler in handlers {
|
||||
handler.dispatch_abort(abort);
|
||||
handler.end_scope();
|
||||
}
|
||||
self.current_block = None;
|
||||
}
|
||||
|
||||
fn get_block_handlers_mut(
|
||||
&mut self,
|
||||
block_type: BlockType,
|
||||
) -> &mut Vec<Box<dyn ErasedBlockHandler>> {
|
||||
match block_type {
|
||||
BlockType::Text => &mut self.text_block_handlers,
|
||||
BlockType::Thinking => &mut self.thinking_block_handlers,
|
||||
BlockType::ToolUse => &mut self.tool_use_block_handlers,
|
||||
BlockType::ToolResult => &mut self.text_block_handlers, // ToolResultはTextとして扱う
|
||||
}
|
||||
}
|
||||
|
||||
/// 現在アクティブなブロックタイプを取得
|
||||
pub fn current_block(&self) -> Option<BlockType> {
|
||||
self.current_block
|
||||
}
|
||||
|
||||
/// 現在アクティブなブロックを中断する
|
||||
///
|
||||
/// キャンセルやエラー時に呼び出し、進行中のブロックに対して
|
||||
/// BlockAbortイベントを発火してスコープをクリーンアップする。
|
||||
pub fn abort_current_block(&mut self) {
|
||||
if let Some(block_type) = self.current_block {
|
||||
let abort = crate::timeline::event::BlockAbort {
|
||||
index: 0, // インデックスは不明なので0
|
||||
block_type,
|
||||
reason: "Cancelled".to_string(),
|
||||
};
|
||||
self.handle_block_abort(&abort);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
use std::sync::{Arc, Mutex};
|
||||
|
||||
#[test]
|
||||
fn test_timeline_creation() {
|
||||
let timeline = Timeline::new();
|
||||
assert!(timeline.current_block().is_none());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn unhandled_sse_is_ignored_by_timeline_handlers() {
|
||||
struct TestTextHandler {
|
||||
calls: Arc<Mutex<Vec<TextBlockEvent>>>,
|
||||
}
|
||||
|
||||
impl Handler<TextBlockKind> for TestTextHandler {
|
||||
type Scope = ();
|
||||
fn on_event(&mut self, _scope: &mut (), event: &TextBlockEvent) {
|
||||
self.calls.lock().unwrap().push(event.clone());
|
||||
}
|
||||
}
|
||||
|
||||
let calls = Arc::new(Mutex::new(Vec::new()));
|
||||
let mut timeline = Timeline::new();
|
||||
timeline.on_text_block(TestTextHandler {
|
||||
calls: calls.clone(),
|
||||
});
|
||||
|
||||
timeline.dispatch(&Event::UnhandledSse(UnhandledSseEvent {
|
||||
provider: "openai_responses".to_string(),
|
||||
event_type: "response.mystery".to_string(),
|
||||
data_preview: "{}".to_string(),
|
||||
data_len: 2,
|
||||
}));
|
||||
|
||||
assert!(timeline.current_block().is_none());
|
||||
assert!(calls.lock().unwrap().is_empty());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_meta_event_dispatch() {
|
||||
// シンプルなテスト用構造体
|
||||
struct TestUsageHandler {
|
||||
calls: Arc<Mutex<Vec<UsageEvent>>>,
|
||||
}
|
||||
|
||||
impl Handler<UsageKind> for TestUsageHandler {
|
||||
type Scope = ();
|
||||
fn on_event(&mut self, _scope: &mut (), event: &UsageEvent) {
|
||||
self.calls.lock().unwrap().push(event.clone());
|
||||
}
|
||||
}
|
||||
|
||||
let calls = Arc::new(Mutex::new(Vec::new()));
|
||||
let handler = TestUsageHandler {
|
||||
calls: calls.clone(),
|
||||
};
|
||||
|
||||
let mut timeline = Timeline::new();
|
||||
timeline.on_usage(handler);
|
||||
|
||||
timeline.dispatch(&Event::usage(100, 50));
|
||||
// pending_usage に積まれているだけなのでまだ未発火
|
||||
assert_eq!(calls.lock().unwrap().len(), 0);
|
||||
|
||||
// flush で 1 度だけ発火
|
||||
timeline.flush_usage();
|
||||
let recorded = calls.lock().unwrap();
|
||||
assert_eq!(recorded.len(), 1);
|
||||
assert_eq!(recorded[0].input_tokens, Some(100));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_usage_aggregation_and_flush() {
|
||||
struct TestUsageHandler {
|
||||
calls: Arc<Mutex<Vec<UsageEvent>>>,
|
||||
}
|
||||
impl Handler<UsageKind> for TestUsageHandler {
|
||||
type Scope = ();
|
||||
fn on_event(&mut self, _scope: &mut (), event: &UsageEvent) {
|
||||
self.calls.lock().unwrap().push(event.clone());
|
||||
}
|
||||
}
|
||||
|
||||
let calls = Arc::new(Mutex::new(Vec::new()));
|
||||
let mut timeline = Timeline::new();
|
||||
timeline.on_usage(TestUsageHandler {
|
||||
calls: calls.clone(),
|
||||
});
|
||||
|
||||
// Anthropic 風: message_start で input + 暫定 output
|
||||
timeline.dispatch(&Event::Usage(UsageEvent {
|
||||
input_tokens: Some(409),
|
||||
output_tokens: Some(1),
|
||||
total_tokens: Some(410),
|
||||
cache_read_input_tokens: Some(0),
|
||||
cache_creation_input_tokens: Some(0),
|
||||
}));
|
||||
// message_delta で最終 output
|
||||
timeline.dispatch(&Event::Usage(UsageEvent {
|
||||
input_tokens: Some(409),
|
||||
output_tokens: Some(71),
|
||||
total_tokens: Some(480),
|
||||
cache_read_input_tokens: Some(0),
|
||||
cache_creation_input_tokens: Some(0),
|
||||
}));
|
||||
|
||||
// 未 flush の段階では発火しない
|
||||
assert_eq!(calls.lock().unwrap().len(), 0);
|
||||
|
||||
timeline.flush_usage();
|
||||
let recorded = calls.lock().unwrap();
|
||||
assert_eq!(recorded.len(), 1);
|
||||
assert_eq!(recorded[0].input_tokens, Some(409));
|
||||
assert_eq!(recorded[0].output_tokens, Some(71));
|
||||
|
||||
// flush 後にもう一度 flush しても何も起きない
|
||||
drop(recorded);
|
||||
timeline.flush_usage();
|
||||
assert_eq!(calls.lock().unwrap().len(), 1);
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,168 @@
|
||||
//! ToolCallCollector - ツール呼び出し収集用ハンドラ
|
||||
//!
|
||||
//! TimelineのToolUseBlockHandler として登録され、
|
||||
//! ストリーム中のToolUseブロックを収集する。
|
||||
|
||||
use crate::{
|
||||
handler::{Handler, ToolUseBlockEvent, ToolUseBlockKind},
|
||||
llm_client::types::parse_tool_arguments,
|
||||
tool::ToolCall,
|
||||
};
|
||||
use std::sync::{Arc, Mutex};
|
||||
|
||||
/// ToolUseブロックから収集したツール呼び出し情報を保持
|
||||
///
|
||||
/// ToolCallCollectorのHandler実装で使用するスコープ型
|
||||
#[derive(Debug, Default)]
|
||||
pub struct CollectorState {
|
||||
/// 現在のツール呼び出し情報 (ブロック進行中)
|
||||
current_id: Option<String>,
|
||||
current_name: Option<String>,
|
||||
/// 蓄積中のJSON入力
|
||||
input_json_buffer: String,
|
||||
}
|
||||
|
||||
/// ToolCallCollector - ToolUseブロックハンドラ
|
||||
///
|
||||
/// Timelineに登録してToolUseブロックイベントを受信し、
|
||||
/// 完了したToolCallを収集する。
|
||||
#[derive(Clone)]
|
||||
pub struct ToolCallCollector {
|
||||
/// 収集されたToolCall
|
||||
collected: Arc<Mutex<Vec<ToolCall>>>,
|
||||
}
|
||||
|
||||
impl ToolCallCollector {
|
||||
/// 新しいToolCallCollectorを作成
|
||||
pub fn new() -> Self {
|
||||
Self {
|
||||
collected: Arc::new(Mutex::new(Vec::new())),
|
||||
}
|
||||
}
|
||||
|
||||
/// 収集されたToolCallを取得してクリア
|
||||
pub fn take_collected(&self) -> Vec<ToolCall> {
|
||||
let mut guard = self.collected.lock().unwrap();
|
||||
std::mem::take(&mut *guard)
|
||||
}
|
||||
|
||||
/// 収集されたToolCallの参照を取得
|
||||
pub fn collected(&self) -> Vec<ToolCall> {
|
||||
self.collected.lock().unwrap().clone()
|
||||
}
|
||||
|
||||
/// 収集されたToolCallがあるかどうか
|
||||
pub fn has_pending_calls(&self) -> bool {
|
||||
!self.collected.lock().unwrap().is_empty()
|
||||
}
|
||||
|
||||
/// 収集をクリア
|
||||
pub fn clear(&self) {
|
||||
self.collected.lock().unwrap().clear();
|
||||
}
|
||||
}
|
||||
|
||||
impl Default for ToolCallCollector {
|
||||
fn default() -> Self {
|
||||
Self::new()
|
||||
}
|
||||
}
|
||||
|
||||
impl Handler<ToolUseBlockKind> for ToolCallCollector {
|
||||
type Scope = CollectorState;
|
||||
|
||||
fn on_event(&mut self, scope: &mut Self::Scope, event: &ToolUseBlockEvent) {
|
||||
match event {
|
||||
ToolUseBlockEvent::Start(start) => {
|
||||
scope.current_id = Some(start.id.clone());
|
||||
scope.current_name = Some(start.name.clone());
|
||||
scope.input_json_buffer.clear();
|
||||
}
|
||||
ToolUseBlockEvent::InputJsonDelta(delta) => {
|
||||
scope.input_json_buffer.push_str(delta);
|
||||
}
|
||||
ToolUseBlockEvent::Stop(_stop) => {
|
||||
// ブロック完了時にToolCallを確定
|
||||
if let (Some(id), Some(name)) = (scope.current_id.take(), scope.current_name.take())
|
||||
{
|
||||
let input = parse_tool_arguments(&scope.input_json_buffer);
|
||||
|
||||
let tool_call = ToolCall { id, name, input };
|
||||
|
||||
self.collected.lock().unwrap().push(tool_call);
|
||||
}
|
||||
scope.input_json_buffer.clear();
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
use crate::timeline::Timeline;
|
||||
use crate::timeline::event::Event;
|
||||
|
||||
#[test]
|
||||
fn test_collect_single_tool_call() {
|
||||
let collector = ToolCallCollector::new();
|
||||
let mut timeline = Timeline::new();
|
||||
timeline.on_tool_use_block(collector.clone());
|
||||
|
||||
// ToolUseブロックのイベントシーケンスをディスパッチ
|
||||
timeline.dispatch(&Event::tool_use_start(0, "tool_123", "get_weather"));
|
||||
timeline.dispatch(&Event::tool_input_delta(0, r#"{"city":"#));
|
||||
timeline.dispatch(&Event::tool_input_delta(0, r#""Tokyo"}"#));
|
||||
timeline.dispatch(&Event::tool_use_stop(0));
|
||||
|
||||
// 収集されたToolCallを確認
|
||||
let calls = collector.take_collected();
|
||||
assert_eq!(calls.len(), 1);
|
||||
assert_eq!(calls[0].id, "tool_123");
|
||||
assert_eq!(calls[0].name, "get_weather");
|
||||
assert_eq!(calls[0].input["city"], "Tokyo");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_collect_empty_buffer_returns_object() {
|
||||
// 引数なしツール呼び出し: input_json_delta が一度も来ないケース
|
||||
let collector = ToolCallCollector::new();
|
||||
let mut timeline = Timeline::new();
|
||||
timeline.on_tool_use_block(collector.clone());
|
||||
|
||||
timeline.dispatch(&Event::tool_use_start(0, "tool_empty", "ListPods"));
|
||||
timeline.dispatch(&Event::tool_use_stop(0));
|
||||
|
||||
let calls = collector.take_collected();
|
||||
assert_eq!(calls.len(), 1);
|
||||
assert_eq!(calls[0].id, "tool_empty");
|
||||
assert_eq!(calls[0].name, "ListPods");
|
||||
assert!(calls[0].input.is_object());
|
||||
assert_eq!(
|
||||
calls[0].input,
|
||||
serde_json::Value::Object(serde_json::Map::new())
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_collect_multiple_tool_calls() {
|
||||
let collector = ToolCallCollector::new();
|
||||
let mut timeline = Timeline::new();
|
||||
timeline.on_tool_use_block(collector.clone());
|
||||
|
||||
// 1つ目のToolCall
|
||||
timeline.dispatch(&Event::tool_use_start(0, "call_1", "tool_a"));
|
||||
timeline.dispatch(&Event::tool_input_delta(0, r#"{"a":1}"#));
|
||||
timeline.dispatch(&Event::tool_use_stop(0));
|
||||
|
||||
// 2つ目のToolCall
|
||||
timeline.dispatch(&Event::tool_use_start(1, "call_2", "tool_b"));
|
||||
timeline.dispatch(&Event::tool_input_delta(1, r#"{"b":2}"#));
|
||||
timeline.dispatch(&Event::tool_use_stop(1));
|
||||
|
||||
let calls = collector.take_collected();
|
||||
assert_eq!(calls.len(), 2);
|
||||
assert_eq!(calls[0].name, "tool_a");
|
||||
assert_eq!(calls[1].name, "tool_b");
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,273 @@
|
||||
//! Usage 履歴ベースのトークン会計（汎用部分）。
|
||||
//!
|
||||
//! `UsageRecord` の列（プロバイダ実測値）と現在の history から、
|
||||
//! 任意の history index 時点のプロンプト全長トークン数を pure に計算する。
|
||||
//!
|
||||
//! # 方針
|
||||
//!
|
||||
//! - ローカルトークナイザは持たない。実測値があればそれを採用し、
|
||||
//! measurement 間はバイト数で按分、最新 measurement より先は byte/4
|
||||
//! fallback で外挿する
|
||||
//! - 推定の出どころは [`EstimateSource`] で呼び出し側に明示する。
|
||||
//! 課金判断には使えないが、compact / prune / memory extract trigger 等の
|
||||
//! 閾値判定には十分な精度
|
||||
//! - `records` は `history_len` 昇順を仮定する（呼び出し側がそのように積む）
|
||||
|
||||
use crate::{Item, UsageRecord};
|
||||
|
||||
/// 推定の出どころ。
|
||||
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
|
||||
pub enum EstimateSource {
|
||||
/// measurement の境界にちょうど一致（実測値そのもの）
|
||||
Measured,
|
||||
/// 連続する 2 つの measurement の間をバイト按分で計算
|
||||
Interpolated,
|
||||
/// 最後の measurement より新しい区間を byte/4 fallback で外挿
|
||||
Extrapolated,
|
||||
/// measurement が 1 件も無く、バイト数のみのフォールバック
|
||||
NoData,
|
||||
}
|
||||
|
||||
/// トークン数の推定値。
|
||||
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
|
||||
pub struct TokenEstimate {
|
||||
pub tokens: u64,
|
||||
pub source: EstimateSource,
|
||||
}
|
||||
|
||||
/// `items[..i]` までの累積バイト数（`prefix[i]`）を返す。長さは `items.len()+1`。
|
||||
pub fn prefix_bytes(items: &[Item]) -> Vec<u64> {
|
||||
let mut prefix = Vec::with_capacity(items.len() + 1);
|
||||
let mut acc: u64 = 0;
|
||||
prefix.push(0);
|
||||
for item in items {
|
||||
acc = acc.saturating_add(item_bytes(item));
|
||||
prefix.push(acc);
|
||||
}
|
||||
prefix
|
||||
}
|
||||
|
||||
/// 1 Item の大きさ。JSON シリアライズ長を使う粗い近似。
|
||||
/// トークン数との絶対変換ではなく区間の按分にしか使わないので、
|
||||
/// プロバイダごとの overhead は比率でキャンセルされる。
|
||||
pub fn item_bytes(item: &Item) -> u64 {
|
||||
serde_json::to_string(item)
|
||||
.map(|s| s.len() as u64)
|
||||
.unwrap_or(0)
|
||||
}
|
||||
|
||||
/// `history[..index]` までのトークン数を推定する。
|
||||
///
|
||||
/// `prefix` は [`prefix_bytes`] で得た `history.len() + 1` 長の累積バイト列。
|
||||
/// 呼び出し側が 1 度だけ計算して使い回すことで、線形探索や複数回の推定が
|
||||
/// O(n) シリアライズで済む（内部で毎回再計算すると O(n²) になる）。
|
||||
pub fn tokens_at(
|
||||
history: &[Item],
|
||||
records: &[UsageRecord],
|
||||
index: usize,
|
||||
prefix: &[u64],
|
||||
) -> TokenEstimate {
|
||||
debug_assert!(index <= history.len());
|
||||
debug_assert_eq!(prefix.len(), history.len() + 1);
|
||||
|
||||
if index == 0 {
|
||||
return TokenEstimate {
|
||||
tokens: 0,
|
||||
source: EstimateSource::Measured,
|
||||
};
|
||||
}
|
||||
|
||||
if records.is_empty() {
|
||||
return TokenEstimate {
|
||||
tokens: prefix[index] / 4,
|
||||
source: EstimateSource::NoData,
|
||||
};
|
||||
}
|
||||
|
||||
// exact match（rev 走査で一番新しい record を採用）
|
||||
if let Some(r) = records.iter().rev().find(|r| r.history_len == index) {
|
||||
return TokenEstimate {
|
||||
tokens: r.input_total_tokens,
|
||||
source: EstimateSource::Measured,
|
||||
};
|
||||
}
|
||||
|
||||
let lower = records.iter().rev().find(|r| r.history_len < index);
|
||||
let upper = records.iter().find(|r| r.history_len > index);
|
||||
let cap = history.len();
|
||||
|
||||
match (lower, upper) {
|
||||
(Some(lo), Some(up)) => {
|
||||
let lo_bytes = prefix[lo.history_len.min(cap)];
|
||||
let up_bytes = prefix[up.history_len.min(cap)];
|
||||
let at_bytes = prefix[index];
|
||||
let span_bytes = up_bytes.saturating_sub(lo_bytes);
|
||||
let span_tokens = up.input_total_tokens.saturating_sub(lo.input_total_tokens);
|
||||
if span_bytes == 0 || span_tokens == 0 {
|
||||
return TokenEstimate {
|
||||
tokens: lo.input_total_tokens,
|
||||
source: EstimateSource::Interpolated,
|
||||
};
|
||||
}
|
||||
let delta_bytes = at_bytes.saturating_sub(lo_bytes);
|
||||
let delta_tokens =
|
||||
(delta_bytes as u128 * span_tokens as u128 / span_bytes as u128) as u64;
|
||||
TokenEstimate {
|
||||
tokens: lo.input_total_tokens + delta_tokens,
|
||||
source: EstimateSource::Interpolated,
|
||||
}
|
||||
}
|
||||
(Some(lo), None) => {
|
||||
let lo_bytes = prefix[lo.history_len.min(cap)];
|
||||
let at_bytes = prefix[index];
|
||||
let delta_bytes = at_bytes.saturating_sub(lo_bytes);
|
||||
|
||||
TokenEstimate {
|
||||
tokens: lo.input_total_tokens.saturating_add(delta_bytes / 4),
|
||||
source: EstimateSource::Extrapolated,
|
||||
}
|
||||
}
|
||||
(None, Some(up)) => {
|
||||
let up_bytes = prefix[up.history_len.min(cap)];
|
||||
let at_bytes = prefix[index];
|
||||
if up_bytes == 0 {
|
||||
return TokenEstimate {
|
||||
tokens: 0,
|
||||
source: EstimateSource::Interpolated,
|
||||
};
|
||||
}
|
||||
let t = (at_bytes as u128 * up.input_total_tokens as u128 / up_bytes as u128) as u64;
|
||||
TokenEstimate {
|
||||
tokens: t,
|
||||
source: EstimateSource::Interpolated,
|
||||
}
|
||||
}
|
||||
(None, None) => unreachable!("records non-empty but neither lower nor upper matched"),
|
||||
}
|
||||
}
|
||||
|
||||
/// 現在の history 全体の推定トークン数。
|
||||
pub fn total_tokens(history: &[Item], records: &[UsageRecord]) -> TokenEstimate {
|
||||
let prefix = prefix_bytes(history);
|
||||
tokens_at(history, records, history.len(), &prefix)
|
||||
}
|
||||
|
||||
/// 任意の history index 時点でのプロンプト全長推定。
|
||||
/// `history_len == 0` で 0 を返す。delta 計算 (extract trigger 等) で
|
||||
/// `total_tokens_at(now) - total_tokens_at(pointer)` の形で使う。
|
||||
pub fn total_tokens_at(
|
||||
history: &[Item],
|
||||
records: &[UsageRecord],
|
||||
history_len: usize,
|
||||
) -> TokenEstimate {
|
||||
let prefix = prefix_bytes(history);
|
||||
tokens_at(history, records, history_len.min(history.len()), &prefix)
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
|
||||
fn msg(text: &str) -> Item {
|
||||
Item::user_message(text)
|
||||
}
|
||||
|
||||
fn record(history_len: usize, tokens: u64) -> UsageRecord {
|
||||
UsageRecord {
|
||||
history_len,
|
||||
input_total_tokens: tokens,
|
||||
cache_read_tokens: 0,
|
||||
cache_write_tokens: 0,
|
||||
output_tokens: 0,
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn total_no_data_falls_back_to_byte_estimate() {
|
||||
let history = vec![msg("hello world")];
|
||||
let est = total_tokens(&history, &[]);
|
||||
assert_eq!(est.source, EstimateSource::NoData);
|
||||
assert!(est.tokens > 0);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn total_measured_when_last_record_matches_history_len() {
|
||||
let history = vec![msg("a"), msg("b"), msg("c")];
|
||||
let records = vec![record(3, 120)];
|
||||
let est = total_tokens(&history, &records);
|
||||
assert_eq!(est.source, EstimateSource::Measured);
|
||||
assert_eq!(est.tokens, 120);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn total_extrapolated_when_history_grew_past_last_measurement() {
|
||||
let history = vec![msg("a"), msg("b"), msg("c"), msg("d")];
|
||||
let records = vec![record(3, 100)];
|
||||
let est = total_tokens(&history, &records);
|
||||
assert_eq!(est.source, EstimateSource::Extrapolated);
|
||||
assert!(est.tokens > 100);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn extrapolation_after_single_measurement_uses_byte_fallback_not_total_prompt_rate() {
|
||||
let history = vec![msg("first"), msg(&"tool output ".repeat(400))];
|
||||
let records = vec![record(1, 11_124)];
|
||||
let prefix = prefix_bytes(&history);
|
||||
let delta_bytes = prefix[2].saturating_sub(prefix[1]);
|
||||
|
||||
let est = total_tokens(&history, &records);
|
||||
|
||||
assert_eq!(est.source, EstimateSource::Extrapolated);
|
||||
assert_eq!(est.tokens, 11_124 + delta_bytes / 4);
|
||||
|
||||
let old_projection =
|
||||
11_124 + (delta_bytes as u128 * 11_124_u128 / prefix[1] as u128) as u64;
|
||||
assert!(
|
||||
old_projection > est.tokens.saturating_mul(10),
|
||||
"old_projection={old_projection}, corrected={}",
|
||||
est.tokens
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn extrapolation_after_multiple_measurements_uses_byte_fallback_for_unmeasured_delta() {
|
||||
let history = vec![
|
||||
msg("first"),
|
||||
msg(&"measured increment ".repeat(20)),
|
||||
msg(&"unmeasured increment ".repeat(30)),
|
||||
];
|
||||
let records = vec![record(1, 10_000), record(2, 10_200)];
|
||||
let prefix = prefix_bytes(&history);
|
||||
let delta_bytes = prefix[3].saturating_sub(prefix[2]);
|
||||
|
||||
let est = total_tokens(&history, &records);
|
||||
|
||||
assert_eq!(est.source, EstimateSource::Extrapolated);
|
||||
assert_eq!(est.tokens, 10_200 + delta_bytes / 4);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn extrapolation_does_not_reuse_measured_rate_after_context_projection() {
|
||||
let compacted_span = msg("x");
|
||||
let projected = vec![
|
||||
msg("first"),
|
||||
msg("summary only"),
|
||||
compacted_span,
|
||||
msg("new user input"),
|
||||
];
|
||||
let records = vec![record(1, 10_000), record(3, 30_000)];
|
||||
let prefix = prefix_bytes(&projected);
|
||||
let delta_bytes = prefix[4].saturating_sub(prefix[3]);
|
||||
|
||||
let est = total_tokens(&projected, &records);
|
||||
|
||||
assert_eq!(est.source, EstimateSource::Extrapolated);
|
||||
assert_eq!(est.tokens, 30_000 + delta_bytes / 4);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn total_zero_history_is_zero() {
|
||||
let est = total_tokens(&[], &[]);
|
||||
assert_eq!(est.tokens, 0);
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,450 @@
|
||||
//! Tool Definition
|
||||
//!
|
||||
//! Traits for defining tools callable by LLM.
|
||||
//! Usually auto-implemented using the `#[tool]` macro.
|
||||
|
||||
use std::collections::HashMap;
|
||||
use std::sync::Arc;
|
||||
|
||||
use async_trait::async_trait;
|
||||
use serde::{Deserialize, Serialize};
|
||||
use serde_json::Value;
|
||||
use thiserror::Error;
|
||||
|
||||
/// Error during tool execution
|
||||
#[derive(Debug, Error)]
|
||||
pub enum ToolError {
|
||||
/// Invalid argument
|
||||
#[error("Invalid argument: {0}")]
|
||||
InvalidArgument(String),
|
||||
/// Execution failed
|
||||
#[error("Execution failed: {0}")]
|
||||
ExecutionFailed(String),
|
||||
/// Internal error
|
||||
#[error("Internal error: {0}")]
|
||||
Internal(String),
|
||||
}
|
||||
|
||||
// =============================================================================
|
||||
// ToolOutput - Tool execution result with summary + content
|
||||
// =============================================================================
|
||||
|
||||
/// Threshold below which tool output is treated as summary-only (no content).
|
||||
/// Outputs this small don't benefit from pruning.
|
||||
pub const SUMMARY_THRESHOLD: usize = 200;
|
||||
|
||||
/// Byte-size caps applied to tool execution `content` at the Engine's
|
||||
/// tool-execution boundary, before results enter conversation history.
|
||||
///
|
||||
/// Exists so a single oversized tool result (e.g. a wide `Glob` scan)
|
||||
/// cannot blow past the provider's per-minute input-token rate limit.
|
||||
/// Individual tools are not trusted to self-limit — this is the single
|
||||
/// chokepoint.
|
||||
///
|
||||
/// The unit is bytes rather than tokens because accurate pre-send token
|
||||
/// estimation is not available. The limits can be migrated to token
|
||||
/// units later without changing callers.
|
||||
#[derive(Debug, Clone)]
|
||||
pub struct ToolOutputLimits {
|
||||
/// Cap applied to any tool not listed in `per_tool`.
|
||||
pub default_max_bytes: usize,
|
||||
/// Per-tool overrides, keyed by tool registration name.
|
||||
pub per_tool: HashMap<String, usize>,
|
||||
}
|
||||
|
||||
impl ToolOutputLimits {
|
||||
/// Resolve the cap for a given tool name.
|
||||
pub fn limit_for(&self, tool_name: &str) -> usize {
|
||||
self.per_tool
|
||||
.get(tool_name)
|
||||
.copied()
|
||||
.unwrap_or(self.default_max_bytes)
|
||||
}
|
||||
}
|
||||
|
||||
/// Truncate `content` in-place if it exceeds `limit` bytes, replacing
|
||||
/// the dropped tail with a short human- and LLM-readable marker so the
|
||||
/// model can self-correct by narrowing its query.
|
||||
///
|
||||
/// The cut point is walked back to the nearest UTF-8 char boundary so
|
||||
/// multibyte characters are never split.
|
||||
pub(crate) fn truncate_content(content: &mut String, limit: usize) {
|
||||
let original_len = content.len();
|
||||
if original_len <= limit {
|
||||
return;
|
||||
}
|
||||
|
||||
let suffix_template = "\n\n[truncated: %BYTES% bytes dropped, refine your query]";
|
||||
// Reserve enough headroom for the suffix (upper bound on the byte length
|
||||
// of the number substitution). usize::MAX fits in 20 digits.
|
||||
let reserved = suffix_template.len() + 20 - "%BYTES%".len();
|
||||
let body_budget = limit.saturating_sub(reserved);
|
||||
|
||||
let mut cut = body_budget.min(original_len);
|
||||
while cut > 0 && !content.is_char_boundary(cut) {
|
||||
cut -= 1;
|
||||
}
|
||||
content.truncate(cut);
|
||||
let dropped = original_len - cut;
|
||||
content.push_str(&suffix_template.replace("%BYTES%", &dropped.to_string()));
|
||||
}
|
||||
|
||||
/// Tool execution result.
|
||||
///
|
||||
/// Every output has a mandatory `summary` (1-2 lines) that persists in
|
||||
/// conversation history even after pruning. The optional `content` carries
|
||||
/// full details and is removed by the Prune mechanism when the context
|
||||
/// grows too large.
|
||||
#[derive(Debug, Clone, Serialize, Deserialize)]
|
||||
pub struct ToolOutput {
|
||||
/// Short summary (1-2 lines). Always remains in history.
|
||||
pub summary: String,
|
||||
/// Detailed output. Removed by Prune when old enough.
|
||||
#[serde(default, skip_serializing_if = "Option::is_none")]
|
||||
pub content: Option<String>,
|
||||
}
|
||||
|
||||
impl From<String> for ToolOutput {
|
||||
fn from(s: String) -> Self {
|
||||
if s.len() <= SUMMARY_THRESHOLD {
|
||||
ToolOutput {
|
||||
summary: s,
|
||||
content: None,
|
||||
}
|
||||
} else {
|
||||
let lines = s.lines().count();
|
||||
let first_line: String = s.lines().next().unwrap_or("").chars().take(80).collect();
|
||||
let summary = format!("{lines} lines | {first_line}…");
|
||||
ToolOutput {
|
||||
summary,
|
||||
content: Some(s),
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// =============================================================================
|
||||
// ToolMeta - Immutable Meta Information
|
||||
// =============================================================================
|
||||
|
||||
/// Origin metadata for a registered tool.
|
||||
///
|
||||
/// This metadata is intentionally not part of the provider-facing tool schema.
|
||||
/// It lets host layers audit where a model-visible tool definition came from
|
||||
/// while keeping execution and permission semantics in the normal Engine path.
|
||||
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
|
||||
pub struct ToolOrigin {
|
||||
/// Origin kind, for example `plugin` or `builtin`.
|
||||
pub kind: String,
|
||||
/// Package-local plugin id.
|
||||
pub plugin_id: String,
|
||||
/// Source-qualified plugin/package reference when `kind == "plugin"`.
|
||||
pub plugin_ref: String,
|
||||
/// Plugin source such as `user`, `project`, or `builtin`.
|
||||
pub source: String,
|
||||
/// Resolved package digest.
|
||||
pub digest: String,
|
||||
/// Resolved package version.
|
||||
pub package_version: String,
|
||||
/// Plugin API/schema version declared by the package.
|
||||
pub package_api_version: u32,
|
||||
/// Surface that contributed this tool. Plugin tools use `tool`.
|
||||
pub surface: String,
|
||||
}
|
||||
|
||||
/// Tool meta information (fixed at registration, immutable)
|
||||
///
|
||||
/// Generated from `ToolDefinition` factory and does not change after registration with Engine.
|
||||
/// Used for sending tool definitions to LLM.
|
||||
#[derive(Debug, Clone, PartialEq, Eq)]
|
||||
pub struct ToolMeta {
|
||||
/// Tool name (used by LLM for identification)
|
||||
pub name: String,
|
||||
/// Tool description (included in prompt to LLM)
|
||||
pub description: String,
|
||||
/// JSON Schema for arguments
|
||||
pub input_schema: Value,
|
||||
/// Optional host-side origin metadata. This is not exposed to the LLM.
|
||||
pub origin: Option<ToolOrigin>,
|
||||
}
|
||||
|
||||
impl ToolMeta {
|
||||
/// Create a new ToolMeta
|
||||
pub fn new(name: impl Into<String>) -> Self {
|
||||
Self {
|
||||
name: name.into(),
|
||||
description: String::new(),
|
||||
input_schema: Value::Object(Default::default()),
|
||||
origin: None,
|
||||
}
|
||||
}
|
||||
|
||||
/// Set the description
|
||||
pub fn description(mut self, desc: impl Into<String>) -> Self {
|
||||
self.description = desc.into();
|
||||
self
|
||||
}
|
||||
|
||||
/// Set the argument schema
|
||||
pub fn input_schema(mut self, schema: Value) -> Self {
|
||||
self.input_schema = schema;
|
||||
self
|
||||
}
|
||||
|
||||
/// Set host-side origin metadata.
|
||||
pub fn origin(mut self, origin: ToolOrigin) -> Self {
|
||||
self.origin = Some(origin);
|
||||
self
|
||||
}
|
||||
}
|
||||
|
||||
// =============================================================================
|
||||
// ToolDefinition - Factory Type
|
||||
// =============================================================================
|
||||
|
||||
/// Tool definition factory
|
||||
///
|
||||
/// When called, returns `(ToolMeta, Arc<dyn Tool>)`.
|
||||
/// Called once during Engine registration, and the meta information and instance
|
||||
/// are cached at session scope.
|
||||
///
|
||||
/// # Examples
|
||||
///
|
||||
/// ```ignore
|
||||
/// let def: ToolDefinition = Arc::new(|| {
|
||||
/// (
|
||||
/// ToolMeta::new("my_tool")
|
||||
/// .description("My tool description")
|
||||
/// .input_schema(json!({"type": "object"})),
|
||||
/// Arc::new(MyToolImpl { state: 0 }) as Arc<dyn Tool>,
|
||||
/// )
|
||||
/// });
|
||||
/// engine.register_tool(def)?;
|
||||
/// ```
|
||||
pub type ToolDefinition = Arc<dyn Fn() -> (ToolMeta, Arc<dyn Tool>) + Send + Sync>;
|
||||
|
||||
/// Per-call context supplied by the engine when executing a tool call.
|
||||
///
|
||||
/// The context identifies a tool call within one assistant response's tool-call
|
||||
/// batch without imposing any scheduling policy on the engine. Tool
|
||||
/// implementations may use it for response-local ordering, diagnostics, or
|
||||
/// correlation, but it is intentionally not a handle to engine state, history,
|
||||
/// or session mutation.
|
||||
#[derive(Debug, Clone, PartialEq, Eq)]
|
||||
pub struct ToolExecutionContext {
|
||||
/// Provider/tool-call id for the call being executed.
|
||||
pub call_id: String,
|
||||
/// Engine-local identity shared by all tool calls from one execution batch.
|
||||
pub batch_id: String,
|
||||
/// Zero-based order of this call in the model-returned tool-call list.
|
||||
pub call_index: usize,
|
||||
}
|
||||
|
||||
impl ToolExecutionContext {
|
||||
pub fn new(call_id: impl Into<String>, batch_id: impl Into<String>, call_index: usize) -> Self {
|
||||
Self {
|
||||
call_id: call_id.into(),
|
||||
batch_id: batch_id.into(),
|
||||
call_index,
|
||||
}
|
||||
}
|
||||
|
||||
/// Context for direct, non-engine calls in unit tests and low-level callers.
|
||||
pub fn direct() -> Self {
|
||||
Self::new("direct", "direct", 0)
|
||||
}
|
||||
}
|
||||
|
||||
impl Default for ToolExecutionContext {
|
||||
fn default() -> Self {
|
||||
Self::direct()
|
||||
}
|
||||
}
|
||||
|
||||
// =============================================================================
|
||||
// Tool trait
|
||||
// =============================================================================
|
||||
|
||||
/// Trait for defining tools callable by LLM
|
||||
///
|
||||
/// Tools are used by LLM to access external resources
|
||||
/// or execute computations.
|
||||
/// Can maintain state during the session.
|
||||
///
|
||||
/// # How to Implement
|
||||
///
|
||||
/// Usually auto-implemented using the `#[tool_registry]` macro:
|
||||
///
|
||||
/// ```ignore
|
||||
/// #[tool_registry]
|
||||
/// impl MyApp {
|
||||
/// #[tool]
|
||||
/// async fn search(&self, query: String) -> String {
|
||||
/// format!("Results for: {}", query)
|
||||
/// }
|
||||
/// }
|
||||
///
|
||||
/// // Register
|
||||
/// engine.register_tool(app.search_definition())?;
|
||||
/// ```
|
||||
///
|
||||
/// # Manual Implementation
|
||||
///
|
||||
/// ```ignore
|
||||
/// use llm_engine::tool::{Tool, ToolError, ToolExecutionContext, ToolMeta, ToolDefinition, ToolOutput};
|
||||
/// use std::sync::Arc;
|
||||
///
|
||||
/// struct MyTool { counter: std::sync::atomic::AtomicUsize }
|
||||
///
|
||||
/// #[async_trait::async_trait]
|
||||
/// impl Tool for MyTool {
|
||||
/// async fn execute(&self, input: &str, ctx: ToolExecutionContext) -> Result<ToolOutput, ToolError> {
|
||||
/// self.counter.fetch_add(1, std::sync::atomic::Ordering::SeqCst);
|
||||
/// Ok(format!("call {}: {}", ctx.call_index, input).into())
|
||||
/// }
|
||||
/// }
|
||||
///
|
||||
/// let def: ToolDefinition = Arc::new(|| {
|
||||
/// (
|
||||
/// ToolMeta::new("my_tool")
|
||||
/// .description("My custom tool")
|
||||
/// .input_schema(serde_json::json!({"type": "object"})),
|
||||
/// Arc::new(MyTool { counter: Default::default() }) as Arc<dyn Tool>,
|
||||
/// )
|
||||
/// });
|
||||
/// ```
|
||||
#[async_trait]
|
||||
pub trait Tool: Send + Sync {
|
||||
/// Execute the tool.
|
||||
///
|
||||
/// # Arguments
|
||||
/// * `input_json` - JSON-formatted arguments generated by LLM
|
||||
/// * `ctx` - response-local call identity and ordering context
|
||||
///
|
||||
/// # Returns
|
||||
/// A [`ToolOutput`] with summary and optional detailed content.
|
||||
/// For simple cases, use `From<String>`: `Ok("done".to_string().into())`
|
||||
async fn execute(
|
||||
&self,
|
||||
input_json: &str,
|
||||
ctx: ToolExecutionContext,
|
||||
) -> Result<ToolOutput, ToolError>;
|
||||
}
|
||||
|
||||
// =============================================================================
|
||||
// Tool Call / Result Types
|
||||
// =============================================================================
|
||||
|
||||
/// Tool call information
|
||||
///
|
||||
/// Represents a ToolUse block from LLM.
|
||||
#[derive(Debug, Clone, Serialize, Deserialize)]
|
||||
pub struct ToolCall {
|
||||
/// Tool call ID (used for linking with response)
|
||||
pub id: String,
|
||||
/// Tool name
|
||||
pub name: String,
|
||||
/// Input arguments (JSON)
|
||||
pub input: Value,
|
||||
}
|
||||
|
||||
/// Tool execution result
|
||||
///
|
||||
/// Intermediate representation between tool execution and history.
|
||||
/// Carries `summary` + optional `content` from [`ToolOutput`].
|
||||
#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)]
|
||||
pub struct ToolResult {
|
||||
/// Corresponding tool call ID
|
||||
pub tool_use_id: String,
|
||||
/// Short summary (always kept in history)
|
||||
pub summary: String,
|
||||
/// Detailed output (prunable)
|
||||
#[serde(default, skip_serializing_if = "Option::is_none")]
|
||||
pub content: Option<String>,
|
||||
/// Whether this is an error
|
||||
#[serde(default)]
|
||||
pub is_error: bool,
|
||||
}
|
||||
|
||||
impl ToolResult {
|
||||
/// Create a success result from a [`ToolOutput`].
|
||||
pub fn from_output(tool_use_id: impl Into<String>, output: ToolOutput) -> Self {
|
||||
Self {
|
||||
tool_use_id: tool_use_id.into(),
|
||||
summary: output.summary,
|
||||
content: output.content,
|
||||
is_error: false,
|
||||
}
|
||||
}
|
||||
|
||||
/// Create an error result.
|
||||
pub fn error(tool_use_id: impl Into<String>, message: impl Into<String>) -> Self {
|
||||
Self {
|
||||
tool_use_id: tool_use_id.into(),
|
||||
summary: message.into(),
|
||||
content: None,
|
||||
is_error: true,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod truncate_tests {
|
||||
use super::*;
|
||||
|
||||
#[test]
|
||||
fn noop_when_within_limit() {
|
||||
let mut s = "hello world".to_string();
|
||||
truncate_content(&mut s, 1024);
|
||||
assert_eq!(s, "hello world");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn noop_at_exact_limit() {
|
||||
let mut s = "a".repeat(100);
|
||||
truncate_content(&mut s, 100);
|
||||
assert_eq!(s.len(), 100);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn truncates_oversized_ascii_with_marker() {
|
||||
let mut s = "a".repeat(1000);
|
||||
truncate_content(&mut s, 200);
|
||||
assert!(s.contains("[truncated:"));
|
||||
assert!(s.contains("refine your query"));
|
||||
assert!(s.len() <= 200, "result was {} bytes", s.len());
|
||||
let dropped: usize = s
|
||||
.split("[truncated: ")
|
||||
.nth(1)
|
||||
.unwrap()
|
||||
.split(' ')
|
||||
.next()
|
||||
.unwrap()
|
||||
.parse()
|
||||
.unwrap();
|
||||
let body_len = s.find("\n\n[truncated:").unwrap();
|
||||
assert_eq!(body_len + dropped, 1000);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn respects_utf8_char_boundaries() {
|
||||
// 100 copies of "あ" (3 bytes each) = 300 bytes.
|
||||
let mut s = "あ".repeat(100);
|
||||
truncate_content(&mut s, 120);
|
||||
// Truncation must not split a multibyte character.
|
||||
assert!(s.is_char_boundary(s.find("\n\n[truncated:").unwrap_or(s.len())));
|
||||
// And the result must still be valid UTF-8 (implicitly true for String).
|
||||
assert!(s.contains("[truncated:"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn limits_per_tool_override() {
|
||||
let mut limits = ToolOutputLimits {
|
||||
default_max_bytes: 1024,
|
||||
per_tool: HashMap::new(),
|
||||
};
|
||||
limits.per_tool.insert("Read".to_string(), 4096);
|
||||
assert_eq!(limits.limit_for("Read"), 4096);
|
||||
assert_eq!(limits.limit_for("Grep"), 1024);
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,542 @@
|
||||
use std::collections::HashMap;
|
||||
use std::sync::{Arc, Mutex};
|
||||
|
||||
use thiserror::Error;
|
||||
|
||||
use crate::llm_client::ToolDefinition as LlmToolDefinition;
|
||||
use crate::tool::{
|
||||
Tool, ToolDefinition as EngineToolDefinition, ToolExecutionContext, ToolMeta, ToolOutput,
|
||||
};
|
||||
|
||||
type ToolMap = HashMap<String, (ToolMeta, Arc<dyn Tool>)>;
|
||||
|
||||
/// Errors produced by ToolServer operations.
|
||||
#[derive(Debug, Error, PartialEq, Eq)]
|
||||
pub enum ToolServerError {
|
||||
/// A tool with the same name already exists.
|
||||
#[error("Tool with name '{0}' already registered")]
|
||||
DuplicateName(String),
|
||||
/// Requested tool was not found.
|
||||
#[error("Tool '{0}' not found")]
|
||||
ToolNotFound(String),
|
||||
/// Tool execution failed.
|
||||
#[error("Tool execution failed: {0}")]
|
||||
ToolExecution(String),
|
||||
}
|
||||
|
||||
/// In-memory tool server.
|
||||
#[derive(Clone, Default)]
|
||||
pub struct ToolServer {
|
||||
tools: Arc<Mutex<ToolMap>>,
|
||||
pending: Arc<Mutex<Vec<EngineToolDefinition>>>,
|
||||
}
|
||||
|
||||
impl ToolServer {
|
||||
/// Create a new empty tool server.
|
||||
pub fn new() -> Self {
|
||||
Self::default()
|
||||
}
|
||||
|
||||
/// Create a handle for shared access.
|
||||
pub fn handle(&self) -> ToolServerHandle {
|
||||
ToolServerHandle {
|
||||
tools: Arc::clone(&self.tools),
|
||||
pending: Arc::clone(&self.pending),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// Shareable handle to a tool server.
|
||||
#[derive(Clone, Default)]
|
||||
pub struct ToolServerHandle {
|
||||
tools: Arc<Mutex<ToolMap>>,
|
||||
pending: Arc<Mutex<Vec<EngineToolDefinition>>>,
|
||||
}
|
||||
|
||||
impl ToolServerHandle {
|
||||
/// Queue a tool factory for deferred initialization.
|
||||
///
|
||||
/// The factory is **not** called here; it is stored and executed
|
||||
/// when [`flush_pending`](Self::flush_pending) is called (typically
|
||||
/// at the start of `Engine::run()`).
|
||||
pub(crate) fn register_tool(&self, factory: EngineToolDefinition) {
|
||||
self.pending
|
||||
.lock()
|
||||
.unwrap_or_else(|e| e.into_inner())
|
||||
.push(factory);
|
||||
}
|
||||
|
||||
/// Queue many tool factories for deferred initialization.
|
||||
pub(crate) fn register_tools(&self, factories: impl IntoIterator<Item = EngineToolDefinition>) {
|
||||
let mut guard = self.pending.lock().unwrap_or_else(|e| e.into_inner());
|
||||
guard.extend(factories);
|
||||
}
|
||||
|
||||
/// Execute all pending factories and register the resulting tools.
|
||||
///
|
||||
/// Called implicitly by `Engine::lock()` before the first turn.
|
||||
/// Exposed as `pub` so higher layers (e.g. Pod) can force-materialise
|
||||
/// tools earlier — for example when building a system-prompt template
|
||||
/// context that needs the list of registered tool names. Redundant
|
||||
/// calls are no-ops.
|
||||
///
|
||||
/// # Panics
|
||||
///
|
||||
/// Panics if any factory produces a tool whose name collides with
|
||||
/// an already-registered tool. Duplicate names are a programming
|
||||
/// error and should be caught during development.
|
||||
pub fn flush_pending(&self) {
|
||||
let pending: Vec<_> = {
|
||||
let mut guard = self.pending.lock().unwrap_or_else(|e| e.into_inner());
|
||||
std::mem::take(&mut *guard)
|
||||
};
|
||||
if pending.is_empty() {
|
||||
return;
|
||||
}
|
||||
// Execute all factories first, then validate and insert atomically.
|
||||
let materialized: Vec<_> = pending.into_iter().map(|f| f()).collect();
|
||||
let mut tools = self.tools.lock().unwrap_or_else(|e| e.into_inner());
|
||||
for (meta, instance) in materialized {
|
||||
assert!(
|
||||
!tools.contains_key(&meta.name),
|
||||
"duplicate tool name: '{}'",
|
||||
meta.name,
|
||||
);
|
||||
tools.insert(meta.name.clone(), (meta, instance));
|
||||
}
|
||||
}
|
||||
|
||||
/// Get a tool by name for hook contexts.
|
||||
pub fn get_tool(&self, name: &str) -> Option<(ToolMeta, Arc<dyn Tool>)> {
|
||||
let guard = self.tools.lock().unwrap_or_else(|e| e.into_inner());
|
||||
guard
|
||||
.get(name)
|
||||
.map(|(meta, tool)| (meta.clone(), Arc::clone(tool)))
|
||||
}
|
||||
|
||||
/// Execute a tool by name.
|
||||
pub async fn call_tool(
|
||||
&self,
|
||||
name: &str,
|
||||
input_json: &str,
|
||||
ctx: ToolExecutionContext,
|
||||
) -> Result<ToolOutput, ToolServerError> {
|
||||
let tool = {
|
||||
let guard = self.tools.lock().unwrap_or_else(|e| e.into_inner());
|
||||
let (_, tool) = guard
|
||||
.get(name)
|
||||
.ok_or_else(|| ToolServerError::ToolNotFound(name.to_string()))?;
|
||||
Arc::clone(tool)
|
||||
};
|
||||
tool.execute(input_json, ctx)
|
||||
.await
|
||||
.map_err(|e| ToolServerError::ToolExecution(e.to_string()))
|
||||
}
|
||||
|
||||
/// Remove a registered tool by name.
|
||||
///
|
||||
/// In-flight calls that already obtained an `Arc<dyn Tool>` clone are
|
||||
/// unaffected and will run to completion.
|
||||
pub fn unregister(&self, name: &str) -> Result<(), ToolServerError> {
|
||||
let mut guard = self.tools.lock().unwrap_or_else(|e| e.into_inner());
|
||||
guard
|
||||
.remove(name)
|
||||
.map(|_| ())
|
||||
.ok_or_else(|| ToolServerError::ToolNotFound(name.to_string()))
|
||||
}
|
||||
|
||||
/// Replace an existing tool with a new implementation.
|
||||
///
|
||||
/// The factory is called immediately and the resulting tool overwrites
|
||||
/// the entry with the same name. Returns `ToolNotFound` if the name
|
||||
/// produced by the factory does not match any registered tool.
|
||||
pub fn replace(&self, factory: EngineToolDefinition) -> Result<(), ToolServerError> {
|
||||
let (meta, instance) = factory();
|
||||
let mut guard = self.tools.lock().unwrap_or_else(|e| e.into_inner());
|
||||
if !guard.contains_key(&meta.name) {
|
||||
return Err(ToolServerError::ToolNotFound(meta.name));
|
||||
}
|
||||
guard.insert(meta.name.clone(), (meta, instance));
|
||||
Ok(())
|
||||
}
|
||||
|
||||
/// Build deterministic tool definitions sorted by tool name.
|
||||
pub fn tool_definitions_sorted(&self) -> Vec<LlmToolDefinition> {
|
||||
let guard = self.tools.lock().unwrap_or_else(|e| e.into_inner());
|
||||
let mut defs: Vec<_> = guard
|
||||
.values()
|
||||
.map(|(meta, _)| {
|
||||
LlmToolDefinition::new(&meta.name)
|
||||
.description(&meta.description)
|
||||
.input_schema(meta.input_schema.clone())
|
||||
})
|
||||
.collect();
|
||||
defs.sort_by(|a, b| a.name.cmp(&b.name));
|
||||
defs
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use std::sync::Arc;
|
||||
|
||||
use async_trait::async_trait;
|
||||
use serde_json::json;
|
||||
|
||||
use super::*;
|
||||
use crate::tool::{Tool, ToolDefinition, ToolError, ToolMeta};
|
||||
|
||||
struct EchoTool;
|
||||
|
||||
#[async_trait]
|
||||
impl Tool for EchoTool {
|
||||
async fn execute(
|
||||
&self,
|
||||
input_json: &str,
|
||||
_ctx: crate::tool::ToolExecutionContext,
|
||||
) -> Result<ToolOutput, ToolError> {
|
||||
Ok(input_json.to_string().into())
|
||||
}
|
||||
}
|
||||
|
||||
fn def(name: &'static str) -> ToolDefinition {
|
||||
Arc::new(move || {
|
||||
(
|
||||
ToolMeta::new(name)
|
||||
.description(format!("desc-{name}"))
|
||||
.input_schema(json!({"type":"object"})),
|
||||
Arc::new(EchoTool) as Arc<dyn Tool>,
|
||||
)
|
||||
})
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn flush_pending_registers_tools() {
|
||||
let handle = ToolServer::new().handle();
|
||||
handle.register_tool(def("alpha"));
|
||||
handle.register_tool(def("beta"));
|
||||
|
||||
// Before flush, no tools are available
|
||||
assert!(handle.get_tool("alpha").is_none());
|
||||
|
||||
handle.flush_pending();
|
||||
|
||||
// After flush, tools are available
|
||||
assert!(handle.get_tool("alpha").is_some());
|
||||
assert!(handle.get_tool("beta").is_some());
|
||||
}
|
||||
|
||||
#[test]
|
||||
#[should_panic(expected = "duplicate tool name: 'alpha'")]
|
||||
fn flush_pending_duplicate_name_panics() {
|
||||
let handle = ToolServer::new().handle();
|
||||
handle.register_tool(def("alpha"));
|
||||
handle.flush_pending();
|
||||
|
||||
handle.register_tool(def("alpha"));
|
||||
handle.flush_pending(); // panics
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn call_tool_success_and_not_found() {
|
||||
let handle = ToolServer::new().handle();
|
||||
handle.register_tool(def("echo"));
|
||||
handle.flush_pending();
|
||||
|
||||
let out = handle
|
||||
.call_tool("echo", r#"{"x":1}"#, Default::default())
|
||||
.await
|
||||
.expect("call");
|
||||
assert_eq!(out.summary, r#"{"x":1}"#);
|
||||
assert!(out.content.is_none());
|
||||
|
||||
let err = handle
|
||||
.call_tool("missing", "{}", Default::default())
|
||||
.await
|
||||
.expect_err("missing tool");
|
||||
assert_eq!(err, ToolServerError::ToolNotFound("missing".to_string()));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn tool_definitions_are_sorted() {
|
||||
let handle = ToolServer::new().handle();
|
||||
handle.register_tool(def("zeta"));
|
||||
handle.register_tool(def("alpha"));
|
||||
handle.register_tool(def("beta"));
|
||||
handle.flush_pending();
|
||||
|
||||
let names: Vec<_> = handle
|
||||
.tool_definitions_sorted()
|
||||
.into_iter()
|
||||
.map(|d| d.name)
|
||||
.collect();
|
||||
assert_eq!(names, vec!["alpha", "beta", "zeta"]);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn flush_pending_is_noop_when_empty() {
|
||||
let handle = ToolServer::new().handle();
|
||||
handle.flush_pending();
|
||||
handle.flush_pending();
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn unregister_removes_tool() {
|
||||
let handle = ToolServer::new().handle();
|
||||
handle.register_tool(def("alpha"));
|
||||
handle.flush_pending();
|
||||
|
||||
handle.unregister("alpha").expect("unregister");
|
||||
assert!(handle.get_tool("alpha").is_none());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn unregister_not_found() {
|
||||
let handle = ToolServer::new().handle();
|
||||
let err = handle.unregister("ghost").expect_err("should fail");
|
||||
assert_eq!(err, ToolServerError::ToolNotFound("ghost".to_string()));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn replace_swaps_implementation() {
|
||||
let handle = ToolServer::new().handle();
|
||||
handle.register_tool(def("alpha"));
|
||||
handle.flush_pending();
|
||||
|
||||
// Replace with a tool that returns a fixed string.
|
||||
struct FixedTool;
|
||||
|
||||
#[async_trait]
|
||||
impl Tool for FixedTool {
|
||||
async fn execute(
|
||||
&self,
|
||||
_input_json: &str,
|
||||
_ctx: crate::tool::ToolExecutionContext,
|
||||
) -> Result<ToolOutput, ToolError> {
|
||||
Ok("replaced".to_string().into())
|
||||
}
|
||||
}
|
||||
|
||||
let replacement: ToolDefinition = Arc::new(|| {
|
||||
(
|
||||
ToolMeta::new("alpha")
|
||||
.description("replaced-desc")
|
||||
.input_schema(json!({"type":"object"})),
|
||||
Arc::new(FixedTool) as Arc<dyn Tool>,
|
||||
)
|
||||
});
|
||||
handle.replace(replacement).expect("replace");
|
||||
|
||||
let (meta, _) = handle.get_tool("alpha").expect("exists");
|
||||
assert_eq!(meta.description, "replaced-desc");
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn replace_updates_call_result() {
|
||||
let handle = ToolServer::new().handle();
|
||||
handle.register_tool(def("echo"));
|
||||
handle.flush_pending();
|
||||
|
||||
struct ConstTool;
|
||||
|
||||
#[async_trait]
|
||||
impl Tool for ConstTool {
|
||||
async fn execute(
|
||||
&self,
|
||||
_input_json: &str,
|
||||
_ctx: crate::tool::ToolExecutionContext,
|
||||
) -> Result<ToolOutput, ToolError> {
|
||||
Ok("const".to_string().into())
|
||||
}
|
||||
}
|
||||
|
||||
let replacement: ToolDefinition = Arc::new(|| {
|
||||
(
|
||||
ToolMeta::new("echo")
|
||||
.description("const")
|
||||
.input_schema(json!({"type":"object"})),
|
||||
Arc::new(ConstTool) as Arc<dyn Tool>,
|
||||
)
|
||||
});
|
||||
handle.replace(replacement).expect("replace");
|
||||
|
||||
let out = handle
|
||||
.call_tool("echo", "{}", Default::default())
|
||||
.await
|
||||
.expect("call");
|
||||
assert_eq!(out.summary, "const");
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn unregister_during_execution_does_not_affect_inflight() {
|
||||
use tokio::sync::Notify;
|
||||
|
||||
let started = Arc::new(Notify::new());
|
||||
let finish = Arc::new(Notify::new());
|
||||
|
||||
struct GatedTool {
|
||||
started: Arc<Notify>,
|
||||
finish: Arc<Notify>,
|
||||
}
|
||||
|
||||
#[async_trait]
|
||||
impl Tool for GatedTool {
|
||||
async fn execute(
|
||||
&self,
|
||||
_input_json: &str,
|
||||
_ctx: crate::tool::ToolExecutionContext,
|
||||
) -> Result<ToolOutput, ToolError> {
|
||||
self.started.notify_one();
|
||||
self.finish.notified().await;
|
||||
Ok("done".to_string().into())
|
||||
}
|
||||
}
|
||||
|
||||
let handle = ToolServer::new().handle();
|
||||
let s = Arc::clone(&started);
|
||||
let f = Arc::clone(&finish);
|
||||
handle.register_tool(Arc::new(move || {
|
||||
(
|
||||
ToolMeta::new("slow")
|
||||
.description("slow")
|
||||
.input_schema(json!({"type":"object"})),
|
||||
Arc::new(GatedTool {
|
||||
started: Arc::clone(&s),
|
||||
finish: Arc::clone(&f),
|
||||
}) as Arc<dyn Tool>,
|
||||
)
|
||||
}));
|
||||
handle.flush_pending();
|
||||
|
||||
let h = handle.clone();
|
||||
let call = tokio::spawn(async move { h.call_tool("slow", "{}", Default::default()).await });
|
||||
|
||||
// Wait until the tool is actually executing.
|
||||
started.notified().await;
|
||||
|
||||
// Unregister while the tool is mid-execution.
|
||||
handle.unregister("slow").expect("unregister");
|
||||
assert!(handle.get_tool("slow").is_none());
|
||||
|
||||
// Let the in-flight call finish.
|
||||
finish.notify_one();
|
||||
let result = call.await.expect("join");
|
||||
assert_eq!(result.expect("call").summary, "done");
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn replace_during_execution_inflight_uses_old_impl() {
|
||||
use tokio::sync::Notify;
|
||||
|
||||
let started = Arc::new(Notify::new());
|
||||
let finish = Arc::new(Notify::new());
|
||||
|
||||
struct OldTool {
|
||||
started: Arc<Notify>,
|
||||
finish: Arc<Notify>,
|
||||
}
|
||||
|
||||
#[async_trait]
|
||||
impl Tool for OldTool {
|
||||
async fn execute(
|
||||
&self,
|
||||
_input_json: &str,
|
||||
_ctx: crate::tool::ToolExecutionContext,
|
||||
) -> Result<ToolOutput, ToolError> {
|
||||
self.started.notify_one();
|
||||
self.finish.notified().await;
|
||||
Ok("old".to_string().into())
|
||||
}
|
||||
}
|
||||
|
||||
let handle = ToolServer::new().handle();
|
||||
let s = Arc::clone(&started);
|
||||
let f = Arc::clone(&finish);
|
||||
handle.register_tool(Arc::new(move || {
|
||||
(
|
||||
ToolMeta::new("t")
|
||||
.description("d")
|
||||
.input_schema(json!({"type":"object"})),
|
||||
Arc::new(OldTool {
|
||||
started: Arc::clone(&s),
|
||||
finish: Arc::clone(&f),
|
||||
}) as Arc<dyn Tool>,
|
||||
)
|
||||
}));
|
||||
handle.flush_pending();
|
||||
|
||||
let h = handle.clone();
|
||||
let call = tokio::spawn(async move { h.call_tool("t", "{}", Default::default()).await });
|
||||
|
||||
// Wait until the old tool is mid-execution.
|
||||
started.notified().await;
|
||||
|
||||
// Replace while the old tool is executing.
|
||||
struct NewTool;
|
||||
|
||||
#[async_trait]
|
||||
impl Tool for NewTool {
|
||||
async fn execute(
|
||||
&self,
|
||||
_input_json: &str,
|
||||
_ctx: crate::tool::ToolExecutionContext,
|
||||
) -> Result<ToolOutput, ToolError> {
|
||||
Ok("new".to_string().into())
|
||||
}
|
||||
}
|
||||
|
||||
handle
|
||||
.replace(Arc::new(|| {
|
||||
(
|
||||
ToolMeta::new("t")
|
||||
.description("d")
|
||||
.input_schema(json!({"type":"object"})),
|
||||
Arc::new(NewTool) as Arc<dyn Tool>,
|
||||
)
|
||||
}))
|
||||
.expect("replace");
|
||||
|
||||
// Let the old in-flight call finish — it should return "old".
|
||||
finish.notify_one();
|
||||
let result = call.await.expect("join");
|
||||
assert_eq!(result.expect("call").summary, "old");
|
||||
|
||||
// New calls use the replacement.
|
||||
let out = handle
|
||||
.call_tool("t", "{}", Default::default())
|
||||
.await
|
||||
.expect("call");
|
||||
assert_eq!(out.summary, "new");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn unregister_reflects_in_tool_definitions() {
|
||||
let handle = ToolServer::new().handle();
|
||||
handle.register_tool(def("alpha"));
|
||||
handle.register_tool(def("beta"));
|
||||
handle.flush_pending();
|
||||
|
||||
handle.unregister("alpha").expect("unregister");
|
||||
let names: Vec<_> = handle
|
||||
.tool_definitions_sorted()
|
||||
.into_iter()
|
||||
.map(|d| d.name)
|
||||
.collect();
|
||||
assert_eq!(names, vec!["beta"]);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn replace_not_found() {
|
||||
let handle = ToolServer::new().handle();
|
||||
let factory: ToolDefinition = Arc::new(|| {
|
||||
(
|
||||
ToolMeta::new("ghost")
|
||||
.description("x")
|
||||
.input_schema(json!({"type":"object"})),
|
||||
Arc::new(EchoTool) as Arc<dyn Tool>,
|
||||
)
|
||||
});
|
||||
let err = handle.replace(factory).expect_err("should fail");
|
||||
assert_eq!(err, ToolServerError::ToolNotFound("ghost".to_string()));
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,22 @@
|
||||
//! Per-LLM-request Usage measurement snapshot.
|
||||
//!
|
||||
//! 1 リクエストの送信時点での「ある history prefix 長で計測した占有量」を
|
||||
//! 1 件分にまとめたもの。`UsageEvent` (provider stream イベント) を
|
||||
//! 受けて呼び出し側 (typically Pod) が組み立て、永続化層
|
||||
//! (session-store) に流したり、token accounting (`token_counter`) で
|
||||
//! 履歴として参照したりする。
|
||||
|
||||
/// LLM リクエスト送信時点での占有量スナップショット。
|
||||
#[derive(Debug, Clone, PartialEq, Eq)]
|
||||
pub struct UsageRecord {
|
||||
/// 送信時の history.len()
|
||||
pub history_len: usize,
|
||||
/// history[..history_len] の占有量（プロンプト全長、実測）
|
||||
pub input_total_tokens: u64,
|
||||
/// 上記のうちキャッシュから読み出された分
|
||||
pub cache_read_tokens: u64,
|
||||
/// 上記のうちこのリクエストでキャッシュに書かれた分
|
||||
pub cache_write_tokens: u64,
|
||||
/// このリクエストで生成された出力トークン数
|
||||
pub output_tokens: u64,
|
||||
}
|
||||
Reference in New Issue
Block a user