llm_worker_rs/docs/spec/worker_design.md

9.1 KiB
Raw Blame History

Worker & Tool/Hook 設計

概要

Workerはアプリケーションの「ターン」を制御する高レベルコンポーネントです。 LlmClientTimelineを内包し、ユーザー定義のToolHookを用いて自律的なインタラクションを行います。

アーキテクチャ

graph TD
    User[Application / User] -->|1. Run| Worker
    Worker -->|2. Event Loop| Timeline
    Timeline -->|3. Dispatch| Handler[Handlers (inc. ToolExecutor)]
    
    subgraph "Worker Layer"
        Worker
        Hook[Hooks]
    end
    
    subgraph "Core Layer"
        Timeline
        LlmClient
    end
    
    Worker -.->|Intervene| Hook
    Handler -.->|Execute| Tool[User Defined Tools]

ライフサイクル (ターン制御)

Workerは以下のループターンを実行します。

  1. Start Turn: Worker::run(messages) 呼び出し
  2. Hook: OnMessageSend:
    • ユーザーメッセージの改変、バリデーション、キャンセルが可能。
    • コンテキストへのシステムプロンプト注入などもここで行う想定。
  3. Request & Stream:
    • LLMへリクエスト送信。イベントストリーム開始。
    • Timelineによるイベント処理。
  4. Tool Handling (Parallel):
    • レスポンス内に含まれる全てのTool Callを収集。
    • 各Toolに対して Hook: BeforeToolCall を実行(実行可否、引数改変)。
    • 許可されたToolを並列実行 (join_all)
    • 各Tool実行後に Hook: AfterToolCall を実行(結果の確認、加工)。
  5. Next Request Decision:
    • Tool実行結果がある場合 -> 結果をMessageとしてContextに追加し、Step 3へ戻る (自動ループ)。
    • Tool実行がない場合 -> Step 6へ。
  6. Hook: OnTurnEnd:
    • 最終的な応答に対するチェックLint/Fmt
    • エラーがある場合、エラーメッセージをContextに追加して Step 3へ戻る ことで自己修正を促せる。
    • 問題なければターン終了。

Tool 設計

アーキテクチャ概要

Rustの静的型付けシステムとLLMの動的なツール呼び出し文字列による指定を、Trait Object動的ディスパッチ を用いて接続します。

  1. 共通インターフェース (Tool Trait): 全てのツールが実装すべき共通の振る舞い(メタデータ取得と実行)を定義します。
  2. ラッパー生成 (#[tool] Macro): ユーザー定義のメソッドをラップし、Tool Traitを実装した構造体を自動生成します。
  3. レジストリ (HashMap): Workerは動的ディスパッチ用に HashMap<String, Box<dyn Tool>> でツールを管理します。

この仕組みにより、「名前からツールを探し、JSON引数を型変換して関数を実行する」フローを安全に実現します。

1. Tool Trait 定義

ツールが最低限持つべきインターフェースです。Send + Sync を必須とし、マルチスレッド(並列実行)に対応します。

#[async_trait]
pub trait Tool: Send + Sync {
    /// ツール名 (LLMが識別に使用)
    fn name(&self) -> &str;
    
    /// ツールの説明 (LLMへのプロンプトに含まれる)
    fn description(&self) -> &str;
    
    /// 引数のJSON Schema (schemars等で生成)
    fn input_schema(&self) -> serde_json::Value;
    
    /// 実行関数
    /// JSON文字列を受け取り、デシリアライズして元のメソッドを実行し、結果を返す
    async fn execute(&self, input_json: &str) -> Result<String, ToolError>;
}

2. マクロと実装モデル

ユーザーは「状態を持つ構造体」とその「メソッド」としてツールを定義します。

ユーザーコード:

#[derive(Clone)] // 状態はClone (Arc推奨) で共有される想定
struct MyApp {
    db: Arc<Database>,
}

impl MyApp {
    /// ユーザー情報を取得する
    /// 指定されたIDのユーザーをDBから検索します。
    #[tool]
    async fn get_user(
        &self, 
        #[description = "取得したいユーザーのID"] user_id: String
    ) -> Result<User, Error> {
        let user = self.db.find(&user_id).await?;
        Ok(user)
    }
}

マクロ展開後のイメージ (擬似コード):

マクロは、元のメソッドに対応するラッパー構造体を生成します。このラッパーが Tool Trait を実装します。

// 1. 引数をデシリアライズ用の中間構造体に変換
#[derive(serde::Deserialize, schemars::JsonSchema)]
struct GetUserArgs {
    /// 取得したいユーザーのID
    user_id: String,
}

// 2. ラッパー構造体 (元のコンテキストを持つ)
struct GetUserTool {
    ctx: MyApp, // コンテキストを保持 (Clone)
}

#[async_trait]
impl Tool for GetUserTool {
    fn name(&self) -> &str { "get_user" }
    
    fn description(&self) -> &str { "ユーザー情報を取得する\n指定されたIDのユーザーをDBから検索します。" }
    
    fn input_schema(&self) -> serde_json::Value {
        schemars::schema_for!(GetUserArgs)
    }
    
    async fn execute(&self, input_json: &str) -> Result<String, ToolError> {
        // A. JSONを引数構造体に変換
        let args: GetUserArgs = serde_json::from_str(input_json)
            .map_err(|e| ToolError::InvalidArgument(e.to_string()))?;
            
        // B. 元のメソッド呼び出し (self.ctx 経由)
        let result = self.ctx.get_user(args.user_id).await
            .map_err(|e| ToolError::ExecutionFailed(e.to_string()))?;
            
        // C. 結果を文字列化
        Ok(format!("{:?}", result)) // または serde_json::to_string(&result)
    }
}

3. Workerによる実行フロー

Workerは生成されたラッパー構造体を Box<dyn Tool> として保持し、以下のフローで実行します。

  1. 登録: アプリケーション開始時、コンテキスト(MyApp)から各ツールのラッパー(GetUserTool)を生成し、WorkerのMapに登録。
  2. 解決: LLMからのレスポンスに含まれる ToolUse { name: "get_user", ... } を受け取る。
  3. 検索: name をキーに Map から Box<dyn Tool> を取得。
  4. 実行:
    • tool.execute(json) を呼び出す。
    • 内部で serde_json による型変換とメソッド実行が行われる。
    • 結果が返る。

これにより、型安全性を保ちつつ、動的なツール実行が可能になります。

Hook 設計

コンセプト

  • 制御の介入: ターンの進行、メッセージの内容、ツールの実行に対して介入します。
  • Contextへのアクセス: メッセージ履歴Contextを読み書きできます。

Hook Trait

#[async_trait]
pub trait WorkerHook: Send + Sync {
    /// メッセージ送信前。
    /// リクエストに含まれるメッセージリストを改変できる。
    async fn on_message_send(&self, context: &mut Vec<Message>) -> Result<ControlFlow, Error> {
        Ok(ControlFlow::Continue)
    }

    /// ツール実行前。
    /// 実行をキャンセルしたり、引数を書き換えることができる。
    async fn before_tool_call(&self, tool_call: &mut ToolCall) -> Result<ControlFlow, Error> {
        Ok(ControlFlow::Continue)
    }

    /// ツール実行後。
    /// 結果を書き換えたり、隠蔽したりできる。
    async fn after_tool_call(&self, tool_result: &mut ToolResult) -> Result<ControlFlow, Error> {
        Ok(ControlFlow::Continue)
    }

    /// ターン終了時。
    /// 生成されたメッセージを検査し、必要ならリトライContinueWithMessagesを指示できる。
    async fn on_turn_end(&self, messages: &[Message]) -> Result<TurnResult, Error> {
        Ok(TurnResult::Finish)
    }
}

pub enum ControlFlow {
    Continue,
    Skip, // Tool実行などをスキップ
    Abort(String), // 処理中断
}

pub enum TurnResult {
    Finish,
    ContinueWithMessages(Vec<Message>), // メッセージを追加してターン継続(自己修正など)
}

実装方針

  1. Worker Struct:

    • Timelineを所有。
    • Handlerとして「ToolCallCollector」をTimelineに登録。
    • stream終了後に収集したToolCallを処理するロジックを持つ。
  2. Tool Executor Handler:

    • Timeline上ではツール実行を行わず、あくまで「ToolCallブロックの収集」に徹するToolの実行は非同期かつ並列で、ストリーム終了後あるいはブロック確定後に行うため
    • ただし、リアルタイム性を重視する場合ストリーミング中にToolを実行開始等は将来的な拡張とするが、現状は「結果が揃うのを待って」という要件に従い、収集フェーズと実行フェーズを分ける。
  3. worker-macros:

    • syn, quote を用いて、関数定義から Tool トレイト実装と InputInputSchema (schemars利用) を生成。