12 Commits
79 changed files with 3084 additions and 1572 deletions
Generated
+255 -44
View File
@@ -104,9 +104,9 @@ checksum = "b35204fbdc0b3f4446b89fc1ac2cf84a8a68971995d0bf2e925ec7cd960f9cb3"
[[package]] [[package]]
name = "cc" name = "cc"
version = "1.2.51" version = "1.2.52"
source = "registry+https://github.com/rust-lang/crates.io-index" source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "7a0aeaff4ff1a90589618835a598e545176939b97874f7abc7851caa0618f203" checksum = "cd4932aefd12402b36c60956a4fe0035421f544799057659ff86f923657aada3"
dependencies = [ dependencies = [
"find-msvc-tools", "find-msvc-tools",
"shlex", "shlex",
@@ -203,6 +203,12 @@ version = "1.0.20"
source = "registry+https://github.com/rust-lang/crates.io-index" source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "d0881ea181b1df73ff77ffaaf9c7544ecc11e82fba9b5f27b262a3c73a332555" checksum = "d0881ea181b1df73ff77ffaaf9c7544ecc11e82fba9b5f27b262a3c73a332555"
[[package]]
name = "equivalent"
version = "1.0.2"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "877a4ace8713b0bcf2a4e7eec82529c029f1d0619886d18145fea96c3ffe5c0f"
[[package]] [[package]]
name = "errno" name = "errno"
version = "0.3.14" version = "0.3.14"
@@ -232,9 +238,15 @@ checksum = "37909eebbb50d72f9059c3b6d82c0463f2ff062c9e95845c43a6c9c0355411be"
[[package]] [[package]]
name = "find-msvc-tools" name = "find-msvc-tools"
version = "0.1.6" version = "0.1.7"
source = "registry+https://github.com/rust-lang/crates.io-index" source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "645cbb3a84e60b7531617d5ae4e57f7e27308f6445f5abf653209ea76dec8dff" checksum = "f449e6c6c08c865631d4890cfacf252b3d396c9bcc83adb6623cdb02a8336c41"
[[package]]
name = "fnv"
version = "1.0.7"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "3f9eec918d3f24069decb9af1554cad7c880e2da24a9afd88aca000531ab82c1"
[[package]] [[package]]
name = "foreign-types" name = "foreign-types"
@@ -349,6 +361,17 @@ dependencies = [
"slab", "slab",
] ]
[[package]]
name = "getrandom"
version = "0.2.16"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "335ff9f135e4384c8150d6f27c6daed433577f86b4750418338c01a1a2528592"
dependencies = [
"cfg-if",
"libc",
"wasi",
]
[[package]] [[package]]
name = "getrandom" name = "getrandom"
version = "0.3.4" version = "0.3.4"
@@ -361,6 +384,31 @@ dependencies = [
"wasip2", "wasip2",
] ]
[[package]]
name = "h2"
version = "0.4.13"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "2f44da3a8150a6703ed5d34e164b875fd14c2cdab9af1252a9a1020bde2bdc54"
dependencies = [
"atomic-waker",
"bytes",
"fnv",
"futures-core",
"futures-sink",
"http",
"indexmap",
"slab",
"tokio",
"tokio-util",
"tracing",
]
[[package]]
name = "hashbrown"
version = "0.16.1"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "841d1cc9bed7f9236f321df977030373f4a4163ae1a7dbfe1a51a2c1a51d9100"
[[package]] [[package]]
name = "heck" name = "heck"
version = "0.5.0" version = "0.5.0"
@@ -416,6 +464,7 @@ dependencies = [
"bytes", "bytes",
"futures-channel", "futures-channel",
"futures-core", "futures-core",
"h2",
"http", "http",
"http-body", "http-body",
"httparse", "httparse",
@@ -427,6 +476,22 @@ dependencies = [
"want", "want",
] ]
[[package]]
name = "hyper-rustls"
version = "0.27.7"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "e3c93eb611681b207e1fe55d5a71ecf91572ec8a6705cdb6857f7d8d5242cf58"
dependencies = [
"http",
"hyper",
"hyper-util",
"rustls",
"rustls-pki-types",
"tokio",
"tokio-rustls",
"tower-service",
]
[[package]] [[package]]
name = "hyper-tls" name = "hyper-tls"
version = "0.6.0" version = "0.6.0"
@@ -569,6 +634,16 @@ dependencies = [
"icu_properties", "icu_properties",
] ]
[[package]]
name = "indexmap"
version = "2.13.0"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "7714e70437a7dc3ac8eb7e6f8df75fd8eb422675fc7678aff7364301092b1017"
dependencies = [
"equivalent",
"hashbrown",
]
[[package]] [[package]]
name = "ipnet" name = "ipnet"
version = "2.11.0" version = "2.11.0"
@@ -631,6 +706,37 @@ version = "0.8.1"
source = "registry+https://github.com/rust-lang/crates.io-index" source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "6373607a59f0be73a39b6fe456b8192fcc3585f602af20751600e974dd455e77" checksum = "6373607a59f0be73a39b6fe456b8192fcc3585f602af20751600e974dd455e77"
[[package]]
name = "llm-worker"
version = "0.2.0"
dependencies = [
"async-trait",
"clap",
"dotenv",
"eventsource-stream",
"futures",
"llm-worker-macros",
"reqwest",
"schemars",
"serde",
"serde_json",
"tempfile",
"thiserror",
"tokio",
"tokio-util",
"tracing",
"tracing-subscriber",
]
[[package]]
name = "llm-worker-macros"
version = "0.2.0"
dependencies = [
"proc-macro2",
"quote",
"syn",
]
[[package]] [[package]]
name = "log" name = "log"
version = "0.4.29" version = "0.4.29"
@@ -865,10 +971,12 @@ dependencies = [
"bytes", "bytes",
"futures-core", "futures-core",
"futures-util", "futures-util",
"h2",
"http", "http",
"http-body", "http-body",
"http-body-util", "http-body-util",
"hyper", "hyper",
"hyper-rustls",
"hyper-tls", "hyper-tls",
"hyper-util", "hyper-util",
"js-sys", "js-sys",
@@ -893,6 +1001,20 @@ dependencies = [
"web-sys", "web-sys",
] ]
[[package]]
name = "ring"
version = "0.17.14"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "a4689e6c2294d81e88dc6261c768b63bc4fcdb852be6d1352498b114f61383b7"
dependencies = [
"cc",
"cfg-if",
"getrandom 0.2.16",
"libc",
"untrusted",
"windows-sys 0.52.0",
]
[[package]] [[package]]
name = "rustix" name = "rustix"
version = "1.1.3" version = "1.1.3"
@@ -906,6 +1028,19 @@ dependencies = [
"windows-sys 0.61.2", "windows-sys 0.61.2",
] ]
[[package]]
name = "rustls"
version = "0.23.36"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "c665f33d38cea657d9614f766881e4d510e0eda4239891eea56b4cadcf01801b"
dependencies = [
"once_cell",
"rustls-pki-types",
"rustls-webpki",
"subtle",
"zeroize",
]
[[package]] [[package]]
name = "rustls-pki-types" name = "rustls-pki-types"
version = "1.13.2" version = "1.13.2"
@@ -915,6 +1050,17 @@ dependencies = [
"zeroize", "zeroize",
] ]
[[package]]
name = "rustls-webpki"
version = "0.103.8"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "2ffdfa2f5286e2247234e03f680868ac2815974dc39e00ea15adc445d0aafe52"
dependencies = [
"ring",
"rustls-pki-types",
"untrusted",
]
[[package]] [[package]]
name = "rustversion" name = "rustversion"
version = "1.0.22" version = "1.0.22"
@@ -1081,6 +1227,12 @@ version = "0.11.1"
source = "registry+https://github.com/rust-lang/crates.io-index" source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "7da8b5736845d9f2fcb837ea5d9e2628564b3b043a70948a3f0b778838c5fb4f" checksum = "7da8b5736845d9f2fcb837ea5d9e2628564b3b043a70948a3f0b778838c5fb4f"
[[package]]
name = "subtle"
version = "2.6.1"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "13c2bddecc57b384dee18652358fb23172facb8a2c51ccc10d74c157bdea3292"
[[package]] [[package]]
name = "syn" name = "syn"
version = "2.0.114" version = "2.0.114"
@@ -1119,7 +1271,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "655da9c7eb6305c55742045d5a8d2037996d61d8de95806335c7c86ce0f82e9c" checksum = "655da9c7eb6305c55742045d5a8d2037996d61d8de95806335c7c86ce0f82e9c"
dependencies = [ dependencies = [
"fastrand", "fastrand",
"getrandom", "getrandom 0.3.4",
"once_cell", "once_cell",
"rustix", "rustix",
"windows-sys 0.61.2", "windows-sys 0.61.2",
@@ -1200,6 +1352,16 @@ dependencies = [
"tokio", "tokio",
] ]
[[package]]
name = "tokio-rustls"
version = "0.26.4"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "1729aa945f29d91ba541258c8df89027d5792d85a8841fb65e8bf0f4ede4ef61"
dependencies = [
"rustls",
"tokio",
]
[[package]] [[package]]
name = "tokio-util" name = "tokio-util"
version = "0.7.18" version = "0.7.18"
@@ -1331,6 +1493,12 @@ version = "1.0.22"
source = "registry+https://github.com/rust-lang/crates.io-index" source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "9312f7c4f6ff9069b165498234ce8be658059c6728633667c526e27dc2cf1df5" checksum = "9312f7c4f6ff9069b165498234ce8be658059c6728633667c526e27dc2cf1df5"
[[package]]
name = "untrusted"
version = "0.9.0"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "8ecb6da28b8a351d773b68d5825ac39017e680750f980f3a1a85cd8dd28a47c1"
[[package]] [[package]]
name = "url" name = "url"
version = "2.5.8" version = "2.5.8"
@@ -1478,13 +1646,22 @@ version = "0.2.1"
source = "registry+https://github.com/rust-lang/crates.io-index" source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "f0805222e57f7521d6a62e36fa9163bc891acd422f971defe97d64e70d0a4fe5" checksum = "f0805222e57f7521d6a62e36fa9163bc891acd422f971defe97d64e70d0a4fe5"
[[package]]
name = "windows-sys"
version = "0.52.0"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "282be5f36a8ce781fad8c8ae18fa3f9beff57ec1b52cb3de0789201425d9a33d"
dependencies = [
"windows-targets 0.52.6",
]
[[package]] [[package]]
name = "windows-sys" name = "windows-sys"
version = "0.60.2" version = "0.60.2"
source = "registry+https://github.com/rust-lang/crates.io-index" source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "f2f500e4d28234f72040990ec9d39e3a6b950f9f22d3dba18416c35882612bcb" checksum = "f2f500e4d28234f72040990ec9d39e3a6b950f9f22d3dba18416c35882612bcb"
dependencies = [ dependencies = [
"windows-targets", "windows-targets 0.53.5",
] ]
[[package]] [[package]]
@@ -1496,6 +1673,22 @@ dependencies = [
"windows-link", "windows-link",
] ]
[[package]]
name = "windows-targets"
version = "0.52.6"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "9b724f72796e036ab90c1021d4780d4d3d648aca59e491e6b98e725b84e99973"
dependencies = [
"windows_aarch64_gnullvm 0.52.6",
"windows_aarch64_msvc 0.52.6",
"windows_i686_gnu 0.52.6",
"windows_i686_gnullvm 0.52.6",
"windows_i686_msvc 0.52.6",
"windows_x86_64_gnu 0.52.6",
"windows_x86_64_gnullvm 0.52.6",
"windows_x86_64_msvc 0.52.6",
]
[[package]] [[package]]
name = "windows-targets" name = "windows-targets"
version = "0.53.5" version = "0.53.5"
@@ -1503,58 +1696,106 @@ source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "4945f9f551b88e0d65f3db0bc25c33b8acea4d9e41163edf90dcd0b19f9069f3" checksum = "4945f9f551b88e0d65f3db0bc25c33b8acea4d9e41163edf90dcd0b19f9069f3"
dependencies = [ dependencies = [
"windows-link", "windows-link",
"windows_aarch64_gnullvm", "windows_aarch64_gnullvm 0.53.1",
"windows_aarch64_msvc", "windows_aarch64_msvc 0.53.1",
"windows_i686_gnu", "windows_i686_gnu 0.53.1",
"windows_i686_gnullvm", "windows_i686_gnullvm 0.53.1",
"windows_i686_msvc", "windows_i686_msvc 0.53.1",
"windows_x86_64_gnu", "windows_x86_64_gnu 0.53.1",
"windows_x86_64_gnullvm", "windows_x86_64_gnullvm 0.53.1",
"windows_x86_64_msvc", "windows_x86_64_msvc 0.53.1",
] ]
[[package]]
name = "windows_aarch64_gnullvm"
version = "0.52.6"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "32a4622180e7a0ec044bb555404c800bc9fd9ec262ec147edd5989ccd0c02cd3"
[[package]] [[package]]
name = "windows_aarch64_gnullvm" name = "windows_aarch64_gnullvm"
version = "0.53.1" version = "0.53.1"
source = "registry+https://github.com/rust-lang/crates.io-index" source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "a9d8416fa8b42f5c947f8482c43e7d89e73a173cead56d044f6a56104a6d1b53" checksum = "a9d8416fa8b42f5c947f8482c43e7d89e73a173cead56d044f6a56104a6d1b53"
[[package]]
name = "windows_aarch64_msvc"
version = "0.52.6"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "09ec2a7bb152e2252b53fa7803150007879548bc709c039df7627cabbd05d469"
[[package]] [[package]]
name = "windows_aarch64_msvc" name = "windows_aarch64_msvc"
version = "0.53.1" version = "0.53.1"
source = "registry+https://github.com/rust-lang/crates.io-index" source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "b9d782e804c2f632e395708e99a94275910eb9100b2114651e04744e9b125006" checksum = "b9d782e804c2f632e395708e99a94275910eb9100b2114651e04744e9b125006"
[[package]]
name = "windows_i686_gnu"
version = "0.52.6"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "8e9b5ad5ab802e97eb8e295ac6720e509ee4c243f69d781394014ebfe8bbfa0b"
[[package]] [[package]]
name = "windows_i686_gnu" name = "windows_i686_gnu"
version = "0.53.1" version = "0.53.1"
source = "registry+https://github.com/rust-lang/crates.io-index" source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "960e6da069d81e09becb0ca57a65220ddff016ff2d6af6a223cf372a506593a3" checksum = "960e6da069d81e09becb0ca57a65220ddff016ff2d6af6a223cf372a506593a3"
[[package]]
name = "windows_i686_gnullvm"
version = "0.52.6"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "0eee52d38c090b3caa76c563b86c3a4bd71ef1a819287c19d586d7334ae8ed66"
[[package]] [[package]]
name = "windows_i686_gnullvm" name = "windows_i686_gnullvm"
version = "0.53.1" version = "0.53.1"
source = "registry+https://github.com/rust-lang/crates.io-index" source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "fa7359d10048f68ab8b09fa71c3daccfb0e9b559aed648a8f95469c27057180c" checksum = "fa7359d10048f68ab8b09fa71c3daccfb0e9b559aed648a8f95469c27057180c"
[[package]]
name = "windows_i686_msvc"
version = "0.52.6"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "240948bc05c5e7c6dabba28bf89d89ffce3e303022809e73deaefe4f6ec56c66"
[[package]] [[package]]
name = "windows_i686_msvc" name = "windows_i686_msvc"
version = "0.53.1" version = "0.53.1"
source = "registry+https://github.com/rust-lang/crates.io-index" source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "1e7ac75179f18232fe9c285163565a57ef8d3c89254a30685b57d83a38d326c2" checksum = "1e7ac75179f18232fe9c285163565a57ef8d3c89254a30685b57d83a38d326c2"
[[package]]
name = "windows_x86_64_gnu"
version = "0.52.6"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "147a5c80aabfbf0c7d901cb5895d1de30ef2907eb21fbbab29ca94c5b08b1a78"
[[package]] [[package]]
name = "windows_x86_64_gnu" name = "windows_x86_64_gnu"
version = "0.53.1" version = "0.53.1"
source = "registry+https://github.com/rust-lang/crates.io-index" source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "9c3842cdd74a865a8066ab39c8a7a473c0778a3f29370b5fd6b4b9aa7df4a499" checksum = "9c3842cdd74a865a8066ab39c8a7a473c0778a3f29370b5fd6b4b9aa7df4a499"
[[package]]
name = "windows_x86_64_gnullvm"
version = "0.52.6"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "24d5b23dc417412679681396f2b49f3de8c1473deb516bd34410872eff51ed0d"
[[package]] [[package]]
name = "windows_x86_64_gnullvm" name = "windows_x86_64_gnullvm"
version = "0.53.1" version = "0.53.1"
source = "registry+https://github.com/rust-lang/crates.io-index" source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "0ffa179e2d07eee8ad8f57493436566c7cc30ac536a3379fdf008f47f6bb7ae1" checksum = "0ffa179e2d07eee8ad8f57493436566c7cc30ac536a3379fdf008f47f6bb7ae1"
[[package]]
name = "windows_x86_64_msvc"
version = "0.52.6"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "589f6da84c646204747d1270a2a5661ea66ed1cced2631d546fdfb155959f9ec"
[[package]] [[package]]
name = "windows_x86_64_msvc" name = "windows_x86_64_msvc"
version = "0.53.1" version = "0.53.1"
@@ -1567,36 +1808,6 @@ version = "0.46.0"
source = "registry+https://github.com/rust-lang/crates.io-index" source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "f17a85883d4e6d00e8a97c586de764dabcc06133f7f1d55dce5cdc070ad7fe59" checksum = "f17a85883d4e6d00e8a97c586de764dabcc06133f7f1d55dce5cdc070ad7fe59"
[[package]]
name = "worker"
version = "0.1.0"
dependencies = [
"async-trait",
"clap",
"dotenv",
"eventsource-stream",
"futures",
"reqwest",
"schemars",
"serde",
"serde_json",
"tempfile",
"thiserror",
"tokio",
"tracing",
"tracing-subscriber",
"worker-macros",
]
[[package]]
name = "worker-macros"
version = "0.1.0"
dependencies = [
"proc-macro2",
"quote",
"syn",
]
[[package]] [[package]]
name = "writeable" name = "writeable"
version = "0.6.2" version = "0.6.2"
+2 -2
View File
@@ -1,8 +1,8 @@
[workspace] [workspace]
resolver = "2" resolver = "2"
members = [ members = [
"worker", "llm-worker",
"worker-macros", "llm-worker-macros",
] ]
[workspace.package] [workspace.package]
+32 -2
View File
@@ -1,5 +1,35 @@
# llm-worker-rs # llm-worker
Rusty, Efficient, and Agentic LLM Client Library Rusty, Efficient, and Agentic LLM Client Library
`llm-worker-rs` is a Rust library designed for building robust and efficient LLM applications. It unifies interactions with multiple LLM providers (Anthropic, Gemini, OpenAI, Ollama) under a single abstraction, with type-safe state management, efficient context caching, and a powerful event-driven architecture. `llm-worker` is a Rust library for building autonomous LLM-powered systems. Define tools, register hooks, and let the Worker handle the agentic loop — tool calls are executed automatically until the task completes.
## Features
- Autonomous Execution: The `Worker` manages the full request-response-tool cycle. You provide tools and input; it loops until done.
- Multi-Provider Support: Unified interface for Anthropic, Gemini, OpenAI, and Ollama.
- Tool System: Define tools as async functions. The Worker automatically parses LLM tool calls, executes them in parallel, and feeds results back.
- Hook System: Intercept execution flow with `before_tool_call`, `after_tool_call`, and `on_turn_end` hooks for validation, logging, or self-correction.
- Event-Driven Streaming: Subscribe to real-time events (text deltas, tool calls, usage) for responsive UIs.
- Cache-Aware State Management: Type-state pattern (`Mutable``CacheLocked`) ensures KV cache efficiency by protecting the conversation prefix.
## Quick Start
```rust
use llm_worker::{Worker, Message};
// Create a Worker with your LLM client
let mut worker = Worker::new(client)
.system_prompt("You are a helpful assistant.");
// Register tools (optional)
worker.register_tool(SearchTool::new());
worker.register_tool(CalculatorTool::new());
// Run — the Worker handles tool calls automatically
let history = worker.run("What is 2+2?").await?;
```
## License
MIT
+1 -1
View File
@@ -17,7 +17,7 @@ LLMを用いたワーカーを作成する小型のSDK・ライブラリ。
module構成概念図 module構成概念図
``` ```plaintext
worker worker
├── context ├── context
├── llm_client ├── llm_client
+3 -3
View File
@@ -27,7 +27,7 @@ RustのType-stateパターンを利用し、Workerの状態によって利用可
* 自由な編集が可能な状態。 * 自由な編集が可能な状態。
* システムプロンプトの設定・変更が可能。 * システムプロンプトの設定・変更が可能。
* メッセージ履歴の初期構築(ロード、編集)が可能。 * メッセージ履歴の初期構築(ロード、編集)が可能。
* **`Locked` (キャッシュ保護状態)** * **`CacheLocked` (キャッシュ保護状態)**
* キャッシュの有効活用を目的とした、前方不変状態。 * キャッシュの有効活用を目的とした、前方不変状態。
* **システムプロンプトの変更不可** * **システムプロンプトの変更不可**
* **既存メッセージ履歴の変更不可**(追記のみ許可)。 * **既存メッセージ履歴の変更不可**(追記のみ許可)。
@@ -47,7 +47,7 @@ worker.history_mut().push(initial_message);
// 3. ロックしてLocked状態へ遷移 // 3. ロックしてLocked状態へ遷移
// これにより、ここまでのコンテキストが "Fixed Prefix" として扱われる // これにより、ここまでのコンテキストが "Fixed Prefix" として扱われる
let mut locked_worker: Worker<Locked> = worker.lock(); let mut locked_worker: Worker<CacheLocked> = worker.lock();
// 4. 利用 (Locked状態) // 4. 利用 (Locked状態)
// 実行は可能。新しいメッセージは履歴の末尾に追記される。 // 実行は可能。新しいメッセージは履歴の末尾に追記される。
@@ -65,4 +65,4 @@ locked_worker.run(new_user_input).await?;
* **状態パラメータの導入**: `Worker<S: WorkerState>` の導入。 * **状態パラメータの導入**: `Worker<S: WorkerState>` の導入。
* **コンテキスト所有権の委譲**: `run` メソッドの引数でコンテキストを受け取るのではなく、`Worker` 内部に `history: Vec<Message>` を保持し管理する形へ移行する。 * **コンテキスト所有権の委譲**: `run` メソッドの引数でコンテキストを受け取るのではなく、`Worker` 内部に `history: Vec<Message>` を保持し管理する形へ移行する。
* **APIの分離**: `Mutable` 特有のメソッド(setter等)と、`Locked` でも使えるメソッド(実行、参照等)をトレイト境界で分離する。 * **APIの分離**: `Mutable` 特有のメソッド(setter等)と、`CacheLocked` でも使えるメソッド(実行、参照等)をトレイト境界で分離する。
+70
View File
@@ -0,0 +1,70 @@
# 非同期キャンセル設計
Workerの非同期キャンセル機構についての設計ドキュメント。
## 概要
`tokio::sync::mpsc`の通知チャネルを用いて、別タスクからWorkerの実行を安全にキャンセルできる。
```rust
let worker = Arc::new(Mutex::new(Worker::new(client)));
// 実行タスク
let w = worker.clone();
let handle = tokio::spawn(async move {
w.lock().await.run("prompt").await
});
// キャンセル
worker.lock().await.cancel();
```
## キャンセル時の処理フロー
```
キャンセル検知
timeline.abort_current_block() // 進行中ブロックの終端処理
run_on_abort_hooks("Cancelled") // on_abort フック呼び出し
Err(WorkerError::Cancelled) // エラー返却
```
## API
| メソッド | 説明 |
| ----------------- | ------------------------------ |
| `cancel()` | キャンセルをトリガー |
| `cancel_sender()` | キャンセル通知用のSenderを取得 |
## on_abort フック
`Hook::on_abort(&self, reason: &str)`がキャンセル時に呼ばれる。
クリーンアップ処理やログ記録に使用できる。
```rust
async fn on_abort(&self, reason: &str) -> Result<(), HookError> {
log::info!("Aborted: {}", reason);
Ok(())
}
```
呼び出しタイミング:
- `WorkerError::Cancelled` — reason: `"Cancelled"`
- `ControlFlow::Abort(reason)` — reason: フックが指定した理由
---
## 既知の問題
### on_abort の発火基準
`on_abort`**interrupt(中断)** された場合に必ず発火する。
interrupt の例:
- `WorkerError::Cancelled`(キャンセル)
- `WorkerError::Aborted`(フックによるAbort
- ストリーム/ツール/クライアント/Hook の各種エラーで処理が中断された場合
+205 -125
View File
@@ -3,7 +3,8 @@
## 概要 ## 概要
HookはWorker層でのターン制御に介入するためのメカニズムです。 HookはWorker層でのターン制御に介入するためのメカニズムです。
Claude CodeのHooks機能に着想を得ており、メッセージ送信・ツール実行・ターン終了の各ポイントで処理を差し込むことができます。
メッセージ送信・ツール実行・ターン終了等の各ポイントで処理を差し込むことができます。
## コンセプト ## コンセプト
@@ -11,120 +12,184 @@ Claude CodeのHooks機能に着想を得ており、メッセージ送信・ツ
- **Contextへのアクセス**: メッセージ履歴を読み書き可能 - **Contextへのアクセス**: メッセージ履歴を読み書き可能
- **非破壊的チェーン**: 複数のHookを登録順に実行、後続Hookへの影響を制御 - **非破壊的チェーン**: 複数のHookを登録順に実行、後続Hookへの影響を制御
## Hook一覧
| 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 ## Hook Trait
```rust ```rust
#[async_trait] #[async_trait]
pub trait WorkerHook: Send + Sync { pub trait Hook<E: HookEventKind>: Send + Sync {
/// メッセージ送信前 async fn call(&self, input: &mut E::Input) -> Result<E::Output, HookError>;
/// リクエストに含まれるメッセージリストを改変できる
async fn on_message_send(
&self,
context: &mut Vec<Message>,
) -> Result<ControlFlow, HookError> {
Ok(ControlFlow::Continue)
}
/// ツール実行前
/// 実行をキャンセルしたり、引数を書き換えることができる
async fn before_tool_call(
&self,
tool_call: &mut ToolCall,
) -> Result<ControlFlow, HookError> {
Ok(ControlFlow::Continue)
}
/// ツール実行後
/// 結果を書き換えたり、隠蔽したりできる
async fn after_tool_call(
&self,
tool_result: &mut ToolResult,
) -> Result<ControlFlow, HookError> {
Ok(ControlFlow::Continue)
}
/// ターン終了時
/// 生成されたメッセージを検査し、必要ならリトライを指示できる
async fn on_turn_end(
&self,
messages: &[Message],
) -> Result<TurnResult, HookError> {
Ok(TurnResult::Finish)
}
} }
``` ```
## 制御フロー型 ## 制御フロー型
### ControlFlow ### HookEventKind / Result
Hook処理の継続/中断を制御する列挙型 Hookイベントごとに入力/出力型を分離し、意味のない制御フローを排除する
```rust ```rust
pub enum ControlFlow { pub trait HookEventKind {
/// 処理を続行(後続Hookも実行) type Input;
type Output;
}
pub struct OnPromptSubmit;
pub struct PreLlmRequest;
pub struct PreToolCall;
pub struct PostToolCall;
pub struct OnTurnEnd;
pub struct OnAbort;
pub enum OnPromptSubmitResult {
Continue,
Cancel(String),
}
pub enum PreLlmRequestResult {
Continue,
Cancel(String),
}
pub enum PreToolCallResult {
Continue, Continue,
/// 現在の処理をスキップ(ツール実行をスキップ等)
Skip, Skip,
/// 処理全体を中断(エラーとして扱う)
Abort(String), Abort(String),
Pause,
}
pub enum PostToolCallResult {
Continue,
Abort(String),
}
pub enum OnTurnEndResult {
Finish,
ContinueWithMessages(Vec<Message>),
Paused,
} }
``` ```
### TurnResult ### Tool Call Context
ターン終了時の判定結果を表す列挙型 `pre_tool_call` / `post_tool_call` は、ツール実行の文脈を含む入力を受け取る
```rust ```rust
pub enum TurnResult { pub struct ToolCallContext {
/// ターンを正常終了 pub call: ToolCall,
Finish, pub meta: ToolMeta, // 不変メタデータ
/// メッセージを追加してターン継続(自己修正など) pub tool: Arc<dyn Tool>, // 状態アクセス用
ContinueWithMessages(Vec<Message>), }
pub struct PostToolCallContext {
pub call: ToolCall,
pub result: ToolResult,
pub meta: ToolMeta,
pub tool: Arc<dyn Tool>,
} }
``` ```
## 呼び出しタイミング ## 呼び出しタイミング
``` ```
Worker::run() ループ Worker::run(user_input)
├─▶ on_message_send ──────────────────────────────┐ ├─▶ on_prompt_submit ───────────────────────────┐
│ ユーザーメッセージの前処理・検証 │
│ (最初の1回のみ) │
│ │
└─▶ loop {
├─▶ pre_llm_request ──────────────────────│
│ コンテキストの改変、バリデーション、 │ │ コンテキストの改変、バリデーション、 │
│ システムプロンプト注入などが可能 │ │ システムプロンプト注入などが可能 │
│ (毎ターン実行) │
│ │ │ │
├─▶ LLMリクエスト送信 & ストリーム処理 │ ├─▶ LLMリクエスト送信 & ストリーム処理 │
│ │ │ │
├─▶ ツール呼び出しがある場合: │ ├─▶ ツール呼び出しがある場合: │
│ │ │ │ │ │
│ ├─▶ before_tool_call (各ツールごと・逐次) │ ├─▶ pre_tool_call (各ツールごと・逐次) │
│ │ 実行可否の判定、引数の改変 │ │ │ 実行可否の判定、引数の改変 │
│ │ │ │ │ │
│ ├─▶ ツール並列実行 (join_all) │ │ ├─▶ ツール並列実行 (join_all) │
│ │ │ │ │ │
│ └─▶ after_tool_call (各結果ごと・逐次) │ └─▶ post_tool_call (各結果ごと・逐次) │
│ 結果の確認、加工、ログ出力 │ │ 結果の確認、加工、ログ出力 │
│ │ │ │
├─▶ ツール結果をコンテキストに追加 → ループ先頭へ ├─▶ ツール結果をコンテキストに追加
│ → ループ先頭へ │
│ │ │ │
└─▶ ツールなしの場合: │ └─▶ ツールなしの場合: │
│ │ │ │
└─▶ on_turn_end ─────────────────────────────┘ └─▶ on_turn_end ───────────────────┘
最終応答のチェック(Lint/Fmt等) 最終応答のチェック(Lint/Fmt等)
エラーがあればContinueWithMessagesでリトライ エラーがあればContinueWithMessagesでリトライ
}
※ 中断時は on_abort が呼ばれる
``` ```
## 各Hookの詳細 ## 各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リクエスト送信前(ループの毎回)
**用途**:
- コンテキストへのシステムメッセージ注入 - コンテキストへのシステムメッセージ注入
- メッセージのバリデーション - メッセージのバリデーション
- 機密情報のフィルタリング - 機密情報のフィルタリング
- リクエスト内容のログ出力 - リクエスト内容のログ出力
- `PreLlmRequestResult::Cancel` による送信キャンセル
**入力**: `&mut Vec<Message>` - コンテキスト全体(改変可能)
**例**: メッセージにタイムスタンプを追加 **例**: メッセージにタイムスタンプを追加
@@ -132,27 +197,33 @@ Worker::run() ループ
struct TimestampHook; struct TimestampHook;
#[async_trait] #[async_trait]
impl WorkerHook for TimestampHook { impl Hook<PreLlmRequest> for TimestampHook {
async fn on_message_send( async fn call(
&self, &self,
context: &mut Vec<Message>, context: &mut Vec<Message>,
) -> Result<ControlFlow, HookError> { ) -> Result<PreLlmRequestResult, HookError> {
let timestamp = chrono::Local::now().to_rfc3339(); let timestamp = chrono::Local::now().to_rfc3339();
context.insert(0, Message::user(format!("[{}]", timestamp))); context.insert(0, Message::user(format!("[{}]", timestamp)));
Ok(ControlFlow::Continue) Ok(PreLlmRequestResult::Continue)
} }
} }
``` ```
### before_tool_call ### pre_tool_call
**呼び出しタイミング**: 各ツール実行前(並列実行フェーズの前) **呼び出しタイミング**: 各ツール実行前(並列実行フェーズの前)
**用途**: **用途**:
- 危険なツールのブロック - 危険なツールのブロック
- 引数のサニタイズ - 引数のサニタイズ
- 確認プロンプトの表示(UIとの連携) - 確認プロンプトの表示(UIとの連携)
- 実行ログの記録 - 実行ログの記録
- `PreToolCallResult::Pause` による一時停止
**入力**:
- `ToolCallContext``ToolCall` + `ToolMeta` + `Arc<dyn Tool>`
**例**: 特定ツールをブロック **例**: 特定ツールをブロック
@@ -162,46 +233,52 @@ struct ToolBlocker {
} }
#[async_trait] #[async_trait]
impl WorkerHook for ToolBlocker { impl Hook<PreToolCall> for ToolBlocker {
async fn before_tool_call( async fn call(
&self, &self,
tool_call: &mut ToolCall, ctx: &mut ToolCallContext,
) -> Result<ControlFlow, HookError> { ) -> Result<PreToolCallResult, HookError> {
if self.blocked_tools.contains(&tool_call.name) { if self.blocked_tools.contains(&ctx.call.name) {
println!("Blocked tool: {}", tool_call.name); println!("Blocked tool: {}", ctx.call.name);
Ok(ControlFlow::Skip) Ok(PreToolCallResult::Skip)
} else { } else {
Ok(ControlFlow::Continue) Ok(PreToolCallResult::Continue)
} }
} }
} }
``` ```
### after_tool_call ### post_tool_call
**呼び出しタイミング**: 各ツール実行後(並列実行フェーズの後) **呼び出しタイミング**: 各ツール実行後(並列実行フェーズの後)
**用途**: **用途**:
- 結果の加工・フォーマット - 結果の加工・フォーマット
- 機密情報のマスキング - 機密情報のマスキング
- 結果のキャッシュ - 結果のキャッシュ
- 実行結果のログ出力 - 実行結果のログ出力
**入力**:
- `PostToolCallContext``ToolCall` + `ToolResult` + `ToolMeta` +
`Arc<dyn Tool>`
**例**: 結果にプレフィックスを追加 **例**: 結果にプレフィックスを追加
```rust ```rust
struct ResultFormatter; struct ResultFormatter;
#[async_trait] #[async_trait]
impl WorkerHook for ResultFormatter { impl Hook<PostToolCall> for ResultFormatter {
async fn after_tool_call( async fn call(
&self, &self,
tool_result: &mut ToolResult, ctx: &mut PostToolCallContext,
) -> Result<ControlFlow, HookError> { ) -> Result<PostToolCallResult, HookError> {
if !tool_result.is_error { if !ctx.result.is_error {
tool_result.content = format!("[OK] {}", tool_result.content); ctx.result.content = format!("[OK] {}", ctx.result.content);
} }
Ok(ControlFlow::Continue) Ok(PostToolCallResult::Continue)
} }
} }
``` ```
@@ -211,10 +288,22 @@ impl WorkerHook for ResultFormatter {
**呼び出しタイミング**: ツール呼び出しなしでターンが終了する直前 **呼び出しタイミング**: ツール呼び出しなしでターンが終了する直前
**用途**: **用途**:
- 生成されたコードのLint/Fmt - 生成されたコードのLint/Fmt
- 出力形式のバリデーション - 出力形式のバリデーション
- 自己修正のためのリトライ指示 - 自己修正のためのリトライ指示
- 最終結果のログ出力 - 最終結果のログ出力
- `OnTurnEndResult::Paused` による一時停止
### on_abort
**呼び出しタイミング**: キャンセル/エラー/AbortなどでWorkerが中断された時
**用途**:
- クリーンアップ処理
- 中断理由のログ出力
- 外部システムへの通知
**例**: JSON形式のバリデーション **例**: JSON形式のバリデーション
@@ -222,11 +311,11 @@ impl WorkerHook for ResultFormatter {
struct JsonValidator; struct JsonValidator;
#[async_trait] #[async_trait]
impl WorkerHook for JsonValidator { impl Hook<OnTurnEnd> for JsonValidator {
async fn on_turn_end( async fn call(
&self, &self,
messages: &[Message], messages: &mut Vec<Message>,
) -> Result<TurnResult, HookError> { ) -> Result<OnTurnEndResult, HookError> {
// 最後のアシスタントメッセージを取得 // 最後のアシスタントメッセージを取得
let last = messages.iter().rev() let last = messages.iter().rev()
.find(|m| m.role == Role::Assistant); .find(|m| m.role == Role::Assistant);
@@ -236,25 +325,25 @@ impl WorkerHook for JsonValidator {
// JSONとしてパースを試みる // JSONとしてパースを試みる
if serde_json::from_str::<serde_json::Value>(text).is_err() { if serde_json::from_str::<serde_json::Value>(text).is_err() {
// 失敗したらリトライ指示 // 失敗したらリトライ指示
return Ok(TurnResult::ContinueWithMessages(vec![ return Ok(OnTurnEndResult::ContinueWithMessages(vec![
Message::user("Invalid JSON. Please fix and try again.") Message::user("Invalid JSON. Please fix and try again.")
])); ]));
} }
} }
} }
Ok(TurnResult::Finish) Ok(OnTurnEndResult::Finish)
} }
} }
``` ```
## 複数Hookの実行順序 ## 複数Hookの実行順序
Hookは**登録順**に実行されます。 Hookは**イベントごとに登録順**に実行されます。
```rust ```rust
worker.add_hook(HookA); // 1番目に実行 worker.add_pre_tool_call_hook(HookA); // 1番目に実行
worker.add_hook(HookB); // 2番目に実行 worker.add_pre_tool_call_hook(HookB); // 2番目に実行
worker.add_hook(HookC); // 3番目に実行 worker.add_pre_tool_call_hook(HookC); // 3番目に実行
``` ```
### 制御フローの伝播 ### 制御フローの伝播
@@ -262,6 +351,7 @@ worker.add_hook(HookC); // 3番目に実行
- `Continue`: 後続Hookも実行 - `Continue`: 後続Hookも実行
- `Skip`: 現在の処理をスキップし、後続Hookは実行しない - `Skip`: 現在の処理をスキップし、後続Hookは実行しない
- `Abort`: 即座にエラーを返し、処理全体を中断 - `Abort`: 即座にエラーを返し、処理全体を中断
- `Pause`: Workerを一時停止(再開は`resume`
``` ```
Hook A: Continue → Hook B: Skip → (Hook Cは実行されない) Hook A: Continue → Hook B: Skip → (Hook Cは実行されない)
@@ -271,52 +361,39 @@ Hook A: Continue → Hook B: Skip → (Hook Cは実行されない)
Hook A: Continue → Hook B: Abort("reason") Hook A: Continue → Hook B: Abort("reason")
WorkerError::Aborted WorkerError::Aborted
Hook A: Continue → Hook B: Pause
WorkerResult::Paused
``` ```
## 設計上のポイント ## 設計上のポイント
### 1. デフォルト実装 ### 1. イベントごとの実装
全メソッドにデフォルト実装があるため、必要なメソッドだけオーバーライドすれば良い 必要なイベントのみ `Hook<Event>` を実装する
```rust
struct SimpleLogger;
#[async_trait]
impl WorkerHook for SimpleLogger {
// on_message_send だけ実装
async fn on_message_send(
&self,
context: &mut Vec<Message>,
) -> Result<ControlFlow, HookError> {
println!("Sending {} messages", context.len());
Ok(ControlFlow::Continue)
}
// 他のメソッドはデフォルト(Continue/Finish
}
```
### 2. 可変参照による改変 ### 2. 可変参照による改変
`&mut`で引数を受け取るため、直接改変が可能。 `&mut`で引数を受け取るため、直接改変が可能。
```rust ```rust
async fn before_tool_call(&self, tool_call: &mut ToolCall) -> ... { async fn call(&self, ctx: &mut ToolCallContext) -> ... {
// 引数を直接書き換え // 引数を直接書き換え
tool_call.input["sanitized"] = json!(true); ctx.call.input["sanitized"] = json!(true);
Ok(ControlFlow::Continue) Ok(PreToolCallResult::Continue)
} }
``` ```
### 3. 並列実行との統合 ### 3. 並列実行との統合
- `before_tool_call`: 並列実行**前**に逐次実行(許可判定のため) - `pre_tool_call`: 並列実行**前**に逐次実行(許可判定のため)
- ツール実行: `join_all`で**並列**実行 - ツール実行: `join_all`で**並列**実行
- `after_tool_call`: 並列実行**後**に逐次実行(結果加工のため) - `post_tool_call`: 並列実行**後**に逐次実行(結果加工のため)
### 4. Send + Sync 要件 ### 4. Send + Sync 要件
`WorkerHook``Send + Sync`を要求するため、スレッドセーフな実装が必要。 `Hook``Send + Sync`を要求するため、スレッドセーフな実装が必要。
状態を持つ場合は`Arc<Mutex<T>>``AtomicUsize`などを使用する。 状態を持つ場合は`Arc<Mutex<T>>``AtomicUsize`などを使用する。
```rust ```rust
@@ -325,10 +402,10 @@ struct CountingHook {
} }
#[async_trait] #[async_trait]
impl WorkerHook for CountingHook { impl Hook<PreToolCall> for CountingHook {
async fn before_tool_call(&self, _: &mut ToolCall) -> Result<ControlFlow, HookError> { async fn call(&self, _: &mut ToolCallContext) -> Result<PreToolCallResult, HookError> {
self.count.fetch_add(1, Ordering::SeqCst); self.count.fetch_add(1, Ordering::SeqCst);
Ok(ControlFlow::Continue) Ok(PreToolCallResult::Continue)
} }
} }
``` ```
@@ -336,13 +413,13 @@ impl WorkerHook for CountingHook {
## 典型的なユースケース ## 典型的なユースケース
| ユースケース | 使用Hook | 処理内容 | | ユースケース | 使用Hook | 処理内容 |
|-------------|----------|----------| | ------------------ | -------------------- | -------------------------- |
| ツール許可制御 | `before_tool_call` | 危険なツールをSkip | | ツール許可制御 | `pre_tool_call` | 危険なツールをSkip |
| 実行ログ | `before/after_tool_call` | 呼び出しと結果を記録 | | 実行ログ | `pre/post_tool_call` | 呼び出しと結果を記録 |
| 出力バリデーション | `on_turn_end` | 形式チェック、リトライ指示 | | 出力バリデーション | `on_turn_end` | 形式チェック、リトライ指示 |
| コンテキスト注入 | `on_message_send` | システムメッセージ追加 | | コンテキスト注入 | `on_message_send` | システムメッセージ追加 |
| 結果のサニタイズ | `after_tool_call` | 機密情報のマスキング | | 結果のサニタイズ | `post_tool_call` | 機密情報のマスキング |
| レート制限 | `before_tool_call` | 呼び出し頻度の制御 | | レート制限 | `pre_tool_call` | 呼び出し頻度の制御 |
## TODO ## TODO
@@ -350,11 +427,14 @@ impl WorkerHook for CountingHook {
現在のHooks実装は基本的なユースケースをカバーしているが、以下の点について将来的に厳密な仕様を定義する必要がある: 現在のHooks実装は基本的なユースケースをカバーしているが、以下の点について将来的に厳密な仕様を定義する必要がある:
- **エラーハンドリングの明確化**: `HookError`発生時のリカバリー戦略、部分的な失敗の扱い - **エラーハンドリングの明確化**:
`HookError`発生時のリカバリー戦略、部分的な失敗の扱い
- **Hook間の依存関係**: 複数Hookの実行順序が結果に影響する場合のセマンティクス - **Hook間の依存関係**: 複数Hookの実行順序が結果に影響する場合のセマンティクス
- **非同期キャンセル**: Hook実行中のキャンセル(タイムアウト等)の振る舞い - **非同期キャンセル**: Hook実行中のキャンセル(タイムアウト等)の振る舞い
- **状態の一貫性**: `on_message_send`で改変されたコンテキストが後続処理で期待通りに反映される保証 - **状態の一貫性**:
- **リトライ制限**: `on_turn_end`での`ContinueWithMessages`による無限ループ防止策 `on_message_send`で改変されたコンテキストが後続処理で期待通りに反映される保証
- **リトライ制限**:
`on_turn_end`での`ContinueWithMessages`による無限ループ防止策
- **Hook優先度**: 登録順以外の優先度指定メカニズムの必要性 - **Hook優先度**: 登録順以外の優先度指定メカニズムの必要性
- **条件付きHook**: 特定条件でのみ有効化されるHookパターン - **条件付きHook**: 特定条件でのみ有効化されるHookパターン
- **テスト容易性**: Hookのモック/スタブ作成のためのユーティリティ - **テスト容易性**: Hookのモック/スタブ作成のためのユーティリティ
+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 を含む | 柔軟な介入を可能に |
+41 -23
View File
@@ -178,41 +178,60 @@ Workerは生成されたラッパー構造体を `Box<dyn Tool>` として保持
```rust ```rust
#[async_trait] #[async_trait]
pub trait WorkerHook: Send + Sync { pub trait Hook<E: HookEventKind>: Send + Sync {
/// メッセージ送信前。 async fn call(&self, input: &mut E::Input) -> Result<E::Output, Error>;
/// リクエストに含まれるメッセージリストを改変できる。
async fn on_message_send(&self, context: &mut Vec<Message>) -> Result<ControlFlow, Error> {
Ok(ControlFlow::Continue)
} }
/// ツール実行前。 pub trait HookEventKind {
/// 実行をキャンセルしたり、引数を書き換えることができる。 type Input;
async fn before_tool_call(&self, tool_call: &mut ToolCall) -> Result<ControlFlow, Error> { type Output;
Ok(ControlFlow::Continue)
} }
/// ツール実行後。 pub struct OnMessageSend;
/// 結果を書き換えたり、隠蔽したりできる。 pub struct BeforeToolCall;
async fn after_tool_call(&self, tool_result: &mut ToolResult) -> Result<ControlFlow, Error> { pub struct AfterToolCall;
Ok(ControlFlow::Continue) pub struct OnTurnEnd;
pub struct OnAbort;
pub enum OnMessageSendResult {
Continue,
Cancel(String),
} }
/// ターン終了時。 pub enum BeforeToolCallResult {
/// 生成されたメッセージを検査し、必要ならリトライ(ContinueWithMessages)を指示できる。
async fn on_turn_end(&self, messages: &[Message]) -> Result<TurnResult, Error> {
Ok(TurnResult::Finish)
}
}
pub enum ControlFlow {
Continue, Continue,
Skip, // Tool実行などをスキップ Skip, // Tool実行などをスキップ
Abort(String), // 処理中断 Abort(String), // 処理中断
Pause,
} }
pub enum TurnResult { pub enum AfterToolCallResult {
Continue,
Abort(String),
}
pub enum OnTurnEndResult {
Finish, Finish,
ContinueWithMessages(Vec<Message>), // メッセージを追加してターン継続(自己修正など) ContinueWithMessages(Vec<Message>), // メッセージを追加してターン継続(自己修正など)
Paused,
}
```
### Tool Call Context
`before_tool_call` / `after_tool_call` は、ツール実行の文脈を含む入力を受け取る。
```rust
pub struct ToolCallContext {
pub call: ToolCall,
pub meta: ToolMeta, // 不変メタデータ
pub tool: Arc<dyn Tool>, // 状態アクセス用
}
pub struct ToolResultContext {
pub result: ToolResult,
pub meta: ToolMeta,
pub tool: Arc<dyn Tool>,
} }
``` ```
@@ -433,4 +452,3 @@ impl<C: LlmClient> Worker<C> {
3. **選択的購読**: on_*で必要なイベントだけ、またはSubscriberで一括 3. **選択的購読**: on_*で必要なイベントだけ、またはSubscriberで一括
4. **累積イベントの追加**: Worker層でComplete系イベントを追加提供 4. **累積イベントの追加**: Worker層でComplete系イベントを追加提供
5. **後方互換性**: 従来の`run()`も引き続き使用可能 5. **後方互換性**: 従来の`run()`も引き続き使用可能
@@ -1,6 +1,7 @@
[package] [package]
name = "worker-macros" name = "llm-worker-macros"
version = "0.1.0" description = "llm-worker's proc macros"
version = "0.2.0"
publish.workspace = true publish.workspace = true
edition.workspace = true edition.workspace = true
license.workspace = true license.workspace = true
@@ -1,4 +1,4 @@
//! worker-macros - Tool生成用手続きマクロ //! llm-worker-macros - Tool生成用手続きマクロ
//! //!
//! `#[tool_registry]` と `#[tool]` マクロを提供し、 //! `#[tool_registry]` と `#[tool]` マクロを提供し、
//! ユーザー定義のメソッドから `Tool` トレイト実装を自動生成する。 //! ユーザー定義のメソッドから `Tool` トレイト実装を自動生成する。
@@ -113,7 +113,7 @@ fn generate_tool_impl(self_ty: &Type, method: &syn::ImplItemFn) -> proc_macro2::
let pascal_name = to_pascal_case(&method_name.to_string()); let pascal_name = to_pascal_case(&method_name.to_string());
let tool_struct_name = format_ident!("Tool{}", pascal_name); let tool_struct_name = format_ident!("Tool{}", pascal_name);
let args_struct_name = format_ident!("{}Args", pascal_name); let args_struct_name = format_ident!("{}Args", pascal_name);
let factory_name = format_ident!("{}_tool", method_name); let definition_name = format_ident!("{}_definition", method_name);
// ドキュメントコメントから説明を取得 // ドキュメントコメントから説明を取得
let description = extract_doc_comment(&method.attrs); let description = extract_doc_comment(&method.attrs);
@@ -193,7 +193,7 @@ fn generate_tool_impl(self_ty: &Type, method: &syn::ImplItemFn) -> proc_macro2::
quote! { quote! {
match result { match result {
Ok(val) => Ok(format!("{:?}", val)), Ok(val) => Ok(format!("{:?}", val)),
Err(e) => Err(::worker::tool::ToolError::ExecutionFailed(format!("{}", e))), Err(e) => Err(::llm_worker::tool::ToolError::ExecutionFailed(format!("{}", e))),
} }
} }
} else { } else {
@@ -230,7 +230,7 @@ fn generate_tool_impl(self_ty: &Type, method: &syn::ImplItemFn) -> proc_macro2::
} else { } else {
quote! { quote! {
let args: #args_struct_name = serde_json::from_str(input_json) let args: #args_struct_name = serde_json::from_str(input_json)
.map_err(|e| ::worker::tool::ToolError::InvalidArgument(e.to_string()))?; .map_err(|e| ::llm_worker::tool::ToolError::InvalidArgument(e.to_string()))?;
let result = self.ctx.#method_name(#(#arg_names),*)#awaiter; let result = self.ctx.#method_name(#(#arg_names),*)#awaiter;
#result_handling #result_handling
@@ -246,30 +246,25 @@ fn generate_tool_impl(self_ty: &Type, method: &syn::ImplItemFn) -> proc_macro2::
} }
#[async_trait::async_trait] #[async_trait::async_trait]
impl ::worker::tool::Tool for #tool_struct_name { impl ::llm_worker::tool::Tool for #tool_struct_name {
fn name(&self) -> &str { async fn execute(&self, input_json: &str) -> Result<String, ::llm_worker::tool::ToolError> {
#tool_name
}
fn description(&self) -> &str {
#description
}
fn input_schema(&self) -> serde_json::Value {
let schema = schemars::schema_for!(#args_struct_name);
serde_json::to_value(schema).unwrap_or(serde_json::json!({}))
}
async fn execute(&self, input_json: &str) -> Result<String, ::worker::tool::ToolError> {
#execute_body #execute_body
} }
} }
impl #self_ty { impl #self_ty {
pub fn #factory_name(&self) -> #tool_struct_name { /// ToolDefinition を取得(Worker への登録用)
#tool_struct_name { pub fn #definition_name(&self) -> ::llm_worker::tool::ToolDefinition {
ctx: self.clone() let ctx = self.clone();
} ::std::sync::Arc::new(move || {
let schema = schemars::schema_for!(#args_struct_name);
let meta = ::llm_worker::tool::ToolMeta::new(#tool_name)
.description(#description)
.input_schema(serde_json::to_value(schema).unwrap_or(serde_json::json!({})));
let tool: ::std::sync::Arc<dyn ::llm_worker::tool::Tool> =
::std::sync::Arc::new(#tool_struct_name { ctx: ctx.clone() });
(meta, tool)
})
} }
} }
} }
+6 -5
View File
@@ -1,7 +1,7 @@
[package] [package]
name = "worker" name = "llm-worker"
description = "" description = "A library for building autonomous LLM-powered systems"
version = "0.1.0" version = "0.2.0"
publish.workspace = true publish.workspace = true
edition.workspace = true edition.workspace = true
license.workspace = true license.workspace = true
@@ -15,9 +15,10 @@ tracing = "0.1"
async-trait = "0.1" async-trait = "0.1"
futures = "0.3" futures = "0.3"
tokio = { version = "1.49", features = ["macros", "rt-multi-thread"] } tokio = { version = "1.49", features = ["macros", "rt-multi-thread"] }
reqwest = { version = "0.13.1", default-features = false, features = ["stream", "json", "native-tls"] } tokio-util = "0.7"
reqwest = { version = "0.13.1", default-features = false, features = ["stream", "json", "native-tls", "http2"] }
eventsource-stream = "0.2" eventsource-stream = "0.2"
worker-macros = { path = "../worker-macros", version = "0.1" } llm-worker-macros = { path = "../llm-worker-macros", version = "0.2" }
[dev-dependencies] [dev-dependencies]
clap = { version = "4.5", features = ["derive", "env"] } clap = { version = "4.5", features = ["derive", "env"] }
@@ -20,9 +20,9 @@ mod recorder;
mod scenarios; mod scenarios;
use clap::{Parser, ValueEnum}; use clap::{Parser, ValueEnum};
use worker::llm_client::providers::anthropic::AnthropicClient; use llm_worker::llm_client::providers::anthropic::AnthropicClient;
use worker::llm_client::providers::gemini::GeminiClient; use llm_worker::llm_client::providers::gemini::GeminiClient;
use worker::llm_client::providers::openai::OpenAIClient; use llm_worker::llm_client::providers::openai::OpenAIClient;
#[derive(Parser, Debug)] #[derive(Parser, Debug)]
#[command(author, version, about, long_about = None)] #[command(author, version, about, long_about = None)]
@@ -101,7 +101,7 @@ async fn run_scenario_with_ollama(
subdir: &str, subdir: &str,
model: Option<String>, model: Option<String>,
) -> Result<(), Box<dyn std::error::Error>> { ) -> Result<(), Box<dyn std::error::Error>> {
use worker::llm_client::providers::ollama::OllamaClient; use llm_worker::llm_client::providers::ollama::OllamaClient;
// Ollama typically runs local, no key needed or placeholder // Ollama typically runs local, no key needed or placeholder
let model = model.as_deref().unwrap_or("llama3"); // default example let model = model.as_deref().unwrap_or("llama3"); // default example
let client = OllamaClient::new(model); // base_url placeholder, handled by client default let client = OllamaClient::new(model); // base_url placeholder, handled by client default
@@ -8,7 +8,7 @@ use std::path::Path;
use std::time::{Instant, SystemTime, UNIX_EPOCH}; use std::time::{Instant, SystemTime, UNIX_EPOCH};
use futures::StreamExt; use futures::StreamExt;
use worker::llm_client::{LlmClient, Request}; use llm_worker::llm_client::{LlmClient, Request};
/// 記録されたイベント /// 記録されたイベント
#[derive(Debug, serde::Serialize, serde::Deserialize)] #[derive(Debug, serde::Serialize, serde::Deserialize)]
@@ -2,7 +2,7 @@
//! //!
//! 各シナリオのリクエストと出力ファイル名を定義 //! 各シナリオのリクエストと出力ファイル名を定義
use worker::llm_client::{Request, ToolDefinition}; use llm_worker::llm_client::{Request, ToolDefinition};
/// テストシナリオ /// テストシナリオ
pub struct TestScenario { pub struct TestScenario {
+71
View File
@@ -0,0 +1,71 @@
//! Worker のキャンセル機能のデモンストレーション
//!
//! ストリーミング受信中に別スレッドからキャンセルする例
use llm_worker::llm_client::providers::anthropic::AnthropicClient;
use llm_worker::{Worker, WorkerResult};
use std::sync::Arc;
use std::time::Duration;
use tokio::sync::Mutex;
#[tokio::main]
async fn main() -> Result<(), Box<dyn std::error::Error>> {
// .envファイルを読み込む
dotenv::dotenv().ok();
// ロギング初期化
tracing_subscriber::fmt()
.with_env_filter(
tracing_subscriber::EnvFilter::try_from_default_env()
.unwrap_or_else(|_| tracing_subscriber::EnvFilter::new("info")),
)
.init();
let api_key =
std::env::var("ANTHROPIC_API_KEY").expect("ANTHROPIC_API_KEY environment variable not set");
let client = AnthropicClient::new(&api_key, "claude-sonnet-4-20250514");
let worker = Arc::new(Mutex::new(Worker::new(client)));
println!("🚀 Starting Worker...");
println!("💡 Will cancel after 2 seconds\n");
// キャンセルSenderを先に取得(ロックを保持しない)
let cancel_tx = {
let w = worker.lock().await;
w.cancel_sender()
};
// タスク1: Workerを実行
let worker_clone = worker.clone();
let task = tokio::spawn(async move {
let mut w = worker_clone.lock().await;
println!("📡 Sending request to LLM...");
match w.run("Tell me a very long story about a brave knight. Make it as detailed as possible with many paragraphs.").await {
Ok(WorkerResult::Finished) => {
println!("✅ Task completed normally");
}
Ok(WorkerResult::Paused) => {
println!("⏸️ Task paused");
}
Err(e) => {
println!("❌ Task error: {}", e);
}
}
});
// タスク2: 2秒後にキャンセル
tokio::spawn(async move {
tokio::time::sleep(Duration::from_secs(2)).await;
println!("\n🛑 Cancelling worker...");
let _ = cancel_tx.send(()).await;
});
// タスク完了を待つ
task.await?;
println!("\n✨ Demo complete!");
Ok(())
}
@@ -39,9 +39,9 @@ use tracing::info;
use tracing_subscriber::EnvFilter; use tracing_subscriber::EnvFilter;
use clap::{Parser, ValueEnum}; use clap::{Parser, ValueEnum};
use worker::{ use llm_worker::{
Worker, Worker,
hook::{ControlFlow, HookError, ToolResult, WorkerHook}, hook::{Hook, HookError, PostToolCall, PostToolCallContext, PostToolCallResult},
llm_client::{ llm_client::{
LlmClient, LlmClient,
providers::{ providers::{
@@ -51,7 +51,7 @@ use worker::{
}, },
timeline::{Handler, TextBlockEvent, TextBlockKind, ToolUseBlockEvent, ToolUseBlockKind}, timeline::{Handler, TextBlockEvent, TextBlockKind, ToolUseBlockEvent, ToolUseBlockKind},
}; };
use worker_macros::tool_registry; use llm_worker_macros::tool_registry;
// 必要なマクロ展開用インポート // 必要なマクロ展開用インポート
use schemars; use schemars;
@@ -282,25 +282,22 @@ impl ToolResultPrinterHook {
} }
#[async_trait] #[async_trait]
impl WorkerHook for ToolResultPrinterHook { impl Hook<PostToolCall> for ToolResultPrinterHook {
async fn after_tool_call( async fn call(&self, ctx: &mut PostToolCallContext) -> Result<PostToolCallResult, HookError> {
&self,
tool_result: &mut ToolResult,
) -> Result<ControlFlow, HookError> {
let name = self let name = self
.call_names .call_names
.lock() .lock()
.unwrap() .unwrap()
.remove(&tool_result.tool_use_id) .remove(&ctx.result.tool_use_id)
.unwrap_or_else(|| tool_result.tool_use_id.clone()); .unwrap_or_else(|| ctx.result.tool_use_id.clone());
if tool_result.is_error { if ctx.result.is_error {
println!(" Result ({}): ❌ {}", name, tool_result.content); println!(" Result ({}): ❌ {}", name, ctx.result.content);
} else { } else {
println!(" Result ({}): ✅ {}", name, tool_result.content); println!(" Result ({}): ✅ {}", name, ctx.result.content);
} }
Ok(ControlFlow::Continue) Ok(PostToolCallResult::Continue)
} }
} }
@@ -441,8 +438,10 @@ async fn main() -> Result<(), Box<dyn std::error::Error>> {
// ツール登録(--no-tools でなければ) // ツール登録(--no-tools でなければ)
if !args.no_tools { if !args.no_tools {
let app = AppContext; let app = AppContext;
worker.register_tool(app.get_current_time_tool()); worker
worker.register_tool(app.calculate_tool()); .register_tool(app.get_current_time_definition())
.unwrap();
worker.register_tool(app.calculate_definition()).unwrap();
} }
// ストリーミング表示用ハンドラーを登録 // ストリーミング表示用ハンドラーを登録
@@ -451,7 +450,7 @@ async fn main() -> Result<(), Box<dyn std::error::Error>> {
.on_text_block(StreamingPrinter::new()) .on_text_block(StreamingPrinter::new())
.on_tool_use_block(ToolCallPrinter::new(tool_call_names.clone())); .on_tool_use_block(ToolCallPrinter::new(tool_call_names.clone()));
worker.add_hook(ToolResultPrinterHook::new(tool_call_names)); worker.add_post_tool_call_hook(ToolResultPrinterHook::new(tool_call_names));
// ワンショットモード // ワンショットモード
if let Some(prompt) = args.prompt { if let Some(prompt) = args.prompt {
@@ -32,7 +32,7 @@ pub trait Kind {
/// # Examples /// # Examples
/// ///
/// ```ignore /// ```ignore
/// use worker::timeline::{Handler, TextBlockEvent, TextBlockKind}; /// use llm_worker::timeline::{Handler, TextBlockEvent, TextBlockKind};
/// ///
/// struct TextCollector { /// struct TextCollector {
/// texts: Vec<String>, /// texts: Vec<String>,
+233
View File
@@ -0,0 +1,233 @@
//! Hook関連の型定義
//!
//! Worker層でのターン制御・介入に使用される型
use async_trait::async_trait;
use serde::{Deserialize, Serialize};
use serde_json::Value;
use thiserror::Error;
// =============================================================================
// Hook Event Kinds
// =============================================================================
pub trait HookEventKind: Send + Sync + 'static {
type Input;
type Output;
}
pub struct OnPromptSubmit;
pub struct PreLlmRequest;
pub struct PreToolCall;
pub struct PostToolCall;
pub struct OnTurnEnd;
pub struct OnAbort;
#[derive(Debug, Clone, PartialEq, Eq)]
pub enum OnPromptSubmitResult {
Continue,
Cancel(String),
}
#[derive(Debug, Clone, PartialEq, Eq)]
pub enum PreLlmRequestResult {
Continue,
Cancel(String),
}
#[derive(Debug, Clone, PartialEq, Eq)]
pub enum PreToolCallResult {
Continue,
Skip,
Abort(String),
Pause,
}
#[derive(Debug, Clone, PartialEq, Eq)]
pub enum PostToolCallResult {
Continue,
Abort(String),
}
#[derive(Debug, Clone)]
pub enum OnTurnEndResult {
Finish,
ContinueWithMessages(Vec<crate::Message>),
Paused,
}
use std::sync::Arc;
use crate::tool::{Tool, ToolMeta};
/// PreToolCall の入力コンテキスト
pub struct ToolCallContext {
/// ツール呼び出し情報(改変可能)
pub call: ToolCall,
/// ツールメタ情報(不変)
pub meta: ToolMeta,
/// ツールインスタンス(状態アクセス用)
pub tool: Arc<dyn Tool>,
}
/// PostToolCall の入力コンテキスト
pub struct PostToolCallContext {
/// ツール呼び出し情報
pub call: ToolCall,
/// ツール実行結果(改変可能)
pub result: ToolResult,
/// ツールメタ情報(不変)
pub meta: ToolMeta,
/// ツールインスタンス(状態アクセス用)
pub tool: Arc<dyn Tool>,
}
impl HookEventKind for OnPromptSubmit {
type Input = crate::Message;
type Output = OnPromptSubmitResult;
}
impl HookEventKind for PreLlmRequest {
type Input = Vec<crate::Message>;
type Output = PreLlmRequestResult;
}
impl HookEventKind for PreToolCall {
type Input = ToolCallContext;
type Output = PreToolCallResult;
}
impl HookEventKind for PostToolCall {
type Input = PostToolCallContext;
type Output = PostToolCallResult;
}
impl HookEventKind for OnTurnEnd {
type Input = Vec<crate::Message>;
type Output = OnTurnEndResult;
}
impl HookEventKind for OnAbort {
type Input = String;
type Output = ();
}
// =============================================================================
// Tool Call / Result Types
// =============================================================================
/// ツール呼び出し情報
///
/// LLMからのToolUseブロックを表現し、Hook処理で改変可能
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct ToolCall {
/// ツール呼び出しID(レスポンスとの紐付けに使用)
pub id: String,
/// ツール名
pub name: String,
/// 入力引数(JSON
pub input: Value,
}
/// ツール実行結果
///
/// ツール実行後の結果を表現し、Hook処理で改変可能
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct ToolResult {
/// 対応するツール呼び出しID
pub tool_use_id: String,
/// 結果コンテンツ
pub content: String,
/// エラーかどうか
#[serde(default)]
pub is_error: bool,
}
impl ToolResult {
/// 成功結果を作成
pub fn success(tool_use_id: impl Into<String>, content: impl Into<String>) -> Self {
Self {
tool_use_id: tool_use_id.into(),
content: content.into(),
is_error: false,
}
}
/// エラー結果を作成
pub fn error(tool_use_id: impl Into<String>, content: impl Into<String>) -> Self {
Self {
tool_use_id: tool_use_id.into(),
content: content.into(),
is_error: true,
}
}
}
// =============================================================================
// Hook Error
// =============================================================================
/// Hookエラー
#[derive(Debug, Error)]
pub enum HookError {
/// 処理が中断された
#[error("Aborted: {0}")]
Aborted(String),
/// 内部エラー
#[error("Hook error: {0}")]
Internal(String),
}
// =============================================================================
// Hook Trait
// =============================================================================
/// Hookイベントの処理を行うトレイト
///
/// 各イベント種別は戻り値型が異なるため、`HookEventKind`を介して型を制約する。
#[async_trait]
pub trait Hook<E: HookEventKind>: Send + Sync {
async fn call(&self, input: &mut E::Input) -> Result<E::Output, HookError>;
}
// =============================================================================
// Hook Registry
// =============================================================================
/// 全 Hook を保持するレジストリ
///
/// Worker 内部で使用され、各種 Hook を一括管理する。
pub struct HookRegistry {
/// on_prompt_submit Hook
pub(crate) on_prompt_submit: Vec<Box<dyn Hook<OnPromptSubmit>>>,
/// pre_llm_request Hook
pub(crate) pre_llm_request: Vec<Box<dyn Hook<PreLlmRequest>>>,
/// pre_tool_call Hook
pub(crate) pre_tool_call: Vec<Box<dyn Hook<PreToolCall>>>,
/// post_tool_call Hook
pub(crate) post_tool_call: Vec<Box<dyn Hook<PostToolCall>>>,
/// on_turn_end Hook
pub(crate) on_turn_end: Vec<Box<dyn Hook<OnTurnEnd>>>,
/// on_abort Hook
pub(crate) on_abort: Vec<Box<dyn Hook<OnAbort>>>,
}
impl Default for HookRegistry {
fn default() -> Self {
Self::new()
}
}
impl HookRegistry {
/// 空の HookRegistry を作成
pub fn new() -> Self {
Self {
on_prompt_submit: Vec::new(),
pre_llm_request: Vec::new(),
pre_tool_call: Vec::new(),
post_tool_call: Vec::new(),
on_turn_end: Vec::new(),
on_abort: Vec::new(),
}
}
}
+12 -17
View File
@@ -1,4 +1,4 @@
//! worker - LLMワーカーライブラリ //! llm-worker - LLMワーカーライブラリ
//! //!
//! LLMとの対話を管理するコンポーネントを提供します。 //! LLMとの対話を管理するコンポーネントを提供します。
//! //!
@@ -6,21 +6,20 @@
//! //!
//! - [`Worker`] - LLMとの対話を管理する中心コンポーネント //! - [`Worker`] - LLMとの対話を管理する中心コンポーネント
//! - [`tool::Tool`] - LLMから呼び出し可能なツール //! - [`tool::Tool`] - LLMから呼び出し可能なツール
//! - [`hook::WorkerHook`] - ターン進行への介入 //! - [`hook::Hook`] - ターン進行への介入
//! - [`subscriber::WorkerSubscriber`] - ストリーミングイベントの購読 //! - [`subscriber::WorkerSubscriber`] - ストリーミングイベントの購読
//! //!
//! # Quick Start //! # Quick Start
//! //!
//! ```ignore //! ```ignore
//! use worker::{Worker, Message}; //! use llm_worker::{Worker, Message};
//! //!
//! // Workerを作成 //! // Workerを作成
//! let mut worker = Worker::new(client) //! let mut worker = Worker::new(client)
//! .system_prompt("You are a helpful assistant."); //! .system_prompt("You are a helpful assistant.");
//! //!
//! // ツールを登録(オプション) //! // ツールを登録(オプション)
//! use worker::tool::Tool; //! // worker.register_tool(my_tool_definition)?;
//! worker.register_tool(my_tool);
//! //!
//! // 対話を実行 //! // 対話を実行
//! let history = worker.run("Hello!").await?; //! let history = worker.run("Hello!").await?;
@@ -36,21 +35,17 @@
//! locked.run("user input").await?; //! locked.run("user input").await?;
//! ``` //! ```
pub mod llm_client;
pub mod timeline;
pub mod event;
mod handler; mod handler;
pub mod hook;
mod message; mod message;
pub mod state;
pub mod subscriber;
pub mod tool;
mod worker; mod worker;
// ============================================================================= pub mod event;
// トップレベル公開(最も頻繁に使う型) pub mod hook;
// ============================================================================= pub mod llm_client;
pub mod state;
pub mod subscriber;
pub mod timeline;
pub mod tool;
pub use message::{ContentPart, Message, MessageContent, Role}; pub use message::{ContentPart, Message, MessageContent, Role};
pub use worker::{Worker, WorkerConfig, WorkerError}; pub use worker::{ToolRegistryError, Worker, WorkerConfig, WorkerError, WorkerResult};
+86
View File
@@ -0,0 +1,86 @@
//! LLMクライアント共通trait定義
use std::pin::Pin;
use crate::llm_client::{ClientError, Request, RequestConfig, event::Event};
use async_trait::async_trait;
use futures::Stream;
/// 設定に関する警告
///
/// プロバイダがサポートしていない設定を使用した場合に返される。
#[derive(Debug, Clone)]
pub struct ConfigWarning {
/// 設定オプション名
pub option_name: &'static str,
/// 警告メッセージ
pub message: String,
}
impl ConfigWarning {
/// 新しい警告を作成
pub fn unsupported(option_name: &'static str, provider_name: &str) -> Self {
Self {
option_name,
message: format!(
"'{}' is not supported by {} and will be ignored",
option_name, provider_name
),
}
}
}
impl std::fmt::Display for ConfigWarning {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
write!(f, "{}: {}", self.option_name, self.message)
}
}
/// LLMクライアントのtrait
///
/// 各プロバイダはこのtraitを実装し、統一されたインターフェースを提供する。
#[async_trait]
pub trait LlmClient: Send + Sync {
/// ストリーミングリクエストを送信し、Eventストリームを返す
///
/// # Arguments
/// * `request` - リクエスト情報
///
/// # Returns
/// * `Ok(Stream)` - イベントストリーム
/// * `Err(ClientError)` - エラー
async fn stream(
&self,
request: Request,
) -> Result<Pin<Box<dyn Stream<Item = Result<Event, ClientError>> + Send>>, ClientError>;
/// 設定をバリデーションし、未サポートの設定があれば警告を返す
///
/// # Arguments
/// * `config` - バリデーション対象の設定
///
/// # Returns
/// サポートされていない設定に対する警告のリスト
fn validate_config(&self, config: &RequestConfig) -> Vec<ConfigWarning> {
// デフォルト実装: 全ての設定をサポート
let _ = config;
Vec::new()
}
}
/// `Box<dyn LlmClient>` に対する `LlmClient` の実装
///
/// これにより、動的ディスパッチを使用するクライアントも `Worker` で利用可能になる。
#[async_trait]
impl LlmClient for Box<dyn LlmClient> {
async fn stream(
&self,
request: Request,
) -> Result<Pin<Box<dyn Stream<Item = Result<Event, ClientError>> + Send>>, ClientError> {
(**self).stream(request).await
}
fn validate_config(&self, config: &RequestConfig) -> Vec<ConfigWarning> {
(**self).validate_config(config)
}
}
@@ -1,6 +1,6 @@
//! LLMクライアント層 //! LLMクライアント層
//! //!
//! 各LLMプロバイダと通信し、統一された[`Event`](crate::llm_client::event::Event) //! 各LLMプロバイダと通信し、統一された[`Event`]
//! ストリームを出力します。 //! ストリームを出力します。
//! //!
//! # サポートするプロバイダ //! # サポートするプロバイダ
@@ -5,7 +5,8 @@
use std::pin::Pin; use std::pin::Pin;
use crate::llm_client::{ use crate::llm_client::{
ClientError, LlmClient, Request, event::Event, scheme::openai::OpenAIScheme, ClientError, ConfigWarning, LlmClient, Request, RequestConfig, event::Event,
scheme::openai::OpenAIScheme,
}; };
use async_trait::async_trait; use async_trait::async_trait;
use eventsource_stream::Eventsource; use eventsource_stream::Eventsource;
@@ -197,4 +198,15 @@ impl LlmClient for OpenAIClient {
Ok(Box::pin(stream)) Ok(Box::pin(stream))
} }
fn validate_config(&self, config: &RequestConfig) -> Vec<ConfigWarning> {
let mut warnings = Vec::new();
// OpenAI does not support top_k
if config.top_k.is_some() {
warnings.push(ConfigWarning::unsupported("top_k", "OpenAI"));
}
warnings
}
} }
@@ -23,6 +23,8 @@ pub(crate) struct AnthropicRequest {
pub temperature: Option<f32>, pub temperature: Option<f32>,
#[serde(skip_serializing_if = "Option::is_none")] #[serde(skip_serializing_if = "Option::is_none")]
pub top_p: Option<f32>, pub top_p: Option<f32>,
#[serde(skip_serializing_if = "Option::is_none")]
pub top_k: Option<u32>,
#[serde(skip_serializing_if = "Vec::is_empty")] #[serde(skip_serializing_if = "Vec::is_empty")]
pub stop_sequences: Vec<String>, pub stop_sequences: Vec<String>,
pub stream: bool, pub stream: bool,
@@ -90,6 +92,7 @@ impl AnthropicScheme {
tools, tools,
temperature: request.config.temperature, temperature: request.config.temperature,
top_p: request.config.top_p, top_p: request.config.top_p,
top_k: request.config.top_k,
stop_sequences: request.config.stop_sequences.clone(), stop_sequences: request.config.stop_sequences.clone(),
stream: true, stream: true,
} }
@@ -133,6 +133,9 @@ pub(crate) struct GeminiGenerationConfig {
/// Top P /// Top P
#[serde(skip_serializing_if = "Option::is_none")] #[serde(skip_serializing_if = "Option::is_none")]
pub top_p: Option<f32>, pub top_p: Option<f32>,
/// Top K
#[serde(skip_serializing_if = "Option::is_none")]
pub top_k: Option<u32>,
/// ストップシーケンス /// ストップシーケンス
#[serde(skip_serializing_if = "Vec::is_empty")] #[serde(skip_serializing_if = "Vec::is_empty")]
pub stop_sequences: Vec<String>, pub stop_sequences: Vec<String>,
@@ -183,6 +186,7 @@ impl GeminiScheme {
max_output_tokens: request.config.max_tokens, max_output_tokens: request.config.max_tokens,
temperature: request.config.temperature, temperature: request.config.temperature,
top_p: request.config.top_p, top_p: request.config.top_p,
top_k: request.config.top_k,
stop_sequences: request.config.stop_sequences.clone(), stop_sequences: request.config.stop_sequences.clone(),
}); });
@@ -62,6 +62,30 @@ impl Request {
self.config.max_tokens = Some(max_tokens); self.config.max_tokens = Some(max_tokens);
self self
} }
/// temperatureを設定
pub fn temperature(mut self, temperature: f32) -> Self {
self.config.temperature = Some(temperature);
self
}
/// top_pを設定
pub fn top_p(mut self, top_p: f32) -> Self {
self.config.top_p = Some(top_p);
self
}
/// top_kを設定
pub fn top_k(mut self, top_k: u32) -> Self {
self.config.top_k = Some(top_k);
self
}
/// ストップシーケンスを追加
pub fn stop_sequence(mut self, sequence: impl Into<String>) -> Self {
self.config.stop_sequences.push(sequence.into());
self
}
} }
/// メッセージ /// メッセージ
@@ -191,8 +215,47 @@ pub struct RequestConfig {
pub max_tokens: Option<u32>, pub max_tokens: Option<u32>,
/// Temperature /// Temperature
pub temperature: Option<f32>, pub temperature: Option<f32>,
/// Top P /// Top P (nucleus sampling)
pub top_p: Option<f32>, pub top_p: Option<f32>,
/// Top K
pub top_k: Option<u32>,
/// ストップシーケンス /// ストップシーケンス
pub stop_sequences: Vec<String>, pub stop_sequences: Vec<String>,
} }
impl RequestConfig {
/// 新しいデフォルト設定を作成
pub fn new() -> Self {
Self::default()
}
/// 最大トークン数を設定
pub fn with_max_tokens(mut self, max_tokens: u32) -> Self {
self.max_tokens = Some(max_tokens);
self
}
/// temperatureを設定
pub fn with_temperature(mut self, temperature: f32) -> Self {
self.temperature = Some(temperature);
self
}
/// top_pを設定
pub fn with_top_p(mut self, top_p: f32) -> Self {
self.top_p = Some(top_p);
self
}
/// top_kを設定
pub fn with_top_k(mut self, top_k: u32) -> Self {
self.top_k = Some(top_k);
self
}
/// ストップシーケンスを追加
pub fn with_stop_sequence(mut self, sequence: impl Into<String>) -> Self {
self.stop_sequences.push(sequence.into());
self
}
}
@@ -20,7 +20,7 @@ pub enum Role {
/// # Examples /// # Examples
/// ///
/// ```ignore /// ```ignore
/// use worker::Message; /// use llm_worker::Message;
/// ///
/// // ユーザーメッセージ /// // ユーザーメッセージ
/// let user_msg = Message::user("Hello!"); /// let user_msg = Message::user("Hello!");
@@ -79,7 +79,7 @@ impl Message {
/// # Examples /// # Examples
/// ///
/// ```ignore /// ```ignore
/// use worker::Message; /// use llm_worker::Message;
/// let msg = Message::user("こんにちは"); /// let msg = Message::user("こんにちは");
/// ``` /// ```
pub fn user(content: impl Into<String>) -> Self { pub fn user(content: impl Into<String>) -> Self {
@@ -1,7 +1,7 @@
//! Worker状態 //! Worker状態
//! //!
//! Type-stateパターンによるキャッシュ保護のための状態マーカー型。 //! Type-stateパターンによるキャッシュ保護のための状態マーカー型。
//! Workerは`Mutable` → `Locked`の状態遷移を持ちます。 //! Workerは`Mutable` → `CacheLocked`の状態遷移を持ちます。
/// Worker状態を表すマーカートレイト /// Worker状態を表すマーカートレイト
/// ///
@@ -19,12 +19,12 @@ mod private {
/// - メッセージ履歴の編集(追加、削除、クリア) /// - メッセージ履歴の編集(追加、削除、クリア)
/// - ツール・Hookの登録 /// - ツール・Hookの登録
/// ///
/// `Worker::lock()`により[`Locked`]状態へ遷移できます。 /// `Worker::lock()`により[`CacheLocked`]状態へ遷移できます。
/// ///
/// # Examples /// # Examples
/// ///
/// ```ignore /// ```ignore
/// use worker::Worker; /// use llm_worker::Worker;
/// ///
/// let mut worker = Worker::new(client) /// let mut worker = Worker::new(client)
/// .system_prompt("You are helpful."); /// .system_prompt("You are helpful.");
@@ -42,7 +42,7 @@ pub struct Mutable;
impl private::Sealed for Mutable {} impl private::Sealed for Mutable {}
impl WorkerState for Mutable {} impl WorkerState for Mutable {}
/// ロック状態(キャッシュ保護) /// キャッシュロック状態(キャッシュ保護)
/// ///
/// この状態では以下の制限があります: /// この状態では以下の制限があります:
/// - システムプロンプトの変更不可 /// - システムプロンプトの変更不可
@@ -54,7 +54,7 @@ impl WorkerState for Mutable {}
/// `Worker::unlock()`により[`Mutable`]状態へ戻せますが、 /// `Worker::unlock()`により[`Mutable`]状態へ戻せますが、
/// キャッシュ保護が解除されることに注意してください。 /// キャッシュ保護が解除されることに注意してください。
#[derive(Debug, Clone, Copy, Default)] #[derive(Debug, Clone, Copy, Default)]
pub struct Locked; pub struct CacheLocked;
impl private::Sealed for Locked {} impl private::Sealed for CacheLocked {}
impl WorkerState for Locked {} impl WorkerState for CacheLocked {}
@@ -33,8 +33,8 @@ use crate::{
/// # Examples /// # Examples
/// ///
/// ```ignore /// ```ignore
/// use worker::subscriber::WorkerSubscriber; /// use llm_worker::subscriber::WorkerSubscriber;
/// use worker::timeline::TextBlockEvent; /// use llm_worker::timeline::TextBlockEvent;
/// ///
/// struct StreamPrinter; /// struct StreamPrinter;
/// ///
@@ -65,10 +65,10 @@ pub trait WorkerSubscriber: Send {
/// ///
/// ブロック開始時にDefault::default()で生成され、 /// ブロック開始時にDefault::default()で生成され、
/// ブロック終了時に破棄される。 /// ブロック終了時に破棄される。
type TextBlockScope: Default + Send; type TextBlockScope: Default + Send + Sync;
/// ツール使用ブロック処理用のスコープ型 /// ツール使用ブロック処理用のスコープ型
type ToolUseBlockScope: Default + Send; type ToolUseBlockScope: Default + Send + Sync;
// ========================================================================= // =========================================================================
// ブロックイベント(スコープ管理あり) // ブロックイベント(スコープ管理あり)
@@ -17,7 +17,7 @@ use crate::handler::*;
/// 各Handlerは独自のScope型を持つため、Timelineで保持するには型消去が必要です。 /// 各Handlerは独自のScope型を持つため、Timelineで保持するには型消去が必要です。
/// 通常は直接使用せず、`Timeline::on_text_block()`などのメソッド経由で /// 通常は直接使用せず、`Timeline::on_text_block()`などのメソッド経由で
/// 自動的にラップされます。 /// 自動的にラップされます。
pub trait ErasedHandler<K: Kind>: Send { pub trait ErasedHandler<K: Kind>: Send + Sync {
/// イベントをディスパッチ /// イベントをディスパッチ
fn dispatch(&mut self, event: &K::Event); fn dispatch(&mut self, event: &K::Event);
/// スコープを開始(Block開始時) /// スコープを開始(Block開始時)
@@ -54,9 +54,9 @@ where
impl<H, K> ErasedHandler<K> for HandlerWrapper<H, K> impl<H, K> ErasedHandler<K> for HandlerWrapper<H, K>
where where
H: Handler<K> + Send, H: Handler<K> + Send + Sync,
K: Kind, K: Kind,
H::Scope: Send, H::Scope: Send + Sync,
{ {
fn dispatch(&mut self, event: &K::Event) { fn dispatch(&mut self, event: &K::Event) {
if let Some(scope) = &mut self.scope { if let Some(scope) = &mut self.scope {
@@ -78,7 +78,7 @@ where
// ============================================================================= // =============================================================================
/// ブロックハンドラーの型消去trait /// ブロックハンドラーの型消去trait
trait ErasedBlockHandler: Send { trait ErasedBlockHandler: Send + Sync {
fn dispatch_start(&mut self, start: &BlockStart); fn dispatch_start(&mut self, start: &BlockStart);
fn dispatch_delta(&mut self, delta: &BlockDelta); fn dispatch_delta(&mut self, delta: &BlockDelta);
fn dispatch_stop(&mut self, stop: &BlockStop); fn dispatch_stop(&mut self, stop: &BlockStop);
@@ -112,8 +112,8 @@ where
impl<H> ErasedBlockHandler for TextBlockHandlerWrapper<H> impl<H> ErasedBlockHandler for TextBlockHandlerWrapper<H>
where where
H: Handler<TextBlockKind> + Send, H: Handler<TextBlockKind> + Send + Sync,
H::Scope: Send, H::Scope: Send + Sync,
{ {
fn dispatch_start(&mut self, start: &BlockStart) { fn dispatch_start(&mut self, start: &BlockStart) {
if let Some(scope) = &mut self.scope { if let Some(scope) = &mut self.scope {
@@ -185,8 +185,8 @@ where
impl<H> ErasedBlockHandler for ThinkingBlockHandlerWrapper<H> impl<H> ErasedBlockHandler for ThinkingBlockHandlerWrapper<H>
where where
H: Handler<ThinkingBlockKind> + Send, H: Handler<ThinkingBlockKind> + Send + Sync,
H::Scope: Send, H::Scope: Send + Sync,
{ {
fn dispatch_start(&mut self, start: &BlockStart) { fn dispatch_start(&mut self, start: &BlockStart) {
if let Some(scope) = &mut self.scope { if let Some(scope) = &mut self.scope {
@@ -255,8 +255,8 @@ where
impl<H> ErasedBlockHandler for ToolUseBlockHandlerWrapper<H> impl<H> ErasedBlockHandler for ToolUseBlockHandlerWrapper<H>
where where
H: Handler<ToolUseBlockKind> + Send, H: Handler<ToolUseBlockKind> + Send + Sync,
H::Scope: Send, H::Scope: Send + Sync,
{ {
fn dispatch_start(&mut self, start: &BlockStart) { fn dispatch_start(&mut self, start: &BlockStart) {
if let Some(scope) = &mut self.scope { if let Some(scope) = &mut self.scope {
@@ -328,7 +328,7 @@ where
/// # Examples /// # Examples
/// ///
/// ```ignore /// ```ignore
/// use worker::{Timeline, Handler, TextBlockKind, TextBlockEvent}; /// use llm_worker::{Timeline, Handler, TextBlockKind, TextBlockEvent};
/// ///
/// struct MyHandler; /// struct MyHandler;
/// impl Handler<TextBlockKind> for MyHandler { /// impl Handler<TextBlockKind> for MyHandler {
@@ -391,8 +391,8 @@ impl Timeline {
/// UsageKind用のHandlerを登録 /// UsageKind用のHandlerを登録
pub fn on_usage<H>(&mut self, handler: H) -> &mut Self pub fn on_usage<H>(&mut self, handler: H) -> &mut Self
where where
H: Handler<UsageKind> + Send + 'static, H: Handler<UsageKind> + Send + Sync + 'static,
H::Scope: Send, H::Scope: Send + Sync,
{ {
// Meta系はデフォルトでスコープを開始しておく // Meta系はデフォルトでスコープを開始しておく
let mut wrapper = HandlerWrapper::new(handler); let mut wrapper = HandlerWrapper::new(handler);
@@ -404,8 +404,8 @@ impl Timeline {
/// PingKind用のHandlerを登録 /// PingKind用のHandlerを登録
pub fn on_ping<H>(&mut self, handler: H) -> &mut Self pub fn on_ping<H>(&mut self, handler: H) -> &mut Self
where where
H: Handler<PingKind> + Send + 'static, H: Handler<PingKind> + Send + Sync + 'static,
H::Scope: Send, H::Scope: Send + Sync,
{ {
let mut wrapper = HandlerWrapper::new(handler); let mut wrapper = HandlerWrapper::new(handler);
wrapper.start_scope(); wrapper.start_scope();
@@ -416,8 +416,8 @@ impl Timeline {
/// StatusKind用のHandlerを登録 /// StatusKind用のHandlerを登録
pub fn on_status<H>(&mut self, handler: H) -> &mut Self pub fn on_status<H>(&mut self, handler: H) -> &mut Self
where where
H: Handler<StatusKind> + Send + 'static, H: Handler<StatusKind> + Send + Sync + 'static,
H::Scope: Send, H::Scope: Send + Sync,
{ {
let mut wrapper = HandlerWrapper::new(handler); let mut wrapper = HandlerWrapper::new(handler);
wrapper.start_scope(); wrapper.start_scope();
@@ -428,8 +428,8 @@ impl Timeline {
/// ErrorKind用のHandlerを登録 /// ErrorKind用のHandlerを登録
pub fn on_error<H>(&mut self, handler: H) -> &mut Self pub fn on_error<H>(&mut self, handler: H) -> &mut Self
where where
H: Handler<ErrorKind> + Send + 'static, H: Handler<ErrorKind> + Send + Sync + 'static,
H::Scope: Send, H::Scope: Send + Sync,
{ {
let mut wrapper = HandlerWrapper::new(handler); let mut wrapper = HandlerWrapper::new(handler);
wrapper.start_scope(); wrapper.start_scope();
@@ -440,8 +440,8 @@ impl Timeline {
/// TextBlockKind用のHandlerを登録 /// TextBlockKind用のHandlerを登録
pub fn on_text_block<H>(&mut self, handler: H) -> &mut Self pub fn on_text_block<H>(&mut self, handler: H) -> &mut Self
where where
H: Handler<TextBlockKind> + Send + 'static, H: Handler<TextBlockKind> + Send + Sync + 'static,
H::Scope: Send, H::Scope: Send + Sync,
{ {
self.text_block_handlers self.text_block_handlers
.push(Box::new(TextBlockHandlerWrapper::new(handler))); .push(Box::new(TextBlockHandlerWrapper::new(handler)));
@@ -451,8 +451,8 @@ impl Timeline {
/// ThinkingBlockKind用のHandlerを登録 /// ThinkingBlockKind用のHandlerを登録
pub fn on_thinking_block<H>(&mut self, handler: H) -> &mut Self pub fn on_thinking_block<H>(&mut self, handler: H) -> &mut Self
where where
H: Handler<ThinkingBlockKind> + Send + 'static, H: Handler<ThinkingBlockKind> + Send + Sync + 'static,
H::Scope: Send, H::Scope: Send + Sync,
{ {
self.thinking_block_handlers self.thinking_block_handlers
.push(Box::new(ThinkingBlockHandlerWrapper::new(handler))); .push(Box::new(ThinkingBlockHandlerWrapper::new(handler)));
@@ -462,8 +462,8 @@ impl Timeline {
/// ToolUseBlockKind用のHandlerを登録 /// ToolUseBlockKind用のHandlerを登録
pub fn on_tool_use_block<H>(&mut self, handler: H) -> &mut Self pub fn on_tool_use_block<H>(&mut self, handler: H) -> &mut Self
where where
H: Handler<ToolUseBlockKind> + Send + 'static, H: Handler<ToolUseBlockKind> + Send + Sync + 'static,
H::Scope: Send, H::Scope: Send + Sync,
{ {
self.tool_use_block_handlers self.tool_use_block_handlers
.push(Box::new(ToolUseBlockHandlerWrapper::new(handler))); .push(Box::new(ToolUseBlockHandlerWrapper::new(handler)));
@@ -578,6 +578,21 @@ impl Timeline {
pub fn current_block(&self) -> Option<BlockType> { pub fn current_block(&self) -> Option<BlockType> {
self.current_block self.current_block
} }
/// 現在アクティブなブロックを中断する
///
/// キャンセルやエラー時に呼び出し、進行中のブロックに対して
/// BlockAbortイベントを発火してスコープをクリーンアップする。
pub fn abort_current_block(&mut self) {
if let Some(block_type) = self.current_block {
let abort = crate::timeline::event::BlockAbort {
index: 0, // インデックスは不明なので0
block_type,
reason: "Cancelled".to_string(),
};
self.handle_block_abort(&abort);
}
}
} }
#[cfg(test)] #[cfg(test)]
+154
View File
@@ -0,0 +1,154 @@
//! ツール定義
//!
//! 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>;
}
File diff suppressed because it is too large Load Diff
@@ -8,9 +8,9 @@ use std::sync::{Arc, Mutex};
use async_trait::async_trait; use async_trait::async_trait;
use futures::Stream; use futures::Stream;
use worker::llm_client::event::{BlockType, DeltaContent, Event}; use llm_worker::llm_client::event::{BlockType, DeltaContent, Event};
use worker::llm_client::{ClientError, LlmClient, Request}; use llm_worker::llm_client::{ClientError, LlmClient, Request};
use worker::timeline::{Handler, TextBlockEvent, TextBlockKind, Timeline}; use llm_worker::timeline::{Handler, TextBlockEvent, TextBlockKind, Timeline};
use std::sync::atomic::{AtomicUsize, Ordering}; use std::sync::atomic::{AtomicUsize, Ordering};
@@ -267,7 +267,7 @@ pub fn assert_timeline_integration(subdir: &str) {
}); });
for event in &events { for event in &events {
let timeline_event: worker::timeline::event::Event = event.clone().into(); let timeline_event: llm_worker::timeline::event::Event = event.clone().into();
timeline.dispatch(&timeline_event); timeline.dispatch(&timeline_event);
} }
@@ -7,10 +7,13 @@ use std::sync::atomic::{AtomicUsize, Ordering};
use std::time::{Duration, Instant}; use std::time::{Duration, Instant};
use async_trait::async_trait; use async_trait::async_trait;
use worker::Worker; use llm_worker::Worker;
use worker::hook::{ControlFlow, HookError, ToolCall, ToolResult, WorkerHook}; use llm_worker::hook::{
use worker::llm_client::event::{Event, ResponseStatus, StatusEvent}; Hook, HookError, PostToolCall, PostToolCallContext, PostToolCallResult, PreToolCall,
use worker::tool::{Tool, ToolError}; PreToolCallResult, ToolCallContext,
};
use llm_worker::llm_client::event::{Event, ResponseStatus, StatusEvent};
use llm_worker::tool::{Tool, ToolDefinition, ToolError, ToolMeta};
mod common; mod common;
use common::MockLlmClient; use common::MockLlmClient;
@@ -39,25 +42,24 @@ impl SlowTool {
fn call_count(&self) -> usize { fn call_count(&self) -> usize {
self.call_count.load(Ordering::SeqCst) self.call_count.load(Ordering::SeqCst)
} }
/// ToolDefinition を作成
fn definition(&self) -> ToolDefinition {
let tool = self.clone();
Arc::new(move || {
let meta = ToolMeta::new(&tool.name)
.description("A tool that waits before responding")
.input_schema(serde_json::json!({
"type": "object",
"properties": {}
}));
(meta, Arc::new(tool.clone()) as Arc<dyn Tool>)
})
}
} }
#[async_trait] #[async_trait]
impl Tool for SlowTool { impl Tool for SlowTool {
fn name(&self) -> &str {
&self.name
}
fn description(&self) -> &str {
"A tool that waits before responding"
}
fn input_schema(&self) -> serde_json::Value {
serde_json::json!({
"type": "object",
"properties": {}
})
}
async fn execute(&self, _input_json: &str) -> Result<String, ToolError> { async fn execute(&self, _input_json: &str) -> Result<String, ToolError> {
self.call_count.fetch_add(1, Ordering::SeqCst); self.call_count.fetch_add(1, Ordering::SeqCst);
tokio::time::sleep(Duration::from_millis(self.delay_ms)).await; tokio::time::sleep(Duration::from_millis(self.delay_ms)).await;
@@ -103,9 +105,9 @@ async fn test_parallel_tool_execution() {
let tool2_clone = tool2.clone(); let tool2_clone = tool2.clone();
let tool3_clone = tool3.clone(); let tool3_clone = tool3.clone();
worker.register_tool(tool1); worker.register_tool(tool1.definition()).unwrap();
worker.register_tool(tool2); worker.register_tool(tool2.definition()).unwrap();
worker.register_tool(tool3); worker.register_tool(tool3.definition()).unwrap();
let start = Instant::now(); let start = Instant::now();
let _result = worker.run("Run all tools").await; let _result = worker.run("Run all tools").await;
@@ -127,7 +129,7 @@ async fn test_parallel_tool_execution() {
println!("Parallel execution completed in {:?}", elapsed); println!("Parallel execution completed in {:?}", elapsed);
} }
/// Hook: before_tool_call でスキップされたツールは実行されないことを確認 /// Hook: pre_tool_call でスキップされたツールは実行されないことを確認
#[tokio::test] #[tokio::test]
async fn test_before_tool_call_skip() { async fn test_before_tool_call_skip() {
let events = vec![ let events = vec![
@@ -151,27 +153,24 @@ async fn test_before_tool_call_skip() {
let allowed_clone = allowed_tool.clone(); let allowed_clone = allowed_tool.clone();
let blocked_clone = blocked_tool.clone(); let blocked_clone = blocked_tool.clone();
worker.register_tool(allowed_tool); worker.register_tool(allowed_tool.definition()).unwrap();
worker.register_tool(blocked_tool); worker.register_tool(blocked_tool.definition()).unwrap();
// "blocked_tool" をスキップするHook // "blocked_tool" をスキップするHook
struct BlockingHook; struct BlockingHook;
#[async_trait] #[async_trait]
impl WorkerHook for BlockingHook { impl Hook<PreToolCall> for BlockingHook {
async fn before_tool_call( async fn call(&self, ctx: &mut ToolCallContext) -> Result<PreToolCallResult, HookError> {
&self, if ctx.call.name == "blocked_tool" {
tool_call: &mut ToolCall, Ok(PreToolCallResult::Skip)
) -> Result<ControlFlow, HookError> {
if tool_call.name == "blocked_tool" {
Ok(ControlFlow::Skip)
} else { } else {
Ok(ControlFlow::Continue) Ok(PreToolCallResult::Continue)
} }
} }
} }
worker.add_hook(BlockingHook); worker.add_pre_tool_call_hook(BlockingHook);
let _result = worker.run("Test hook").await; let _result = worker.run("Test hook").await;
@@ -188,9 +187,9 @@ async fn test_before_tool_call_skip() {
); );
} }
/// Hook: after_tool_call で結果が改変されることを確認 /// Hook: post_tool_call で結果が改変されることを確認
#[tokio::test] #[tokio::test]
async fn test_after_tool_call_modification() { async fn test_post_tool_call_modification() {
// 複数リクエストに対応するレスポンスを準備 // 複数リクエストに対応するレスポンスを準備
let client = MockLlmClient::with_responses(vec![ let client = MockLlmClient::with_responses(vec![
// 1回目のリクエスト: ツール呼び出し // 1回目のリクエスト: ツール呼び出し
@@ -220,21 +219,21 @@ async fn test_after_tool_call_modification() {
#[async_trait] #[async_trait]
impl Tool for SimpleTool { 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!({})
}
async fn execute(&self, _: &str) -> Result<String, ToolError> { async fn execute(&self, _: &str) -> Result<String, ToolError> {
Ok("Original Result".to_string()) Ok("Original Result".to_string())
} }
} }
worker.register_tool(SimpleTool); fn simple_tool_definition() -> ToolDefinition {
Arc::new(|| {
let meta = ToolMeta::new("test_tool")
.description("Test")
.input_schema(serde_json::json!({}));
(meta, Arc::new(SimpleTool) as Arc<dyn Tool>)
})
}
worker.register_tool(simple_tool_definition()).unwrap();
// 結果を改変するHook // 結果を改変するHook
struct ModifyingHook { struct ModifyingHook {
@@ -242,19 +241,19 @@ async fn test_after_tool_call_modification() {
} }
#[async_trait] #[async_trait]
impl WorkerHook for ModifyingHook { impl Hook<PostToolCall> for ModifyingHook {
async fn after_tool_call( async fn call(
&self, &self,
tool_result: &mut ToolResult, ctx: &mut PostToolCallContext,
) -> Result<ControlFlow, HookError> { ) -> Result<PostToolCallResult, HookError> {
tool_result.content = format!("[Modified] {}", tool_result.content); ctx.result.content = format!("[Modified] {}", ctx.result.content);
*self.modified_content.lock().unwrap() = Some(tool_result.content.clone()); *self.modified_content.lock().unwrap() = Some(ctx.result.content.clone());
Ok(ControlFlow::Continue) Ok(PostToolCallResult::Continue)
} }
} }
let modified_content = Arc::new(std::sync::Mutex::new(None)); let modified_content = Arc::new(std::sync::Mutex::new(None));
worker.add_hook(ModifyingHook { worker.add_post_tool_call_hook(ModifyingHook {
modified_content: modified_content.clone(), modified_content: modified_content.clone(),
}); });
@@ -7,12 +7,12 @@ mod common;
use std::sync::{Arc, Mutex}; use std::sync::{Arc, Mutex};
use common::MockLlmClient; use common::MockLlmClient;
use worker::Worker; use llm_worker::Worker;
use worker::hook::ToolCall; use llm_worker::hook::ToolCall;
use worker::llm_client::event::{Event, ResponseStatus, StatusEvent as ClientStatusEvent}; use llm_worker::llm_client::event::{Event, ResponseStatus, StatusEvent as ClientStatusEvent};
use worker::subscriber::WorkerSubscriber; use llm_worker::subscriber::WorkerSubscriber;
use worker::timeline::event::{ErrorEvent, StatusEvent, UsageEvent}; use llm_worker::timeline::event::{ErrorEvent, StatusEvent, UsageEvent};
use worker::timeline::{TextBlockEvent, ToolUseBlockEvent}; use llm_worker::timeline::{TextBlockEvent, ToolUseBlockEvent};
// ============================================================================= // =============================================================================
// Test Subscriber // Test Subscriber
@@ -9,8 +9,7 @@ use std::sync::atomic::{AtomicUsize, Ordering};
use schemars; use schemars;
use serde; use serde;
use worker::tool::Tool; use llm_worker_macros::tool_registry;
use worker_macros::tool_registry;
// ============================================================================= // =============================================================================
// Test: Basic Tool Generation // Test: Basic Tool Generation
@@ -51,30 +50,31 @@ async fn test_basic_tool_generation() {
prefix: "Hello".to_string(), prefix: "Hello".to_string(),
}; };
// ファクトリメソッドでツールを取得 // ファクトリメソッドでToolDefinitionを取得
let greet_tool = ctx.greet_tool(); let greet_definition = ctx.greet_definition();
// 名前の確認 // ファクトリを呼び出してMetaとToolを取得
assert_eq!(greet_tool.name(), "greet"); let (meta, tool) = greet_definition();
// 説明の確認(docコメントから取得) // メタ情報の確認
let desc = greet_tool.description(); assert_eq!(meta.name, "greet");
assert!( assert!(
desc.contains("メッセージに挨拶を追加する"), meta.description.contains("メッセージに挨拶を追加する"),
"Description should contain doc comment: {}", "Description should contain doc comment: {}",
desc meta.description
); );
// スキーマの確認
let schema = greet_tool.input_schema();
println!("Schema: {}", serde_json::to_string_pretty(&schema).unwrap());
assert!( assert!(
schema.get("properties").is_some(), meta.input_schema.get("properties").is_some(),
"Schema should have properties" "Schema should have properties"
); );
println!(
"Schema: {}",
serde_json::to_string_pretty(&meta.input_schema).unwrap()
);
// 実行テスト // 実行テスト
let result = greet_tool.execute(r#"{"message": "World"}"#).await; let result = tool.execute(r#"{"message": "World"}"#).await;
assert!(result.is_ok(), "Should execute successfully"); assert!(result.is_ok(), "Should execute successfully");
let output = result.unwrap(); let output = result.unwrap();
assert!(output.contains("Hello"), "Output should contain prefix"); assert!(output.contains("Hello"), "Output should contain prefix");
@@ -87,11 +87,11 @@ async fn test_multiple_arguments() {
prefix: "".to_string(), prefix: "".to_string(),
}; };
let add_tool = ctx.add_tool(); let (meta, tool) = ctx.add_definition()();
assert_eq!(add_tool.name(), "add"); assert_eq!(meta.name, "add");
let result = add_tool.execute(r#"{"a": 10, "b": 20}"#).await; let result = tool.execute(r#"{"a": 10, "b": 20}"#).await;
assert!(result.is_ok()); assert!(result.is_ok());
let output = result.unwrap(); let output = result.unwrap();
assert!(output.contains("30"), "Should contain sum: {}", output); assert!(output.contains("30"), "Should contain sum: {}", output);
@@ -103,12 +103,12 @@ async fn test_no_arguments() {
prefix: "TestPrefix".to_string(), prefix: "TestPrefix".to_string(),
}; };
let get_prefix_tool = ctx.get_prefix_tool(); let (meta, tool) = ctx.get_prefix_definition()();
assert_eq!(get_prefix_tool.name(), "get_prefix"); assert_eq!(meta.name, "get_prefix");
// 空のJSONオブジェクトで呼び出し // 空のJSONオブジェクトで呼び出し
let result = get_prefix_tool.execute(r#"{}"#).await; let result = tool.execute(r#"{}"#).await;
assert!(result.is_ok()); assert!(result.is_ok());
let output = result.unwrap(); let output = result.unwrap();
assert!( assert!(
@@ -124,10 +124,10 @@ async fn test_invalid_arguments() {
prefix: "".to_string(), prefix: "".to_string(),
}; };
let greet_tool = ctx.greet_tool(); let (_, tool) = ctx.greet_definition()();
// 不正なJSON // 不正なJSON
let result = greet_tool.execute(r#"{"wrong_field": "value"}"#).await; let result = tool.execute(r#"{"wrong_field": "value"}"#).await;
assert!(result.is_err(), "Should fail with invalid arguments"); assert!(result.is_err(), "Should fail with invalid arguments");
} }
@@ -163,9 +163,9 @@ impl FallibleContext {
#[tokio::test] #[tokio::test]
async fn test_result_return_type_success() { async fn test_result_return_type_success() {
let ctx = FallibleContext; let ctx = FallibleContext;
let validate_tool = ctx.validate_tool(); let (_, tool) = ctx.validate_definition()();
let result = validate_tool.execute(r#"{"value": 42}"#).await; let result = tool.execute(r#"{"value": 42}"#).await;
assert!(result.is_ok(), "Should succeed for positive value"); assert!(result.is_ok(), "Should succeed for positive value");
let output = result.unwrap(); let output = result.unwrap();
assert!(output.contains("Valid"), "Should contain Valid: {}", output); assert!(output.contains("Valid"), "Should contain Valid: {}", output);
@@ -174,9 +174,9 @@ async fn test_result_return_type_success() {
#[tokio::test] #[tokio::test]
async fn test_result_return_type_error() { async fn test_result_return_type_error() {
let ctx = FallibleContext; let ctx = FallibleContext;
let validate_tool = ctx.validate_tool(); let (_, tool) = ctx.validate_definition()();
let result = validate_tool.execute(r#"{"value": -1}"#).await; let result = tool.execute(r#"{"value": -1}"#).await;
assert!(result.is_err(), "Should fail for negative value"); assert!(result.is_err(), "Should fail for negative value");
let err = result.unwrap_err(); let err = result.unwrap_err();
@@ -211,12 +211,12 @@ async fn test_sync_method() {
counter: Arc::new(AtomicUsize::new(0)), counter: Arc::new(AtomicUsize::new(0)),
}; };
let increment_tool = ctx.increment_tool(); let (_, tool) = ctx.increment_definition()();
// 3回実行 // 3回実行
let result1 = increment_tool.execute(r#"{}"#).await; let result1 = tool.execute(r#"{}"#).await;
let result2 = increment_tool.execute(r#"{}"#).await; let result2 = tool.execute(r#"{}"#).await;
let result3 = increment_tool.execute(r#"{}"#).await; let result3 = tool.execute(r#"{}"#).await;
assert!(result1.is_ok()); assert!(result1.is_ok());
assert!(result2.is_ok()); assert!(result2.is_ok());
@@ -225,3 +225,22 @@ async fn test_sync_method() {
// カウンターは3になっているはず // カウンターは3になっているはず
assert_eq!(ctx.counter.load(Ordering::SeqCst), 3); assert_eq!(ctx.counter.load(Ordering::SeqCst), 3);
} }
// =============================================================================
// Test: ToolMeta Immutability
// =============================================================================
#[tokio::test]
async fn test_tool_meta_immutability() {
let ctx = SimpleContext {
prefix: "Test".to_string(),
};
// 2回取得しても同じメタ情報が得られることを確認
let (meta1, _) = ctx.greet_definition()();
let (meta2, _) = ctx.greet_definition()();
assert_eq!(meta1.name, meta2.name);
assert_eq!(meta1.description, meta2.description);
assert_eq!(meta1.input_schema, meta2.input_schema);
}
+39
View File
@@ -0,0 +1,39 @@
use llm_worker::llm_client::providers::openai::OpenAIClient;
use llm_worker::{Worker, WorkerError};
#[test]
fn test_openai_top_k_warning() {
// ダミーキーでクライアント作成(validate_configは通信しないため安全)
let client = OpenAIClient::new("dummy-key", "gpt-4o");
// top_kを設定したWorkerを作成
let worker = Worker::new(client).top_k(50); // OpenAIはtop_k非対応
// validate()を実行
let result = worker.validate();
// エラーが返り、ConfigWarningsが含まれていることを確認
match result {
Err(WorkerError::ConfigWarnings(warnings)) => {
assert_eq!(warnings.len(), 1);
assert_eq!(warnings[0].option_name, "top_k");
println!("Got expected warning: {}", warnings[0]);
}
Ok(_) => panic!("Should have returned validation error"),
Err(e) => panic!("Unexpected error type: {:?}", e),
}
}
#[test]
fn test_openai_valid_config() {
let client = OpenAIClient::new("dummy-key", "gpt-4o");
// validな設定(temperatureのみ)
let worker = Worker::new(client).temperature(0.7);
// validate()を実行
let result = worker.validate();
// 成功を確認
assert!(result.is_ok());
}
@@ -11,8 +11,8 @@ use std::sync::atomic::{AtomicUsize, Ordering};
use async_trait::async_trait; use async_trait::async_trait;
use common::MockLlmClient; use common::MockLlmClient;
use worker::Worker; use llm_worker::Worker;
use worker::tool::{Tool, ToolError}; use llm_worker::tool::{Tool, ToolDefinition, ToolError, ToolMeta};
/// フィクスチャディレクトリのパス /// フィクスチャディレクトリのパス
fn fixtures_dir() -> std::path::PathBuf { fn fixtures_dir() -> std::path::PathBuf {
@@ -35,20 +35,13 @@ impl MockWeatherTool {
fn get_call_count(&self) -> usize { fn get_call_count(&self) -> usize {
self.call_count.load(Ordering::SeqCst) self.call_count.load(Ordering::SeqCst)
} }
}
#[async_trait] fn definition(&self) -> ToolDefinition {
impl Tool for MockWeatherTool { let tool = self.clone();
fn name(&self) -> &str { Arc::new(move || {
"get_weather" let meta = ToolMeta::new("get_weather")
} .description("Get the current weather for a city")
.input_schema(serde_json::json!({
fn description(&self) -> &str {
"Get the current weather for a city"
}
fn input_schema(&self) -> serde_json::Value {
serde_json::json!({
"type": "object", "type": "object",
"properties": { "properties": {
"city": { "city": {
@@ -57,9 +50,14 @@ impl Tool for MockWeatherTool {
} }
}, },
"required": ["city"] "required": ["city"]
}));
(meta, Arc::new(tool.clone()) as Arc<dyn Tool>)
}) })
} }
}
#[async_trait]
impl Tool for MockWeatherTool {
async fn execute(&self, input_json: &str) -> Result<String, ToolError> { async fn execute(&self, input_json: &str) -> Result<String, ToolError> {
self.call_count.fetch_add(1, Ordering::SeqCst); self.call_count.fetch_add(1, Ordering::SeqCst);
@@ -100,7 +98,7 @@ fn test_mock_client_from_fixture() {
/// fixtureファイルを使わず、プログラムでイベントを構築してクライアントを作成する。 /// fixtureファイルを使わず、プログラムでイベントを構築してクライアントを作成する。
#[test] #[test]
fn test_mock_client_from_events() { fn test_mock_client_from_events() {
use worker::llm_client::event::Event; use llm_worker::llm_client::event::Event;
// 直接イベントを指定 // 直接イベントを指定
let events = vec![ let events = vec![
@@ -158,7 +156,7 @@ async fn test_worker_tool_call() {
// ツールを登録 // ツールを登録
let weather_tool = MockWeatherTool::new(); let weather_tool = MockWeatherTool::new();
let tool_for_check = weather_tool.clone(); let tool_for_check = weather_tool.clone();
worker.register_tool(weather_tool); worker.register_tool(weather_tool.definition()).unwrap();
// メッセージを送信 // メッセージを送信
let _result = worker.run("What's the weather in Tokyo?").await; let _result = worker.run("What's the weather in Tokyo?").await;
@@ -178,7 +176,7 @@ async fn test_worker_tool_call() {
/// テストの独立性を高め、外部ファイルへの依存を排除したい場合に有用。 /// テストの独立性を高め、外部ファイルへの依存を排除したい場合に有用。
#[tokio::test] #[tokio::test]
async fn test_worker_with_programmatic_events() { async fn test_worker_with_programmatic_events() {
use worker::llm_client::event::{Event, ResponseStatus, StatusEvent}; use llm_worker::llm_client::event::{Event, ResponseStatus, StatusEvent};
// プログラムでイベントシーケンスを構築 // プログラムでイベントシーケンスを構築
let events = vec![ let events = vec![
@@ -205,8 +203,8 @@ async fn test_worker_with_programmatic_events() {
/// id, name, inputJSON)を正しく抽出できることを検証する。 /// id, name, inputJSON)を正しく抽出できることを検証する。
#[tokio::test] #[tokio::test]
async fn test_tool_call_collector_integration() { async fn test_tool_call_collector_integration() {
use worker::llm_client::event::Event; use llm_worker::llm_client::event::Event;
use worker::timeline::{Timeline, ToolCallCollector}; use llm_worker::timeline::{Timeline, ToolCallCollector};
// ToolUseブロックを含むイベントシーケンス // ToolUseブロックを含むイベントシーケンス
let events = vec![ let events = vec![
@@ -222,7 +220,7 @@ async fn test_tool_call_collector_integration() {
// イベントをディスパッチ // イベントをディスパッチ
for event in &events { for event in &events {
let timeline_event: worker::timeline::event::Event = event.clone().into(); let timeline_event: llm_worker::timeline::event::Event = event.clone().into();
timeline.dispatch(&timeline_event); timeline.dispatch(&timeline_event);
} }
@@ -1,14 +1,14 @@
//! Worker状態管理のテスト //! Worker状態管理のテスト
//! //!
//! Type-stateパターン(Mutable/Locked)による状態遷移と //! Type-stateパターン(Mutable/CacheLocked)による状態遷移と
//! ターン間の状態保持をテストする。 //! ターン間の状態保持をテストする。
mod common; mod common;
use common::MockLlmClient; use common::MockLlmClient;
use worker::Worker; use llm_worker::Worker;
use worker::llm_client::event::{Event, ResponseStatus, StatusEvent}; use llm_worker::llm_client::event::{Event, ResponseStatus, StatusEvent};
use worker::{Message, MessageContent}; use llm_worker::{Message, MessageContent};
// ============================================================================= // =============================================================================
// Mutable状態のテスト // Mutable状態のテスト
@@ -95,7 +95,7 @@ fn test_mutable_extend_history() {
// 状態遷移テスト // 状態遷移テスト
// ============================================================================= // =============================================================================
/// lock()でMutable -> Locked状態に遷移することを確認 /// lock()でMutable -> CacheLocked状態に遷移することを確認
#[test] #[test]
fn test_lock_transition() { fn test_lock_transition() {
let client = MockLlmClient::new(vec![]); let client = MockLlmClient::new(vec![]);
@@ -108,13 +108,13 @@ fn test_lock_transition() {
// ロック // ロック
let locked_worker = worker.lock(); let locked_worker = worker.lock();
// Locked状態でも履歴とシステムプロンプトにアクセス可能 // CacheLocked状態でも履歴とシステムプロンプトにアクセス可能
assert_eq!(locked_worker.get_system_prompt(), Some("System")); assert_eq!(locked_worker.get_system_prompt(), Some("System"));
assert_eq!(locked_worker.history().len(), 2); assert_eq!(locked_worker.history().len(), 2);
assert_eq!(locked_worker.locked_prefix_len(), 2); assert_eq!(locked_worker.locked_prefix_len(), 2);
} }
/// unlock()でLocked -> Mutable状態に遷移することを確認 /// unlock()でCacheLocked -> Mutable状態に遷移することを確認
#[test] #[test]
fn test_unlock_transition() { fn test_unlock_transition() {
let client = MockLlmClient::new(vec![]); let client = MockLlmClient::new(vec![]);
@@ -172,7 +172,7 @@ async fn test_mutable_run_updates_history() {
)); ));
} }
/// Locked状態で複数ターンを実行し、履歴が正しく累積することを確認 /// CacheLocked状態で複数ターンを実行し、履歴が正しく累積することを確認
#[tokio::test] #[tokio::test]
async fn test_locked_multi_turn_history_accumulation() { async fn test_locked_multi_turn_history_accumulation() {
// 2回のリクエストに対応するレスポンスを準備 // 2回のリクエストに対応するレスポンスを準備
@@ -340,7 +340,7 @@ async fn test_unlock_edit_relock() {
// システムプロンプト保持のテスト // システムプロンプト保持のテスト
// ============================================================================= // =============================================================================
/// Locked状態でもシステムプロンプトが保持されることを確認 /// CacheLocked状態でもシステムプロンプトが保持されることを確認
#[test] #[test]
fn test_system_prompt_preserved_in_locked_state() { fn test_system_prompt_preserved_in_locked_state() {
let client = MockLlmClient::new(vec![]); let client = MockLlmClient::new(vec![]);
-182
View File
@@ -1,182 +0,0 @@
//! Hook関連の型定義
//!
//! Worker層でのターン制御・介入に使用される型
use async_trait::async_trait;
use serde::{Deserialize, Serialize};
use serde_json::Value;
use thiserror::Error;
// =============================================================================
// Control Flow Types
// =============================================================================
/// Hook処理の制御フロー
#[derive(Debug, Clone, PartialEq, Eq)]
pub enum ControlFlow {
/// 処理を続行
Continue,
/// 現在の処理をスキップ(Tool実行など)
Skip,
/// 処理を中断
Abort(String),
}
/// ターン終了時の判定結果
#[derive(Debug, Clone)]
pub enum TurnResult {
/// ターンを終了
Finish,
/// メッセージを追加してターン継続(自己修正など)
ContinueWithMessages(Vec<crate::Message>),
}
// =============================================================================
// Tool Call / Result Types
// =============================================================================
/// ツール呼び出し情報
///
/// LLMからのToolUseブロックを表現し、Hook処理で改変可能
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct ToolCall {
/// ツール呼び出しID(レスポンスとの紐付けに使用)
pub id: String,
/// ツール名
pub name: String,
/// 入力引数(JSON
pub input: Value,
}
/// ツール実行結果
///
/// ツール実行後の結果を表現し、Hook処理で改変可能
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct ToolResult {
/// 対応するツール呼び出しID
pub tool_use_id: String,
/// 結果コンテンツ
pub content: String,
/// エラーかどうか
#[serde(default)]
pub is_error: bool,
}
impl ToolResult {
/// 成功結果を作成
pub fn success(tool_use_id: impl Into<String>, content: impl Into<String>) -> Self {
Self {
tool_use_id: tool_use_id.into(),
content: content.into(),
is_error: false,
}
}
/// エラー結果を作成
pub fn error(tool_use_id: impl Into<String>, content: impl Into<String>) -> Self {
Self {
tool_use_id: tool_use_id.into(),
content: content.into(),
is_error: true,
}
}
}
// =============================================================================
// Hook Error
// =============================================================================
/// Hookエラー
#[derive(Debug, Error)]
pub enum HookError {
/// 処理が中断された
#[error("Aborted: {0}")]
Aborted(String),
/// 内部エラー
#[error("Hook error: {0}")]
Internal(String),
}
// =============================================================================
// WorkerHook Trait
// =============================================================================
/// ターンの進行・ツール実行に介入するためのトレイト
///
/// Hookを使うと、メッセージ送信前、ツール実行前後、ターン終了時に
/// 処理を挟んだり、実行をキャンセルしたりできます。
///
/// # Examples
///
/// ```ignore
/// use worker::hook::{ControlFlow, HookError, ToolCall, TurnResult, WorkerHook};
/// use worker::Message;
///
/// struct ValidationHook;
///
/// #[async_trait::async_trait]
/// impl WorkerHook for ValidationHook {
/// async fn before_tool_call(&self, call: &mut ToolCall) -> Result<ControlFlow, HookError> {
/// // 危険なツールをブロック
/// if call.name == "delete_all" {
/// return Ok(ControlFlow::Skip);
/// }
/// Ok(ControlFlow::Continue)
/// }
///
/// async fn on_turn_end(&self, messages: &[Message]) -> Result<TurnResult, HookError> {
/// // 条件を満たさなければ追加メッセージで継続
/// if messages.len() < 3 {
/// return Ok(TurnResult::ContinueWithMessages(vec![
/// Message::user("Please elaborate.")
/// ]));
/// }
/// Ok(TurnResult::Finish)
/// }
/// }
/// ```
///
/// # デフォルト実装
///
/// すべてのメソッドにはデフォルト実装があり、何も行わず`Continue`を返します。
/// 必要なメソッドのみオーバーライドしてください。
#[async_trait]
pub trait WorkerHook: Send + Sync {
/// メッセージ送信前に呼ばれる
///
/// リクエストに含まれるメッセージリストを参照・改変できます。
/// `ControlFlow::Abort`を返すとターンが中断されます。
async fn on_message_send(
&self,
_context: &mut Vec<crate::Message>,
) -> Result<ControlFlow, HookError> {
Ok(ControlFlow::Continue)
}
/// ツール実行前に呼ばれる
///
/// ツール呼び出しの引数を書き換えたり、実行をスキップしたりできます。
/// `ControlFlow::Skip`を返すとこのツールの実行がスキップされます。
async fn before_tool_call(&self, _tool_call: &mut ToolCall) -> Result<ControlFlow, HookError> {
Ok(ControlFlow::Continue)
}
/// ツール実行後に呼ばれる
///
/// ツールの実行結果を書き換えたり、隠蔽したりできます。
async fn after_tool_call(
&self,
_tool_result: &mut ToolResult,
) -> Result<ControlFlow, HookError> {
Ok(ControlFlow::Continue)
}
/// ターン終了時に呼ばれる
///
/// 生成されたメッセージを検査し、必要なら追加メッセージで継続を指示できます。
/// `TurnResult::ContinueWithMessages`を返すと、指定したメッセージを追加して
/// 次のターンに進みます。
async fn on_turn_end(&self, _messages: &[crate::Message]) -> Result<TurnResult, HookError> {
Ok(TurnResult::Finish)
}
}
-39
View File
@@ -1,39 +0,0 @@
//! LLMクライアント共通trait定義
use std::pin::Pin;
use crate::llm_client::{ClientError, Request, event::Event};
use async_trait::async_trait;
use futures::Stream;
/// LLMクライアントのtrait
///
/// 各プロバイダはこのtraitを実装し、統一されたインターフェースを提供する。
#[async_trait]
pub trait LlmClient: Send + Sync {
/// ストリーミングリクエストを送信し、Eventストリームを返す
///
/// # Arguments
/// * `request` - リクエスト情報
///
/// # Returns
/// * `Ok(Stream)` - イベントストリーム
/// * `Err(ClientError)` - エラー
async fn stream(
&self,
request: Request,
) -> Result<Pin<Box<dyn Stream<Item = Result<Event, ClientError>> + Send>>, ClientError>;
}
/// `Box<dyn LlmClient>` に対する `LlmClient` の実装
///
/// これにより、動的ディスパッチを使用するクライアントも `Worker` で利用可能になる。
#[async_trait]
impl LlmClient for Box<dyn LlmClient> {
async fn stream(
&self,
request: Request,
) -> Result<Pin<Box<dyn Stream<Item = Result<Event, ClientError>> + Send>>, ClientError> {
(**self).stream(request).await
}
}
-90
View File
@@ -1,90 +0,0 @@
//! ツール定義
//!
//! LLMから呼び出し可能なツールを定義するためのトレイト。
//! 通常は`#[tool]`マクロを使用して自動実装します。
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),
}
/// LLMから呼び出し可能なツールを定義するトレイト
///
/// ツールはLLMが外部リソースにアクセスしたり、
/// 計算を実行したりするために使用します。
///
/// # 実装方法
///
/// 通常は`#[tool]`マクロを使用して自動実装します:
///
/// ```ignore
/// use worker::tool;
///
/// #[tool(description = "Search the web for information")]
/// async fn search(query: String) -> String {
/// // 検索処理
/// format!("Results for: {}", query)
/// }
/// ```
///
/// # 手動実装
///
/// ```ignore
/// use worker::tool::{Tool, ToolError};
/// use serde_json::{json, Value};
///
/// struct MyTool;
///
/// #[async_trait::async_trait]
/// impl Tool for MyTool {
/// fn name(&self) -> &str { "my_tool" }
/// fn description(&self) -> &str { "My custom tool" }
/// fn input_schema(&self) -> Value {
/// json!({
/// "type": "object",
/// "properties": {
/// "query": { "type": "string" }
/// },
/// "required": ["query"]
/// })
/// }
/// async fn execute(&self, input: &str) -> Result<String, ToolError> {
/// Ok("result".to_string())
/// }
/// }
/// ```
#[async_trait]
pub trait Tool: Send + Sync {
/// ツール名(LLMが識別に使用)
fn name(&self) -> &str;
/// ツールの説明(LLMへのプロンプトに含まれる)
fn description(&self) -> &str;
/// 引数のJSON Schema
///
/// LLMはこのスキーマに従って引数を生成します。
fn input_schema(&self) -> Value;
/// ツールを実行する
///
/// # Arguments
/// * `input_json` - LLMが生成したJSON形式の引数
///
/// # Returns
/// 実行結果の文字列。この内容がLLMに返されます。
async fn execute(&self, input_json: &str) -> Result<String, ToolError>;
}
-788
View File
@@ -1,788 +0,0 @@
use std::collections::HashMap;
use std::marker::PhantomData;
use std::sync::{Arc, Mutex};
use futures::StreamExt;
use tracing::{debug, info, trace, warn};
use crate::{
ContentPart, Message, MessageContent, Role,
hook::{ControlFlow, HookError, ToolCall, ToolResult, TurnResult, WorkerHook},
llm_client::{ClientError, LlmClient, Request, ToolDefinition},
state::{Locked, Mutable, WorkerState},
subscriber::{
ErrorSubscriberAdapter, StatusSubscriberAdapter, TextBlockSubscriberAdapter,
ToolUseBlockSubscriberAdapter, UsageSubscriberAdapter, WorkerSubscriber,
},
timeline::{TextBlockCollector, Timeline, ToolCallCollector},
tool::{Tool, ToolError},
};
// =============================================================================
// Worker Error
// =============================================================================
/// Workerエラー
#[derive(Debug, thiserror::Error)]
pub enum WorkerError {
/// クライアントエラー
#[error("Client error: {0}")]
Client(#[from] ClientError),
/// ツールエラー
#[error("Tool error: {0}")]
Tool(#[from] ToolError),
/// Hookエラー
#[error("Hook error: {0}")]
Hook(#[from] HookError),
/// 処理が中断された
#[error("Aborted: {0}")]
Aborted(String),
}
// =============================================================================
// Worker Config
// =============================================================================
/// Worker設定
#[derive(Debug, Clone, Default)]
pub struct WorkerConfig {
// 将来の拡張用(現在は空)
_private: (),
}
// =============================================================================
// ターン制御用コールバック保持
// =============================================================================
/// ターンイベントを通知するためのコールバック (型消去)
trait TurnNotifier: Send {
fn on_turn_start(&self, turn: usize);
fn on_turn_end(&self, turn: usize);
}
struct SubscriberTurnNotifier<S: WorkerSubscriber + 'static> {
subscriber: Arc<Mutex<S>>,
}
impl<S: WorkerSubscriber + 'static> TurnNotifier for SubscriberTurnNotifier<S> {
fn on_turn_start(&self, turn: usize) {
if let Ok(mut s) = self.subscriber.lock() {
s.on_turn_start(turn);
}
}
fn on_turn_end(&self, turn: usize) {
if let Ok(mut s) = self.subscriber.lock() {
s.on_turn_end(turn);
}
}
}
// =============================================================================
// Worker
// =============================================================================
/// LLMとの対話を管理する中心コンポーネント
///
/// ユーザーからの入力を受け取り、LLMにリクエストを送信し、
/// ツール呼び出しがあれば自動的に実行してターンを進行させます。
///
/// # 状態遷移(Type-state
///
/// - [`Mutable`]: 初期状態。システムプロンプトや履歴を自由に編集可能。
/// - [`Locked`]: キャッシュ保護状態。`lock()`で遷移。前方コンテキストは不変。
///
/// # Examples
///
/// ```ignore
/// use worker::{Worker, Message};
///
/// // Workerを作成してツールを登録
/// let mut worker = Worker::new(client)
/// .system_prompt("You are a helpful assistant.");
/// worker.register_tool(my_tool);
///
/// // 対話を実行
/// let history = worker.run("Hello!").await?;
/// ```
///
/// # キャッシュ保護が必要な場合
///
/// ```ignore
/// let mut worker = Worker::new(client)
/// .system_prompt("...");
///
/// // 履歴を設定後、ロックしてキャッシュを保護
/// let mut locked = worker.lock();
/// locked.run("user input").await?;
/// ```
pub struct Worker<C: LlmClient, S: WorkerState = Mutable> {
/// LLMクライアント
client: C,
/// イベントタイムライン
timeline: Timeline,
/// テキストブロックコレクター(Timeline用ハンドラ)
text_block_collector: TextBlockCollector,
/// ツールコールコレクター(Timeline用ハンドラ)
tool_call_collector: ToolCallCollector,
/// 登録されたツール
tools: HashMap<String, Arc<dyn Tool>>,
/// 登録されたHook
hooks: Vec<Box<dyn WorkerHook>>,
/// システムプロンプト
system_prompt: Option<String>,
/// メッセージ履歴(Workerが所有)
history: Vec<Message>,
/// ロック時点での履歴長(Locked状態でのみ意味を持つ)
locked_prefix_len: usize,
/// ターンカウント
turn_count: usize,
/// ターン通知用のコールバック
turn_notifiers: Vec<Box<dyn TurnNotifier>>,
/// 状態マーカー
_state: PhantomData<S>,
}
// =============================================================================
// 共通実装(全状態で利用可能)
// =============================================================================
impl<C: LlmClient, S: WorkerState> Worker<C, S> {
/// イベント購読者を登録する
///
/// 登録したSubscriberは、LLMからのストリーミングイベントを
/// リアルタイムで受信できます。UIへのストリーム表示などに利用します。
///
/// # 受信できるイベント
///
/// - **ブロックイベント**: `on_text_block`, `on_tool_use_block`
/// - **メタイベント**: `on_usage`, `on_status`, `on_error`
/// - **完了イベント**: `on_text_complete`, `on_tool_call_complete`
/// - **ターン制御**: `on_turn_start`, `on_turn_end`
///
/// # Examples
///
/// ```ignore
/// use worker::{Worker, WorkerSubscriber, TextBlockEvent};
///
/// struct MyPrinter;
/// impl WorkerSubscriber for MyPrinter {
/// type TextBlockScope = ();
/// type ToolUseBlockScope = ();
///
/// fn on_text_block(&mut self, _: &mut (), event: &TextBlockEvent) {
/// if let TextBlockEvent::Delta(text) = event {
/// print!("{}", text);
/// }
/// }
/// }
///
/// worker.subscribe(MyPrinter);
/// ```
pub fn subscribe<Sub: WorkerSubscriber + 'static>(&mut self, subscriber: Sub) {
let subscriber = Arc::new(Mutex::new(subscriber));
// TextBlock用ハンドラを登録
self.timeline
.on_text_block(TextBlockSubscriberAdapter::new(subscriber.clone()));
// ToolUseBlock用ハンドラを登録
self.timeline
.on_tool_use_block(ToolUseBlockSubscriberAdapter::new(subscriber.clone()));
// Meta系ハンドラを登録
self.timeline
.on_usage(UsageSubscriberAdapter::new(subscriber.clone()));
self.timeline
.on_status(StatusSubscriberAdapter::new(subscriber.clone()));
self.timeline
.on_error(ErrorSubscriberAdapter::new(subscriber.clone()));
// ターン制御用コールバックを登録
self.turn_notifiers
.push(Box::new(SubscriberTurnNotifier { subscriber }));
}
/// ツールを登録する
///
/// 登録されたツールはLLMからの呼び出しで自動的に実行されます。
/// 同名のツールを登録した場合、後から登録したものが優先されます。
///
/// # Examples
///
/// ```ignore
/// use worker::Worker;
/// use my_tools::SearchTool;
///
/// worker.register_tool(SearchTool::new());
/// ```
pub fn register_tool(&mut self, tool: impl Tool + 'static) {
let name = tool.name().to_string();
self.tools.insert(name, Arc::new(tool));
}
/// 複数のツールを登録
pub fn register_tools(&mut self, tools: impl IntoIterator<Item = impl Tool + 'static>) {
for tool in tools {
self.register_tool(tool);
}
}
/// Hookを追加する
///
/// Hookはターンの進行・ツール実行に介入できます。
/// 複数のHookを登録した場合、登録順に実行されます。
///
/// # Examples
///
/// ```ignore
/// use worker::{Worker, WorkerHook, ControlFlow, ToolCall};
///
/// struct LoggingHook;
///
/// #[async_trait::async_trait]
/// impl WorkerHook for LoggingHook {
/// async fn before_tool_call(&self, call: &mut ToolCall) -> Result<ControlFlow, HookError> {
/// println!("Calling tool: {}", call.name);
/// Ok(ControlFlow::Continue)
/// }
/// }
///
/// worker.add_hook(LoggingHook);
/// ```
pub fn add_hook(&mut self, hook: impl WorkerHook + 'static) {
self.hooks.push(Box::new(hook));
}
/// タイムラインへの可変参照を取得(追加ハンドラ登録用)
pub fn timeline_mut(&mut self) -> &mut Timeline {
&mut self.timeline
}
/// 履歴への参照を取得
pub fn history(&self) -> &[Message] {
&self.history
}
/// システムプロンプトへの参照を取得
pub fn get_system_prompt(&self) -> Option<&str> {
self.system_prompt.as_deref()
}
/// 現在のターンカウントを取得
pub fn turn_count(&self) -> usize {
self.turn_count
}
/// 登録されたツールからToolDefinitionのリストを生成
fn build_tool_definitions(&self) -> Vec<ToolDefinition> {
self.tools
.values()
.map(|tool| {
ToolDefinition::new(tool.name())
.description(tool.description())
.input_schema(tool.input_schema())
})
.collect()
}
/// テキストブロックとツール呼び出しからアシスタントメッセージを構築
fn build_assistant_message(
&self,
text_blocks: &[String],
tool_calls: &[ToolCall],
) -> Option<Message> {
// テキストもツール呼び出しもない場合はNone
if text_blocks.is_empty() && tool_calls.is_empty() {
return None;
}
// テキストのみの場合はシンプルなテキストメッセージ
if tool_calls.is_empty() {
let text = text_blocks.join("");
return Some(Message::assistant(text));
}
// ツール呼び出しがある場合は Parts として構築
let mut parts = Vec::new();
// テキストパーツを追加
for text in text_blocks {
if !text.is_empty() {
parts.push(ContentPart::Text { text: text.clone() });
}
}
// ツール呼び出しパーツを追加
for call in tool_calls {
parts.push(ContentPart::ToolUse {
id: call.id.clone(),
name: call.name.clone(),
input: call.input.clone(),
});
}
Some(Message {
role: Role::Assistant,
content: MessageContent::Parts(parts),
})
}
/// リクエストを構築
fn build_request(&self, tool_definitions: &[ToolDefinition]) -> Request {
let mut request = Request::new();
// システムプロンプトを設定
if let Some(ref system) = self.system_prompt {
request = request.system(system);
}
// メッセージを追加
for msg in &self.history {
// Message から llm_client::Message への変換
request = request.message(crate::llm_client::Message {
role: match msg.role {
Role::User => crate::llm_client::Role::User,
Role::Assistant => crate::llm_client::Role::Assistant,
},
content: match &msg.content {
MessageContent::Text(t) => crate::llm_client::MessageContent::Text(t.clone()),
MessageContent::ToolResult {
tool_use_id,
content,
} => crate::llm_client::MessageContent::ToolResult {
tool_use_id: tool_use_id.clone(),
content: content.clone(),
},
MessageContent::Parts(parts) => crate::llm_client::MessageContent::Parts(
parts
.iter()
.map(|p| match p {
ContentPart::Text { text } => {
crate::llm_client::ContentPart::Text { text: text.clone() }
}
ContentPart::ToolUse { id, name, input } => {
crate::llm_client::ContentPart::ToolUse {
id: id.clone(),
name: name.clone(),
input: input.clone(),
}
}
ContentPart::ToolResult {
tool_use_id,
content,
} => crate::llm_client::ContentPart::ToolResult {
tool_use_id: tool_use_id.clone(),
content: content.clone(),
},
})
.collect(),
),
},
});
}
// ツール定義を追加
for tool_def in tool_definitions {
request = request.tool(tool_def.clone());
}
request
}
/// Hooks: on_message_send
async fn run_on_message_send_hooks(&self) -> Result<ControlFlow, WorkerError> {
for hook in &self.hooks {
// Note: Locked状態でも履歴全体を参照として渡す(変更は不可)
// HookのAPIを変更し、immutable参照のみを渡すようにする必要があるかもしれない
// 現在は空のVecを渡して回避(要検討)
let mut temp_context = self.history.clone();
let result = hook.on_message_send(&mut temp_context).await?;
match result {
ControlFlow::Continue => continue,
ControlFlow::Skip => return Ok(ControlFlow::Skip),
ControlFlow::Abort(reason) => return Ok(ControlFlow::Abort(reason)),
}
}
Ok(ControlFlow::Continue)
}
/// Hooks: on_turn_end
async fn run_on_turn_end_hooks(&self) -> Result<TurnResult, WorkerError> {
for hook in &self.hooks {
let result = hook.on_turn_end(&self.history).await?;
match result {
TurnResult::Finish => continue,
TurnResult::ContinueWithMessages(msgs) => {
return Ok(TurnResult::ContinueWithMessages(msgs));
}
}
}
Ok(TurnResult::Finish)
}
/// ツールを並列実行
///
/// 全てのツールに対してbefore_tool_callフックを実行後、
/// 許可されたツールを並列に実行し、結果にafter_tool_callフックを適用する。
async fn execute_tools(
&self,
tool_calls: Vec<ToolCall>,
) -> Result<Vec<ToolResult>, WorkerError> {
use futures::future::join_all;
// Phase 1: before_tool_call フックを適用(スキップ/中断を判定)
let mut approved_calls = Vec::new();
for mut tool_call in tool_calls {
let mut skip = false;
for hook in &self.hooks {
let result = hook.before_tool_call(&mut tool_call).await?;
match result {
ControlFlow::Continue => {}
ControlFlow::Skip => {
skip = true;
break;
}
ControlFlow::Abort(reason) => {
return Err(WorkerError::Aborted(reason));
}
}
}
if !skip {
approved_calls.push(tool_call);
}
}
// Phase 2: 許可されたツールを並列実行
let futures: Vec<_> = approved_calls
.into_iter()
.map(|tool_call| {
let tools = &self.tools;
async move {
if let Some(tool) = tools.get(&tool_call.name) {
let input_json =
serde_json::to_string(&tool_call.input).unwrap_or_default();
match tool.execute(&input_json).await {
Ok(content) => ToolResult::success(&tool_call.id, content),
Err(e) => ToolResult::error(&tool_call.id, e.to_string()),
}
} else {
ToolResult::error(
&tool_call.id,
format!("Tool '{}' not found", tool_call.name),
)
}
}
})
.collect();
let mut results = join_all(futures).await;
// Phase 3: after_tool_call フックを適用
for tool_result in &mut results {
for hook in &self.hooks {
let result = hook.after_tool_call(tool_result).await?;
match result {
ControlFlow::Continue => {}
ControlFlow::Skip => break,
ControlFlow::Abort(reason) => {
return Err(WorkerError::Aborted(reason));
}
}
}
}
Ok(results)
}
/// 内部で使用するターン実行ロジック
async fn run_turn_loop(&mut self) -> Result<(), WorkerError> {
let tool_definitions = self.build_tool_definitions();
info!(
message_count = self.history.len(),
tool_count = tool_definitions.len(),
"Starting worker run"
);
loop {
// ターン開始を通知
let current_turn = self.turn_count;
debug!(turn = current_turn, "Turn start");
for notifier in &self.turn_notifiers {
notifier.on_turn_start(current_turn);
}
// Hook: on_message_send
let control = self.run_on_message_send_hooks().await?;
if let ControlFlow::Abort(reason) = control {
warn!(reason = %reason, "Aborted by hook");
// ターン終了を通知(異常終了)
for notifier in &self.turn_notifiers {
notifier.on_turn_end(current_turn);
}
return Err(WorkerError::Aborted(reason));
}
// リクエスト構築
let request = self.build_request(&tool_definitions);
debug!(
message_count = request.messages.len(),
tool_count = request.tools.len(),
has_system = request.system_prompt.is_some(),
"Sending request to LLM"
);
// ストリーム処理
debug!("Starting stream...");
let mut stream = self.client.stream(request).await?;
let mut event_count = 0;
while let Some(event_result) = stream.next().await {
match &event_result {
Ok(event) => {
trace!(event = ?event, "Received event");
event_count += 1;
}
Err(e) => {
warn!(error = %e, "Stream error");
}
}
let event = event_result?;
let timeline_event: crate::timeline::event::Event = event.into();
self.timeline.dispatch(&timeline_event);
}
debug!(event_count = event_count, "Stream completed");
// ターン終了を通知
for notifier in &self.turn_notifiers {
notifier.on_turn_end(current_turn);
}
self.turn_count += 1;
// 収集結果を取得
let text_blocks = self.text_block_collector.take_collected();
let tool_calls = self.tool_call_collector.take_collected();
// アシスタントメッセージを履歴に追加
let assistant_message = self.build_assistant_message(&text_blocks, &tool_calls);
if let Some(msg) = assistant_message {
self.history.push(msg);
}
if tool_calls.is_empty() {
// ツール呼び出しなし → ターン終了判定
let turn_result = self.run_on_turn_end_hooks().await?;
match turn_result {
TurnResult::Finish => {
return Ok(());
}
TurnResult::ContinueWithMessages(additional) => {
self.history.extend(additional);
continue;
}
}
}
// ツール実行
let tool_results = self.execute_tools(tool_calls).await?;
// ツール結果を履歴に追加
for result in tool_results {
self.history
.push(Message::tool_result(&result.tool_use_id, &result.content));
}
}
}
}
// =============================================================================
// Mutable状態専用の実装
// =============================================================================
impl<C: LlmClient> Worker<C, Mutable> {
/// 新しいWorkerを作成(Mutable状態)
pub fn new(client: C) -> Self {
let text_block_collector = TextBlockCollector::new();
let tool_call_collector = ToolCallCollector::new();
let mut timeline = Timeline::new();
// コレクターをTimelineに登録
timeline.on_text_block(text_block_collector.clone());
timeline.on_tool_use_block(tool_call_collector.clone());
Self {
client,
timeline,
text_block_collector,
tool_call_collector,
tools: HashMap::new(),
hooks: Vec::new(),
system_prompt: None,
history: Vec::new(),
locked_prefix_len: 0,
turn_count: 0,
turn_notifiers: Vec::new(),
_state: PhantomData,
}
}
/// システムプロンプトを設定(ビルダーパターン)
pub fn system_prompt(mut self, prompt: impl Into<String>) -> Self {
self.system_prompt = Some(prompt.into());
self
}
/// システムプロンプトを設定(可変参照版)
pub fn set_system_prompt(&mut self, prompt: impl Into<String>) {
self.system_prompt = Some(prompt.into());
}
/// 履歴への可変参照を取得
///
/// Mutable状態でのみ利用可能。履歴を自由に編集できる。
pub fn history_mut(&mut self) -> &mut Vec<Message> {
&mut self.history
}
/// 履歴を設定
pub fn set_history(&mut self, messages: Vec<Message>) {
self.history = messages;
}
/// 履歴にメッセージを追加(ビルダーパターン)
pub fn with_message(mut self, message: Message) -> Self {
self.history.push(message);
self
}
/// 履歴にメッセージを追加
pub fn push_message(&mut self, message: Message) {
self.history.push(message);
}
/// 複数のメッセージを履歴に追加(ビルダーパターン)
pub fn with_messages(mut self, messages: impl IntoIterator<Item = Message>) -> Self {
self.history.extend(messages);
self
}
/// 複数のメッセージを履歴に追加
pub fn extend_history(&mut self, messages: impl IntoIterator<Item = Message>) {
self.history.extend(messages);
}
/// 履歴をクリア
pub fn clear_history(&mut self) {
self.history.clear();
}
/// 設定を適用(将来の拡張用)
#[allow(dead_code)]
pub fn config(self, _config: WorkerConfig) -> Self {
self
}
/// ロックしてLocked状態へ遷移
///
/// この操作により、現在のシステムプロンプトと履歴が「確定済みプレフィックス」として
/// 固定される。以降は履歴への追記のみが可能となり、キャッシュヒットが保証される。
pub fn lock(self) -> Worker<C, Locked> {
let locked_prefix_len = self.history.len();
Worker {
client: self.client,
timeline: self.timeline,
text_block_collector: self.text_block_collector,
tool_call_collector: self.tool_call_collector,
tools: self.tools,
hooks: self.hooks,
system_prompt: self.system_prompt,
history: self.history,
locked_prefix_len,
turn_count: self.turn_count,
turn_notifiers: self.turn_notifiers,
_state: PhantomData,
}
}
/// ターンを実行(Mutable状態)
///
/// 新しいユーザーメッセージを履歴に追加し、LLMにリクエストを送信する。
/// ツール呼び出しがある場合は自動的にループする。
///
/// 注意: この関数は履歴を変更するため、キャッシュ保護が必要な場合は
/// `lock()` を呼んでからLocked状態で `run` を使用すること。
pub async fn run(&mut self, user_input: impl Into<String>) -> Result<&[Message], WorkerError> {
self.history.push(Message::user(user_input));
self.run_turn_loop().await?;
Ok(&self.history)
}
/// 複数メッセージでターンを実行(Mutable状態)
///
/// 指定されたメッセージを履歴に追加してから実行する。
pub async fn run_with_messages(
&mut self,
messages: Vec<Message>,
) -> Result<&[Message], WorkerError> {
self.history.extend(messages);
self.run_turn_loop().await?;
Ok(&self.history)
}
}
// =============================================================================
// Locked状態専用の実装
// =============================================================================
impl<C: LlmClient> Worker<C, Locked> {
/// ターンを実行(Locked状態)
///
/// 新しいユーザーメッセージを履歴の末尾に追加し、LLMにリクエストを送信する。
/// ロック時点より前の履歴(プレフィックス)は不変であるため、キャッシュヒットが保証される。
pub async fn run(&mut self, user_input: impl Into<String>) -> Result<&[Message], WorkerError> {
self.history.push(Message::user(user_input));
self.run_turn_loop().await?;
Ok(&self.history)
}
/// 複数メッセージでターンを実行(Locked状態)
pub async fn run_with_messages(
&mut self,
messages: Vec<Message>,
) -> Result<&[Message], WorkerError> {
self.history.extend(messages);
self.run_turn_loop().await?;
Ok(&self.history)
}
/// ロック時点のプレフィックス長を取得
pub fn locked_prefix_len(&self) -> usize {
self.locked_prefix_len
}
/// ロックを解除してMutable状態へ戻す
///
/// 注意: この操作を行うと、以降のリクエストでキャッシュがヒットしなくなる可能性がある。
/// 履歴を編集する必要がある場合にのみ使用すること。
pub fn unlock(self) -> Worker<C, Mutable> {
Worker {
client: self.client,
timeline: self.timeline,
text_block_collector: self.text_block_collector,
tool_call_collector: self.tool_call_collector,
tools: self.tools,
hooks: self.hooks,
system_prompt: self.system_prompt,
history: self.history,
locked_prefix_len: 0,
turn_count: self.turn_count,
turn_notifiers: self.turn_notifiers,
_state: PhantomData,
}
}
}
#[cfg(test)]
mod tests {
// 基本的なテストのみ。LlmClientを使ったテストは統合テストで行う。
}