Hare/llm_worker_rs
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

feat: Implement HookEventKind

Browse Source
This commit is contained in:
Hare
2026-01-09 19:18:20 +09:00
parent 33f1c218f2
commit 5691b09fc8
15 changed files with 916 additions and 416 deletions
Download Patch File Download Diff File Expand all files Collapse all files
llm-worker/Cargo.toml
+2 -1
View File
@@ -15,7 +15,8 @@ tracing = "0.1"
async-trait = "0.1"
futures = "0.3"
tokio = { version = "1.49", features = ["macros", "rt-multi-thread"] }
reqwest = { version = "0.13.1", default-features = false, features = ["stream", "json", "native-tls"] }
tokio-util = "0.7"
reqwest = { version = "0.13.1", default-features = false, features = ["stream", "json", "native-tls", "http2"] }
eventsource-stream = "0.2"
llm-worker-macros = { path = "../llm-worker-macros", version = "0.1" }
llm-worker/examples/worker_cancel_demo.rs
+71
View File
@@ -0,0 +1,71 @@
//! Worker のキャンセル機能のデモンストレーション
//!
//! ストリーミング受信中に別スレッドからキャンセルする例
use std::sync::Arc;
use std::time::Duration;
use tokio::sync::Mutex;
use llm_worker::{Worker, WorkerResult};
use llm_worker::llm_client::providers::anthropic::AnthropicClient;
#[tokio::main]
async fn main() -> Result<(), Box<dyn std::error::Error>> {
// .envファイルを読み込む
dotenv::dotenv().ok();
// ロギング初期化
tracing_subscriber::fmt()
.with_env_filter(
tracing_subscriber::EnvFilter::try_from_default_env()
.unwrap_or_else(|_| tracing_subscriber::EnvFilter::new("info")),
)
.init();
let api_key = std::env::var("ANTHROPIC_API_KEY")
.expect("ANTHROPIC_API_KEY environment variable not set");
let client = AnthropicClient::new(&api_key, "claude-sonnet-4-20250514");
let worker = Arc::new(Mutex::new(Worker::new(client)));
println!("🚀 Starting Worker...");
println!("💡 Will cancel after 2 seconds\n");
// キャンセルトークンを先に取得（ロックを保持しない）
let cancel_token = {
let w = worker.lock().await;
w.cancellation_token().clone()
};
// タスク1: Workerを実行
let worker_clone = worker.clone();
let task = tokio::spawn(async move {
let mut w = worker_clone.lock().await;
println!("📡 Sending request to LLM...");
match w.run("Tell me a very long story about a brave knight. Make it as detailed as possible with many paragraphs.").await {
Ok(WorkerResult::Finished(_)) => {
println!("✅ Task completed normally");
}
Ok(WorkerResult::Paused(_)) => {
println!("⏸️ Task paused");
}
Err(e) => {
println!("❌ Task error: {}", e);
}
}
});
// タスク2: 2秒後にキャンセル
tokio::spawn(async move {
tokio::time::sleep(Duration::from_secs(2)).await;
println!("\n🛑 Cancelling worker...");
cancel_token.cancel();
});
// タスク完了を待つ
task.await?;
println!("\n✨ Demo complete!");
Ok(())
}
llm-worker/examples/worker_cli.rs
+6 -6
View File
@@ -41,7 +41,7 @@ use tracing_subscriber::EnvFilter;
use clap::{Parser, ValueEnum};
use llm_worker::{
Worker,
hook::{ControlFlow, HookError, ToolResult, WorkerHook},
hook::{AfterToolCall, AfterToolCallResult, Hook, HookError, ToolResult},
llm_client::{
LlmClient,
providers::{
@@ -282,11 +282,11 @@ impl ToolResultPrinterHook {
}
#[async_trait]
impl WorkerHook for ToolResultPrinterHook {
async fn after_tool_call(
impl Hook<AfterToolCall> for ToolResultPrinterHook {
async fn call(
&self,
tool_result: &mut ToolResult,
) -> Result<ControlFlow, HookError> {
) -> Result<AfterToolCallResult, HookError> {
let name = self
.call_names
.lock()
@@ -300,7 +300,7 @@ impl WorkerHook for ToolResultPrinterHook {
println!(" Result ({}): ✅ {}", name, tool_result.content);
}
Ok(ControlFlow::Continue)
Ok(AfterToolCallResult::Continue)
}
}
@@ -451,7 +451,7 @@ async fn main() -> Result<(), Box<dyn std::error::Error>> {
.on_text_block(StreamingPrinter::new())
.on_tool_use_block(ToolCallPrinter::new(tool_call_names.clone()));
worker.add_hook(ToolResultPrinterHook::new(tool_call_names));
worker.add_after_tool_call_hook(ToolResultPrinterHook::new(tool_call_names));
// ワンショットモード
if let Some(prompt) = args.prompt {
llm-worker/src/hook.rs
+56 -89
View File
@@ -8,33 +8,72 @@ use serde_json::Value;
use thiserror::Error;
// =============================================================================
// Control Flow Types
// Hook Event Kinds
// =============================================================================
/// Hook処理の制御フロー
pub trait HookEventKind: Send + Sync + 'static {
type Input;
type Output;
}
pub struct OnMessageSend;
pub struct BeforeToolCall;
pub struct AfterToolCall;
pub struct OnTurnEnd;
pub struct OnAbort;
#[derive(Debug, Clone, PartialEq, Eq)]
pub enum ControlFlow {
/// 処理を続行
pub enum OnMessageSendResult {
Continue,
Cancel(String),
}
#[derive(Debug, Clone, PartialEq, Eq)]
pub enum BeforeToolCallResult {
Continue,
/// 現在の処理をスキップ（Tool実行など）
Skip,
/// 処理を中断
Abort(String),
/// 処理を一時停止（再開可能）
Pause,
}
/// ターン終了時の判定結果
#[derive(Debug, Clone, PartialEq, Eq)]
pub enum AfterToolCallResult {
Continue,
Abort(String),
}
#[derive(Debug, Clone)]
pub enum TurnResult {
/// ターンを終了
pub enum OnTurnEndResult {
Finish,
/// メッセージを追加してターン継続（自己修正など）
ContinueWithMessages(Vec<crate::Message>),
/// ターンを一時停止
Paused,
}
impl HookEventKind for OnMessageSend {
type Input = Vec<crate::Message>;
type Output = OnMessageSendResult;
}
impl HookEventKind for BeforeToolCall {
type Input = ToolCall;
type Output = BeforeToolCallResult;
}
impl HookEventKind for AfterToolCall {
type Input = ToolResult;
type Output = AfterToolCallResult;
}
impl HookEventKind for OnTurnEnd {
type Input = Vec<crate::Message>;
type Output = OnTurnEndResult;
}
impl HookEventKind for OnAbort {
type Input = String;
type Output = ();
}
// =============================================================================
// Tool Call / Result Types
// =============================================================================
@@ -102,85 +141,13 @@ pub enum HookError {
}
// =============================================================================
// WorkerHook Trait
// Hook Trait
// =============================================================================
/// ターンの進行・ツール実行に介入するためのトレイト
/// Hookイベントの処理を行うトレイト
///
/// Hookを使うと、メッセージ送信前、ツール実行前後、ターン終了時に
/// 処理を挟んだり、実行をキャンセルしたりできます。
///
/// # Examples
///
/// ```ignore
/// use llm_worker::hook::{ControlFlow, HookError, ToolCall, TurnResult, WorkerHook};
/// use llm_worker::Message;
///
/// struct ValidationHook;
///
/// #[async_trait::async_trait]
/// impl WorkerHook for ValidationHook {
/// async fn before_tool_call(&self, call: &mut ToolCall) -> Result<ControlFlow, HookError> {
/// // 危険なツールをブロック
/// if call.name == "delete_all" {
/// return Ok(ControlFlow::Skip);
/// }
/// Ok(ControlFlow::Continue)
/// }
///
/// async fn on_turn_end(&self, messages: &[Message]) -> Result<TurnResult, HookError> {
/// // 条件を満たさなければ追加メッセージで継続
/// if messages.len() < 3 {
/// return Ok(TurnResult::ContinueWithMessages(vec![
/// Message::user("Please elaborate.")
/// ]));
/// }
/// Ok(TurnResult::Finish)
/// }
/// }
/// ```
///
/// # デフォルト実装
///
/// すべてのメソッドにはデフォルト実装があり、何も行わず`Continue`を返します。
/// 必要なメソッドのみオーバーライドしてください。
/// 各イベント種別は戻り値型が異なるため、`HookEventKind`を介して型を制約する。
#[async_trait]
pub trait WorkerHook: Send + Sync {
/// メッセージ送信前に呼ばれる
///
/// リクエストに含まれるメッセージリストを参照・改変できます。
/// `ControlFlow::Abort`を返すとターンが中断されます。
async fn on_message_send(
&self,
_context: &mut Vec<crate::Message>,
) -> Result<ControlFlow, HookError> {
Ok(ControlFlow::Continue)
}
/// ツール実行前に呼ばれる
///
/// ツール呼び出しの引数を書き換えたり、実行をスキップしたりできます。
/// `ControlFlow::Skip`を返すとこのツールの実行がスキップされます。
async fn before_tool_call(&self, _tool_call: &mut ToolCall) -> Result<ControlFlow, HookError> {
Ok(ControlFlow::Continue)
}
/// ツール実行後に呼ばれる
///
/// ツールの実行結果を書き換えたり、隠蔽したりできます。
async fn after_tool_call(
&self,
_tool_result: &mut ToolResult,
) -> Result<ControlFlow, HookError> {
Ok(ControlFlow::Continue)
}
/// ターン終了時に呼ばれる
///
/// 生成されたメッセージを検査し、必要なら追加メッセージで継続を指示できます。
/// `TurnResult::ContinueWithMessages`を返すと、指定したメッセージを追加して
/// 次のターンに進みます。
async fn on_turn_end(&self, _messages: &[crate::Message]) -> Result<TurnResult, HookError> {
Ok(TurnResult::Finish)
}
pub trait Hook<E: HookEventKind>: Send + Sync {
async fn call(&self, input: &mut E::Input) -> Result<E::Output, HookError>;
}
llm-worker/src/lib.rs
+2 -6
View File
@@ -6,7 +6,7 @@
//!
//! - [`Worker`] - LLMとの対話を管理する中心コンポーネント
//! - [`tool::Tool`] - LLMから呼び出し可能なツール
//! - [`hook::WorkerHook`] - ターン進行への介入
//! - [`hook::Hook`] - ターン進行への介入
//! - [`subscriber::WorkerSubscriber`] - ストリーミングイベントの購読
//!
//! # Quick Start
@@ -48,9 +48,5 @@ pub mod subscriber;
pub mod timeline;
pub mod tool;
// =============================================================================
// トップレベル公開（最も頻繁に使う型）
// =============================================================================
pub use message::{ContentPart, Message, MessageContent, Role};
pub use worker::{Worker, WorkerConfig, WorkerError};
pub use worker::{Worker, WorkerConfig, WorkerError, WorkerResult};
llm-worker/src/subscriber.rs
+2 -2
View File
@@ -65,10 +65,10 @@ pub trait WorkerSubscriber: Send {
///
/// ブロック開始時にDefault::default()で生成され、
/// ブロック終了時に破棄される。
type TextBlockScope: Default + Send;
type TextBlockScope: Default + Send + Sync;
/// ツール使用ブロック処理用のスコープ型
type ToolUseBlockScope: Default + Send;
type ToolUseBlockScope: Default + Send + Sync;
// =========================================================================
// ブロックイベント（スコープ管理あり）
llm-worker/src/timeline/timeline.rs
+39 -24
View File
@@ -17,7 +17,7 @@ use crate::handler::*;
/// 各Handlerは独自のScope型を持つため、Timelineで保持するには型消去が必要です。
/// 通常は直接使用せず、`Timeline::on_text_block()`などのメソッド経由で
/// 自動的にラップされます。
pub trait ErasedHandler<K: Kind>: Send {
pub trait ErasedHandler<K: Kind>: Send + Sync {
/// イベントをディスパッチ
fn dispatch(&mut self, event: &K::Event);
/// スコープを開始（Block開始時）
@@ -54,9 +54,9 @@ where
impl<H, K> ErasedHandler<K> for HandlerWrapper<H, K>
where
H: Handler<K> + Send,
H: Handler<K> + Send + Sync,
K: Kind,
H::Scope: Send,
H::Scope: Send + Sync,
{
fn dispatch(&mut self, event: &K::Event) {
if let Some(scope) = &mut self.scope {
@@ -78,7 +78,7 @@ where
// =============================================================================
/// ブロックハンドラーの型消去trait
trait ErasedBlockHandler: Send {
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);
@@ -112,8 +112,8 @@ where
impl<H> ErasedBlockHandler for TextBlockHandlerWrapper<H>
where
H: Handler<TextBlockKind> + Send,
H::Scope: Send,
H: Handler<TextBlockKind> + Send + Sync,
H::Scope: Send + Sync,
{
fn dispatch_start(&mut self, start: &BlockStart) {
if let Some(scope) = &mut self.scope {
@@ -185,8 +185,8 @@ where
impl<H> ErasedBlockHandler for ThinkingBlockHandlerWrapper<H>
where
H: Handler<ThinkingBlockKind> + Send,
H::Scope: Send,
H: Handler<ThinkingBlockKind> + Send + Sync,
H::Scope: Send + Sync,
{
fn dispatch_start(&mut self, start: &BlockStart) {
if let Some(scope) = &mut self.scope {
@@ -255,8 +255,8 @@ where
impl<H> ErasedBlockHandler for ToolUseBlockHandlerWrapper<H>
where
H: Handler<ToolUseBlockKind> + Send,
H::Scope: Send,
H: Handler<ToolUseBlockKind> + Send + Sync,
H::Scope: Send + Sync,
{
fn dispatch_start(&mut self, start: &BlockStart) {
if let Some(scope) = &mut self.scope {
@@ -391,8 +391,8 @@ impl Timeline {
/// UsageKind用のHandlerを登録
pub fn on_usage<H>(&mut self, handler: H) -> &mut Self
where
H: Handler<UsageKind> + Send + 'static,
H::Scope: Send,
H: Handler<UsageKind> + Send + Sync + 'static,
H::Scope: Send + Sync,
{
// Meta系はデフォルトでスコープを開始しておく
let mut wrapper = HandlerWrapper::new(handler);
@@ -404,8 +404,8 @@ impl Timeline {
/// PingKind用のHandlerを登録
pub fn on_ping<H>(&mut self, handler: H) -> &mut Self
where
H: Handler<PingKind> + Send + 'static,
H::Scope: Send,
H: Handler<PingKind> + Send + Sync + 'static,
H::Scope: Send + Sync,
{
let mut wrapper = HandlerWrapper::new(handler);
wrapper.start_scope();
@@ -416,8 +416,8 @@ impl Timeline {
/// StatusKind用のHandlerを登録
pub fn on_status<H>(&mut self, handler: H) -> &mut Self
where
H: Handler<StatusKind> + Send + 'static,
H::Scope: Send,
H: Handler<StatusKind> + Send + Sync + 'static,
H::Scope: Send + Sync,
{
let mut wrapper = HandlerWrapper::new(handler);
wrapper.start_scope();
@@ -428,8 +428,8 @@ impl Timeline {
/// ErrorKind用のHandlerを登録
pub fn on_error<H>(&mut self, handler: H) -> &mut Self
where
H: Handler<ErrorKind> + Send + 'static,
H::Scope: Send,
H: Handler<ErrorKind> + Send + Sync + 'static,
H::Scope: Send + Sync,
{
let mut wrapper = HandlerWrapper::new(handler);
wrapper.start_scope();
@@ -440,8 +440,8 @@ impl Timeline {
/// TextBlockKind用のHandlerを登録
pub fn on_text_block<H>(&mut self, handler: H) -> &mut Self
where
H: Handler<TextBlockKind> + Send + 'static,
H::Scope: Send,
H: Handler<TextBlockKind> + Send + Sync + 'static,
H::Scope: Send + Sync,
{
self.text_block_handlers
.push(Box::new(TextBlockHandlerWrapper::new(handler)));
@@ -451,8 +451,8 @@ impl Timeline {
/// ThinkingBlockKind用のHandlerを登録
pub fn on_thinking_block<H>(&mut self, handler: H) -> &mut Self
where
H: Handler<ThinkingBlockKind> + Send + 'static,
H::Scope: Send,
H: Handler<ThinkingBlockKind> + Send + Sync + 'static,
H::Scope: Send + Sync,
{
self.thinking_block_handlers
.push(Box::new(ThinkingBlockHandlerWrapper::new(handler)));
@@ -462,8 +462,8 @@ impl Timeline {
/// ToolUseBlockKind用のHandlerを登録
pub fn on_tool_use_block<H>(&mut self, handler: H) -> &mut Self
where
H: Handler<ToolUseBlockKind> + Send + 'static,
H::Scope: Send,
H: Handler<ToolUseBlockKind> + Send + Sync + 'static,
H::Scope: Send + Sync,
{
self.tool_use_block_handlers
.push(Box::new(ToolUseBlockHandlerWrapper::new(handler)));
@@ -578,6 +578,21 @@ impl Timeline {
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)]
llm-worker/src/worker.rs
+207 -97
View File
@@ -3,11 +3,16 @@ use std::marker::PhantomData;
use std::sync::{Arc, Mutex};
use futures::StreamExt;
use tokio_util::sync::CancellationToken;
use tracing::{debug, info, trace, warn};
use crate::{
ContentPart, Message, MessageContent, Role,
hook::{ControlFlow, HookError, ToolCall, ToolResult, TurnResult, WorkerHook},
hook::{
AfterToolCall, AfterToolCallResult, BeforeToolCall, BeforeToolCallResult, Hook, HookError,
OnAbort, OnMessageSend, OnMessageSendResult, OnTurnEnd, OnTurnEndResult, ToolCall,
ToolResult,
},
llm_client::{ClientError, ConfigWarning, LlmClient, Request, RequestConfig, ToolDefinition},
state::{Locked, Mutable, WorkerState},
subscriber::{
@@ -37,6 +42,9 @@ pub enum WorkerError {
/// 処理が中断された
#[error("Aborted: {0}")]
Aborted(String),
/// Cancellation Tokenによって中断された
#[error("Cancelled")]
Cancelled,
/// 設定に関する警告（未サポートのオプション）
#[error("Config warnings: {}", .0.iter().map(|w| w.to_string()).collect::<Vec<_>>().join(", "))]
ConfigWarnings(Vec<ConfigWarning>),
@@ -77,7 +85,7 @@ enum ToolExecutionResult {
// =============================================================================
/// ターンイベントを通知するためのコールバック (型消去)
trait TurnNotifier: Send {
trait TurnNotifier: Send + Sync {
fn on_turn_start(&self, turn: usize);
fn on_turn_end(&self, turn: usize);
}
@@ -149,8 +157,16 @@ pub struct Worker<C: LlmClient, S: WorkerState = Mutable> {
tool_call_collector: ToolCallCollector,
/// 登録されたツール
tools: HashMap<String, Arc<dyn Tool>>,
/// 登録されたHook
hooks: Vec<Box<dyn WorkerHook>>,
/// on_message_send Hook
hooks_on_message_send: Vec<Box<dyn Hook<OnMessageSend>>>,
/// before_tool_call Hook
hooks_before_tool_call: Vec<Box<dyn Hook<BeforeToolCall>>>,
/// after_tool_call Hook
hooks_after_tool_call: Vec<Box<dyn Hook<AfterToolCall>>>,
/// on_turn_end Hook
hooks_on_turn_end: Vec<Box<dyn Hook<OnTurnEnd>>>,
/// on_abort Hook
hooks_on_abort: Vec<Box<dyn Hook<OnAbort>>>,
/// システムプロンプト
system_prompt: Option<String>,
/// メッセージ履歴（Workerが所有）
@@ -163,6 +179,8 @@ pub struct Worker<C: LlmClient, S: WorkerState = Mutable> {
turn_notifiers: Vec<Box<dyn TurnNotifier>>,
/// リクエスト設定（max_tokens, temperature等）
request_config: RequestConfig,
/// キャンセレーショントークン（実行中断用）
cancellation_token: CancellationToken,
/// 状態マーカー
_state: PhantomData<S>,
}
@@ -252,30 +270,29 @@ impl<C: LlmClient, S: WorkerState> Worker<C, S> {
}
}
/// Hookを追加する
///
/// Hookはターンの進行・ツール実行に介入できます。
/// 複数のHookを登録した場合、登録順に実行されます。
///
/// # Examples
///
/// ```ignore
/// use llm_worker::{Worker, WorkerHook, ControlFlow, ToolCall};
///
/// struct LoggingHook;
///
/// #[async_trait::async_trait]
/// impl WorkerHook for LoggingHook {
/// async fn before_tool_call(&self, call: &mut ToolCall) -> Result<ControlFlow, HookError> {
/// println!("Calling tool: {}", call.name);
/// Ok(ControlFlow::Continue)
/// }
/// }
///
/// worker.add_hook(LoggingHook);
/// ```
pub fn add_hook(&mut self, hook: impl WorkerHook + 'static) {
self.hooks.push(Box::new(hook));
/// on_message_send Hookを追加する
pub fn add_on_message_send_hook(&mut self, hook: impl Hook<OnMessageSend> + 'static) {
self.hooks_on_message_send.push(Box::new(hook));
}
/// before_tool_call Hookを追加する
pub fn add_before_tool_call_hook(&mut self, hook: impl Hook<BeforeToolCall> + 'static) {
self.hooks_before_tool_call.push(Box::new(hook));
}
/// after_tool_call Hookを追加する
pub fn add_after_tool_call_hook(&mut self, hook: impl Hook<AfterToolCall> + 'static) {
self.hooks_after_tool_call.push(Box::new(hook));
}
/// on_turn_end Hookを追加する
pub fn add_on_turn_end_hook(&mut self, hook: impl Hook<OnTurnEnd> + 'static) {
self.hooks_on_turn_end.push(Box::new(hook));
}
/// on_abort Hookを追加する
pub fn add_on_abort_hook(&mut self, hook: impl Hook<OnAbort> + 'static) {
self.hooks_on_abort.push(Box::new(hook));
}
/// タイムラインへの可変参照を取得（追加ハンドラ登録用）
@@ -375,6 +392,41 @@ impl<C: LlmClient, S: WorkerState> Worker<C, S> {
self.request_config = config;
}
/// 実行をキャンセルする
///
/// 現在実行中のストリーミングやツール実行を中断します。
/// 次のイベントループのチェックポイントでWorkerError::Cancelledが返されます。
///
/// # Examples
///
/// ```ignore
/// use std::sync::Arc;
/// let worker = Arc::new(Mutex::new(Worker::new(client)));
///
/// // 別スレッドで実行
/// let worker_clone = worker.clone();
/// tokio::spawn(async move {
/// let mut w = worker_clone.lock().unwrap();
/// w.run("Long task...").await
/// });
///
/// // キャンセル
/// worker.lock().unwrap().cancel();
/// ```
pub fn cancel(&self) {
self.cancellation_token.cancel();
}
/// キャンセルされているかチェック
pub fn is_cancelled(&self) -> bool {
self.cancellation_token.is_cancelled()
}
/// キャンセレーショントークンへの参照を取得
pub fn cancellation_token(&self) -> &CancellationToken {
&self.cancellation_token
}
/// 登録されたツールからToolDefinitionのリストを生成
fn build_tool_definitions(&self) -> Vec<ToolDefinition> {
self.tools
@@ -430,7 +482,7 @@ impl<C: LlmClient, S: WorkerState> Worker<C, S> {
}
/// リクエストを構築
fn build_request(&self, tool_definitions: &[ToolDefinition]) -> Request {
fn build_request(&self, tool_definitions: &[ToolDefinition], context: &[Message]) -> Request {
let mut request = Request::new();
// システムプロンプトを設定
@@ -439,7 +491,7 @@ impl<C: LlmClient, S: WorkerState> Worker<C, S> {
}
// メッセージを追加
for msg in &self.history {
for msg in context {
// Message から llm_client::Message への変換
request = request.message(crate::llm_client::Message {
role: match msg.role {
@@ -495,36 +547,45 @@ impl<C: LlmClient, S: WorkerState> Worker<C, S> {
}
/// Hooks: on_message_send
async fn run_on_message_send_hooks(&self) -> Result<ControlFlow, WorkerError> {
for hook in &self.hooks {
// Note: Locked状態でも履歴全体を参照として渡す（変更は不可）
// HookのAPIを変更し、immutable参照のみを渡すようにする必要があるかもしれない
// 現在は空のVecを渡して回避（要検討）
let mut temp_context = self.history.clone();
let result = hook.on_message_send(&mut temp_context).await?;
async fn run_on_message_send_hooks(
&self,
) -> Result<(OnMessageSendResult, Vec<Message>), WorkerError> {
let mut temp_context = self.history.clone();
for hook in &self.hooks_on_message_send {
let result = hook.call(&mut temp_context).await?;
match result {
ControlFlow::Continue => continue,
ControlFlow::Skip => return Ok(ControlFlow::Skip),
ControlFlow::Abort(reason) => return Ok(ControlFlow::Abort(reason)),
ControlFlow::Pause => return Ok(ControlFlow::Pause),
OnMessageSendResult::Continue => continue,
OnMessageSendResult::Cancel(reason) => {
return Ok((OnMessageSendResult::Cancel(reason), temp_context));
}
}
}
Ok(ControlFlow::Continue)
Ok((OnMessageSendResult::Continue, temp_context))
}
/// Hooks: on_turn_end
async fn run_on_turn_end_hooks(&self) -> Result<TurnResult, WorkerError> {
for hook in &self.hooks {
let result = hook.on_turn_end(&self.history).await?;
async fn run_on_turn_end_hooks(&self) -> Result<OnTurnEndResult, WorkerError> {
let mut temp_messages = self.history.clone();
for hook in &self.hooks_on_turn_end {
let result = hook.call(&mut temp_messages).await?;
match result {
TurnResult::Finish => continue,
TurnResult::ContinueWithMessages(msgs) => {
return Ok(TurnResult::ContinueWithMessages(msgs));
OnTurnEndResult::Finish => continue,
OnTurnEndResult::ContinueWithMessages(msgs) => {
return Ok(OnTurnEndResult::ContinueWithMessages(msgs));
}
TurnResult::Paused => return Ok(TurnResult::Paused),
OnTurnEndResult::Paused => return Ok(OnTurnEndResult::Paused),
}
}
Ok(TurnResult::Finish)
Ok(OnTurnEndResult::Finish)
}
/// Hooks: on_abort
async fn run_on_abort_hooks(&self, reason: &str) -> Result<(), WorkerError> {
let mut reason = reason.to_string();
for hook in &self.hooks_on_abort {
hook.call(&mut reason).await?;
}
Ok(())
}
/// 未実行のツール呼び出しがあるかチェック（Pauseからの復帰用）
@@ -559,7 +620,7 @@ impl<C: LlmClient, S: WorkerState> Worker<C, S> {
/// 全てのツールに対してbefore_tool_callフックを実行後、
/// 許可されたツールを並列に実行し、結果にafter_tool_callフックを適用する。
async fn execute_tools(
&self,
&mut self,
tool_calls: Vec<ToolCall>,
) -> Result<ToolExecutionResult, WorkerError> {
use futures::future::join_all;
@@ -568,18 +629,18 @@ impl<C: LlmClient, S: WorkerState> Worker<C, S> {
let mut approved_calls = Vec::new();
for mut tool_call in tool_calls {
let mut skip = false;
for hook in &self.hooks {
let result = hook.before_tool_call(&mut tool_call).await?;
for hook in &self.hooks_before_tool_call {
let result = hook.call(&mut tool_call).await?;
match result {
ControlFlow::Continue => {}
ControlFlow::Skip => {
BeforeToolCallResult::Continue => {}
BeforeToolCallResult::Skip => {
skip = true;
break;
}
ControlFlow::Abort(reason) => {
BeforeToolCallResult::Abort(reason) => {
return Err(WorkerError::Aborted(reason));
}
ControlFlow::Pause => {
BeforeToolCallResult::Pause => {
return Ok(ToolExecutionResult::Paused);
}
}
@@ -589,7 +650,7 @@ impl<C: LlmClient, S: WorkerState> Worker<C, S> {
}
}
// Phase 2: 許可されたツールを並列実行
// Phase 2: 許可されたツールを並列実行（キャンセル可能）
let futures: Vec<_> = approved_calls
.into_iter()
.map(|tool_call| {
@@ -612,25 +673,26 @@ impl<C: LlmClient, S: WorkerState> Worker<C, S> {
})
.collect();
let mut results = join_all(futures).await;
// ツール実行をキャンセル可能にする
let mut results = tokio::select! {
results = join_all(futures) => results,
_ = self.cancellation_token.cancelled() => {
info!("Tool execution cancelled");
self.timeline.abort_current_block();
self.run_on_abort_hooks("Cancelled").await?;
return Err(WorkerError::Cancelled);
}
};
// Phase 3: after_tool_call フックを適用
for tool_result in &mut results {
for hook in &self.hooks {
let result = hook.after_tool_call(tool_result).await?;
for hook in &self.hooks_after_tool_call {
let result = hook.call(tool_result).await?;
match result {
ControlFlow::Continue => {}
ControlFlow::Skip => break,
ControlFlow::Abort(reason) => {
AfterToolCallResult::Continue => {}
AfterToolCallResult::Abort(reason) => {
return Err(WorkerError::Aborted(reason));
}
ControlFlow::Pause => {
// after_tool_callでのPauseは結果を受け入れた後、次の処理前に止まる動作とする
// ここではContinue扱いとし、on_message_send等でPauseすることを期待する
// あるいはここでのPauseをサポートする場合は戻り値を調整する必要がある
// 現状はログを出してContinue
warn!("ControlFlow::Pause in after_tool_call is treated as Continue");
}
}
}
}
@@ -663,6 +725,14 @@ impl<C: LlmClient, S: WorkerState> Worker<C, S> {
}
loop {
// キャンセルチェック
if self.cancellation_token.is_cancelled() {
info!("Execution cancelled");
self.timeline.abort_current_block();
self.run_on_abort_hooks("Cancelled").await?;
return Err(WorkerError::Cancelled);
}
// ターン開始を通知
let current_turn = self.turn_count;
debug!(turn = current_turn, "Turn start");
@@ -671,24 +741,21 @@ impl<C: LlmClient, S: WorkerState> Worker<C, S> {
}
// Hook: on_message_send
let control = self.run_on_message_send_hooks().await?;
let (control, request_context) = self.run_on_message_send_hooks().await?;
match control {
ControlFlow::Abort(reason) => {
warn!(reason = %reason, "Aborted by hook");
OnMessageSendResult::Cancel(reason) => {
info!(reason = %reason, "Aborted by hook");
for notifier in &self.turn_notifiers {
notifier.on_turn_end(current_turn);
}
self.run_on_abort_hooks(&reason).await?;
return Err(WorkerError::Aborted(reason));
}
ControlFlow::Pause | ControlFlow::Skip => {
// Skip or Pause -> Pause the worker
return Ok(WorkerResult::Paused(&self.history));
}
ControlFlow::Continue => {}
OnMessageSendResult::Continue => {}
}
// リクエスト構築
let request = self.build_request(&tool_definitions);
let request = self.build_request(&tool_definitions, &request_context);
debug!(
message_count = request.messages.len(),
tool_count = request.tools.len(),
@@ -698,21 +765,49 @@ impl<C: LlmClient, S: WorkerState> Worker<C, S> {
// ストリーム処理
debug!("Starting stream...");
let mut stream = self.client.stream(request).await?;
let mut event_count = 0;
while let Some(event_result) = stream.next().await {
match &event_result {
Ok(event) => {
trace!(event = ?event, "Received event");
event_count += 1;
// ストリームを取得（キャンセル可能）
let mut stream = tokio::select! {
stream_result = self.client.stream(request) => stream_result?,
_ = self.cancellation_token.cancelled() => {
info!("Cancelled before stream started");
self.timeline.abort_current_block();
self.run_on_abort_hooks("Cancelled").await?;
return Err(WorkerError::Cancelled);
}
};
loop {
tokio::select! {
// ストリームからイベントを受信
event_result = stream.next() => {
match event_result {
Some(result) => {
match &result {
Ok(event) => {
trace!(event = ?event, "Received event");
event_count += 1;
}
Err(e) => {
warn!(error = %e, "Stream error");
}
}
let event = result?;
let timeline_event: crate::timeline::event::Event = event.into();
self.timeline.dispatch(&timeline_event);
}
None => break, // ストリーム終了
}
}
Err(e) => {
warn!(error = %e, "Stream error");
// キャンセル待機
_ = self.cancellation_token.cancelled() => {
info!("Stream cancelled");
self.timeline.abort_current_block();
self.run_on_abort_hooks("Cancelled").await?;
return Err(WorkerError::Cancelled);
}
}
let event = event_result?;
let timeline_event: crate::timeline::event::Event = event.into();
self.timeline.dispatch(&timeline_event);
}
debug!(event_count = event_count, "Stream completed");
@@ -736,14 +831,14 @@ impl<C: LlmClient, S: WorkerState> Worker<C, S> {
// ツール呼び出しなし → ターン終了判定
let turn_result = self.run_on_turn_end_hooks().await?;
match turn_result {
TurnResult::Finish => {
OnTurnEndResult::Finish => {
return Ok(WorkerResult::Finished(&self.history));
}
TurnResult::ContinueWithMessages(additional) => {
OnTurnEndResult::ContinueWithMessages(additional) => {
self.history.extend(additional);
continue;
}
TurnResult::Paused => {
OnTurnEndResult::Paused => {
return Ok(WorkerResult::Paused(&self.history));
}
}
@@ -790,13 +885,18 @@ impl<C: LlmClient> Worker<C, Mutable> {
text_block_collector,
tool_call_collector,
tools: HashMap::new(),
hooks: Vec::new(),
hooks_on_message_send: Vec::new(),
hooks_before_tool_call: Vec::new(),
hooks_after_tool_call: Vec::new(),
hooks_on_turn_end: Vec::new(),
hooks_on_abort: Vec::new(),
system_prompt: None,
history: Vec::new(),
locked_prefix_len: 0,
turn_count: 0,
turn_notifiers: Vec::new(),
request_config: RequestConfig::default(),
cancellation_token: CancellationToken::new(),
_state: PhantomData,
}
}
@@ -958,13 +1058,18 @@ impl<C: LlmClient> Worker<C, Mutable> {
text_block_collector: self.text_block_collector,
tool_call_collector: self.tool_call_collector,
tools: self.tools,
hooks: self.hooks,
hooks_on_message_send: self.hooks_on_message_send,
hooks_before_tool_call: self.hooks_before_tool_call,
hooks_after_tool_call: self.hooks_after_tool_call,
hooks_on_turn_end: self.hooks_on_turn_end,
hooks_on_abort: self.hooks_on_abort,
system_prompt: self.system_prompt,
history: self.history,
locked_prefix_len,
turn_count: self.turn_count,
turn_notifiers: self.turn_notifiers,
request_config: self.request_config,
cancellation_token: self.cancellation_token,
_state: PhantomData,
}
}
@@ -1032,13 +1137,18 @@ impl<C: LlmClient> Worker<C, Locked> {
text_block_collector: self.text_block_collector,
tool_call_collector: self.tool_call_collector,
tools: self.tools,
hooks: self.hooks,
hooks_on_message_send: self.hooks_on_message_send,
hooks_before_tool_call: self.hooks_before_tool_call,
hooks_after_tool_call: self.hooks_after_tool_call,
hooks_on_turn_end: self.hooks_on_turn_end,
hooks_on_abort: self.hooks_on_abort,
system_prompt: self.system_prompt,
history: self.history,
locked_prefix_len: 0,
turn_count: self.turn_count,
turn_notifiers: self.turn_notifiers,
request_config: self.request_config,
cancellation_token: self.cancellation_token,
_state: PhantomData,
}
}
llm-worker/tests/parallel_execution_test.rs
+14 -14
View File
@@ -8,7 +8,10 @@ use std::time::{Duration, Instant};
use async_trait::async_trait;
use llm_worker::Worker;
use llm_worker::hook::{ControlFlow, HookError, ToolCall, ToolResult, WorkerHook};
use llm_worker::hook::{
AfterToolCall, AfterToolCallResult, BeforeToolCall, BeforeToolCallResult, Hook, HookError,
ToolCall, ToolResult,
};
use llm_worker::llm_client::event::{Event, ResponseStatus, StatusEvent};
use llm_worker::tool::{Tool, ToolError};
@@ -158,20 +161,17 @@ async fn test_before_tool_call_skip() {
struct BlockingHook;
#[async_trait]
impl WorkerHook for BlockingHook {
async fn before_tool_call(
&self,
tool_call: &mut ToolCall,
) -> Result<ControlFlow, HookError> {
impl Hook<BeforeToolCall> for BlockingHook {
async fn call(&self, tool_call: &mut ToolCall) -> Result<BeforeToolCallResult, HookError> {
if tool_call.name == "blocked_tool" {
Ok(ControlFlow::Skip)
Ok(BeforeToolCallResult::Skip)
} else {
Ok(ControlFlow::Continue)
Ok(BeforeToolCallResult::Continue)
}
}
}
worker.add_hook(BlockingHook);
worker.add_before_tool_call_hook(BlockingHook);
let _result = worker.run("Test hook").await;
@@ -242,19 +242,19 @@ async fn test_after_tool_call_modification() {
}
#[async_trait]
impl WorkerHook for ModifyingHook {
async fn after_tool_call(
impl Hook<AfterToolCall> for ModifyingHook {
async fn call(
&self,
tool_result: &mut ToolResult,
) -> Result<ControlFlow, HookError> {
) -> Result<AfterToolCallResult, HookError> {
tool_result.content = format!("[Modified] {}", tool_result.content);
*self.modified_content.lock().unwrap() = Some(tool_result.content.clone());
Ok(ControlFlow::Continue)
Ok(AfterToolCallResult::Continue)
}
}
let modified_content = Arc::new(std::sync::Mutex::new(None));
worker.add_hook(ModifyingHook {
worker.add_after_tool_call_hook(ModifyingHook {
modified_content: modified_content.clone(),
});