feat: #[tool_registry] and #[tool] macros
This commit is contained in:
@@ -16,4 +16,5 @@ worker-macros = { path = "../worker-macros" }
|
||||
worker-types = { path = "../worker-types" }
|
||||
|
||||
[dev-dependencies]
|
||||
schemars = "1.2.0"
|
||||
tempfile = "3.24.0"
|
||||
|
||||
@@ -7,10 +7,12 @@
|
||||
//! - 型消去されたHandler実装
|
||||
|
||||
pub mod llm_client;
|
||||
mod text_block_collector;
|
||||
mod timeline;
|
||||
mod tool_call_collector;
|
||||
mod worker;
|
||||
|
||||
pub use text_block_collector::TextBlockCollector;
|
||||
pub use timeline::*;
|
||||
pub use tool_call_collector::ToolCallCollector;
|
||||
pub use worker::*;
|
||||
|
||||
@@ -0,0 +1,131 @@
|
||||
//! TextBlockCollector - テキストブロック収集用ハンドラ
|
||||
//!
|
||||
//! TimelineのTextBlockHandler として登録され、
|
||||
//! ストリーム中のテキストブロックを収集する。
|
||||
|
||||
use std::sync::{Arc, Mutex};
|
||||
use worker_types::{Handler, TextBlockEvent, TextBlockKind};
|
||||
|
||||
/// 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;
|
||||
use worker_types::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");
|
||||
}
|
||||
}
|
||||
+103
-50
@@ -8,10 +8,12 @@ use std::sync::Arc;
|
||||
use futures::StreamExt;
|
||||
|
||||
use crate::llm_client::{ClientError, LlmClient, Request, ToolDefinition};
|
||||
use crate::text_block_collector::TextBlockCollector;
|
||||
use crate::tool_call_collector::ToolCallCollector;
|
||||
use crate::Timeline;
|
||||
use worker_types::{
|
||||
ControlFlow, HookError, Message, Tool, ToolCall, ToolError, ToolResult, TurnResult, WorkerHook,
|
||||
ContentPart, ControlFlow, HookError, Message, MessageContent, Tool, ToolCall, ToolError,
|
||||
ToolResult, TurnResult, WorkerHook,
|
||||
};
|
||||
|
||||
// =============================================================================
|
||||
@@ -40,16 +42,10 @@ pub enum WorkerError {
|
||||
// =============================================================================
|
||||
|
||||
/// Worker設定
|
||||
#[derive(Debug, Clone)]
|
||||
#[derive(Debug, Clone, Default)]
|
||||
pub struct WorkerConfig {
|
||||
/// 最大ターン数(無限ループ防止)
|
||||
pub max_turns: usize,
|
||||
}
|
||||
|
||||
impl Default for WorkerConfig {
|
||||
fn default() -> Self {
|
||||
Self { max_turns: 10 }
|
||||
}
|
||||
// 将来の拡張用(現在は空)
|
||||
_private: (),
|
||||
}
|
||||
|
||||
// =============================================================================
|
||||
@@ -68,38 +64,40 @@ pub struct Worker<C: LlmClient> {
|
||||
client: C,
|
||||
/// イベントタイムライン
|
||||
timeline: Timeline,
|
||||
/// ツールコレクター(Timeline用ハンドラ)
|
||||
/// テキストブロックコレクター(Timeline用ハンドラ)
|
||||
text_block_collector: TextBlockCollector,
|
||||
/// ツールコールコレクター(Timeline用ハンドラ)
|
||||
tool_call_collector: ToolCallCollector,
|
||||
/// 登録されたツール
|
||||
tools: HashMap<String, Arc<dyn Tool>>,
|
||||
/// 登録されたHook
|
||||
hooks: Vec<Box<dyn WorkerHook>>,
|
||||
/// 設定
|
||||
config: WorkerConfig,
|
||||
}
|
||||
|
||||
impl<C: LlmClient> Worker<C> {
|
||||
/// 新しいWorkerを作成
|
||||
pub fn new(client: C) -> Self {
|
||||
let text_block_collector = TextBlockCollector::new();
|
||||
let tool_call_collector = ToolCallCollector::new();
|
||||
let mut timeline = Timeline::new();
|
||||
|
||||
// ToolCallCollectorをTimelineに登録
|
||||
// コレクターをTimelineに登録
|
||||
timeline.on_text_block(text_block_collector.clone());
|
||||
timeline.on_tool_use_block(tool_call_collector.clone());
|
||||
|
||||
Self {
|
||||
client,
|
||||
timeline,
|
||||
text_block_collector,
|
||||
tool_call_collector,
|
||||
tools: HashMap::new(),
|
||||
hooks: Vec::new(),
|
||||
config: WorkerConfig::default(),
|
||||
}
|
||||
}
|
||||
|
||||
/// 設定を適用
|
||||
pub fn config(mut self, config: WorkerConfig) -> Self {
|
||||
self.config = config;
|
||||
/// 設定を適用(将来の拡張用)
|
||||
#[allow(dead_code)]
|
||||
pub fn config(self, _config: WorkerConfig) -> Self {
|
||||
self
|
||||
}
|
||||
|
||||
@@ -146,7 +144,7 @@ impl<C: LlmClient> Worker<C> {
|
||||
let mut context = messages;
|
||||
let tool_definitions = self.build_tool_definitions();
|
||||
|
||||
for _turn in 0..self.config.max_turns {
|
||||
loop {
|
||||
// Hook: on_message_send
|
||||
let control = self.run_on_message_send_hooks(&mut context).await?;
|
||||
if let ControlFlow::Abort(reason) = control {
|
||||
@@ -163,9 +161,16 @@ impl<C: LlmClient> Worker<C> {
|
||||
self.timeline.dispatch(&event);
|
||||
}
|
||||
|
||||
// ツール呼び出しの収集結果を取得
|
||||
// 収集結果を取得
|
||||
let text_blocks = self.text_block_collector.take_collected();
|
||||
let tool_calls = self.tool_call_collector.take_collected();
|
||||
|
||||
// アシスタントメッセージをコンテキストに追加
|
||||
let assistant_message = self.build_assistant_message(&text_blocks, &tool_calls);
|
||||
if let Some(msg) = assistant_message {
|
||||
context.push(msg);
|
||||
}
|
||||
|
||||
if tool_calls.is_empty() {
|
||||
// ツール呼び出しなし → ターン終了判定
|
||||
let turn_result = self.run_on_turn_end_hooks(&context).await?;
|
||||
@@ -188,12 +193,48 @@ impl<C: LlmClient> Worker<C> {
|
||||
context.push(Message::tool_result(&result.tool_use_id, &result.content));
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// 最大ターン数到達
|
||||
Err(WorkerError::Aborted(format!(
|
||||
"Maximum turns ({}) reached",
|
||||
self.config.max_turns
|
||||
)))
|
||||
/// テキストブロックとツール呼び出しからアシスタントメッセージを構築
|
||||
fn build_assistant_message(
|
||||
&self,
|
||||
text_blocks: &[String],
|
||||
tool_calls: &[ToolCall],
|
||||
) -> Option<Message> {
|
||||
// テキストもツール呼び出しもない場合はNone
|
||||
if text_blocks.is_empty() && tool_calls.is_empty() {
|
||||
return None;
|
||||
}
|
||||
|
||||
// テキストのみの場合はシンプルなテキストメッセージ
|
||||
if tool_calls.is_empty() {
|
||||
let text = text_blocks.join("");
|
||||
return Some(Message::assistant(text));
|
||||
}
|
||||
|
||||
// ツール呼び出しがある場合は Parts として構築
|
||||
let mut parts = Vec::new();
|
||||
|
||||
// テキストパーツを追加
|
||||
for text in text_blocks {
|
||||
if !text.is_empty() {
|
||||
parts.push(ContentPart::Text { text: text.clone() });
|
||||
}
|
||||
}
|
||||
|
||||
// ツール呼び出しパーツを追加
|
||||
for call in tool_calls {
|
||||
parts.push(ContentPart::ToolUse {
|
||||
id: call.id.clone(),
|
||||
name: call.name.clone(),
|
||||
input: call.input.clone(),
|
||||
});
|
||||
}
|
||||
|
||||
Some(Message {
|
||||
role: worker_types::Role::Assistant,
|
||||
content: MessageContent::Parts(parts),
|
||||
})
|
||||
}
|
||||
|
||||
/// リクエストを構築
|
||||
@@ -291,16 +332,18 @@ impl<C: LlmClient> Worker<C> {
|
||||
}
|
||||
|
||||
/// ツールを並列実行
|
||||
///
|
||||
/// 全てのツールに対してbefore_tool_callフックを実行後、
|
||||
/// 許可されたツールを並列に実行し、結果にafter_tool_callフックを適用する。
|
||||
async fn execute_tools(
|
||||
&self,
|
||||
mut tool_calls: Vec<ToolCall>,
|
||||
tool_calls: Vec<ToolCall>,
|
||||
) -> Result<Vec<ToolResult>, WorkerError> {
|
||||
let mut results = Vec::new();
|
||||
use futures::future::join_all;
|
||||
|
||||
// TODO: 将来的には join_all で並列実行
|
||||
// 現在は逐次実行
|
||||
for mut tool_call in tool_calls.drain(..) {
|
||||
// Hook: before_tool_call
|
||||
// Phase 1: before_tool_call フックを適用(スキップ/中断を判定)
|
||||
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?;
|
||||
@@ -315,28 +358,40 @@ impl<C: LlmClient> Worker<C> {
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
if skip {
|
||||
continue;
|
||||
if !skip {
|
||||
approved_calls.push(tool_call);
|
||||
}
|
||||
}
|
||||
|
||||
// ツール実行
|
||||
let mut tool_result = if let Some(tool) = self.tools.get(&tool_call.name) {
|
||||
let input_json = serde_json::to_string(&tool_call.input).unwrap_or_default();
|
||||
match tool.execute(&input_json).await {
|
||||
Ok(content) => ToolResult::success(&tool_call.id, content),
|
||||
Err(e) => ToolResult::error(&tool_call.id, e.to_string()),
|
||||
// Phase 2: 許可されたツールを並列実行
|
||||
let futures: Vec<_> = approved_calls
|
||||
.into_iter()
|
||||
.map(|tool_call| {
|
||||
let tools = &self.tools;
|
||||
async move {
|
||||
if let Some(tool) = tools.get(&tool_call.name) {
|
||||
let input_json =
|
||||
serde_json::to_string(&tool_call.input).unwrap_or_default();
|
||||
match tool.execute(&input_json).await {
|
||||
Ok(content) => ToolResult::success(&tool_call.id, content),
|
||||
Err(e) => ToolResult::error(&tool_call.id, e.to_string()),
|
||||
}
|
||||
} else {
|
||||
ToolResult::error(
|
||||
&tool_call.id,
|
||||
format!("Tool '{}' not found", tool_call.name),
|
||||
)
|
||||
}
|
||||
}
|
||||
} else {
|
||||
ToolResult::error(
|
||||
&tool_call.id,
|
||||
format!("Tool '{}' not found", tool_call.name),
|
||||
)
|
||||
};
|
||||
})
|
||||
.collect();
|
||||
|
||||
// Hook: after_tool_call
|
||||
let mut results = join_all(futures).await;
|
||||
|
||||
// Phase 3: after_tool_call フックを適用
|
||||
for tool_result in &mut results {
|
||||
for hook in &self.hooks {
|
||||
let result = hook.after_tool_call(&mut tool_result).await?;
|
||||
let result = hook.after_tool_call(tool_result).await?;
|
||||
match result {
|
||||
ControlFlow::Continue => {}
|
||||
ControlFlow::Skip => break,
|
||||
@@ -345,8 +400,6 @@ impl<C: LlmClient> Worker<C> {
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
results.push(tool_result);
|
||||
}
|
||||
|
||||
Ok(results)
|
||||
|
||||
@@ -204,26 +204,81 @@ impl EventPlayer {
|
||||
///
|
||||
/// 事前に定義されたイベントシーケンスをストリームとして返す。
|
||||
/// fixtureファイルからロードすることも、直接イベントを渡すこともできる。
|
||||
///
|
||||
/// # 複数リクエスト対応
|
||||
///
|
||||
/// `with_responses()`を使用して、複数回のリクエストに対して異なるレスポンスを設定できる。
|
||||
/// リクエスト回数が設定されたレスポンス数を超えた場合は空のストリームを返す。
|
||||
pub struct MockLlmClient {
|
||||
events: Vec<Event>,
|
||||
/// 各リクエストに対するレスポンス(イベントシーケンス)
|
||||
responses: std::sync::Arc<std::sync::Mutex<Vec<Vec<Event>>>>,
|
||||
/// 現在のリクエストインデックス
|
||||
request_index: std::sync::Arc<std::sync::atomic::AtomicUsize>,
|
||||
}
|
||||
|
||||
#[allow(dead_code)]
|
||||
impl MockLlmClient {
|
||||
/// イベントリストから直接作成
|
||||
/// イベントリストから直接作成(単一レスポンス)
|
||||
///
|
||||
/// すべてのリクエストに対して同じイベントシーケンスを返す(従来の動作)
|
||||
pub fn new(events: Vec<Event>) -> Self {
|
||||
Self { events }
|
||||
Self {
|
||||
responses: std::sync::Arc::new(std::sync::Mutex::new(vec![events])),
|
||||
request_index: std::sync::Arc::new(std::sync::atomic::AtomicUsize::new(0)),
|
||||
}
|
||||
}
|
||||
|
||||
/// fixtureファイルからロード
|
||||
/// 複数のレスポンスを設定
|
||||
///
|
||||
/// 各リクエストに対して順番にイベントシーケンスを返す。
|
||||
/// N回目のリクエストにはN番目のレスポンスが使用される。
|
||||
///
|
||||
/// # Example
|
||||
/// ```ignore
|
||||
/// let client = MockLlmClient::with_responses(vec![
|
||||
/// // 1回目のリクエスト: ツール呼び出し
|
||||
/// vec![Event::tool_use_start(0, "call_1", "my_tool"), ...],
|
||||
/// // 2回目のリクエスト: テキストレスポンス
|
||||
/// vec![Event::text_block_start(0), ...],
|
||||
/// ]);
|
||||
/// ```
|
||||
pub fn with_responses(responses: Vec<Vec<Event>>) -> Self {
|
||||
Self {
|
||||
responses: std::sync::Arc::new(std::sync::Mutex::new(responses)),
|
||||
request_index: std::sync::Arc::new(std::sync::atomic::AtomicUsize::new(0)),
|
||||
}
|
||||
}
|
||||
|
||||
/// fixtureファイルからロード(単一レスポンス)
|
||||
pub fn from_fixture(path: impl AsRef<Path>) -> std::io::Result<Self> {
|
||||
let player = EventPlayer::load(path)?;
|
||||
let events = player.parse_events();
|
||||
Ok(Self { events })
|
||||
Ok(Self::new(events))
|
||||
}
|
||||
|
||||
/// 保持しているイベント数を取得
|
||||
/// 保持しているレスポンス数を取得
|
||||
pub fn response_count(&self) -> usize {
|
||||
self.responses.lock().unwrap().len()
|
||||
}
|
||||
|
||||
/// 最初のレスポンスのイベント数を取得(後方互換性)
|
||||
pub fn event_count(&self) -> usize {
|
||||
self.events.len()
|
||||
self.responses
|
||||
.lock()
|
||||
.unwrap()
|
||||
.first()
|
||||
.map(|v| v.len())
|
||||
.unwrap_or(0)
|
||||
}
|
||||
|
||||
/// 現在のリクエストインデックスを取得
|
||||
pub fn current_request_index(&self) -> usize {
|
||||
self.request_index.load(std::sync::atomic::Ordering::SeqCst)
|
||||
}
|
||||
|
||||
/// リクエストインデックスをリセット
|
||||
pub fn reset(&self) {
|
||||
self.request_index.store(0, std::sync::atomic::Ordering::SeqCst);
|
||||
}
|
||||
}
|
||||
|
||||
@@ -233,8 +288,20 @@ impl LlmClient for MockLlmClient {
|
||||
&self,
|
||||
_request: Request,
|
||||
) -> Result<Pin<Box<dyn Stream<Item = Result<Event, ClientError>> + Send>>, ClientError> {
|
||||
let events = self.events.clone();
|
||||
let index = self.request_index.fetch_add(1, std::sync::atomic::Ordering::SeqCst);
|
||||
|
||||
let events = {
|
||||
let responses = self.responses.lock().unwrap();
|
||||
if index < responses.len() {
|
||||
responses[index].clone()
|
||||
} else {
|
||||
// レスポンスが尽きた場合は空のストリーム
|
||||
Vec::new()
|
||||
}
|
||||
};
|
||||
|
||||
let stream = futures::stream::iter(events.into_iter().map(Ok));
|
||||
Ok(Box::pin(stream))
|
||||
}
|
||||
}
|
||||
|
||||
|
||||
@@ -0,0 +1,254 @@
|
||||
//! 並列ツール実行のテスト
|
||||
//!
|
||||
//! Workerが複数のツールを並列に実行することを確認する。
|
||||
|
||||
use std::sync::atomic::{AtomicUsize, Ordering};
|
||||
use std::sync::Arc;
|
||||
use std::time::{Duration, Instant};
|
||||
|
||||
use async_trait::async_trait;
|
||||
use worker::Worker;
|
||||
use worker_types::{Event, Message, ResponseStatus, StatusEvent, Tool, ToolError, ToolResult, ToolCall, ControlFlow, HookError, WorkerHook};
|
||||
|
||||
mod common;
|
||||
use common::MockLlmClient;
|
||||
|
||||
// =============================================================================
|
||||
// Parallel Execution Test Tools
|
||||
// =============================================================================
|
||||
|
||||
/// 一定時間待機してから応答するツール
|
||||
#[derive(Clone)]
|
||||
struct SlowTool {
|
||||
name: String,
|
||||
delay_ms: u64,
|
||||
call_count: Arc<AtomicUsize>,
|
||||
}
|
||||
|
||||
impl SlowTool {
|
||||
fn new(name: impl Into<String>, delay_ms: u64) -> Self {
|
||||
Self {
|
||||
name: name.into(),
|
||||
delay_ms,
|
||||
call_count: Arc::new(AtomicUsize::new(0)),
|
||||
}
|
||||
}
|
||||
|
||||
fn call_count(&self) -> usize {
|
||||
self.call_count.load(Ordering::SeqCst)
|
||||
}
|
||||
}
|
||||
|
||||
#[async_trait]
|
||||
impl Tool for SlowTool {
|
||||
fn name(&self) -> &str {
|
||||
&self.name
|
||||
}
|
||||
|
||||
fn description(&self) -> &str {
|
||||
"A tool that waits before responding"
|
||||
}
|
||||
|
||||
fn input_schema(&self) -> serde_json::Value {
|
||||
serde_json::json!({
|
||||
"type": "object",
|
||||
"properties": {}
|
||||
})
|
||||
}
|
||||
|
||||
async fn execute(&self, _input_json: &str) -> Result<String, ToolError> {
|
||||
self.call_count.fetch_add(1, Ordering::SeqCst);
|
||||
tokio::time::sleep(Duration::from_millis(self.delay_ms)).await;
|
||||
Ok(format!("Completed after {}ms", self.delay_ms))
|
||||
}
|
||||
}
|
||||
|
||||
// =============================================================================
|
||||
// Tests
|
||||
// =============================================================================
|
||||
|
||||
/// 複数のツールが並列に実行されることを確認
|
||||
///
|
||||
/// 各ツールが100msかかる場合、逐次実行なら300ms以上かかるが、
|
||||
/// 並列実行なら100ms程度で完了するはず。
|
||||
#[tokio::test]
|
||||
async fn test_parallel_tool_execution() {
|
||||
// 3つのツール呼び出しを含むイベントシーケンス
|
||||
let events = vec![
|
||||
Event::tool_use_start(0, "call_1", "slow_tool_1"),
|
||||
Event::tool_input_delta(0, r#"{}"#),
|
||||
Event::tool_use_stop(0),
|
||||
Event::tool_use_start(1, "call_2", "slow_tool_2"),
|
||||
Event::tool_input_delta(1, r#"{}"#),
|
||||
Event::tool_use_stop(1),
|
||||
Event::tool_use_start(2, "call_3", "slow_tool_3"),
|
||||
Event::tool_input_delta(2, r#"{}"#),
|
||||
Event::tool_use_stop(2),
|
||||
Event::Status(StatusEvent {
|
||||
status: ResponseStatus::Completed,
|
||||
}),
|
||||
];
|
||||
|
||||
let client = MockLlmClient::new(events);
|
||||
let mut worker = Worker::new(client);
|
||||
|
||||
// 各ツールは100ms待機
|
||||
let tool1 = SlowTool::new("slow_tool_1", 100);
|
||||
let tool2 = SlowTool::new("slow_tool_2", 100);
|
||||
let tool3 = SlowTool::new("slow_tool_3", 100);
|
||||
|
||||
let tool1_clone = tool1.clone();
|
||||
let tool2_clone = tool2.clone();
|
||||
let tool3_clone = tool3.clone();
|
||||
|
||||
worker.register_tool(tool1);
|
||||
worker.register_tool(tool2);
|
||||
worker.register_tool(tool3);
|
||||
|
||||
|
||||
|
||||
let messages = vec![Message::user("Run all tools")];
|
||||
|
||||
let start = Instant::now();
|
||||
let _result = worker.run(messages).await;
|
||||
let elapsed = start.elapsed();
|
||||
|
||||
// 全ツールが呼び出されたことを確認
|
||||
assert_eq!(tool1_clone.call_count(), 1, "Tool 1 should be called once");
|
||||
assert_eq!(tool2_clone.call_count(), 1, "Tool 2 should be called once");
|
||||
assert_eq!(tool3_clone.call_count(), 1, "Tool 3 should be called once");
|
||||
|
||||
// 並列実行なら200ms以下で完了するはず(逐次なら300ms以上)
|
||||
// マージン込みで250msをしきい値とする
|
||||
assert!(
|
||||
elapsed < Duration::from_millis(250),
|
||||
"Parallel execution should complete in ~100ms, but took {:?}",
|
||||
elapsed
|
||||
);
|
||||
|
||||
println!("Parallel execution completed in {:?}", elapsed);
|
||||
}
|
||||
|
||||
/// Hook: before_tool_call でスキップされたツールは実行されないことを確認
|
||||
#[tokio::test]
|
||||
async fn test_before_tool_call_skip() {
|
||||
let events = vec![
|
||||
Event::tool_use_start(0, "call_1", "allowed_tool"),
|
||||
Event::tool_input_delta(0, r#"{}"#),
|
||||
Event::tool_use_stop(0),
|
||||
Event::tool_use_start(1, "call_2", "blocked_tool"),
|
||||
Event::tool_input_delta(1, r#"{}"#),
|
||||
Event::tool_use_stop(1),
|
||||
Event::Status(StatusEvent {
|
||||
status: ResponseStatus::Completed,
|
||||
}),
|
||||
];
|
||||
|
||||
let client = MockLlmClient::new(events);
|
||||
let mut worker = Worker::new(client);
|
||||
|
||||
let allowed_tool = SlowTool::new("allowed_tool", 10);
|
||||
let blocked_tool = SlowTool::new("blocked_tool", 10);
|
||||
|
||||
let allowed_clone = allowed_tool.clone();
|
||||
let blocked_clone = blocked_tool.clone();
|
||||
|
||||
worker.register_tool(allowed_tool);
|
||||
worker.register_tool(blocked_tool);
|
||||
|
||||
// "blocked_tool" をスキップするHook
|
||||
struct BlockingHook;
|
||||
|
||||
#[async_trait]
|
||||
impl WorkerHook for BlockingHook {
|
||||
async fn before_tool_call(&self, tool_call: &mut ToolCall) -> Result<ControlFlow, HookError> {
|
||||
if tool_call.name == "blocked_tool" {
|
||||
Ok(ControlFlow::Skip)
|
||||
} else {
|
||||
Ok(ControlFlow::Continue)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
worker.add_hook(BlockingHook);
|
||||
|
||||
let messages = vec![Message::user("Test hook")];
|
||||
let _result = worker.run(messages).await;
|
||||
|
||||
// allowed_tool は呼び出されるが、blocked_tool は呼び出されない
|
||||
assert_eq!(allowed_clone.call_count(), 1, "Allowed tool should be called");
|
||||
assert_eq!(blocked_clone.call_count(), 0, "Blocked tool should not be called");
|
||||
}
|
||||
|
||||
/// Hook: after_tool_call で結果が改変されることを確認
|
||||
#[tokio::test]
|
||||
async fn test_after_tool_call_modification() {
|
||||
// 複数リクエストに対応するレスポンスを準備
|
||||
let client = MockLlmClient::with_responses(vec![
|
||||
// 1回目のリクエスト: ツール呼び出し
|
||||
vec![
|
||||
Event::tool_use_start(0, "call_1", "test_tool"),
|
||||
Event::tool_input_delta(0, r#"{}"#),
|
||||
Event::tool_use_stop(0),
|
||||
Event::Status(StatusEvent {
|
||||
status: ResponseStatus::Completed,
|
||||
}),
|
||||
],
|
||||
// 2回目のリクエスト: ツール結果を受けてテキストレスポンス
|
||||
vec![
|
||||
Event::text_block_start(0),
|
||||
Event::text_delta(0, "Done!"),
|
||||
Event::text_block_stop(0, None),
|
||||
Event::Status(StatusEvent {
|
||||
status: ResponseStatus::Completed,
|
||||
}),
|
||||
],
|
||||
]);
|
||||
|
||||
let mut worker = Worker::new(client);
|
||||
|
||||
#[derive(Clone)]
|
||||
struct SimpleTool;
|
||||
|
||||
#[async_trait]
|
||||
impl Tool for SimpleTool {
|
||||
fn name(&self) -> &str { "test_tool" }
|
||||
fn description(&self) -> &str { "Test" }
|
||||
fn input_schema(&self) -> serde_json::Value { serde_json::json!({}) }
|
||||
async fn execute(&self, _: &str) -> Result<String, ToolError> {
|
||||
Ok("Original Result".to_string())
|
||||
}
|
||||
}
|
||||
|
||||
worker.register_tool(SimpleTool);
|
||||
|
||||
// 結果を改変するHook
|
||||
struct ModifyingHook {
|
||||
modified_content: Arc<std::sync::Mutex<Option<String>>>,
|
||||
}
|
||||
|
||||
#[async_trait]
|
||||
impl WorkerHook for ModifyingHook {
|
||||
async fn after_tool_call(&self, tool_result: &mut ToolResult) -> Result<ControlFlow, HookError> {
|
||||
tool_result.content = format!("[Modified] {}", tool_result.content);
|
||||
*self.modified_content.lock().unwrap() = Some(tool_result.content.clone());
|
||||
Ok(ControlFlow::Continue)
|
||||
}
|
||||
}
|
||||
|
||||
let modified_content = Arc::new(std::sync::Mutex::new(None));
|
||||
worker.add_hook(ModifyingHook { modified_content: modified_content.clone() });
|
||||
|
||||
let messages = vec![Message::user("Test modification")];
|
||||
let result = worker.run(messages).await;
|
||||
|
||||
assert!(result.is_ok(), "Worker should complete: {:?}", result);
|
||||
|
||||
// Hookが呼ばれて内容が改変されたことを確認
|
||||
let content = modified_content.lock().unwrap().clone();
|
||||
assert!(content.is_some(), "Hook should have been called");
|
||||
assert!(
|
||||
content.unwrap().contains("[Modified]"),
|
||||
"Result should be modified"
|
||||
);
|
||||
}
|
||||
@@ -0,0 +1,212 @@
|
||||
//! ツールマクロのテスト
|
||||
//!
|
||||
//! `#[tool_registry]` と `#[tool]` マクロの動作を確認する。
|
||||
|
||||
use std::sync::atomic::{AtomicUsize, Ordering};
|
||||
use std::sync::Arc;
|
||||
|
||||
// マクロ展開に必要なインポート
|
||||
use schemars;
|
||||
use serde;
|
||||
|
||||
use worker_macros::tool_registry;
|
||||
use worker_types::Tool;
|
||||
|
||||
// =============================================================================
|
||||
// Test: Basic Tool Generation
|
||||
// =============================================================================
|
||||
|
||||
/// シンプルなコンテキスト構造体
|
||||
#[derive(Clone)]
|
||||
struct SimpleContext {
|
||||
prefix: String,
|
||||
}
|
||||
|
||||
#[tool_registry]
|
||||
impl SimpleContext {
|
||||
/// メッセージに挨拶を追加する
|
||||
///
|
||||
/// 指定されたメッセージにプレフィックスを付けて返します。
|
||||
#[tool]
|
||||
async fn greet(&self, message: String) -> String {
|
||||
format!("{}: {}", self.prefix, message)
|
||||
}
|
||||
|
||||
/// 二つの数を足す
|
||||
#[tool]
|
||||
async fn add(&self, a: i32, b: i32) -> i32 {
|
||||
a + b
|
||||
}
|
||||
|
||||
/// 引数なしのツール
|
||||
#[tool]
|
||||
async fn get_prefix(&self) -> String {
|
||||
self.prefix.clone()
|
||||
}
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_basic_tool_generation() {
|
||||
let ctx = SimpleContext {
|
||||
prefix: "Hello".to_string(),
|
||||
};
|
||||
|
||||
// ファクトリメソッドでツールを取得
|
||||
let greet_tool = ctx.greet_tool();
|
||||
|
||||
// 名前の確認
|
||||
assert_eq!(greet_tool.name(), "greet");
|
||||
|
||||
// 説明の確認(docコメントから取得)
|
||||
let desc = greet_tool.description();
|
||||
assert!(desc.contains("メッセージに挨拶を追加する"), "Description should contain doc comment: {}", desc);
|
||||
|
||||
// スキーマの確認
|
||||
let schema = greet_tool.input_schema();
|
||||
println!("Schema: {}", serde_json::to_string_pretty(&schema).unwrap());
|
||||
assert!(schema.get("properties").is_some(), "Schema should have properties");
|
||||
|
||||
// 実行テスト
|
||||
let result = greet_tool.execute(r#"{"message": "World"}"#).await;
|
||||
assert!(result.is_ok(), "Should execute successfully");
|
||||
let output = result.unwrap();
|
||||
assert!(output.contains("Hello"), "Output should contain prefix");
|
||||
assert!(output.contains("World"), "Output should contain message");
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_multiple_arguments() {
|
||||
let ctx = SimpleContext {
|
||||
prefix: "".to_string(),
|
||||
};
|
||||
|
||||
let add_tool = ctx.add_tool();
|
||||
|
||||
assert_eq!(add_tool.name(), "add");
|
||||
|
||||
let result = add_tool.execute(r#"{"a": 10, "b": 20}"#).await;
|
||||
assert!(result.is_ok());
|
||||
let output = result.unwrap();
|
||||
assert!(output.contains("30"), "Should contain sum: {}", output);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_no_arguments() {
|
||||
let ctx = SimpleContext {
|
||||
prefix: "TestPrefix".to_string(),
|
||||
};
|
||||
|
||||
let get_prefix_tool = ctx.get_prefix_tool();
|
||||
|
||||
assert_eq!(get_prefix_tool.name(), "get_prefix");
|
||||
|
||||
// 空のJSONオブジェクトで呼び出し
|
||||
let result = get_prefix_tool.execute(r#"{}"#).await;
|
||||
assert!(result.is_ok());
|
||||
let output = result.unwrap();
|
||||
assert!(output.contains("TestPrefix"), "Should contain prefix: {}", output);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_invalid_arguments() {
|
||||
let ctx = SimpleContext {
|
||||
prefix: "".to_string(),
|
||||
};
|
||||
|
||||
let greet_tool = ctx.greet_tool();
|
||||
|
||||
// 不正なJSON
|
||||
let result = greet_tool.execute(r#"{"wrong_field": "value"}"#).await;
|
||||
assert!(result.is_err(), "Should fail with invalid arguments");
|
||||
}
|
||||
|
||||
// =============================================================================
|
||||
// Test: Result Return Type
|
||||
// =============================================================================
|
||||
|
||||
#[derive(Clone)]
|
||||
struct FallibleContext;
|
||||
|
||||
#[derive(Debug)]
|
||||
struct MyError(String);
|
||||
|
||||
impl std::fmt::Display for MyError {
|
||||
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
|
||||
write!(f, "{}", self.0)
|
||||
}
|
||||
}
|
||||
|
||||
#[tool_registry]
|
||||
impl FallibleContext {
|
||||
/// 与えられた値を検証する
|
||||
#[tool]
|
||||
async fn validate(&self, value: i32) -> Result<String, MyError> {
|
||||
if value > 0 {
|
||||
Ok(format!("Valid: {}", value))
|
||||
} else {
|
||||
Err(MyError("Value must be positive".to_string()))
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_result_return_type_success() {
|
||||
let ctx = FallibleContext;
|
||||
let validate_tool = ctx.validate_tool();
|
||||
|
||||
let result = validate_tool.execute(r#"{"value": 42}"#).await;
|
||||
assert!(result.is_ok(), "Should succeed for positive value");
|
||||
let output = result.unwrap();
|
||||
assert!(output.contains("Valid"), "Should contain Valid: {}", output);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_result_return_type_error() {
|
||||
let ctx = FallibleContext;
|
||||
let validate_tool = ctx.validate_tool();
|
||||
|
||||
let result = validate_tool.execute(r#"{"value": -1}"#).await;
|
||||
assert!(result.is_err(), "Should fail for negative value");
|
||||
|
||||
let err = result.unwrap_err();
|
||||
assert!(err.to_string().contains("positive"), "Error should mention positive: {}", err);
|
||||
}
|
||||
|
||||
// =============================================================================
|
||||
// Test: Synchronous Methods
|
||||
// =============================================================================
|
||||
|
||||
#[derive(Clone)]
|
||||
struct SyncContext {
|
||||
counter: Arc<AtomicUsize>,
|
||||
}
|
||||
|
||||
#[tool_registry]
|
||||
impl SyncContext {
|
||||
/// カウンターをインクリメントして返す (非async)
|
||||
#[tool]
|
||||
fn increment(&self) -> usize {
|
||||
self.counter.fetch_add(1, Ordering::SeqCst) + 1
|
||||
}
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_sync_method() {
|
||||
let ctx = SyncContext {
|
||||
counter: Arc::new(AtomicUsize::new(0)),
|
||||
};
|
||||
|
||||
let increment_tool = ctx.increment_tool();
|
||||
|
||||
// 3回実行
|
||||
let result1 = increment_tool.execute(r#"{}"#).await;
|
||||
let result2 = increment_tool.execute(r#"{}"#).await;
|
||||
let result3 = increment_tool.execute(r#"{}"#).await;
|
||||
|
||||
assert!(result1.is_ok());
|
||||
assert!(result2.is_ok());
|
||||
assert!(result3.is_ok());
|
||||
|
||||
// カウンターは3になっているはず
|
||||
assert_eq!(ctx.counter.load(Ordering::SeqCst), 3);
|
||||
}
|
||||
@@ -11,7 +11,7 @@ use std::sync::Arc;
|
||||
|
||||
use async_trait::async_trait;
|
||||
use common::MockLlmClient;
|
||||
use worker::{Worker, WorkerConfig};
|
||||
use worker::Worker;
|
||||
use worker_types::{Tool, ToolError};
|
||||
|
||||
/// フィクスチャディレクトリのパス
|
||||
@@ -163,8 +163,7 @@ async fn test_worker_tool_call() {
|
||||
let tool_for_check = weather_tool.clone();
|
||||
worker.register_tool(weather_tool);
|
||||
|
||||
// 設定: ツール実行後はターン終了(ループしない)
|
||||
worker = worker.config(WorkerConfig { max_turns: 1 });
|
||||
|
||||
|
||||
// メッセージを送信
|
||||
let messages = vec![worker_types::Message::user("What's the weather in Tokyo?")];
|
||||
|
||||
Reference in New Issue
Block a user