llm-worker 0.2.0

//! ツール定義
//!
//! LLMから呼び出し可能なツールを定義するためのトレイト。
//! 通常は`#[tool]`マクロを使用して自動実装します。

use std::sync::Arc;

use async_trait::async_trait;
use serde_json::Value;
use thiserror::Error;

/// ツール実行時のエラー
#[derive(Debug, Error)]
pub enum ToolError {
    /// 引数が不正
    #[error("Invalid argument: {0}")]
    InvalidArgument(String),
    /// 実行に失敗
    #[error("Execution failed: {0}")]
    ExecutionFailed(String),
    /// 内部エラー
    #[error("Internal error: {0}")]
    Internal(String),
}

// =============================================================================
// ToolMeta - 不変のメタ情報
// =============================================================================

/// ツールのメタ情報（登録時に固定、不変）
///
/// `ToolDefinition` ファクトリから生成され、Worker に登録後は変更されません。
/// LLM へのツール定義送信に使用されます。
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct ToolMeta {
    /// ツール名（LLMが識別に使用）
    pub name: String,
    /// ツールの説明（LLMへのプロンプトに含まれる）
    pub description: String,
    /// 引数のJSON Schema
    pub input_schema: Value,
}

impl ToolMeta {
    /// 新しい ToolMeta を作成
    pub fn new(name: impl Into<String>) -> Self {
        Self {
            name: name.into(),
            description: String::new(),
            input_schema: Value::Object(Default::default()),
        }
    }

    /// 説明を設定
    pub fn description(mut self, desc: impl Into<String>) -> Self {
        self.description = desc.into();
        self
    }

    /// 引数スキーマを設定
    pub fn input_schema(mut self, schema: Value) -> Self {
        self.input_schema = schema;
        self
    }
}

// =============================================================================
// ToolDefinition - ファクトリ型
// =============================================================================

/// ツール定義ファクトリ
///
/// 呼び出すと `(ToolMeta, Arc<dyn Tool>)` を返します。
/// Worker への登録時に一度だけ呼び出され、メタ情報とインスタンスが
/// セッションスコープでキャッシュされます。
///
/// # Examples
///
/// ```ignore
/// let def: ToolDefinition = Arc::new(|| {
///     (
///         ToolMeta::new("my_tool")
///             .description("My tool description")
///             .input_schema(json!({"type": "object"})),
///         Arc::new(MyToolImpl { state: 0 }) as Arc<dyn Tool>,
///     )
/// });
/// worker.register_tool(def)?;
/// ```
pub type ToolDefinition = Arc<dyn Fn() -> (ToolMeta, Arc<dyn Tool>) + Send + Sync>;

// =============================================================================
// Tool trait
// =============================================================================

/// LLMから呼び出し可能なツールを定義するトレイト
///
/// ツールはLLMが外部リソースにアクセスしたり、
/// 計算を実行したりするために使用します。
/// セッション中の状態を保持できます。
///
/// # 実装方法
///
/// 通常は`#[tool_registry]`マクロを使用して自動実装します:
///
/// ```ignore
/// #[tool_registry]
/// impl MyApp {
///     #[tool]
///     async fn search(&self, query: String) -> String {
///         format!("Results for: {}", query)
///     }
/// }
///
/// // 登録
/// worker.register_tool(app.search_definition())?;
/// ```
///
/// # 手動実装
///
/// ```ignore
/// use llm_worker::tool::{Tool, ToolError, ToolMeta, ToolDefinition};
/// use std::sync::Arc;
///
/// struct MyTool { counter: std::sync::atomic::AtomicUsize }
///
/// #[async_trait::async_trait]
/// impl Tool for MyTool {
///     async fn execute(&self, input: &str) -> Result<String, ToolError> {
///         self.counter.fetch_add(1, std::sync::atomic::Ordering::SeqCst);
///         Ok("result".to_string())
///     }
/// }
///
/// let def: ToolDefinition = Arc::new(|| {
///     (
///         ToolMeta::new("my_tool")
///             .description("My custom tool")
///             .input_schema(serde_json::json!({"type": "object"})),
///         Arc::new(MyTool { counter: Default::default() }) as Arc<dyn Tool>,
///     )
/// });
/// ```
#[async_trait]
pub trait Tool: Send + Sync {
    /// ツールを実行する
    ///
    /// # Arguments
    /// * `input_json` - LLMが生成したJSON形式の引数
    ///
    /// # Returns
    /// 実行結果の文字列。この内容がLLMに返されます。
    async fn execute(&self, input_json: &str) -> Result<String, ToolError>;
}