Hare/llm_worker_rs
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

fmt: cargo fmt

Browse Source
This commit is contained in:
Hare
2026-01-07 22:04:44 +09:00
parent bb73dc6a45
commit 1e126c1698
20 changed files with 263 additions and 227 deletions
Download Patch File Download Diff File Expand all files Collapse all files
worker/tests/common/mod.rs
+32 -18
View File
@@ -3,13 +3,13 @@
use std::fs::File;
use std::io::{BufRead, BufReader};
use std::path::{Path, PathBuf};
use std::sync::{Arc, Mutex};
use std::pin::Pin;
use std::sync::{Arc, Mutex};
use async_trait::async_trait;
use futures::Stream;
use worker::{Handler, TextBlockEvent, TextBlockKind, Timeline};
use worker::llm_client::{ClientError, LlmClient, Request};
use worker::{Handler, TextBlockEvent, TextBlockKind, Timeline};
use worker_types::{BlockType, DeltaContent, Event};
use std::sync::atomic::{AtomicUsize, Ordering};
@@ -51,11 +51,11 @@ impl LlmClient for MockLlmClient {
) -> Result<Pin<Box<dyn Stream<Item = Result<Event, ClientError>> + Send>>, ClientError> {
let count = self.call_count.fetch_add(1, Ordering::SeqCst);
if count >= self.responses.len() {
return Err(ClientError::Api {
status: Some(500),
code: Some("mock_error".to_string()),
message: "No more mock responses".to_string(),
});
return Err(ClientError::Api {
status: Some(500),
code: Some("mock_error".to_string()),
message: "No more mock responses".to_string(),
});
}
let events = self.responses[count].clone();
let stream = futures::stream::iter(events.into_iter().map(Ok));
@@ -135,7 +135,8 @@ pub fn assert_event_sequence(subdir: &str) {
}
// Find a text-based fixture
let fixture_path = fixtures.iter()
let fixture_path = fixtures
.iter()
.find(|p| p.to_string_lossy().contains("text"))
.unwrap_or(&fixtures[0]);
@@ -156,9 +157,9 @@ pub fn assert_event_sequence(subdir: &str) {
}
}
Event::BlockDelta(delta) => {
if let DeltaContent::Text(_) = &delta.delta {
delta_found = true;
}
if let DeltaContent::Text(_) = &delta.delta {
delta_found = true;
}
}
Event::BlockStop(stop) => {
if stop.block_type == BlockType::Text {
@@ -173,9 +174,9 @@ pub fn assert_event_sequence(subdir: &str) {
// Check for BlockStart (Warn only for OpenAI/Ollama as it might be missing for text)
if !start_found {
println!("Warning: No BlockStart found. This is common for OpenAI/Ollama text streams.");
// For Anthropic, strict start is usually expected, but to keep common logic simple we allow warning.
// If specific strictness is needed, we could add a `strict: bool` arg.
println!("Warning: No BlockStart found. This is common for OpenAI/Ollama text streams.");
// For Anthropic, strict start is usually expected, but to keep common logic simple we allow warning.
// If specific strictness is needed, we could add a `strict: bool` arg.
}
assert!(delta_found, "Should contain BlockDelta");
@@ -184,7 +185,9 @@ pub fn assert_event_sequence(subdir: &str) {
assert!(stop_found, "Should contain BlockStop for Text block");
} else {
if !stop_found {
println!(" [Type: ToolUse] BlockStop detection skipped (not explicitly emitted by scheme)");
println!(
" [Type: ToolUse] BlockStop detection skipped (not explicitly emitted by scheme)"
);
}
}
}
@@ -200,13 +203,23 @@ pub fn assert_usage_tokens(subdir: &str) {
let events = load_events_from_fixture(&fixture);
let usage_events: Vec<_> = events
.iter()
.filter_map(|e| if let Event::Usage(u) = e { Some(u) } else { None })
.filter_map(|e| {
if let Event::Usage(u) = e {
Some(u)
} else {
None
}
})
.collect();
if !usage_events.is_empty() {
let last_usage = usage_events.last().unwrap();
if last_usage.input_tokens.is_some() || last_usage.output_tokens.is_some() {
println!(" Fixture {:?} Usage: {:?}", fixture.file_name(), last_usage);
println!(
" Fixture {:?} Usage: {:?}",
fixture.file_name(),
last_usage
);
return; // Found valid usage
}
}
@@ -221,7 +234,8 @@ pub fn assert_timeline_integration(subdir: &str) {
return;
}
let fixture_path = fixtures.iter()
let fixture_path = fixtures
.iter()
.find(|p| p.to_string_lossy().contains("text"))
.unwrap_or(&fixtures[0]);
worker/tests/parallel_execution_test.rs
+35 -12
View File
@@ -2,13 +2,16 @@
//!
//! Workerが複数のツールを並列に実行することを確認する。
use std::sync::atomic::{AtomicUsize, Ordering};
use std::sync::Arc;
use std::sync::atomic::{AtomicUsize, Ordering};
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};
use worker_types::{
ControlFlow, Event, HookError, Message, ResponseStatus, StatusEvent, Tool, ToolCall, ToolError,
ToolResult, WorkerHook,
};
mod common;
use common::MockLlmClient;
@@ -105,8 +108,6 @@ async fn test_parallel_tool_execution() {
worker.register_tool(tool2);
worker.register_tool(tool3);
let messages = vec![Message::user("Run all tools")];
let start = Instant::now();
@@ -161,7 +162,10 @@ async fn test_before_tool_call_skip() {
#[async_trait]
impl WorkerHook for BlockingHook {
async fn before_tool_call(&self, tool_call: &mut ToolCall) -> Result<ControlFlow, HookError> {
async fn before_tool_call(
&self,
tool_call: &mut ToolCall,
) -> Result<ControlFlow, HookError> {
if tool_call.name == "blocked_tool" {
Ok(ControlFlow::Skip)
} else {
@@ -176,8 +180,16 @@ async fn test_before_tool_call_skip() {
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");
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 で結果が改変されることを確認
@@ -212,9 +224,15 @@ async fn test_after_tool_call_modification() {
#[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!({}) }
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())
}
@@ -229,7 +247,10 @@ async fn test_after_tool_call_modification() {
#[async_trait]
impl WorkerHook for ModifyingHook {
async fn after_tool_call(&self, tool_result: &mut ToolResult) -> Result<ControlFlow, HookError> {
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)
@@ -237,7 +258,9 @@ async fn test_after_tool_call_modification() {
}
let modified_content = Arc::new(std::sync::Mutex::new(None));
worker.add_hook(ModifyingHook { modified_content: modified_content.clone() });
worker.add_hook(ModifyingHook {
modified_content: modified_content.clone(),
});
let messages = vec![Message::user("Test modification")];
let result = worker.run(messages).await;
worker/tests/tool_macro_test.rs
+20 -5
View File
@@ -2,8 +2,8 @@
//!
//! `#[tool_registry]` と `#[tool]` マクロの動作を確認する。
use std::sync::atomic::{AtomicUsize, Ordering};
use std::sync::Arc;
use std::sync::atomic::{AtomicUsize, Ordering};
// マクロ展開に必要なインポート
use schemars;
@@ -59,12 +59,19 @@ async fn test_basic_tool_generation() {
// 説明の確認（docコメントから取得）
let desc = greet_tool.description();
assert!(desc.contains("メッセージに挨拶を追加する"), "Description should contain doc comment: {}", desc);
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");
assert!(
schema.get("properties").is_some(),
"Schema should have properties"
);
// 実行テスト
let result = greet_tool.execute(r#"{"message": "World"}"#).await;
@@ -104,7 +111,11 @@ async fn test_no_arguments() {
let result = get_prefix_tool.execute(r#"{}"#).await;
assert!(result.is_ok());
let output = result.unwrap();
assert!(output.contains("TestPrefix"), "Should contain prefix: {}", output);
assert!(
output.contains("TestPrefix"),
"Should contain prefix: {}",
output
);
}
#[tokio::test]
@@ -169,7 +180,11 @@ async fn test_result_return_type_error() {
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);
assert!(
err.to_string().contains("positive"),
"Error should mention positive: {}",
err
);
}
// =============================================================================
worker/tests/worker_fixtures.rs
+3 -7
View File
@@ -6,8 +6,8 @@
mod common;
use std::path::Path;
use std::sync::atomic::{AtomicUsize, Ordering};
use std::sync::Arc;
use std::sync::atomic::{AtomicUsize, Ordering};
use async_trait::async_trait;
use common::MockLlmClient;
@@ -67,9 +67,7 @@ impl Tool for MockWeatherTool {
let input: serde_json::Value = serde_json::from_str(input_json)
.map_err(|e| ToolError::InvalidArgument(e.to_string()))?;
let city = input["city"]
.as_str()
.unwrap_or("Unknown");
let city = input["city"].as_str().unwrap_or("Unknown");
// モックのレスポンスを返す
Ok(format!("Weather in {}: Sunny, 22°C", city))
@@ -163,8 +161,6 @@ async fn test_worker_tool_call() {
let tool_for_check = weather_tool.clone();
worker.register_tool(weather_tool);
// メッセージを送信
let messages = vec![worker_types::Message::user("What's the weather in Tokyo?")];
let _result = worker.run(messages).await;
@@ -212,8 +208,8 @@ async fn test_worker_with_programmatic_events() {
/// id, name, input（JSON）を正しく抽出できることを検証する。
#[tokio::test]
async fn test_tool_call_collector_integration() {
use worker::ToolCallCollector;
use worker::Timeline;
use worker::ToolCallCollector;
use worker_types::Event;
// ToolUseブロックを含むイベントシーケンス