638 lines
18 KiB
Rust
638 lines
18 KiB
Rust
//! Timeline層
|
||
//!
|
||
//! LLMからのイベントストリームを受信し、登録されたHandlerにディスパッチします。
|
||
//! 通常はWorker経由で使用しますが、直接使用することも可能です。
|
||
|
||
use std::marker::PhantomData;
|
||
|
||
use super::event::*;
|
||
use crate::handler::*;
|
||
|
||
// =============================================================================
|
||
// 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);
|
||
/// スコープを開始（Block開始時）
|
||
fn start_scope(&mut self);
|
||
/// スコープを終了（Block終了時）
|
||
fn end_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());
|
||
}
|
||
|
||
fn end_scope(&mut self) {
|
||
self.scope = None;
|
||
}
|
||
}
|
||
|
||
// =============================================================================
|
||
// 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,
|
||
scope: Option<H::Scope>,
|
||
}
|
||
|
||
impl<H> ThinkingBlockHandlerWrapper<H>
|
||
where
|
||
H: Handler<ThinkingBlockKind>,
|
||
{
|
||
fn new(handler: H) -> Self {
|
||
Self {
|
||
handler,
|
||
scope: None,
|
||
}
|
||
}
|
||
}
|
||
|
||
impl<H> ErasedBlockHandler for ThinkingBlockHandlerWrapper<H>
|
||
where
|
||
H: Handler<ThinkingBlockKind> + 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,
|
||
&ThinkingBlockEvent::Start(ThinkingBlockStart { index: start.index }),
|
||
);
|
||
}
|
||
}
|
||
|
||
fn dispatch_delta(&mut self, delta: &BlockDelta) {
|
||
if let Some(scope) = &mut self.scope {
|
||
if let DeltaContent::Thinking(text) = &delta.delta {
|
||
self.handler
|
||
.on_event(scope, &ThinkingBlockEvent::Delta(text.clone()));
|
||
}
|
||
}
|
||
}
|
||
|
||
fn dispatch_stop(&mut self, stop: &BlockStop) {
|
||
if let Some(scope) = &mut self.scope {
|
||
self.handler.on_event(
|
||
scope,
|
||
&ThinkingBlockEvent::Stop(ThinkingBlockStop { index: stop.index }),
|
||
);
|
||
}
|
||
}
|
||
|
||
fn dispatch_abort(&mut self, _abort: &BlockAbort) {}
|
||
|
||
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()
|
||
}
|
||
}
|
||
|
||
/// 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_worker::{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>,
|
||
}
|
||
|
||
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,
|
||
}
|
||
}
|
||
|
||
// =========================================================================
|
||
// 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),
|
||
|
||
// 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),
|
||
}
|
||
}
|
||
|
||
fn dispatch_usage(&mut self, event: &UsageEvent) {
|
||
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 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));
|
||
|
||
let recorded = calls.lock().unwrap();
|
||
assert_eq!(recorded.len(), 1);
|
||
assert_eq!(recorded[0].input_tokens, Some(100));
|
||
}
|
||
}
|