llm_worker_rs/llm-worker/src/handler.rs

//! Handler/Kind型
//!
//! Timeline層でイベントを処理するためのトレイト。
//! カスタムハンドラを実装してTimelineに登録することで、
//! ストリームイベントを受信できます。

use crate::timeline::event::*;

// =============================================================================
// Kind Trait
// =============================================================================

/// イベント種別を定義するマーカートレイト
///
/// 各Kindは対応するイベント型を指定します。
/// HandlerはこのKindに対して実装され、同じKindに対して
/// 異なるScope型を持つ複数のHandlerを登録できます。
pub trait Kind {
    /// このKindに対応するイベント型
    type Event;
}

// =============================================================================
// Handler Trait
// =============================================================================

/// イベントを処理するハンドラトレイト
///
/// 特定の`Kind`に対するイベント処理を定義します。
/// `Scope`はブロックのライフサイクル中に保持される状態です。
///
/// # Examples
///
/// ```ignore
/// use llm_worker::timeline::{Handler, TextBlockEvent, TextBlockKind};
///
/// struct TextCollector {
///     texts: Vec<String>,
/// }
///
/// impl Handler<TextBlockKind> for TextCollector {
///     type Scope = String;  // ブロックごとのバッファ
///
///     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固有のスコープ型
    ///
    /// ブロック開始時に`Default::default()`で生成され、
    /// ブロック終了時に破棄されます。
    type Scope: Default;

    /// イベントを処理する
    fn on_event(&mut self, scope: &mut Self::Scope, event: &K::Event);
}

// =============================================================================
// Meta Kind Definitions
// =============================================================================

/// Usage Kind - 使用量イベント用
pub struct UsageKind;
impl Kind for UsageKind {
    type Event = UsageEvent;
}

/// Ping Kind - Pingイベント用
pub struct PingKind;
impl Kind for PingKind {
    type Event = PingEvent;
}

/// Status Kind - ステータスイベント用
pub struct StatusKind;
impl Kind for StatusKind {
    type Event = StatusEvent;
}

/// Error Kind - エラーイベント用
pub struct ErrorKind;
impl Kind for ErrorKind {
    type Event = ErrorEvent;
}

// =============================================================================
// Block Kind Definitions
// =============================================================================

/// TextBlock Kind - テキストブロック用
pub struct TextBlockKind;
impl Kind for TextBlockKind {
    type Event = TextBlockEvent;
}

/// テキストブロックのイベント
#[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 - 思考ブロック用
pub struct ThinkingBlockKind;
impl Kind for ThinkingBlockKind {
    type Event = ThinkingBlockEvent;
}

/// 思考ブロックのイベント
#[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,
}

/// ToolUseBlock Kind - ツール使用ブロック用
pub struct ToolUseBlockKind;
impl Kind for ToolUseBlockKind {
    type Event = ToolUseBlockEvent;
}

/// ツール使用ブロックのイベント
#[derive(Debug, Clone, PartialEq)]
pub enum ToolUseBlockEvent {
    Start(ToolUseBlockStart),
    /// ツール引数のJSON部分文字列
    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,
}