llm_worker_rs/worker/examples/worker_cli.rs

284 lines
8.7 KiB
Rust
Raw Blame History

This file contains ambiguous Unicode characters

This file contains Unicode characters that might be confused with other characters. If you think that this is intentional, you can safely ignore this warning. Use the Escape button to reveal them.

//! Worker を用いた対話型 CLI クライアント
//!
//! Anthropic Claude API と対話するシンプルなCLIアプリケーション。
//! ツールの登録と実行、ストリーミングレスポンスの表示をデモする。
//!
//! ## 使用方法
//!
//! ```bash
//! # .envファイルにAPIキーを設定
//! echo "ANTHROPIC_API_KEY=your-api-key" > .env
//!
//! # 基本的な実行
//! cargo run --example worker_cli
//!
//! # オプション指定
//! cargo run --example worker_cli -- --model claude-3-haiku-20240307 --system "You are a helpful assistant."
//!
//! # ヘルプ表示
//! cargo run --example worker_cli -- --help
//! ```
use std::io::{self, Write};
use std::sync::{Arc, Mutex};
use clap::Parser;
use worker::{
llm_client::providers::anthropic::AnthropicClient, Handler, TextBlockEvent, TextBlockKind,
ToolUseBlockEvent, ToolUseBlockKind, Worker,
};
use worker_macros::tool_registry;
use worker_types::Message;
// 必要なマクロ展開用インポート
use schemars;
use serde;
// =============================================================================
// CLI引数定義
// =============================================================================
/// Anthropic Claude API を使った対話型CLIクライアント
#[derive(Parser, Debug)]
#[command(name = "worker-cli")]
#[command(about = "Interactive CLI client for Anthropic Claude API using Worker")]
#[command(version)]
struct Args {
/// 使用するモデル名
#[arg(short, long, default_value = "claude-sonnet-4-20250514")]
model: String,
/// システムプロンプト
#[arg(short, long)]
system: Option<String>,
/// ツールを無効化
#[arg(long, default_value = "false")]
no_tools: bool,
/// 最初のメッセージ(指定するとそれを送信して終了)
#[arg(short = 'p', long)]
prompt: Option<String>,
/// APIキー環境変数 ANTHROPIC_API_KEY より優先)
#[arg(long, env = "ANTHROPIC_API_KEY")]
api_key: String,
}
// =============================================================================
// ツール定義
// =============================================================================
/// アプリケーションコンテキスト
#[derive(Clone)]
struct AppContext;
#[tool_registry]
impl AppContext {
/// 現在の日時を取得する
///
/// システムの現在の日付と時刻を返します。
#[tool]
fn get_current_time(&self) -> String {
let now = std::time::SystemTime::now()
.duration_since(std::time::UNIX_EPOCH)
.unwrap()
.as_secs();
// シンプルなUnixタイムスタンプからの変換
format!("Current Unix timestamp: {}", now)
}
/// 簡単な計算を行う
///
/// 2つの数値の四則演算を実行します。
#[tool]
fn calculate(&self, a: f64, b: f64, operation: String) -> Result<String, String> {
let result = match operation.as_str() {
"add" | "+" => a + b,
"subtract" | "-" => a - b,
"multiply" | "*" => a * b,
"divide" | "/" => {
if b == 0.0 {
return Err("Cannot divide by zero".to_string());
}
a / b
}
_ => return Err(format!("Unknown operation: {}", operation)),
};
Ok(format!("{} {} {} = {}", a, operation, b, result))
}
}
// =============================================================================
// ストリーミング表示用ハンドラー
// =============================================================================
/// テキストをリアルタイムで出力するハンドラー
struct StreamingPrinter {
is_first_delta: Arc<Mutex<bool>>,
}
impl StreamingPrinter {
fn new() -> Self {
Self {
is_first_delta: Arc::new(Mutex::new(true)),
}
}
}
impl Handler<TextBlockKind> for StreamingPrinter {
type Scope = ();
fn on_event(&mut self, _scope: &mut (), event: &TextBlockEvent) {
match event {
TextBlockEvent::Start(_) => {
let mut first = self.is_first_delta.lock().unwrap();
if *first {
print!("\n🤖 ");
*first = false;
}
}
TextBlockEvent::Delta(text) => {
print!("{}", text);
io::stdout().flush().ok();
}
TextBlockEvent::Stop(_) => {
println!();
}
}
}
}
/// ツール呼び出しを表示するハンドラー
struct ToolCallPrinter;
impl Handler<ToolUseBlockKind> for ToolCallPrinter {
type Scope = String;
fn on_event(&mut self, json_buffer: &mut String, event: &ToolUseBlockEvent) {
match event {
ToolUseBlockEvent::Start(start) => {
println!("\n🔧 Calling tool: {}", start.name);
}
ToolUseBlockEvent::InputJsonDelta(json) => {
json_buffer.push_str(json);
}
ToolUseBlockEvent::Stop(_) => {
println!(" Args: {}", json_buffer);
}
}
}
}
// =============================================================================
// メイン
// =============================================================================
#[tokio::main]
async fn main() -> Result<(), Box<dyn std::error::Error>> {
// CLI引数をパース
let args = Args::parse();
// 対話モードかワンショットモードか
let is_interactive = args.prompt.is_none();
if is_interactive {
println!("╔════════════════════════════════════════════════╗");
println!("║ Worker CLI - Anthropic Claude Client ║");
println!("╚════════════════════════════════════════════════╝");
println!();
println!("Model: {}", args.model);
if let Some(ref system) = args.system {
println!("System: {}", system);
}
if args.no_tools {
println!("Tools: disabled");
} else {
println!("Tools:");
println!(" • get_current_time - Get the current timestamp");
println!(" • calculate - Perform arithmetic (add, subtract, multiply, divide)");
}
println!();
println!("Type 'quit' or 'exit' to end the session.");
println!("─────────────────────────────────────────────────");
}
// クライアント作成
let client = AnthropicClient::new(&args.api_key, &args.model);
// Worker作成
let mut worker = Worker::new(client);
// システムプロンプトを設定
if let Some(ref system_prompt) = args.system {
worker.set_system_prompt(system_prompt);
}
// ツール登録(--no-tools でなければ)
if !args.no_tools {
let app = AppContext;
worker.register_tool(app.get_current_time_tool());
worker.register_tool(app.calculate_tool());
}
// ストリーミング表示用ハンドラーを登録
worker
.timeline_mut()
.on_text_block(StreamingPrinter::new())
.on_tool_use_block(ToolCallPrinter);
// 会話履歴
let mut history: Vec<Message> = Vec::new();
// ワンショットモード
if let Some(prompt) = args.prompt {
history.push(Message::user(&prompt));
match worker.run(history).await {
Ok(_) => {}
Err(e) => {
eprintln!("\n❌ Error: {}", e);
std::process::exit(1);
}
}
return Ok(());
}
// 対話ループ
loop {
print!("\n👤 You: ");
io::stdout().flush()?;
let mut input = String::new();
io::stdin().read_line(&mut input)?;
let input = input.trim();
if input.is_empty() {
continue;
}
if input == "quit" || input == "exit" {
println!("\n👋 Goodbye!");
break;
}
// ユーザーメッセージを履歴に追加
history.push(Message::user(input));
// Workerを実行
match worker.run(history.clone()).await {
Ok(new_history) => {
history = new_history;
}
Err(e) => {
eprintln!("\n❌ Error: {}", e);
// エラー時は最後のユーザーメッセージを削除
history.pop();
}
}
}
Ok(())
}