feat: Redesign the tool system

This commit is contained in:
2026-01-10 00:31:14 +09:00
parent 5691b09fc8
commit 16fda38039
13 changed files with 897 additions and 396 deletions
+132 -74
View File
@@ -14,13 +14,14 @@ HookはWorker層でのターン制御に介入するためのメカニズムで
## Hook一覧
| Hook | タイミング | 主な用途 | 戻り値 |
| ------------------ | -------------------------- | --------------------- | ---------------------- |
| `on_message_send` | LLM送信前 | コンテキスト改変/検証 | `OnMessageSendResult` |
| `before_tool_call` | ツール実行前 | 実行許可/引数改変 | `BeforeToolCallResult` |
| `after_tool_call` | ツール実行 | 結果加工/マスキング | `AfterToolCallResult` |
| `on_turn_end` | ツールなしでターン終了直前 | 検証/リトライ指示 | `OnTurnEndResult` |
| `on_abort` | 中断時 | クリーンアップ/通知 | `()` |
| Hook | タイミング | 主な用途 | 戻り値 |
| ------------------ | -------------------------- | -------------------------- | ---------------------- |
| `on_prompt_submit` | `run()` 呼び出し時 | ユーザーメッセージの前処理 | `OnPromptSubmitResult` |
| `pre_llm_request` | 各ターンのLLM送信前 | コンテキスト改変/検証 | `PreLlmRequestResult` |
| `pre_tool_call` | ツール実行 | 実行許可/引数改変 | `PreToolCallResult` |
| `post_tool_call` | ツール実行後 | 結果加工/マスキング | `PostToolCallResult` |
| `on_turn_end` | ツールなしでターン終了直前 | 検証/リトライ指示 | `OnTurnEndResult` |
| `on_abort` | 中断時 | クリーンアップ/通知 | `()` |
## Hook Trait
@@ -43,25 +44,31 @@ pub trait HookEventKind {
type Output;
}
pub struct OnMessageSend;
pub struct BeforeToolCall;
pub struct AfterToolCall;
pub struct OnPromptSubmit;
pub struct PreLlmRequest;
pub struct PreToolCall;
pub struct PostToolCall;
pub struct OnTurnEnd;
pub struct OnAbort;
pub enum OnMessageSendResult {
pub enum OnPromptSubmitResult {
Continue,
Cancel(String),
}
pub enum BeforeToolCallResult {
pub enum PreLlmRequestResult {
Continue,
Cancel(String),
}
pub enum PreToolCallResult {
Continue,
Skip,
Abort(String),
Pause,
}
pub enum AfterToolCallResult {
pub enum PostToolCallResult {
Continue,
Abort(String),
}
@@ -75,7 +82,7 @@ pub enum OnTurnEndResult {
### Tool Call Context
`before_tool_call` / `after_tool_call` は、ツール実行の文脈を含む入力を受け取る。
`pre_tool_call` / `post_tool_call` は、ツール実行の文脈を含む入力を受け取る。
```rust
pub struct ToolCallContext {
@@ -84,7 +91,8 @@ pub struct ToolCallContext {
pub tool: Arc<dyn Tool>, // 状態アクセス用
}
pub struct ToolResultContext {
pub struct PostToolCallContext {
pub call: ToolCall,
pub result: ToolResult,
pub meta: ToolMeta,
pub tool: Arc<dyn Tool>,
@@ -94,40 +102,84 @@ pub struct ToolResultContext {
## 呼び出しタイミング
```
Worker::run() ループ
Worker::run(user_input)
├─▶ on_message_send ──────────────────────────────┐
コンテキストの改変、バリデーション、
システムプロンプト注入などが可能
─▶ LLMリクエスト送信 & ストリーム処理 │
├─▶ ツール呼び出しがある場合:
│ │
│ ├─▶ before_tool_call (各ツールごと・逐次)
実行可否の判定、引数の改変
├─▶ ツール並列実行 (join_all)
─▶ after_tool_call (各結果ごと・逐次)
結果の確認、加工、ログ出力
├─▶ ツール結果をコンテキストに追加 → ループ先頭へ
└─▶ ツールなしの場合:
└─▶ on_turn_end ─────────────────────────────┘
最終応答のチェック(Lint/Fmt等)
エラーがあればContinueWithMessagesでリトライ
├─▶ on_prompt_submit ───────────────────────────┐
ユーザーメッセージの前処理・検証
(最初の1回のみ)
│ │
─▶ loop {
├─▶ pre_llm_request ──────────────────────
│ コンテキストの改変、バリデーション、
│ システムプロンプト注入などが可能
(毎ターン実行)
│ │
├─▶ LLMリクエスト送信 & ストリーム処理
│ │
─▶ ツール呼び出しがある場合:
├─▶ pre_tool_call (各ツールごと・逐次)
│ │ 実行可否の判定、引数の改変
├─▶ ツール並列実行 (join_all)
│ └─▶ post_tool_call (各結果ごと・逐次) │
│ 結果の確認、加工、ログ出力 │
│ │
├─▶ ツール結果をコンテキストに追加 │
│ → ループ先頭へ │
│ │
└─▶ ツールなしの場合: │
│ │
└─▶ on_turn_end ───────────────────┘
最終応答のチェック(Lint/Fmt等)
エラーがあればContinueWithMessagesでリトライ
}
※ 中断時は on_abort が呼ばれる
```
## 各Hookの詳細
### on_message_send
### on_prompt_submit
**呼び出しタイミング**: LLMへリクエスト送信前(ターンループの冒頭)
**呼び出しタイミング**: `run()`
でユーザーメッセージを受け取った直後(最初の1回のみ)
**用途**:
- ユーザー入力のバリデーション
- 入力のサニタイズ・フィルタリング
- ログ出力
- `OnPromptSubmitResult::Cancel` による実行キャンセル
**入力**: `&mut Message` - ユーザーメッセージ(改変可能)
**例**: 入力のバリデーション
```rust
struct InputValidator;
#[async_trait]
impl Hook<OnPromptSubmit> for InputValidator {
async fn call(
&self,
message: &mut Message,
) -> Result<OnPromptSubmitResult, HookError> {
if let MessageContent::Text(text) = &message.content {
if text.trim().is_empty() {
return Ok(OnPromptSubmitResult::Cancel("Empty input".to_string()));
}
}
Ok(OnPromptSubmitResult::Continue)
}
}
```
### pre_llm_request
**呼び出しタイミング**: 各ターンのLLMリクエスト送信前(ループの毎回)
**用途**:
@@ -135,7 +187,9 @@ Worker::run() ループ
- メッセージのバリデーション
- 機密情報のフィルタリング
- リクエスト内容のログ出力
- `OnMessageSendResult::Cancel` による送信キャンセル
- `PreLlmRequestResult::Cancel` による送信キャンセル
**入力**: `&mut Vec<Message>` - コンテキスト全体(改変可能)
**例**: メッセージにタイムスタンプを追加
@@ -143,19 +197,19 @@ Worker::run() ループ
struct TimestampHook;
#[async_trait]
impl Hook<OnMessageSend> for TimestampHook {
impl Hook<PreLlmRequest> for TimestampHook {
async fn call(
&self,
context: &mut Vec<Message>,
) -> Result<OnMessageSendResult, HookError> {
) -> Result<PreLlmRequestResult, HookError> {
let timestamp = chrono::Local::now().to_rfc3339();
context.insert(0, Message::user(format!("[{}]", timestamp)));
Ok(OnMessageSendResult::Continue)
Ok(PreLlmRequestResult::Continue)
}
}
```
### before_tool_call
### pre_tool_call
**呼び出しタイミング**: 各ツール実行前(並列実行フェーズの前)
@@ -165,9 +219,10 @@ impl Hook<OnMessageSend> for TimestampHook {
- 引数のサニタイズ
- 確認プロンプトの表示(UIとの連携)
- 実行ログの記録
- `BeforeToolCallResult::Pause` による一時停止
- `PreToolCallResult::Pause` による一時停止
**入力**:
- `ToolCallContext``ToolCall` + `ToolMeta` + `Arc<dyn Tool>`
**例**: 特定ツールをブロック
@@ -178,22 +233,22 @@ struct ToolBlocker {
}
#[async_trait]
impl Hook<BeforeToolCall> for ToolBlocker {
impl Hook<PreToolCall> for ToolBlocker {
async fn call(
&self,
ctx: &mut ToolCallContext,
) -> Result<BeforeToolCallResult, HookError> {
) -> Result<PreToolCallResult, HookError> {
if self.blocked_tools.contains(&ctx.call.name) {
println!("Blocked tool: {}", ctx.call.name);
Ok(BeforeToolCallResult::Skip)
Ok(PreToolCallResult::Skip)
} else {
Ok(BeforeToolCallResult::Continue)
Ok(PreToolCallResult::Continue)
}
}
}
```
### after_tool_call
### post_tool_call
**呼び出しタイミング**: 各ツール実行後(並列実行フェーズの後)
@@ -203,8 +258,11 @@ impl Hook<BeforeToolCall> for ToolBlocker {
- 機密情報のマスキング
- 結果のキャッシュ
- 実行結果のログ出力
**入力**:
- `ToolResultContext``ToolResult` + `ToolMeta` + `Arc<dyn Tool>`
- `PostToolCallContext``ToolCall` + `ToolResult` + `ToolMeta` +
`Arc<dyn Tool>`
**例**: 結果にプレフィックスを追加
@@ -212,15 +270,15 @@ impl Hook<BeforeToolCall> for ToolBlocker {
struct ResultFormatter;
#[async_trait]
impl Hook<AfterToolCall> for ResultFormatter {
impl Hook<PostToolCall> for ResultFormatter {
async fn call(
&self,
ctx: &mut ToolResultContext,
) -> Result<AfterToolCallResult, HookError> {
ctx: &mut PostToolCallContext,
) -> Result<PostToolCallResult, HookError> {
if !ctx.result.is_error {
ctx.result.content = format!("[OK] {}", ctx.result.content);
}
Ok(AfterToolCallResult::Continue)
Ok(PostToolCallResult::Continue)
}
}
```
@@ -283,9 +341,9 @@ impl Hook<OnTurnEnd> for JsonValidator {
Hookは**イベントごとに登録順**に実行されます。
```rust
worker.add_before_tool_call_hook(HookA); // 1番目に実行
worker.add_before_tool_call_hook(HookB); // 2番目に実行
worker.add_before_tool_call_hook(HookC); // 3番目に実行
worker.add_pre_tool_call_hook(HookA); // 1番目に実行
worker.add_pre_tool_call_hook(HookB); // 2番目に実行
worker.add_pre_tool_call_hook(HookC); // 3番目に実行
```
### 制御フローの伝播
@@ -323,15 +381,15 @@ Hook A: Continue → Hook B: Pause
async fn call(&self, ctx: &mut ToolCallContext) -> ... {
// 引数を直接書き換え
ctx.call.input["sanitized"] = json!(true);
Ok(BeforeToolCallResult::Continue)
Ok(PreToolCallResult::Continue)
}
```
### 3. 並列実行との統合
- `before_tool_call`: 並列実行**前**に逐次実行(許可判定のため)
- `pre_tool_call`: 並列実行**前**に逐次実行(許可判定のため)
- ツール実行: `join_all`で**並列**実行
- `after_tool_call`: 並列実行**後**に逐次実行(結果加工のため)
- `post_tool_call`: 並列実行**後**に逐次実行(結果加工のため)
### 4. Send + Sync 要件
@@ -344,24 +402,24 @@ struct CountingHook {
}
#[async_trait]
impl Hook<BeforeToolCall> for CountingHook {
async fn call(&self, _: &mut ToolCallContext) -> Result<BeforeToolCallResult, HookError> {
impl Hook<PreToolCall> for CountingHook {
async fn call(&self, _: &mut ToolCallContext) -> Result<PreToolCallResult, HookError> {
self.count.fetch_add(1, Ordering::SeqCst);
Ok(BeforeToolCallResult::Continue)
Ok(PreToolCallResult::Continue)
}
}
```
## 典型的なユースケース
| ユースケース | 使用Hook | 処理内容 |
| ------------------ | ------------------------ | -------------------------- |
| ツール許可制御 | `before_tool_call` | 危険なツールをSkip |
| 実行ログ | `before/after_tool_call` | 呼び出しと結果を記録 |
| 出力バリデーション | `on_turn_end` | 形式チェック、リトライ指示 |
| コンテキスト注入 | `on_message_send` | システムメッセージ追加 |
| 結果のサニタイズ | `after_tool_call` | 機密情報のマスキング |
| レート制限 | `before_tool_call` | 呼び出し頻度の制御 |
| ユースケース | 使用Hook | 処理内容 |
| ------------------ | -------------------- | -------------------------- |
| ツール許可制御 | `pre_tool_call` | 危険なツールをSkip |
| 実行ログ | `pre/post_tool_call` | 呼び出しと結果を記録 |
| 出力バリデーション | `on_turn_end` | 形式チェック、リトライ指示 |
| コンテキスト注入 | `on_message_send` | システムメッセージ追加 |
| 結果のサニタイズ | `post_tool_call` | 機密情報のマスキング |
| レート制限 | `pre_tool_call` | 呼び出し頻度の制御 |
## TODO
+191
View File
@@ -0,0 +1,191 @@
# Tool 設計
## 概要
`llm-worker`のツールシステムは、LLMが外部リソースにアクセスしたり計算を実行するための仕組みを提供する。
メタ情報の不変性とセッションスコープの状態管理を両立させる設計となっている。
## 主要な型
```
type ToolDefinition
Fn() -> (ToolMeta, Arc<dyn Tool>)
worker.register_tool() で呼び出し
- struct ToolMeta (name, desc, schema)
不変・登録時固定
- trait Tool (executer)
登録時生成・セッション中再利用
```
### ToolMeta
ツールのメタ情報を保持する不変構造体。登録時に固定され、Worker内で変更されない。
```rust
pub struct ToolMeta {
pub name: String,
pub description: String,
pub input_schema: Value,
}
```
**目的:**
- LLM へのツール定義として送信
- Hook からの参照(読み取り専用)
- 登録後の不変性を保証
### Tool trait
ツールの実行ロジックのみを定義するトレイト。
```rust
#[async_trait]
pub trait Tool: Send + Sync {
async fn execute(&self, input_json: &str) -> Result<String, ToolError>;
}
```
**設計方針:**
- メタ情報(name, description, schema)は含まない
- 状態を持つことが可能(セッション中のカウンターなど)
- `Send + Sync` で並列実行に対応
**インスタンスのライフサイクル:**
1. `register_tool()` 呼び出し時にファクトリが実行され、インスタンスが生成される
2. LLM がツールを呼び出すと、既存インスタンスの `execute()` が実行される
3. 同じセッション中は同一インスタンスが再利用される
※ 「最初に呼ばれたとき」の遅延初期化ではなく、**登録時の即時初期化**である。
### ToolDefinition
メタ情報とツールインスタンスを生成するファクトリ。
```rust
pub type ToolDefinition = Arc<dyn Fn() -> (ToolMeta, Arc<dyn Tool>) + Send + Sync>;
```
**なぜファクトリか:**
- Worker への登録時に一度だけ呼び出される
- メタ情報とインスタンスを同時に生成し、整合性を保証
- クロージャでコンテキスト(`self.clone()`)をキャプチャ可能
## Worker でのツール管理
```rust
// Worker 内部
tools: HashMap<String, (ToolMeta, Arc<dyn Tool>)>
// 登録 API
pub fn register_tool(&mut self, factory: ToolDefinition) -> Result<(), ToolRegistryError>
```
登録時の処理:
1. ファクトリを呼び出し `(meta, instance)` を取得
2. 同名ツールが既に登録されていればエラー
3. HashMap に `(meta, instance)` を保存
## マクロによる自動生成
`#[tool_registry]` マクロは `{method}_definition()` メソッドを生成する。
```rust
#[tool_registry]
impl MyApp {
/// 検索を実行する
#[tool]
async fn search(&self, query: String) -> String {
// 実装
}
}
// 生成されるコード:
impl MyApp {
pub fn search_definition(&self) -> ToolDefinition {
let ctx = self.clone();
Arc::new(move || {
let meta = ToolMeta::new("search")
.description("検索を実行する")
.input_schema(/* schemars で生成 */);
let tool = Arc::new(ToolSearch { ctx: ctx.clone() });
(meta, tool)
})
}
}
```
## Hook との連携
Hook は `ToolCallContext` / `AfterToolCallContext`
を通じてメタ情報とインスタンスにアクセスできる。
```rust
pub struct ToolCallContext {
pub call: ToolCall, // 呼び出し情報(改変可能)
pub meta: ToolMeta, // メタ情報(読み取り専用)
pub tool: Arc<dyn Tool>, // インスタンス(状態アクセス用)
}
```
**用途:**
- `meta` で名前やスキーマを確認
- `tool` でツールの内部状態を読み取り(ダウンキャスト必要)
- `call` の引数を改変してツールに渡す
## 使用例
### 手動実装
```rust
struct Counter { count: AtomicUsize }
impl Tool for Counter {
async fn execute(&self, _: &str) -> Result<String, ToolError> {
let n = self.count.fetch_add(1, Ordering::SeqCst);
Ok(format!("count: {}", n))
}
}
let def: ToolDefinition = Arc::new(|| {
let meta = ToolMeta::new("counter")
.description("カウンターを増加")
.input_schema(json!({"type": "object"}));
(meta, Arc::new(Counter { count: AtomicUsize::new(0) }))
});
worker.register_tool(def)?;
```
### マクロ使用(推奨)
```rust
#[tool_registry]
impl App {
#[tool]
async fn greet(&self, name: String) -> String {
format!("Hello, {}!", name)
}
}
let app = App;
worker.register_tool(app.greet_definition())?;
```
## 設計上の決定
| 問題 | 決定 | 理由 |
| -------------------- | ------------------------------ | ---------------------------------------------- |
| メタ情報の変更可能性 | ToolMeta を分離・不変化 | 登録後の整合性を保証 |
| 状態管理 | 登録時にインスタンス生成 | セッション中の状態保持、同一インスタンス再利用 |
| Factory vs Instance | Factory + 登録時即時呼び出し | コンテキストキャプチャと登録時検証 |
| Hook からのアクセス | Context に meta と tool を含む | 柔軟な介入を可能に |