Compare commits
340 Commits
| Author | SHA1 | Date | |
|---|---|---|---|
| e69352ee2d | |||
| ee4e3bc13f | |||
| a576961bd6 | |||
| 59c7fc1276 | |||
| bcf512fcfc | |||
| 195401c496 | |||
| 34d8d20ec6 | |||
| 08ba6f0446 | |||
| 26984892af | |||
| 526a426073 | |||
| c53e0546d4 | |||
| 349b3748bd | |||
| e23e5f9f7b | |||
| 8d02782de6 | |||
| 27ceefdb40 | |||
| 5168eb6781 | |||
| ddb73a9a33 | |||
| 53eff10d75 | |||
| 1df6114ff3 | |||
| 975484cc2b | |||
| 0421c9b643 | |||
| fb69c21252 | |||
| 0cb9122d16 | |||
| c164ad3cbb | |||
| 9b4171a468 | |||
| 5cae4e44fb | |||
| a145a42b2b | |||
| 715807645a | |||
| 1259c6865f | |||
| ff42460cb4 | |||
| 39a16f8d56 | |||
| 83de60f59c | |||
| cf60e090a5 | |||
| 0fb37c33ab | |||
| d81508c22a | |||
| 883ac659b2 | |||
| c6c10b5e24 | |||
| a4e5bef1b7 | |||
| f72c7b03f9 | |||
| bd6f709374 | |||
| 00f2201157 | |||
| b3f0d66071 | |||
| 8730d413bc | |||
| 79140fda3c | |||
| 67e749ea3a | |||
| 7bcfc133ae | |||
| e3e246607e | |||
| 16104cb2c5 | |||
| 224e51c386 | |||
| b022ca089c | |||
| 0ebb761c09 | |||
| c8067828d5 | |||
| 30eedd9b8c | |||
| d701b45057 | |||
| 722c9c101e | |||
| 86aa45f0c4 | |||
| cf45dc4820 | |||
| db77034431 | |||
| abdaec11b0 | |||
| 95fb349656 | |||
| d0b6b6c324 | |||
| d74c23ccf5 | |||
| ea1cfda0d6 | |||
| 5623f47f9a | |||
| e4df9ec193 | |||
| a6306d6b76 | |||
| 64529ba5cc | |||
| cc7f963b89 | |||
| 0ce86af116 | |||
| 2cb0ed3f64 | |||
| fb61854f11 | |||
| 53ba3344b1 | |||
| e20c8be8bb | |||
| 894dcb1d3c | |||
| 9a9e890f8a | |||
| 818ea634f0 | |||
| 780460f8d8 | |||
| e19483a920 | |||
| aca93f1cae | |||
| 1371a4aad2 | |||
| db4a45c0f6 | |||
| e95b1e5f82 | |||
| 15f4008f4b | |||
| f45f81fb45 | |||
| 2220fd2542 | |||
| 564480e165 | |||
| 297c63d91a | |||
| 26e2cd3f65 | |||
| 9f899466d4 | |||
| 38393ea4cf | |||
| a4f25826e3 | |||
| 93484fb33f | |||
| c90f003f92 | |||
| 24793b9b8d | |||
| 78e772f455 | |||
| 1e0d269aad | |||
| f6b1d408fc | |||
| 442b318b6c | |||
| a7c97aedb7 | |||
| 746f9e7b24 | |||
| 0d6c61af5c | |||
| 673f31c059 | |||
| 369a4f0a89 | |||
| 8d54eae4d0 | |||
| a805d5beab | |||
| dbb2aec8b6 | |||
| 1a98b76a1f | |||
| 51d10ab2b5 | |||
| 1aad750395 | |||
| e0aab6bd02 | |||
| 6cb93132b7 | |||
| 04126b99d6 | |||
| 0794eb960d | |||
| d619ad1d48 | |||
| 5b147e07b3 | |||
| 944ce441d8 | |||
| a7dcb8519b | |||
| d912d44fb3 | |||
| 4f7254a634 | |||
| bf923cb296 | |||
| d9f737e1bf | |||
| 59690d045e | |||
| 5d95acba53 | |||
| d46225d2a9 | |||
| 3af30a0e62 | |||
| 69eca4d96d | |||
| 7b2e4a83c9 | |||
| 344b80872a | |||
| ddf828ff5f | |||
| 4e170b069b | |||
| 22c75fb578 | |||
| 11ab9eb6b8 | |||
| 29b232f407 | |||
| 53e8c920e5 | |||
| 78d19bed4d | |||
| 10f4160635 | |||
| 7622836e8b | |||
| 4d4713a9fa | |||
| 25008599f9 | |||
| c00ab074f8 | |||
| aed1f1957f | |||
| c6a959e2e1 | |||
| 02b7ed37f6 | |||
| 0d84aaabb9 | |||
| 6efdcf9610 | |||
| 4266d317d8 | |||
| 4ce7aafcbd | |||
| 35d8b69f92 | |||
| 562057e608 | |||
| b7024e5340 | |||
| 088588231b | |||
| eff117d3d9 | |||
| 968c535709 | |||
| c8b6fa7b11 | |||
| 0aa334b54e | |||
| 78a49f841d | |||
| 43b2bd937e | |||
| a4326875ba | |||
| eb31a58346 | |||
| a6b0acc35d | |||
| cc7fcd0b5b | |||
| 02fe59b913 | |||
| 6fd5f47089 | |||
| 2a2922760e | |||
| a3793460fd | |||
| e0927a04d9 | |||
| 8665604bab | |||
| d4c3c135b3 | |||
| 60bd5e493c | |||
| 0753b2d841 | |||
| 17e6fbd692 | |||
| 0710441650 | |||
| 20a76cee3e | |||
| cb64785867 | |||
| e6e26103c4 | |||
| 15529a14f1 | |||
| 86839188e0 | |||
| 39701b378b | |||
| 45ff6da737 | |||
| a260dd1503 | |||
| 57859301df | |||
| 8c968d3f53 | |||
| 0034bfbe46 | |||
| a733b9247a | |||
| e0afa349b9 | |||
| 7d0ce94907 | |||
| 9045763c35 | |||
| 29898552d7 | |||
| 9d7c2f5c2f | |||
| 5c0fa42351 | |||
| ab045b0ef3 | |||
| 41e6843db1 | |||
| 911ec3c9b9 | |||
| fc6f0a1a7b | |||
| 21873da278 | |||
| d1cd6be2c9 | |||
| 0c0ae41bca | |||
| c9ed7a904a | |||
| d200a8f554 | |||
| 3d04c8fcf1 | |||
| f53f165d91 | |||
| e5645e4064 | |||
| 95e15ca8c4 | |||
| dbf7329e87 | |||
| ed6c3ae431 | |||
| 214d2ecc67 | |||
| 29c95671de | |||
| 238f93a096 | |||
| c76877e7b3 | |||
| 12e5a9c5aa | |||
| 7f4be2ca3f | |||
| 29ffe12d8c | |||
| d34bed4f15 | |||
| aec7ea7e80 | |||
| 5938e1af29 | |||
| 60902297c5 | |||
| 12a95aa6fa | |||
| 78fc459a97 | |||
| 281565804c | |||
| 33a32fd9c8 | |||
| b64aad55e9 | |||
| 2392958114 | |||
| ec04e8e24a | |||
| 4e14ee7f50 | |||
| 7ba4ab0608 | |||
| fd816112fb | |||
| d0ee85be40 | |||
| 9448704af3 | |||
| 9dad9d6ca8 | |||
| 3f41abed7c | |||
| debcbab445 | |||
| 7fcabf1de7 | |||
| e116a1841d | |||
| cd3103ca14 | |||
| 50d07a4b13 | |||
| ed1352936e | |||
| f4b4156a0c | |||
| 5cf2cce0e3 | |||
| 249453d829 | |||
| c14939cecc | |||
| 72f516abb1 | |||
| 66478ed264 | |||
| 6b10dff41d | |||
| f8cc736482 | |||
| a0794fecfc | |||
| c68059e5b3 | |||
| 832ca6b0de | |||
| 89ee43830e | |||
| f7cf13901e | |||
| ad41fa93fb | |||
| 617b7dcd49 | |||
| 417ea032c4 | |||
| b77bb6e200 | |||
| 1fa3b4a600 | |||
| 99bd502f62 | |||
| 25a271dc95 | |||
| 5002ac7716 | |||
| d92a559460 | |||
| 3d571e1a31 | |||
| d338daa4b6 | |||
| 6f802c2a58 | |||
| a3f0168817 | |||
| 677702655f | |||
| b0bbd0c083 | |||
| 5cbf23a1f4 | |||
| 39eb9b34ec | |||
| 5da8616518 | |||
| b267fe05cd | |||
| 29f7ebe559 | |||
| bbffaca511 | |||
| 80532836c3 | |||
| 9474f4f322 | |||
| 93a09d3a9f | |||
| e3935ce699 | |||
| 58c15e7833 | |||
| fd2b7f3aa0 | |||
| 5ccbc629d1 | |||
| e98ff5e8e5 | |||
| a6fffa7b57 | |||
| 3ac153dd06 | |||
| 8db3108c94 | |||
| e25ff4ad19 | |||
| 21e76c6461 | |||
| 103aa1a432 | |||
| d2f4fefcf3 | |||
| 629527988d | |||
| 7f520f1346 | |||
| e28619b55a | |||
| f474e6130e | |||
| 4b5bcb45ac | |||
| 50565a0f17 | |||
| cf37db4fa2 | |||
| ad9b4097ef | |||
| c22c01c6c3 | |||
| 31f7f50c4a | |||
| a7f6ed4b16 | |||
| 73ada5a221 | |||
| 2f96256893 | |||
| 23d9e0775f | |||
| 72ade39144 | |||
| ec64c68777 | |||
| 80932e069f | |||
| 2f9b154b07 | |||
| 20bf911732 | |||
| 65a3dbb228 | |||
| 5844cc93ca | |||
| 4d23ce58c4 | |||
| 2bb592d5f6 | |||
| 3146b20c15 | |||
| 455cf67750 | |||
| a6d6a877b0 | |||
| a7bd54471c | |||
| fe5f803163 | |||
| 66a9b5362a | |||
| f3569cf68b | |||
| 2573f14726 | |||
| f1fb2d6abf | |||
| 4934e0ff0a | |||
| f772a80501 | |||
| 8950843be2 | |||
| 9b89e68908 | |||
| ba134ca53f | |||
| 21dbd9c057 | |||
| 40a68f8e05 | |||
| 37d861a631 | |||
| 31f3e885ce | |||
| 7ffaab2012 | |||
| 35b7946b0d | |||
| 3a05a8e712 | |||
| 294a1149ef | |||
| 8d80370014 | |||
| 1cbdef36cf | |||
| 4c8accbfc1 | |||
| c4c2d9cb93 | |||
| 7aed112326 | |||
| 216a3d53cd | |||
| e0823b343b | |||
| cb0bc65ee4 | |||
| 5b9ab6636f | |||
| 9fd77feebb |
Generated
+354
-773
File diff suppressed because it is too large
Load Diff
+17
-15
@@ -10,7 +10,7 @@ repository = "https://github.com/Dark-Alex-17/loki"
|
|||||||
categories = ["command-line-utilities"]
|
categories = ["command-line-utilities"]
|
||||||
readme = "README.md"
|
readme = "README.md"
|
||||||
license = "MIT"
|
license = "MIT"
|
||||||
rust-version = "1.95.0"
|
rust-version = "1.89.0"
|
||||||
exclude = [".github", "CONTRIBUTING.md"]
|
exclude = [".github", "CONTRIBUTING.md"]
|
||||||
|
|
||||||
[dependencies]
|
[dependencies]
|
||||||
@@ -22,7 +22,7 @@ dunce = "1.0.5"
|
|||||||
futures-util = "0.3.29"
|
futures-util = "0.3.29"
|
||||||
inquire = "0.9.4"
|
inquire = "0.9.4"
|
||||||
is-terminal = "0.4.9"
|
is-terminal = "0.4.9"
|
||||||
reedline = "0.47.0"
|
reedline = "0.46.0"
|
||||||
serde = { version = "1.0.152", features = ["derive"] }
|
serde = { version = "1.0.152", features = ["derive"] }
|
||||||
serde_json = { version = "1.0.93", features = ["preserve_order"] }
|
serde_json = { version = "1.0.93", features = ["preserve_order"] }
|
||||||
serde_yaml = "0.9.17"
|
serde_yaml = "0.9.17"
|
||||||
@@ -34,6 +34,10 @@ tokio = { version = "1.34.0", features = [
|
|||||||
"rt-multi-thread",
|
"rt-multi-thread",
|
||||||
"full",
|
"full",
|
||||||
] }
|
] }
|
||||||
|
tokio-graceful = "0.2.2"
|
||||||
|
tokio-stream = { version = "0.1.15", default-features = false, features = [
|
||||||
|
"sync",
|
||||||
|
] }
|
||||||
crossterm = "0.29.0"
|
crossterm = "0.29.0"
|
||||||
chrono = "0.4.23"
|
chrono = "0.4.23"
|
||||||
bincode = { version = "2.0.0", features = [
|
bincode = { version = "2.0.0", features = [
|
||||||
@@ -47,7 +51,7 @@ nu-ansi-term = "0.50.0"
|
|||||||
async-trait = "0.1.74"
|
async-trait = "0.1.74"
|
||||||
textwrap = "0.16.0"
|
textwrap = "0.16.0"
|
||||||
ansi_colours = "1.2.2"
|
ansi_colours = "1.2.2"
|
||||||
eventsource-stream = "0.2.3"
|
reqwest-eventsource = "0.6.0"
|
||||||
log = "0.4.28"
|
log = "0.4.28"
|
||||||
log4rs = { version = "1.4.0", features = ["file_appender"] }
|
log4rs = { version = "1.4.0", features = ["file_appender"] }
|
||||||
shell-words = "1.1.0"
|
shell-words = "1.1.0"
|
||||||
@@ -55,14 +59,20 @@ sha2 = "0.10.8"
|
|||||||
unicode-width = "0.2.0"
|
unicode-width = "0.2.0"
|
||||||
async-recursion = "1.1.1"
|
async-recursion = "1.1.1"
|
||||||
http = "1.1.0"
|
http = "1.1.0"
|
||||||
|
http-body-util = "0.1"
|
||||||
|
hyper = { version = "1.0", features = ["full"] }
|
||||||
|
hyper-util = { version = "0.1", features = ["server-auto", "client-legacy"] }
|
||||||
|
time = { version = "0.3.36", features = ["macros"] }
|
||||||
indexmap = { version = "2.2.6", features = ["serde"] }
|
indexmap = { version = "2.2.6", features = ["serde"] }
|
||||||
hmac = "0.12.1"
|
hmac = "0.12.1"
|
||||||
aws-smithy-eventstream = "0.60.4"
|
aws-smithy-eventstream = "0.60.4"
|
||||||
urlencoding = "2.1.3"
|
urlencoding = "2.1.3"
|
||||||
|
unicode-segmentation = "1.11.0"
|
||||||
json-patch = { version = "4.0.0", default-features = false }
|
json-patch = { version = "4.0.0", default-features = false }
|
||||||
bitflags = "2.5.0"
|
bitflags = "2.5.0"
|
||||||
path-absolutize = "3.1.1"
|
path-absolutize = "3.1.1"
|
||||||
hnsw_rs = "0.3.0"
|
hnsw_rs = "0.3.0"
|
||||||
|
rayon = "1.10.0"
|
||||||
uuid = { version = "1.9.1", features = ["v4"] }
|
uuid = { version = "1.9.1", features = ["v4"] }
|
||||||
scraper = { version = "0.23.1", default-features = false, features = [
|
scraper = { version = "0.23.1", default-features = false, features = [
|
||||||
"deterministic",
|
"deterministic",
|
||||||
@@ -87,6 +97,7 @@ rmcp = { version = "1.5.0", features = [
|
|||||||
] }
|
] }
|
||||||
num_cpus = "1.17.0"
|
num_cpus = "1.17.0"
|
||||||
tree-sitter = "0.26.8"
|
tree-sitter = "0.26.8"
|
||||||
|
tree-sitter-language = "0.1"
|
||||||
tree-sitter-python = "0.25.0"
|
tree-sitter-python = "0.25.0"
|
||||||
tree-sitter-typescript = "0.23"
|
tree-sitter-typescript = "0.23"
|
||||||
colored = "3.0.0"
|
colored = "3.0.0"
|
||||||
@@ -96,24 +107,15 @@ clap_complete_nushell = "4.5.9"
|
|||||||
open = "5"
|
open = "5"
|
||||||
rand = { version = "0.10.0", features = ["default"] }
|
rand = { version = "0.10.0", features = ["default"] }
|
||||||
url = "2.5.8"
|
url = "2.5.8"
|
||||||
self_update = { version = "0.44", default-features = false, features = [
|
|
||||||
"reqwest",
|
|
||||||
"rustls",
|
|
||||||
"archive-tar",
|
|
||||||
"compression-flate2",
|
|
||||||
"archive-zip",
|
|
||||||
"compression-zip-deflate",
|
|
||||||
] }
|
|
||||||
|
|
||||||
[dependencies.reqwest]
|
[dependencies.reqwest]
|
||||||
version = "0.13.3"
|
version = "0.12.0"
|
||||||
features = [
|
features = [
|
||||||
"json",
|
"json",
|
||||||
"multipart",
|
"multipart",
|
||||||
"stream",
|
|
||||||
"form",
|
|
||||||
"socks",
|
"socks",
|
||||||
"rustls",
|
"rustls-tls",
|
||||||
|
"rustls-tls-native-roots",
|
||||||
]
|
]
|
||||||
default-features = false
|
default-features = false
|
||||||
|
|
||||||
|
|||||||
@@ -10,8 +10,7 @@ Loki is an all-in-one, batteries-included, LLM CLI tool featuring Shell Assistan
|
|||||||
Agents, and More.
|
Agents, and More.
|
||||||
|
|
||||||
It is designed to include a number of useful agents, roles, macros, and more so users can get up and running with Loki
|
It is designed to include a number of useful agents, roles, macros, and more so users can get up and running with Loki
|
||||||
in as little time as possible. You can also install entire bundles of agents, roles, macros, tools, and MCP servers from
|
in as little time as possible.
|
||||||
any git repository — see [Sharing Configurations](#sharing-configurations).
|
|
||||||
|
|
||||||

|

|
||||||
|
|
||||||
@@ -21,7 +20,6 @@ Coming from [AIChat](https://github.com/sigoden/aichat)? Follow the [migration g
|
|||||||
* [AIChat Migration Guide](https://github.com/Dark-Alex-17/loki/wiki/AIChat-Migration): Coming from AIChat? Follow the migration guide to get started.
|
* [AIChat Migration Guide](https://github.com/Dark-Alex-17/loki/wiki/AIChat-Migration): Coming from AIChat? Follow the migration guide to get started.
|
||||||
* [Installation](#install): Install Loki
|
* [Installation](#install): Install Loki
|
||||||
* [Getting Started](#getting-started): Get started with Loki by doing first-run setup steps.
|
* [Getting Started](#getting-started): Get started with Loki by doing first-run setup steps.
|
||||||
* [Sharing Configurations](https://github.com/Dark-Alex-17/loki/wiki/Sharing-Configurations): Install bundles of agents, roles, macros, tools, and MCP servers from any git repo, and share your own.
|
|
||||||
* [REPL](https://github.com/Dark-Alex-17/loki/wiki/REPL): Interactive Read-Eval-Print Loop for conversational interactions with LLMs and Loki.
|
* [REPL](https://github.com/Dark-Alex-17/loki/wiki/REPL): Interactive Read-Eval-Print Loop for conversational interactions with LLMs and Loki.
|
||||||
* [Custom REPL Prompt](https://github.com/Dark-Alex-17/loki/wiki/REPL-Prompt): Customize the REPL prompt to provide useful contextual information.
|
* [Custom REPL Prompt](https://github.com/Dark-Alex-17/loki/wiki/REPL-Prompt): Customize the REPL prompt to provide useful contextual information.
|
||||||
* [Vault](https://github.com/Dark-Alex-17/loki/wiki/Vault): Securely store and manage sensitive information such as API keys and credentials.
|
* [Vault](https://github.com/Dark-Alex-17/loki/wiki/Vault): Securely store and manage sensitive information such as API keys and credentials.
|
||||||
@@ -38,8 +36,7 @@ Coming from [AIChat](https://github.com/sigoden/aichat)? Follow the [migration g
|
|||||||
* [Sessions](https://github.com/Dark-Alex-17/loki/wiki/Sessions): Manage and persist conversational contexts and settings across multiple interactions.
|
* [Sessions](https://github.com/Dark-Alex-17/loki/wiki/Sessions): Manage and persist conversational contexts and settings across multiple interactions.
|
||||||
* [Roles](https://github.com/Dark-Alex-17/loki/wiki/Roles): Customize model behavior for specific tasks or domains.
|
* [Roles](https://github.com/Dark-Alex-17/loki/wiki/Roles): Customize model behavior for specific tasks or domains.
|
||||||
* [Agents](https://github.com/Dark-Alex-17/loki/wiki/Agents): Leverage AI agents to perform complex tasks and workflows, including sub-agent spawning, teammate messaging, and user interaction tools.
|
* [Agents](https://github.com/Dark-Alex-17/loki/wiki/Agents): Leverage AI agents to perform complex tasks and workflows, including sub-agent spawning, teammate messaging, and user interaction tools.
|
||||||
* [Graph Agents](https://github.com/Dark-Alex-17/loki/wiki/Graph-Agents): Define an agent as a declarative, YAML-driven workflow. A directed graph of typed nodes (LLM calls, scripts, approvals, user input, RAG retrieval, sub-agent spawns).
|
* [Todo System](https://github.com/Dark-Alex-17/loki/wiki/TODO-System): Built-in task tracking for improved agent reliability with smaller models.
|
||||||
* [Todo System](https://github.com/Dark-Alex-17/loki/wiki/TODO-System): Built-in task tracking for improved LLM reliability with smaller models.
|
|
||||||
* [Environment Variables](https://github.com/Dark-Alex-17/loki/wiki/Environment-Variables): Override and customize your Loki configuration at runtime with environment variables.
|
* [Environment Variables](https://github.com/Dark-Alex-17/loki/wiki/Environment-Variables): Override and customize your Loki configuration at runtime with environment variables.
|
||||||
* [Client Configurations](https://github.com/Dark-Alex-17/loki/wiki/Clients): Configuration instructions for various LLM providers.
|
* [Client Configurations](https://github.com/Dark-Alex-17/loki/wiki/Clients): Configuration instructions for various LLM providers.
|
||||||
* [Authentication (API Key & OAuth)](https://github.com/Dark-Alex-17/loki/wiki/Clients#authentication): Authenticate with API keys or OAuth for subscription-based access.
|
* [Authentication (API Key & OAuth)](https://github.com/Dark-Alex-17/loki/wiki/Clients#authentication): Authenticate with API keys or OAuth for subscription-based access.
|
||||||
@@ -51,6 +48,16 @@ Coming from [AIChat](https://github.com/sigoden/aichat)? Follow the [migration g
|
|||||||
Loki requires the following tools to be installed on your system:
|
Loki requires the following tools to be installed on your system:
|
||||||
* [jq](https://github.com/jqlang/jq)
|
* [jq](https://github.com/jqlang/jq)
|
||||||
* `brew install jq`
|
* `brew install jq`
|
||||||
|
* [jira (optional)](https://github.com/ankitpokhrel/jira-cli/wiki/Installation) (For the `query_jira_issues` tool)
|
||||||
|
* `brew tap ankitpokhrel/jira-cli && brew install jira-cli`
|
||||||
|
* You'll need to [create a JIRA API token](https://id.atlassian.com/manage-profile/security/api-tokens) for authentication
|
||||||
|
* Then, save it as an environment variable to your shell profile:
|
||||||
|
```sh
|
||||||
|
# ~/.bashrc or ~/.zshrc
|
||||||
|
export JIRA_API_TOKEN="your_jira_api_token_here"
|
||||||
|
```
|
||||||
|
* Then run `jira init`, select installation type as `cloud`, and provide the required details to generate a config
|
||||||
|
file for the Jira CLI.
|
||||||
* [usql](https://github.com/xo/usql) (For the `sql` agent)
|
* [usql](https://github.com/xo/usql) (For the `sql` agent)
|
||||||
* `brew install xo/xo/usql`
|
* `brew install xo/xo/usql`
|
||||||
* [docker](https://docs.docker.com/engine/install/)
|
* [docker](https://docs.docker.com/engine/install/)
|
||||||
@@ -58,7 +65,7 @@ Loki requires the following tools to be installed on your system:
|
|||||||
* `curl -LsSf https://astral.sh/uv/install.sh | sh`
|
* `curl -LsSf https://astral.sh/uv/install.sh | sh`
|
||||||
|
|
||||||
These tools are used to provide various functionalities within Loki, such as document processing, JSON manipulation,
|
These tools are used to provide various functionalities within Loki, such as document processing, JSON manipulation,
|
||||||
etc., and they are used within agents and tools.
|
interaction with Jira, and they are used within agents and tools.
|
||||||
|
|
||||||
## Install
|
## Install
|
||||||
|
|
||||||
@@ -130,29 +137,6 @@ To use a binary from the releases page on Linux/MacOS, do the following:
|
|||||||
3. Extract the binary with `tar -C /usr/local/bin -xzf loki-<arch>.tar.gz` (Note: This may require `sudo`)
|
3. Extract the binary with `tar -C /usr/local/bin -xzf loki-<arch>.tar.gz` (Note: This may require `sudo`)
|
||||||
4. Now you can run `loki`!
|
4. Now you can run `loki`!
|
||||||
|
|
||||||
## Updating
|
|
||||||
Loki can update itself in place to the latest GitHub release. Run `loki --update`
|
|
||||||
for the newest release, or `loki --update v0.4.0` for a specific version:
|
|
||||||
|
|
||||||
```shell
|
|
||||||
loki --update
|
|
||||||
loki --update v0.4.0
|
|
||||||
```
|
|
||||||
|
|
||||||
The same is available from within the REPL via `.update` and `.update v0.4.0`.
|
|
||||||
|
|
||||||
If Loki was installed with a package manager, prefer that package manager so its
|
|
||||||
records stay in sync with the binary on disk; i.e. `brew upgrade loki` for Homebrew,
|
|
||||||
or `cargo install --locked loki-ai` for Cargo.
|
|
||||||
|
|
||||||
When Loki detects a package-manager install it prints a warning and asks for
|
|
||||||
confirmation. In a non-interactive shell (no TTY), pass `--force` to update
|
|
||||||
anyway:
|
|
||||||
|
|
||||||
```shell
|
|
||||||
loki --update --force
|
|
||||||
```
|
|
||||||
|
|
||||||
## Getting Started
|
## Getting Started
|
||||||
After installation, you can generate the configuration files and directories by simply running:
|
After installation, you can generate the configuration files and directories by simply running:
|
||||||
|
|
||||||
@@ -176,11 +160,12 @@ subscribers, Google Gemini), you can authenticate with your existing subscriptio
|
|||||||
# In your config.yaml
|
# In your config.yaml
|
||||||
clients:
|
clients:
|
||||||
- type: claude
|
- type: claude
|
||||||
|
name: my-claude-oauth
|
||||||
auth: oauth # Indicate you want to authenticate with OAuth instead of an API key
|
auth: oauth # Indicate you want to authenticate with OAuth instead of an API key
|
||||||
```
|
```
|
||||||
|
|
||||||
```sh
|
```sh
|
||||||
loki --authenticate claude
|
loki --authenticate my-claude-oauth
|
||||||
# Or via the REPL: .authenticate
|
# Or via the REPL: .authenticate
|
||||||
```
|
```
|
||||||
|
|
||||||
|
|||||||
@@ -1,82 +1,40 @@
|
|||||||
# Coder
|
# Coder
|
||||||
|
|
||||||
A graph-based implementation agent. Plans, implements, and runs build +
|
An AI agent that assists you with your coding tasks.
|
||||||
tests in a bounded fix-loop until verified. Designed to be delegated to by
|
|
||||||
the **[Sisyphus](../sisyphus/README.md)** agent.
|
|
||||||
|
|
||||||
Coder is a [graph agent](https://github.com/Dark-Alex-17/loki/wiki/Graph-Agents): its workflow is
|
This agent is designed to be delegated to by the **[Sisyphus](../sisyphus/README.md)** agent to implement code specifications. Sisyphus
|
||||||
defined declaratively in `graph.yaml`, with verification and the
|
acts as the coordinator/architect, while Coder handles the implementation details.
|
||||||
implement-fix loop enforced as graph edges rather than prose.
|
|
||||||
|
|
||||||
## Workflow
|
## Features
|
||||||
|
|
||||||
```
|
- 🏗️ Intelligent project structure creation and management
|
||||||
analyze_request (llm + output_schema) plan + complexity extraction
|
- 🖼️ Convert screenshots into clean, functional code
|
||||||
↓
|
- 📁 Comprehensive file system operations (create folders, files, read/write files)
|
||||||
route_complexity (script) opt-out approval gate (complexity ≥ 7)
|
- 🧐 Advanced code analysis and improvement suggestions
|
||||||
↓
|
- 📊 Precise diff-based file editing for controlled code modifications
|
||||||
gate_approval (approval, optional)
|
|
||||||
↓
|
|
||||||
implement (llm + fs tools) actual file edits
|
|
||||||
↓
|
|
||||||
verify_build (script)
|
|
||||||
↓
|
|
||||||
verify_tests (script)
|
|
||||||
↓
|
|
||||||
fix_loop_gate (script) back-edge to implement (bounded)
|
|
||||||
↓
|
|
||||||
end_success / end_rejected / end_failure
|
|
||||||
```
|
|
||||||
|
|
||||||
End nodes emit one of three sentinel outcomes for the caller:
|
It can also be used as a standalone tool for direct coding assistance.
|
||||||
|
|
||||||
- `CODER_COMPLETE` — build and tests passed.
|
## Pro-Tip: Use an IDE MCP Server for Improved Performance
|
||||||
- `CODER_REJECTED` — user rejected the plan at the approval gate.
|
Many modern IDEs now include MCP servers that let LLMs perform operations within the IDE itself and use IDE tools. Using
|
||||||
- `CODER_FAILED` — fix-loop exhausted; build/tests still failing.
|
an IDE's MCP server dramatically improves the performance of coding agents. So if you have an IDE, try adding that MCP
|
||||||
|
server to your config (see the [MCP Server docs](../../../docs/function-calling/MCP-SERVERS.md) to see how to configure
|
||||||
## Tuning
|
them), and modify the agent definition to look like this:
|
||||||
|
|
||||||
The agent's `project_dir` is exposed via the standard `variables:` block,
|
|
||||||
so it accepts the runtime override flag:
|
|
||||||
|
|
||||||
```sh
|
|
||||||
# Invoke from inside the project (project_dir defaults to ".")
|
|
||||||
cd /path/to/your/project
|
|
||||||
loki -a coder "Add a foo() function..."
|
|
||||||
|
|
||||||
# Or invoke from anywhere with an explicit override
|
|
||||||
loki -a coder --agent-variable project_dir /path/to/your/project "Add..."
|
|
||||||
```
|
|
||||||
|
|
||||||
`graph.yaml` `initial_state` exposes:
|
|
||||||
|
|
||||||
- `max_fix_attempts` (default `3`) — fix-loop budget before `end_failure`.
|
|
||||||
|
|
||||||
Environment overrides honored by the script nodes:
|
|
||||||
|
|
||||||
- `BUILD_CMD` — skip project-type detection for the build/check command.
|
|
||||||
- `TEST_CMD` — skip detection for tests.
|
|
||||||
- `CODER_AUTOAPPROVE=1` — bypass the approval gate (for non-interactive runs
|
|
||||||
where complexity might trip the gate).
|
|
||||||
|
|
||||||
## Pro-Tip: IDE MCP Server
|
|
||||||
|
|
||||||
Modern IDEs (JetBrains, VS Code, Cursor, Zed, etc.) expose MCP servers
|
|
||||||
that let LLMs use IDE tools directly. To wire one in, edit `graph.yaml`:
|
|
||||||
|
|
||||||
```yaml
|
```yaml
|
||||||
|
# ...
|
||||||
|
|
||||||
mcp_servers:
|
mcp_servers:
|
||||||
- your-ide-mcp-server
|
- jetbrains # The name of your configured IDE MCP server
|
||||||
|
|
||||||
global_tools:
|
global_tools:
|
||||||
# Keep read-only fs tools for files outside the IDE project
|
# Keep useful read-only tools for reading files in other non-project directories
|
||||||
- fs_read.sh
|
- fs_read.sh
|
||||||
- fs_grep.sh
|
- fs_grep.sh
|
||||||
- fs_glob.sh
|
- fs_glob.sh
|
||||||
# - fs_write.sh
|
# - fs_write.sh
|
||||||
# - fs_patch.sh
|
# - fs_patch.sh
|
||||||
- execute_command.sh
|
- execute_command.sh
|
||||||
```
|
|
||||||
|
|
||||||
Then add the MCP server's write/patch tools to the `implement` node's
|
# ...
|
||||||
`tools:` whitelist.
|
```
|
||||||
@@ -0,0 +1,129 @@
|
|||||||
|
name: coder
|
||||||
|
description: Implementation agent - writes code, follows patterns, verifies with builds
|
||||||
|
version: 1.0.0
|
||||||
|
temperature: 0.1
|
||||||
|
|
||||||
|
auto_continue: true
|
||||||
|
max_auto_continues: 15
|
||||||
|
inject_todo_instructions: true
|
||||||
|
|
||||||
|
variables:
|
||||||
|
- name: project_dir
|
||||||
|
description: Project directory to work in
|
||||||
|
default: '.'
|
||||||
|
- name: auto_confirm
|
||||||
|
description: Auto-confirm command execution
|
||||||
|
default: '1'
|
||||||
|
|
||||||
|
global_tools:
|
||||||
|
- fs_read.sh
|
||||||
|
- fs_grep.sh
|
||||||
|
- fs_glob.sh
|
||||||
|
- fs_write.sh
|
||||||
|
- fs_patch.sh
|
||||||
|
- execute_command.sh
|
||||||
|
|
||||||
|
instructions: |
|
||||||
|
You are a senior engineer. You write code that works on the first try.
|
||||||
|
|
||||||
|
## Your Mission
|
||||||
|
|
||||||
|
Given an implementation task:
|
||||||
|
1. Check for orchestrator context first (see below)
|
||||||
|
2. Fill gaps only. Read files NOT already covered in context
|
||||||
|
3. Write the code (using tools, NOT chat output)
|
||||||
|
4. Verify it compiles/builds
|
||||||
|
5. Signal completion with a summary
|
||||||
|
|
||||||
|
## Using Orchestrator Context (IMPORTANT)
|
||||||
|
|
||||||
|
When spawned by sisyphus, your prompt will often contain a `<context>` block
|
||||||
|
with prior findings: file paths, code patterns, and conventions discovered by
|
||||||
|
explore agents.
|
||||||
|
|
||||||
|
**If context is provided:**
|
||||||
|
1. Use it as your primary reference. Don't re-read files already summarized
|
||||||
|
2. Follow the code patterns shown. Snippets in context ARE the style guide
|
||||||
|
3. Read the referenced files ONLY IF you need more detail (e.g. full function
|
||||||
|
signature, import list, or adjacent code not included in the snippet)
|
||||||
|
4. If context includes a "Conventions" section, follow it exactly
|
||||||
|
|
||||||
|
**If context is NOT provided or is too vague to act on:**
|
||||||
|
Fall back to self-exploration: grep for similar files, read 1-2 examples,
|
||||||
|
match their style.
|
||||||
|
|
||||||
|
**Never ignore provided context.** It represents work already done upstream.
|
||||||
|
|
||||||
|
## Todo System
|
||||||
|
|
||||||
|
For multi-file changes:
|
||||||
|
1. `todo__init` with the implementation goal
|
||||||
|
2. `todo__add` for each file to create/modify
|
||||||
|
3. Implement each, calling `todo__done` immediately after
|
||||||
|
|
||||||
|
## Writing Code
|
||||||
|
|
||||||
|
**CRITICAL**: Write code using `write_file` tool, NEVER paste code in chat.
|
||||||
|
|
||||||
|
Correct:
|
||||||
|
```
|
||||||
|
write_file --path "src/user.rs" --content "pub struct User { ... }"
|
||||||
|
```
|
||||||
|
|
||||||
|
Wrong:
|
||||||
|
```
|
||||||
|
Here's the implementation:
|
||||||
|
\`\`\`rust
|
||||||
|
pub struct User { ... }
|
||||||
|
\`\`\`
|
||||||
|
```
|
||||||
|
|
||||||
|
## File Reading Strategy (IMPORTANT - minimize token usage)
|
||||||
|
|
||||||
|
1. **Use grep to find relevant code** - `fs_grep --pattern "fn handle_request" --include "*.rs"` finds where things are
|
||||||
|
2. **Read only what you need** - `fs_read --path "src/main.rs" --offset 50 --limit 30` reads lines 50-79
|
||||||
|
3. **Never cat entire large files** - If 500+ lines, read the relevant section after grepping for it
|
||||||
|
4. **Use glob to find files** - `fs_glob --pattern "*.rs" --path src/` discovers files by name
|
||||||
|
|
||||||
|
## Pattern Matching
|
||||||
|
|
||||||
|
Before writing ANY file:
|
||||||
|
1. Find a similar existing file (use `fs_grep` to locate, then `fs_read` to examine)
|
||||||
|
2. Match its style: imports, naming, structure
|
||||||
|
3. Follow the same patterns exactly
|
||||||
|
|
||||||
|
## Verification
|
||||||
|
|
||||||
|
After writing files:
|
||||||
|
1. Run `verify_build` to check compilation
|
||||||
|
2. If it fails, fix the error (minimal change)
|
||||||
|
3. Don't move on until build passes
|
||||||
|
|
||||||
|
## Completion Signal
|
||||||
|
|
||||||
|
When done, end your response with a summary so the parent agent knows what happened:
|
||||||
|
|
||||||
|
```
|
||||||
|
CODER_COMPLETE: [summary of what was implemented, which files were created/modified, and build status]
|
||||||
|
```
|
||||||
|
|
||||||
|
Or if something went wrong:
|
||||||
|
```
|
||||||
|
CODER_FAILED: [what went wrong]
|
||||||
|
```
|
||||||
|
|
||||||
|
## Rules
|
||||||
|
|
||||||
|
1. **Write code via tools** - Never output code to chat
|
||||||
|
2. **Follow patterns** - Read existing files first
|
||||||
|
3. **Verify builds** - Don't finish without checking
|
||||||
|
4. **Minimal fixes** - If build fails, fix precisely
|
||||||
|
5. **No refactoring** - Only implement what's asked
|
||||||
|
|
||||||
|
## Context
|
||||||
|
- Project: {{project_dir}}
|
||||||
|
- CWD: {{__cwd__}}
|
||||||
|
- Shell: {{__shell__}}
|
||||||
|
|
||||||
|
## Available tools:
|
||||||
|
{{__tools__}}
|
||||||
@@ -1,278 +0,0 @@
|
|||||||
name: coder
|
|
||||||
description: |
|
|
||||||
Implementation agent. Plans, implements, and runs build + tests in a
|
|
||||||
bounded fix-loop until verified. Designed to be delegated to by sisyphus.
|
|
||||||
version: "1.0"
|
|
||||||
|
|
||||||
temperature: 0.1
|
|
||||||
|
|
||||||
global_tools:
|
|
||||||
- fs_cat.sh
|
|
||||||
- fs_ls.sh
|
|
||||||
- fs_write.sh
|
|
||||||
- fs_patch.sh
|
|
||||||
- execute_command.sh
|
|
||||||
|
|
||||||
variables:
|
|
||||||
- name: project_dir
|
|
||||||
description: |
|
|
||||||
Absolute path to the project directory. Defaults to "." which is the
|
|
||||||
directory you invoked `loki` from. Override at runtime with
|
|
||||||
`loki -a coder --agent-variable project_dir /abs/path "..."`.
|
|
||||||
default: "."
|
|
||||||
|
|
||||||
settings:
|
|
||||||
max_loop_iterations: 20
|
|
||||||
log_state_snapshots: true
|
|
||||||
validate_before_run: true
|
|
||||||
timeout: 1800
|
|
||||||
|
|
||||||
initial_state:
|
|
||||||
project_dir: ""
|
|
||||||
fix_attempts: 0
|
|
||||||
max_fix_attempts: 3
|
|
||||||
fix_instructions: ""
|
|
||||||
build_output: ""
|
|
||||||
tests_output: ""
|
|
||||||
last_node_output: ""
|
|
||||||
plan_summary: ""
|
|
||||||
files_to_modify: []
|
|
||||||
files_to_create: []
|
|
||||||
risks: []
|
|
||||||
complexity_score: 0
|
|
||||||
|
|
||||||
start: resolve_paths
|
|
||||||
|
|
||||||
nodes:
|
|
||||||
resolve_paths:
|
|
||||||
id: resolve_paths
|
|
||||||
type: script
|
|
||||||
description: Resolve project_dir to an absolute path from the agent variable
|
|
||||||
script: scripts/resolve_paths.sh
|
|
||||||
timeout: 5
|
|
||||||
fallback: end_failure
|
|
||||||
|
|
||||||
analyze_request:
|
|
||||||
id: analyze_request
|
|
||||||
type: llm
|
|
||||||
description: Extract a structured plan and complexity score from the orchestrator's prompt
|
|
||||||
instructions: |
|
|
||||||
You are a senior engineer's planning assistant. Read the orchestrator's
|
|
||||||
request and emit a structured plan. You only plan. You never edit files.
|
|
||||||
|
|
||||||
Score complexity from 1 to 10:
|
|
||||||
1-3: trivial - single file, <=20 lines changed, obvious approach
|
|
||||||
4-6: moderate - 2-5 files, clear approach, some pattern matching
|
|
||||||
7-10: complex - multi-component, ambiguous tradeoffs, refactoring,
|
|
||||||
or wide blast radius
|
|
||||||
|
|
||||||
Be specific in `files_to_modify` and `files_to_create`. All paths
|
|
||||||
MUST be absolute. The project root is {{project_dir}}. Prefer paths
|
|
||||||
like "{{project_dir}}/src/foo.rs" over "src/foo.rs". The implementer
|
|
||||||
uses these paths directly with fs_write and fs_patch tools, which
|
|
||||||
resolve relative paths against the loki invocation directory (NOT
|
|
||||||
the project dir). Empty arrays are fine if no files in that category.
|
|
||||||
|
|
||||||
`risks` is a list of short strings. Anything that could derail the
|
|
||||||
implementation: unknown dependencies, brittle tests, blast radius,
|
|
||||||
etc. Empty list is fine.
|
|
||||||
|
|
||||||
Project directory: {{project_dir}}
|
|
||||||
prompt: "{{initial_prompt}}"
|
|
||||||
tools: []
|
|
||||||
output_schema:
|
|
||||||
type: object
|
|
||||||
properties:
|
|
||||||
plan_summary:
|
|
||||||
type: string
|
|
||||||
description: 1-3 sentences summarizing what will be done
|
|
||||||
files_to_modify:
|
|
||||||
type: array
|
|
||||||
items: {type: string}
|
|
||||||
files_to_create:
|
|
||||||
type: array
|
|
||||||
items: {type: string}
|
|
||||||
complexity_score:
|
|
||||||
type: integer
|
|
||||||
minimum: 1
|
|
||||||
maximum: 10
|
|
||||||
risks:
|
|
||||||
type: array
|
|
||||||
items: {type: string}
|
|
||||||
required: [plan_summary, files_to_modify, files_to_create, complexity_score, risks]
|
|
||||||
state_updates:
|
|
||||||
last_node_output: "{{output}}"
|
|
||||||
fallback: end_failure
|
|
||||||
next: route_complexity
|
|
||||||
|
|
||||||
route_complexity:
|
|
||||||
id: route_complexity
|
|
||||||
type: script
|
|
||||||
description: Route to approval gate for complex plans; skip otherwise
|
|
||||||
script: scripts/route_complexity.sh
|
|
||||||
timeout: 5
|
|
||||||
fallback: implement
|
|
||||||
|
|
||||||
gate_approval:
|
|
||||||
id: gate_approval
|
|
||||||
type: approval
|
|
||||||
description: Optional human checkpoint for high-complexity plans
|
|
||||||
question: |
|
|
||||||
## Plan
|
|
||||||
{{plan_summary}}
|
|
||||||
|
|
||||||
## Files to modify
|
|
||||||
{{files_to_modify}}
|
|
||||||
|
|
||||||
## Files to create
|
|
||||||
{{files_to_create}}
|
|
||||||
|
|
||||||
## Risks
|
|
||||||
{{risks}}
|
|
||||||
|
|
||||||
Complexity: {{complexity_score}}/10
|
|
||||||
|
|
||||||
Approve this plan?
|
|
||||||
options:
|
|
||||||
- "yes"
|
|
||||||
- "no"
|
|
||||||
routes:
|
|
||||||
"yes": implement
|
|
||||||
"no": end_rejected
|
|
||||||
on_other: end_rejected
|
|
||||||
|
|
||||||
implement:
|
|
||||||
id: implement
|
|
||||||
type: llm
|
|
||||||
description: Write code via fs tools. Bounded tool-call loop.
|
|
||||||
instructions: |
|
|
||||||
You are a senior engineer. Implement the plan by writing code via
|
|
||||||
tools. Follow existing patterns in the codebase.
|
|
||||||
|
|
||||||
## Writing code
|
|
||||||
|
|
||||||
1. Use `fs_patch` for surgical edits to existing files.
|
|
||||||
2. Use `fs_write` for new files or full rewrites.
|
|
||||||
3. NEVER output code to chat. Always use tools.
|
|
||||||
4. ALWAYS pass ABSOLUTE paths to fs_write and fs_patch. Relative
|
|
||||||
paths resolve against the loki invocation directory (not the
|
|
||||||
project dir), which is rarely what you want. The project root
|
|
||||||
is {{project_dir}}.
|
|
||||||
|
|
||||||
## File reading
|
|
||||||
|
|
||||||
1. Use `execute_command` to grep/find:
|
|
||||||
`execute_command --command "grep -rn 'fn handle_request' --include='*.rs' ."`
|
|
||||||
`execute_command --command "find . -name '*.rs' -not -path '*/target/*'"`
|
|
||||||
2. Read only what you need:
|
|
||||||
`fs_cat --path "src/main.rs" --offset 50 --limit 30`
|
|
||||||
3. Never read entire large files. Use offset/limit.
|
|
||||||
4. Use `fs_ls` to list directory contents.
|
|
||||||
|
|
||||||
## Pattern matching
|
|
||||||
|
|
||||||
Before writing ANY file:
|
|
||||||
1. Find a similar existing file (grep, then read).
|
|
||||||
2. Match its style: imports, naming, structure, error handling.
|
|
||||||
3. Follow the same patterns exactly. Do not invent new ones.
|
|
||||||
|
|
||||||
## Fix loop
|
|
||||||
|
|
||||||
If the "Fix loop status" section in your user prompt is non-empty,
|
|
||||||
the previous attempt failed verification. Read the error, identify
|
|
||||||
the minimal fix, apply it. Do not refactor while fixing.
|
|
||||||
|
|
||||||
## Rules
|
|
||||||
|
|
||||||
1. Match existing patterns - read examples first.
|
|
||||||
2. Minimal changes - implement only what's asked.
|
|
||||||
3. Never suppress errors (`as any`, `@ts-ignore`, `#[allow(...)]`
|
|
||||||
on unfamiliar lints, etc.).
|
|
||||||
4. No dead code, no commented-out blocks, no premature abstractions.
|
|
||||||
5. End your turn when editing is done. The graph runs verification next.
|
|
||||||
|
|
||||||
Project directory: {{project_dir}}
|
|
||||||
prompt: |
|
|
||||||
## Plan summary
|
|
||||||
{{plan_summary}}
|
|
||||||
|
|
||||||
## Files involved
|
|
||||||
- Modify: {{files_to_modify}}
|
|
||||||
- Create: {{files_to_create}}
|
|
||||||
|
|
||||||
## Original request from the orchestrator
|
|
||||||
{{initial_prompt}}
|
|
||||||
|
|
||||||
## Fix loop status
|
|
||||||
{{fix_instructions}}
|
|
||||||
tools:
|
|
||||||
- fs_cat
|
|
||||||
- fs_ls
|
|
||||||
- fs_write
|
|
||||||
- fs_patch
|
|
||||||
- execute_command
|
|
||||||
max_iterations: 30
|
|
||||||
state_updates:
|
|
||||||
last_node_output: "{{output}}"
|
|
||||||
fallback: end_failure
|
|
||||||
next: verify_build
|
|
||||||
|
|
||||||
verify_build:
|
|
||||||
id: verify_build
|
|
||||||
type: script
|
|
||||||
description: Run the project's check/build command. Routes to verify_tests on success, fix_loop_gate on failure.
|
|
||||||
script: scripts/verify_build.sh
|
|
||||||
timeout: 300
|
|
||||||
fallback: fix_loop_gate
|
|
||||||
|
|
||||||
verify_tests:
|
|
||||||
id: verify_tests
|
|
||||||
type: script
|
|
||||||
description: Run the project's test command. Routes to end_success on pass, fix_loop_gate on failure.
|
|
||||||
script: scripts/verify_tests.sh
|
|
||||||
timeout: 600
|
|
||||||
fallback: fix_loop_gate
|
|
||||||
|
|
||||||
fix_loop_gate:
|
|
||||||
id: fix_loop_gate
|
|
||||||
type: script
|
|
||||||
description: Budget gate. Loops back to implement with fix_instructions populated, or terminates as end_failure.
|
|
||||||
script: scripts/fix_loop_gate.sh
|
|
||||||
timeout: 5
|
|
||||||
fallback: end_failure
|
|
||||||
|
|
||||||
end_success:
|
|
||||||
id: end_success
|
|
||||||
type: end
|
|
||||||
output: |
|
|
||||||
CODER_COMPLETE
|
|
||||||
Plan: {{plan_summary}}
|
|
||||||
Files modified: {{files_to_modify}}
|
|
||||||
Files created: {{files_to_create}}
|
|
||||||
Build: passed
|
|
||||||
Tests: passed
|
|
||||||
|
|
||||||
end_rejected:
|
|
||||||
id: end_rejected
|
|
||||||
type: end
|
|
||||||
output: |
|
|
||||||
CODER_REJECTED
|
|
||||||
Plan was rejected at the approval gate.
|
|
||||||
Plan: {{plan_summary}}
|
|
||||||
|
|
||||||
end_failure:
|
|
||||||
id: end_failure
|
|
||||||
type: end
|
|
||||||
output: |
|
|
||||||
CODER_FAILED
|
|
||||||
Plan: {{plan_summary}}
|
|
||||||
Attempts: {{fix_attempts}}/{{max_fix_attempts}}
|
|
||||||
|
|
||||||
Last node output:
|
|
||||||
{{last_node_output}}
|
|
||||||
|
|
||||||
Last build output:
|
|
||||||
{{build_output}}
|
|
||||||
|
|
||||||
Last tests output:
|
|
||||||
{{tests_output}}
|
|
||||||
@@ -1,49 +0,0 @@
|
|||||||
#!/usr/bin/env bash
|
|
||||||
set -euo pipefail
|
|
||||||
|
|
||||||
if [[ -n "${GRAPH_STATE_FILE:-}" ]]; then
|
|
||||||
state=$(cat "$GRAPH_STATE_FILE")
|
|
||||||
elif [[ -n "${GRAPH_STATE:-}" ]]; then
|
|
||||||
state="$GRAPH_STATE"
|
|
||||||
else
|
|
||||||
state='{}'
|
|
||||||
fi
|
|
||||||
|
|
||||||
fix_attempts=$(echo "$state" | jq -r '.fix_attempts // 0')
|
|
||||||
max_fix_attempts=$(echo "$state" | jq -r '.max_fix_attempts // 3')
|
|
||||||
build_ok=$(echo "$state" | jq -r '.build_ok | if . == null then "true" else (. | tostring) end')
|
|
||||||
tests_ok=$(echo "$state" | jq -r '.tests_ok | if . == null then "true" else (. | tostring) end')
|
|
||||||
build_output=$(echo "$state" | jq -r '.build_output // ""')
|
|
||||||
tests_output=$(echo "$state" | jq -r '.tests_output // ""')
|
|
||||||
|
|
||||||
if (( fix_attempts >= max_fix_attempts )); then
|
|
||||||
jq -nc \
|
|
||||||
--argjson n "$fix_attempts" \
|
|
||||||
'{
|
|
||||||
"fix_attempts": $n,
|
|
||||||
"_next": "end_failure"
|
|
||||||
}'
|
|
||||||
exit 0
|
|
||||||
fi
|
|
||||||
|
|
||||||
next_attempts=$((fix_attempts + 1))
|
|
||||||
|
|
||||||
if [[ "$build_ok" != "true" ]]; then
|
|
||||||
fix_instructions=$(printf '## Fix loop status (attempt %d of %d)\n\nThe previous attempt failed the build.\n\nBuild output:\n```\n%s\n```\n\nIdentify the minimal fix and apply it. Do not refactor.' \
|
|
||||||
"$next_attempts" "$max_fix_attempts" "$build_output")
|
|
||||||
elif [[ "$tests_ok" != "true" ]]; then
|
|
||||||
fix_instructions=$(printf '## Fix loop status (attempt %d of %d)\n\nBuild passed but tests failed.\n\nTest output:\n```\n%s\n```\n\nIdentify the minimal fix and apply it. Do not refactor.' \
|
|
||||||
"$next_attempts" "$max_fix_attempts" "$tests_output")
|
|
||||||
else
|
|
||||||
fix_instructions=$(printf '## Fix loop status (attempt %d of %d)\n\nfix_loop_gate was reached but no failure was detected in state. Re-run the verification step.' \
|
|
||||||
"$next_attempts" "$max_fix_attempts")
|
|
||||||
fi
|
|
||||||
|
|
||||||
jq -nc \
|
|
||||||
--argjson n "$next_attempts" \
|
|
||||||
--arg fi "$fix_instructions" \
|
|
||||||
'{
|
|
||||||
"fix_attempts": $n,
|
|
||||||
"fix_instructions": $fi,
|
|
||||||
"_next": "implement"
|
|
||||||
}'
|
|
||||||
@@ -1,12 +0,0 @@
|
|||||||
#!/usr/bin/env bash
|
|
||||||
set -euo pipefail
|
|
||||||
|
|
||||||
project_dir="${LLM_AGENT_VAR_PROJECT_DIR:-.}"
|
|
||||||
resolved=$(cd "$project_dir" 2>/dev/null && pwd) || resolved="$project_dir"
|
|
||||||
|
|
||||||
jq -nc \
|
|
||||||
--arg pd "$resolved" \
|
|
||||||
'{
|
|
||||||
"project_dir": $pd,
|
|
||||||
"_next": "analyze_request"
|
|
||||||
}'
|
|
||||||
@@ -1,23 +0,0 @@
|
|||||||
#!/usr/bin/env bash
|
|
||||||
set -euo pipefail
|
|
||||||
|
|
||||||
if [[ -n "${GRAPH_STATE_FILE:-}" ]]; then
|
|
||||||
state=$(cat "$GRAPH_STATE_FILE")
|
|
||||||
elif [[ -n "${GRAPH_STATE:-}" ]]; then
|
|
||||||
state="$GRAPH_STATE"
|
|
||||||
else
|
|
||||||
state='{}'
|
|
||||||
fi
|
|
||||||
|
|
||||||
complexity=$(echo "$state" | jq -r '.complexity_score // 0')
|
|
||||||
|
|
||||||
if [[ "${CODER_AUTOAPPROVE:-0}" == "1" ]]; then
|
|
||||||
jq -nc '{"_next": "implement"}'
|
|
||||||
exit 0
|
|
||||||
fi
|
|
||||||
|
|
||||||
if (( complexity >= 7 )); then
|
|
||||||
jq -nc '{"_next": "gate_approval"}'
|
|
||||||
else
|
|
||||||
jq -nc '{"_next": "implement"}'
|
|
||||||
fi
|
|
||||||
@@ -1,55 +0,0 @@
|
|||||||
#!/usr/bin/env bash
|
|
||||||
set -uo pipefail
|
|
||||||
|
|
||||||
# shellcheck disable=SC1091
|
|
||||||
source "$(dirname "$0")/../../.shared/utils.sh"
|
|
||||||
|
|
||||||
if [[ -n "${GRAPH_STATE_FILE:-}" ]]; then
|
|
||||||
state=$(cat "$GRAPH_STATE_FILE")
|
|
||||||
elif [[ -n "${GRAPH_STATE:-}" ]]; then
|
|
||||||
state="$GRAPH_STATE"
|
|
||||||
else
|
|
||||||
state='{}'
|
|
||||||
fi
|
|
||||||
|
|
||||||
project_dir=$(echo "$state" | jq -r '.project_dir // "."')
|
|
||||||
|
|
||||||
if [[ -n "${BUILD_CMD:-}" ]]; then
|
|
||||||
cmd="$BUILD_CMD"
|
|
||||||
else
|
|
||||||
project_info=$(detect_project "$project_dir")
|
|
||||||
cmd=$(echo "$project_info" | jq -r '.check // .build // ""')
|
|
||||||
fi
|
|
||||||
|
|
||||||
if [[ -z "$cmd" || "$cmd" == "null" ]]; then
|
|
||||||
jq -nc '{
|
|
||||||
"build_ok": true,
|
|
||||||
"build_output": "(no build/check command available for this project type)",
|
|
||||||
"_next": "verify_tests"
|
|
||||||
}'
|
|
||||||
exit 0
|
|
||||||
fi
|
|
||||||
|
|
||||||
exit_code=0
|
|
||||||
output=$(cd "$project_dir" && eval "$cmd" 2>&1) || exit_code=$?
|
|
||||||
|
|
||||||
if (( exit_code == 0 )); then
|
|
||||||
jq -nc \
|
|
||||||
--arg out "$output" \
|
|
||||||
--arg cmd "$cmd" \
|
|
||||||
'{
|
|
||||||
"build_ok": true,
|
|
||||||
"build_output": ("Ran: " + $cmd + "\n\n" + $out),
|
|
||||||
"_next": "verify_tests"
|
|
||||||
}'
|
|
||||||
else
|
|
||||||
jq -nc \
|
|
||||||
--arg out "$output" \
|
|
||||||
--arg cmd "$cmd" \
|
|
||||||
--argjson rc "$exit_code" \
|
|
||||||
'{
|
|
||||||
"build_ok": false,
|
|
||||||
"build_output": ("Ran: " + $cmd + "\nExit code: " + ($rc | tostring) + "\n\n" + $out),
|
|
||||||
"_next": "fix_loop_gate"
|
|
||||||
}'
|
|
||||||
fi
|
|
||||||
@@ -1,55 +0,0 @@
|
|||||||
#!/usr/bin/env bash
|
|
||||||
set -uo pipefail
|
|
||||||
|
|
||||||
# shellcheck disable=SC1091
|
|
||||||
source "$(dirname "$0")/../../.shared/utils.sh"
|
|
||||||
|
|
||||||
if [[ -n "${GRAPH_STATE_FILE:-}" ]]; then
|
|
||||||
state=$(cat "$GRAPH_STATE_FILE")
|
|
||||||
elif [[ -n "${GRAPH_STATE:-}" ]]; then
|
|
||||||
state="$GRAPH_STATE"
|
|
||||||
else
|
|
||||||
state='{}'
|
|
||||||
fi
|
|
||||||
|
|
||||||
project_dir=$(echo "$state" | jq -r '.project_dir // "."')
|
|
||||||
|
|
||||||
if [[ -n "${TEST_CMD:-}" ]]; then
|
|
||||||
cmd="$TEST_CMD"
|
|
||||||
else
|
|
||||||
project_info=$(detect_project "$project_dir")
|
|
||||||
cmd=$(echo "$project_info" | jq -r '.test // ""')
|
|
||||||
fi
|
|
||||||
|
|
||||||
if [[ -z "$cmd" || "$cmd" == "null" ]]; then
|
|
||||||
jq -nc '{
|
|
||||||
"tests_ok": true,
|
|
||||||
"tests_output": "(no test command available for this project type)",
|
|
||||||
"_next": "end_success"
|
|
||||||
}'
|
|
||||||
exit 0
|
|
||||||
fi
|
|
||||||
|
|
||||||
exit_code=0
|
|
||||||
output=$(cd "$project_dir" && eval "$cmd" 2>&1) || exit_code=$?
|
|
||||||
|
|
||||||
if (( exit_code == 0 )); then
|
|
||||||
jq -nc \
|
|
||||||
--arg out "$output" \
|
|
||||||
--arg cmd "$cmd" \
|
|
||||||
'{
|
|
||||||
"tests_ok": true,
|
|
||||||
"tests_output": ("Ran: " + $cmd + "\n\n" + $out),
|
|
||||||
"_next": "end_success"
|
|
||||||
}'
|
|
||||||
else
|
|
||||||
jq -nc \
|
|
||||||
--arg out "$output" \
|
|
||||||
--arg cmd "$cmd" \
|
|
||||||
--argjson rc "$exit_code" \
|
|
||||||
'{
|
|
||||||
"tests_ok": false,
|
|
||||||
"tests_output": ("Ran: " + $cmd + "\nExit code: " + ($rc | tostring) + "\n\n" + $out),
|
|
||||||
"_next": "fix_loop_gate"
|
|
||||||
}'
|
|
||||||
fi
|
|
||||||
@@ -14,6 +14,99 @@ _project_dir() {
|
|||||||
(cd "${dir}" 2>/dev/null && pwd) || echo "${dir}"
|
(cd "${dir}" 2>/dev/null && pwd) || echo "${dir}"
|
||||||
}
|
}
|
||||||
|
|
||||||
|
# Normalize a path to be relative to project root.
|
||||||
|
# Strips the project_dir prefix if the LLM passes an absolute path.
|
||||||
|
# Usage: local rel_path; rel_path=$(_normalize_path "/abs/or/rel/path")
|
||||||
|
_normalize_path() {
|
||||||
|
local input_path="$1"
|
||||||
|
local project_dir
|
||||||
|
project_dir=$(_project_dir)
|
||||||
|
|
||||||
|
if [[ "${input_path}" == /* ]]; then
|
||||||
|
input_path="${input_path#"${project_dir}"/}"
|
||||||
|
fi
|
||||||
|
|
||||||
|
input_path="${input_path#./}"
|
||||||
|
echo "${input_path}"
|
||||||
|
}
|
||||||
|
|
||||||
|
# @cmd Read a file's contents before modifying
|
||||||
|
# @option --path! Path to the file (relative to project root)
|
||||||
|
read_file() {
|
||||||
|
local file_path
|
||||||
|
# shellcheck disable=SC2154
|
||||||
|
file_path=$(_normalize_path "${argc_path}")
|
||||||
|
local project_dir
|
||||||
|
project_dir=$(_project_dir)
|
||||||
|
local full_path="${project_dir}/${file_path}"
|
||||||
|
|
||||||
|
if [[ ! -f "${full_path}" ]]; then
|
||||||
|
warn "File not found: ${file_path}" >> "$LLM_OUTPUT"
|
||||||
|
return 0
|
||||||
|
fi
|
||||||
|
|
||||||
|
{
|
||||||
|
info "Reading: ${file_path}"
|
||||||
|
echo ""
|
||||||
|
cat "${full_path}"
|
||||||
|
} >> "$LLM_OUTPUT"
|
||||||
|
}
|
||||||
|
|
||||||
|
# @cmd Write complete file contents
|
||||||
|
# @option --path! Path for the file (relative to project root)
|
||||||
|
# @option --content! Complete file contents to write
|
||||||
|
write_file() {
|
||||||
|
local file_path
|
||||||
|
file_path=$(_normalize_path "${argc_path}")
|
||||||
|
# shellcheck disable=SC2154
|
||||||
|
local content="${argc_content}"
|
||||||
|
local project_dir
|
||||||
|
project_dir=$(_project_dir)
|
||||||
|
local full_path="${project_dir}/${file_path}"
|
||||||
|
|
||||||
|
mkdir -p "$(dirname "${full_path}")"
|
||||||
|
printf '%s' "${content}" > "${full_path}"
|
||||||
|
|
||||||
|
green "Wrote: ${file_path}" >> "$LLM_OUTPUT"
|
||||||
|
}
|
||||||
|
|
||||||
|
# @cmd Find files similar to a given path (for pattern matching)
|
||||||
|
# @option --path! Path to find similar files for
|
||||||
|
find_similar_files() {
|
||||||
|
local file_path
|
||||||
|
file_path=$(_normalize_path "${argc_path}")
|
||||||
|
local project_dir
|
||||||
|
project_dir=$(_project_dir)
|
||||||
|
|
||||||
|
local ext="${file_path##*.}"
|
||||||
|
local dir
|
||||||
|
dir=$(dirname "${file_path}")
|
||||||
|
|
||||||
|
info "Similar files to: ${file_path}" >> "$LLM_OUTPUT"
|
||||||
|
echo "" >> "$LLM_OUTPUT"
|
||||||
|
|
||||||
|
local results
|
||||||
|
results=$(find "${project_dir}/${dir}" -maxdepth 1 -type f -name "*.${ext}" \
|
||||||
|
! -name "$(basename "${file_path}")" \
|
||||||
|
! -name "*test*" \
|
||||||
|
! -name "*spec*" \
|
||||||
|
2>/dev/null | sed "s|^${project_dir}/||" | head -3)
|
||||||
|
|
||||||
|
if [[ -z "${results}" ]]; then
|
||||||
|
results=$(find "${project_dir}/src" -type f -name "*.${ext}" \
|
||||||
|
! -name "*test*" \
|
||||||
|
! -name "*spec*" \
|
||||||
|
-not -path '*/target/*' \
|
||||||
|
2>/dev/null | sed "s|^${project_dir}/||" | head -3)
|
||||||
|
fi
|
||||||
|
|
||||||
|
if [[ -n "${results}" ]]; then
|
||||||
|
echo "${results}" >> "$LLM_OUTPUT"
|
||||||
|
else
|
||||||
|
warn "No similar files found" >> "$LLM_OUTPUT"
|
||||||
|
fi
|
||||||
|
}
|
||||||
|
|
||||||
# @cmd Verify the project builds successfully
|
# @cmd Verify the project builds successfully
|
||||||
verify_build() {
|
verify_build() {
|
||||||
local project_dir
|
local project_dir
|
||||||
@@ -96,3 +189,28 @@ get_project_structure() {
|
|||||||
} >> "$LLM_OUTPUT"
|
} >> "$LLM_OUTPUT"
|
||||||
}
|
}
|
||||||
|
|
||||||
|
# @cmd Search for content in the codebase
|
||||||
|
# @option --pattern! Pattern to search for
|
||||||
|
search_code() {
|
||||||
|
# shellcheck disable=SC2154
|
||||||
|
local pattern="${argc_pattern}"
|
||||||
|
local project_dir
|
||||||
|
project_dir=$(_project_dir)
|
||||||
|
|
||||||
|
info "Searching: ${pattern}" >> "$LLM_OUTPUT"
|
||||||
|
echo "" >> "$LLM_OUTPUT"
|
||||||
|
|
||||||
|
local results
|
||||||
|
results=$(grep -rn "${pattern}" "${project_dir}" 2>/dev/null | \
|
||||||
|
grep -v '/target/' | \
|
||||||
|
grep -v '/node_modules/' | \
|
||||||
|
grep -v '/.git/' | \
|
||||||
|
sed "s|^${project_dir}/||" | \
|
||||||
|
head -20) || true
|
||||||
|
|
||||||
|
if [[ -n "${results}" ]]; then
|
||||||
|
echo "${results}" >> "$LLM_OUTPUT"
|
||||||
|
else
|
||||||
|
warn "No matches" >> "$LLM_OUTPUT"
|
||||||
|
fi
|
||||||
|
}
|
||||||
@@ -1,274 +0,0 @@
|
|||||||
# deep-research
|
|
||||||
|
|
||||||
A deep web research agent, built as a Loki graph agent. It plans an
|
|
||||||
investigation, decomposes it into sub-questions researched in
|
|
||||||
parallel, grounds the work in a local knowledge corpus, vets the
|
|
||||||
credibility of cited sources, runs a reflexion self-critique loop to
|
|
||||||
revise weak findings, delegates the final write-up to a focused
|
|
||||||
sub-agent, checks that the cited sources are reachable, and gates the
|
|
||||||
result behind human approval.
|
|
||||||
|
|
||||||
Unlike a regular agent (which takes a goal and improvises the steps),
|
|
||||||
this agent runs a fixed graph: every request goes through the same
|
|
||||||
`plan -> parallel research -> vet -> critique -> synthesize -> verify -> approve`
|
|
||||||
pipeline.
|
|
||||||
|
|
||||||
This agent is also the **canonical reference for the Loki graph
|
|
||||||
system**: it exercises every node type (`script`, `llm`, `rag`, `map`,
|
|
||||||
`agent`, `input`, `approval`, `end`) and both static fan-out and
|
|
||||||
dynamic `map` fan-out. If you are learning how to build a graph
|
|
||||||
agent, this is the file to read alongside the
|
|
||||||
[Graph-Agents wiki](https://github.com/Dark-Alex-17/loki/wiki/Graph-Agents).
|
|
||||||
|
|
||||||
## Workflow
|
|
||||||
|
|
||||||
17 nodes. `->` is the static route; a script node can also route
|
|
||||||
dynamically via `_next`. The `▶▶` line is a parallel super-step —
|
|
||||||
those branches run concurrently:
|
|
||||||
|
|
||||||
```
|
|
||||||
parse_request (script) -> bootstrap_research (or -> ask_topic if no topic)
|
|
||||||
ask_topic (input) -> bootstrap_research
|
|
||||||
bootstrap_research (script) -> [plan, knowledge_lookup] ▶▶ parallel
|
|
||||||
plan (llm + output_schema) -> research_each_question
|
|
||||||
knowledge_lookup (rag) -> research_each_question
|
|
||||||
research_each_question (map) -> combine_findings (spawns one branch per question)
|
|
||||||
└─ research_one_question (llm) (atomic; runs N×, joins at map)
|
|
||||||
combine_findings (script) -> vet_sources
|
|
||||||
vet_sources (llm + custom tool) -> critique
|
|
||||||
critique (llm) -> reflexion_gate
|
|
||||||
reflexion_gate (script) -> synthesize (or -> research_each_question: reflexion loop)
|
|
||||||
synthesize (agent: report-writer) -> verify_sources
|
|
||||||
verify_sources (script) -> approve
|
|
||||||
approve (approval) -> end_accepted ("accept")
|
|
||||||
-> end_rejected ("reject")
|
|
||||||
-> incorporate_feedback (any free-form answer)
|
|
||||||
incorporate_feedback (script) -> research_each_question (the human-feedback loop)
|
|
||||||
```
|
|
||||||
|
|
||||||
### Node-type breakdown
|
|
||||||
|
|
||||||
| Type | Nodes |
|
|
||||||
|---|---|
|
|
||||||
| `script` (Python) | `parse_request`, `bootstrap_research`, `combine_findings`, `reflexion_gate`, `verify_sources`, `incorporate_feedback` |
|
|
||||||
| `llm` (tools: `[]`) | `plan`, `critique` |
|
|
||||||
| `llm` (with tool whitelist) | `research_one_question`, `vet_sources` |
|
|
||||||
| `rag` | `knowledge_lookup` — local corpus retrieval |
|
|
||||||
| `map` | `research_each_question` — dynamic fan-out per sub-question |
|
|
||||||
| `agent` | `synthesize` — spawns the `report-writer` sub-agent |
|
|
||||||
| `input` | `ask_topic` |
|
|
||||||
| `approval` | `approve` |
|
|
||||||
| `end` | `end_accepted`, `end_rejected` |
|
|
||||||
|
|
||||||
## Parallel execution
|
|
||||||
|
|
||||||
The graph has two parallel super-steps where Loki's BSP scheduler runs
|
|
||||||
branches concurrently.
|
|
||||||
|
|
||||||
**1. Context loading (`plan` ‖ `knowledge_lookup`)** — after
|
|
||||||
`bootstrap_research`, the LLM planner (which decomposes the topic into
|
|
||||||
sub-questions) and the RAG retrieval over the local `knowledge/`
|
|
||||||
corpus run side by side. They write disjoint state keys (`plan` writes
|
|
||||||
`research_plan` and `questions`; `knowledge_lookup` writes
|
|
||||||
`local_context` and `local_sources`) so no reducer is needed.
|
|
||||||
|
|
||||||
**2. Per-question research (`research_each_question` map)** — the
|
|
||||||
plan emits a `questions` array (3-5 entries, enforced by its
|
|
||||||
`output_schema`). The `map` node spawns one parallel branch per
|
|
||||||
question (`max_concurrency: 3`). Each branch is an isolated
|
|
||||||
`research_one_question` LLM invocation with web tools, instructed to
|
|
||||||
investigate exactly its assigned question. Outputs collect into
|
|
||||||
`question_findings` in input order, then `combine_findings` joins
|
|
||||||
them into a single `findings` Markdown document for downstream nodes.
|
|
||||||
|
|
||||||
`settings.max_concurrency: 4` is the graph-wide cap; the per-`map`
|
|
||||||
override (`max_concurrency: 3` on `research_each_question`) is
|
|
||||||
deliberately lower to leave headroom for the planner's tool calls
|
|
||||||
running alongside RAG.
|
|
||||||
|
|
||||||
## Local knowledge corpus
|
|
||||||
|
|
||||||
`knowledge_lookup` is a `rag` node — it runs hybrid (vector + keyword)
|
|
||||||
retrieval over every file in `knowledge/`. The directory ships with a
|
|
||||||
small `research-style-notes.md` so the RAG node has something to
|
|
||||||
retrieve against on a clean install; drop your own Markdown notes,
|
|
||||||
PDFs, or text files into `knowledge/` to bias the research toward
|
|
||||||
your local context.
|
|
||||||
|
|
||||||
The knowledge base is built once, at agent-load time, into
|
|
||||||
`~/.config/loki/agents/deep-research/knowledge_lookup.yaml`. Because
|
|
||||||
the node fully specifies its build config (`embedding_model`,
|
|
||||||
`chunk_size`, `chunk_overlap`), the build is non-interactive. Delete
|
|
||||||
that cached file after adding or changing knowledge to force a
|
|
||||||
rebuild.
|
|
||||||
|
|
||||||
## Sub-agent: report-writer
|
|
||||||
|
|
||||||
The `synthesize` node is an `agent` node that spawns the
|
|
||||||
`report-writer` sub-agent (`assets/agents/report-writer/`). This is
|
|
||||||
the agent-as-tool pattern: the orchestrating graph delegates the
|
|
||||||
writing phase to a focused sub-agent dedicated to coherent prose,
|
|
||||||
while the research phase uses different (typically cheaper) LLM nodes
|
|
||||||
for fast-and-many-question investigation.
|
|
||||||
|
|
||||||
The `report-writer` sub-agent has no tools — it cannot access the
|
|
||||||
web, cannot search, and cannot invent facts. It reads only the
|
|
||||||
findings it is given and produces a final Markdown report preserving
|
|
||||||
every inline citation. See `assets/agents/report-writer/README.md`
|
|
||||||
for details.
|
|
||||||
|
|
||||||
## Tools and tool scoping
|
|
||||||
|
|
||||||
This agent demonstrates Loki's three tool sources and how an `llm`
|
|
||||||
node's `tools:` whitelist scopes them per node.
|
|
||||||
|
|
||||||
The agent's full tool universe, declared in `graph.yaml`:
|
|
||||||
|
|
||||||
- **Global tools** (`global_tools`): `web_search_loki`,
|
|
||||||
`fetch_url_via_curl`, `search_arxiv` - Loki's built-in tool scripts.
|
|
||||||
- **MCP server** (`mcp_servers`): `ddg-search` - a DuckDuckGo web
|
|
||||||
search MCP server. Referenced in a whitelist as `mcp:ddg-search`.
|
|
||||||
- **Custom agent tool** (`tools.sh`): `classify_source` - a
|
|
||||||
deterministic source-credibility classifier shipped with this agent.
|
|
||||||
|
|
||||||
No node receives all of these. Each `llm` node's `tools:` whitelist
|
|
||||||
narrows the universe to exactly what that step needs:
|
|
||||||
|
|
||||||
| Node | `tools:` whitelist | Draws from |
|
|
||||||
|---|---|---|
|
|
||||||
| `plan`, `critique` | `[]` | nothing - pure reasoning |
|
|
||||||
| `research_one_question` | `web_search_loki`, `fetch_url_via_curl`, `search_arxiv`, `mcp:ddg-search` | global tools + MCP |
|
|
||||||
| `vet_sources` | `classify_source` | the custom tool only |
|
|
||||||
|
|
||||||
`research_one_question` (each parallel branch of the map) can search
|
|
||||||
and fetch but cannot classify sources; `vet_sources` can classify
|
|
||||||
sources but cannot touch the web. That separation is the point of the
|
|
||||||
`tools:` whitelist: a node gets only the tools its job calls for,
|
|
||||||
never the agent's full set.
|
|
||||||
|
|
||||||
The `classify_source` custom tool (`tools.sh`) takes a URL and returns
|
|
||||||
a credibility tier (government, academic, preprint, organization,
|
|
||||||
unverified) derived from the host and top-level domain. It is
|
|
||||||
deterministic - exactly the kind of logic a tool should own rather than
|
|
||||||
the LLM guessing.
|
|
||||||
|
|
||||||
Web search may require API-key configuration; see the
|
|
||||||
[Tools](https://github.com/Dark-Alex-17/loki/wiki/Tools) docs.
|
|
||||||
`fetch_url_via_curl`, `search_arxiv`, and `classify_source` work
|
|
||||||
without a key.
|
|
||||||
|
|
||||||
## Setup
|
|
||||||
|
|
||||||
`research_one_question` (each parallel branch of the `map`) uses the
|
|
||||||
`ddg-search` MCP server via `mcp:ddg-search`. It is one of Loki's
|
|
||||||
default MCP servers; make sure it is registered in
|
|
||||||
`~/.config/loki/mcp.json` (run `loki --install mcp_config` to restore
|
|
||||||
the default template if it is missing). If `ddg-search` is unavailable,
|
|
||||||
the branches still have their global web-search tools to fall back on.
|
|
||||||
|
|
||||||
The `synthesize` node spawns the `report-writer` sub-agent. Both
|
|
||||||
agents ship with `loki agents install`; if you install one manually,
|
|
||||||
install both so the agent reference resolves.
|
|
||||||
|
|
||||||
## Reflexion
|
|
||||||
|
|
||||||
The agent has two loops, both built with script nodes that route via
|
|
||||||
`_next`. The engine allows back-edges at runtime; the validator only
|
|
||||||
rejects cycles built from static `next` / `routes` edges, so script
|
|
||||||
`_next` loops are always allowed.
|
|
||||||
|
|
||||||
**Automated reflexion loop.** After the parallel research map and
|
|
||||||
`vet_sources`, the `critique` node reviews the merged findings
|
|
||||||
against the research plan and the source credibility assessment, and
|
|
||||||
emits `VERDICT: PASS` or `VERDICT: REVISE` with specific feedback.
|
|
||||||
`reflexion_gate.py` then:
|
|
||||||
|
|
||||||
- `PASS` -> continue to `synthesize`.
|
|
||||||
- `REVISE`, budget remaining -> loop back to `research_each_question`,
|
|
||||||
with the critique injected as `research_feedback` so every parallel
|
|
||||||
branch sees it on the retry.
|
|
||||||
- `REVISE`, budget spent -> continue to `synthesize` anyway (the human
|
|
||||||
approval step is the final backstop).
|
|
||||||
|
|
||||||
The budget is `MAX_REFLEXION_REVISIONS` in `reflexion_gate.py`
|
|
||||||
(default 2, so the research map runs at most 3 times per pass).
|
|
||||||
|
|
||||||
**Human-feedback loop.** At `approve` the user answers `accept`,
|
|
||||||
`reject`, or types their own feedback. A free-form answer routes via
|
|
||||||
the approval node's `on_other` to `incorporate_feedback.py`, which
|
|
||||||
folds that text into `research_feedback` and loops back to
|
|
||||||
`research_each_question` for another parallel pass.
|
|
||||||
|
|
||||||
`settings.max_loop_iterations` (40) is the engine's infinite-loop
|
|
||||||
backstop: it caps the total visits to any single node.
|
|
||||||
|
|
||||||
## Running
|
|
||||||
|
|
||||||
```sh
|
|
||||||
loki agents install # ships deep-research
|
|
||||||
loki -a deep-research "How does HTTP/3 differ from HTTP/2?"
|
|
||||||
loki -a deep-research "Recent advances in solid-state batteries"
|
|
||||||
loki -a deep-research # no prompt -> triggers ask_topic
|
|
||||||
```
|
|
||||||
|
|
||||||
## Anti-hallucination
|
|
||||||
|
|
||||||
- `research_one_question` (each map branch) is instructed to back
|
|
||||||
every claim with a real retrieved source and never to fabricate
|
|
||||||
URLs, titles, or DOIs.
|
|
||||||
- `vet_sources` classifies every cited source so weak sources are
|
|
||||||
visible to the critique step.
|
|
||||||
- `critique` independently reviews the merged findings and sends weak
|
|
||||||
or uncited work back for another parallel research pass.
|
|
||||||
- `synthesize` (the `report-writer` sub-agent) is grounded: it may use
|
|
||||||
only the gathered findings and must keep each claim's inline source.
|
|
||||||
It has no tools and cannot browse the web.
|
|
||||||
- `verify_sources` probes every cited URL / DOI with an HTTP HEAD
|
|
||||||
request and reports which are unreachable, so the human reviewer
|
|
||||||
sees broken citations before approving.
|
|
||||||
|
|
||||||
## Customizing
|
|
||||||
|
|
||||||
- **Loop budget.** `MAX_REFLEXION_REVISIONS` in `reflexion_gate.py`.
|
|
||||||
- **Map concurrency.** The `research_each_question` node's
|
|
||||||
`max_concurrency: 3` caps simultaneous web-research branches.
|
|
||||||
Raise to investigate more questions in parallel; lower to be gentle
|
|
||||||
on rate-limited providers.
|
|
||||||
- **Per-node model.** Add `model: anthropic:...` to any `llm` node.
|
|
||||||
Cheap models work well for `plan` / `critique` / `vet_sources`; the
|
|
||||||
heavy intelligence is needed in `research_one_question` and the
|
|
||||||
`report-writer` sub-agent.
|
|
||||||
- **Tool scope.** Narrow the `research_one_question` node's `tools:`
|
|
||||||
list to constrain where each branch looks (for example, drop
|
|
||||||
`web_search_loki` and `mcp:ddg-search` to force arXiv-only
|
|
||||||
research).
|
|
||||||
- **Local knowledge.** Drop files into `knowledge/` to bias every
|
|
||||||
research branch toward your local context (see the *Local
|
|
||||||
knowledge corpus* section above).
|
|
||||||
- **Different writer.** Replace `agent: report-writer` on the
|
|
||||||
`synthesize` node with the name of any other agent. The
|
|
||||||
orchestrator does not care what kind of agent the writer is.
|
|
||||||
- **Skip approval.** Point both `approve` routes at `end_accepted`,
|
|
||||||
or wire `verify_sources` straight to an `end` node.
|
|
||||||
|
|
||||||
## Files
|
|
||||||
|
|
||||||
```
|
|
||||||
assets/agents/deep-research/
|
|
||||||
graph.yaml - agent config + 17-node workflow
|
|
||||||
tools.sh - classify_source custom tool
|
|
||||||
README.md - this file
|
|
||||||
knowledge/
|
|
||||||
README.md - corpus-format notes
|
|
||||||
research-style-notes.md - starter knowledge file (replace with your notes)
|
|
||||||
scripts/
|
|
||||||
parse_request.py - _next: bootstrap_research, or ask_topic if no topic
|
|
||||||
bootstrap_research.py - fan-out source: next [plan, knowledge_lookup]
|
|
||||||
combine_findings.py - joins map output (question_findings) into findings
|
|
||||||
reflexion_gate.py - _next: research_each_question (revise) or synthesize
|
|
||||||
verify_sources.py - HTTP HEAD on cited URLs / DOIs
|
|
||||||
incorporate_feedback.py - _next: research_each_question, with user feedback
|
|
||||||
```
|
|
||||||
|
|
||||||
See also `assets/agents/report-writer/` — the sub-agent the
|
|
||||||
`synthesize` node spawns.
|
|
||||||
@@ -1,293 +0,0 @@
|
|||||||
name: deep-research
|
|
||||||
description: |
|
|
||||||
Deep web research workflow. Plans an investigation, decomposes it
|
|
||||||
into sub-questions researched in parallel, grounds the work in a
|
|
||||||
local knowledge corpus, vets the credibility of cited sources, runs
|
|
||||||
a reflexion self-critique loop to revise weak or incomplete findings,
|
|
||||||
delegates the final write-up to a focused sub-agent, checks that the
|
|
||||||
cited sources are reachable, and gates the result behind human
|
|
||||||
approval. A reviewer's free-form feedback at the approval step feeds
|
|
||||||
back into another research pass.
|
|
||||||
|
|
||||||
This is the canonical Loki graph-agent reference: it exercises every
|
|
||||||
node type (script, llm, rag, map, agent, input, approval, end) and
|
|
||||||
both static fan-out and dynamic map fan-out.
|
|
||||||
|
|
||||||
version: "1.0"
|
|
||||||
|
|
||||||
temperature: 0.0
|
|
||||||
|
|
||||||
global_tools:
|
|
||||||
- web_search_loki.sh
|
|
||||||
- fetch_url_via_curl.sh
|
|
||||||
- search_arxiv.sh
|
|
||||||
|
|
||||||
mcp_servers:
|
|
||||||
- ddg-search
|
|
||||||
|
|
||||||
conversation_starters:
|
|
||||||
- "How does HTTP/3 differ from HTTP/2?"
|
|
||||||
- "Summarize recent advances in solid-state battery chemistry"
|
|
||||||
|
|
||||||
settings:
|
|
||||||
max_loop_iterations: 40
|
|
||||||
log_state_snapshots: false
|
|
||||||
validate_before_run: true
|
|
||||||
max_concurrency: 4
|
|
||||||
|
|
||||||
initial_state:
|
|
||||||
research_feedback: ""
|
|
||||||
research_attempts: 0
|
|
||||||
local_context: ""
|
|
||||||
local_sources: ""
|
|
||||||
|
|
||||||
start: parse_request
|
|
||||||
|
|
||||||
nodes:
|
|
||||||
|
|
||||||
parse_request:
|
|
||||||
id: parse_request
|
|
||||||
type: script
|
|
||||||
script: scripts/parse_request.py
|
|
||||||
next: bootstrap_research
|
|
||||||
|
|
||||||
ask_topic:
|
|
||||||
id: ask_topic
|
|
||||||
type: input
|
|
||||||
question: "What would you like me to research?"
|
|
||||||
validation: "len(input) > 0"
|
|
||||||
state_updates:
|
|
||||||
topic: "{{input}}"
|
|
||||||
next: bootstrap_research
|
|
||||||
|
|
||||||
bootstrap_research:
|
|
||||||
id: bootstrap_research
|
|
||||||
type: script
|
|
||||||
script: scripts/bootstrap_research.py
|
|
||||||
next: [plan, knowledge_lookup]
|
|
||||||
|
|
||||||
plan:
|
|
||||||
id: plan
|
|
||||||
type: llm
|
|
||||||
instructions: |
|
|
||||||
You are a research planner. Given a topic, produce a focused
|
|
||||||
research plan and decompose it into 3-5 specific sub-questions
|
|
||||||
that can each be researched independently in parallel.
|
|
||||||
|
|
||||||
The plan is a short narrative naming the key questions and the
|
|
||||||
kinds of sources that would be authoritative. The sub-questions
|
|
||||||
are precise, self-contained queries (each one is sent on its own
|
|
||||||
to a separate research worker, so they must be answerable
|
|
||||||
without each other's context).
|
|
||||||
prompt: "Research topic: {{topic}}"
|
|
||||||
tools: []
|
|
||||||
output_schema:
|
|
||||||
type: object
|
|
||||||
properties:
|
|
||||||
research_plan:
|
|
||||||
type: string
|
|
||||||
description: A short plan narrative.
|
|
||||||
questions:
|
|
||||||
type: array
|
|
||||||
items: { type: string }
|
|
||||||
minItems: 1
|
|
||||||
maxItems: 6
|
|
||||||
description: 3-5 specific, self-contained sub-questions.
|
|
||||||
required: [research_plan, questions]
|
|
||||||
next: research_each_question
|
|
||||||
|
|
||||||
knowledge_lookup:
|
|
||||||
id: knowledge_lookup
|
|
||||||
type: rag
|
|
||||||
documents:
|
|
||||||
- ./knowledge/
|
|
||||||
query: "{{topic}}"
|
|
||||||
top_k: 6
|
|
||||||
chunk_size: 1000
|
|
||||||
chunk_overlap: 100
|
|
||||||
state_updates:
|
|
||||||
local_context: "{{output.context}}"
|
|
||||||
local_sources: "{{output.sources}}"
|
|
||||||
next: research_each_question
|
|
||||||
|
|
||||||
research_each_question:
|
|
||||||
id: research_each_question
|
|
||||||
type: map
|
|
||||||
over: "{{questions}}"
|
|
||||||
as: question
|
|
||||||
branch: research_one_question
|
|
||||||
collect_into: question_findings
|
|
||||||
max_concurrency: 3
|
|
||||||
next: combine_findings
|
|
||||||
|
|
||||||
research_one_question:
|
|
||||||
id: research_one_question
|
|
||||||
type: llm
|
|
||||||
instructions: |
|
|
||||||
You are a web research assistant. Investigate the SINGLE question
|
|
||||||
given to you using your tools: search the web, fetch and read
|
|
||||||
pages, and search arXiv for academic sources.
|
|
||||||
|
|
||||||
Rules:
|
|
||||||
- Every factual claim must be backed by a real source you
|
|
||||||
actually retrieved. Never fabricate URLs, page titles,
|
|
||||||
authors, or DOIs.
|
|
||||||
- Prefer primary and authoritative sources over aggregators.
|
|
||||||
- Where sources disagree, report the disagreement rather than
|
|
||||||
papering over it.
|
|
||||||
- Put the URL (or DOI) inline next to each claim it supports.
|
|
||||||
|
|
||||||
Return organized findings in plain text. Do not include
|
|
||||||
meta-commentary about the process.
|
|
||||||
prompt: |
|
|
||||||
Research question: {{question}}
|
|
||||||
|
|
||||||
Local context that may help:
|
|
||||||
{{local_context}}
|
|
||||||
|
|
||||||
{{research_feedback}}
|
|
||||||
tools:
|
|
||||||
- web_search_loki
|
|
||||||
- fetch_url_via_curl
|
|
||||||
- search_arxiv
|
|
||||||
- mcp:ddg-search
|
|
||||||
max_iterations: 10
|
|
||||||
max_attempts: 2
|
|
||||||
temperature: 0.1
|
|
||||||
|
|
||||||
combine_findings:
|
|
||||||
id: combine_findings
|
|
||||||
type: script
|
|
||||||
script: scripts/combine_findings.py
|
|
||||||
next: vet_sources
|
|
||||||
|
|
||||||
vet_sources:
|
|
||||||
id: vet_sources
|
|
||||||
type: llm
|
|
||||||
instructions: |
|
|
||||||
You assess the credibility of the sources cited in a set of
|
|
||||||
research findings. For every distinct source URL in the findings,
|
|
||||||
call the `classify_source` tool to get its credibility tier. Then
|
|
||||||
summarize: which claims rest on HIGH-credibility sources, and
|
|
||||||
which rest on PREPRINT or UNVERIFIED sources and so need
|
|
||||||
corroboration. Do NOT do any new research -- assess only what is
|
|
||||||
already cited.
|
|
||||||
prompt: |
|
|
||||||
Findings to assess:
|
|
||||||
{{findings}}
|
|
||||||
tools:
|
|
||||||
- classify_source
|
|
||||||
max_iterations: 15
|
|
||||||
state_updates:
|
|
||||||
source_assessment: "{{output}}"
|
|
||||||
next: critique
|
|
||||||
|
|
||||||
critique:
|
|
||||||
id: critique
|
|
||||||
type: llm
|
|
||||||
instructions: |
|
|
||||||
You are a meticulous research reviewer. Judge whether the
|
|
||||||
findings below are good enough to synthesize a complete,
|
|
||||||
well-supported report that answers the research plan.
|
|
||||||
|
|
||||||
Mark the findings REVISE if ANY of these hold:
|
|
||||||
- A research-plan question is unanswered or only weakly
|
|
||||||
addressed.
|
|
||||||
- A factual claim has no source, or cites a source that looks
|
|
||||||
fabricated.
|
|
||||||
- The findings lean on a single source where corroboration is
|
|
||||||
needed.
|
|
||||||
- A key claim rests only on a PREPRINT or UNVERIFIED source,
|
|
||||||
per the source credibility assessment below.
|
|
||||||
- An obvious counter-perspective or recent development is
|
|
||||||
missing.
|
|
||||||
Otherwise mark them PASS.
|
|
||||||
|
|
||||||
Respond in EXACTLY this format, nothing else:
|
|
||||||
|
|
||||||
VERDICT: <PASS or REVISE>
|
|
||||||
FEEDBACK: <if REVISE, be specific and actionable -- name the gaps
|
|
||||||
and what kind of source would close them; if PASS, write "none">
|
|
||||||
prompt: |
|
|
||||||
Research plan:
|
|
||||||
{{research_plan}}
|
|
||||||
|
|
||||||
Findings under review:
|
|
||||||
{{findings}}
|
|
||||||
|
|
||||||
Source credibility assessment:
|
|
||||||
{{source_assessment}}
|
|
||||||
tools: []
|
|
||||||
state_updates:
|
|
||||||
critique: "{{output}}"
|
|
||||||
next: reflexion_gate
|
|
||||||
|
|
||||||
reflexion_gate:
|
|
||||||
id: reflexion_gate
|
|
||||||
type: script
|
|
||||||
script: scripts/reflexion_gate.py
|
|
||||||
next: synthesize
|
|
||||||
|
|
||||||
synthesize:
|
|
||||||
id: synthesize
|
|
||||||
type: agent
|
|
||||||
agent: report-writer
|
|
||||||
prompt: |
|
|
||||||
Research topic: {{topic}}
|
|
||||||
|
|
||||||
Findings (organized by sub-question, with inline citations):
|
|
||||||
{{findings}}
|
|
||||||
|
|
||||||
Source credibility assessment:
|
|
||||||
{{source_assessment}}
|
|
||||||
|
|
||||||
Produce the final report following your instructions.
|
|
||||||
timeout: 300
|
|
||||||
state_updates:
|
|
||||||
report: "{{output}}"
|
|
||||||
next: verify_sources
|
|
||||||
|
|
||||||
verify_sources:
|
|
||||||
id: verify_sources
|
|
||||||
type: script
|
|
||||||
script: scripts/verify_sources.py
|
|
||||||
next: approve
|
|
||||||
|
|
||||||
approve:
|
|
||||||
id: approve
|
|
||||||
type: approval
|
|
||||||
question: |
|
|
||||||
Research report on: {{topic}}
|
|
||||||
|
|
||||||
{{report}}
|
|
||||||
|
|
||||||
----
|
|
||||||
{{source_check}}
|
|
||||||
----
|
|
||||||
|
|
||||||
Accept this report? Pick "accept" or "reject", or type specific
|
|
||||||
feedback to send the research back for another pass.
|
|
||||||
options:
|
|
||||||
- "accept"
|
|
||||||
- "reject"
|
|
||||||
routes:
|
|
||||||
"accept": end_accepted
|
|
||||||
"reject": end_rejected
|
|
||||||
on_other: incorporate_feedback
|
|
||||||
state_updates:
|
|
||||||
decision: "{{choice}}"
|
|
||||||
|
|
||||||
incorporate_feedback:
|
|
||||||
id: incorporate_feedback
|
|
||||||
type: script
|
|
||||||
script: scripts/incorporate_feedback.py
|
|
||||||
|
|
||||||
end_accepted:
|
|
||||||
id: end_accepted
|
|
||||||
type: end
|
|
||||||
output: "{{report}}"
|
|
||||||
|
|
||||||
end_rejected:
|
|
||||||
id: end_rejected
|
|
||||||
type: end
|
|
||||||
output: "Research on '{{topic}}' was rejected and discarded."
|
|
||||||
@@ -1,23 +0,0 @@
|
|||||||
# Local knowledge corpus for deep-research
|
|
||||||
|
|
||||||
The `knowledge_lookup` node in `graph.yaml` is a `rag` node that runs
|
|
||||||
hybrid (vector + keyword) retrieval over every file in this directory.
|
|
||||||
Drop your own notes, papers (PDFs), Markdown docs, or text files here
|
|
||||||
and they will be indexed into a per-agent knowledge base on first run.
|
|
||||||
|
|
||||||
Loki supports common file types out of the box: `.md`, `.txt`, `.pdf`,
|
|
||||||
`.html`, and others. Subdirectories are walked recursively.
|
|
||||||
|
|
||||||
A small starter file (`research-style-notes.md`) ships so the RAG
|
|
||||||
node has something non-empty to retrieve against on a clean install.
|
|
||||||
Replace or extend it with your own materials to bias the research
|
|
||||||
phase toward your local context.
|
|
||||||
|
|
||||||
To force the knowledge base to rebuild after you add or change files,
|
|
||||||
delete the cached index:
|
|
||||||
|
|
||||||
```sh
|
|
||||||
rm ~/.config/loki/agents/deep-research/knowledge_lookup.yaml
|
|
||||||
```
|
|
||||||
|
|
||||||
The next run will rebuild from the current contents of this directory.
|
|
||||||
@@ -1,49 +0,0 @@
|
|||||||
# Research style notes
|
|
||||||
|
|
||||||
These are general principles the `deep-research` agent should keep in
|
|
||||||
mind regardless of topic. Replace this file with your own notes if you
|
|
||||||
want to bias retrieval toward your local context.
|
|
||||||
|
|
||||||
## What "good research" means here
|
|
||||||
|
|
||||||
- **Every factual claim cites a source you actually retrieved.** Never
|
|
||||||
fabricate URLs, page titles, authors, or DOIs.
|
|
||||||
- **Primary sources beat aggregators.** Prefer the original paper, the
|
|
||||||
RFC, the standards body, or the manufacturer over a blog summarizing
|
|
||||||
them.
|
|
||||||
- **Corroboration matters where stakes are high.** If a single source
|
|
||||||
makes a strong claim, look for a second independent source before
|
|
||||||
taking it as established.
|
|
||||||
- **Disagreement is information, not noise.** If two credible sources
|
|
||||||
disagree, report the disagreement and the reasoning on each side.
|
|
||||||
- **Old does not mean wrong.** A 2014 RFC is still authoritative if no
|
|
||||||
newer one has obsoleted it; check before assuming a source is stale.
|
|
||||||
|
|
||||||
## Source-tier heuristics
|
|
||||||
|
|
||||||
The `vet_sources` node uses these rough tiers to weigh credibility.
|
|
||||||
The custom tool `classify_source` (see `tools.sh`) implements this
|
|
||||||
deterministically by hostname / TLD.
|
|
||||||
|
|
||||||
- **HIGH:** government domains (`.gov`, `.mil`), academic institutions
|
|
||||||
(`.edu`, university subdomains), peer-reviewed journals, standards
|
|
||||||
bodies (IETF/RFCs, W3C, ISO, IEEE, NIST), and primary documents from
|
|
||||||
the entities being researched (e.g. a vendor's official spec page).
|
|
||||||
- **PREPRINT:** arXiv, bioRxiv, medRxiv, SSRN. Useful but not yet
|
|
||||||
peer-reviewed; treat numeric claims with extra caution.
|
|
||||||
- **ORGANIZATION:** established nonprofits, standards-adjacent groups,
|
|
||||||
industry consortia. Reliable for their stated mission but may have a
|
|
||||||
perspective.
|
|
||||||
- **UNVERIFIED:** general web pages, blogs, news aggregators, social
|
|
||||||
media. Useful for leads but should not be the only source for a
|
|
||||||
factual claim.
|
|
||||||
|
|
||||||
## Common pitfalls to flag in critique
|
|
||||||
|
|
||||||
- A claim cited only to a PREPRINT or UNVERIFIED source on a numeric
|
|
||||||
or contested point.
|
|
||||||
- A research-plan question that the findings address only obliquely.
|
|
||||||
- "Findings" that paraphrase a single source three times rather than
|
|
||||||
triangulating.
|
|
||||||
- Citation collisions where two sources are listed but turn out to
|
|
||||||
be the same study reported via different aggregators.
|
|
||||||
@@ -1,18 +0,0 @@
|
|||||||
#!/usr/bin/env python3
|
|
||||||
"""Fan-out source for context loading.
|
|
||||||
|
|
||||||
Has no logic of its own. Exists so the static `next: [plan, knowledge_lookup]`
|
|
||||||
list on this node fans out into two parallel branches (the LLM planner and
|
|
||||||
the RAG knowledge lookup) as a single super-step. The validator requires
|
|
||||||
declared parallel-branch script outputs, so we emit an empty JSON object
|
|
||||||
explicitly here.
|
|
||||||
"""
|
|
||||||
import json
|
|
||||||
|
|
||||||
|
|
||||||
def main():
|
|
||||||
print(json.dumps({}))
|
|
||||||
|
|
||||||
|
|
||||||
if __name__ == "__main__":
|
|
||||||
main()
|
|
||||||
@@ -1,39 +0,0 @@
|
|||||||
#!/usr/bin/env python3
|
|
||||||
"""Join the per-question map outputs into a single `findings` string.
|
|
||||||
|
|
||||||
The `research_each_question` map writes `question_findings` (an array,
|
|
||||||
one entry per sub-question, in input order). Downstream nodes
|
|
||||||
(`vet_sources`, `critique`, `synthesize`) read `{{findings}}` as a
|
|
||||||
single block, so this script renders the array as a Markdown document
|
|
||||||
with one section per question.
|
|
||||||
"""
|
|
||||||
import json
|
|
||||||
import os
|
|
||||||
|
|
||||||
|
|
||||||
def load_state():
|
|
||||||
path = os.environ.get("GRAPH_STATE_FILE")
|
|
||||||
if path:
|
|
||||||
with open(path) as f:
|
|
||||||
return json.load(f)
|
|
||||||
return json.loads(os.environ.get("GRAPH_STATE", "{}"))
|
|
||||||
|
|
||||||
|
|
||||||
def main():
|
|
||||||
state = load_state()
|
|
||||||
questions = state.get("questions") or []
|
|
||||||
per_question = state.get("question_findings") or []
|
|
||||||
|
|
||||||
sections = []
|
|
||||||
for idx, q in enumerate(questions):
|
|
||||||
body = per_question[idx] if idx < len(per_question) else ""
|
|
||||||
if isinstance(body, dict) or isinstance(body, list):
|
|
||||||
body = json.dumps(body, indent=2)
|
|
||||||
sections.append(f"## {q}\n\n{body}")
|
|
||||||
|
|
||||||
findings = "\n\n".join(sections) if sections else "No findings gathered."
|
|
||||||
print(json.dumps({"findings": findings}))
|
|
||||||
|
|
||||||
|
|
||||||
if __name__ == "__main__":
|
|
||||||
main()
|
|
||||||
@@ -1,41 +0,0 @@
|
|||||||
#!/usr/bin/env python3
|
|
||||||
"""Fold a reviewer's free-form feedback back into the research loop.
|
|
||||||
|
|
||||||
Runs when the user answers the approval step with their own text
|
|
||||||
instead of "accept" or "reject". That text (saved by the approval node
|
|
||||||
as `decision`) becomes `research_feedback`, and the graph loops back to
|
|
||||||
`research_each_question` for another informed pass (each sub-question is
|
|
||||||
re-researched in parallel with the new feedback in context). The
|
|
||||||
reflexion counter is reset so the user-driven pass gets a fresh revision
|
|
||||||
budget.
|
|
||||||
|
|
||||||
Routing (`_next`): always research_each_question.
|
|
||||||
"""
|
|
||||||
import json
|
|
||||||
import os
|
|
||||||
|
|
||||||
|
|
||||||
def load_state():
|
|
||||||
path = os.environ.get("GRAPH_STATE_FILE")
|
|
||||||
if path:
|
|
||||||
with open(path) as f:
|
|
||||||
return json.load(f)
|
|
||||||
return json.loads(os.environ.get("GRAPH_STATE", "{}"))
|
|
||||||
|
|
||||||
|
|
||||||
def main():
|
|
||||||
state = load_state()
|
|
||||||
feedback = (state.get("decision") or "").strip()
|
|
||||||
output = {
|
|
||||||
"_next": "research_each_question",
|
|
||||||
"research_attempts": 0,
|
|
||||||
"research_feedback": (
|
|
||||||
"The user reviewed the report and asked for changes. Treat "
|
|
||||||
"this as the top priority for the next pass:\n\n" + feedback
|
|
||||||
),
|
|
||||||
}
|
|
||||||
print(json.dumps(output))
|
|
||||||
|
|
||||||
|
|
||||||
if __name__ == "__main__":
|
|
||||||
main()
|
|
||||||
@@ -1,35 +0,0 @@
|
|||||||
#!/usr/bin/env python3
|
|
||||||
"""Entry router for deep-research.
|
|
||||||
|
|
||||||
Reads the caller's prompt from state. If it contains a usable research
|
|
||||||
topic, stores it as `topic` and falls through to the static `next`
|
|
||||||
(plan). If the prompt is empty, routes to `ask_topic` so the user can
|
|
||||||
supply one interactively.
|
|
||||||
|
|
||||||
Routing (`_next`):
|
|
||||||
- prompt present -> (no _next; static next: plan)
|
|
||||||
- prompt empty -> ask_topic
|
|
||||||
"""
|
|
||||||
import json
|
|
||||||
import os
|
|
||||||
|
|
||||||
|
|
||||||
def load_state():
|
|
||||||
path = os.environ.get("GRAPH_STATE_FILE")
|
|
||||||
if path:
|
|
||||||
with open(path) as f:
|
|
||||||
return json.load(f)
|
|
||||||
return json.loads(os.environ.get("GRAPH_STATE", "{}"))
|
|
||||||
|
|
||||||
|
|
||||||
def main():
|
|
||||||
state = load_state()
|
|
||||||
prompt = (state.get("initial_prompt") or "").strip()
|
|
||||||
if prompt:
|
|
||||||
print(json.dumps({"topic": prompt}))
|
|
||||||
else:
|
|
||||||
print(json.dumps({"_next": "ask_topic"}))
|
|
||||||
|
|
||||||
|
|
||||||
if __name__ == "__main__":
|
|
||||||
main()
|
|
||||||
@@ -1,76 +0,0 @@
|
|||||||
#!/usr/bin/env python3
|
|
||||||
"""Reflexion gate for deep-research.
|
|
||||||
|
|
||||||
Runs after `critique` has reviewed the current research findings. If the
|
|
||||||
critique's verdict is REVISE and the reflexion budget is not spent,
|
|
||||||
loops back to `research` with the critique attached as
|
|
||||||
`research_feedback`, so the retry is informed rather than a blind
|
|
||||||
re-run. Otherwise it proceeds to `synthesize`.
|
|
||||||
|
|
||||||
Routing (`_next`):
|
|
||||||
- verdict PASS -> synthesize
|
|
||||||
- verdict REVISE, budget remaining -> research_each_question (+ research_feedback)
|
|
||||||
- verdict REVISE, budget spent -> synthesize
|
|
||||||
|
|
||||||
Reflexion is a best-effort quality booster, not a hard gate: once the
|
|
||||||
budget is spent the workflow proceeds anyway, and the human approval
|
|
||||||
step is the final backstop.
|
|
||||||
"""
|
|
||||||
import json
|
|
||||||
import os
|
|
||||||
import re
|
|
||||||
|
|
||||||
# Automated revision passes allowed. `research` runs at most
|
|
||||||
# MAX_REFLEXION_REVISIONS + 1 times per user pass. Bump to allow more.
|
|
||||||
MAX_REFLEXION_REVISIONS = 2
|
|
||||||
|
|
||||||
|
|
||||||
def load_state():
|
|
||||||
path = os.environ.get("GRAPH_STATE_FILE")
|
|
||||||
if path:
|
|
||||||
with open(path) as f:
|
|
||||||
return json.load(f)
|
|
||||||
return json.loads(os.environ.get("GRAPH_STATE", "{}"))
|
|
||||||
|
|
||||||
|
|
||||||
def as_int(value, default=0):
|
|
||||||
try:
|
|
||||||
return int(value)
|
|
||||||
except (TypeError, ValueError):
|
|
||||||
return default
|
|
||||||
|
|
||||||
|
|
||||||
def parse_verdict(critique):
|
|
||||||
"""Pull PASS/REVISE from the critique's `VERDICT:` line. Defaults to
|
|
||||||
PASS when no verdict line is found, so a malformed critique lets the
|
|
||||||
workflow proceed instead of burning the whole revision budget."""
|
|
||||||
match = re.search(r"VERDICT:\s*([A-Za-z]+)", critique, re.IGNORECASE)
|
|
||||||
if not match:
|
|
||||||
return "PASS"
|
|
||||||
return match.group(1).upper()
|
|
||||||
|
|
||||||
|
|
||||||
def main():
|
|
||||||
state = load_state()
|
|
||||||
critique = state.get("critique") or ""
|
|
||||||
verdict = parse_verdict(critique)
|
|
||||||
attempts = as_int(state.get("research_attempts"))
|
|
||||||
|
|
||||||
if verdict == "REVISE" and attempts < MAX_REFLEXION_REVISIONS:
|
|
||||||
feedback = (
|
|
||||||
"A reviewer judged the previous research pass incomplete. "
|
|
||||||
"Address every point in the critique below:\n\n" + critique
|
|
||||||
)
|
|
||||||
output = {
|
|
||||||
"_next": "research_each_question",
|
|
||||||
"research_attempts": attempts + 1,
|
|
||||||
"research_feedback": feedback,
|
|
||||||
}
|
|
||||||
else:
|
|
||||||
output = {"_next": "synthesize"}
|
|
||||||
|
|
||||||
print(json.dumps(output))
|
|
||||||
|
|
||||||
|
|
||||||
if __name__ == "__main__":
|
|
||||||
main()
|
|
||||||
@@ -1,69 +0,0 @@
|
|||||||
#!/usr/bin/env python3
|
|
||||||
"""Check that the sources cited in the research report are reachable.
|
|
||||||
|
|
||||||
Scans the final report for URLs and DOIs, probes each with a HEAD
|
|
||||||
request, and writes a `source_check` summary into state so the human
|
|
||||||
reviewer sees broken citations at the approval step.
|
|
||||||
|
|
||||||
Times out per request so a slow source cannot stall the graph.
|
|
||||||
"""
|
|
||||||
import json
|
|
||||||
import os
|
|
||||||
import re
|
|
||||||
import urllib.error
|
|
||||||
import urllib.request
|
|
||||||
|
|
||||||
DOI_RE = re.compile(r"\b(10\.\d{4,9}/[-._;()/:A-Z0-9]+)", re.IGNORECASE)
|
|
||||||
URL_RE = re.compile(r"https?://[^\s)\]\}\"'>]+")
|
|
||||||
|
|
||||||
|
|
||||||
def load_state():
|
|
||||||
path = os.environ.get("GRAPH_STATE_FILE")
|
|
||||||
if path:
|
|
||||||
with open(path) as f:
|
|
||||||
return json.load(f)
|
|
||||||
return json.loads(os.environ.get("GRAPH_STATE", "{}"))
|
|
||||||
|
|
||||||
|
|
||||||
def reachable(url, timeout=5.0):
|
|
||||||
req = urllib.request.Request(url, method="HEAD")
|
|
||||||
try:
|
|
||||||
with urllib.request.urlopen(req, timeout=timeout) as resp:
|
|
||||||
return 200 <= resp.status < 400
|
|
||||||
except urllib.error.HTTPError as e:
|
|
||||||
return 200 <= e.code < 400
|
|
||||||
except Exception:
|
|
||||||
return False
|
|
||||||
|
|
||||||
|
|
||||||
def main():
|
|
||||||
state = load_state()
|
|
||||||
report = state.get("report") or ""
|
|
||||||
|
|
||||||
urls = sorted({u.rstrip(".,;)") for u in URL_RE.findall(report)})
|
|
||||||
dois = sorted(set(DOI_RE.findall(report)))
|
|
||||||
|
|
||||||
results = []
|
|
||||||
for url in urls:
|
|
||||||
ok = reachable(url)
|
|
||||||
results.append(f" {'OK' if ok else 'UNREACHABLE'} {url}")
|
|
||||||
for doi in dois:
|
|
||||||
url = f"https://doi.org/{doi}"
|
|
||||||
if url in urls:
|
|
||||||
continue
|
|
||||||
ok = reachable(url)
|
|
||||||
results.append(f" {'OK' if ok else 'UNREACHABLE'} DOI {doi} ({url})")
|
|
||||||
|
|
||||||
if not results:
|
|
||||||
summary = "No web sources were cited in the report."
|
|
||||||
else:
|
|
||||||
summary = (
|
|
||||||
f"Source reachability ({len(results)} checked):\n"
|
|
||||||
+ "\n".join(results)
|
|
||||||
)
|
|
||||||
|
|
||||||
print(json.dumps({"source_check": summary}))
|
|
||||||
|
|
||||||
|
|
||||||
if __name__ == "__main__":
|
|
||||||
main()
|
|
||||||
@@ -1,39 +0,0 @@
|
|||||||
#!/usr/bin/env bash
|
|
||||||
|
|
||||||
set -e
|
|
||||||
|
|
||||||
# @env LLM_OUTPUT=/dev/stdout The output path
|
|
||||||
|
|
||||||
# @cmd Classify the credibility tier of a web source from its URL.
|
|
||||||
# A deterministic check based on the host and top-level domain. Use it
|
|
||||||
# to weigh how much trust to place in a source before relying on it.
|
|
||||||
# @option --url! The full source URL to classify
|
|
||||||
classify_source() {
|
|
||||||
# shellcheck disable=SC2154
|
|
||||||
local url="$argc_url"
|
|
||||||
local host="${url#*://}"
|
|
||||||
host="${host%%/*}"
|
|
||||||
host="${host##*@}"
|
|
||||||
host="${host%%:*}"
|
|
||||||
host="$(printf '%s' "$host" | tr '[:upper:]' '[:lower:]')"
|
|
||||||
|
|
||||||
local tier
|
|
||||||
case "$host" in
|
|
||||||
'')
|
|
||||||
tier="UNKNOWN - no host could be parsed from the URL" ;;
|
|
||||||
*.gov | *.gov.* | *.mil)
|
|
||||||
tier="HIGH - government source" ;;
|
|
||||||
*.edu | *.edu.* | *.ac.*)
|
|
||||||
tier="HIGH - academic institution" ;;
|
|
||||||
arxiv.org | *.arxiv.org | biorxiv.org | *.biorxiv.org | medrxiv.org | *.medrxiv.org | ssrn.com | *.ssrn.com)
|
|
||||||
tier="PREPRINT - not yet peer reviewed, corroborate before citing" ;;
|
|
||||||
wikipedia.org | *.wikipedia.org)
|
|
||||||
tier="TERTIARY - encyclopedia, good for orientation not citation" ;;
|
|
||||||
*.org | *.org.*)
|
|
||||||
tier="MEDIUM - organization site, check for institutional bias" ;;
|
|
||||||
*)
|
|
||||||
tier="UNVERIFIED - general web source, corroborate before citing" ;;
|
|
||||||
esac
|
|
||||||
|
|
||||||
printf '%s: %s\n' "${host:-<none>}" "$tier" >> "$LLM_OUTPUT"
|
|
||||||
}
|
|
||||||
@@ -0,0 +1,14 @@
|
|||||||
|
# Jira AI Agent
|
||||||
|
|
||||||
|
## Overview
|
||||||
|
|
||||||
|
The Jira AI Agent is designed to assist with managing tasks within Jira projects, providing capabilities such as
|
||||||
|
creating, searching, updating, assigning, linking, and commenting on issues. Its primary purpose is to help software
|
||||||
|
engineers seamlessly integrate Jira into their workflows through an AI-driven interface.
|
||||||
|
|
||||||
|
## Configuration
|
||||||
|
This agent uses the official [Atlassian MCP Server](https://github.com/atlassian/atlassian-mcp-server). To use it,
|
||||||
|
ensure you have Node.js v18+ installed to run the local MCP proxy (`mcp-remote`).
|
||||||
|
|
||||||
|
The server uses OAuth 2.0 so it will automatically open your browser for you to sign in to your account. No manual
|
||||||
|
configuration is necessary!
|
||||||
@@ -0,0 +1,37 @@
|
|||||||
|
name: Jira Agent
|
||||||
|
description: An AI agent that can assist with Jira tasks such as creating issues, searching for issues, and updating issues.
|
||||||
|
version: 0.1.0
|
||||||
|
agent_session: temp
|
||||||
|
mcp_servers:
|
||||||
|
- atlassian
|
||||||
|
instructions: |
|
||||||
|
You are a AI agent designed to assist with managing Jira tasks and helping software engineers utilize and integrate
|
||||||
|
Jira into their workflows. You can create, search, update, assign, link, and comment on issues in Jira.
|
||||||
|
|
||||||
|
## Create Issue (MANDATORY when creating a issue)
|
||||||
|
When a user prompts you to create a Jira issue:
|
||||||
|
1. Prompt the user for what Jira project they want the ticket created in
|
||||||
|
2. If the ticket type requires a parent issue:
|
||||||
|
a. Query Jira for potentially relevant parents
|
||||||
|
b. Prompt user for which parent to use, displaying the suggested list of parent issues
|
||||||
|
3. Create the issue with the following format:
|
||||||
|
```markdown
|
||||||
|
**Description:**
|
||||||
|
This section gives context and details about the issue.
|
||||||
|
**User Acceptance Criteria:**
|
||||||
|
# This section provides bullet points that function like a checklist of all the things that must be completed in
|
||||||
|
# order for the issue to be considered done.
|
||||||
|
* Example criteria one
|
||||||
|
* Example criteria two
|
||||||
|
```
|
||||||
|
4. Ask the user if the issue should be assigned to them
|
||||||
|
a. If yes, then assign the user to the newly created issue
|
||||||
|
|
||||||
|
|
||||||
|
Available tools:
|
||||||
|
{{__tools__}}
|
||||||
|
conversation_starters:
|
||||||
|
- What are the latest issues in my Jira project?
|
||||||
|
- Can you create a new Jira issue for me?
|
||||||
|
- What are my open Jira issues?
|
||||||
|
- Can you search for issues with the label "bug" in my Jira project?
|
||||||
@@ -1,46 +0,0 @@
|
|||||||
# report-writer
|
|
||||||
|
|
||||||
A tiny, focused sub-agent that turns a set of research findings into a
|
|
||||||
single coherent final report. Reads only what it is given — does not
|
|
||||||
do independent research, does not access the web, does not invent
|
|
||||||
facts. It exists as a focused tool for orchestrating agents to
|
|
||||||
delegate the writing phase to.
|
|
||||||
|
|
||||||
## Why a separate agent?
|
|
||||||
|
|
||||||
This is an example of the **agent-as-tool** pattern in graph agents.
|
|
||||||
The `deep-research` graph agent's `synthesize` node is an `agent` node
|
|
||||||
that spawns this one (see `assets/agents/deep-research/graph.yaml`).
|
|
||||||
Separating the role has two practical benefits:
|
|
||||||
|
|
||||||
- The orchestrating agent can use a cheap model (or a high-temperature
|
|
||||||
exploratory one) for the research phase, while letting the writing
|
|
||||||
phase use a different (typically lower-temperature, possibly larger)
|
|
||||||
model dedicated to coherent prose.
|
|
||||||
- The writing prompt is owned by this agent's `config.yaml` rather
|
|
||||||
than buried inside another agent's graph. You can polish it
|
|
||||||
independently without touching the research flow.
|
|
||||||
|
|
||||||
## Standalone use
|
|
||||||
|
|
||||||
You can also use this agent directly if you have a set of findings you
|
|
||||||
want polished:
|
|
||||||
|
|
||||||
```sh
|
|
||||||
loki -a report-writer "Topic: X. Findings: <paste findings here>"
|
|
||||||
```
|
|
||||||
|
|
||||||
It will produce a single Markdown report following the rules in its
|
|
||||||
system prompt: executive summary at the top, grouped sections by
|
|
||||||
related sub-questions, every inline citation preserved verbatim, and a
|
|
||||||
final "Open questions / disagreements" section.
|
|
||||||
|
|
||||||
## What it will NOT do
|
|
||||||
|
|
||||||
- Search the web, fetch URLs, query an MCP server, or use any tool.
|
|
||||||
It has no tools configured.
|
|
||||||
- Invent facts beyond what is in the findings you give it.
|
|
||||||
- Strip or rewrite citations.
|
|
||||||
|
|
||||||
These constraints are the point of the agent existing: a writer that
|
|
||||||
the orchestrator can trust to stay in its lane.
|
|
||||||
@@ -1,34 +0,0 @@
|
|||||||
name: report-writer
|
|
||||||
description: Polishes research findings into a clear, citation-preserving final report
|
|
||||||
version: 1.0.0
|
|
||||||
temperature: 0.2
|
|
||||||
|
|
||||||
instructions: |
|
|
||||||
You are a technical writer. You will be given:
|
|
||||||
- a research topic
|
|
||||||
- a set of findings, organized per sub-question, with inline
|
|
||||||
citations next to each claim
|
|
||||||
- a source-credibility assessment of the cited sources
|
|
||||||
|
|
||||||
Your job is to produce a single, well-organized final report:
|
|
||||||
|
|
||||||
Rules:
|
|
||||||
- Use ONLY the findings provided. Do not introduce facts from
|
|
||||||
your own memory. Do not speculate beyond what the findings
|
|
||||||
support.
|
|
||||||
- Preserve every inline citation. If a sentence in the findings
|
|
||||||
had a URL or DOI, the equivalent sentence in your report must
|
|
||||||
keep the same citation.
|
|
||||||
- Lead with a 2-3 sentence executive summary at the top.
|
|
||||||
- Organize the body so that related sub-questions are grouped,
|
|
||||||
not strictly one section per question. The findings are raw
|
|
||||||
material; the report should read as a single coherent answer
|
|
||||||
to the original topic.
|
|
||||||
- End with a short "Open questions / disagreements" section
|
|
||||||
naming anything the findings flagged as unresolved or
|
|
||||||
contested.
|
|
||||||
|
|
||||||
Output plain Markdown. No metadata, no JSON wrapper.
|
|
||||||
|
|
||||||
conversation_starters:
|
|
||||||
- "Polish these findings into a cited report"
|
|
||||||
@@ -18,15 +18,16 @@ Sisyphus acts as the primary entry point, capable of handling complex tasks by c
|
|||||||
- 🛠️ **Tool Integration**: Seamlessly uses system tools for building, testing, and file manipulation.
|
- 🛠️ **Tool Integration**: Seamlessly uses system tools for building, testing, and file manipulation.
|
||||||
|
|
||||||
## Pro-Tip: Use an IDE MCP Server for Improved Performance
|
## Pro-Tip: Use an IDE MCP Server for Improved Performance
|
||||||
Many modern IDEs (JetBrains, VS Code, Cursor, Zed, etc.) expose MCP servers that let LLMs use IDE tools directly. Using
|
Many modern IDEs now include MCP servers that let LLMs perform operations within the IDE itself and use IDE tools. Using
|
||||||
one dramatically improves the performance of coding agents. If you have one, add it to your loki config (see the
|
an IDE's MCP server dramatically improves the performance of coding agents. So if you have an IDE, try adding that MCP
|
||||||
[MCP Server docs](../../../docs/function-calling/MCP-SERVERS.md)) and reference it in this agent's `mcp_servers:` list:
|
server to your config (see the [MCP Server docs](../../../docs/function-calling/MCP-SERVERS.md) to see how to configure
|
||||||
|
them), and modify the agent definition to look like this:
|
||||||
|
|
||||||
```yaml
|
```yaml
|
||||||
# ...
|
# ...
|
||||||
|
|
||||||
mcp_servers:
|
mcp_servers:
|
||||||
- your-ide-mcp-server
|
- jetbrains
|
||||||
|
|
||||||
global_tools:
|
global_tools:
|
||||||
- fs_read.sh
|
- fs_read.sh
|
||||||
|
|||||||
@@ -119,21 +119,20 @@ instructions: |
|
|||||||
1. todo__init --goal "Add user profiles API endpoint"
|
1. todo__init --goal "Add user profiles API endpoint"
|
||||||
2. todo__add --task "Explore existing API patterns"
|
2. todo__add --task "Explore existing API patterns"
|
||||||
3. todo__add --task "Implement profile endpoint"
|
3. todo__add --task "Implement profile endpoint"
|
||||||
4. agent__spawn --agent explore --prompt "Find existing API endpoint patterns, route structures, and controller conventions. Include code snippets."
|
4. todo__add --task "Verify with build/test"
|
||||||
5. agent__spawn --agent explore --prompt "Find existing data models and database query patterns. Include code snippets."
|
5. agent__spawn --agent explore --prompt "Find existing API endpoint patterns, route structures, and controller conventions. Include code snippets."
|
||||||
6. agent__collect --id <id1>
|
6. agent__spawn --agent explore --prompt "Find existing data models and database query patterns. Include code snippets."
|
||||||
7. agent__collect --id <id2>
|
7. agent__collect --id <id1>
|
||||||
8. todo__done --id 1
|
8. agent__collect --id <id2>
|
||||||
9. agent__spawn --agent coder --prompt "<structured prompt using Coder Delegation Format above, including code snippets from explore results>"
|
9. todo__done --id 1
|
||||||
10. agent__collect --id <coder_id>
|
10. agent__spawn --agent coder --prompt "<structured prompt using Coder Delegation Format above, including code snippets from explore results>"
|
||||||
11. todo__done --id 2
|
11. agent__collect --id <coder_id>
|
||||||
|
12. todo__done --id 2
|
||||||
|
13. run_build
|
||||||
|
14. run_tests
|
||||||
|
15. todo__done --id 3
|
||||||
```
|
```
|
||||||
|
|
||||||
Note: the `coder` agent is a graph agent that runs verification (build +
|
|
||||||
tests) and a bounded fix-loop internally. You do NOT need to spawn a
|
|
||||||
separate build/test step. A `CODER_COMPLETE` outcome means build and
|
|
||||||
tests already passed.
|
|
||||||
|
|
||||||
### Example 2: Architecture/design question (explore + oracle in parallel)
|
### Example 2: Architecture/design question (explore + oracle in parallel)
|
||||||
|
|
||||||
User: "How should I structure the authentication for this app?"
|
User: "How should I structure the authentication for this app?"
|
||||||
@@ -173,22 +172,6 @@ instructions: |
|
|||||||
10. **Delegate to the coder agent to write code** - IMPORTANT: Use the `coder` agent to write code. Do not try to write code yourself except for trivial changes
|
10. **Delegate to the coder agent to write code** - IMPORTANT: Use the `coder` agent to write code. Do not try to write code yourself except for trivial changes
|
||||||
11. **Always output a summary of changes when finished** - Make it clear to user's that you've completed your tasks
|
11. **Always output a summary of changes when finished** - Make it clear to user's that you've completed your tasks
|
||||||
|
|
||||||
## Coder Outcomes
|
|
||||||
|
|
||||||
The `coder` agent is a graph agent that runs the implement -> verify_build
|
|
||||||
-> verify_tests -> fix_loop pipeline internally. It always returns one of
|
|
||||||
three sentinel outcomes:
|
|
||||||
|
|
||||||
- `CODER_COMPLETE` - implementation succeeded with build + tests green.
|
|
||||||
Continue with any follow-up todos.
|
|
||||||
- `CODER_REJECTED` - user rejected the plan at the approval gate (only
|
|
||||||
triggered for high-complexity plans). Do NOT re-spawn coder blindly;
|
|
||||||
ask the user what to change first.
|
|
||||||
- `CODER_FAILED` - the fix-loop exhausted its budget without producing
|
|
||||||
green build/tests. The failure output includes the last build and tests
|
|
||||||
output. Surface this to the user; consider spawning `oracle` for
|
|
||||||
diagnosis if the failure is unclear.
|
|
||||||
|
|
||||||
## When to Do It Yourself
|
## When to Do It Yourself
|
||||||
|
|
||||||
- Simple command execution
|
- Simple command execution
|
||||||
|
|||||||
@@ -73,11 +73,11 @@ def to_args:
|
|||||||
to_entries | .[] |
|
to_entries | .[] |
|
||||||
(.key | split("_") | join("-")) as $key |
|
(.key | split("_") | join("-")) as $key |
|
||||||
if .value | type == "array" then
|
if .value | type == "array" then
|
||||||
.value | .[] | "--\($key)=\(. | escape_shell_word)"
|
.value | .[] | "--\($key) \(. | escape_shell_word)"
|
||||||
elif .value | type == "boolean" then
|
elif .value | type == "boolean" then
|
||||||
if .value then "--\($key)" else "" end
|
if .value then "--\($key)" else "" end
|
||||||
else
|
else
|
||||||
"--\($key)=\(.value | escape_shell_word)"
|
"--\($key) \(.value | escape_shell_word)"
|
||||||
end;
|
end;
|
||||||
[ to_args ] | join(" ")
|
[ to_args ] | join(" ")
|
||||||
EOF
|
EOF
|
||||||
|
|||||||
@@ -70,11 +70,11 @@ def to_args:
|
|||||||
to_entries | .[] |
|
to_entries | .[] |
|
||||||
(.key | split("_") | join("-")) as $key |
|
(.key | split("_") | join("-")) as $key |
|
||||||
if .value | type == "array" then
|
if .value | type == "array" then
|
||||||
.value | .[] | "--\($key)=\(. | escape_shell_word)"
|
.value | .[] | "--\($key) \(. | escape_shell_word)"
|
||||||
elif .value | type == "boolean" then
|
elif .value | type == "boolean" then
|
||||||
if .value then "--\($key)" else "" end
|
if .value then "--\($key)" else "" end
|
||||||
else
|
else
|
||||||
"--\($key)=\(.value | escape_shell_word)"
|
"--\($key) \(.value | escape_shell_word)"
|
||||||
end;
|
end;
|
||||||
[ to_args ] | join(" ")
|
[ to_args ] | join(" ")
|
||||||
EOF
|
EOF
|
||||||
|
|||||||
@@ -1,8 +0,0 @@
|
|||||||
---
|
|
||||||
enabled_mcp_servers: atlassian
|
|
||||||
---
|
|
||||||
You are the librarian for the company's Confluence and Jira knowledge bases. Your job is to help users find and retrieve
|
|
||||||
information from these platforms. Use all tools at your disposal to answer user queries.
|
|
||||||
|
|
||||||
Available Tools:
|
|
||||||
{{__tools__}}
|
|
||||||
@@ -17,18 +17,16 @@ agent_session: null # Set a session to use when starting the agent.
|
|||||||
name: <agent-name> # Name of the agent, used in the UI and logs
|
name: <agent-name> # Name of the agent, used in the UI and logs
|
||||||
description: <description> # Description of the agent, used in the UI
|
description: <description> # Description of the agent, used in the UI
|
||||||
version: 1 # Version of the agent
|
version: 1 # Version of the agent
|
||||||
# Auto-Continue (Todo System)
|
# Todo System & Auto-Continuation
|
||||||
# The auto-continue system provides built-in task tracking for improved reliability.
|
# These settings help smaller models handle multi-step tasks more reliably.
|
||||||
# When enabled, the model can create todo lists and the system will automatically
|
# See docs/TODO-SYSTEM.md for detailed documentation.
|
||||||
# prompt it to continue when incomplete tasks remain.
|
|
||||||
# See the [Todo System documentation](https://github.com/Dark-Alex-17/loki/wiki/TODO-System) for more information
|
|
||||||
auto_continue: false # Enable automatic continuation when incomplete todos remain
|
auto_continue: false # Enable automatic continuation when incomplete todos remain
|
||||||
max_auto_continues: 10 # Maximum number of automatic continuations before stopping
|
max_auto_continues: 10 # Maximum number of automatic continuations before stopping
|
||||||
inject_todo_instructions: true # Inject the default todo tool usage instructions into the agent's system prompt
|
inject_todo_instructions: true # Inject the default todo tool usage instructions into the agent's system prompt
|
||||||
continuation_prompt: null # Custom prompt used when auto-continuing (optional; uses default if null)
|
continuation_prompt: null # Custom prompt used when auto-continuing (optional; uses default if null)
|
||||||
# Sub-Agent Spawning System
|
# Sub-Agent Spawning System
|
||||||
# Enable this agent to spawn and manage child agents in parallel.
|
# Enable this agent to spawn and manage child agents in parallel.
|
||||||
# See https://github.com/Dark-Alex-17/loki/wiki/Agents for detailed documentation.
|
# See docs/AGENTS.md for detailed documentation.
|
||||||
can_spawn_agents: false # Enable the agent to spawn child agents
|
can_spawn_agents: false # Enable the agent to spawn child agents
|
||||||
max_concurrent_agents: 4 # Maximum number of agents that can run simultaneously
|
max_concurrent_agents: 4 # Maximum number of agents that can run simultaneously
|
||||||
max_agent_depth: 3 # Maximum nesting depth for sub-agents (prevents runaway spawning)
|
max_agent_depth: 3 # Maximum nesting depth for sub-agents (prevents runaway spawning)
|
||||||
|
|||||||
+8
-17
@@ -27,18 +27,18 @@ sync_models_url: > # URL to sync model changes from
|
|||||||
https://raw.githubusercontent.com/Dark-Alex-17/loki/refs/heads/main/models.yaml
|
https://raw.githubusercontent.com/Dark-Alex-17/loki/refs/heads/main/models.yaml
|
||||||
|
|
||||||
# ---- REPL Prompt ----
|
# ---- REPL Prompt ----
|
||||||
# Custom REPL left/right prompts; see the [REPL Prompt Documentation](https://github.com/Dark-Alex-17/loki/wiki/REPL-Prompt) for more information
|
# Custom REPL left/right prompts; see the [REPL Prompt Documentation](./docs/REPL-PROMPT.md) for more information
|
||||||
left_prompt:
|
left_prompt:
|
||||||
'{color.red}{model}){color.green}{?session {?agent {agent}>}{session}{?role /}}{!session {?agent {agent}>}}{role}{?rag @{rag}}{color.cyan}{?session )}{!session >}{color.reset} '
|
'{color.red}{model}){color.green}{?session {?agent {agent}>}{session}{?role /}}{!session {?agent {agent}>}}{role}{?rag @{rag}}{color.cyan}{?session )}{!session >}{color.reset} '
|
||||||
right_prompt:
|
right_prompt:
|
||||||
'{color.purple}{?session {?consume_tokens {consume_tokens}({consume_percent}%)}{!consume_tokens {consume_tokens}}}{color.reset}'
|
'{color.purple}{?session {?consume_tokens {consume_tokens}({consume_percent}%)}{!consume_tokens {consume_tokens}}}{color.reset}'
|
||||||
|
|
||||||
# ---- Vault ----
|
# ---- Vault ----
|
||||||
# See the [Vault documentation](https://github.com/Dark-Alex-17/loki/wiki/Vault) for more information on the Loki vault
|
# See the [Vault documentation](./docs/VAULT.md) for more information on the Loki vault
|
||||||
vault_password_file: null # Path to a file containing the password for the Loki vault (cannot be a secret template)
|
vault_password_file: null # Path to a file containing the password for the Loki vault (cannot be a secret template)
|
||||||
|
|
||||||
# ---- Function Calling ----
|
# ---- Function Calling ----
|
||||||
# See the [Tools documentation](https://github.com/Dark-Alex-17/loki/wiki/Tools) for more details
|
# See the [Tools documentation](./docs/function-calling/TOOLS.md) for more details
|
||||||
function_calling: true # Enables or disables function calling (Globally).
|
function_calling: true # Enables or disables function calling (Globally).
|
||||||
mapping_tools: # Alias for a tool or toolset
|
mapping_tools: # Alias for a tool or toolset
|
||||||
fs: 'fs_cat,fs_ls,fs_mkdir,fs_rm,fs_write,fs_read,fs_glob,fs_grep'
|
fs: 'fs_cat,fs_ls,fs_mkdir,fs_rm,fs_write,fs_read,fs_glob,fs_grep'
|
||||||
@@ -64,6 +64,7 @@ visible_tools: # Which tools are visible to be compiled (and a
|
|||||||
# - get_current_weather.py
|
# - get_current_weather.py
|
||||||
# - get_current_weather.ts
|
# - get_current_weather.ts
|
||||||
- get_current_weather.sh
|
- get_current_weather.sh
|
||||||
|
- query_jira_issues.sh
|
||||||
# - search_arxiv.sh
|
# - search_arxiv.sh
|
||||||
# - search_wikipedia.sh
|
# - search_wikipedia.sh
|
||||||
# - search_wolframalpha.sh
|
# - search_wolframalpha.sh
|
||||||
@@ -74,24 +75,14 @@ visible_tools: # Which tools are visible to be compiled (and a
|
|||||||
# - web_search_tavily.sh
|
# - web_search_tavily.sh
|
||||||
|
|
||||||
# ---- MCP Servers ----
|
# ---- MCP Servers ----
|
||||||
# See the [MCP Servers documentation](https://github.com/Dark-Alex-17/loki/wiki/MCP-Servers) for more details
|
# See the [MCP Servers documentation](./docs/MCP-SERVERS.md) for more details
|
||||||
mcp_server_support: true # Enables or disables MCP servers (globally).
|
mcp_server_support: true # Enables or disables MCP servers (globally).
|
||||||
mapping_mcp_servers: # Alias for an MCP server or set of servers
|
mapping_mcp_servers: # Alias for an MCP server or set of servers
|
||||||
git: github,gitmcp
|
git: github,gitmcp
|
||||||
enabled_mcp_servers: null # Which MCP servers to enable by default (e.g. 'github,slack,ddg-search')
|
enabled_mcp_servers: null # Which MCP servers to enable by default (e.g. 'github,slack,ddg-search')
|
||||||
|
|
||||||
# ---- Auto-Continue (Todo System) ----
|
|
||||||
# The auto-continue system provides built-in task tracking for improved reliability.
|
|
||||||
# When enabled, the model can create todo lists and the system will automatically
|
|
||||||
# prompt it to continue when incomplete tasks remain.
|
|
||||||
# See the [Todo System documentation](https://github.com/Dark-Alex-17/loki/wiki/TODO-System) for more information
|
|
||||||
auto_continue: false # Enable automatic continuation when incomplete todos remain (default: false)
|
|
||||||
max_auto_continues: 10 # Maximum number of automatic continuations before stopping (default: 10)
|
|
||||||
inject_todo_instructions: true # Inject default todo usage instructions into the system prompt (default: true)
|
|
||||||
continuation_prompt: null # Custom prompt used when auto-continuing. If null, uses built-in default
|
|
||||||
|
|
||||||
# ---- Session ----
|
# ---- Session ----
|
||||||
# See the [Session documentation](https://github.com/Dark-Alex-17/loki/wiki/Sessions) for more information
|
# See the [Session documentation](./docs/SESSIONS.md) for more information
|
||||||
save_session: null # Controls the persistence of the session. If true, auto save; if false, don't auto-save save; if null, ask the user what to do
|
save_session: null # Controls the persistence of the session. If true, auto save; if false, don't auto-save save; if null, ask the user what to do
|
||||||
compression_threshold: 4000 # Compress the session when the token count reaches or exceeds this threshold
|
compression_threshold: 4000 # Compress the session when the token count reaches or exceeds this threshold
|
||||||
summarization_prompt: > # The text prompt used for creating a concise summary of session message
|
summarization_prompt: > # The text prompt used for creating a concise summary of session message
|
||||||
@@ -100,7 +91,7 @@ summary_context_prompt: > # The text prompt used for including the summar
|
|||||||
'This is a summary of the chat history as a recap: '
|
'This is a summary of the chat history as a recap: '
|
||||||
|
|
||||||
# ---- RAG ----
|
# ---- RAG ----
|
||||||
# See the [RAG Docs](https://github.com/Dark-Alex-17/loki/wiki/RAG) for more details.
|
# See the [RAG Docs](./docs/RAG.md) for more details.
|
||||||
rag_embedding_model: null # Specifies the embedding model used for context retrieval
|
rag_embedding_model: null # Specifies the embedding model used for context retrieval
|
||||||
rag_reranker_model: null # Specifies the reranker model used for sorting retrieved documents; Loki uses Reciprocal Rank Fusion by default
|
rag_reranker_model: null # Specifies the reranker model used for sorting retrieved documents; Loki uses Reciprocal Rank Fusion by default
|
||||||
rag_top_k: 5 # Specifies the number of documents to retrieve for answering queries
|
rag_top_k: 5 # Specifies the number of documents to retrieve for answering queries
|
||||||
@@ -146,7 +137,7 @@ document_loaders:
|
|||||||
sh -c "yek $1 --json | jq 'map({ path: .filename, contents: .content })'"
|
sh -c "yek $1 --json | jq 'map({ path: .filename, contents: .content })'"
|
||||||
|
|
||||||
# ---- Clients ----
|
# ---- Clients ----
|
||||||
# See the [Clients documentation](https://github.com/Dark-Alex-17/loki/wiki/Clients) for more details
|
# See the [Clients documentation](./docs/clients/CLIENTS.md) for more details
|
||||||
clients:
|
clients:
|
||||||
# All clients have the following configuration:
|
# All clients have the following configuration:
|
||||||
# - type: xxxx
|
# - type: xxxx
|
||||||
|
|||||||
+1
-14
@@ -1,9 +1,5 @@
|
|||||||
---
|
---
|
||||||
############################################
|
# Everything in this section is optional
|
||||||
## Everything in this section is optional ##
|
|
||||||
############################################
|
|
||||||
|
|
||||||
# Role Configuration
|
|
||||||
name: <role-name> # The name of the role
|
name: <role-name> # The name of the role
|
||||||
model: openai:gpt-4o # The model to use for this role
|
model: openai:gpt-4o # The model to use for this role
|
||||||
temperature: 0.2 # The temperature to use for this role when querying the model
|
temperature: 0.2 # The temperature to use for this role when querying the model
|
||||||
@@ -12,14 +8,5 @@ enabled_tools: fs_ls,fs_cat # A comma-separated list of tools to enabl
|
|||||||
enabled_mcp_servers: github,gitmcp # A comma-separated list of MCP servers to enable for this role
|
enabled_mcp_servers: github,gitmcp # A comma-separated list of MCP servers to enable for this role
|
||||||
prompt: null # A custom prompt to use for this role that will immediately query
|
prompt: null # A custom prompt to use for this role that will immediately query
|
||||||
# the model for output instead of using the instructions below
|
# the model for output instead of using the instructions below
|
||||||
# Auto-Continue (Todo System)
|
|
||||||
# The auto-continue system provides built-in task tracking for improved reliability.
|
|
||||||
# When enabled, the model can create todo lists and the system will automatically
|
|
||||||
# prompt it to continue when incomplete tasks remain.
|
|
||||||
# See the [Todo System documentation](https://github.com/Dark-Alex-17/loki/wiki/TODO-System) for more information
|
|
||||||
auto_continue: false # Enable automatic continuation when incomplete todos remain (default: false)
|
|
||||||
max_auto_continues: 10 # Maximum number of automatic continuations before stopping (default: 10)
|
|
||||||
inject_todo_instructions: true # Inject default todo tool usage instructions into the system prompt (default: true)
|
|
||||||
continuation_prompt: null # Custom prompt used when auto-continuing. If null, uses built-in default
|
|
||||||
---
|
---
|
||||||
You are an expert at doing things. This is where you write the instructions for the role.
|
You are an expert at doing things. This is where you write the instructions for the role.
|
||||||
|
|||||||
@@ -1,427 +0,0 @@
|
|||||||
# Graph-based agent definition (full-featured reference)
|
|
||||||
# Location: <loki-config-dir>/agents/<agent-name>/graph.yaml
|
|
||||||
#
|
|
||||||
# A graph agent is defined by this file alone. An agent directory contains
|
|
||||||
# either a config.yaml (a normal LLM-loop agent) or a graph.yaml (a graph
|
|
||||||
# agent), never both. The presence of graph.yaml is what makes the agent
|
|
||||||
# a graph agent.
|
|
||||||
#
|
|
||||||
# This file is a reference: it documents every available field, themed
|
|
||||||
# around a deep web research workflow with parallel retrieval. It is not
|
|
||||||
# a runnable agent as-is. The `agent:`, `script:`, and `documents:` values
|
|
||||||
# point at things that would need to exist for a real agent. For a real,
|
|
||||||
# runnable deep-research graph agent, see assets/agents/deep-research/.
|
|
||||||
#
|
|
||||||
# Full documentation:
|
|
||||||
# https://github.com/Dark-Alex-17/loki/wiki/Graph-Agents
|
|
||||||
|
|
||||||
# ---------------------------------------------------------------------------
|
|
||||||
# Identity
|
|
||||||
# ---------------------------------------------------------------------------
|
|
||||||
name: deep-research-example # Agent name (should match the directory name)
|
|
||||||
description: | # Free-form prose describing the workflow
|
|
||||||
A reference workflow: triage a research request, retrieve local
|
|
||||||
context, branch on a script decision, run either a sub-agent or an
|
|
||||||
LLM research step, then gate the result behind human approval.
|
|
||||||
version: "1.0" # Graph schema version. Only "1.0" is accepted.
|
|
||||||
|
|
||||||
# ---------------------------------------------------------------------------
|
|
||||||
# Agent-level config (all optional)
|
|
||||||
# The same knobs a normal agent's config.yaml carries. In a graph agent they
|
|
||||||
# live here instead of in a config.yaml.
|
|
||||||
# ---------------------------------------------------------------------------
|
|
||||||
model: claude:claude-sonnet-4-6 # Default model for `llm` nodes that don't override it
|
|
||||||
temperature: 0.0 # Default sampling temperature for `llm` nodes
|
|
||||||
top_p: null # Default sampling top-p for `llm` nodes
|
|
||||||
|
|
||||||
global_tools: # Tool universe an `llm` node's `tools:` whitelist draws from
|
|
||||||
- web_search_loki.sh
|
|
||||||
- fetch_url_via_curl.sh
|
|
||||||
|
|
||||||
mcp_servers: # MCP servers an `llm` node may reference via `mcp:<server>`
|
|
||||||
- ddg-search
|
|
||||||
|
|
||||||
conversation_starters: # Suggested prompts surfaced in the UI
|
|
||||||
- "Research the current state of WebAssembly outside the browser"
|
|
||||||
|
|
||||||
# ---------------------------------------------------------------------------
|
|
||||||
# Agent variables (optional)
|
|
||||||
# Declared the same way as a normal agent's config.yaml `variables:` block.
|
|
||||||
# Each variable becomes available to:
|
|
||||||
# - LLM nodes via the template form `{{name}}` once seeded into state
|
|
||||||
# (see initial_state below).
|
|
||||||
# - Script nodes via the env var `LLM_AGENT_VAR_<UPPER_NAME>`.
|
|
||||||
# Values may be overridden at runtime with
|
|
||||||
# `loki -a <agent> --agent-variable <name> <value> "..."`.
|
|
||||||
# ---------------------------------------------------------------------------
|
|
||||||
variables:
|
|
||||||
- name: project_dir
|
|
||||||
description: |
|
|
||||||
Absolute path to the project directory.
|
|
||||||
default: "."
|
|
||||||
|
|
||||||
# ---------------------------------------------------------------------------
|
|
||||||
# Execution settings (all optional)
|
|
||||||
# ---------------------------------------------------------------------------
|
|
||||||
settings:
|
|
||||||
max_loop_iterations: 100 # Per-node visit cap. If one node id is entered more
|
|
||||||
# than this many times, execution aborts. Default 100.
|
|
||||||
timeout: 600 # Optional wall-clock cap (seconds) on the whole run,
|
|
||||||
# checked between node transitions.
|
|
||||||
log_state_snapshots: true # Log state before each node (debug/trace). Default true.
|
|
||||||
validate_before_run: true # Run the graph validator at startup. Default true.
|
|
||||||
max_concurrency: 4 # Cap on simultaneously running branches in any
|
|
||||||
# super-step (static fan-out OR a `map` node).
|
|
||||||
# Default 4. Per-`map` overrides this. See Parallel
|
|
||||||
# Execution below.
|
|
||||||
|
|
||||||
# ---------------------------------------------------------------------------
|
|
||||||
# Reducers (optional, required whenever two parallel branches write the same
|
|
||||||
# state key in the same super-step; otherwise the validator errors at load).
|
|
||||||
#
|
|
||||||
# A reducer says how two values for the same key get merged. Built-ins:
|
|
||||||
# append list += [value] (single value appended to a list)
|
|
||||||
# extend list += value (a list) (list-of-lists flattened by one level)
|
|
||||||
# concat "a\nb" (string join with newline separator)
|
|
||||||
# sum a + b (numeric add; ints stay ints)
|
|
||||||
# max max(a, b)
|
|
||||||
# min min(a, b)
|
|
||||||
# merge {**a, **b} (dict union, RHS wins on key collision)
|
|
||||||
# overwrite last-write-wins (explicit opt-in; B's value replaces A's)
|
|
||||||
#
|
|
||||||
# Keys not listed here have an implicit "single writer per super-step" rule:
|
|
||||||
# the validator rejects any graph where two parallel branches both write a
|
|
||||||
# key with no reducer.
|
|
||||||
# ---------------------------------------------------------------------------
|
|
||||||
reducers:
|
|
||||||
sources: append # The diamond below writes `sources` from both
|
|
||||||
# branches; append accumulates them into a list.
|
|
||||||
context: concat # Each branch contributes prose; concat joins them.
|
|
||||||
|
|
||||||
# ---------------------------------------------------------------------------
|
|
||||||
# Seed state (optional)
|
|
||||||
# Values placed into graph state before any node runs; reference anywhere via
|
|
||||||
# {{key}}.
|
|
||||||
#
|
|
||||||
# Note: `initial_prompt` is seeded automatically by Loki with the
|
|
||||||
# caller's prompt. So there's no need to set it here.
|
|
||||||
# ---------------------------------------------------------------------------
|
|
||||||
initial_state:
|
|
||||||
audience: "general reader"
|
|
||||||
# Seed an empty default for any key that a strict field (a node prompt /
|
|
||||||
# instructions / question / End output) references but that is only set on
|
|
||||||
# some paths. `refinement` is set only if the `refine` input node runs;
|
|
||||||
# seeding it "" keeps `finalize`'s strict prompt from failing on the
|
|
||||||
# approve-directly path.
|
|
||||||
refinement: ""
|
|
||||||
|
|
||||||
# ---------------------------------------------------------------------------
|
|
||||||
# Entry point
|
|
||||||
# ---------------------------------------------------------------------------
|
|
||||||
start: triage # ID of the first node to run (must exist in `nodes`)
|
|
||||||
|
|
||||||
# ---------------------------------------------------------------------------
|
|
||||||
# Nodes
|
|
||||||
# Each node is keyed by its id. The `id:` inside a node must match its key
|
|
||||||
# (it may also be omitted and thus Loki fills it in from the key).
|
|
||||||
#
|
|
||||||
# Node types: agent | script | approval | input | llm | rag | map | end
|
|
||||||
# ---------------------------------------------------------------------------
|
|
||||||
nodes:
|
|
||||||
|
|
||||||
# --- llm node -----------------------------------------------------------
|
|
||||||
# A one-shot LLM call (with an optional bounded tool-call loop). Runs in a
|
|
||||||
# fresh isolated context. Tools are strictly opt-in (see `tools`).
|
|
||||||
triage:
|
|
||||||
id: triage
|
|
||||||
type: llm
|
|
||||||
description: Classify the research request and extract its topic.
|
|
||||||
instructions: | # Optional system prompt (templated against state)
|
|
||||||
You triage research requests for a {{audience}} audience.
|
|
||||||
prompt: | # Required user prompt (templated against state)
|
|
||||||
Classify this request and extract the core research topic:
|
|
||||||
{{initial_prompt}}
|
|
||||||
tools: [] # Tool whitelist. Omitted or [] = no tools at all.
|
|
||||||
# A list narrows to exactly those entries.
|
|
||||||
output_schema: # Optional JSON Schema. The output is parsed to JSON
|
|
||||||
type: object # and its top-level object keys auto-merge into state
|
|
||||||
properties: # (so `topic` / `needs_deep_dive` become {{topic}} etc).
|
|
||||||
topic: { type: string }
|
|
||||||
needs_deep_dive: { type: boolean }
|
|
||||||
required: [topic, needs_deep_dive]
|
|
||||||
state_updates: # {{output}} = this node's result (here, the parsed object)
|
|
||||||
triage_result: "{{output}}"
|
|
||||||
# --- Polymorphic `next` -----------------------------------------------
|
|
||||||
# A single string runs the next node sequentially (e.g. `next: retrieve`).
|
|
||||||
# A list runs all listed nodes in parallel as one BSP super-step
|
|
||||||
# (for more info on BSP, see https://en.wikipedia.org/wiki/Bulk_synchronous_parallel).
|
|
||||||
# Their writes are merged via `reducers:` at the join. Branches converge
|
|
||||||
# implicitly when they all route to the same downstream node (here,
|
|
||||||
# `synthesize`). See the diamond:
|
|
||||||
#
|
|
||||||
# triage
|
|
||||||
# / \
|
|
||||||
# retrieve web_search (run concurrently)
|
|
||||||
# \ /
|
|
||||||
# synthesize (join; fires once after both finish)
|
|
||||||
next: [retrieve, web_search]
|
|
||||||
|
|
||||||
# --- rag node (parallel branch 1 of the diamond) ------------------------
|
|
||||||
# Hybrid (vector + keyword) retrieval against a per-node knowledge base.
|
|
||||||
# The knowledge base is built once, at agent load time, into
|
|
||||||
# <agent-dir>/retrieve.yaml (named after this node's id).
|
|
||||||
retrieve:
|
|
||||||
id: retrieve
|
|
||||||
type: rag
|
|
||||||
documents: # Required. Files, directories, URLs, loader paths.
|
|
||||||
- ./knowledge/ # relative paths resolve against the agent directory
|
|
||||||
- https://example.com/reference
|
|
||||||
query: "{{topic}}" # Retrieval query (templated). Default: {{initial_prompt}}.
|
|
||||||
top_k: 5 # Chunks to retrieve. Default = the KB's own top_k.
|
|
||||||
timeout: 120 # Retrieval timeout in seconds. Default 120.
|
|
||||||
# Knowledge-base build config (optional; used only when the KB is first
|
|
||||||
# built). When embedding_model + chunk_size + chunk_overlap are all set,
|
|
||||||
# the KB builds with no interactive prompts (works in non-interactive runs).
|
|
||||||
embedding_model: openai:text-embedding-3-small
|
|
||||||
chunk_size: 1000
|
|
||||||
chunk_overlap: 100
|
|
||||||
reranker_model: null # Optional reranker for hybrid-search results
|
|
||||||
batch_size: 100 # Optional embedding-request batch size
|
|
||||||
state_updates: # {{output}} = { context: <str>, sources: [<path>, ...] }
|
|
||||||
context: "{{output.context}}" # writes `context` -> `reducers.context = concat`
|
|
||||||
sources: "{{output.sources}}" # writes `sources` -> `reducers.sources = append`
|
|
||||||
next: synthesize # Joins with web_search at `synthesize`.
|
|
||||||
|
|
||||||
# --- llm node (parallel branch 2 of the diamond) ------------------------
|
|
||||||
# Runs concurrently with `retrieve`. Both branches write `context` and
|
|
||||||
# `sources`; the validator confirms both keys have a reducer declared, and
|
|
||||||
# the BSP scheduler merges them at the join.
|
|
||||||
web_search:
|
|
||||||
id: web_search
|
|
||||||
type: llm
|
|
||||||
instructions: "You are a web researcher. Cite every claim."
|
|
||||||
prompt: "Web research: {{topic}}. Return findings and sources."
|
|
||||||
tools:
|
|
||||||
- web_search_loki
|
|
||||||
- mcp:ddg-search
|
|
||||||
output_schema:
|
|
||||||
type: object
|
|
||||||
properties:
|
|
||||||
context: { type: string }
|
|
||||||
sources:
|
|
||||||
type: array
|
|
||||||
items: { type: string }
|
|
||||||
required: [context, sources]
|
|
||||||
# When `output_schema` is set, top-level keys auto-merge into state, so
|
|
||||||
# `context` and `sources` are produced without needing `state_updates`.
|
|
||||||
next: synthesize # Joins with retrieve at `synthesize`.
|
|
||||||
|
|
||||||
# --- script node (the diamond's join; also dispatches) -----------------
|
|
||||||
# Runs a .sh / .py / .ts script. The script receives state via the
|
|
||||||
# GRAPH_STATE env var (inline JSON) or GRAPH_STATE_FILE (path to a JSON
|
|
||||||
# file, used when state exceeds 32 KiB). Exactly one is set. It must print
|
|
||||||
# a single JSON object on stdout: keys merge into state, and the reserved
|
|
||||||
# `_next` key (if present) overrides routing.
|
|
||||||
#
|
|
||||||
# The script also receives these env vars (parity with bash tools called
|
|
||||||
# from normal agents):
|
|
||||||
# GRAPH_STATE / GRAPH_STATE_FILE state payload (one of the two is set)
|
|
||||||
# LLM_ROOT_DIR loki config dir
|
|
||||||
# LLM_PROMPT_UTILS_FILE path to .shared/prompt-utils.sh
|
|
||||||
# LLM_AGENT_DATA_DIR this agent's data directory
|
|
||||||
# LLM_AGENT_VAR_<NAME> one per declared `variables:` entry
|
|
||||||
# PATH with loki's functions bin dir prepended
|
|
||||||
# CLICOLOR_FORCE / FORCE_COLOR so child tools emit ANSI colors
|
|
||||||
# The script's working directory is loki's invocation CWD (not the agent
|
|
||||||
# directory), matching the behavior of bash tools.
|
|
||||||
#
|
|
||||||
# This node fires once: after both `retrieve` and `web_search` finish.
|
|
||||||
# The BSP scheduler dedups the two incoming edges into a single frontier
|
|
||||||
# entry, applies the staged branch writes through the reducers, then runs
|
|
||||||
# this node against the merged state. Inside the script, `context` is the
|
|
||||||
# concatenated text of both branches and `sources` is the combined list.
|
|
||||||
synthesize:
|
|
||||||
id: synthesize
|
|
||||||
type: script
|
|
||||||
script: scripts/synthesize.py # Path relative to the agent directory
|
|
||||||
timeout: 30 # Seconds. Default 30.
|
|
||||||
state_updates: # Applied after the stdout JSON is merged
|
|
||||||
decided_for: "{{topic}}"
|
|
||||||
next: summarize # Default route if the script emits no `_next`
|
|
||||||
fallback: summarize # Route taken if the script fails (crash / bad JSON)
|
|
||||||
# This script is expected to emit `_next: deep_dive` (or `_next: subjects_map`
|
|
||||||
# to demonstrate the map node below), or no `_next` (then `next` is used).
|
|
||||||
# Targets reached only via the script's dynamic `_next` get an
|
|
||||||
# "unreachable" warning from the validator. This is expected for `_next`-routed
|
|
||||||
# targets.
|
|
||||||
|
|
||||||
# --- agent node ---------------------------------------------------------
|
|
||||||
# Spawns a full Loki sub-agent and waits for it. The child uses its own
|
|
||||||
# tool stack. Agent nodes have no `tools:` field. No schema hint is
|
|
||||||
# injected even when `output_schema` is set (unlike llm nodes).
|
|
||||||
deep_dive:
|
|
||||||
id: deep_dive
|
|
||||||
type: agent
|
|
||||||
agent: deep-research # Name of an existing Loki agent to spawn
|
|
||||||
prompt: | # User message sent to the child (templated)
|
|
||||||
Research {{topic}} in depth. Existing context:
|
|
||||||
{{context}}
|
|
||||||
timeout: 600 # Optional wall-clock cap, seconds. Default 300.
|
|
||||||
output_schema: # Optional. Same extraction as llm nodes
|
|
||||||
type: object
|
|
||||||
properties:
|
|
||||||
summary: { type: string }
|
|
||||||
findings:
|
|
||||||
type: array
|
|
||||||
items: { type: string }
|
|
||||||
required: [summary, findings]
|
|
||||||
state_updates:
|
|
||||||
research: "{{output}}"
|
|
||||||
next: review # Required for agent nodes
|
|
||||||
|
|
||||||
# --- map node (Dynamic fan-out. Think: LangGraph's `Send` API) ----------------
|
|
||||||
# Spawns one parallel sub-branch per item in `over`. Each sub-branch runs
|
|
||||||
# the node referenced by `branch:` with the item bound to `as:`. Outputs
|
|
||||||
# collect into the array named by `collect_into:`, preserving input order.
|
|
||||||
#
|
|
||||||
# Reach via `synthesize`'s `_next: subjects_map`. The producer is expected
|
|
||||||
# to have written a list at `subjects` (e.g. an upstream LLM node with an
|
|
||||||
# `output_schema` returning {"subjects": ["a", "b", "c"]}).
|
|
||||||
subjects_map:
|
|
||||||
id: subjects_map
|
|
||||||
type: map
|
|
||||||
over: "{{subjects}}" # Required. List expression resolved from state.
|
|
||||||
# Empty list is allowed. It means no branches spawn,
|
|
||||||
# and thus `collect_into` is written as [].
|
|
||||||
as: subject # Required. Per-branch state key holding the
|
|
||||||
# current item. Read with {{subject}} inside
|
|
||||||
# the branch node's prompt.
|
|
||||||
branch: research_subject # Required. Node id to invoke per item.
|
|
||||||
# Must point to an llm | agent | rag | script
|
|
||||||
# node satisfying the map branch contract:
|
|
||||||
# - no `next:` (atomic, joined at map exit)
|
|
||||||
# - no `state_updates:` other than via the
|
|
||||||
# map's `collect_into` channel
|
|
||||||
# - no `output_schema:` (top-level merge
|
|
||||||
# would clash with collect_into)
|
|
||||||
# Validator enforces all three.
|
|
||||||
collect_into: subject_findings # Required. State key for the array of
|
|
||||||
# per-branch outputs, in input order
|
|
||||||
# (not spawn-finish order).
|
|
||||||
max_concurrency: 3 # Optional per-map cap. Defaults to
|
|
||||||
# settings.max_concurrency above.
|
|
||||||
output_key: output # Optional. State key the branch's output
|
|
||||||
# appears under. Default "output". Useful
|
|
||||||
# only if the branch reads its own bound
|
|
||||||
# name back (rare).
|
|
||||||
next: aggregate_subjects # Where to go after all sub-branches finish.
|
|
||||||
|
|
||||||
# Branch node for subjects_map. Each invocation receives a different
|
|
||||||
# `subject` in state. The branch is "atomic", meaning it cannot route on
|
|
||||||
# its own; the surrounding `map` joins after all invocations finish.
|
|
||||||
research_subject:
|
|
||||||
id: research_subject
|
|
||||||
type: llm
|
|
||||||
instructions: "Research one subject deeply for a {{audience}} audience."
|
|
||||||
prompt: "Research {{subject}}: pull the key facts and one citation."
|
|
||||||
tools:
|
|
||||||
- web_search_loki
|
|
||||||
# No `next:`, `state_updates:`, or `output_schema:` here. Map branches
|
|
||||||
# have a strict contract (see `subjects_map.branch` comment).
|
|
||||||
|
|
||||||
# Aggregator that runs after the map joins. Reads the collected list.
|
|
||||||
aggregate_subjects:
|
|
||||||
id: aggregate_subjects
|
|
||||||
type: llm
|
|
||||||
instructions: "Combine N per-subject reports into one cohesive summary."
|
|
||||||
prompt: |
|
|
||||||
Per-subject reports (in original input order):
|
|
||||||
{{subject_findings}}
|
|
||||||
state_updates:
|
|
||||||
research: "{{output}}"
|
|
||||||
next: review
|
|
||||||
|
|
||||||
# --- llm node with a narrowed tool whitelist ----------------------------
|
|
||||||
summarize:
|
|
||||||
id: summarize
|
|
||||||
type: llm
|
|
||||||
instructions: "You write concise research summaries for a {{audience}} audience."
|
|
||||||
prompt: "Summarize the topic {{topic}}, using your tools as needed."
|
|
||||||
tools: # Narrow whitelist: exactly these entries, nothing else
|
|
||||||
- web_search_loki # an exact global-tool / custom-tool name
|
|
||||||
- mcp:ddg-search # `mcp:<server>` includes that server's functions
|
|
||||||
model: claude:claude-haiku-4-5 # Optional per-node model override
|
|
||||||
temperature: 0.3 # Optional per-node sampling override
|
|
||||||
max_attempts: 2 # Retry count on transient errors only. Default 1.
|
|
||||||
max_iterations: 10 # Tool-call-loop turn cap. Default 10.
|
|
||||||
fallback: review # Route here if all attempts fail
|
|
||||||
timeout: 300 # Optional node wall-clock cap, seconds (unset = no timeout)
|
|
||||||
state_updates:
|
|
||||||
research: "{{output}}"
|
|
||||||
next: review # Required for llm nodes: the success route
|
|
||||||
|
|
||||||
# --- approval node ------------------------------------------------------
|
|
||||||
# Human-in-the-loop checkpoint. `user__ask` always offers a free-form
|
|
||||||
# "type your own answer" option, so `on_other` is required.
|
|
||||||
review:
|
|
||||||
id: review
|
|
||||||
type: approval
|
|
||||||
question: |
|
|
||||||
Proposed research result for {{topic}}:
|
|
||||||
{{research}}
|
|
||||||
|
|
||||||
Approve?
|
|
||||||
options: # The listed choices shown to the user
|
|
||||||
- "yes"
|
|
||||||
- "no"
|
|
||||||
routes: # Map each listed option to its next node
|
|
||||||
"yes": finalize
|
|
||||||
"no": rejected_end
|
|
||||||
on_other: refine # Required: route for ANY answer not in `routes`
|
|
||||||
state_updates:
|
|
||||||
decision: "{{choice}}" # {{choice}} = the chosen option or the free-form text
|
|
||||||
|
|
||||||
# --- input node ---------------------------------------------------------
|
|
||||||
# Collects a free-form string from the user.
|
|
||||||
refine:
|
|
||||||
id: refine
|
|
||||||
type: input
|
|
||||||
question: "What should be changed about the research result?"
|
|
||||||
default: "tighten the summary" # Optional: used if the user submits empty input.
|
|
||||||
# Note: a substituted default is not re-validated,
|
|
||||||
# so make sure it would satisfy `validation`.
|
|
||||||
validation: "len(input) > 0" # Optional length predicate: len(input) <op> N,
|
|
||||||
# <op> in > >= < <= == . Length only -- no regex.
|
|
||||||
state_updates:
|
|
||||||
refinement: "{{input}}" # {{input}} = the user's text
|
|
||||||
next: finalize # Required for input nodes: the success route
|
|
||||||
|
|
||||||
# --- llm node (final synthesis) -----------------------------------------
|
|
||||||
finalize:
|
|
||||||
id: finalize
|
|
||||||
type: llm
|
|
||||||
prompt: |
|
|
||||||
Produce the final research report for {{topic}}.
|
|
||||||
Result so far: {{research}}
|
|
||||||
Requested refinement (if any): {{refinement}}
|
|
||||||
state_updates:
|
|
||||||
final_report: "{{output}}"
|
|
||||||
next: done
|
|
||||||
|
|
||||||
# --- end nodes ----------------------------------------------------------
|
|
||||||
# Terminate the graph. `output` (templated, lenient interpolation) becomes
|
|
||||||
# the graph's final result. A graph needs at least one `end` node.
|
|
||||||
done:
|
|
||||||
id: done
|
|
||||||
type: end
|
|
||||||
state_updates: # Optional: applied before `output` is rendered
|
|
||||||
status: "completed"
|
|
||||||
output: |
|
|
||||||
[{{status}}] {{final_report}}
|
|
||||||
|
|
||||||
Sources: {{sources}}
|
|
||||||
|
|
||||||
rejected_end:
|
|
||||||
id: rejected_end
|
|
||||||
type: end
|
|
||||||
output: "Research on {{topic}} was not approved."
|
|
||||||
@@ -487,6 +487,14 @@
|
|||||||
thinking:
|
thinking:
|
||||||
type: enabled
|
type: enabled
|
||||||
budget_tokens: 16000
|
budget_tokens: 16000
|
||||||
|
- name: claude-3-5-haiku-20241022
|
||||||
|
max_input_tokens: 200000
|
||||||
|
max_output_tokens: 8192
|
||||||
|
require_max_tokens: true
|
||||||
|
input_price: 0.8
|
||||||
|
output_price: 4
|
||||||
|
supports_vision: true
|
||||||
|
supports_function_calling: true
|
||||||
|
|
||||||
# Links:
|
# Links:
|
||||||
# - https://docs.mistral.ai/getting-started/models/models_overview/
|
# - https://docs.mistral.ai/getting-started/models/models_overview/
|
||||||
|
|||||||
+1
-30
@@ -1,11 +1,9 @@
|
|||||||
use crate::client::{ModelType, list_models};
|
use crate::client::{ModelType, list_models};
|
||||||
use crate::config::paths;
|
use crate::config::paths;
|
||||||
use crate::config::{AppConfig, Config, list_agents, list_sessions};
|
use crate::config::{AppConfig, Config, list_agents, list_sessions};
|
||||||
use crate::utils::list_file_names;
|
|
||||||
use crate::vault::Vault;
|
use crate::vault::Vault;
|
||||||
use clap_complete::{CompletionCandidate, Shell, generate};
|
use clap_complete::{CompletionCandidate, Shell, generate};
|
||||||
use clap_complete_nushell::Nushell;
|
use clap_complete_nushell::Nushell;
|
||||||
use std::env;
|
|
||||||
use std::ffi::OsStr;
|
use std::ffi::OsStr;
|
||||||
use std::io;
|
use std::io;
|
||||||
|
|
||||||
@@ -96,36 +94,9 @@ pub(super) fn macro_completer(current: &OsStr) -> Vec<CompletionCandidate> {
|
|||||||
.collect()
|
.collect()
|
||||||
}
|
}
|
||||||
|
|
||||||
fn extract_agent_from_args() -> Option<String> {
|
|
||||||
let args: Vec<String> = env::args().collect();
|
|
||||||
let mut i = 0;
|
|
||||||
while i < args.len() {
|
|
||||||
let arg = &args[i];
|
|
||||||
|
|
||||||
if let Some(value) = arg.strip_prefix("--agent=") {
|
|
||||||
return Some(value.to_string());
|
|
||||||
}
|
|
||||||
|
|
||||||
if (arg == "--agent" || arg == "-a") && i + 1 < args.len() {
|
|
||||||
return Some(args[i + 1].clone());
|
|
||||||
}
|
|
||||||
|
|
||||||
i += 1;
|
|
||||||
}
|
|
||||||
None
|
|
||||||
}
|
|
||||||
|
|
||||||
pub(super) fn session_completer(current: &OsStr) -> Vec<CompletionCandidate> {
|
pub(super) fn session_completer(current: &OsStr) -> Vec<CompletionCandidate> {
|
||||||
let cur = current.to_string_lossy();
|
let cur = current.to_string_lossy();
|
||||||
|
list_sessions()
|
||||||
let sessions = if let Some(agent_name) = extract_agent_from_args() {
|
|
||||||
let sessions_dir = paths::agent_data_dir(&agent_name).join("sessions");
|
|
||||||
list_file_names(sessions_dir, ".yaml")
|
|
||||||
} else {
|
|
||||||
list_sessions()
|
|
||||||
};
|
|
||||||
|
|
||||||
sessions
|
|
||||||
.into_iter()
|
.into_iter()
|
||||||
.filter(|s| s.starts_with(&*cur))
|
.filter(|s| s.starts_with(&*cur))
|
||||||
.map(CompletionCandidate::new)
|
.map(CompletionCandidate::new)
|
||||||
|
|||||||
+2
-48
@@ -4,10 +4,9 @@ use crate::cli::completer::{
|
|||||||
ShellCompletion, agent_completer, macro_completer, model_completer, rag_completer,
|
ShellCompletion, agent_completer, macro_completer, model_completer, rag_completer,
|
||||||
role_completer, secrets_completer, session_completer,
|
role_completer, secrets_completer, session_completer,
|
||||||
};
|
};
|
||||||
use crate::config::{AssetCategory, InstallFilter};
|
|
||||||
use anyhow::{Context, Result};
|
use anyhow::{Context, Result};
|
||||||
use clap::ValueHint;
|
use clap::ValueHint;
|
||||||
use clap::{Parser, crate_authors, crate_description, crate_version};
|
use clap::{Parser, crate_authors, crate_description, crate_name, crate_version};
|
||||||
use clap_complete::ArgValueCompleter;
|
use clap_complete::ArgValueCompleter;
|
||||||
use is_terminal::IsTerminal;
|
use is_terminal::IsTerminal;
|
||||||
use std::io::{Read, stdin};
|
use std::io::{Read, stdin};
|
||||||
@@ -15,7 +14,7 @@ use std::io::{Read, stdin};
|
|||||||
#[derive(Parser, Debug)]
|
#[derive(Parser, Debug)]
|
||||||
#[command(author, version, about, long_about = None)]
|
#[command(author, version, about, long_about = None)]
|
||||||
#[command(
|
#[command(
|
||||||
name = "loki",
|
name = crate_name!(),
|
||||||
author = crate_authors!(),
|
author = crate_authors!(),
|
||||||
version = crate_version!(),
|
version = crate_version!(),
|
||||||
about = crate_description!(),
|
about = crate_description!(),
|
||||||
@@ -83,18 +82,6 @@ pub struct Cli {
|
|||||||
/// Build all configured Bash tool scripts
|
/// Build all configured Bash tool scripts
|
||||||
#[arg(long)]
|
#[arg(long)]
|
||||||
pub build_tools: bool,
|
pub build_tools: bool,
|
||||||
/// Reinstall bundled assets, overwriting any local changes
|
|
||||||
#[arg(long, value_name = "CATEGORY", value_enum)]
|
|
||||||
pub install: Option<AssetCategory>,
|
|
||||||
/// Install assets from a remote git repository (URL may be suffixed with #<ref>)
|
|
||||||
#[arg(long, value_name = "GIT_URL")]
|
|
||||||
pub install_from: Option<String>,
|
|
||||||
/// Restrict --install-from to a single asset category
|
|
||||||
#[arg(long, value_name = "CATEGORY", value_enum, requires = "install_from")]
|
|
||||||
pub filter: Option<InstallFilter>,
|
|
||||||
/// Overwrite all conflicts without prompting (used with --install-from)
|
|
||||||
#[arg(long, requires = "install_from")]
|
|
||||||
pub install_force: bool,
|
|
||||||
/// Sync models updates
|
/// Sync models updates
|
||||||
#[arg(long)]
|
#[arg(long)]
|
||||||
pub sync_models: bool,
|
pub sync_models: bool,
|
||||||
@@ -146,12 +133,6 @@ pub struct Cli {
|
|||||||
/// Generate static shell completion scripts
|
/// Generate static shell completion scripts
|
||||||
#[arg(long, value_name = "SHELL", value_enum)]
|
#[arg(long, value_name = "SHELL", value_enum)]
|
||||||
pub completions: Option<ShellCompletion>,
|
pub completions: Option<ShellCompletion>,
|
||||||
/// Update Loki to the latest release, or to a specific version
|
|
||||||
#[arg(long, value_name = "VERSION")]
|
|
||||||
pub update: Option<Option<String>>,
|
|
||||||
/// With --update, update even if Loki was installed via a package manager
|
|
||||||
#[arg(long, requires = "update")]
|
|
||||||
pub force: bool,
|
|
||||||
}
|
}
|
||||||
|
|
||||||
impl Cli {
|
impl Cli {
|
||||||
@@ -411,31 +392,4 @@ mod tests {
|
|||||||
let cli = parse(&["--macro", "my-macro"]);
|
let cli = parse(&["--macro", "my-macro"]);
|
||||||
assert_eq!(cli.macro_name, Some("my-macro".to_string()));
|
assert_eq!(cli.macro_name, Some("my-macro".to_string()));
|
||||||
}
|
}
|
||||||
|
|
||||||
#[test]
|
|
||||||
fn parse_update_flag_no_value() {
|
|
||||||
let cli = parse(&["--update"]);
|
|
||||||
|
|
||||||
assert_eq!(cli.update, Some(None));
|
|
||||||
}
|
|
||||||
|
|
||||||
#[test]
|
|
||||||
fn parse_update_flag_with_version() {
|
|
||||||
let cli = parse(&["--update", "v0.4.0"]);
|
|
||||||
|
|
||||||
assert_eq!(cli.update, Some(Some("v0.4.0".to_string())));
|
|
||||||
}
|
|
||||||
|
|
||||||
#[test]
|
|
||||||
fn parse_update_with_force() {
|
|
||||||
let cli = parse(&["--update", "--force"]);
|
|
||||||
|
|
||||||
assert_eq!(cli.update, Some(None));
|
|
||||||
assert!(cli.force);
|
|
||||||
}
|
|
||||||
|
|
||||||
#[test]
|
|
||||||
fn parse_force_without_update_fails() {
|
|
||||||
assert!(Cli::try_parse_from(["loki", "--force"]).is_err());
|
|
||||||
}
|
|
||||||
}
|
}
|
||||||
|
|||||||
+25
-30
@@ -117,38 +117,33 @@ async fn prepare_chat_completions(
|
|||||||
/// So this function injects the Claude Code system prompt into the request
|
/// So this function injects the Claude Code system prompt into the request
|
||||||
/// body to make it a valid request.
|
/// body to make it a valid request.
|
||||||
fn inject_oauth_system_prompt(body: &mut Value) {
|
fn inject_oauth_system_prompt(body: &mut Value) {
|
||||||
let existing_text = match body.get("system") {
|
let prefix_block = json!({
|
||||||
Some(Value::String(s)) => {
|
"type": "text",
|
||||||
if s.starts_with(CLAUDE_CODE_PREFIX) {
|
"text": CLAUDE_CODE_PREFIX,
|
||||||
return;
|
});
|
||||||
}
|
|
||||||
(!s.is_empty()).then(|| s.clone())
|
match body.get("system") {
|
||||||
}
|
Some(Value::String(existing)) => {
|
||||||
Some(Value::Array(blocks)) => {
|
let existing_block = json!({
|
||||||
let already_injected = blocks.iter().any(|b| {
|
"type": "text",
|
||||||
b.get("text")
|
"text": existing,
|
||||||
.and_then(|t| t.as_str())
|
|
||||||
.map(|t| t.starts_with(CLAUDE_CODE_PREFIX))
|
|
||||||
.unwrap_or(false)
|
|
||||||
});
|
});
|
||||||
if already_injected {
|
body["system"] = json!([prefix_block, existing_block]);
|
||||||
return;
|
|
||||||
}
|
|
||||||
let joined: Vec<String> = blocks
|
|
||||||
.iter()
|
|
||||||
.filter_map(|b| b.get("text").and_then(|t| t.as_str()).map(String::from))
|
|
||||||
.collect();
|
|
||||||
(!joined.is_empty()).then(|| joined.join("\n\n"))
|
|
||||||
}
|
}
|
||||||
_ => None,
|
Some(Value::Array(_)) => {
|
||||||
};
|
if let Some(arr) = body["system"].as_array_mut() {
|
||||||
|
let already_injected = arr
|
||||||
let merged = match existing_text {
|
.iter()
|
||||||
Some(rest) => format!("{}\n\n{}", CLAUDE_CODE_PREFIX, rest),
|
.any(|block| block["text"].as_str() == Some(CLAUDE_CODE_PREFIX));
|
||||||
None => CLAUDE_CODE_PREFIX.to_string(),
|
if !already_injected {
|
||||||
};
|
arr.insert(0, prefix_block);
|
||||||
|
}
|
||||||
body["system"] = json!([{ "type": "text", "text": merged }]);
|
}
|
||||||
|
}
|
||||||
|
_ => {
|
||||||
|
body["system"] = json!([prefix_block]);
|
||||||
|
}
|
||||||
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
pub async fn claude_chat_completions(
|
pub async fn claude_chat_completions(
|
||||||
|
|||||||
@@ -1,6 +1,6 @@
|
|||||||
use super::*;
|
use super::*;
|
||||||
|
|
||||||
use crate::config::{RenderMode, paths};
|
use crate::config::paths;
|
||||||
use crate::{
|
use crate::{
|
||||||
config::{AppConfig, Input, RequestContext},
|
config::{AppConfig, Input, RequestContext},
|
||||||
function::{FunctionDeclaration, ToolCall, ToolResult, eval_tool_calls},
|
function::{FunctionDeclaration, ToolCall, ToolResult, eval_tool_calls},
|
||||||
@@ -418,8 +418,7 @@ pub async fn call_chat_completions(
|
|||||||
abort_signal: AbortSignal,
|
abort_signal: AbortSignal,
|
||||||
) -> Result<(String, Vec<ToolResult>)> {
|
) -> Result<(String, Vec<ToolResult>)> {
|
||||||
let is_child_agent = ctx.current_depth > 0;
|
let is_child_agent = ctx.current_depth > 0;
|
||||||
let suppress_spinner = is_child_agent || ctx.render_mode == RenderMode::Silent;
|
let spinner_message = if is_child_agent { "" } else { "Generating" };
|
||||||
let spinner_message = if suppress_spinner { "" } else { "Generating" };
|
|
||||||
let ret = abortable_run_with_spinner(
|
let ret = abortable_run_with_spinner(
|
||||||
client.chat_completions(input.clone()),
|
client.chat_completions(input.clone()),
|
||||||
spinner_message,
|
spinner_message,
|
||||||
@@ -460,14 +459,10 @@ pub async fn call_chat_completions_streaming(
|
|||||||
) -> Result<(String, Vec<ToolResult>)> {
|
) -> Result<(String, Vec<ToolResult>)> {
|
||||||
let (tx, rx) = unbounded_channel();
|
let (tx, rx) = unbounded_channel();
|
||||||
let mut handler = SseHandler::new(tx, abort_signal.clone());
|
let mut handler = SseHandler::new(tx, abort_signal.clone());
|
||||||
let silent = ctx.render_mode == RenderMode::Silent;
|
|
||||||
if silent {
|
|
||||||
handler.set_silent(true);
|
|
||||||
}
|
|
||||||
|
|
||||||
let (send_ret, render_ret) = tokio::join!(
|
let (send_ret, render_ret) = tokio::join!(
|
||||||
client.chat_completions_streaming(input, &mut handler),
|
client.chat_completions_streaming(input, &mut handler),
|
||||||
render_stream(rx, client.app_config(), abort_signal.clone(), silent),
|
render_stream(rx, client.app_config(), abort_signal.clone()),
|
||||||
);
|
);
|
||||||
|
|
||||||
if handler.abort().aborted() {
|
if handler.abort().aborted() {
|
||||||
|
|||||||
+5
-10
@@ -94,21 +94,21 @@ impl MessageContent {
|
|||||||
match self {
|
match self {
|
||||||
MessageContent::Text(text) => multiline_text(text),
|
MessageContent::Text(text) => multiline_text(text),
|
||||||
MessageContent::Array(list) => {
|
MessageContent::Array(list) => {
|
||||||
let (mut concatenated_text, mut files) = (String::new(), vec![]);
|
let (mut concated_text, mut files) = (String::new(), vec![]);
|
||||||
for item in list {
|
for item in list {
|
||||||
match item {
|
match item {
|
||||||
MessageContentPart::Text { text } => {
|
MessageContentPart::Text { text } => {
|
||||||
concatenated_text = format!("{concatenated_text} {text}")
|
concated_text = format!("{concated_text} {text}")
|
||||||
}
|
}
|
||||||
MessageContentPart::ImageUrl { image_url } => {
|
MessageContentPart::ImageUrl { image_url } => {
|
||||||
files.push(resolve_url_fn(&image_url.url))
|
files.push(resolve_url_fn(&image_url.url))
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
if !concatenated_text.is_empty() {
|
if !concated_text.is_empty() {
|
||||||
concatenated_text = format!(" -- {}", multiline_text(&concatenated_text))
|
concated_text = format!(" -- {}", multiline_text(&concated_text))
|
||||||
}
|
}
|
||||||
format!(".file {}{}", files.join(" "), concatenated_text)
|
format!(".file {}{}", files.join(" "), concated_text)
|
||||||
}
|
}
|
||||||
MessageContent::ToolCalls(MessageContentToolCalls {
|
MessageContent::ToolCalls(MessageContentToolCalls {
|
||||||
tool_results, text, ..
|
tool_results, text, ..
|
||||||
@@ -227,14 +227,9 @@ pub fn patch_messages(messages: &mut Vec<Message>, model: &Model) {
|
|||||||
}
|
}
|
||||||
|
|
||||||
pub fn extract_system_message(messages: &mut Vec<Message>) -> Option<String> {
|
pub fn extract_system_message(messages: &mut Vec<Message>) -> Option<String> {
|
||||||
if messages.is_empty() {
|
|
||||||
return None;
|
|
||||||
}
|
|
||||||
|
|
||||||
if messages[0].role.is_system() {
|
if messages[0].role.is_system() {
|
||||||
let system_message = messages.remove(0);
|
let system_message = messages.remove(0);
|
||||||
return Some(system_message.content.to_text());
|
return Some(system_message.content.to_text());
|
||||||
}
|
}
|
||||||
|
|
||||||
None
|
None
|
||||||
}
|
}
|
||||||
|
|||||||
+32
-52
@@ -2,9 +2,9 @@ use super::{ToolCall, catch_error};
|
|||||||
use crate::utils::AbortSignal;
|
use crate::utils::AbortSignal;
|
||||||
|
|
||||||
use anyhow::{Context, Result, anyhow, bail};
|
use anyhow::{Context, Result, anyhow, bail};
|
||||||
use eventsource_stream::Eventsource;
|
|
||||||
use futures_util::{Stream, StreamExt};
|
use futures_util::{Stream, StreamExt};
|
||||||
use reqwest::{RequestBuilder, header};
|
use reqwest::RequestBuilder;
|
||||||
|
use reqwest_eventsource::{Error as EventSourceError, Event, RequestBuilderExt};
|
||||||
use serde_json::Value;
|
use serde_json::Value;
|
||||||
use tokio::sync::mpsc::UnboundedSender;
|
use tokio::sync::mpsc::UnboundedSender;
|
||||||
|
|
||||||
@@ -16,7 +16,6 @@ pub struct SseHandler {
|
|||||||
last_tool_calls: Vec<ToolCall>,
|
last_tool_calls: Vec<ToolCall>,
|
||||||
max_call_repeats: usize,
|
max_call_repeats: usize,
|
||||||
call_repeat_chain_len: usize,
|
call_repeat_chain_len: usize,
|
||||||
silent: bool,
|
|
||||||
}
|
}
|
||||||
|
|
||||||
impl SseHandler {
|
impl SseHandler {
|
||||||
@@ -29,24 +28,14 @@ impl SseHandler {
|
|||||||
last_tool_calls: Vec::new(),
|
last_tool_calls: Vec::new(),
|
||||||
max_call_repeats: 2,
|
max_call_repeats: 2,
|
||||||
call_repeat_chain_len: 3,
|
call_repeat_chain_len: 3,
|
||||||
silent: false,
|
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
pub fn set_silent(&mut self, silent: bool) {
|
|
||||||
self.silent = silent;
|
|
||||||
}
|
|
||||||
|
|
||||||
pub fn text(&mut self, text: &str) -> Result<()> {
|
pub fn text(&mut self, text: &str) -> Result<()> {
|
||||||
if text.is_empty() {
|
if text.is_empty() {
|
||||||
return Ok(());
|
return Ok(());
|
||||||
}
|
}
|
||||||
self.buffer.push_str(text);
|
self.buffer.push_str(text);
|
||||||
|
|
||||||
if self.silent {
|
|
||||||
return Ok(());
|
|
||||||
}
|
|
||||||
|
|
||||||
let ret = self
|
let ret = self
|
||||||
.sender
|
.sender
|
||||||
.send(SseEvent::Text(text.to_string()))
|
.send(SseEvent::Text(text.to_string()))
|
||||||
@@ -204,46 +193,11 @@ pub async fn sse_stream<F>(builder: RequestBuilder, mut handle: F) -> Result<()>
|
|||||||
where
|
where
|
||||||
F: FnMut(SseMessage) -> Result<bool>,
|
F: FnMut(SseMessage) -> Result<bool>,
|
||||||
{
|
{
|
||||||
let res = builder
|
let mut es = builder.eventsource()?;
|
||||||
.header(header::ACCEPT, "text/event-stream")
|
|
||||||
.header(header::CACHE_CONTROL, "no-store")
|
|
||||||
.send()
|
|
||||||
.await?;
|
|
||||||
let status = res.status();
|
|
||||||
if !status.is_success() {
|
|
||||||
let text = res.text().await?;
|
|
||||||
let data: Value = match text.parse() {
|
|
||||||
Ok(data) => data,
|
|
||||||
Err(_) => {
|
|
||||||
bail!(
|
|
||||||
"Invalid response data: {text} (status: {})",
|
|
||||||
status.as_u16()
|
|
||||||
);
|
|
||||||
}
|
|
||||||
};
|
|
||||||
catch_error(&data, status.as_u16())?;
|
|
||||||
return Ok(());
|
|
||||||
}
|
|
||||||
|
|
||||||
let content_type = res
|
|
||||||
.headers()
|
|
||||||
.get(header::CONTENT_TYPE)
|
|
||||||
.and_then(|value| value.to_str().ok())
|
|
||||||
.map(|value| value.to_string());
|
|
||||||
let is_event_stream = content_type
|
|
||||||
.as_deref()
|
|
||||||
.map(|ct| ct.starts_with("text/event-stream"))
|
|
||||||
.unwrap_or(false);
|
|
||||||
if !is_event_stream {
|
|
||||||
let header_value = content_type.unwrap_or_default();
|
|
||||||
let text = res.text().await?;
|
|
||||||
bail!("Invalid response event-stream. content-type: {header_value}, data: {text}");
|
|
||||||
}
|
|
||||||
|
|
||||||
let mut es = res.bytes_stream().boxed().eventsource();
|
|
||||||
while let Some(event) = es.next().await {
|
while let Some(event) = es.next().await {
|
||||||
match event {
|
match event {
|
||||||
Ok(message) => {
|
Ok(Event::Open) => {}
|
||||||
|
Ok(Event::Message(message)) => {
|
||||||
let message = SseMessage {
|
let message = SseMessage {
|
||||||
event: message.event,
|
event: message.event,
|
||||||
data: message.data,
|
data: message.data,
|
||||||
@@ -253,7 +207,33 @@ where
|
|||||||
}
|
}
|
||||||
}
|
}
|
||||||
Err(err) => {
|
Err(err) => {
|
||||||
bail!("{err}");
|
match err {
|
||||||
|
EventSourceError::StreamEnded => {}
|
||||||
|
EventSourceError::InvalidStatusCode(status, res) => {
|
||||||
|
let text = res.text().await?;
|
||||||
|
let data: Value = match text.parse() {
|
||||||
|
Ok(data) => data,
|
||||||
|
Err(_) => {
|
||||||
|
bail!(
|
||||||
|
"Invalid response data: {text} (status: {})",
|
||||||
|
status.as_u16()
|
||||||
|
);
|
||||||
|
}
|
||||||
|
};
|
||||||
|
catch_error(&data, status.as_u16())?;
|
||||||
|
}
|
||||||
|
EventSourceError::InvalidContentType(header_value, res) => {
|
||||||
|
let text = res.text().await?;
|
||||||
|
bail!(
|
||||||
|
"Invalid response event-stream. content-type: {}, data: {text}",
|
||||||
|
header_value.to_str().unwrap_or_default()
|
||||||
|
);
|
||||||
|
}
|
||||||
|
_ => {
|
||||||
|
bail!("{}", err);
|
||||||
|
}
|
||||||
|
}
|
||||||
|
es.close();
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|||||||
+54
-301
@@ -11,8 +11,6 @@ use crate::config::prompts::{
|
|||||||
DEFAULT_SPAWN_INSTRUCTIONS, DEFAULT_TEAMMATE_INSTRUCTIONS, DEFAULT_TODO_INSTRUCTIONS,
|
DEFAULT_SPAWN_INSTRUCTIONS, DEFAULT_TEAMMATE_INSTRUCTIONS, DEFAULT_TODO_INSTRUCTIONS,
|
||||||
DEFAULT_USER_INTERACTION_INSTRUCTIONS,
|
DEFAULT_USER_INTERACTION_INSTRUCTIONS,
|
||||||
};
|
};
|
||||||
use crate::graph::{Graph, GraphParser, NodeType};
|
|
||||||
use crate::rag::RagInitConfig;
|
|
||||||
use crate::vault::SECRET_RE;
|
use crate::vault::SECRET_RE;
|
||||||
use anyhow::{Context, Result};
|
use anyhow::{Context, Result};
|
||||||
use fancy_regex::Captures;
|
use fancy_regex::Captures;
|
||||||
@@ -39,13 +37,12 @@ pub struct Agent {
|
|||||||
session_dynamic_instructions: Option<String>,
|
session_dynamic_instructions: Option<String>,
|
||||||
functions: Functions,
|
functions: Functions,
|
||||||
rag: Option<Arc<Rag>>,
|
rag: Option<Arc<Rag>>,
|
||||||
graph_rags: HashMap<String, Arc<Rag>>,
|
|
||||||
model: Model,
|
model: Model,
|
||||||
vault: GlobalVault,
|
vault: GlobalVault,
|
||||||
}
|
}
|
||||||
|
|
||||||
impl Agent {
|
impl Agent {
|
||||||
pub fn install_builtin_agents(force: bool) -> Result<()> {
|
pub fn install_builtin_agents() -> Result<()> {
|
||||||
info!(
|
info!(
|
||||||
"Installing built-in agents in {}",
|
"Installing built-in agents in {}",
|
||||||
paths::agents_data_dir().display()
|
paths::agents_data_dir().display()
|
||||||
@@ -65,7 +62,7 @@ impl Agent {
|
|||||||
#[cfg_attr(not(unix), expect(unused))]
|
#[cfg_attr(not(unix), expect(unused))]
|
||||||
let is_script = matches!(file_extension.as_deref(), Some("sh") | Some("py"));
|
let is_script = matches!(file_extension.as_deref(), Some("sh") | Some("py"));
|
||||||
|
|
||||||
if file_path.exists() && !force {
|
if file_path.exists() {
|
||||||
debug!(
|
debug!(
|
||||||
"Agent file already exists, skipping: {}",
|
"Agent file already exists, skipping: {}",
|
||||||
file_path.display()
|
file_path.display()
|
||||||
@@ -100,28 +97,10 @@ impl Agent {
|
|||||||
let loaders = app.document_loaders.clone();
|
let loaders = app.document_loaders.clone();
|
||||||
let rag_path = paths::agent_rag_file(name, DEFAULT_AGENT_NAME);
|
let rag_path = paths::agent_rag_file(name, DEFAULT_AGENT_NAME);
|
||||||
let config_path = paths::agent_config_file(name);
|
let config_path = paths::agent_config_file(name);
|
||||||
let graph_path = paths::agent_graph_file(name);
|
let mut agent_config = if config_path.exists() {
|
||||||
let mut graph_for_rag: Option<Graph> = None;
|
AgentConfig::load(&config_path)?
|
||||||
let mut agent_config = match (config_path.exists(), graph_path.exists()) {
|
} else {
|
||||||
(true, true) => bail!(
|
bail!("Agent config file not found at '{}'", config_path.display())
|
||||||
"Agent '{name}' has both config.yaml and graph.yaml. A graph agent \
|
|
||||||
is defined by graph.yaml alone; a normal agent by config.yaml alone. \
|
|
||||||
Remove one of the two files."
|
|
||||||
),
|
|
||||||
(true, false) => AgentConfig::load(&config_path)?,
|
|
||||||
(false, true) => {
|
|
||||||
let parser = GraphParser::new(&agent_data_dir);
|
|
||||||
let graph = parser
|
|
||||||
.load_from_file(&graph_path)
|
|
||||||
.with_context(|| format!("Failed to load graph.yaml for agent '{name}'"))?;
|
|
||||||
let config = AgentConfig::from_graph(name, &graph);
|
|
||||||
graph_for_rag = Some(graph);
|
|
||||||
config
|
|
||||||
}
|
|
||||||
(false, false) => bail!(
|
|
||||||
"Agent '{name}' has neither a config.yaml nor a graph.yaml at '{}'",
|
|
||||||
agent_data_dir.display()
|
|
||||||
),
|
|
||||||
};
|
};
|
||||||
let mut functions = Functions::init_agent(name, &agent_config.global_tools)?;
|
let mut functions = Functions::init_agent(name, &agent_config.global_tools)?;
|
||||||
|
|
||||||
@@ -159,16 +138,44 @@ impl Agent {
|
|||||||
.prompt()?;
|
.prompt()?;
|
||||||
}
|
}
|
||||||
if ans {
|
if ans {
|
||||||
let document_paths =
|
let mut document_paths = vec![];
|
||||||
resolve_document_paths(&agent_config.documents, &loaders, &agent_data_dir)?;
|
for path in &agent_config.documents {
|
||||||
|
if is_url(path) {
|
||||||
|
document_paths.push(path.to_string());
|
||||||
|
} else if is_loader_protocol(&loaders, path) {
|
||||||
|
let (protocol, document_path) = path
|
||||||
|
.split_once(':')
|
||||||
|
.with_context(|| "Invalid loader protocol path")?;
|
||||||
|
let resolved_path = resolve_home_dir(document_path);
|
||||||
|
let new_path = if Path::new(&resolved_path).is_relative() {
|
||||||
|
safe_join_path(&agent_data_dir, resolved_path)
|
||||||
|
.ok_or_else(|| anyhow!("Invalid document path: '{path}'"))?
|
||||||
|
} else {
|
||||||
|
PathBuf::from(&resolved_path)
|
||||||
|
};
|
||||||
|
document_paths.push(format!("{}:{}", protocol, new_path.display()));
|
||||||
|
} else if Path::new(&resolve_home_dir(path)).is_relative() {
|
||||||
|
let new_path = safe_join_path(&agent_data_dir, path)
|
||||||
|
.ok_or_else(|| anyhow!("Invalid document path: '{path}'"))?;
|
||||||
|
document_paths.push(new_path.display().to_string())
|
||||||
|
} else {
|
||||||
|
document_paths.push(path.to_string())
|
||||||
|
}
|
||||||
|
}
|
||||||
let key = RagKey::Agent(name.to_string());
|
let key = RagKey::Agent(name.to_string());
|
||||||
let app_clone = app.clone();
|
let app_clone = app.clone();
|
||||||
let rag_path_clone = rag_path.clone();
|
let rag_path_clone = rag_path.clone();
|
||||||
let abort = abort_signal.clone();
|
|
||||||
let rag = app_state
|
let rag = app_state
|
||||||
.rag_cache
|
.rag_cache
|
||||||
.load_with(key, || async move {
|
.load_with(key, || async move {
|
||||||
Rag::init(&app_clone, "rag", &rag_path_clone, &document_paths, abort).await
|
Rag::init(
|
||||||
|
&app_clone,
|
||||||
|
"rag",
|
||||||
|
&rag_path_clone,
|
||||||
|
&document_paths,
|
||||||
|
abort_signal,
|
||||||
|
)
|
||||||
|
.await
|
||||||
})
|
})
|
||||||
.await?;
|
.await?;
|
||||||
Some(rag)
|
Some(rag)
|
||||||
@@ -179,23 +186,6 @@ impl Agent {
|
|||||||
None
|
None
|
||||||
};
|
};
|
||||||
|
|
||||||
let graph_rags = match &graph_for_rag {
|
|
||||||
Some(graph) => {
|
|
||||||
init_graph_rags(
|
|
||||||
app,
|
|
||||||
app_state,
|
|
||||||
name,
|
|
||||||
graph,
|
|
||||||
&agent_data_dir,
|
|
||||||
&loaders,
|
|
||||||
info_flag,
|
|
||||||
abort_signal.clone(),
|
|
||||||
)
|
|
||||||
.await?
|
|
||||||
}
|
|
||||||
None => HashMap::new(),
|
|
||||||
};
|
|
||||||
|
|
||||||
if agent_config.auto_continue {
|
if agent_config.auto_continue {
|
||||||
functions.append_todo_functions();
|
functions.append_todo_functions();
|
||||||
}
|
}
|
||||||
@@ -218,7 +208,6 @@ impl Agent {
|
|||||||
session_dynamic_instructions: None,
|
session_dynamic_instructions: None,
|
||||||
functions,
|
functions,
|
||||||
rag,
|
rag,
|
||||||
graph_rags,
|
|
||||||
model,
|
model,
|
||||||
vault: app_state.vault.clone(),
|
vault: app_state.vault.clone(),
|
||||||
})
|
})
|
||||||
@@ -298,13 +287,10 @@ impl Agent {
|
|||||||
.display()
|
.display()
|
||||||
.to_string()
|
.to_string()
|
||||||
.into();
|
.into();
|
||||||
let config_path = paths::agent_config_file(&self.name);
|
value["config_file"] = paths::agent_config_file(&self.name)
|
||||||
let definition_file = if config_path.exists() {
|
.display()
|
||||||
config_path
|
.to_string()
|
||||||
} else {
|
.into();
|
||||||
paths::agent_graph_file(&self.name)
|
|
||||||
};
|
|
||||||
value["config_file"] = definition_file.display().to_string().into();
|
|
||||||
let data = serde_yaml::to_string(&value)?;
|
let data = serde_yaml::to_string(&value)?;
|
||||||
Ok(data)
|
Ok(data)
|
||||||
}
|
}
|
||||||
@@ -325,10 +311,6 @@ impl Agent {
|
|||||||
self.rag.clone()
|
self.rag.clone()
|
||||||
}
|
}
|
||||||
|
|
||||||
pub fn graph_rag(&self, node_id: &str) -> Option<Arc<Rag>> {
|
|
||||||
self.graph_rags.get(node_id).cloned()
|
|
||||||
}
|
|
||||||
|
|
||||||
pub fn append_mcp_meta_functions(&mut self, mcp_servers: Vec<String>) {
|
pub fn append_mcp_meta_functions(&mut self, mcp_servers: Vec<String>) {
|
||||||
self.functions.append_mcp_meta_functions(mcp_servers);
|
self.functions.append_mcp_meta_functions(mcp_servers);
|
||||||
}
|
}
|
||||||
@@ -433,14 +415,6 @@ impl Agent {
|
|||||||
self.config.max_auto_continues
|
self.config.max_auto_continues
|
||||||
}
|
}
|
||||||
|
|
||||||
pub fn inject_todo_instructions(&self) -> bool {
|
|
||||||
self.config.inject_todo_instructions
|
|
||||||
}
|
|
||||||
|
|
||||||
pub fn continuation_prompt_value(&self) -> Option<String> {
|
|
||||||
self.config.continuation_prompt.clone()
|
|
||||||
}
|
|
||||||
|
|
||||||
pub fn can_spawn_agents(&self) -> bool {
|
pub fn can_spawn_agents(&self) -> bool {
|
||||||
self.config.can_spawn_agents
|
self.config.can_spawn_agents
|
||||||
}
|
}
|
||||||
@@ -465,6 +439,18 @@ impl Agent {
|
|||||||
self.config.escalation_timeout
|
self.config.escalation_timeout
|
||||||
}
|
}
|
||||||
|
|
||||||
|
pub fn continuation_prompt(&self) -> String {
|
||||||
|
self.config.continuation_prompt.clone().unwrap_or_else(|| {
|
||||||
|
formatdoc! {"
|
||||||
|
[SYSTEM REMINDER - TODO CONTINUATION]
|
||||||
|
You have incomplete tasks. Rules:
|
||||||
|
1. BEFORE marking a todo done: verify the work compiles/works. No premature completion.
|
||||||
|
2. If a todo is broad (e.g. \"implement X and implement Y\"): break it into specific subtasks FIRST using todo__add, then work on those.\n\
|
||||||
|
3. Each todo should be atomic and be \"single responsibility\" - completable in one focused action.
|
||||||
|
4. Continue with the next pending item now. Call tools immediately."}
|
||||||
|
})
|
||||||
|
}
|
||||||
|
|
||||||
pub fn compression_threshold(&self) -> Option<usize> {
|
pub fn compression_threshold(&self) -> Option<usize> {
|
||||||
self.config.compression_threshold
|
self.config.compression_threshold
|
||||||
}
|
}
|
||||||
@@ -668,25 +654,6 @@ impl AgentConfig {
|
|||||||
Ok(agent_config)
|
Ok(agent_config)
|
||||||
}
|
}
|
||||||
|
|
||||||
pub fn from_graph(dir_name: &str, graph: &Graph) -> Self {
|
|
||||||
AgentConfig {
|
|
||||||
name: dir_name.to_string(),
|
|
||||||
model_id: graph.model.clone(),
|
|
||||||
temperature: graph.temperature,
|
|
||||||
top_p: graph.top_p,
|
|
||||||
description: graph.description.clone(),
|
|
||||||
global_tools: graph.global_tools.clone(),
|
|
||||||
mcp_servers: graph.mcp_servers.clone(),
|
|
||||||
conversation_starters: graph.conversation_starters.clone(),
|
|
||||||
variables: graph.variables.clone(),
|
|
||||||
can_spawn_agents: graph.has_agent_node(),
|
|
||||||
max_concurrent_agents: default_max_concurrent_agents(),
|
|
||||||
max_agent_depth: default_max_agent_depth(),
|
|
||||||
escalation_timeout: default_escalation_timeout(),
|
|
||||||
..AgentConfig::default()
|
|
||||||
}
|
|
||||||
}
|
|
||||||
|
|
||||||
fn load_envs(&mut self, app: &AppConfig) {
|
fn load_envs(&mut self, app: &AppConfig) {
|
||||||
let name = &self.name;
|
let name = &self.name;
|
||||||
let with_prefix = |v: &str| normalize_env_name(&format!("{name}_{v}"));
|
let with_prefix = |v: &str| normalize_env_name(&format!("{name}_{v}"));
|
||||||
@@ -783,136 +750,6 @@ pub struct AgentVariable {
|
|||||||
pub value: String,
|
pub value: String,
|
||||||
}
|
}
|
||||||
|
|
||||||
fn resolve_document_paths(
|
|
||||||
documents: &[String],
|
|
||||||
loaders: &HashMap<String, String>,
|
|
||||||
agent_data_dir: &Path,
|
|
||||||
) -> Result<Vec<String>> {
|
|
||||||
let mut document_paths = vec![];
|
|
||||||
for path in documents {
|
|
||||||
if is_url(path) {
|
|
||||||
document_paths.push(path.to_string());
|
|
||||||
} else if is_loader_protocol(loaders, path) {
|
|
||||||
let (protocol, document_path) = path
|
|
||||||
.split_once(':')
|
|
||||||
.with_context(|| "Invalid loader protocol path")?;
|
|
||||||
let resolved_path = resolve_home_dir(document_path);
|
|
||||||
let new_path = if Path::new(&resolved_path).is_relative() {
|
|
||||||
safe_join_path(agent_data_dir, resolved_path)
|
|
||||||
.ok_or_else(|| anyhow!("Invalid document path: '{path}'"))?
|
|
||||||
} else {
|
|
||||||
PathBuf::from(&resolved_path)
|
|
||||||
};
|
|
||||||
|
|
||||||
document_paths.push(format!("{}:{}", protocol, new_path.display()));
|
|
||||||
} else if Path::new(&resolve_home_dir(path)).is_relative() {
|
|
||||||
let new_path = safe_join_path(agent_data_dir, path)
|
|
||||||
.ok_or_else(|| anyhow!("Invalid document path: '{path}'"))?;
|
|
||||||
document_paths.push(new_path.display().to_string())
|
|
||||||
} else {
|
|
||||||
document_paths.push(path.to_string())
|
|
||||||
}
|
|
||||||
}
|
|
||||||
Ok(document_paths)
|
|
||||||
}
|
|
||||||
|
|
||||||
#[allow(clippy::too_many_arguments)]
|
|
||||||
async fn init_graph_rags(
|
|
||||||
app: &AppConfig,
|
|
||||||
app_state: &AppState,
|
|
||||||
agent_name: &str,
|
|
||||||
graph: &Graph,
|
|
||||||
agent_data_dir: &Path,
|
|
||||||
loaders: &HashMap<String, String>,
|
|
||||||
info_flag: bool,
|
|
||||||
abort_signal: AbortSignal,
|
|
||||||
) -> Result<HashMap<String, Arc<Rag>>> {
|
|
||||||
let mut rags = HashMap::new();
|
|
||||||
if info_flag {
|
|
||||||
return Ok(rags);
|
|
||||||
}
|
|
||||||
|
|
||||||
for (node_id, node) in &graph.nodes {
|
|
||||||
let NodeType::Rag(rag_node) = &node.node_type else {
|
|
||||||
continue;
|
|
||||||
};
|
|
||||||
let rag_path = paths::agent_rag_file(agent_name, node_id);
|
|
||||||
let key = RagKey::GraphNode {
|
|
||||||
agent: agent_name.to_string(),
|
|
||||||
node: node_id.clone(),
|
|
||||||
};
|
|
||||||
let rag = if rag_path.exists() {
|
|
||||||
let app_clone = app.clone();
|
|
||||||
let path_clone = rag_path.clone();
|
|
||||||
let name_clone = node_id.clone();
|
|
||||||
app_state
|
|
||||||
.rag_cache
|
|
||||||
.load_with(key, || async move {
|
|
||||||
Rag::load(&app_clone, &name_clone, &path_clone)
|
|
||||||
})
|
|
||||||
.await?
|
|
||||||
} else {
|
|
||||||
let config = RagInitConfig {
|
|
||||||
embedding_model: rag_node.embedding_model.clone(),
|
|
||||||
chunk_size: rag_node.chunk_size,
|
|
||||||
chunk_overlap: rag_node.chunk_overlap,
|
|
||||||
reranker_model: rag_node.reranker_model.clone(),
|
|
||||||
top_k: rag_node.top_k,
|
|
||||||
batch_size: rag_node.batch_size,
|
|
||||||
};
|
|
||||||
let fully_specified = config.embedding_model.is_some()
|
|
||||||
&& config.chunk_size.is_some()
|
|
||||||
&& config.chunk_overlap.is_some();
|
|
||||||
if !fully_specified {
|
|
||||||
if !*IS_STDOUT_TERMINAL {
|
|
||||||
bail!(
|
|
||||||
"Agent '{agent_name}' requires RAG for rag node '{node_id}', but its \
|
|
||||||
knowledge base is not built and the node does not fully specify how \
|
|
||||||
to build it. Set `embedding_model`, `chunk_size`, and `chunk_overlap` \
|
|
||||||
on the node, or run the agent once interactively."
|
|
||||||
);
|
|
||||||
}
|
|
||||||
|
|
||||||
let ans = Confirm::new(&format!(
|
|
||||||
"Initialize RAG knowledge base for rag node '{node_id}'?"
|
|
||||||
))
|
|
||||||
.with_default(true)
|
|
||||||
.prompt()?;
|
|
||||||
|
|
||||||
if !ans {
|
|
||||||
bail!(
|
|
||||||
"Agent '{agent_name}' has rag node '{node_id}' but its RAG was not \
|
|
||||||
initialized. RAG initialization is required for this agent."
|
|
||||||
);
|
|
||||||
}
|
|
||||||
}
|
|
||||||
|
|
||||||
let document_paths =
|
|
||||||
resolve_document_paths(&rag_node.documents, loaders, agent_data_dir)?;
|
|
||||||
let app_clone = app.clone();
|
|
||||||
let path_clone = rag_path.clone();
|
|
||||||
let name_clone = node_id.clone();
|
|
||||||
let abort = abort_signal.clone();
|
|
||||||
app_state
|
|
||||||
.rag_cache
|
|
||||||
.load_with(key, || async move {
|
|
||||||
Rag::init_with_config(
|
|
||||||
&app_clone,
|
|
||||||
&name_clone,
|
|
||||||
&path_clone,
|
|
||||||
&document_paths,
|
|
||||||
&config,
|
|
||||||
abort,
|
|
||||||
)
|
|
||||||
.await
|
|
||||||
})
|
|
||||||
.await?
|
|
||||||
};
|
|
||||||
rags.insert(node_id.clone(), rag);
|
|
||||||
}
|
|
||||||
Ok(rags)
|
|
||||||
}
|
|
||||||
|
|
||||||
pub fn list_agents() -> Vec<String> {
|
pub fn list_agents() -> Vec<String> {
|
||||||
let agents_data_dir = paths::agents_data_dir();
|
let agents_data_dir = paths::agents_data_dir();
|
||||||
if !agents_data_dir.exists() {
|
if !agents_data_dir.exists() {
|
||||||
@@ -1039,88 +876,4 @@ variables:
|
|||||||
assert!(config.inject_todo_instructions);
|
assert!(config.inject_todo_instructions);
|
||||||
assert!(config.inject_spawn_instructions);
|
assert!(config.inject_spawn_instructions);
|
||||||
}
|
}
|
||||||
|
|
||||||
#[test]
|
|
||||||
fn from_graph_maps_agent_level_fields() {
|
|
||||||
let yaml = formatdoc! {r#"
|
|
||||||
name: graph_name_ignored
|
|
||||||
description: A graph agent
|
|
||||||
model: claude:claude-sonnet-4-6
|
|
||||||
temperature: 0.3
|
|
||||||
top_p: 0.8
|
|
||||||
global_tools:
|
|
||||||
- fetch_pdf.sh
|
|
||||||
mcp_servers:
|
|
||||||
- pubmed-search
|
|
||||||
conversation_starters:
|
|
||||||
- "Start here"
|
|
||||||
start: e
|
|
||||||
nodes:
|
|
||||||
e:
|
|
||||||
id: e
|
|
||||||
type: end
|
|
||||||
output: done
|
|
||||||
"#};
|
|
||||||
let graph: Graph = serde_yaml::from_str(&yaml).unwrap();
|
|
||||||
|
|
||||||
let config = AgentConfig::from_graph("my-agent-dir", &graph);
|
|
||||||
|
|
||||||
assert_eq!(config.name, "my-agent-dir");
|
|
||||||
assert_eq!(config.description, "A graph agent");
|
|
||||||
assert_eq!(config.model_id.as_deref(), Some("claude:claude-sonnet-4-6"));
|
|
||||||
assert_eq!(config.temperature, Some(0.3));
|
|
||||||
assert_eq!(config.top_p, Some(0.8));
|
|
||||||
assert_eq!(config.global_tools, vec!["fetch_pdf.sh"]);
|
|
||||||
assert_eq!(config.mcp_servers, vec!["pubmed-search"]);
|
|
||||||
assert_eq!(config.conversation_starters, vec!["Start here"]);
|
|
||||||
}
|
|
||||||
|
|
||||||
#[test]
|
|
||||||
fn from_graph_derives_can_spawn_agents_from_agent_nodes() {
|
|
||||||
let with_agent = formatdoc! {r#"
|
|
||||||
name: g
|
|
||||||
start: a
|
|
||||||
nodes:
|
|
||||||
a:
|
|
||||||
id: a
|
|
||||||
type: agent
|
|
||||||
agent: helper
|
|
||||||
prompt: hi
|
|
||||||
next: e
|
|
||||||
e:
|
|
||||||
id: e
|
|
||||||
type: end
|
|
||||||
output: done
|
|
||||||
"#};
|
|
||||||
let graph: Graph = serde_yaml::from_str(&with_agent).unwrap();
|
|
||||||
assert!(AgentConfig::from_graph("d", &graph).can_spawn_agents);
|
|
||||||
|
|
||||||
let no_agent =
|
|
||||||
"name: g\nstart: x\nnodes:\n x:\n id: x\n type: end\n output: ok\n";
|
|
||||||
let graph: Graph = serde_yaml::from_str(no_agent).unwrap();
|
|
||||||
assert!(!AgentConfig::from_graph("d", &graph).can_spawn_agents);
|
|
||||||
}
|
|
||||||
|
|
||||||
#[test]
|
|
||||||
fn from_graph_keeps_defaults_for_llm_loop_fields() {
|
|
||||||
let yaml = "name: g\nstart: x\nnodes:\n x:\n id: x\n type: end\n output: ok\n";
|
|
||||||
let graph: Graph = serde_yaml::from_str(yaml).unwrap();
|
|
||||||
|
|
||||||
let config = AgentConfig::from_graph("d", &graph);
|
|
||||||
|
|
||||||
assert!(!config.auto_continue);
|
|
||||||
assert!(config.instructions.is_empty());
|
|
||||||
assert!(config.documents.is_empty());
|
|
||||||
assert!(!config.inject_todo_instructions);
|
|
||||||
assert!(!config.inject_spawn_instructions);
|
|
||||||
assert_eq!(config.max_auto_continues, 0);
|
|
||||||
assert_eq!(config.summarization_threshold, 0);
|
|
||||||
|
|
||||||
assert_eq!(
|
|
||||||
config.max_concurrent_agents,
|
|
||||||
default_max_concurrent_agents()
|
|
||||||
);
|
|
||||||
assert_eq!(config.max_agent_depth, default_max_agent_depth());
|
|
||||||
assert_eq!(config.escalation_timeout, default_escalation_timeout());
|
|
||||||
}
|
|
||||||
}
|
}
|
||||||
|
|||||||
@@ -39,11 +39,6 @@ pub struct AppConfig {
|
|||||||
pub mapping_mcp_servers: IndexMap<String, String>,
|
pub mapping_mcp_servers: IndexMap<String, String>,
|
||||||
pub enabled_mcp_servers: Option<String>,
|
pub enabled_mcp_servers: Option<String>,
|
||||||
|
|
||||||
pub auto_continue: bool,
|
|
||||||
pub max_auto_continues: usize,
|
|
||||||
pub inject_todo_instructions: bool,
|
|
||||||
pub continuation_prompt: Option<String>,
|
|
||||||
|
|
||||||
pub repl_prelude: Option<String>,
|
pub repl_prelude: Option<String>,
|
||||||
pub cmd_prelude: Option<String>,
|
pub cmd_prelude: Option<String>,
|
||||||
pub agent_session: Option<String>,
|
pub agent_session: Option<String>,
|
||||||
@@ -100,11 +95,6 @@ impl Default for AppConfig {
|
|||||||
mapping_mcp_servers: Default::default(),
|
mapping_mcp_servers: Default::default(),
|
||||||
enabled_mcp_servers: None,
|
enabled_mcp_servers: None,
|
||||||
|
|
||||||
auto_continue: false,
|
|
||||||
max_auto_continues: 10,
|
|
||||||
inject_todo_instructions: true,
|
|
||||||
continuation_prompt: None,
|
|
||||||
|
|
||||||
repl_prelude: None,
|
repl_prelude: None,
|
||||||
cmd_prelude: None,
|
cmd_prelude: None,
|
||||||
agent_session: None,
|
agent_session: None,
|
||||||
@@ -162,11 +152,6 @@ impl AppConfig {
|
|||||||
mapping_mcp_servers: config.mapping_mcp_servers,
|
mapping_mcp_servers: config.mapping_mcp_servers,
|
||||||
enabled_mcp_servers: config.enabled_mcp_servers,
|
enabled_mcp_servers: config.enabled_mcp_servers,
|
||||||
|
|
||||||
auto_continue: config.auto_continue,
|
|
||||||
max_auto_continues: config.max_auto_continues,
|
|
||||||
inject_todo_instructions: config.inject_todo_instructions,
|
|
||||||
continuation_prompt: config.continuation_prompt,
|
|
||||||
|
|
||||||
repl_prelude: config.repl_prelude,
|
repl_prelude: config.repl_prelude,
|
||||||
cmd_prelude: config.cmd_prelude,
|
cmd_prelude: config.cmd_prelude,
|
||||||
agent_session: config.agent_session,
|
agent_session: config.agent_session,
|
||||||
|
|||||||
File diff suppressed because it is too large
Load Diff
@@ -85,7 +85,7 @@ impl Macro {
|
|||||||
Ok(value)
|
Ok(value)
|
||||||
}
|
}
|
||||||
|
|
||||||
pub fn install_macros(force: bool) -> Result<()> {
|
pub fn install_macros() -> Result<()> {
|
||||||
info!(
|
info!(
|
||||||
"Installing built-in macros in {}",
|
"Installing built-in macros in {}",
|
||||||
paths::macros_dir().display()
|
paths::macros_dir().display()
|
||||||
@@ -98,7 +98,7 @@ impl Macro {
|
|||||||
let content = unsafe { std::str::from_utf8_unchecked(&embedded_file.data) };
|
let content = unsafe { std::str::from_utf8_unchecked(&embedded_file.data) };
|
||||||
let file_path = paths::macros_dir().join(file.as_ref());
|
let file_path = paths::macros_dir().join(file.as_ref());
|
||||||
|
|
||||||
if file_path.exists() && !force {
|
if file_path.exists() {
|
||||||
debug!(
|
debug!(
|
||||||
"Macro file already exists, skipping: {}",
|
"Macro file already exists, skipping: {}",
|
||||||
file_path.display()
|
file_path.display()
|
||||||
|
|||||||
@@ -109,13 +109,12 @@ impl McpFactory {
|
|||||||
mod tests {
|
mod tests {
|
||||||
use super::*;
|
use super::*;
|
||||||
use crate::mcp::{JsonField, McpServer, McpTransportType};
|
use crate::mcp::{JsonField, McpServer, McpTransportType};
|
||||||
use indexmap::IndexMap;
|
|
||||||
use std::collections::HashMap;
|
use std::collections::HashMap;
|
||||||
|
|
||||||
fn stdio_spec(
|
fn stdio_spec(
|
||||||
command: &str,
|
command: &str,
|
||||||
args: Option<Vec<String>>,
|
args: Option<Vec<String>>,
|
||||||
env: Option<IndexMap<String, JsonField>>,
|
env: Option<HashMap<String, JsonField>>,
|
||||||
) -> McpServer {
|
) -> McpServer {
|
||||||
McpServer {
|
McpServer {
|
||||||
transport_type: McpTransportType::Stdio,
|
transport_type: McpTransportType::Stdio,
|
||||||
@@ -131,7 +130,7 @@ mod tests {
|
|||||||
fn remote_spec(
|
fn remote_spec(
|
||||||
transport: McpTransportType,
|
transport: McpTransportType,
|
||||||
url: &str,
|
url: &str,
|
||||||
headers: Option<IndexMap<String, String>>,
|
headers: Option<HashMap<String, String>>,
|
||||||
) -> McpServer {
|
) -> McpServer {
|
||||||
McpServer {
|
McpServer {
|
||||||
transport_type: transport,
|
transport_type: transport,
|
||||||
@@ -146,7 +145,7 @@ mod tests {
|
|||||||
|
|
||||||
#[test]
|
#[test]
|
||||||
fn key_from_stdio_spec_captures_command_args_env() {
|
fn key_from_stdio_spec_captures_command_args_env() {
|
||||||
let mut env = IndexMap::new();
|
let mut env = HashMap::new();
|
||||||
env.insert("TOKEN".into(), JsonField::Str("abc".into()));
|
env.insert("TOKEN".into(), JsonField::Str("abc".into()));
|
||||||
let spec = stdio_spec("npx", Some(vec!["-y".into(), "server".into()]), Some(env));
|
let spec = stdio_spec("npx", Some(vec!["-y".into(), "server".into()]), Some(env));
|
||||||
let key = McpServerKey::from_spec("my-server", &spec);
|
let key = McpServerKey::from_spec("my-server", &spec);
|
||||||
@@ -164,7 +163,7 @@ mod tests {
|
|||||||
|
|
||||||
#[test]
|
#[test]
|
||||||
fn key_from_stdio_spec_sorts_args_and_env() {
|
fn key_from_stdio_spec_sorts_args_and_env() {
|
||||||
let mut env = IndexMap::new();
|
let mut env = HashMap::new();
|
||||||
env.insert("Z_VAR".into(), JsonField::Str("z".into()));
|
env.insert("Z_VAR".into(), JsonField::Str("z".into()));
|
||||||
env.insert("A_VAR".into(), JsonField::Int(42));
|
env.insert("A_VAR".into(), JsonField::Int(42));
|
||||||
let spec = stdio_spec(
|
let spec = stdio_spec(
|
||||||
@@ -223,7 +222,7 @@ mod tests {
|
|||||||
|
|
||||||
#[test]
|
#[test]
|
||||||
fn key_from_remote_sse_spec_with_sorted_headers() {
|
fn key_from_remote_sse_spec_with_sorted_headers() {
|
||||||
let mut hdrs = IndexMap::new();
|
let mut hdrs = HashMap::new();
|
||||||
hdrs.insert("Z-Key".into(), "z-val".into());
|
hdrs.insert("Z-Key".into(), "z-val".into());
|
||||||
hdrs.insert("A-Key".into(), "a-val".into());
|
hdrs.insert("A-Key".into(), "a-val".into());
|
||||||
let spec = remote_spec(McpTransportType::Sse, "http://sse.example.com", Some(hdrs));
|
let spec = remote_spec(McpTransportType::Sse, "http://sse.example.com", Some(hdrs));
|
||||||
@@ -265,7 +264,7 @@ mod tests {
|
|||||||
|
|
||||||
#[test]
|
#[test]
|
||||||
fn key_env_bool_and_int_coerce_to_string() {
|
fn key_env_bool_and_int_coerce_to_string() {
|
||||||
let mut env = IndexMap::new();
|
let mut env = HashMap::new();
|
||||||
env.insert("FLAG".into(), JsonField::Bool(true));
|
env.insert("FLAG".into(), JsonField::Bool(true));
|
||||||
env.insert("PORT".into(), JsonField::Int(3000));
|
env.insert("PORT".into(), JsonField::Int(3000));
|
||||||
let spec = stdio_spec("cmd", None, Some(env));
|
let spec = stdio_spec("cmd", None, Some(env));
|
||||||
|
|||||||
+5
-152
@@ -2,7 +2,6 @@ mod agent;
|
|||||||
mod app_config;
|
mod app_config;
|
||||||
mod app_state;
|
mod app_state;
|
||||||
mod input;
|
mod input;
|
||||||
mod install_remote;
|
|
||||||
mod macros;
|
mod macros;
|
||||||
mod mcp_factory;
|
mod mcp_factory;
|
||||||
pub(crate) mod paths;
|
pub(crate) mod paths;
|
||||||
@@ -13,24 +12,19 @@ mod role;
|
|||||||
mod session;
|
mod session;
|
||||||
pub(crate) mod todo;
|
pub(crate) mod todo;
|
||||||
mod tool_scope;
|
mod tool_scope;
|
||||||
mod update;
|
|
||||||
|
|
||||||
pub use self::agent::{
|
pub use self::agent::{Agent, AgentVariables, complete_agent_variables, list_agents};
|
||||||
Agent, AgentVariable, AgentVariables, complete_agent_variables, list_agents,
|
|
||||||
};
|
|
||||||
#[allow(unused_imports)]
|
#[allow(unused_imports)]
|
||||||
pub use self::app_config::AppConfig;
|
pub use self::app_config::AppConfig;
|
||||||
#[allow(unused_imports)]
|
#[allow(unused_imports)]
|
||||||
pub use self::app_state::AppState;
|
pub use self::app_state::AppState;
|
||||||
pub use self::input::Input;
|
pub use self::input::Input;
|
||||||
pub use self::install_remote::{install_remote, install_remote_from_repl_args};
|
|
||||||
#[allow(unused_imports)]
|
#[allow(unused_imports)]
|
||||||
pub use self::request_context::{RenderMode, RequestContext};
|
pub use self::request_context::RequestContext;
|
||||||
pub use self::role::{
|
pub use self::role::{
|
||||||
CODE_ROLE, CREATE_TITLE_ROLE, EXPLAIN_SHELL_ROLE, Role, RoleLike, SHELL_ROLE,
|
CODE_ROLE, CREATE_TITLE_ROLE, EXPLAIN_SHELL_ROLE, Role, RoleLike, SHELL_ROLE,
|
||||||
};
|
};
|
||||||
use self::session::Session;
|
use self::session::Session;
|
||||||
pub use self::update::run_self_update;
|
|
||||||
use crate::client::{
|
use crate::client::{
|
||||||
ClientConfig, MessageContentToolCalls, Model, ModelType, OPENAI_COMPATIBLE_PROVIDERS,
|
ClientConfig, MessageContentToolCalls, Model, ModelType, OPENAI_COMPATIBLE_PROVIDERS,
|
||||||
ProviderModels, create_client_config, list_client_types,
|
ProviderModels, create_client_config, list_client_types,
|
||||||
@@ -72,7 +66,6 @@ const DARK_THEME: &[u8] = include_bytes!("../../assets/monokai-extended.theme.bi
|
|||||||
const LIGHT_THEME: &[u8] = include_bytes!("../../assets/monokai-extended-light.theme.bin");
|
const LIGHT_THEME: &[u8] = include_bytes!("../../assets/monokai-extended-light.theme.bin");
|
||||||
|
|
||||||
const CONFIG_FILE_NAME: &str = "config.yaml";
|
const CONFIG_FILE_NAME: &str = "config.yaml";
|
||||||
const AGENT_GRAPH_FILE_NAME: &str = "graph.yaml";
|
|
||||||
const ROLES_DIR_NAME: &str = "roles";
|
const ROLES_DIR_NAME: &str = "roles";
|
||||||
const MACROS_DIR_NAME: &str = "macros";
|
const MACROS_DIR_NAME: &str = "macros";
|
||||||
const ENV_FILE_NAME: &str = ".env";
|
const ENV_FILE_NAME: &str = ".env";
|
||||||
@@ -86,26 +79,6 @@ const GLOBAL_TOOLS_DIR_NAME: &str = "tools";
|
|||||||
const GLOBAL_TOOLS_UTILS_DIR_NAME: &str = "utils";
|
const GLOBAL_TOOLS_UTILS_DIR_NAME: &str = "utils";
|
||||||
const BASH_PROMPT_UTILS_FILE_NAME: &str = "prompt-utils.sh";
|
const BASH_PROMPT_UTILS_FILE_NAME: &str = "prompt-utils.sh";
|
||||||
const MCP_FILE_NAME: &str = "mcp.json";
|
const MCP_FILE_NAME: &str = "mcp.json";
|
||||||
const DEFAULT_VISIBLE_TOOLS: [&str; 18] = [
|
|
||||||
"execute_command.sh",
|
|
||||||
"execute_py_code.py",
|
|
||||||
"execute_sql_code.sh",
|
|
||||||
"fetch_url_via_curl.sh",
|
|
||||||
"fs_cat.sh",
|
|
||||||
"fs_glob.sh",
|
|
||||||
"fs_grep.sh",
|
|
||||||
"fs_ls.sh",
|
|
||||||
"fs_mkdir.sh",
|
|
||||||
"fs_patch.sh",
|
|
||||||
"fs_read.sh",
|
|
||||||
"fs_rm.sh",
|
|
||||||
"fs_write.sh",
|
|
||||||
"get_current_time.sh",
|
|
||||||
"get_current_weather.sh",
|
|
||||||
"search_wikipedia.sh",
|
|
||||||
"search_arxiv.sh",
|
|
||||||
"web_search_loki.sh",
|
|
||||||
];
|
|
||||||
|
|
||||||
const CLIENTS_FIELD: &str = "clients";
|
const CLIENTS_FIELD: &str = "clients";
|
||||||
|
|
||||||
@@ -148,11 +121,6 @@ pub struct Config {
|
|||||||
pub mapping_mcp_servers: IndexMap<String, String>,
|
pub mapping_mcp_servers: IndexMap<String, String>,
|
||||||
pub enabled_mcp_servers: Option<String>,
|
pub enabled_mcp_servers: Option<String>,
|
||||||
|
|
||||||
pub auto_continue: bool,
|
|
||||||
pub max_auto_continues: usize,
|
|
||||||
pub inject_todo_instructions: bool,
|
|
||||||
pub continuation_prompt: Option<String>,
|
|
||||||
|
|
||||||
pub repl_prelude: Option<String>,
|
pub repl_prelude: Option<String>,
|
||||||
pub cmd_prelude: Option<String>,
|
pub cmd_prelude: Option<String>,
|
||||||
pub agent_session: Option<String>,
|
pub agent_session: Option<String>,
|
||||||
@@ -209,11 +177,6 @@ impl Default for Config {
|
|||||||
mapping_mcp_servers: Default::default(),
|
mapping_mcp_servers: Default::default(),
|
||||||
enabled_mcp_servers: None,
|
enabled_mcp_servers: None,
|
||||||
|
|
||||||
auto_continue: false,
|
|
||||||
max_auto_continues: 10,
|
|
||||||
inject_todo_instructions: true,
|
|
||||||
continuation_prompt: None,
|
|
||||||
|
|
||||||
repl_prelude: None,
|
repl_prelude: None,
|
||||||
cmd_prelude: None,
|
cmd_prelude: None,
|
||||||
agent_session: None,
|
agent_session: None,
|
||||||
@@ -247,110 +210,12 @@ impl Default for Config {
|
|||||||
}
|
}
|
||||||
|
|
||||||
pub fn install_builtins() -> Result<()> {
|
pub fn install_builtins() -> Result<()> {
|
||||||
Functions::install_builtin_global_tools(false)?;
|
Functions::install_builtin_global_tools()?;
|
||||||
Agent::install_builtin_agents(false)?;
|
Agent::install_builtin_agents()?;
|
||||||
Macro::install_macros(false)?;
|
Macro::install_macros()?;
|
||||||
Ok(())
|
Ok(())
|
||||||
}
|
}
|
||||||
|
|
||||||
#[derive(Debug, Clone, Copy, PartialEq, Eq, clap::ValueEnum)]
|
|
||||||
pub enum AssetCategory {
|
|
||||||
Agents,
|
|
||||||
Macros,
|
|
||||||
Functions,
|
|
||||||
#[value(name = "mcp_config")]
|
|
||||||
McpConfig,
|
|
||||||
}
|
|
||||||
|
|
||||||
impl AssetCategory {
|
|
||||||
pub const NAMES: [&'static str; 4] = ["agents", "macros", "functions", "mcp_config"];
|
|
||||||
|
|
||||||
pub fn parse(name: &str) -> Option<Self> {
|
|
||||||
match name {
|
|
||||||
"agents" => Some(Self::Agents),
|
|
||||||
"macros" => Some(Self::Macros),
|
|
||||||
"functions" => Some(Self::Functions),
|
|
||||||
"mcp_config" => Some(Self::McpConfig),
|
|
||||||
_ => None,
|
|
||||||
}
|
|
||||||
}
|
|
||||||
}
|
|
||||||
|
|
||||||
#[derive(Debug, Clone, Copy, PartialEq, Eq, clap::ValueEnum)]
|
|
||||||
pub enum InstallFilter {
|
|
||||||
Agents,
|
|
||||||
Roles,
|
|
||||||
Macros,
|
|
||||||
Functions,
|
|
||||||
#[value(name = "mcp_config")]
|
|
||||||
McpConfig,
|
|
||||||
}
|
|
||||||
|
|
||||||
impl InstallFilter {
|
|
||||||
pub const NAMES: [&'static str; 5] = ["agents", "roles", "macros", "functions", "mcp_config"];
|
|
||||||
|
|
||||||
pub fn parse(name: &str) -> Option<Self> {
|
|
||||||
match name {
|
|
||||||
"agents" => Some(Self::Agents),
|
|
||||||
"roles" => Some(Self::Roles),
|
|
||||||
"macros" => Some(Self::Macros),
|
|
||||||
"functions" => Some(Self::Functions),
|
|
||||||
"mcp_config" => Some(Self::McpConfig),
|
|
||||||
_ => None,
|
|
||||||
}
|
|
||||||
}
|
|
||||||
}
|
|
||||||
|
|
||||||
pub fn install_assets(category: AssetCategory) -> Result<()> {
|
|
||||||
let (label, target) = match category {
|
|
||||||
AssetCategory::Agents => ("agents", paths::agents_data_dir()),
|
|
||||||
AssetCategory::Macros => ("macros", paths::macros_dir()),
|
|
||||||
AssetCategory::Functions => ("functions", paths::functions_dir()),
|
|
||||||
AssetCategory::McpConfig => ("MCP config", paths::mcp_config_file()),
|
|
||||||
};
|
|
||||||
|
|
||||||
if !confirm_asset_overwrite(category, label, &target)? {
|
|
||||||
println!("Aborted. No files were changed.");
|
|
||||||
return Ok(());
|
|
||||||
}
|
|
||||||
|
|
||||||
match category {
|
|
||||||
AssetCategory::Agents => Agent::install_builtin_agents(true)?,
|
|
||||||
AssetCategory::Macros => Macro::install_macros(true)?,
|
|
||||||
AssetCategory::Functions => Functions::install_builtin_global_tools(true)?,
|
|
||||||
AssetCategory::McpConfig => Functions::install_mcp_config()?,
|
|
||||||
}
|
|
||||||
|
|
||||||
println!("Reinstalled bundled {label} ({})", target.display());
|
|
||||||
|
|
||||||
Ok(())
|
|
||||||
}
|
|
||||||
|
|
||||||
fn confirm_asset_overwrite(category: AssetCategory, label: &str, target: &Path) -> Result<bool> {
|
|
||||||
if !*IS_STDOUT_TERMINAL {
|
|
||||||
return Ok(true);
|
|
||||||
}
|
|
||||||
let body = match category {
|
|
||||||
AssetCategory::McpConfig => format!(
|
|
||||||
"This replaces your MCP server configuration at {} with this \
|
|
||||||
build's bundled template. Your configured MCP servers (and any \
|
|
||||||
custom secret references they contain) will be lost.",
|
|
||||||
target.display()
|
|
||||||
),
|
|
||||||
_ => format!(
|
|
||||||
"Reinstalling bundled {label} overwrites every bundled {label} in \
|
|
||||||
{} with this build's packaged versions. Local changes to bundled \
|
|
||||||
{label} will be lost; {label} you created yourself are left \
|
|
||||||
untouched.",
|
|
||||||
target.display()
|
|
||||||
),
|
|
||||||
};
|
|
||||||
let prompt = format!("{} {body}\nContinue? [y/N] ", warning_text("WARNING:"));
|
|
||||||
let answer = read_single_key(&['y', 'Y', 'n', 'N'], 'n', &prompt)?;
|
|
||||||
|
|
||||||
Ok(matches!(answer, 'y' | 'Y'))
|
|
||||||
}
|
|
||||||
|
|
||||||
pub fn default_sessions_dir() -> PathBuf {
|
pub fn default_sessions_dir() -> PathBuf {
|
||||||
match env::var(get_env_name("sessions_dir")) {
|
match env::var(get_env_name("sessions_dir")) {
|
||||||
Ok(value) => PathBuf::from(value),
|
Ok(value) => PathBuf::from(value),
|
||||||
@@ -609,18 +474,6 @@ pub async fn create_config_file(config_path: &Path) -> Result<()> {
|
|||||||
let (model, clients_config) = create_client_config(client, &vault).await?;
|
let (model, clients_config) = create_client_config(client, &vault).await?;
|
||||||
config["model"] = model.into();
|
config["model"] = model.into();
|
||||||
config["vault_password_file"] = vault.password_file()?.display().to_string().into();
|
config["vault_password_file"] = vault.password_file()?.display().to_string().into();
|
||||||
config["stream"] = json!(true);
|
|
||||||
config["save"] = json!(true);
|
|
||||||
config["keybindings"] = json!("vi");
|
|
||||||
config["wrap"] = json!("auto");
|
|
||||||
config["wrap_code"] = json!(false);
|
|
||||||
config["function_calling_support"] = json!(true);
|
|
||||||
config["enabled_tools"] = json!(null);
|
|
||||||
config["visible_tools"] = json!(DEFAULT_VISIBLE_TOOLS);
|
|
||||||
config["mcp_server_support"] = json!(true);
|
|
||||||
config["enabled_mcp_servers"] = json!(null);
|
|
||||||
config["highlight"] = json!(true);
|
|
||||||
config["light_theme"] = json!(false);
|
|
||||||
config[CLIENTS_FIELD] = clients_config;
|
config[CLIENTS_FIELD] = clients_config;
|
||||||
|
|
||||||
let config_data = serde_yaml::to_string(&config).with_context(|| "Failed to create config")?;
|
let config_data = serde_yaml::to_string(&config).with_context(|| "Failed to create config")?;
|
||||||
|
|||||||
+3
-8
@@ -1,9 +1,8 @@
|
|||||||
use super::role::Role;
|
use super::role::Role;
|
||||||
use super::{
|
use super::{
|
||||||
AGENT_GRAPH_FILE_NAME, AGENTS_DIR_NAME, BASH_PROMPT_UTILS_FILE_NAME, CONFIG_FILE_NAME,
|
AGENTS_DIR_NAME, BASH_PROMPT_UTILS_FILE_NAME, CONFIG_FILE_NAME, ENV_FILE_NAME,
|
||||||
ENV_FILE_NAME, FUNCTIONS_BIN_DIR_NAME, FUNCTIONS_DIR_NAME, GLOBAL_TOOLS_DIR_NAME,
|
FUNCTIONS_BIN_DIR_NAME, FUNCTIONS_DIR_NAME, GLOBAL_TOOLS_DIR_NAME, GLOBAL_TOOLS_UTILS_DIR_NAME,
|
||||||
GLOBAL_TOOLS_UTILS_DIR_NAME, MACROS_DIR_NAME, MCP_FILE_NAME, ModelsOverride, RAGS_DIR_NAME,
|
MACROS_DIR_NAME, MCP_FILE_NAME, ModelsOverride, RAGS_DIR_NAME, ROLES_DIR_NAME,
|
||||||
ROLES_DIR_NAME,
|
|
||||||
};
|
};
|
||||||
use crate::client::ProviderModels;
|
use crate::client::ProviderModels;
|
||||||
use crate::utils::{get_env_name, list_file_names, normalize_env_name};
|
use crate::utils::{get_env_name, list_file_names, normalize_env_name};
|
||||||
@@ -128,10 +127,6 @@ pub fn agent_data_dir(name: &str) -> PathBuf {
|
|||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
pub fn agent_graph_file(agent_name: &str) -> PathBuf {
|
|
||||||
agent_data_dir(agent_name).join(AGENT_GRAPH_FILE_NAME)
|
|
||||||
}
|
|
||||||
|
|
||||||
pub fn agent_config_file(name: &str) -> PathBuf {
|
pub fn agent_config_file(name: &str) -> PathBuf {
|
||||||
match env::var(format!("{}_CONFIG_FILE", normalize_env_name(name))) {
|
match env::var(format!("{}_CONFIG_FILE", normalize_env_name(name))) {
|
||||||
Ok(value) => PathBuf::from(value),
|
Ok(value) => PathBuf::from(value),
|
||||||
|
|||||||
@@ -9,7 +9,6 @@ use std::sync::{Arc, Weak};
|
|||||||
pub enum RagKey {
|
pub enum RagKey {
|
||||||
Named(String),
|
Named(String),
|
||||||
Agent(String),
|
Agent(String),
|
||||||
GraphNode { agent: String, node: String },
|
|
||||||
}
|
}
|
||||||
|
|
||||||
#[derive(Default)]
|
#[derive(Default)]
|
||||||
|
|||||||
+81
-658
@@ -1,14 +1,14 @@
|
|||||||
|
use super::MessageContentToolCalls;
|
||||||
use super::rag_cache::{RagCache, RagKey};
|
use super::rag_cache::{RagCache, RagKey};
|
||||||
use super::session::Session;
|
use super::session::Session;
|
||||||
use super::todo::TodoList;
|
use super::todo::TodoList;
|
||||||
use super::tool_scope::{McpRuntime, ToolScope};
|
use super::tool_scope::{McpRuntime, ToolScope};
|
||||||
use super::{
|
use super::{
|
||||||
AGENTS_DIR_NAME, Agent, AgentVariables, AppConfig, AppState, AssetCategory, CREATE_TITLE_ROLE,
|
AGENTS_DIR_NAME, Agent, AgentVariables, AppConfig, AppState, CREATE_TITLE_ROLE, Input,
|
||||||
Input, InstallFilter, LEFT_PROMPT, LastMessage, MESSAGES_FILE_NAME, RIGHT_PROMPT, Role,
|
LEFT_PROMPT, LastMessage, MESSAGES_FILE_NAME, RIGHT_PROMPT, Role, RoleLike, SESSIONS_DIR_NAME,
|
||||||
RoleLike, SESSIONS_DIR_NAME, SUMMARIZATION_PROMPT, SUMMARY_CONTEXT_PROMPT, StateFlags,
|
SUMMARIZATION_PROMPT, SUMMARY_CONTEXT_PROMPT, StateFlags, TEMP_ROLE_NAME, TEMP_SESSION_NAME,
|
||||||
TEMP_ROLE_NAME, TEMP_SESSION_NAME, WorkingMode, ensure_parent_exists, list_agents, paths,
|
WorkingMode, ensure_parent_exists, list_agents, paths,
|
||||||
};
|
};
|
||||||
use super::{MessageContentToolCalls, prompts};
|
|
||||||
use crate::client::{Model, ModelType, list_models};
|
use crate::client::{Model, ModelType, list_models};
|
||||||
use crate::function::{
|
use crate::function::{
|
||||||
FunctionDeclaration, Functions, ToolCallTracker, ToolResult,
|
FunctionDeclaration, Functions, ToolCallTracker, ToolResult,
|
||||||
@@ -27,33 +27,16 @@ use crate::utils::{
|
|||||||
list_file_names, now, render_prompt, temp_file,
|
list_file_names, now, render_prompt, temp_file,
|
||||||
};
|
};
|
||||||
|
|
||||||
use crate::graph;
|
|
||||||
use anyhow::{Context, Error, Result, bail};
|
use anyhow::{Context, Error, Result, bail};
|
||||||
#[cfg(test)]
|
|
||||||
use indexmap::IndexMap;
|
|
||||||
use indoc::formatdoc;
|
use indoc::formatdoc;
|
||||||
use inquire::{Confirm, MultiSelect, Text, list_option::ListOption, validator::Validation};
|
use inquire::{Confirm, MultiSelect, Text, list_option::ListOption, validator::Validation};
|
||||||
use parking_lot::RwLock;
|
use parking_lot::RwLock;
|
||||||
use std::collections::{HashMap, HashSet};
|
use std::collections::{HashMap, HashSet};
|
||||||
|
use std::env;
|
||||||
use std::fs::{File, OpenOptions, read_dir, read_to_string, remove_dir_all, remove_file};
|
use std::fs::{File, OpenOptions, read_dir, read_to_string, remove_dir_all, remove_file};
|
||||||
use std::io::Write;
|
use std::io::Write;
|
||||||
use std::path::{Path, PathBuf};
|
use std::path::{Path, PathBuf};
|
||||||
use std::sync::Arc;
|
use std::sync::Arc;
|
||||||
use std::{env, fs};
|
|
||||||
|
|
||||||
pub struct AutoContinueConfig {
|
|
||||||
pub enabled: bool,
|
|
||||||
pub max_continues: usize,
|
|
||||||
pub inject_instructions: bool,
|
|
||||||
pub continuation_prompt: Option<String>,
|
|
||||||
}
|
|
||||||
|
|
||||||
#[derive(Debug, Clone, Copy, PartialEq, Eq, Default)]
|
|
||||||
pub enum RenderMode {
|
|
||||||
#[default]
|
|
||||||
Streaming,
|
|
||||||
Silent,
|
|
||||||
}
|
|
||||||
|
|
||||||
pub struct RequestContext {
|
pub struct RequestContext {
|
||||||
pub app: Arc<AppState>,
|
pub app: Arc<AppState>,
|
||||||
@@ -83,8 +66,6 @@ pub struct RequestContext {
|
|||||||
pub auto_continue_count: usize,
|
pub auto_continue_count: usize,
|
||||||
pub todo_list: TodoList,
|
pub todo_list: TodoList,
|
||||||
pub last_continuation_response: Option<String>,
|
pub last_continuation_response: Option<String>,
|
||||||
|
|
||||||
pub render_mode: RenderMode,
|
|
||||||
}
|
}
|
||||||
|
|
||||||
impl RequestContext {
|
impl RequestContext {
|
||||||
@@ -111,7 +92,6 @@ impl RequestContext {
|
|||||||
auto_continue_count: 0,
|
auto_continue_count: 0,
|
||||||
todo_list: TodoList::default(),
|
todo_list: TodoList::default(),
|
||||||
last_continuation_response: None,
|
last_continuation_response: None,
|
||||||
render_mode: RenderMode::default(),
|
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
@@ -158,51 +138,9 @@ impl RequestContext {
|
|||||||
auto_continue_count: 0,
|
auto_continue_count: 0,
|
||||||
todo_list: TodoList::default(),
|
todo_list: TodoList::default(),
|
||||||
last_continuation_response: None,
|
last_continuation_response: None,
|
||||||
render_mode: RenderMode::default(),
|
|
||||||
})
|
})
|
||||||
}
|
}
|
||||||
|
|
||||||
/// Forks the context for one parallel branch of a graph super-step.
|
|
||||||
///
|
|
||||||
/// Each branch gets a fresh, owned clone. Mutations (role swap,
|
|
||||||
/// `before/after_chat_completion`, tool tracker, last_message, etc.) are
|
|
||||||
/// scoped to the branch and discarded when the branch finishes. The
|
|
||||||
/// user-visible state communication happens through the graph's
|
|
||||||
/// `StateManager` (via `fork_for_branch_state` + `diff_against` +
|
|
||||||
/// `apply_branch_writes` reducers), and not through `RequestContext`.
|
|
||||||
///
|
|
||||||
/// Distinction from `new_for_child`: `new_for_child` builds a fresh context
|
|
||||||
/// for a spawned sub-agent (different agent identity, different supervisor
|
|
||||||
/// hierarchy, depth+1, fresh tool tracker). `fork_for_branch` keeps the
|
|
||||||
/// caller's identity and supervisor hierarchy; it's a sibling clone of the
|
|
||||||
/// same logical agent, running one of N parallel work items.
|
|
||||||
pub fn fork_for_branch(&self) -> Self {
|
|
||||||
Self {
|
|
||||||
app: Arc::clone(&self.app),
|
|
||||||
macro_flag: self.macro_flag,
|
|
||||||
info_flag: self.info_flag,
|
|
||||||
working_mode: self.working_mode,
|
|
||||||
model: self.model.clone(),
|
|
||||||
agent_variables: self.agent_variables.clone(),
|
|
||||||
role: self.role.clone(),
|
|
||||||
session: self.session.clone(),
|
|
||||||
rag: self.rag.clone(),
|
|
||||||
agent: self.agent.clone(),
|
|
||||||
last_message: self.last_message.clone(),
|
|
||||||
tool_scope: self.tool_scope.clone(),
|
|
||||||
supervisor: self.supervisor.clone(),
|
|
||||||
parent_supervisor: self.parent_supervisor.clone(),
|
|
||||||
self_agent_id: self.self_agent_id.clone(),
|
|
||||||
inbox: self.inbox.clone(),
|
|
||||||
escalation_queue: self.escalation_queue.clone(),
|
|
||||||
current_depth: self.current_depth,
|
|
||||||
auto_continue_count: 0,
|
|
||||||
todo_list: self.todo_list.clone(),
|
|
||||||
last_continuation_response: None,
|
|
||||||
render_mode: self.render_mode,
|
|
||||||
}
|
|
||||||
}
|
|
||||||
|
|
||||||
pub fn new_for_child(
|
pub fn new_for_child(
|
||||||
app: Arc<AppState>,
|
app: Arc<AppState>,
|
||||||
parent: &Self,
|
parent: &Self,
|
||||||
@@ -238,7 +176,6 @@ impl RequestContext {
|
|||||||
auto_continue_count: 0,
|
auto_continue_count: 0,
|
||||||
todo_list: TodoList::default(),
|
todo_list: TodoList::default(),
|
||||||
last_continuation_response: None,
|
last_continuation_response: None,
|
||||||
render_mode: parent.render_mode,
|
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
@@ -586,7 +523,7 @@ impl RequestContext {
|
|||||||
}
|
}
|
||||||
|
|
||||||
pub fn extract_role(&self, app: &AppConfig) -> Role {
|
pub fn extract_role(&self, app: &AppConfig) -> Role {
|
||||||
let mut role = if let Some(session) = self.session.as_ref() {
|
if let Some(session) = self.session.as_ref() {
|
||||||
session.to_role()
|
session.to_role()
|
||||||
} else if let Some(agent) = self.agent.as_ref() {
|
} else if let Some(agent) = self.agent.as_ref() {
|
||||||
agent.to_role()
|
agent.to_role()
|
||||||
@@ -602,65 +539,6 @@ impl RequestContext {
|
|||||||
app.enabled_mcp_servers.clone(),
|
app.enabled_mcp_servers.clone(),
|
||||||
);
|
);
|
||||||
role
|
role
|
||||||
};
|
|
||||||
|
|
||||||
if self.agent.is_none() && self.app.config.function_calling_support {
|
|
||||||
let config = self.auto_continue_config();
|
|
||||||
if config.enabled && config.inject_instructions {
|
|
||||||
role.append_to_prompt(prompts::DEFAULT_TODO_INSTRUCTIONS);
|
|
||||||
}
|
|
||||||
}
|
|
||||||
|
|
||||||
role
|
|
||||||
}
|
|
||||||
|
|
||||||
pub fn auto_continue_config(&self) -> AutoContinueConfig {
|
|
||||||
if let Some(agent) = &self.agent {
|
|
||||||
return AutoContinueConfig {
|
|
||||||
enabled: agent.auto_continue_enabled(),
|
|
||||||
max_continues: agent.max_auto_continues(),
|
|
||||||
inject_instructions: agent.inject_todo_instructions(),
|
|
||||||
continuation_prompt: agent.continuation_prompt_value(),
|
|
||||||
};
|
|
||||||
}
|
|
||||||
let app = &self.app.config;
|
|
||||||
let enabled = self
|
|
||||||
.session
|
|
||||||
.as_ref()
|
|
||||||
.and_then(|s| s.auto_continue())
|
|
||||||
.or_else(|| self.role.as_ref().and_then(|r| r.auto_continue()))
|
|
||||||
.unwrap_or(app.auto_continue);
|
|
||||||
let max = self
|
|
||||||
.session
|
|
||||||
.as_ref()
|
|
||||||
.and_then(|s| s.max_auto_continues())
|
|
||||||
.or_else(|| self.role.as_ref().and_then(|r| r.max_auto_continues()))
|
|
||||||
.unwrap_or(app.max_auto_continues);
|
|
||||||
let inject = self
|
|
||||||
.session
|
|
||||||
.as_ref()
|
|
||||||
.and_then(|s| s.inject_todo_instructions())
|
|
||||||
.or_else(|| {
|
|
||||||
self.role
|
|
||||||
.as_ref()
|
|
||||||
.and_then(|r| r.inject_todo_instructions())
|
|
||||||
})
|
|
||||||
.unwrap_or(app.inject_todo_instructions);
|
|
||||||
let prompt = self
|
|
||||||
.session
|
|
||||||
.as_ref()
|
|
||||||
.and_then(|s| s.continuation_prompt().map(|v| v.to_string()))
|
|
||||||
.or_else(|| {
|
|
||||||
self.role
|
|
||||||
.as_ref()
|
|
||||||
.and_then(|r| r.continuation_prompt().map(|v| v.to_string()))
|
|
||||||
})
|
|
||||||
.or_else(|| app.continuation_prompt.clone());
|
|
||||||
AutoContinueConfig {
|
|
||||||
enabled,
|
|
||||||
max_continues: max,
|
|
||||||
inject_instructions: inject,
|
|
||||||
continuation_prompt: prompt,
|
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
@@ -869,8 +747,6 @@ impl RequestContext {
|
|||||||
app.function_calling_support.to_string(),
|
app.function_calling_support.to_string(),
|
||||||
),
|
),
|
||||||
("mcp_server_support", app.mcp_server_support.to_string()),
|
("mcp_server_support", app.mcp_server_support.to_string()),
|
||||||
("auto_continue", app.auto_continue.to_string()),
|
|
||||||
("max_auto_continues", app.max_auto_continues.to_string()),
|
|
||||||
("stream", app.stream.to_string()),
|
("stream", app.stream.to_string()),
|
||||||
("save", app.save.to_string()),
|
("save", app.save.to_string()),
|
||||||
("keybindings", app.keybindings.clone()),
|
("keybindings", app.keybindings.clone()),
|
||||||
@@ -1046,12 +922,9 @@ impl RequestContext {
|
|||||||
let app = self.app.config.as_ref();
|
let app = self.app.config.as_ref();
|
||||||
let mut functions = vec![];
|
let mut functions = vec![];
|
||||||
if app.function_calling_support {
|
if app.function_calling_support {
|
||||||
// Compute the set of tool names enabled by the role filter, drawn
|
if let Some(enabled_tools) = role.enabled_tools() {
|
||||||
// from BOTH the tool_scope pool and the agent's pool so that an
|
let mut tool_names: HashSet<String> = Default::default();
|
||||||
// explicit `enabled_tools` list (e.g. from a graph LLM node) can
|
let declaration_names: HashSet<String> = self
|
||||||
// narrow the agent's own custom tools too.
|
|
||||||
let role_filter: Option<HashSet<String>> = role.enabled_tools().map(|enabled_tools| {
|
|
||||||
let mut declaration_names: HashSet<String> = self
|
|
||||||
.tool_scope
|
.tool_scope
|
||||||
.functions
|
.functions
|
||||||
.declarations()
|
.declarations()
|
||||||
@@ -1063,32 +936,11 @@ impl RequestContext {
|
|||||||
})
|
})
|
||||||
.map(|v| v.name.to_string())
|
.map(|v| v.name.to_string())
|
||||||
.collect();
|
.collect();
|
||||||
|
|
||||||
if let Some(agent) = &self.agent {
|
|
||||||
declaration_names.extend(
|
|
||||||
agent
|
|
||||||
.functions()
|
|
||||||
.declarations()
|
|
||||||
.iter()
|
|
||||||
.filter(|v| {
|
|
||||||
!v.name.starts_with(MCP_INVOKE_META_FUNCTION_NAME_PREFIX)
|
|
||||||
&& !v.name.starts_with(MCP_SEARCH_META_FUNCTION_NAME_PREFIX)
|
|
||||||
&& !v.name.starts_with(MCP_DESCRIBE_META_FUNCTION_NAME_PREFIX)
|
|
||||||
})
|
|
||||||
.map(|v| v.name.to_string()),
|
|
||||||
);
|
|
||||||
}
|
|
||||||
|
|
||||||
let mut tool_names: HashSet<String> = Default::default();
|
|
||||||
if enabled_tools == "all" {
|
if enabled_tools == "all" {
|
||||||
tool_names.extend(declaration_names);
|
tool_names.extend(declaration_names);
|
||||||
} else {
|
} else {
|
||||||
for item in enabled_tools.split(',') {
|
for item in enabled_tools.split(',') {
|
||||||
let item = item.trim();
|
let item = item.trim();
|
||||||
if item.is_empty() {
|
|
||||||
continue;
|
|
||||||
}
|
|
||||||
|
|
||||||
if let Some(values) = app.mapping_tools.get(item) {
|
if let Some(values) = app.mapping_tools.get(item) {
|
||||||
tool_names.extend(
|
tool_names.extend(
|
||||||
values
|
values
|
||||||
@@ -1101,10 +953,6 @@ impl RequestContext {
|
|||||||
}
|
}
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
tool_names
|
|
||||||
});
|
|
||||||
|
|
||||||
if let Some(ref tool_names) = role_filter {
|
|
||||||
functions = self
|
functions = self
|
||||||
.tool_scope
|
.tool_scope
|
||||||
.functions
|
.functions
|
||||||
@@ -1147,11 +995,6 @@ impl RequestContext {
|
|||||||
&& !v.name.starts_with(MCP_DESCRIBE_META_FUNCTION_NAME_PREFIX)
|
&& !v.name.starts_with(MCP_DESCRIBE_META_FUNCTION_NAME_PREFIX)
|
||||||
})
|
})
|
||||||
.collect();
|
.collect();
|
||||||
|
|
||||||
if let Some(ref tool_names) = role_filter {
|
|
||||||
agent_functions.retain(|v| tool_names.contains(&v.name));
|
|
||||||
}
|
|
||||||
|
|
||||||
let tool_names: HashSet<String> = agent_functions
|
let tool_names: HashSet<String> = agent_functions
|
||||||
.iter()
|
.iter()
|
||||||
.filter_map(|v| {
|
.filter_map(|v| {
|
||||||
@@ -1178,88 +1021,63 @@ impl RequestContext {
|
|||||||
let app = self.app.config.as_ref();
|
let app = self.app.config.as_ref();
|
||||||
let mut mcp_functions = vec![];
|
let mut mcp_functions = vec![];
|
||||||
if app.mcp_server_support {
|
if app.mcp_server_support {
|
||||||
let role_filter: Option<HashSet<String>> =
|
if let Some(enabled_mcp_servers) = role.enabled_mcp_servers() {
|
||||||
role.enabled_mcp_servers().map(|enabled_mcp_servers| {
|
let mut server_names: HashSet<String> = Default::default();
|
||||||
let mut mcp_declaration_names: HashSet<String> = self
|
let mcp_declaration_names: HashSet<String> = self
|
||||||
.tool_scope
|
.tool_scope
|
||||||
.functions
|
.functions
|
||||||
.declarations()
|
.declarations()
|
||||||
.iter()
|
.iter()
|
||||||
.filter(|v| {
|
.filter(|v| {
|
||||||
v.name.starts_with(MCP_INVOKE_META_FUNCTION_NAME_PREFIX)
|
v.name.starts_with(MCP_INVOKE_META_FUNCTION_NAME_PREFIX)
|
||||||
|| v.name.starts_with(MCP_SEARCH_META_FUNCTION_NAME_PREFIX)
|
|| v.name.starts_with(MCP_SEARCH_META_FUNCTION_NAME_PREFIX)
|
||||||
|| v.name.starts_with(MCP_DESCRIBE_META_FUNCTION_NAME_PREFIX)
|
|| v.name.starts_with(MCP_DESCRIBE_META_FUNCTION_NAME_PREFIX)
|
||||||
})
|
})
|
||||||
.map(|v| v.name.to_string())
|
.map(|v| v.name.to_string())
|
||||||
.collect();
|
.collect();
|
||||||
if let Some(agent) = &self.agent {
|
if enabled_mcp_servers == "all" {
|
||||||
mcp_declaration_names.extend(
|
server_names.extend(mcp_declaration_names);
|
||||||
agent
|
} else {
|
||||||
.functions()
|
for item in enabled_mcp_servers.split(',') {
|
||||||
.declarations()
|
let item = item.trim();
|
||||||
.iter()
|
let item_invoke_name =
|
||||||
.filter(|v| {
|
format!("{}_{item}", MCP_INVOKE_META_FUNCTION_NAME_PREFIX);
|
||||||
v.name.starts_with(MCP_INVOKE_META_FUNCTION_NAME_PREFIX)
|
let item_search_name =
|
||||||
|| v.name.starts_with(MCP_SEARCH_META_FUNCTION_NAME_PREFIX)
|
format!("{}_{item}", MCP_SEARCH_META_FUNCTION_NAME_PREFIX);
|
||||||
|| v.name
|
let item_describe_name =
|
||||||
.starts_with(MCP_DESCRIBE_META_FUNCTION_NAME_PREFIX)
|
format!("{}_{item}", MCP_DESCRIBE_META_FUNCTION_NAME_PREFIX);
|
||||||
})
|
if let Some(values) = app.mapping_mcp_servers.get(item) {
|
||||||
.map(|v| v.name.to_string()),
|
server_names.extend(
|
||||||
);
|
values
|
||||||
}
|
.split(',')
|
||||||
|
.flat_map(|v| {
|
||||||
let mut server_names: HashSet<String> = Default::default();
|
vec![
|
||||||
if enabled_mcp_servers == "all" {
|
format!(
|
||||||
server_names.extend(mcp_declaration_names);
|
"{}_{}",
|
||||||
} else {
|
MCP_INVOKE_META_FUNCTION_NAME_PREFIX,
|
||||||
for item in enabled_mcp_servers.split(',') {
|
v.to_string()
|
||||||
let item = item.trim();
|
),
|
||||||
if item.is_empty() {
|
format!(
|
||||||
continue;
|
"{}_{}",
|
||||||
}
|
MCP_SEARCH_META_FUNCTION_NAME_PREFIX,
|
||||||
|
v.to_string()
|
||||||
let item_invoke_name =
|
),
|
||||||
format!("{}_{item}", MCP_INVOKE_META_FUNCTION_NAME_PREFIX);
|
format!(
|
||||||
let item_search_name =
|
"{}_{}",
|
||||||
format!("{}_{item}", MCP_SEARCH_META_FUNCTION_NAME_PREFIX);
|
MCP_DESCRIBE_META_FUNCTION_NAME_PREFIX,
|
||||||
let item_describe_name =
|
v.to_string()
|
||||||
format!("{}_{item}", MCP_DESCRIBE_META_FUNCTION_NAME_PREFIX);
|
),
|
||||||
if let Some(values) = app.mapping_mcp_servers.get(item) {
|
]
|
||||||
server_names.extend(
|
})
|
||||||
values
|
.filter(|v| mcp_declaration_names.contains(v)),
|
||||||
.split(',')
|
)
|
||||||
.flat_map(|v| {
|
} else if mcp_declaration_names.contains(&item_invoke_name) {
|
||||||
vec![
|
server_names.insert(item_invoke_name);
|
||||||
format!(
|
server_names.insert(item_search_name);
|
||||||
"{}_{}",
|
server_names.insert(item_describe_name);
|
||||||
MCP_INVOKE_META_FUNCTION_NAME_PREFIX,
|
|
||||||
v.to_string()
|
|
||||||
),
|
|
||||||
format!(
|
|
||||||
"{}_{}",
|
|
||||||
MCP_SEARCH_META_FUNCTION_NAME_PREFIX,
|
|
||||||
v.to_string()
|
|
||||||
),
|
|
||||||
format!(
|
|
||||||
"{}_{}",
|
|
||||||
MCP_DESCRIBE_META_FUNCTION_NAME_PREFIX,
|
|
||||||
v.to_string()
|
|
||||||
),
|
|
||||||
]
|
|
||||||
})
|
|
||||||
.filter(|v| mcp_declaration_names.contains(v)),
|
|
||||||
)
|
|
||||||
} else if mcp_declaration_names.contains(&item_invoke_name) {
|
|
||||||
server_names.insert(item_invoke_name);
|
|
||||||
server_names.insert(item_search_name);
|
|
||||||
server_names.insert(item_describe_name);
|
|
||||||
}
|
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
server_names
|
}
|
||||||
});
|
|
||||||
|
|
||||||
if let Some(ref server_names) = role_filter {
|
|
||||||
mcp_functions = self
|
mcp_functions = self
|
||||||
.tool_scope
|
.tool_scope
|
||||||
.functions
|
.functions
|
||||||
@@ -1287,11 +1105,6 @@ impl RequestContext {
|
|||||||
|| v.name.starts_with(MCP_DESCRIBE_META_FUNCTION_NAME_PREFIX)
|
|| v.name.starts_with(MCP_DESCRIBE_META_FUNCTION_NAME_PREFIX)
|
||||||
})
|
})
|
||||||
.collect();
|
.collect();
|
||||||
|
|
||||||
if let Some(ref server_names) = role_filter {
|
|
||||||
agent_functions.retain(|v| server_names.contains(&v.name));
|
|
||||||
}
|
|
||||||
|
|
||||||
let tool_names: HashSet<String> = agent_functions
|
let tool_names: HashSet<String> = agent_functions
|
||||||
.iter()
|
.iter()
|
||||||
.filter_map(|v| {
|
.filter_map(|v| {
|
||||||
@@ -1399,19 +1212,6 @@ impl RequestContext {
|
|||||||
Ok(())
|
Ok(())
|
||||||
}
|
}
|
||||||
|
|
||||||
pub fn edit_mcp_config(&self) -> Result<()> {
|
|
||||||
let mcp_path = paths::mcp_config_file();
|
|
||||||
let editor = self.app.config.editor()?;
|
|
||||||
edit_file(&editor, &mcp_path)?;
|
|
||||||
println!(
|
|
||||||
"NOTE: Remember to restart {} for changes to '{}' to take effect",
|
|
||||||
env!("CARGO_CRATE_NAME"),
|
|
||||||
mcp_path.display(),
|
|
||||||
);
|
|
||||||
|
|
||||||
Ok(())
|
|
||||||
}
|
|
||||||
|
|
||||||
pub fn new_role(&self, app: &AppConfig, name: &str) -> Result<()> {
|
pub fn new_role(&self, app: &AppConfig, name: &str) -> Result<()> {
|
||||||
if self.macro_flag {
|
if self.macro_flag {
|
||||||
bail!("No role");
|
bail!("No role");
|
||||||
@@ -1486,30 +1286,21 @@ impl RequestContext {
|
|||||||
Some(agent) => agent.name(),
|
Some(agent) => agent.name(),
|
||||||
None => bail!("No agent"),
|
None => bail!("No agent"),
|
||||||
};
|
};
|
||||||
let config_path = paths::agent_config_file(agent_name);
|
let agent_config_path = paths::agent_config_file(agent_name);
|
||||||
let graph_path = paths::agent_graph_file(agent_name);
|
ensure_parent_exists(&agent_config_path)?;
|
||||||
let target_path = if !config_path.exists() && graph_path.exists() {
|
if !agent_config_path.exists() {
|
||||||
graph_path
|
std::fs::write(
|
||||||
} else {
|
&agent_config_path,
|
||||||
config_path
|
|
||||||
};
|
|
||||||
|
|
||||||
ensure_parent_exists(&target_path)?;
|
|
||||||
if !target_path.exists() {
|
|
||||||
fs::write(
|
|
||||||
&target_path,
|
|
||||||
"# see https://github.com/Dark-Alex-17/loki/blob/main/config.agent.example.yaml\n",
|
"# see https://github.com/Dark-Alex-17/loki/blob/main/config.agent.example.yaml\n",
|
||||||
)
|
)
|
||||||
.with_context(|| format!("Failed to write to '{}'", target_path.display()))?;
|
.with_context(|| format!("Failed to write to '{}'", agent_config_path.display()))?;
|
||||||
}
|
}
|
||||||
|
|
||||||
let editor = app.editor()?;
|
let editor = app.editor()?;
|
||||||
edit_file(&editor, &target_path)?;
|
edit_file(&editor, &agent_config_path)?;
|
||||||
println!(
|
println!(
|
||||||
"NOTE: Remember to reload the agent if there are changes made to '{}'",
|
"NOTE: Remember to reload the agent if there are changes made to '{}'",
|
||||||
target_path.display()
|
agent_config_path.display()
|
||||||
);
|
);
|
||||||
|
|
||||||
Ok(())
|
Ok(())
|
||||||
}
|
}
|
||||||
|
|
||||||
@@ -1611,24 +1402,12 @@ impl RequestContext {
|
|||||||
}
|
}
|
||||||
|
|
||||||
pub async fn update(&mut self, data: &str, abort_signal: AbortSignal) -> Result<()> {
|
pub async fn update(&mut self, data: &str, abort_signal: AbortSignal) -> Result<()> {
|
||||||
let (key, raw_value) = match data.split_once(char::is_whitespace) {
|
let parts: Vec<&str> = data.split_whitespace().collect();
|
||||||
Some((k, v)) => (k, v.trim()),
|
if parts.len() != 2 {
|
||||||
None => bail!("Usage: .set <key> <value>. If value is null, unset key."),
|
|
||||||
};
|
|
||||||
|
|
||||||
if raw_value.is_empty() {
|
|
||||||
bail!("Usage: .set <key> <value>. If value is null, unset key.");
|
bail!("Usage: .set <key> <value>. If value is null, unset key.");
|
||||||
}
|
}
|
||||||
|
let key = parts[0];
|
||||||
let value = match key {
|
let value = parts[1];
|
||||||
"continuation_prompt" => raw_value,
|
|
||||||
_ => {
|
|
||||||
if raw_value.contains(char::is_whitespace) {
|
|
||||||
bail!("Usage: .set <key> <value>. If value is null, unset key.");
|
|
||||||
}
|
|
||||||
raw_value
|
|
||||||
}
|
|
||||||
};
|
|
||||||
match key {
|
match key {
|
||||||
"temperature" => {
|
"temperature" => {
|
||||||
let value = super::parse_value(value)?;
|
let value = super::parse_value(value)?;
|
||||||
@@ -1743,49 +1522,6 @@ impl RequestContext {
|
|||||||
let value = value.parse().with_context(|| "Invalid value")?;
|
let value = value.parse().with_context(|| "Invalid value")?;
|
||||||
self.update_app_config(|app| app.highlight = value);
|
self.update_app_config(|app| app.highlight = value);
|
||||||
}
|
}
|
||||||
"auto_continue" => {
|
|
||||||
let value: bool = value.parse().with_context(|| "Invalid value")?;
|
|
||||||
if value && !self.app.config.function_calling_support {
|
|
||||||
bail!(
|
|
||||||
"Cannot enable auto_continue: function calling is disabled. Set 'function_calling_support: true' first."
|
|
||||||
);
|
|
||||||
}
|
|
||||||
if let Some(session) = self.session.as_mut() {
|
|
||||||
session.set_auto_continue(Some(value));
|
|
||||||
} else {
|
|
||||||
self.update_app_config(|app| app.auto_continue = value);
|
|
||||||
}
|
|
||||||
if value
|
|
||||||
&& self.app.config.function_calling_support
|
|
||||||
&& !self.tool_scope.functions.contains("todo__init")
|
|
||||||
{
|
|
||||||
self.tool_scope.functions.append_todo_functions();
|
|
||||||
}
|
|
||||||
}
|
|
||||||
"max_auto_continues" => {
|
|
||||||
let value: usize = value.parse().with_context(|| "Invalid value")?;
|
|
||||||
if let Some(session) = self.session.as_mut() {
|
|
||||||
session.set_max_auto_continues(Some(value));
|
|
||||||
} else {
|
|
||||||
self.update_app_config(|app| app.max_auto_continues = value);
|
|
||||||
}
|
|
||||||
}
|
|
||||||
"inject_todo_instructions" => {
|
|
||||||
let value: bool = value.parse().with_context(|| "Invalid value")?;
|
|
||||||
if let Some(session) = self.session.as_mut() {
|
|
||||||
session.set_inject_todo_instructions(Some(value));
|
|
||||||
} else {
|
|
||||||
self.update_app_config(|app| app.inject_todo_instructions = value);
|
|
||||||
}
|
|
||||||
}
|
|
||||||
"continuation_prompt" => {
|
|
||||||
let value: Option<String> = super::parse_value(value)?;
|
|
||||||
if let Some(session) = self.session.as_mut() {
|
|
||||||
session.set_continuation_prompt(value);
|
|
||||||
} else {
|
|
||||||
self.update_app_config(|app| app.continuation_prompt = value);
|
|
||||||
}
|
|
||||||
}
|
|
||||||
_ => bail!("Unknown key '{key}'"),
|
_ => bail!("Unknown key '{key}'"),
|
||||||
}
|
}
|
||||||
Ok(())
|
Ok(())
|
||||||
@@ -1855,12 +1591,6 @@ impl RequestContext {
|
|||||||
}
|
}
|
||||||
".rag" => super::map_completion_values(paths::list_rags()),
|
".rag" => super::map_completion_values(paths::list_rags()),
|
||||||
".agent" => super::map_completion_values(list_agents()),
|
".agent" => super::map_completion_values(list_agents()),
|
||||||
".install" => {
|
|
||||||
let mut values: Vec<String> =
|
|
||||||
AssetCategory::NAMES.iter().map(|s| s.to_string()).collect();
|
|
||||||
values.push("remote".to_string());
|
|
||||||
super::map_completion_values(values)
|
|
||||||
}
|
|
||||||
".macro" => super::map_completion_values(paths::list_macros()),
|
".macro" => super::map_completion_values(paths::list_macros()),
|
||||||
".starter" => match &self.agent {
|
".starter" => match &self.agent {
|
||||||
Some(agent) => agent
|
Some(agent) => agent
|
||||||
@@ -1873,14 +1603,10 @@ impl RequestContext {
|
|||||||
},
|
},
|
||||||
".set" => {
|
".set" => {
|
||||||
let mut values = vec![
|
let mut values = vec![
|
||||||
"auto_continue",
|
|
||||||
"continuation_prompt",
|
|
||||||
"temperature",
|
"temperature",
|
||||||
"top_p",
|
"top_p",
|
||||||
"enabled_tools",
|
"enabled_tools",
|
||||||
"enabled_mcp_servers",
|
"enabled_mcp_servers",
|
||||||
"inject_todo_instructions",
|
|
||||||
"max_auto_continues",
|
|
||||||
"save_session",
|
"save_session",
|
||||||
"compression_threshold",
|
"compression_threshold",
|
||||||
"rag_reranker_model",
|
"rag_reranker_model",
|
||||||
@@ -1916,28 +1642,6 @@ impl RequestContext {
|
|||||||
}
|
}
|
||||||
_ => vec![],
|
_ => vec![],
|
||||||
};
|
};
|
||||||
} else if cmd == ".install" && args.first() == Some(&"remote") && args.len() >= 2 {
|
|
||||||
let prev = args.get(args.len() - 2).copied().unwrap_or("");
|
|
||||||
if prev == "--filter" {
|
|
||||||
values = super::map_completion_values(
|
|
||||||
InstallFilter::NAMES.iter().map(|s| s.to_string()).collect(),
|
|
||||||
);
|
|
||||||
} else {
|
|
||||||
let has_filter = args.iter().enumerate().any(|(i, a)| {
|
|
||||||
a.starts_with("--filter=") || (*a == "--filter" && i < args.len() - 1)
|
|
||||||
});
|
|
||||||
let has_force = args.contains(&"--force");
|
|
||||||
let mut available: Vec<&str> = vec![];
|
|
||||||
|
|
||||||
if !has_filter {
|
|
||||||
available.push("--filter");
|
|
||||||
}
|
|
||||||
if !has_force {
|
|
||||||
available.push("--force");
|
|
||||||
}
|
|
||||||
|
|
||||||
values = super::map_completion_values(available);
|
|
||||||
}
|
|
||||||
} else if cmd == ".set" && args.len() == 2 {
|
} else if cmd == ".set" && args.len() == 2 {
|
||||||
let candidates = match args[0] {
|
let candidates = match args[0] {
|
||||||
"max_output_tokens" => match self.current_model().max_output_tokens() {
|
"max_output_tokens" => match self.current_model().max_output_tokens() {
|
||||||
@@ -2017,19 +1721,6 @@ impl RequestContext {
|
|||||||
.map(|v| v.id())
|
.map(|v| v.id())
|
||||||
.collect(),
|
.collect(),
|
||||||
"highlight" => super::complete_bool(app.highlight),
|
"highlight" => super::complete_bool(app.highlight),
|
||||||
"auto_continue" => {
|
|
||||||
let config = self.auto_continue_config();
|
|
||||||
super::complete_bool(config.enabled)
|
|
||||||
}
|
|
||||||
"max_auto_continues" => {
|
|
||||||
let config = self.auto_continue_config();
|
|
||||||
vec![config.max_continues.to_string()]
|
|
||||||
}
|
|
||||||
"inject_todo_instructions" => {
|
|
||||||
let config = self.auto_continue_config();
|
|
||||||
super::complete_bool(config.inject_instructions)
|
|
||||||
}
|
|
||||||
"continuation_prompt" => vec!["null".to_string()],
|
|
||||||
_ => vec![],
|
_ => vec![],
|
||||||
};
|
};
|
||||||
values = candidates.into_iter().map(|v| (v, None)).collect();
|
values = candidates.into_iter().map(|v| (v, None)).collect();
|
||||||
@@ -2044,7 +1735,7 @@ impl RequestContext {
|
|||||||
.collect();
|
.collect();
|
||||||
} else if cmd == ".agent" {
|
} else if cmd == ".agent" {
|
||||||
if args.len() == 2 {
|
if args.len() == 2 {
|
||||||
let dir = paths::agent_data_dir(args[0]).join(SESSIONS_DIR_NAME);
|
let dir = paths::agent_data_dir(args[0]).join(super::SESSIONS_DIR_NAME);
|
||||||
values = list_file_names(dir, ".yaml")
|
values = list_file_names(dir, ".yaml")
|
||||||
.into_iter()
|
.into_iter()
|
||||||
.map(|v| (v, None))
|
.map(|v| (v, None))
|
||||||
@@ -2119,12 +1810,6 @@ impl RequestContext {
|
|||||||
if self.working_mode.is_repl() {
|
if self.working_mode.is_repl() {
|
||||||
functions.append_user_interaction_functions();
|
functions.append_user_interaction_functions();
|
||||||
}
|
}
|
||||||
if self.agent.is_none()
|
|
||||||
&& app.function_calling_support
|
|
||||||
&& self.auto_continue_config().enabled
|
|
||||||
{
|
|
||||||
functions.append_todo_functions();
|
|
||||||
}
|
|
||||||
if !mcp_runtime.is_empty() {
|
if !mcp_runtime.is_empty() {
|
||||||
functions.append_mcp_meta_functions(mcp_runtime.server_names());
|
functions.append_mcp_meta_functions(mcp_runtime.server_names());
|
||||||
}
|
}
|
||||||
@@ -2145,10 +1830,6 @@ impl RequestContext {
|
|||||||
abort_signal: AbortSignal,
|
abort_signal: AbortSignal,
|
||||||
) -> Result<()> {
|
) -> Result<()> {
|
||||||
let role = self.retrieve_role(app, name)?;
|
let role = self.retrieve_role(app, name)?;
|
||||||
if let Some(session) = self.session.as_mut() {
|
|
||||||
session.guard_empty()?;
|
|
||||||
}
|
|
||||||
|
|
||||||
let mcp_servers = if app.mcp_server_support {
|
let mcp_servers = if app.mcp_server_support {
|
||||||
role.enabled_mcp_servers()
|
role.enabled_mcp_servers()
|
||||||
} else {
|
} else {
|
||||||
@@ -2275,14 +1956,6 @@ impl RequestContext {
|
|||||||
)
|
)
|
||||||
.await?;
|
.await?;
|
||||||
|
|
||||||
let is_graph_agent = graph::agent_has_graph(agent_name);
|
|
||||||
if is_graph_agent && session_name.is_some() {
|
|
||||||
bail!(
|
|
||||||
"Graph-based agent '{agent_name}' does not support sessions. \
|
|
||||||
The graph manages its own state; re-run without a session."
|
|
||||||
);
|
|
||||||
}
|
|
||||||
|
|
||||||
let mcp_servers = if app.mcp_server_support {
|
let mcp_servers = if app.mcp_server_support {
|
||||||
(!agent.mcp_server_names().is_empty()).then(|| agent.mcp_server_names().join(","))
|
(!agent.mcp_server_names().is_empty()).then(|| agent.mcp_server_names().join(","))
|
||||||
} else {
|
} else {
|
||||||
@@ -2304,22 +1977,14 @@ impl RequestContext {
|
|||||||
);
|
);
|
||||||
}
|
}
|
||||||
|
|
||||||
// Graph agents manage their own state; never engage a session,
|
|
||||||
// not even an inherited app-level `agent_session` default.
|
|
||||||
let session_name = session_name.map(|v| v.to_string()).or_else(|| {
|
let session_name = session_name.map(|v| v.to_string()).or_else(|| {
|
||||||
if self.macro_flag || is_graph_agent {
|
if self.macro_flag {
|
||||||
None
|
None
|
||||||
} else {
|
} else {
|
||||||
agent.agent_session().map(|v| v.to_string())
|
agent.agent_session().map(|v| v.to_string())
|
||||||
}
|
}
|
||||||
});
|
});
|
||||||
|
|
||||||
if self.session.is_some() {
|
|
||||||
bail!(
|
|
||||||
"Already in a session, please run '.exit session' first to exit the current session."
|
|
||||||
);
|
|
||||||
}
|
|
||||||
|
|
||||||
let should_init_supervisor = agent.can_spawn_agents();
|
let should_init_supervisor = agent.can_spawn_agents();
|
||||||
let max_concurrent = agent.max_concurrent_agents();
|
let max_concurrent = agent.max_concurrent_agents();
|
||||||
let max_depth = agent.max_agent_depth();
|
let max_depth = agent.max_agent_depth();
|
||||||
@@ -2531,7 +2196,7 @@ impl RequestContext {
|
|||||||
.clone()
|
.clone()
|
||||||
.unwrap_or_else(|| SUMMARY_CONTEXT_PROMPT.into());
|
.unwrap_or_else(|| SUMMARY_CONTEXT_PROMPT.into());
|
||||||
|
|
||||||
let todo_prefix = if self.auto_continue_config().enabled && !self.todo_list.is_empty() {
|
let todo_prefix = if self.agent.is_some() && !self.todo_list.is_empty() {
|
||||||
format!(
|
format!(
|
||||||
"[ACTIVE TODO LIST]\n{}\n\n",
|
"[ACTIVE TODO LIST]\n{}\n\n",
|
||||||
self.todo_list.render_for_model()
|
self.todo_list.render_for_model()
|
||||||
@@ -2926,7 +2591,7 @@ mod tests {
|
|||||||
let mcp_config = if server_names.is_empty() {
|
let mcp_config = if server_names.is_empty() {
|
||||||
None
|
None
|
||||||
} else {
|
} else {
|
||||||
let mut servers = IndexMap::new();
|
let mut servers = HashMap::new();
|
||||||
for name in server_names {
|
for name in server_names {
|
||||||
servers.insert(
|
servers.insert(
|
||||||
name.to_string(),
|
name.to_string(),
|
||||||
@@ -3624,246 +3289,4 @@ mod tests {
|
|||||||
create_dir_all(&rags_dir).unwrap();
|
create_dir_all(&rags_dir).unwrap();
|
||||||
assert!(paths::list_rags().is_empty());
|
assert!(paths::list_rags().is_empty());
|
||||||
}
|
}
|
||||||
|
|
||||||
#[test]
|
|
||||||
#[serial]
|
|
||||||
fn use_agent_errors_when_already_in_session() {
|
|
||||||
let _guard = TestConfigDirGuard::new();
|
|
||||||
let mut ctx = create_test_ctx();
|
|
||||||
ctx.session = Some(Session::default());
|
|
||||||
|
|
||||||
let app = ctx.app.config.clone();
|
|
||||||
let agent_name = format!(
|
|
||||||
"test_agent_{}",
|
|
||||||
SystemTime::now()
|
|
||||||
.duration_since(UNIX_EPOCH)
|
|
||||||
.unwrap()
|
|
||||||
.as_nanos()
|
|
||||||
);
|
|
||||||
let agent_dir = paths::agent_data_dir(&agent_name);
|
|
||||||
create_dir_all(&agent_dir).unwrap();
|
|
||||||
write(
|
|
||||||
agent_dir.join("config.yaml"),
|
|
||||||
format!("name: {agent_name}\ninstructions: hi\n"),
|
|
||||||
)
|
|
||||||
.unwrap();
|
|
||||||
|
|
||||||
let abort = utils::create_abort_signal();
|
|
||||||
let result = run_async(ctx.use_agent(&app, &agent_name, Some("test_session"), abort));
|
|
||||||
assert!(result.is_err());
|
|
||||||
assert!(
|
|
||||||
result
|
|
||||||
.unwrap_err()
|
|
||||||
.to_string()
|
|
||||||
.contains("Already in a session")
|
|
||||||
);
|
|
||||||
assert!(
|
|
||||||
ctx.agent.is_none(),
|
|
||||||
"Agent should not be set when session check fails"
|
|
||||||
);
|
|
||||||
}
|
|
||||||
|
|
||||||
#[test]
|
|
||||||
#[serial]
|
|
||||||
fn use_agent_errors_when_already_in_session_even_without_session_name() {
|
|
||||||
let _guard = TestConfigDirGuard::new();
|
|
||||||
let mut ctx = create_test_ctx();
|
|
||||||
ctx.session = Some(Session::default());
|
|
||||||
|
|
||||||
let app = ctx.app.config.clone();
|
|
||||||
let agent_name = format!(
|
|
||||||
"test_agent_{}",
|
|
||||||
SystemTime::now()
|
|
||||||
.duration_since(UNIX_EPOCH)
|
|
||||||
.unwrap()
|
|
||||||
.as_nanos()
|
|
||||||
);
|
|
||||||
let agent_dir = paths::agent_data_dir(&agent_name);
|
|
||||||
create_dir_all(&agent_dir).unwrap();
|
|
||||||
write(
|
|
||||||
agent_dir.join("config.yaml"),
|
|
||||||
format!("name: {agent_name}\ninstructions: hi\n"),
|
|
||||||
)
|
|
||||||
.unwrap();
|
|
||||||
|
|
||||||
let abort = utils::create_abort_signal();
|
|
||||||
let result = run_async(ctx.use_agent(&app, &agent_name, None, abort));
|
|
||||||
assert!(result.is_err());
|
|
||||||
assert!(
|
|
||||||
result
|
|
||||||
.unwrap_err()
|
|
||||||
.to_string()
|
|
||||||
.contains("Already in a session")
|
|
||||||
);
|
|
||||||
assert!(
|
|
||||||
ctx.agent.is_none(),
|
|
||||||
"Agent should not be set when session check fails"
|
|
||||||
);
|
|
||||||
}
|
|
||||||
|
|
||||||
#[test]
|
|
||||||
#[serial]
|
|
||||||
fn use_agent_errors_when_graph_agent_given_explicit_session() {
|
|
||||||
let _guard = TestConfigDirGuard::new();
|
|
||||||
let mut ctx = create_test_ctx();
|
|
||||||
|
|
||||||
let app = ctx.app.config.clone();
|
|
||||||
let agent_name = format!(
|
|
||||||
"test_graph_agent_{}",
|
|
||||||
SystemTime::now()
|
|
||||||
.duration_since(UNIX_EPOCH)
|
|
||||||
.unwrap()
|
|
||||||
.as_nanos()
|
|
||||||
);
|
|
||||||
let agent_dir = paths::agent_data_dir(&agent_name);
|
|
||||||
create_dir_all(&agent_dir).unwrap();
|
|
||||||
write(
|
|
||||||
agent_dir.join("graph.yaml"),
|
|
||||||
format!(
|
|
||||||
"name: {agent_name}\nversion: \"1.0\"\nstart: done\nnodes:\n done:\n type: end\n output: ok\n"
|
|
||||||
),
|
|
||||||
)
|
|
||||||
.unwrap();
|
|
||||||
|
|
||||||
let abort = utils::create_abort_signal();
|
|
||||||
let result = run_async(ctx.use_agent(&app, &agent_name, Some("test_session"), abort));
|
|
||||||
|
|
||||||
assert!(result.is_err());
|
|
||||||
assert!(
|
|
||||||
result
|
|
||||||
.unwrap_err()
|
|
||||||
.to_string()
|
|
||||||
.contains("does not support sessions")
|
|
||||||
);
|
|
||||||
assert!(
|
|
||||||
ctx.agent.is_none(),
|
|
||||||
"Agent should not be set when the graph-agent session guard fails"
|
|
||||||
);
|
|
||||||
}
|
|
||||||
|
|
||||||
#[test]
|
|
||||||
#[serial]
|
|
||||||
fn use_agent_skips_inherited_session_for_graph_agent() {
|
|
||||||
let _guard = TestConfigDirGuard::new();
|
|
||||||
let mut ctx = create_test_ctx();
|
|
||||||
ctx.update_app_config(|app| app.agent_session = Some("inherited".to_string()));
|
|
||||||
|
|
||||||
let app = ctx.app.config.clone();
|
|
||||||
let agent_name = format!(
|
|
||||||
"test_graph_agent_{}",
|
|
||||||
SystemTime::now()
|
|
||||||
.duration_since(UNIX_EPOCH)
|
|
||||||
.unwrap()
|
|
||||||
.as_nanos()
|
|
||||||
);
|
|
||||||
let agent_dir = paths::agent_data_dir(&agent_name);
|
|
||||||
create_dir_all(&agent_dir).unwrap();
|
|
||||||
write(
|
|
||||||
agent_dir.join("graph.yaml"),
|
|
||||||
format!(
|
|
||||||
"name: {agent_name}\nversion: \"1.0\"\nstart: done\nnodes:\n done:\n type: end\n output: ok\n"
|
|
||||||
),
|
|
||||||
)
|
|
||||||
.unwrap();
|
|
||||||
|
|
||||||
let abort = utils::create_abort_signal();
|
|
||||||
run_async(ctx.use_agent(&app, &agent_name, None, abort)).unwrap();
|
|
||||||
|
|
||||||
assert!(ctx.agent.is_some(), "Graph agent should load successfully");
|
|
||||||
assert!(
|
|
||||||
ctx.session.is_none(),
|
|
||||||
"Graph agent must not engage a session, not even an inherited default"
|
|
||||||
);
|
|
||||||
}
|
|
||||||
|
|
||||||
fn first_file(dir: &Path) -> Option<PathBuf> {
|
|
||||||
for entry in read_dir(dir).ok()?.flatten() {
|
|
||||||
let path = entry.path();
|
|
||||||
if path.is_dir() {
|
|
||||||
if let Some(found) = first_file(&path) {
|
|
||||||
return Some(found);
|
|
||||||
}
|
|
||||||
} else {
|
|
||||||
return Some(path);
|
|
||||||
}
|
|
||||||
}
|
|
||||||
None
|
|
||||||
}
|
|
||||||
|
|
||||||
#[test]
|
|
||||||
fn asset_category_parse_maps_known_names() {
|
|
||||||
assert_eq!(AssetCategory::parse("agents"), Some(AssetCategory::Agents));
|
|
||||||
assert_eq!(AssetCategory::parse("macros"), Some(AssetCategory::Macros));
|
|
||||||
assert_eq!(
|
|
||||||
AssetCategory::parse("functions"),
|
|
||||||
Some(AssetCategory::Functions)
|
|
||||||
);
|
|
||||||
assert_eq!(
|
|
||||||
AssetCategory::parse("mcp_config"),
|
|
||||||
Some(AssetCategory::McpConfig)
|
|
||||||
);
|
|
||||||
assert_eq!(AssetCategory::parse("roles"), None);
|
|
||||||
assert_eq!(AssetCategory::parse(""), None);
|
|
||||||
}
|
|
||||||
|
|
||||||
#[test]
|
|
||||||
#[serial]
|
|
||||||
fn install_builtin_agents_force_overwrites_only_with_force() {
|
|
||||||
let _guard = TestConfigDirGuard::new();
|
|
||||||
|
|
||||||
Agent::install_builtin_agents(false).unwrap();
|
|
||||||
let file =
|
|
||||||
first_file(&paths::agents_data_dir()).expect("bundled agents should be installed");
|
|
||||||
|
|
||||||
write(&file, "SENTINEL").unwrap();
|
|
||||||
Agent::install_builtin_agents(false).unwrap();
|
|
||||||
assert_eq!(
|
|
||||||
read_to_string(&file).unwrap(),
|
|
||||||
"SENTINEL",
|
|
||||||
"non-force install must not overwrite an existing file"
|
|
||||||
);
|
|
||||||
|
|
||||||
Agent::install_builtin_agents(true).unwrap();
|
|
||||||
assert_ne!(
|
|
||||||
read_to_string(&file).unwrap(),
|
|
||||||
"SENTINEL",
|
|
||||||
"force install must overwrite the existing file"
|
|
||||||
);
|
|
||||||
}
|
|
||||||
|
|
||||||
#[test]
|
|
||||||
#[serial]
|
|
||||||
fn install_functions_force_preserves_user_mcp_json() {
|
|
||||||
let _guard = TestConfigDirGuard::new();
|
|
||||||
|
|
||||||
Functions::install_builtin_global_tools(false).unwrap();
|
|
||||||
let mcp = paths::mcp_config_file();
|
|
||||||
assert!(mcp.exists(), "mcp.json should be installed on first run");
|
|
||||||
|
|
||||||
write(&mcp, "USER_MCP_CONFIG").unwrap();
|
|
||||||
Functions::install_builtin_global_tools(true).unwrap();
|
|
||||||
assert_eq!(
|
|
||||||
read_to_string(&mcp).unwrap(),
|
|
||||||
"USER_MCP_CONFIG",
|
|
||||||
"force install must NOT overwrite the user's mcp.json"
|
|
||||||
);
|
|
||||||
}
|
|
||||||
|
|
||||||
#[test]
|
|
||||||
#[serial]
|
|
||||||
fn install_mcp_config_overwrites_existing() {
|
|
||||||
let _guard = TestConfigDirGuard::new();
|
|
||||||
|
|
||||||
Functions::install_mcp_config().unwrap();
|
|
||||||
let mcp = paths::mcp_config_file();
|
|
||||||
assert!(mcp.exists(), "install_mcp_config should create mcp.json");
|
|
||||||
|
|
||||||
write(&mcp, "USER_MCP_CONFIG").unwrap();
|
|
||||||
Functions::install_mcp_config().unwrap();
|
|
||||||
assert_ne!(
|
|
||||||
read_to_string(&mcp).unwrap(),
|
|
||||||
"USER_MCP_CONFIG",
|
|
||||||
"install_mcp_config must overwrite the existing mcp.json"
|
|
||||||
);
|
|
||||||
}
|
|
||||||
}
|
}
|
||||||
|
|||||||
@@ -55,14 +55,6 @@ pub struct Role {
|
|||||||
enabled_tools: Option<String>,
|
enabled_tools: Option<String>,
|
||||||
#[serde(skip_serializing_if = "Option::is_none")]
|
#[serde(skip_serializing_if = "Option::is_none")]
|
||||||
enabled_mcp_servers: Option<String>,
|
enabled_mcp_servers: Option<String>,
|
||||||
#[serde(skip_serializing_if = "Option::is_none")]
|
|
||||||
auto_continue: Option<bool>,
|
|
||||||
#[serde(skip_serializing_if = "Option::is_none")]
|
|
||||||
max_auto_continues: Option<usize>,
|
|
||||||
#[serde(skip_serializing_if = "Option::is_none")]
|
|
||||||
inject_todo_instructions: Option<bool>,
|
|
||||||
#[serde(skip_serializing_if = "Option::is_none")]
|
|
||||||
continuation_prompt: Option<String>,
|
|
||||||
|
|
||||||
#[serde(skip)]
|
#[serde(skip)]
|
||||||
model: Model,
|
model: Model,
|
||||||
@@ -98,14 +90,6 @@ impl Role {
|
|||||||
"enabled_mcp_servers" => {
|
"enabled_mcp_servers" => {
|
||||||
role.enabled_mcp_servers = value.as_str().map(|v| v.to_string())
|
role.enabled_mcp_servers = value.as_str().map(|v| v.to_string())
|
||||||
}
|
}
|
||||||
"auto_continue" => role.auto_continue = value.as_bool(),
|
|
||||||
"max_auto_continues" => {
|
|
||||||
role.max_auto_continues = value.as_u64().map(|v| v as usize)
|
|
||||||
}
|
|
||||||
"inject_todo_instructions" => role.inject_todo_instructions = value.as_bool(),
|
|
||||||
"continuation_prompt" => {
|
|
||||||
role.continuation_prompt = value.as_str().map(|v| v.to_string())
|
|
||||||
}
|
|
||||||
_ => (),
|
_ => (),
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
@@ -147,20 +131,6 @@ impl Role {
|
|||||||
if let Some(enabled_mcp_servers) = self.enabled_mcp_servers() {
|
if let Some(enabled_mcp_servers) = self.enabled_mcp_servers() {
|
||||||
metadata.push(format!("enabled_mcp_servers: {enabled_mcp_servers}"));
|
metadata.push(format!("enabled_mcp_servers: {enabled_mcp_servers}"));
|
||||||
}
|
}
|
||||||
if let Some(auto_continue) = self.auto_continue {
|
|
||||||
metadata.push(format!("auto_continue: {auto_continue}"));
|
|
||||||
}
|
|
||||||
if let Some(max_auto_continues) = self.max_auto_continues {
|
|
||||||
metadata.push(format!("max_auto_continues: {max_auto_continues}"));
|
|
||||||
}
|
|
||||||
if let Some(inject_todo_instructions) = self.inject_todo_instructions {
|
|
||||||
metadata.push(format!(
|
|
||||||
"inject_todo_instructions: {inject_todo_instructions}"
|
|
||||||
));
|
|
||||||
}
|
|
||||||
if let Some(continuation_prompt) = &self.continuation_prompt {
|
|
||||||
metadata.push(format!("continuation_prompt: {continuation_prompt}"));
|
|
||||||
}
|
|
||||||
if metadata.is_empty() {
|
if metadata.is_empty() {
|
||||||
format!("{}\n", self.prompt)
|
format!("{}\n", self.prompt)
|
||||||
} else if self.prompt.is_empty() {
|
} else if self.prompt.is_empty() {
|
||||||
@@ -255,26 +225,6 @@ impl Role {
|
|||||||
self.prompt.contains(INPUT_PLACEHOLDER)
|
self.prompt.contains(INPUT_PLACEHOLDER)
|
||||||
}
|
}
|
||||||
|
|
||||||
pub fn auto_continue(&self) -> Option<bool> {
|
|
||||||
self.auto_continue
|
|
||||||
}
|
|
||||||
|
|
||||||
pub fn max_auto_continues(&self) -> Option<usize> {
|
|
||||||
self.max_auto_continues
|
|
||||||
}
|
|
||||||
|
|
||||||
pub fn inject_todo_instructions(&self) -> Option<bool> {
|
|
||||||
self.inject_todo_instructions
|
|
||||||
}
|
|
||||||
|
|
||||||
pub fn continuation_prompt(&self) -> Option<&str> {
|
|
||||||
self.continuation_prompt.as_deref()
|
|
||||||
}
|
|
||||||
|
|
||||||
pub fn append_to_prompt(&mut self, text: &str) {
|
|
||||||
self.prompt.push_str(text);
|
|
||||||
}
|
|
||||||
|
|
||||||
pub fn echo_messages(&self, input: &Input) -> String {
|
pub fn echo_messages(&self, input: &Input) -> String {
|
||||||
let input_markdown = input.render();
|
let input_markdown = input.render();
|
||||||
if self.is_empty_prompt() {
|
if self.is_empty_prompt() {
|
||||||
|
|||||||
+9
-80
@@ -32,14 +32,6 @@ pub struct Session {
|
|||||||
save_session: Option<bool>,
|
save_session: Option<bool>,
|
||||||
#[serde(skip_serializing_if = "Option::is_none")]
|
#[serde(skip_serializing_if = "Option::is_none")]
|
||||||
compression_threshold: Option<usize>,
|
compression_threshold: Option<usize>,
|
||||||
#[serde(skip_serializing_if = "Option::is_none")]
|
|
||||||
auto_continue: Option<bool>,
|
|
||||||
#[serde(skip_serializing_if = "Option::is_none")]
|
|
||||||
max_auto_continues: Option<usize>,
|
|
||||||
#[serde(skip_serializing_if = "Option::is_none")]
|
|
||||||
inject_todo_instructions: Option<bool>,
|
|
||||||
#[serde(skip_serializing_if = "Option::is_none")]
|
|
||||||
continuation_prompt: Option<String>,
|
|
||||||
|
|
||||||
#[serde(skip_serializing_if = "Option::is_none")]
|
#[serde(skip_serializing_if = "Option::is_none")]
|
||||||
role_name: Option<String>,
|
role_name: Option<String>,
|
||||||
@@ -178,18 +170,6 @@ impl Session {
|
|||||||
if let Some(save_session) = self.save_session() {
|
if let Some(save_session) = self.save_session() {
|
||||||
data["save_session"] = save_session.into();
|
data["save_session"] = save_session.into();
|
||||||
}
|
}
|
||||||
if let Some(auto_continue) = self.auto_continue() {
|
|
||||||
data["auto_continue"] = auto_continue.into();
|
|
||||||
}
|
|
||||||
if let Some(max_auto_continues) = self.max_auto_continues() {
|
|
||||||
data["max_auto_continues"] = max_auto_continues.into();
|
|
||||||
}
|
|
||||||
if let Some(inject_todo_instructions) = self.inject_todo_instructions() {
|
|
||||||
data["inject_todo_instructions"] = inject_todo_instructions.into();
|
|
||||||
}
|
|
||||||
if let Some(continuation_prompt) = self.continuation_prompt() {
|
|
||||||
data["continuation_prompt"] = continuation_prompt.into();
|
|
||||||
}
|
|
||||||
let (tokens, percent) = self.tokens_usage();
|
let (tokens, percent) = self.tokens_usage();
|
||||||
data["total_tokens"] = tokens.into();
|
data["total_tokens"] = tokens.into();
|
||||||
if let Some(max_input_tokens) = self.model().max_input_tokens() {
|
if let Some(max_input_tokens) = self.model().max_input_tokens() {
|
||||||
@@ -245,22 +225,6 @@ impl Session {
|
|||||||
items.push(("compression_threshold", compression_threshold.to_string()));
|
items.push(("compression_threshold", compression_threshold.to_string()));
|
||||||
}
|
}
|
||||||
|
|
||||||
if let Some(auto_continue) = self.auto_continue() {
|
|
||||||
items.push(("auto_continue", auto_continue.to_string()));
|
|
||||||
}
|
|
||||||
if let Some(max_auto_continues) = self.max_auto_continues() {
|
|
||||||
items.push(("max_auto_continues", max_auto_continues.to_string()));
|
|
||||||
}
|
|
||||||
if let Some(inject_todo_instructions) = self.inject_todo_instructions() {
|
|
||||||
items.push((
|
|
||||||
"inject_todo_instructions",
|
|
||||||
inject_todo_instructions.to_string(),
|
|
||||||
));
|
|
||||||
}
|
|
||||||
if let Some(continuation_prompt) = self.continuation_prompt() {
|
|
||||||
items.push(("continuation_prompt", continuation_prompt.to_string()));
|
|
||||||
}
|
|
||||||
|
|
||||||
if let Some(max_input_tokens) = self.model().max_input_tokens() {
|
if let Some(max_input_tokens) = self.model().max_input_tokens() {
|
||||||
items.push(("max_input_tokens", max_input_tokens.to_string()));
|
items.push(("max_input_tokens", max_input_tokens.to_string()));
|
||||||
}
|
}
|
||||||
@@ -371,50 +335,6 @@ impl Session {
|
|||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
pub fn auto_continue(&self) -> Option<bool> {
|
|
||||||
self.auto_continue
|
|
||||||
}
|
|
||||||
|
|
||||||
pub fn max_auto_continues(&self) -> Option<usize> {
|
|
||||||
self.max_auto_continues
|
|
||||||
}
|
|
||||||
|
|
||||||
pub fn set_auto_continue(&mut self, value: Option<bool>) {
|
|
||||||
if self.auto_continue != value {
|
|
||||||
self.auto_continue = value;
|
|
||||||
self.dirty = true;
|
|
||||||
}
|
|
||||||
}
|
|
||||||
|
|
||||||
pub fn set_max_auto_continues(&mut self, value: Option<usize>) {
|
|
||||||
if self.max_auto_continues != value {
|
|
||||||
self.max_auto_continues = value;
|
|
||||||
self.dirty = true;
|
|
||||||
}
|
|
||||||
}
|
|
||||||
|
|
||||||
pub fn inject_todo_instructions(&self) -> Option<bool> {
|
|
||||||
self.inject_todo_instructions
|
|
||||||
}
|
|
||||||
|
|
||||||
pub fn continuation_prompt(&self) -> Option<&str> {
|
|
||||||
self.continuation_prompt.as_deref()
|
|
||||||
}
|
|
||||||
|
|
||||||
pub fn set_inject_todo_instructions(&mut self, value: Option<bool>) {
|
|
||||||
if self.inject_todo_instructions != value {
|
|
||||||
self.inject_todo_instructions = value;
|
|
||||||
self.dirty = true;
|
|
||||||
}
|
|
||||||
}
|
|
||||||
|
|
||||||
pub fn set_continuation_prompt(&mut self, value: Option<String>) {
|
|
||||||
if self.continuation_prompt != value {
|
|
||||||
self.continuation_prompt = value;
|
|
||||||
self.dirty = true;
|
|
||||||
}
|
|
||||||
}
|
|
||||||
|
|
||||||
pub fn needs_compression(&self, global_compression_threshold: usize) -> bool {
|
pub fn needs_compression(&self, global_compression_threshold: usize) -> bool {
|
||||||
if self.compressing {
|
if self.compressing {
|
||||||
return false;
|
return false;
|
||||||
@@ -628,6 +548,15 @@ impl Session {
|
|||||||
let mut messages = self.messages.clone();
|
let mut messages = self.messages.clone();
|
||||||
if input.continue_output().is_some() {
|
if input.continue_output().is_some() {
|
||||||
return messages;
|
return messages;
|
||||||
|
} else if input.regenerate() {
|
||||||
|
while let Some(last) = messages.last() {
|
||||||
|
if !last.role.is_user() {
|
||||||
|
messages.pop();
|
||||||
|
} else {
|
||||||
|
break;
|
||||||
|
}
|
||||||
|
}
|
||||||
|
return messages;
|
||||||
}
|
}
|
||||||
let mut need_add_msg = true;
|
let mut need_add_msg = true;
|
||||||
let len = messages.len();
|
let len = messages.len();
|
||||||
|
|||||||
@@ -8,7 +8,6 @@ use serde_json::{Value, json};
|
|||||||
use std::collections::HashMap;
|
use std::collections::HashMap;
|
||||||
use std::sync::Arc;
|
use std::sync::Arc;
|
||||||
|
|
||||||
#[derive(Clone)]
|
|
||||||
pub struct ToolScope {
|
pub struct ToolScope {
|
||||||
pub functions: Functions,
|
pub functions: Functions,
|
||||||
pub mcp_runtime: McpRuntime,
|
pub mcp_runtime: McpRuntime,
|
||||||
@@ -25,7 +24,7 @@ impl Default for ToolScope {
|
|||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
#[derive(Default, Clone)]
|
#[derive(Default)]
|
||||||
pub struct McpRuntime {
|
pub struct McpRuntime {
|
||||||
pub servers: HashMap<String, Arc<ConnectedServer>>,
|
pub servers: HashMap<String, Arc<ConnectedServer>>,
|
||||||
}
|
}
|
||||||
|
|||||||
@@ -1,244 +0,0 @@
|
|||||||
use crate::utils::warning_text;
|
|
||||||
use anyhow::{Context, Result, bail};
|
|
||||||
use dunce::canonicalize;
|
|
||||||
use inquire::Confirm;
|
|
||||||
use is_terminal::IsTerminal;
|
|
||||||
use self_update::Status;
|
|
||||||
use self_update::backends::github::Update;
|
|
||||||
use std::fs::OpenOptions;
|
|
||||||
use std::path::Path;
|
|
||||||
use std::{env, fs, io, process};
|
|
||||||
|
|
||||||
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
|
|
||||||
enum InstallSource {
|
|
||||||
Cargo,
|
|
||||||
Homebrew,
|
|
||||||
Manual,
|
|
||||||
}
|
|
||||||
|
|
||||||
impl InstallSource {
|
|
||||||
fn is_package_managed(self) -> bool {
|
|
||||||
matches!(self, InstallSource::Cargo | InstallSource::Homebrew)
|
|
||||||
}
|
|
||||||
|
|
||||||
fn label(self) -> &'static str {
|
|
||||||
match self {
|
|
||||||
InstallSource::Cargo => "Cargo",
|
|
||||||
InstallSource::Homebrew => "Homebrew",
|
|
||||||
InstallSource::Manual => "manually-installed",
|
|
||||||
}
|
|
||||||
}
|
|
||||||
}
|
|
||||||
|
|
||||||
fn classify_install_path(path: &Path) -> InstallSource {
|
|
||||||
let components: Vec<&str> = path
|
|
||||||
.components()
|
|
||||||
.filter_map(|c| c.as_os_str().to_str())
|
|
||||||
.collect();
|
|
||||||
|
|
||||||
if components
|
|
||||||
.windows(2)
|
|
||||||
.any(|w| w[0] == ".cargo" && w[1] == "bin")
|
|
||||||
{
|
|
||||||
return InstallSource::Cargo;
|
|
||||||
}
|
|
||||||
|
|
||||||
if components.contains(&"Cellar") {
|
|
||||||
return InstallSource::Homebrew;
|
|
||||||
}
|
|
||||||
let path_str = path.to_string_lossy();
|
|
||||||
if path_str.starts_with("/opt/homebrew/") || path_str.starts_with("/home/linuxbrew/.linuxbrew/")
|
|
||||||
{
|
|
||||||
return InstallSource::Homebrew;
|
|
||||||
}
|
|
||||||
|
|
||||||
InstallSource::Manual
|
|
||||||
}
|
|
||||||
|
|
||||||
fn normalize_version(requested: Option<String>) -> Option<String> {
|
|
||||||
let raw = requested?;
|
|
||||||
let trimmed = raw.trim();
|
|
||||||
if trimmed.is_empty() || trimmed.eq_ignore_ascii_case("latest") {
|
|
||||||
return None;
|
|
||||||
}
|
|
||||||
match trimmed.chars().next() {
|
|
||||||
Some('v' | 'V') => Some(trimmed.to_string()),
|
|
||||||
Some(c) if c.is_ascii_digit() => Some(format!("v{trimmed}")),
|
|
||||||
_ => Some(trimmed.to_string()),
|
|
||||||
}
|
|
||||||
}
|
|
||||||
|
|
||||||
fn is_dir_writable(dir: &Path) -> bool {
|
|
||||||
let probe = dir.join(format!(".loki-update-write-test-{}", process::id()));
|
|
||||||
match OpenOptions::new().write(true).create_new(true).open(&probe) {
|
|
||||||
Ok(_) => {
|
|
||||||
let _ = fs::remove_file(&probe);
|
|
||||||
true
|
|
||||||
}
|
|
||||||
Err(_) => false,
|
|
||||||
}
|
|
||||||
}
|
|
||||||
|
|
||||||
pub fn run_self_update(requested: Option<String>, force: bool) -> Result<()> {
|
|
||||||
let target_tag = normalize_version(requested);
|
|
||||||
|
|
||||||
let exe_path = env::current_exe()
|
|
||||||
.context("Could not determine the path of the running loki executable")?;
|
|
||||||
let resolved = canonicalize(&exe_path).unwrap_or_else(|_| exe_path.clone());
|
|
||||||
let source = classify_install_path(&resolved);
|
|
||||||
|
|
||||||
if source.is_package_managed() {
|
|
||||||
let body = match source {
|
|
||||||
InstallSource::Homebrew => format!(
|
|
||||||
"Loki appears to be installed via Homebrew ({}).\n\
|
|
||||||
Updating in place replaces the binary inside Homebrew's Cellar; `brew` will\n\
|
|
||||||
then report a version that no longer matches the file on disk, and a later\n\
|
|
||||||
`brew upgrade`/`brew reinstall` may overwrite it or fail.\n\
|
|
||||||
The clean way to update is: brew upgrade loki",
|
|
||||||
exe_path.display()
|
|
||||||
),
|
|
||||||
InstallSource::Cargo => format!(
|
|
||||||
"Loki appears to be installed via `cargo install` ({}).\n\
|
|
||||||
Updating in place leaves Cargo's records out of sync with the binary on disk.\n\
|
|
||||||
The clean way to update is: cargo install --locked loki-ai",
|
|
||||||
exe_path.display()
|
|
||||||
),
|
|
||||||
InstallSource::Manual => unreachable!("Manual installs are not package-managed"),
|
|
||||||
};
|
|
||||||
println!("{} {body}", warning_text("WARNING:"));
|
|
||||||
|
|
||||||
if force {
|
|
||||||
println!("--force specified; updating anyway.");
|
|
||||||
} else if io::stdin().is_terminal() {
|
|
||||||
let proceed = Confirm::new("Update anyway?")
|
|
||||||
.with_default(false)
|
|
||||||
.prompt()?;
|
|
||||||
if !proceed {
|
|
||||||
println!("Update cancelled.");
|
|
||||||
return Ok(());
|
|
||||||
}
|
|
||||||
} else {
|
|
||||||
bail!(
|
|
||||||
"Refusing to update a {} install. Re-run with --force to override.",
|
|
||||||
source.label()
|
|
||||||
);
|
|
||||||
}
|
|
||||||
}
|
|
||||||
|
|
||||||
if let Some(parent) = exe_path.parent()
|
|
||||||
&& !is_dir_writable(parent)
|
|
||||||
{
|
|
||||||
bail!(
|
|
||||||
"No write permission for '{}'. Re-run with elevated permissions (e.g. sudo), \
|
|
||||||
or update Loki through your package manager.",
|
|
||||||
parent.display()
|
|
||||||
);
|
|
||||||
}
|
|
||||||
|
|
||||||
let interactive = io::stdin().is_terminal();
|
|
||||||
let mut builder = Update::configure();
|
|
||||||
builder
|
|
||||||
.repo_owner("Dark-Alex-17")
|
|
||||||
.repo_name("loki")
|
|
||||||
.bin_name("loki")
|
|
||||||
.current_version(env!("CARGO_PKG_VERSION"))
|
|
||||||
.no_confirm(true)
|
|
||||||
.show_download_progress(interactive);
|
|
||||||
if let Some(tag) = &target_tag {
|
|
||||||
builder.target_version_tag(tag.as_str());
|
|
||||||
}
|
|
||||||
let status = builder
|
|
||||||
.build()
|
|
||||||
.context("Failed to configure the self-update")?
|
|
||||||
.update()
|
|
||||||
.context("Self-update failed")?;
|
|
||||||
|
|
||||||
match status {
|
|
||||||
Status::UpToDate(version) => {
|
|
||||||
println!("Loki is already up to date (v{version}).");
|
|
||||||
}
|
|
||||||
Status::Updated(version) => {
|
|
||||||
println!("Loki updated to v{version}. Restart loki to use the new version.");
|
|
||||||
}
|
|
||||||
}
|
|
||||||
Ok(())
|
|
||||||
}
|
|
||||||
|
|
||||||
#[cfg(test)]
|
|
||||||
mod tests {
|
|
||||||
use super::*;
|
|
||||||
use std::path::PathBuf;
|
|
||||||
|
|
||||||
#[test]
|
|
||||||
fn classify_cargo_install() {
|
|
||||||
assert_eq!(
|
|
||||||
classify_install_path(&PathBuf::from("/home/u/.cargo/bin/loki")),
|
|
||||||
InstallSource::Cargo
|
|
||||||
);
|
|
||||||
}
|
|
||||||
|
|
||||||
#[test]
|
|
||||||
fn classify_homebrew_opt_prefix() {
|
|
||||||
assert_eq!(
|
|
||||||
classify_install_path(&PathBuf::from("/opt/homebrew/bin/loki")),
|
|
||||||
InstallSource::Homebrew
|
|
||||||
);
|
|
||||||
}
|
|
||||||
|
|
||||||
#[test]
|
|
||||||
fn classify_homebrew_cellar() {
|
|
||||||
assert_eq!(
|
|
||||||
classify_install_path(&PathBuf::from("/usr/local/Cellar/loki/0.3.0/bin/loki")),
|
|
||||||
InstallSource::Homebrew
|
|
||||||
);
|
|
||||||
}
|
|
||||||
|
|
||||||
#[test]
|
|
||||||
fn classify_homebrew_linuxbrew() {
|
|
||||||
assert_eq!(
|
|
||||||
classify_install_path(&PathBuf::from("/home/linuxbrew/.linuxbrew/bin/loki")),
|
|
||||||
InstallSource::Homebrew
|
|
||||||
);
|
|
||||||
}
|
|
||||||
|
|
||||||
#[test]
|
|
||||||
fn classify_manual_usr_local_bin() {
|
|
||||||
assert_eq!(
|
|
||||||
classify_install_path(&PathBuf::from("/usr/local/bin/loki")),
|
|
||||||
InstallSource::Manual
|
|
||||||
);
|
|
||||||
}
|
|
||||||
|
|
||||||
#[test]
|
|
||||||
fn classify_manual_local_bin() {
|
|
||||||
assert_eq!(
|
|
||||||
classify_install_path(&PathBuf::from("/home/u/.local/bin/loki")),
|
|
||||||
InstallSource::Manual
|
|
||||||
);
|
|
||||||
}
|
|
||||||
|
|
||||||
#[test]
|
|
||||||
fn normalize_version_latest_and_empty_are_none() {
|
|
||||||
assert_eq!(normalize_version(None), None);
|
|
||||||
assert_eq!(normalize_version(Some(String::new())), None);
|
|
||||||
assert_eq!(normalize_version(Some(" ".to_string())), None);
|
|
||||||
assert_eq!(normalize_version(Some("latest".to_string())), None);
|
|
||||||
assert_eq!(normalize_version(Some("LATEST".to_string())), None);
|
|
||||||
}
|
|
||||||
|
|
||||||
#[test]
|
|
||||||
fn normalize_version_prepends_v_for_bare_semver() {
|
|
||||||
assert_eq!(
|
|
||||||
normalize_version(Some("0.4.0".to_string())),
|
|
||||||
Some("v0.4.0".to_string())
|
|
||||||
);
|
|
||||||
assert_eq!(
|
|
||||||
normalize_version(Some("v0.4.0".to_string())),
|
|
||||||
Some("v0.4.0".to_string())
|
|
||||||
);
|
|
||||||
assert_eq!(
|
|
||||||
normalize_version(Some(" v0.4.0 ".to_string())),
|
|
||||||
Some("v0.4.0".to_string())
|
|
||||||
);
|
|
||||||
}
|
|
||||||
}
|
|
||||||
+8
-58
@@ -51,7 +51,7 @@ enum BinaryType<'a> {
|
|||||||
}
|
}
|
||||||
|
|
||||||
#[derive(Debug, Clone, Copy, PartialEq, Eq, AsRefStr)]
|
#[derive(Debug, Clone, Copy, PartialEq, Eq, AsRefStr)]
|
||||||
pub enum Language {
|
enum Language {
|
||||||
Bash,
|
Bash,
|
||||||
Python,
|
Python,
|
||||||
TypeScript,
|
TypeScript,
|
||||||
@@ -60,13 +60,7 @@ pub enum Language {
|
|||||||
|
|
||||||
impl From<&String> for Language {
|
impl From<&String> for Language {
|
||||||
fn from(s: &String) -> Self {
|
fn from(s: &String) -> Self {
|
||||||
Language::from_extension(s)
|
match s.to_lowercase().as_str() {
|
||||||
}
|
|
||||||
}
|
|
||||||
|
|
||||||
impl Language {
|
|
||||||
pub fn from_extension(ext: &str) -> Self {
|
|
||||||
match ext.to_lowercase().as_str() {
|
|
||||||
"sh" => Language::Bash,
|
"sh" => Language::Bash,
|
||||||
"py" => Language::Python,
|
"py" => Language::Python,
|
||||||
"ts" => Language::TypeScript,
|
"ts" => Language::TypeScript,
|
||||||
@@ -96,17 +90,6 @@ impl Language {
|
|||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
impl Language {
|
|
||||||
pub fn direct_invoker(self) -> Option<(&'static str, &'static [&'static str])> {
|
|
||||||
match self {
|
|
||||||
Language::Bash => Some(("bash", &[])),
|
|
||||||
Language::Python => Some(("python3", &[])),
|
|
||||||
Language::TypeScript => Some(("npx", &["tsx"])),
|
|
||||||
Language::Unsupported => None,
|
|
||||||
}
|
|
||||||
}
|
|
||||||
}
|
|
||||||
|
|
||||||
fn extract_shebang_runtime(path: &Path) -> Option<String> {
|
fn extract_shebang_runtime(path: &Path) -> Option<String> {
|
||||||
let file = File::open(path).ok()?;
|
let file = File::open(path).ok()?;
|
||||||
let reader = io::BufReader::new(file);
|
let reader = io::BufReader::new(file);
|
||||||
@@ -209,7 +192,7 @@ pub struct Functions {
|
|||||||
}
|
}
|
||||||
|
|
||||||
impl Functions {
|
impl Functions {
|
||||||
pub fn install_builtin_global_tools(force: bool) -> Result<()> {
|
pub fn install_builtin_global_tools() -> Result<()> {
|
||||||
info!(
|
info!(
|
||||||
"Installing global built-in functions in {}",
|
"Installing global built-in functions in {}",
|
||||||
paths::functions_dir().display()
|
paths::functions_dir().display()
|
||||||
@@ -227,14 +210,14 @@ impl Functions {
|
|||||||
})?;
|
})?;
|
||||||
let content = unsafe { std::str::from_utf8_unchecked(&embedded_file.data) };
|
let content = unsafe { std::str::from_utf8_unchecked(&embedded_file.data) };
|
||||||
let file_path = paths::functions_dir().join(file.as_ref());
|
let file_path = paths::functions_dir().join(file.as_ref());
|
||||||
#[cfg_attr(not(unix), expect(unused))]
|
let file_extension = file_path
|
||||||
let is_script = file_path
|
|
||||||
.extension()
|
.extension()
|
||||||
.and_then(OsStr::to_str)
|
.and_then(OsStr::to_str)
|
||||||
.is_some_and(|ext| Language::from_extension(ext) != Language::Unsupported);
|
.map(|s| s.to_lowercase());
|
||||||
|
#[cfg_attr(not(unix), expect(unused))]
|
||||||
|
let is_script = matches!(file_extension.as_deref(), Some("sh") | Some("py"));
|
||||||
|
|
||||||
let force_this = force && file.as_ref() != "mcp.json";
|
if file_path.exists() {
|
||||||
if file_path.exists() && !force_this {
|
|
||||||
debug!(
|
debug!(
|
||||||
"Function file already exists, skipping: {}",
|
"Function file already exists, skipping: {}",
|
||||||
file_path.display()
|
file_path.display()
|
||||||
@@ -257,22 +240,6 @@ impl Functions {
|
|||||||
Ok(())
|
Ok(())
|
||||||
}
|
}
|
||||||
|
|
||||||
pub fn install_mcp_config() -> Result<()> {
|
|
||||||
let file_path = paths::mcp_config_file();
|
|
||||||
let embedded = FunctionAssets::get("mcp.json")
|
|
||||||
.ok_or_else(|| anyhow!("Failed to load embedded mcp.json"))?;
|
|
||||||
let content = unsafe { std::str::from_utf8_unchecked(&embedded.data) };
|
|
||||||
|
|
||||||
ensure_parent_exists(&file_path)?;
|
|
||||||
|
|
||||||
info!("Reinstalling MCP config file: {}", file_path.display());
|
|
||||||
|
|
||||||
let mut config_file = File::create(&file_path)?;
|
|
||||||
config_file.write_all(content.as_bytes())?;
|
|
||||||
|
|
||||||
Ok(())
|
|
||||||
}
|
|
||||||
|
|
||||||
pub fn init(visible_tools: &[String]) -> Result<Self> {
|
pub fn init(visible_tools: &[String]) -> Result<Self> {
|
||||||
Self::clear_global_functions_bin_dir()?;
|
Self::clear_global_functions_bin_dir()?;
|
||||||
|
|
||||||
@@ -1448,23 +1415,6 @@ mod tests {
|
|||||||
assert!(tc.thought_signature.is_none());
|
assert!(tc.thought_signature.is_none());
|
||||||
}
|
}
|
||||||
|
|
||||||
#[test]
|
|
||||||
fn direct_invoker_maps_each_language() {
|
|
||||||
assert_eq!(
|
|
||||||
Language::Bash.direct_invoker(),
|
|
||||||
Some(("bash", &[] as &[&str]))
|
|
||||||
);
|
|
||||||
assert_eq!(
|
|
||||||
Language::Python.direct_invoker(),
|
|
||||||
Some(("python3", &[] as &[&str]))
|
|
||||||
);
|
|
||||||
assert_eq!(
|
|
||||||
Language::TypeScript.direct_invoker(),
|
|
||||||
Some(("npx", &["tsx"] as &[&str]))
|
|
||||||
);
|
|
||||||
assert_eq!(Language::Unsupported.direct_invoker(), None);
|
|
||||||
}
|
|
||||||
|
|
||||||
#[test]
|
#[test]
|
||||||
fn toolcall_with_thought_signature() {
|
fn toolcall_with_thought_signature() {
|
||||||
let tc = ToolCall::new("t".into(), json!({}), None)
|
let tc = ToolCall::new("t".into(), json!({}), None)
|
||||||
|
|||||||
+13
-167
@@ -5,7 +5,6 @@ use crate::supervisor::mailbox::{Envelope, EnvelopePayload, Inbox};
|
|||||||
use crate::supervisor::{AgentExitStatus, AgentHandle, AgentResult, Supervisor};
|
use crate::supervisor::{AgentExitStatus, AgentHandle, AgentResult, Supervisor};
|
||||||
use crate::utils::{AbortSignal, create_abort_signal};
|
use crate::utils::{AbortSignal, create_abort_signal};
|
||||||
|
|
||||||
use crate::graph;
|
|
||||||
use anyhow::{Context, Result, anyhow, bail};
|
use anyhow::{Context, Result, anyhow, bail};
|
||||||
use chrono::Utc;
|
use chrono::Utc;
|
||||||
use indexmap::IndexMap;
|
use indexmap::IndexMap;
|
||||||
@@ -14,8 +13,6 @@ use parking_lot::RwLock;
|
|||||||
use serde_json::{Value, json};
|
use serde_json::{Value, json};
|
||||||
use std::pin::Pin;
|
use std::pin::Pin;
|
||||||
use std::sync::Arc;
|
use std::sync::Arc;
|
||||||
use std::time::Duration;
|
|
||||||
use tokio::time;
|
|
||||||
use uuid::Uuid;
|
use uuid::Uuid;
|
||||||
|
|
||||||
pub const SUPERVISOR_FUNCTION_PREFIX: &str = "agent__";
|
pub const SUPERVISOR_FUNCTION_PREFIX: &str = "agent__";
|
||||||
@@ -327,21 +324,12 @@ pub async fn handle_supervisor_tool(
|
|||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
pub fn run_child_agent(
|
fn run_child_agent(
|
||||||
mut child_ctx: RequestContext,
|
mut child_ctx: RequestContext,
|
||||||
initial_input: Input,
|
initial_input: Input,
|
||||||
abort_signal: AbortSignal,
|
abort_signal: AbortSignal,
|
||||||
) -> Pin<Box<dyn Future<Output = Result<String>> + Send>> {
|
) -> Pin<Box<dyn Future<Output = Result<String>> + Send>> {
|
||||||
Box::pin(async move {
|
Box::pin(async move {
|
||||||
if graph::active_agent_graph_name(&child_ctx).is_some() {
|
|
||||||
return graph::run_active_agent_graph(
|
|
||||||
&mut child_ctx,
|
|
||||||
&initial_input.text(),
|
|
||||||
abort_signal,
|
|
||||||
)
|
|
||||||
.await;
|
|
||||||
}
|
|
||||||
|
|
||||||
let mut accumulated_output = String::new();
|
let mut accumulated_output = String::new();
|
||||||
let mut input = initial_input;
|
let mut input = initial_input;
|
||||||
let app = Arc::clone(&child_ctx.app.config);
|
let app = Arc::clone(&child_ctx.app.config);
|
||||||
@@ -384,98 +372,6 @@ pub fn run_child_agent(
|
|||||||
})
|
})
|
||||||
}
|
}
|
||||||
|
|
||||||
/// Spawn an agent synchronously from a graph node and return its accumulated
|
|
||||||
/// output. This is similar to `handle_spawn` but runs the child agent in the
|
|
||||||
/// current task (no tokio::spawn, no supervisor handle registration) so the
|
|
||||||
/// graph executor can sequence agent nodes directly.
|
|
||||||
pub async fn run_agent_for_graph(
|
|
||||||
parent_ctx: &mut RequestContext,
|
|
||||||
agent_name: &str,
|
|
||||||
prompt: &str,
|
|
||||||
) -> Result<String> {
|
|
||||||
let short_uuid = &Uuid::new_v4().to_string()[..8];
|
|
||||||
let agent_id = format!("graph_agent_{agent_name}_{short_uuid}");
|
|
||||||
let current_depth = parent_ctx.current_depth + 1;
|
|
||||||
|
|
||||||
if let Some(supervisor) = parent_ctx.supervisor.as_ref().cloned() {
|
|
||||||
let max_depth = supervisor.read().max_depth();
|
|
||||||
if current_depth > max_depth {
|
|
||||||
bail!("Max agent depth exceeded ({current_depth}/{max_depth})");
|
|
||||||
}
|
|
||||||
}
|
|
||||||
|
|
||||||
if !parent_ctx.app.config.function_calling_support {
|
|
||||||
bail!("Function calling support must be enabled to spawn agents.");
|
|
||||||
}
|
|
||||||
|
|
||||||
let child_inbox = Arc::new(Inbox::new());
|
|
||||||
parent_ctx.ensure_root_escalation_queue();
|
|
||||||
let child_abort = create_abort_signal();
|
|
||||||
|
|
||||||
let app_config = Arc::clone(&parent_ctx.app.config);
|
|
||||||
let current_model = parent_ctx.current_model().clone();
|
|
||||||
let info_flag = parent_ctx.info_flag;
|
|
||||||
let child_app_state = Arc::new(AppState {
|
|
||||||
config: Arc::new(app_config.as_ref().clone()),
|
|
||||||
vault: parent_ctx.app.vault.clone(),
|
|
||||||
mcp_factory: parent_ctx.app.mcp_factory.clone(),
|
|
||||||
rag_cache: parent_ctx.app.rag_cache.clone(),
|
|
||||||
mcp_config: parent_ctx.app.mcp_config.clone(),
|
|
||||||
mcp_log_path: parent_ctx.app.mcp_log_path.clone(),
|
|
||||||
mcp_registry: parent_ctx.app.mcp_registry.clone(),
|
|
||||||
functions: parent_ctx.app.functions.clone(),
|
|
||||||
});
|
|
||||||
|
|
||||||
let agent = Agent::init(
|
|
||||||
app_config.as_ref(),
|
|
||||||
child_app_state.as_ref(),
|
|
||||||
¤t_model,
|
|
||||||
info_flag,
|
|
||||||
agent_name,
|
|
||||||
child_abort.clone(),
|
|
||||||
)
|
|
||||||
.await?;
|
|
||||||
|
|
||||||
let agent_mcp_servers = agent.mcp_server_names().to_vec();
|
|
||||||
let session = agent.agent_session().map(|v| v.to_string());
|
|
||||||
let should_init_supervisor = agent.can_spawn_agents();
|
|
||||||
let agent_max_concurrent = agent.max_concurrent_agents();
|
|
||||||
let agent_max_depth = agent.max_agent_depth();
|
|
||||||
|
|
||||||
let mut child_ctx = RequestContext::new_for_child(
|
|
||||||
Arc::clone(&child_app_state),
|
|
||||||
parent_ctx,
|
|
||||||
current_depth,
|
|
||||||
Arc::clone(&child_inbox),
|
|
||||||
agent_id.clone(),
|
|
||||||
);
|
|
||||||
child_ctx.rag = agent.rag();
|
|
||||||
child_ctx.agent = Some(agent);
|
|
||||||
if should_init_supervisor {
|
|
||||||
child_ctx.supervisor = Some(Arc::new(RwLock::new(Supervisor::new(
|
|
||||||
agent_max_concurrent,
|
|
||||||
agent_max_depth,
|
|
||||||
))));
|
|
||||||
}
|
|
||||||
|
|
||||||
if let Some(session) = session {
|
|
||||||
child_ctx
|
|
||||||
.use_session(app_config.as_ref(), Some(&session), child_abort.clone())
|
|
||||||
.await?;
|
|
||||||
sync_agent_functions_to_ctx(&mut child_ctx)?;
|
|
||||||
} else {
|
|
||||||
populate_agent_mcp_runtime(&mut child_ctx, &agent_mcp_servers).await?;
|
|
||||||
sync_agent_functions_to_ctx(&mut child_ctx)?;
|
|
||||||
child_ctx.init_agent_shared_variables()?;
|
|
||||||
}
|
|
||||||
|
|
||||||
let input = Input::from_str(&child_ctx, prompt, None);
|
|
||||||
|
|
||||||
debug!("Spawning agent '{agent_name}' for graph node as '{agent_id}'");
|
|
||||||
|
|
||||||
run_child_agent(child_ctx, input, child_abort).await
|
|
||||||
}
|
|
||||||
|
|
||||||
async fn populate_agent_mcp_runtime(ctx: &mut RequestContext, server_ids: &[String]) -> Result<()> {
|
async fn populate_agent_mcp_runtime(ctx: &mut RequestContext, server_ids: &[String]) -> Result<()> {
|
||||||
if !ctx.app.config.mcp_server_support {
|
if !ctx.app.config.mcp_server_support {
|
||||||
return Ok(());
|
return Ok(());
|
||||||
@@ -705,25 +601,11 @@ async fn handle_check(ctx: &mut RequestContext, args: &Value) -> Result<Value> {
|
|||||||
|
|
||||||
match is_finished {
|
match is_finished {
|
||||||
Some(true) => handle_collect(ctx, args).await,
|
Some(true) => handle_collect(ctx, args).await,
|
||||||
Some(false) => {
|
Some(false) => Ok(json!({
|
||||||
let mut result = json!({
|
"status": "pending",
|
||||||
"status": "pending",
|
"id": id,
|
||||||
"id": id,
|
"message": "Agent is still running"
|
||||||
"message": "Agent is still running"
|
})),
|
||||||
});
|
|
||||||
|
|
||||||
if let Some(queue) = ctx.root_escalation_queue()
|
|
||||||
&& queue.has_pending()
|
|
||||||
{
|
|
||||||
let summary = queue.pending_summary();
|
|
||||||
result["pending_escalations"] = json!(summary);
|
|
||||||
result["message"] = json!(
|
|
||||||
"Agent is still running. Child agents have pending escalations that need your reply via agent__reply_escalation."
|
|
||||||
);
|
|
||||||
}
|
|
||||||
|
|
||||||
Ok(result)
|
|
||||||
}
|
|
||||||
None => Ok(json!({
|
None => Ok(json!({
|
||||||
"status": "error",
|
"status": "error",
|
||||||
"message": format!("No agent found with id '{id}'")
|
"message": format!("No agent found with id '{id}'")
|
||||||
@@ -737,48 +619,12 @@ async fn handle_collect(ctx: &mut RequestContext, args: &Value) -> Result<Value>
|
|||||||
.and_then(Value::as_str)
|
.and_then(Value::as_str)
|
||||||
.ok_or_else(|| anyhow!("'id' is required"))?;
|
.ok_or_else(|| anyhow!("'id' is required"))?;
|
||||||
|
|
||||||
let supervisor = ctx
|
|
||||||
.supervisor
|
|
||||||
.as_ref()
|
|
||||||
.cloned()
|
|
||||||
.ok_or_else(|| anyhow!("No supervisor active"))?;
|
|
||||||
|
|
||||||
{
|
|
||||||
let sup = supervisor.read();
|
|
||||||
if sup.is_finished(id).is_none() {
|
|
||||||
return Ok(json!({
|
|
||||||
"status": "error",
|
|
||||||
"message": format!("Agent '{id}' not found. Use agent__check to verify it exists and is finished.")
|
|
||||||
}));
|
|
||||||
}
|
|
||||||
}
|
|
||||||
|
|
||||||
loop {
|
|
||||||
let is_finished = {
|
|
||||||
let sup = supervisor.read();
|
|
||||||
sup.is_finished(id).unwrap_or(false)
|
|
||||||
};
|
|
||||||
|
|
||||||
if is_finished {
|
|
||||||
break;
|
|
||||||
}
|
|
||||||
|
|
||||||
if let Some(queue) = ctx.root_escalation_queue()
|
|
||||||
&& queue.has_pending()
|
|
||||||
{
|
|
||||||
let summary = queue.pending_summary();
|
|
||||||
return Ok(json!({
|
|
||||||
"status": "pending",
|
|
||||||
"id": id,
|
|
||||||
"message": format!("Agent '{id}' is still running, but child agents have pending escalations that need your reply. Reply via agent__reply_escalation, then call agent__collect again."),
|
|
||||||
"pending_escalations": summary,
|
|
||||||
}));
|
|
||||||
}
|
|
||||||
|
|
||||||
time::sleep(Duration::from_millis(200)).await;
|
|
||||||
}
|
|
||||||
|
|
||||||
let handle = {
|
let handle = {
|
||||||
|
let supervisor = ctx
|
||||||
|
.supervisor
|
||||||
|
.as_ref()
|
||||||
|
.cloned()
|
||||||
|
.ok_or_else(|| anyhow!("No supervisor active"))?;
|
||||||
let mut sup = supervisor.write();
|
let mut sup = supervisor.write();
|
||||||
sup.take(id)
|
sup.take(id)
|
||||||
};
|
};
|
||||||
@@ -803,7 +649,7 @@ async fn handle_collect(ctx: &mut RequestContext, args: &Value) -> Result<Value>
|
|||||||
}
|
}
|
||||||
None => Ok(json!({
|
None => Ok(json!({
|
||||||
"status": "error",
|
"status": "error",
|
||||||
"message": format!("Agent '{id}' completed but could not be collected. It may have been collected by another call.")
|
"message": format!("Agent '{id}' not found. Use agent__check to verify it exists and is finished.")
|
||||||
})),
|
})),
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
@@ -1347,7 +1193,7 @@ mod tests {
|
|||||||
let inbox = Arc::new(Inbox::new());
|
let inbox = Arc::new(Inbox::new());
|
||||||
let abort = create_abort_signal();
|
let abort = create_abort_signal();
|
||||||
let join_handle = tokio::spawn(async {
|
let join_handle = tokio::spawn(async {
|
||||||
time::sleep(Duration::from_secs(60)).await;
|
tokio::time::sleep(std::time::Duration::from_secs(60)).await;
|
||||||
Ok(AgentResult {
|
Ok(AgentResult {
|
||||||
id: "slow".into(),
|
id: "slow".into(),
|
||||||
agent_name: "test".into(),
|
agent_name: "test".into(),
|
||||||
|
|||||||
@@ -94,14 +94,8 @@ pub fn handle_todo_tool(ctx: &mut RequestContext, cmd_name: &str, args: &Value)
|
|||||||
.strip_prefix(TODO_FUNCTION_PREFIX)
|
.strip_prefix(TODO_FUNCTION_PREFIX)
|
||||||
.unwrap_or(cmd_name);
|
.unwrap_or(cmd_name);
|
||||||
|
|
||||||
if !ctx.app.config.function_calling_support {
|
if ctx.agent.is_none() {
|
||||||
bail!("Cannot use todo tools: function calling is disabled.");
|
bail!("No active agent");
|
||||||
}
|
|
||||||
let auto_config = ctx.auto_continue_config();
|
|
||||||
if !auto_config.enabled {
|
|
||||||
bail!(
|
|
||||||
"Auto-continue is not enabled. Set 'auto_continue: true' in your config to use todo tools."
|
|
||||||
);
|
|
||||||
}
|
}
|
||||||
|
|
||||||
match action {
|
match action {
|
||||||
|
|||||||
@@ -2,7 +2,7 @@ use super::{FunctionDeclaration, JsonSchema};
|
|||||||
use crate::config::RequestContext;
|
use crate::config::RequestContext;
|
||||||
use crate::supervisor::escalation::{EscalationRequest, new_escalation_id};
|
use crate::supervisor::escalation::{EscalationRequest, new_escalation_id};
|
||||||
|
|
||||||
use anyhow::{Result, anyhow, bail};
|
use anyhow::{Result, anyhow};
|
||||||
use indexmap::IndexMap;
|
use indexmap::IndexMap;
|
||||||
use inquire::{Confirm, MultiSelect, Select, Text};
|
use inquire::{Confirm, MultiSelect, Select, Text};
|
||||||
use serde_json::{Value, json};
|
use serde_json::{Value, json};
|
||||||
@@ -155,10 +155,7 @@ fn handle_direct_ask(args: &Value) -> Result<Value> {
|
|||||||
let mut options = parse_options(args)?;
|
let mut options = parse_options(args)?;
|
||||||
options.push(CUSTOM_MULTI_CHOICE_ANSWER_OPTION.to_string());
|
options.push(CUSTOM_MULTI_CHOICE_ANSWER_OPTION.to_string());
|
||||||
|
|
||||||
let mut answer = Select::new(question, options)
|
let mut answer = Select::new(question, options).prompt()?;
|
||||||
.without_filtering()
|
|
||||||
.with_help_message("↑↓ to move, enter to select")
|
|
||||||
.prompt()?;
|
|
||||||
|
|
||||||
if answer == CUSTOM_MULTI_CHOICE_ANSWER_OPTION {
|
if answer == CUSTOM_MULTI_CHOICE_ANSWER_OPTION {
|
||||||
answer = Text::new("Custom response:").prompt()?
|
answer = Text::new("Custom response:").prompt()?
|
||||||
@@ -208,11 +205,12 @@ async fn handle_escalated(ctx: &RequestContext, action: &str, args: &Value) -> R
|
|||||||
.ok_or_else(|| anyhow!("'question' is required"))?
|
.ok_or_else(|| anyhow!("'question' is required"))?
|
||||||
.to_string();
|
.to_string();
|
||||||
|
|
||||||
let options: Option<Vec<String>> = if args.get("options").is_some() {
|
let options: Option<Vec<String>> = args.get("options").and_then(Value::as_array).map(|arr| {
|
||||||
Some(parse_options(args)?)
|
arr.iter()
|
||||||
} else {
|
.filter_map(Value::as_str)
|
||||||
None
|
.map(String::from)
|
||||||
};
|
.collect()
|
||||||
|
});
|
||||||
|
|
||||||
let from_agent_id = ctx
|
let from_agent_id = ctx
|
||||||
.self_agent_id
|
.self_agent_id
|
||||||
@@ -264,24 +262,13 @@ async fn handle_escalated(ctx: &RequestContext, action: &str, args: &Value) -> R
|
|||||||
}
|
}
|
||||||
|
|
||||||
fn parse_options(args: &Value) -> Result<Vec<String>> {
|
fn parse_options(args: &Value) -> Result<Vec<String>> {
|
||||||
let raw = args
|
args.get("options")
|
||||||
.get("options")
|
.and_then(Value::as_array)
|
||||||
.ok_or_else(|| anyhow!("'options' is required and must be an array of strings"))?;
|
.map(|arr| {
|
||||||
|
arr.iter()
|
||||||
let arr: Vec<Value> = match raw {
|
.filter_map(Value::as_str)
|
||||||
Value::Array(arr) => arr.clone(),
|
.map(String::from)
|
||||||
Value::String(s) => serde_json::from_str::<Vec<Value>>(s).map_err(|_| {
|
.collect()
|
||||||
anyhow!(
|
})
|
||||||
"'options' was a string but did not parse as a JSON array. \
|
.ok_or_else(|| anyhow!("'options' is required and must be an array of strings"))
|
||||||
Pass options as a native JSON array, e.g. [\"yes\", \"no\"]."
|
|
||||||
)
|
|
||||||
})?,
|
|
||||||
_ => bail!("'options' is required and must be an array of strings"),
|
|
||||||
};
|
|
||||||
|
|
||||||
Ok(arr
|
|
||||||
.iter()
|
|
||||||
.filter_map(Value::as_str)
|
|
||||||
.map(String::from)
|
|
||||||
.collect())
|
|
||||||
}
|
}
|
||||||
|
|||||||
@@ -1,272 +0,0 @@
|
|||||||
use super::state::StateManager;
|
|
||||||
use super::structured;
|
|
||||||
use super::types::AgentNode;
|
|
||||||
use crate::config::RequestContext;
|
|
||||||
use crate::function::supervisor::run_agent_for_graph;
|
|
||||||
use anyhow::{Context, Result};
|
|
||||||
use serde_json::Value;
|
|
||||||
use std::time::Duration;
|
|
||||||
use tokio::time::timeout;
|
|
||||||
|
|
||||||
const OUTPUT_KEY: &str = "output";
|
|
||||||
const DEFAULT_TIMEOUT_SECS: u64 = 300;
|
|
||||||
|
|
||||||
pub struct AgentNodeExecutor;
|
|
||||||
|
|
||||||
impl AgentNodeExecutor {
|
|
||||||
pub async fn execute(
|
|
||||||
node: &AgentNode,
|
|
||||||
state_manager: &mut StateManager,
|
|
||||||
parent_ctx: &mut RequestContext,
|
|
||||||
) -> Result<String> {
|
|
||||||
let prompt = state_manager
|
|
||||||
.interpolate(&node.prompt)
|
|
||||||
.with_context(|| format!("Failed to interpolate prompt for agent '{}'", node.agent))?;
|
|
||||||
|
|
||||||
let timeout_dur = Duration::from_secs(node.timeout.unwrap_or(DEFAULT_TIMEOUT_SECS));
|
|
||||||
|
|
||||||
let raw = timeout(
|
|
||||||
timeout_dur,
|
|
||||||
run_agent_for_graph(parent_ctx, &node.agent, &prompt),
|
|
||||||
)
|
|
||||||
.await
|
|
||||||
.with_context(|| {
|
|
||||||
format!(
|
|
||||||
"Agent '{}' timed out after {}s",
|
|
||||||
node.agent,
|
|
||||||
timeout_dur.as_secs()
|
|
||||||
)
|
|
||||||
})?
|
|
||||||
.with_context(|| format!("Agent '{}' failed", node.agent))?;
|
|
||||||
|
|
||||||
let output_value = match &node.output_schema {
|
|
||||||
Some(schema) => structured::extract(&raw, schema, parent_ctx)
|
|
||||||
.await
|
|
||||||
.with_context(|| {
|
|
||||||
format!(
|
|
||||||
"Agent '{}' output failed structured-output extraction",
|
|
||||||
node.agent
|
|
||||||
)
|
|
||||||
})?,
|
|
||||||
None => Value::String(raw.clone()),
|
|
||||||
};
|
|
||||||
|
|
||||||
apply_state_updates(node, state_manager, &output_value);
|
|
||||||
|
|
||||||
Ok(raw)
|
|
||||||
}
|
|
||||||
}
|
|
||||||
|
|
||||||
fn apply_state_updates(node: &AgentNode, state_manager: &mut StateManager, output: &Value) {
|
|
||||||
if node.output_schema.is_some()
|
|
||||||
&& let Some(obj) = output.as_object()
|
|
||||||
{
|
|
||||||
for (k, v) in obj {
|
|
||||||
state_manager.state_mut().set(k.clone(), v.clone());
|
|
||||||
}
|
|
||||||
}
|
|
||||||
|
|
||||||
let Some(updates) = &node.state_updates else {
|
|
||||||
return;
|
|
||||||
};
|
|
||||||
let prev_output = state_manager.state().get(OUTPUT_KEY).cloned();
|
|
||||||
state_manager
|
|
||||||
.state_mut()
|
|
||||||
.set(OUTPUT_KEY.into(), output.clone());
|
|
||||||
|
|
||||||
for (key, template) in updates {
|
|
||||||
let value = state_manager.interpolate_lenient(template);
|
|
||||||
state_manager
|
|
||||||
.state_mut()
|
|
||||||
.set(key.clone(), Value::String(value));
|
|
||||||
}
|
|
||||||
|
|
||||||
match prev_output {
|
|
||||||
Some(v) => state_manager.state_mut().set(OUTPUT_KEY.into(), v),
|
|
||||||
None => {
|
|
||||||
state_manager
|
|
||||||
.state_mut()
|
|
||||||
.set(OUTPUT_KEY.into(), Value::Null);
|
|
||||||
}
|
|
||||||
}
|
|
||||||
}
|
|
||||||
|
|
||||||
#[cfg(test)]
|
|
||||||
mod tests {
|
|
||||||
use super::super::types::AgentNode;
|
|
||||||
use super::*;
|
|
||||||
use serde_json::json;
|
|
||||||
use std::collections::HashMap;
|
|
||||||
|
|
||||||
fn manager_with(pairs: &[(&str, Value)]) -> StateManager {
|
|
||||||
let mut map = HashMap::new();
|
|
||||||
for (k, v) in pairs {
|
|
||||||
map.insert((*k).into(), v.clone());
|
|
||||||
}
|
|
||||||
StateManager::new(map)
|
|
||||||
}
|
|
||||||
|
|
||||||
fn node_with(prompt: &str, updates: Option<HashMap<String, String>>) -> AgentNode {
|
|
||||||
AgentNode {
|
|
||||||
agent: "test_agent".into(),
|
|
||||||
prompt: prompt.into(),
|
|
||||||
state_updates: updates,
|
|
||||||
output_schema: None,
|
|
||||||
timeout: None,
|
|
||||||
}
|
|
||||||
}
|
|
||||||
|
|
||||||
#[test]
|
|
||||||
fn state_updates_use_output_placeholder() {
|
|
||||||
let node = {
|
|
||||||
let mut u = HashMap::new();
|
|
||||||
u.insert("findings".into(), "{{output}}".into());
|
|
||||||
node_with("hi", Some(u))
|
|
||||||
};
|
|
||||||
let mut state = manager_with(&[]);
|
|
||||||
|
|
||||||
apply_state_updates(&node, &mut state, &json!("agent finished its work"));
|
|
||||||
|
|
||||||
assert_eq!(
|
|
||||||
state.state().get("findings"),
|
|
||||||
Some(&json!("agent finished its work"))
|
|
||||||
);
|
|
||||||
}
|
|
||||||
|
|
||||||
#[test]
|
|
||||||
fn state_updates_can_reference_existing_keys_and_output() {
|
|
||||||
let node = {
|
|
||||||
let mut u = HashMap::new();
|
|
||||||
u.insert("summary".into(), "{{topic}}: {{output}}".into());
|
|
||||||
node_with("hi", Some(u))
|
|
||||||
};
|
|
||||||
let mut state = manager_with(&[("topic", json!("auth"))]);
|
|
||||||
|
|
||||||
apply_state_updates(&node, &mut state, &json!("JWT vs sessions"));
|
|
||||||
|
|
||||||
assert_eq!(
|
|
||||||
state.state().get("summary"),
|
|
||||||
Some(&json!("auth: JWT vs sessions"))
|
|
||||||
);
|
|
||||||
}
|
|
||||||
|
|
||||||
#[test]
|
|
||||||
fn output_key_is_cleaned_up_after_state_updates() {
|
|
||||||
let node = {
|
|
||||||
let mut u = HashMap::new();
|
|
||||||
u.insert("findings".into(), "{{output}}".into());
|
|
||||||
node_with("hi", Some(u))
|
|
||||||
};
|
|
||||||
let mut state = manager_with(&[]);
|
|
||||||
|
|
||||||
apply_state_updates(&node, &mut state, &json!("anything"));
|
|
||||||
|
|
||||||
assert_eq!(state.state().get("output"), Some(&Value::Null));
|
|
||||||
}
|
|
||||||
|
|
||||||
#[test]
|
|
||||||
fn pre_existing_output_value_is_preserved() {
|
|
||||||
let node = {
|
|
||||||
let mut u = HashMap::new();
|
|
||||||
u.insert("greeting".into(), "{{output}}".into());
|
|
||||||
node_with("hi", Some(u))
|
|
||||||
};
|
|
||||||
let mut state = manager_with(&[("output", json!("preserved"))]);
|
|
||||||
|
|
||||||
apply_state_updates(&node, &mut state, &json!("new agent output"));
|
|
||||||
|
|
||||||
assert_eq!(
|
|
||||||
state.state().get("greeting"),
|
|
||||||
Some(&json!("new agent output"))
|
|
||||||
);
|
|
||||||
assert_eq!(state.state().get("output"), Some(&json!("preserved")));
|
|
||||||
}
|
|
||||||
|
|
||||||
#[test]
|
|
||||||
fn no_state_updates_is_a_noop() {
|
|
||||||
let node = node_with("hi", None);
|
|
||||||
let mut state = manager_with(&[("k", json!("v"))]);
|
|
||||||
|
|
||||||
apply_state_updates(&node, &mut state, &json!("ignored"));
|
|
||||||
|
|
||||||
assert_eq!(state.state().get("k"), Some(&json!("v")));
|
|
||||||
assert!(state.state().get("output").is_none());
|
|
||||||
}
|
|
||||||
|
|
||||||
#[test]
|
|
||||||
fn interpolate_lenient_on_state_updates_handles_missing_keys() {
|
|
||||||
let node = {
|
|
||||||
let mut u = HashMap::new();
|
|
||||||
u.insert("decorated".into(), "[{{missing}}] {{output}}".into());
|
|
||||||
node_with("hi", Some(u))
|
|
||||||
};
|
|
||||||
let mut state = manager_with(&[]);
|
|
||||||
|
|
||||||
apply_state_updates(&node, &mut state, &json!("DATA"));
|
|
||||||
|
|
||||||
assert_eq!(state.state().get("decorated"), Some(&json!("[] DATA")));
|
|
||||||
}
|
|
||||||
|
|
||||||
fn node_with_schema(
|
|
||||||
prompt: &str,
|
|
||||||
updates: Option<HashMap<String, String>>,
|
|
||||||
schema: Value,
|
|
||||||
) -> AgentNode {
|
|
||||||
let mut n = node_with(prompt, updates);
|
|
||||||
n.output_schema = Some(schema);
|
|
||||||
n
|
|
||||||
}
|
|
||||||
|
|
||||||
#[test]
|
|
||||||
fn output_schema_auto_merges_top_level_keys() {
|
|
||||||
let node = node_with_schema("hi", None, json!({"type": "object"}));
|
|
||||||
let mut state = manager_with(&[]);
|
|
||||||
let output = json!({"goal": "do X", "summary": "details"});
|
|
||||||
|
|
||||||
apply_state_updates(&node, &mut state, &output);
|
|
||||||
|
|
||||||
assert_eq!(state.state().get("goal"), Some(&json!("do X")));
|
|
||||||
assert_eq!(state.state().get("summary"), Some(&json!("details")));
|
|
||||||
}
|
|
||||||
|
|
||||||
#[test]
|
|
||||||
fn output_schema_preserves_nested_value_types() {
|
|
||||||
let node = node_with_schema("hi", None, json!({"type": "object"}));
|
|
||||||
let mut state = manager_with(&[]);
|
|
||||||
let output = json!({
|
|
||||||
"tags": ["a", "b"],
|
|
||||||
"config": { "key": "value" },
|
|
||||||
"count": 42
|
|
||||||
});
|
|
||||||
|
|
||||||
apply_state_updates(&node, &mut state, &output);
|
|
||||||
|
|
||||||
assert_eq!(state.state().get("tags"), Some(&json!(["a", "b"])));
|
|
||||||
assert_eq!(state.state().get("config"), Some(&json!({"key": "value"})));
|
|
||||||
assert_eq!(state.state().get("count"), Some(&json!(42)));
|
|
||||||
}
|
|
||||||
|
|
||||||
#[test]
|
|
||||||
fn output_schema_explicit_state_updates_override_auto_merge() {
|
|
||||||
let mut u = HashMap::new();
|
|
||||||
u.insert("goal".into(), "renamed-{{output.goal}}".into());
|
|
||||||
let node = node_with_schema("hi", Some(u), json!({"type": "object"}));
|
|
||||||
let mut state = manager_with(&[]);
|
|
||||||
let output = json!({"goal": "do X"});
|
|
||||||
|
|
||||||
apply_state_updates(&node, &mut state, &output);
|
|
||||||
|
|
||||||
assert_eq!(state.state().get("goal"), Some(&json!("renamed-do X")));
|
|
||||||
}
|
|
||||||
|
|
||||||
#[test]
|
|
||||||
fn no_schema_does_not_auto_merge() {
|
|
||||||
let node = node_with("hi", None);
|
|
||||||
let mut state = manager_with(&[]);
|
|
||||||
let output = json!({"goal": "do X"});
|
|
||||||
|
|
||||||
apply_state_updates(&node, &mut state, &output);
|
|
||||||
|
|
||||||
assert!(state.state().get("goal").is_none());
|
|
||||||
}
|
|
||||||
}
|
|
||||||
@@ -1,47 +0,0 @@
|
|||||||
use super::{GraphExecutor, GraphParser, agent_has_graph};
|
|
||||||
use crate::config::RequestContext;
|
|
||||||
use crate::config::paths;
|
|
||||||
use crate::utils::AbortSignal;
|
|
||||||
use anyhow::{Context, Result, anyhow};
|
|
||||||
use log::info;
|
|
||||||
use serde_json::Value;
|
|
||||||
|
|
||||||
pub fn active_agent_graph_name(ctx: &RequestContext) -> Option<String> {
|
|
||||||
let name = ctx.agent.as_ref()?.name().to_string();
|
|
||||||
agent_has_graph(&name).then_some(name)
|
|
||||||
}
|
|
||||||
|
|
||||||
pub async fn run_active_agent_graph(
|
|
||||||
ctx: &mut RequestContext,
|
|
||||||
prompt: &str,
|
|
||||||
abort_signal: AbortSignal,
|
|
||||||
) -> Result<String> {
|
|
||||||
let agent_name =
|
|
||||||
active_agent_graph_name(ctx).ok_or_else(|| anyhow!("Active agent has no graph.yaml"))?;
|
|
||||||
|
|
||||||
info!("Agent '{agent_name}' has graph.yaml; routing to graph executor");
|
|
||||||
|
|
||||||
let agent_dir = paths::agent_data_dir(&agent_name);
|
|
||||||
let graph_path = paths::agent_graph_file(&agent_name);
|
|
||||||
|
|
||||||
let parser = GraphParser::new(&agent_dir);
|
|
||||||
let mut graph = parser
|
|
||||||
.load_from_file(&graph_path)
|
|
||||||
.with_context(|| format!("Failed to load graph.yaml for agent '{agent_name}'"))?;
|
|
||||||
|
|
||||||
graph
|
|
||||||
.initial_state
|
|
||||||
.insert("initial_prompt".into(), Value::String(prompt.to_string()));
|
|
||||||
|
|
||||||
let executor = GraphExecutor::new(graph, agent_dir);
|
|
||||||
let output = executor
|
|
||||||
.execute(ctx, abort_signal)
|
|
||||||
.await
|
|
||||||
.with_context(|| format!("Graph execution failed for agent '{agent_name}'"))?;
|
|
||||||
|
|
||||||
if let Some(supervisor) = ctx.supervisor.clone() {
|
|
||||||
supervisor.read().cancel_all();
|
|
||||||
}
|
|
||||||
|
|
||||||
Ok(output)
|
|
||||||
}
|
|
||||||
@@ -1,796 +0,0 @@
|
|||||||
use super::agent::AgentNodeExecutor;
|
|
||||||
use super::llm::{LlmExecutionOutcome, LlmNodeExecutor};
|
|
||||||
use super::logging::{GraphLogger, narrate_node_complete, narrate_node_failed};
|
|
||||||
use super::map::MapNodeExecutor;
|
|
||||||
use super::rag::RagNodeExecutor;
|
|
||||||
use super::script::ScriptExecutor;
|
|
||||||
use super::staging::BranchWrites;
|
|
||||||
use super::state::StateManager;
|
|
||||||
use super::types::{EndNode, Graph, Node, NodeType};
|
|
||||||
use super::user_interaction::{ApprovalNodeExecutor, InputNodeExecutor};
|
|
||||||
use super::validator::{AgentValidationContext, GraphValidator};
|
|
||||||
use crate::config::{RenderMode, RequestContext};
|
|
||||||
use crate::utils::AbortSignal;
|
|
||||||
use anyhow::{Context, Result, anyhow, bail};
|
|
||||||
use futures_util::future::join_all;
|
|
||||||
use serde_json::Value;
|
|
||||||
use std::collections::{HashMap, HashSet};
|
|
||||||
use std::path::PathBuf;
|
|
||||||
use std::sync::Arc;
|
|
||||||
use std::time::{Duration, Instant};
|
|
||||||
use tokio::sync::Semaphore;
|
|
||||||
|
|
||||||
pub struct GraphExecutor {
|
|
||||||
graph: Graph,
|
|
||||||
base_dir: PathBuf,
|
|
||||||
}
|
|
||||||
|
|
||||||
impl GraphExecutor {
|
|
||||||
pub fn new(graph: Graph, base_dir: impl Into<PathBuf>) -> Self {
|
|
||||||
Self {
|
|
||||||
graph,
|
|
||||||
base_dir: base_dir.into(),
|
|
||||||
}
|
|
||||||
}
|
|
||||||
|
|
||||||
pub async fn execute(
|
|
||||||
self,
|
|
||||||
ctx: &mut RequestContext,
|
|
||||||
abort_signal: AbortSignal,
|
|
||||||
) -> Result<String> {
|
|
||||||
let is_nested = ctx.current_depth > 0;
|
|
||||||
let mut logger = GraphLogger::with_visibility(
|
|
||||||
&self.graph.name,
|
|
||||||
self.graph.settings.log_state_snapshots,
|
|
||||||
is_nested,
|
|
||||||
);
|
|
||||||
let result = self.run(&mut logger, ctx, abort_signal).await;
|
|
||||||
if let Err(e) = &result {
|
|
||||||
logger.graph_error(e);
|
|
||||||
}
|
|
||||||
result
|
|
||||||
}
|
|
||||||
|
|
||||||
async fn run(
|
|
||||||
self,
|
|
||||||
logger: &mut GraphLogger,
|
|
||||||
ctx: &mut RequestContext,
|
|
||||||
abort_signal: AbortSignal,
|
|
||||||
) -> Result<String> {
|
|
||||||
let GraphExecutor { graph, base_dir } = self;
|
|
||||||
|
|
||||||
if graph.settings.validate_before_run {
|
|
||||||
let mut validator = GraphValidator::new(&base_dir);
|
|
||||||
if let Some(agent) = &ctx.agent {
|
|
||||||
validator = validator.with_agent_context(AgentValidationContext::from_agent(
|
|
||||||
agent,
|
|
||||||
Arc::clone(&ctx.app.config),
|
|
||||||
));
|
|
||||||
}
|
|
||||||
let result = validator.validate(&graph);
|
|
||||||
for w in &result.warnings {
|
|
||||||
logger.validation_warning(w.node_id.as_deref(), &w.message);
|
|
||||||
}
|
|
||||||
result.into_result()?;
|
|
||||||
}
|
|
||||||
|
|
||||||
let mut state = StateManager::new(graph.initial_state.clone());
|
|
||||||
let agent_envs = ctx
|
|
||||||
.agent
|
|
||||||
.as_ref()
|
|
||||||
.map(|a| a.variable_envs())
|
|
||||||
.unwrap_or_default();
|
|
||||||
let script_executor = ScriptExecutor::new(&base_dir).with_envs(agent_envs);
|
|
||||||
let max_iterations = graph.settings.max_loop_iterations;
|
|
||||||
let graph_timeout = graph.settings.timeout.map(Duration::from_secs);
|
|
||||||
let max_concurrency = graph.settings.max_concurrency;
|
|
||||||
let graph = Arc::new(graph);
|
|
||||||
let start = Instant::now();
|
|
||||||
|
|
||||||
let mut frontier: HashSet<String> = HashSet::from([graph.start.clone()]);
|
|
||||||
logger.graph_start(&graph.start, graph.nodes.len());
|
|
||||||
|
|
||||||
loop {
|
|
||||||
if frontier.is_empty() {
|
|
||||||
bail!(
|
|
||||||
"Graph '{}' frontier emptied without reaching an End node",
|
|
||||||
graph.name
|
|
||||||
);
|
|
||||||
}
|
|
||||||
|
|
||||||
if abort_signal.aborted() {
|
|
||||||
bail!(
|
|
||||||
"Graph '{}' aborted before super-step with frontier {:?}",
|
|
||||||
graph.name,
|
|
||||||
sorted_frontier(&frontier)
|
|
||||||
);
|
|
||||||
}
|
|
||||||
if let Some(t) = graph_timeout
|
|
||||||
&& start.elapsed() > t
|
|
||||||
{
|
|
||||||
bail!(
|
|
||||||
"Graph '{}' timed out after {}s before super-step with frontier {:?}",
|
|
||||||
graph.name,
|
|
||||||
t.as_secs(),
|
|
||||||
sorted_frontier(&frontier)
|
|
||||||
);
|
|
||||||
}
|
|
||||||
|
|
||||||
// Loop-count and visit tracking on live state, BEFORE forking.
|
|
||||||
// This counts every entry to a node toward max_loop_iterations
|
|
||||||
// regardless of how many parallel branches converged on it.
|
|
||||||
for node_id in &frontier {
|
|
||||||
state.state_mut().visit_node(node_id);
|
|
||||||
let visits = state.state().loop_count(node_id);
|
|
||||||
if visits > max_iterations {
|
|
||||||
bail!(
|
|
||||||
"Node '{}' visited {} times (max_loop_iterations={}). \
|
|
||||||
Possible infinite loop.",
|
|
||||||
node_id,
|
|
||||||
visits,
|
|
||||||
max_iterations
|
|
||||||
);
|
|
||||||
}
|
|
||||||
}
|
|
||||||
|
|
||||||
for node_id in &frontier {
|
|
||||||
let node = graph.get_node(node_id).ok_or_else(|| {
|
|
||||||
anyhow!("Node '{}' not found in graph '{}'", node_id, graph.name)
|
|
||||||
})?;
|
|
||||||
let visits = state.state().loop_count(node_id);
|
|
||||||
logger.node_entry(node, visits);
|
|
||||||
}
|
|
||||||
let snapshot_label = if frontier.len() == 1 {
|
|
||||||
frontier.iter().next().cloned().unwrap_or_default()
|
|
||||||
} else {
|
|
||||||
format!("super-step {{{}}}", sorted_frontier(&frontier).join(","))
|
|
||||||
};
|
|
||||||
logger.state_snapshot(&snapshot_label, &state);
|
|
||||||
|
|
||||||
let snapshot = state.read_snapshot();
|
|
||||||
let semaphore = Arc::new(Semaphore::new(max_concurrency));
|
|
||||||
|
|
||||||
let frontier_size = frontier.len();
|
|
||||||
let in_super_step = frontier_size > 1;
|
|
||||||
let silent = logger.silent();
|
|
||||||
|
|
||||||
if in_super_step {
|
|
||||||
let mut branches = sorted_frontier(&frontier);
|
|
||||||
branches.sort();
|
|
||||||
logger.super_step_start(&branches);
|
|
||||||
}
|
|
||||||
|
|
||||||
let mut branch_tasks = Vec::with_capacity(frontier_size);
|
|
||||||
for node_id in &frontier {
|
|
||||||
let node = graph
|
|
||||||
.get_node(node_id)
|
|
||||||
.ok_or_else(|| {
|
|
||||||
anyhow!("Node '{}' not found in graph '{}'", node_id, graph.name)
|
|
||||||
})?
|
|
||||||
.clone();
|
|
||||||
logger.node_start(&node, in_super_step);
|
|
||||||
let branch_state = state.fork_for_branch_state();
|
|
||||||
let mut branch_ctx = ctx.fork_for_branch();
|
|
||||||
if in_super_step {
|
|
||||||
branch_ctx.render_mode = RenderMode::Silent;
|
|
||||||
}
|
|
||||||
let script_exec_clone = script_executor.clone();
|
|
||||||
let graph_clone = Arc::clone(&graph);
|
|
||||||
let current = node_id.clone();
|
|
||||||
let sem_clone = semaphore.clone();
|
|
||||||
let abort_clone = abort_signal.clone();
|
|
||||||
|
|
||||||
let task = tokio::spawn(async move {
|
|
||||||
let _permit = sem_clone
|
|
||||||
.acquire()
|
|
||||||
.await
|
|
||||||
.expect("semaphore should not be closed");
|
|
||||||
if abort_clone.aborted() {
|
|
||||||
narrate_node_failed(
|
|
||||||
silent,
|
|
||||||
&node,
|
|
||||||
Duration::default(),
|
|
||||||
"aborted",
|
|
||||||
in_super_step,
|
|
||||||
);
|
|
||||||
return (
|
|
||||||
current.clone(),
|
|
||||||
branch_state,
|
|
||||||
Err(anyhow!("branch aborted")),
|
|
||||||
Duration::default(),
|
|
||||||
);
|
|
||||||
}
|
|
||||||
let node_start = Instant::now();
|
|
||||||
let mut state = branch_state;
|
|
||||||
let mut ctx = branch_ctx;
|
|
||||||
let step_ctx = StepContext {
|
|
||||||
graph: graph_clone.as_ref(),
|
|
||||||
script_executor: &script_exec_clone,
|
|
||||||
max_concurrency,
|
|
||||||
abort_signal: &abort_clone,
|
|
||||||
};
|
|
||||||
let result = step(&node, &mut state, &mut ctx, &step_ctx, ¤t).await;
|
|
||||||
let elapsed = node_start.elapsed();
|
|
||||||
match &result {
|
|
||||||
Ok(StepResult::Continue(targets)) => {
|
|
||||||
let route = if targets.is_empty() {
|
|
||||||
None
|
|
||||||
} else {
|
|
||||||
Some(targets.join(", "))
|
|
||||||
};
|
|
||||||
narrate_node_complete(
|
|
||||||
silent,
|
|
||||||
&node,
|
|
||||||
elapsed,
|
|
||||||
route.as_deref(),
|
|
||||||
in_super_step,
|
|
||||||
);
|
|
||||||
}
|
|
||||||
Ok(StepResult::End(_)) => {
|
|
||||||
narrate_node_complete(
|
|
||||||
silent,
|
|
||||||
&node,
|
|
||||||
elapsed,
|
|
||||||
Some("END"),
|
|
||||||
in_super_step,
|
|
||||||
);
|
|
||||||
}
|
|
||||||
Err(e) => {
|
|
||||||
narrate_node_failed(
|
|
||||||
silent,
|
|
||||||
&node,
|
|
||||||
elapsed,
|
|
||||||
&e.to_string(),
|
|
||||||
in_super_step,
|
|
||||||
);
|
|
||||||
}
|
|
||||||
}
|
|
||||||
(current, state, result, elapsed)
|
|
||||||
});
|
|
||||||
branch_tasks.push(task);
|
|
||||||
}
|
|
||||||
|
|
||||||
let joined = join_all(branch_tasks).await;
|
|
||||||
|
|
||||||
let mut branch_writes: Vec<BranchWrites> = Vec::new();
|
|
||||||
let mut next_frontier: HashSet<String> = HashSet::new();
|
|
||||||
let mut end_results: Vec<(String, StateManager, String)> = Vec::new();
|
|
||||||
|
|
||||||
for join_result in joined {
|
|
||||||
let (node_id, branch_state, step_result, elapsed) =
|
|
||||||
join_result.map_err(|e| anyhow!("Branch task panicked: {e}"))?;
|
|
||||||
logger.record_timing(&node_id, elapsed);
|
|
||||||
|
|
||||||
let step_outcome = step_result.with_context(|| format!("at node '{node_id}'"))?;
|
|
||||||
|
|
||||||
match step_outcome {
|
|
||||||
StepResult::Continue(targets) => {
|
|
||||||
for target in &targets {
|
|
||||||
logger.routing(&node_id, target);
|
|
||||||
}
|
|
||||||
let diff = branch_state.diff_against(snapshot.as_ref());
|
|
||||||
branch_writes.push(BranchWrites {
|
|
||||||
node_id: node_id.clone(),
|
|
||||||
invocation_index: 0,
|
|
||||||
writes: diff,
|
|
||||||
});
|
|
||||||
next_frontier.extend(targets);
|
|
||||||
}
|
|
||||||
StepResult::End(output) => {
|
|
||||||
end_results.push((node_id.clone(), branch_state, output));
|
|
||||||
}
|
|
||||||
}
|
|
||||||
}
|
|
||||||
|
|
||||||
if end_results.len() > 1 {
|
|
||||||
let mut ids: Vec<String> =
|
|
||||||
end_results.iter().map(|(id, _, _)| id.clone()).collect();
|
|
||||||
ids.sort();
|
|
||||||
bail!(
|
|
||||||
"super-step ended with multiple End targets ({}). \
|
|
||||||
Fan-out branches must converge at a join node before \
|
|
||||||
terminating. To fix: route all parallel branches to a \
|
|
||||||
single shared next-node, then terminate from there.",
|
|
||||||
ids.join(", ")
|
|
||||||
);
|
|
||||||
}
|
|
||||||
|
|
||||||
// Sort by (node_id, invocation_index) so non-commutative reducers
|
|
||||||
// like Concat/Merge produce deterministic output across runs.
|
|
||||||
branch_writes.sort_by(|a, b| {
|
|
||||||
a.node_id
|
|
||||||
.cmp(&b.node_id)
|
|
||||||
.then(a.invocation_index.cmp(&b.invocation_index))
|
|
||||||
});
|
|
||||||
state.apply_branch_writes(branch_writes, &graph.reducers)?;
|
|
||||||
|
|
||||||
if let Some((node_id, end_state, output)) = end_results.into_iter().next() {
|
|
||||||
let diff = end_state.diff_against(snapshot.as_ref());
|
|
||||||
state.apply_branch_writes(
|
|
||||||
vec![BranchWrites {
|
|
||||||
node_id: node_id.clone(),
|
|
||||||
invocation_index: 0,
|
|
||||||
writes: diff,
|
|
||||||
}],
|
|
||||||
&graph.reducers,
|
|
||||||
)?;
|
|
||||||
logger.graph_complete(&node_id, start.elapsed());
|
|
||||||
return Ok(output);
|
|
||||||
}
|
|
||||||
|
|
||||||
if in_super_step {
|
|
||||||
logger.super_step_end(&sorted_frontier(&next_frontier));
|
|
||||||
}
|
|
||||||
frontier = next_frontier;
|
|
||||||
}
|
|
||||||
}
|
|
||||||
}
|
|
||||||
|
|
||||||
fn sorted_frontier(frontier: &HashSet<String>) -> Vec<String> {
|
|
||||||
let mut v: Vec<String> = frontier.iter().cloned().collect();
|
|
||||||
v.sort();
|
|
||||||
v
|
|
||||||
}
|
|
||||||
|
|
||||||
pub(super) struct StepContext<'a> {
|
|
||||||
pub graph: &'a Graph,
|
|
||||||
pub script_executor: &'a ScriptExecutor,
|
|
||||||
pub max_concurrency: usize,
|
|
||||||
pub abort_signal: &'a AbortSignal,
|
|
||||||
}
|
|
||||||
|
|
||||||
impl StepContext<'_> {
|
|
||||||
pub fn graph_name(&self) -> &str {
|
|
||||||
&self.graph.name
|
|
||||||
}
|
|
||||||
}
|
|
||||||
|
|
||||||
enum StepResult {
|
|
||||||
// The set of next-node ids the executor should add to the next super-step's
|
|
||||||
// frontier. A `Vec` of length 1 for sequential routing (default) and the
|
|
||||||
// full target list for fan-out (`next: [a, b, ...]`). Dynamic single-route
|
|
||||||
// decisions (script `_next`, approval routes, LLM/RAG fallback) always emit
|
|
||||||
// a single-element vec.
|
|
||||||
Continue(Vec<String>),
|
|
||||||
End(String),
|
|
||||||
}
|
|
||||||
|
|
||||||
async fn step(
|
|
||||||
node: &Node,
|
|
||||||
state: &mut StateManager,
|
|
||||||
ctx: &mut RequestContext,
|
|
||||||
step_ctx: &StepContext<'_>,
|
|
||||||
current: &str,
|
|
||||||
) -> Result<StepResult> {
|
|
||||||
match &node.node_type {
|
|
||||||
NodeType::Agent(agent_node) => {
|
|
||||||
AgentNodeExecutor::execute(agent_node, state, ctx).await?;
|
|
||||||
let targets = static_next_targets(node, current, "agent")?;
|
|
||||||
Ok(StepResult::Continue(targets))
|
|
||||||
}
|
|
||||||
NodeType::Script(script_node) => {
|
|
||||||
let dynamic = match step_ctx.script_executor.execute(script_node, state).await {
|
|
||||||
Ok(n) => n,
|
|
||||||
Err(e) => {
|
|
||||||
if let Some(fallback) = &script_node.fallback {
|
|
||||||
warn!(
|
|
||||||
"[graph:{}] script '{}' failed, routing to fallback '{}': {}",
|
|
||||||
step_ctx.graph_name(),
|
|
||||||
current,
|
|
||||||
fallback,
|
|
||||||
e
|
|
||||||
);
|
|
||||||
return Ok(StepResult::Continue(vec![fallback.clone()]));
|
|
||||||
}
|
|
||||||
return Err(e);
|
|
||||||
}
|
|
||||||
};
|
|
||||||
let targets = match dynamic {
|
|
||||||
Some(n) => vec![n],
|
|
||||||
None => static_next_targets(node, current, "script")?,
|
|
||||||
};
|
|
||||||
Ok(StepResult::Continue(targets))
|
|
||||||
}
|
|
||||||
NodeType::Approval(approval_node) => {
|
|
||||||
let next = ApprovalNodeExecutor::execute(approval_node, state, ctx).await?;
|
|
||||||
Ok(StepResult::Continue(vec![next]))
|
|
||||||
}
|
|
||||||
NodeType::Input(input_node) => {
|
|
||||||
let next_id = first_next_target(node);
|
|
||||||
let next = InputNodeExecutor::execute(input_node, next_id, state, ctx).await?;
|
|
||||||
Ok(StepResult::Continue(vec![next]))
|
|
||||||
}
|
|
||||||
NodeType::Llm(llm_node) => {
|
|
||||||
let outcome = LlmNodeExecutor::execute(llm_node, state, ctx).await?;
|
|
||||||
let targets = match outcome {
|
|
||||||
LlmExecutionOutcome::Continue => static_next_targets(node, current, "llm")?,
|
|
||||||
LlmExecutionOutcome::FellBack(target) => vec![target],
|
|
||||||
};
|
|
||||||
Ok(StepResult::Continue(targets))
|
|
||||||
}
|
|
||||||
NodeType::Rag(rag_node) => {
|
|
||||||
RagNodeExecutor::execute(rag_node, current, state, ctx).await?;
|
|
||||||
let targets = static_next_targets(node, current, "rag")?;
|
|
||||||
Ok(StepResult::Continue(targets))
|
|
||||||
}
|
|
||||||
NodeType::End(end_node) => Ok(StepResult::End(resolve_end_output(end_node, state))),
|
|
||||||
NodeType::Map(map_node) => {
|
|
||||||
let targets = static_next_targets(node, current, "map")?;
|
|
||||||
MapNodeExecutor::execute(map_node, state, ctx, step_ctx, current).await?;
|
|
||||||
Ok(StepResult::Continue(targets))
|
|
||||||
}
|
|
||||||
}
|
|
||||||
}
|
|
||||||
|
|
||||||
fn static_next_targets(node: &Node, current: &str, kind: &str) -> Result<Vec<String>> {
|
|
||||||
node.next
|
|
||||||
.as_ref()
|
|
||||||
.map(|t| t.as_slice().to_vec())
|
|
||||||
.ok_or_else(|| anyhow!("{kind} node '{current}' has no `next` and is not an end node"))
|
|
||||||
}
|
|
||||||
|
|
||||||
fn first_next_target(node: &Node) -> Option<&str> {
|
|
||||||
node.next
|
|
||||||
.as_ref()
|
|
||||||
.and_then(|t| t.as_slice().first().map(|s| s.as_str()))
|
|
||||||
}
|
|
||||||
|
|
||||||
fn resolve_end_output(end_node: &EndNode, state: &mut StateManager) -> String {
|
|
||||||
apply_simple_state_updates(end_node.state_updates.as_ref(), state);
|
|
||||||
state.interpolate_lenient(&end_node.output)
|
|
||||||
}
|
|
||||||
|
|
||||||
fn apply_simple_state_updates(updates: Option<&HashMap<String, String>>, state: &mut StateManager) {
|
|
||||||
let Some(updates) = updates else {
|
|
||||||
return;
|
|
||||||
};
|
|
||||||
for (key, template) in updates {
|
|
||||||
let value = state.interpolate_lenient(template);
|
|
||||||
state.state_mut().set(key.clone(), Value::String(value));
|
|
||||||
}
|
|
||||||
}
|
|
||||||
|
|
||||||
#[cfg(test)]
|
|
||||||
mod tests {
|
|
||||||
use super::*;
|
|
||||||
use serde_json::json;
|
|
||||||
|
|
||||||
fn state_with(pairs: &[(&str, Value)]) -> StateManager {
|
|
||||||
let mut map = HashMap::new();
|
|
||||||
for (k, v) in pairs {
|
|
||||||
map.insert((*k).into(), v.clone());
|
|
||||||
}
|
|
||||||
StateManager::new(map)
|
|
||||||
}
|
|
||||||
|
|
||||||
fn end_node(output: &str, updates: Option<HashMap<String, String>>) -> EndNode {
|
|
||||||
EndNode {
|
|
||||||
output: output.into(),
|
|
||||||
state_updates: updates,
|
|
||||||
}
|
|
||||||
}
|
|
||||||
|
|
||||||
#[test]
|
|
||||||
fn resolve_end_output_interpolates_template_against_state() {
|
|
||||||
let mut state = state_with(&[("name", json!("alice"))]);
|
|
||||||
|
|
||||||
let node = end_node("done: {{name}}", None);
|
|
||||||
|
|
||||||
assert_eq!(resolve_end_output(&node, &mut state), "done: alice");
|
|
||||||
}
|
|
||||||
|
|
||||||
#[test]
|
|
||||||
fn resolve_end_output_applies_state_updates_before_interpolation() {
|
|
||||||
let mut updates = HashMap::new();
|
|
||||||
updates.insert("summary".into(), "completed for {{user}}".into());
|
|
||||||
let node = end_node("RESULT: {{summary}}", Some(updates));
|
|
||||||
|
|
||||||
let mut state = state_with(&[("user", json!("bob"))]);
|
|
||||||
|
|
||||||
assert_eq!(
|
|
||||||
resolve_end_output(&node, &mut state),
|
|
||||||
"RESULT: completed for bob"
|
|
||||||
);
|
|
||||||
assert_eq!(
|
|
||||||
state.state().get("summary"),
|
|
||||||
Some(&json!("completed for bob"))
|
|
||||||
);
|
|
||||||
}
|
|
||||||
|
|
||||||
#[test]
|
|
||||||
fn resolve_end_output_with_empty_template_returns_empty_string() {
|
|
||||||
let mut state = state_with(&[]);
|
|
||||||
|
|
||||||
let node = end_node("", None);
|
|
||||||
|
|
||||||
assert_eq!(resolve_end_output(&node, &mut state), "");
|
|
||||||
}
|
|
||||||
|
|
||||||
#[test]
|
|
||||||
fn resolve_end_output_lenient_on_missing_keys() {
|
|
||||||
let mut state = state_with(&[]);
|
|
||||||
|
|
||||||
let node = end_node("hello {{unknown}}!", None);
|
|
||||||
|
|
||||||
assert_eq!(resolve_end_output(&node, &mut state), "hello !");
|
|
||||||
}
|
|
||||||
|
|
||||||
#[test]
|
|
||||||
fn apply_simple_state_updates_does_nothing_when_none() {
|
|
||||||
let mut state = state_with(&[("k", json!("v"))]);
|
|
||||||
|
|
||||||
apply_simple_state_updates(None, &mut state);
|
|
||||||
|
|
||||||
assert_eq!(state.state().get("k"), Some(&json!("v")));
|
|
||||||
}
|
|
||||||
|
|
||||||
#[test]
|
|
||||||
fn apply_simple_state_updates_overwrites_existing_values() {
|
|
||||||
let mut updates = HashMap::new();
|
|
||||||
updates.insert("k".into(), "new-{{k}}".into());
|
|
||||||
let mut state = state_with(&[("k", json!("old"))]);
|
|
||||||
|
|
||||||
apply_simple_state_updates(Some(&updates), &mut state);
|
|
||||||
|
|
||||||
assert_eq!(state.state().get("k"), Some(&json!("new-old")));
|
|
||||||
}
|
|
||||||
}
|
|
||||||
|
|
||||||
#[cfg(test)]
|
|
||||||
mod integration_tests {
|
|
||||||
use super::*;
|
|
||||||
use crate::config::{AppState, WorkingMode};
|
|
||||||
use crate::utils::{create_abort_signal, temp_file};
|
|
||||||
use std::fs;
|
|
||||||
|
|
||||||
fn cmd_available(name: &str) -> bool {
|
|
||||||
which::which(name).is_ok()
|
|
||||||
}
|
|
||||||
|
|
||||||
struct TestWorkspace {
|
|
||||||
dir: PathBuf,
|
|
||||||
}
|
|
||||||
|
|
||||||
impl TestWorkspace {
|
|
||||||
fn new() -> Self {
|
|
||||||
let dir = temp_file("-graph-integration-", "");
|
|
||||||
fs::create_dir_all(&dir).unwrap();
|
|
||||||
Self { dir }
|
|
||||||
}
|
|
||||||
|
|
||||||
fn write_script(&self, name: &str, contents: &str) {
|
|
||||||
fs::write(self.dir.join(name), contents).unwrap();
|
|
||||||
}
|
|
||||||
}
|
|
||||||
|
|
||||||
impl Drop for TestWorkspace {
|
|
||||||
fn drop(&mut self) {
|
|
||||||
let _ = fs::remove_dir_all(&self.dir);
|
|
||||||
}
|
|
||||||
}
|
|
||||||
|
|
||||||
fn make_ctx() -> RequestContext {
|
|
||||||
RequestContext::new(Arc::new(AppState::test_default()), WorkingMode::Cmd)
|
|
||||||
}
|
|
||||||
|
|
||||||
#[tokio::test]
|
|
||||||
async fn static_fan_out_merges_branch_writes_via_append_reducer() {
|
|
||||||
if !cmd_available("bash") {
|
|
||||||
eprintln!("skipping: bash not available");
|
|
||||||
return;
|
|
||||||
}
|
|
||||||
let ws = TestWorkspace::new();
|
|
||||||
ws.write_script("dispatcher.sh", "#!/bin/bash\necho '{}'\n");
|
|
||||||
ws.write_script(
|
|
||||||
"worker_a.sh",
|
|
||||||
"#!/bin/bash\necho '{\"results\": \"alpha\"}'\n",
|
|
||||||
);
|
|
||||||
ws.write_script(
|
|
||||||
"worker_b.sh",
|
|
||||||
"#!/bin/bash\necho '{\"results\": \"beta\"}'\n",
|
|
||||||
);
|
|
||||||
|
|
||||||
let yaml = r#"
|
|
||||||
name: static_fan_out_test
|
|
||||||
start: dispatcher
|
|
||||||
reducers:
|
|
||||||
results: append
|
|
||||||
nodes:
|
|
||||||
dispatcher:
|
|
||||||
type: script
|
|
||||||
script: dispatcher.sh
|
|
||||||
state_updates: {}
|
|
||||||
next: [worker_a, worker_b]
|
|
||||||
worker_a:
|
|
||||||
type: script
|
|
||||||
script: worker_a.sh
|
|
||||||
state_updates: {}
|
|
||||||
next: join
|
|
||||||
worker_b:
|
|
||||||
type: script
|
|
||||||
script: worker_b.sh
|
|
||||||
state_updates: {}
|
|
||||||
next: join
|
|
||||||
join:
|
|
||||||
type: end
|
|
||||||
output: "{{results}}"
|
|
||||||
"#;
|
|
||||||
let graph: Graph = serde_yaml::from_str(yaml).unwrap();
|
|
||||||
let mut ctx = make_ctx();
|
|
||||||
let abort = create_abort_signal();
|
|
||||||
let result = GraphExecutor::new(graph, &ws.dir)
|
|
||||||
.execute(&mut ctx, abort)
|
|
||||||
.await
|
|
||||||
.unwrap_or_else(|e| panic!("executor failed: {e:#}"));
|
|
||||||
|
|
||||||
let parsed: Value = serde_json::from_str(&result)
|
|
||||||
.unwrap_or_else(|_| panic!("expected JSON array, got: {result}"));
|
|
||||||
let arr = parsed.as_array().expect("results should be an array");
|
|
||||||
assert_eq!(arr.len(), 2, "expected 2 elements, got: {result}");
|
|
||||||
let strs: Vec<&str> = arr.iter().filter_map(|v| v.as_str()).collect();
|
|
||||||
assert!(strs.contains(&"alpha"), "missing 'alpha' in {strs:?}");
|
|
||||||
assert!(strs.contains(&"beta"), "missing 'beta' in {strs:?}");
|
|
||||||
}
|
|
||||||
|
|
||||||
#[tokio::test]
|
|
||||||
async fn map_over_list_collects_outputs_in_input_order() {
|
|
||||||
if !cmd_available("python3") {
|
|
||||||
eprintln!("skipping: python3 not available");
|
|
||||||
return;
|
|
||||||
}
|
|
||||||
let ws = TestWorkspace::new();
|
|
||||||
ws.write_script(
|
|
||||||
"doubler.py",
|
|
||||||
r#"#!/usr/bin/env python3
|
|
||||||
import os, json
|
|
||||||
state = json.loads(os.environ.get("GRAPH_STATE", "{}"))
|
|
||||||
val = state["item"]
|
|
||||||
print(json.dumps({"output": val * 2}))
|
|
||||||
"#,
|
|
||||||
);
|
|
||||||
|
|
||||||
let yaml = r#"
|
|
||||||
name: map_input_order_test
|
|
||||||
start: fan_out
|
|
||||||
initial_state:
|
|
||||||
items: [1, 2, 3, 4, 5]
|
|
||||||
nodes:
|
|
||||||
fan_out:
|
|
||||||
type: map
|
|
||||||
over: "{{items}}"
|
|
||||||
as: item
|
|
||||||
branch: doubler
|
|
||||||
collect_into: doubled
|
|
||||||
next: done
|
|
||||||
doubler:
|
|
||||||
type: script
|
|
||||||
script: doubler.py
|
|
||||||
state_updates: {}
|
|
||||||
done:
|
|
||||||
type: end
|
|
||||||
output: "{{doubled}}"
|
|
||||||
"#;
|
|
||||||
let graph: Graph = serde_yaml::from_str(yaml).unwrap();
|
|
||||||
let mut ctx = make_ctx();
|
|
||||||
let abort = create_abort_signal();
|
|
||||||
let result = GraphExecutor::new(graph, &ws.dir)
|
|
||||||
.execute(&mut ctx, abort)
|
|
||||||
.await
|
|
||||||
.unwrap_or_else(|e| panic!("executor failed: {e:#}"));
|
|
||||||
|
|
||||||
let parsed: Value = serde_json::from_str(&result)
|
|
||||||
.unwrap_or_else(|_| panic!("expected JSON array, got: {result}"));
|
|
||||||
let arr = parsed.as_array().expect("doubled should be an array");
|
|
||||||
let nums: Vec<i64> = arr
|
|
||||||
.iter()
|
|
||||||
.map(|v| v.as_i64().expect("each item should be int"))
|
|
||||||
.collect();
|
|
||||||
|
|
||||||
assert_eq!(
|
|
||||||
nums,
|
|
||||||
vec![2, 4, 6, 8, 10],
|
|
||||||
"map outputs should be in input order, not finish order"
|
|
||||||
);
|
|
||||||
}
|
|
||||||
|
|
||||||
#[tokio::test]
|
|
||||||
async fn parallel_branch_error_aborts_super_step() {
|
|
||||||
if !cmd_available("bash") {
|
|
||||||
eprintln!("skipping: bash not available");
|
|
||||||
return;
|
|
||||||
}
|
|
||||||
let ws = TestWorkspace::new();
|
|
||||||
ws.write_script("dispatcher.sh", "#!/bin/bash\necho '{}'\n");
|
|
||||||
ws.write_script(
|
|
||||||
"worker_ok.sh",
|
|
||||||
"#!/bin/bash\necho '{\"results\": \"ok\"}'\n",
|
|
||||||
);
|
|
||||||
ws.write_script(
|
|
||||||
"worker_fail.sh",
|
|
||||||
"#!/bin/bash\necho 'simulated failure' >&2\nexit 1\n",
|
|
||||||
);
|
|
||||||
|
|
||||||
let yaml = r#"
|
|
||||||
name: branch_error_test
|
|
||||||
start: dispatcher
|
|
||||||
reducers:
|
|
||||||
results: append
|
|
||||||
nodes:
|
|
||||||
dispatcher:
|
|
||||||
type: script
|
|
||||||
script: dispatcher.sh
|
|
||||||
state_updates: {}
|
|
||||||
next: [worker_ok, worker_fail]
|
|
||||||
worker_ok:
|
|
||||||
type: script
|
|
||||||
script: worker_ok.sh
|
|
||||||
state_updates: {}
|
|
||||||
next: join
|
|
||||||
worker_fail:
|
|
||||||
type: script
|
|
||||||
script: worker_fail.sh
|
|
||||||
state_updates: {}
|
|
||||||
next: join
|
|
||||||
join:
|
|
||||||
type: end
|
|
||||||
output: "{{results}}"
|
|
||||||
"#;
|
|
||||||
let graph: Graph = serde_yaml::from_str(yaml).unwrap();
|
|
||||||
let mut ctx = make_ctx();
|
|
||||||
let abort = create_abort_signal();
|
|
||||||
let result = GraphExecutor::new(graph, &ws.dir)
|
|
||||||
.execute(&mut ctx, abort)
|
|
||||||
.await;
|
|
||||||
|
|
||||||
assert!(result.is_err(), "expected branch error to propagate");
|
|
||||||
let err = format!("{:#}", result.unwrap_err());
|
|
||||||
assert!(
|
|
||||||
err.contains("worker_fail"),
|
|
||||||
"error should mention failing node: {err}"
|
|
||||||
);
|
|
||||||
}
|
|
||||||
|
|
||||||
#[tokio::test]
|
|
||||||
async fn multi_end_in_super_step_is_rejected() {
|
|
||||||
if !cmd_available("bash") {
|
|
||||||
eprintln!("skipping: bash not available");
|
|
||||||
return;
|
|
||||||
}
|
|
||||||
let ws = TestWorkspace::new();
|
|
||||||
ws.write_script("dispatcher.sh", "#!/bin/bash\necho '{}'\n");
|
|
||||||
|
|
||||||
let yaml = r#"
|
|
||||||
name: multi_end_test
|
|
||||||
start: dispatcher
|
|
||||||
nodes:
|
|
||||||
dispatcher:
|
|
||||||
type: script
|
|
||||||
script: dispatcher.sh
|
|
||||||
state_updates: {}
|
|
||||||
next: [end_a, end_b]
|
|
||||||
end_a:
|
|
||||||
type: end
|
|
||||||
output: "from a"
|
|
||||||
end_b:
|
|
||||||
type: end
|
|
||||||
output: "from b"
|
|
||||||
"#;
|
|
||||||
let graph: Graph = serde_yaml::from_str(yaml).unwrap();
|
|
||||||
let mut ctx = make_ctx();
|
|
||||||
let abort = create_abort_signal();
|
|
||||||
let result = GraphExecutor::new(graph, &ws.dir)
|
|
||||||
.execute(&mut ctx, abort)
|
|
||||||
.await;
|
|
||||||
|
|
||||||
assert!(result.is_err(), "expected multi-End to be rejected");
|
|
||||||
let err = format!("{:#}", result.unwrap_err());
|
|
||||||
assert!(
|
|
||||||
err.contains("multiple End targets"),
|
|
||||||
"error should explain multi-End cause: {err}"
|
|
||||||
);
|
|
||||||
assert!(
|
|
||||||
err.contains("end_a") && err.contains("end_b"),
|
|
||||||
"error should list both End nodes: {err}"
|
|
||||||
);
|
|
||||||
}
|
|
||||||
}
|
|
||||||
@@ -1,616 +0,0 @@
|
|||||||
use super::state::StateManager;
|
|
||||||
use super::structured;
|
|
||||||
use super::types::LlmNode;
|
|
||||||
use crate::client::{Model, ModelType, call_chat_completions};
|
|
||||||
use crate::config::{Input, RequestContext, Role, RoleLike};
|
|
||||||
use crate::utils::create_abort_signal;
|
|
||||||
use anyhow::{Context, Error, Result, anyhow, bail};
|
|
||||||
use serde_json::Value;
|
|
||||||
use std::collections::HashSet;
|
|
||||||
use std::sync::Arc;
|
|
||||||
use std::time::Duration;
|
|
||||||
use tokio::time::timeout;
|
|
||||||
|
|
||||||
const OUTPUT_KEY: &str = "output";
|
|
||||||
|
|
||||||
#[derive(Debug, PartialEq, Eq)]
|
|
||||||
pub(super) enum LlmExecutionOutcome {
|
|
||||||
Continue,
|
|
||||||
FellBack(String),
|
|
||||||
}
|
|
||||||
|
|
||||||
pub struct LlmNodeExecutor;
|
|
||||||
|
|
||||||
impl LlmNodeExecutor {
|
|
||||||
pub(super) async fn execute(
|
|
||||||
node: &LlmNode,
|
|
||||||
state_manager: &mut StateManager,
|
|
||||||
parent_ctx: &mut RequestContext,
|
|
||||||
) -> Result<LlmExecutionOutcome> {
|
|
||||||
let result = run(node, state_manager, parent_ctx).await;
|
|
||||||
let (output, failure_reason) = match result {
|
|
||||||
Ok(raw) => match &node.output_schema {
|
|
||||||
Some(schema) => match structured::extract(&raw, schema, parent_ctx).await {
|
|
||||||
Ok(value) => (value, None),
|
|
||||||
Err(e) => {
|
|
||||||
warn!("llm node structured extraction failed: {e}");
|
|
||||||
(
|
|
||||||
Value::String(format!("LLM node structured-extraction failed: {e}")),
|
|
||||||
Some(format!("structured-extraction failed: {e}")),
|
|
||||||
)
|
|
||||||
}
|
|
||||||
},
|
|
||||||
None => (Value::String(raw), None),
|
|
||||||
},
|
|
||||||
Err(e) => {
|
|
||||||
warn!("llm node failed: {e}");
|
|
||||||
(
|
|
||||||
Value::String(format!("LLM node failed: {e}")),
|
|
||||||
Some(format!("LLM call failed: {e:#}")),
|
|
||||||
)
|
|
||||||
}
|
|
||||||
};
|
|
||||||
|
|
||||||
apply_state_updates_with_output(node, state_manager, &output);
|
|
||||||
outcome_from(failure_reason.as_deref(), node.fallback.as_deref())
|
|
||||||
}
|
|
||||||
}
|
|
||||||
|
|
||||||
fn outcome_from(
|
|
||||||
failure_reason: Option<&str>,
|
|
||||||
fallback: Option<&str>,
|
|
||||||
) -> Result<LlmExecutionOutcome> {
|
|
||||||
match (failure_reason, fallback) {
|
|
||||||
(None, _) => Ok(LlmExecutionOutcome::Continue),
|
|
||||||
(Some(_), Some(fb)) => Ok(LlmExecutionOutcome::FellBack(fb.to_string())),
|
|
||||||
(Some(reason), None) => bail!(
|
|
||||||
"LLM node failed and no fallback declared: {reason}. \
|
|
||||||
Add a `fallback:` route on the node to route on failure, \
|
|
||||||
or fix the underlying error."
|
|
||||||
),
|
|
||||||
}
|
|
||||||
}
|
|
||||||
|
|
||||||
async fn run(
|
|
||||||
node: &LlmNode,
|
|
||||||
state_manager: &mut StateManager,
|
|
||||||
parent_ctx: &mut RequestContext,
|
|
||||||
) -> Result<String> {
|
|
||||||
let mut instructions: Option<String> = match &node.instructions {
|
|
||||||
Some(s) => Some(
|
|
||||||
state_manager
|
|
||||||
.interpolate(s)
|
|
||||||
.context("Failed to interpolate llm node instructions")?,
|
|
||||||
),
|
|
||||||
None => None,
|
|
||||||
};
|
|
||||||
let mut prompt = state_manager
|
|
||||||
.interpolate(&node.prompt)
|
|
||||||
.context("Failed to interpolate llm node prompt")?;
|
|
||||||
|
|
||||||
if let Some(schema) = &node.output_schema {
|
|
||||||
let hint = format_schema_hint(schema);
|
|
||||||
match instructions.as_mut() {
|
|
||||||
Some(s) => {
|
|
||||||
s.push_str("\n\n");
|
|
||||||
s.push_str(&hint);
|
|
||||||
}
|
|
||||||
None => {
|
|
||||||
prompt.push_str("\n\n");
|
|
||||||
prompt.push_str(&hint);
|
|
||||||
}
|
|
||||||
}
|
|
||||||
}
|
|
||||||
|
|
||||||
let (regular_tools, mcp_servers) = categorize_tools(node.tools.as_deref());
|
|
||||||
validate_tools_subset(®ular_tools, &mcp_servers, parent_ctx)?;
|
|
||||||
|
|
||||||
let role = build_inline_role(
|
|
||||||
node,
|
|
||||||
instructions.as_deref(),
|
|
||||||
®ular_tools,
|
|
||||||
&mcp_servers,
|
|
||||||
parent_ctx,
|
|
||||||
)?;
|
|
||||||
|
|
||||||
let saved_role = parent_ctx.role.clone();
|
|
||||||
parent_ctx.role = Some(role);
|
|
||||||
let result = match node.timeout {
|
|
||||||
Some(secs) => match timeout(
|
|
||||||
Duration::from_secs(secs),
|
|
||||||
run_with_retries(node, &prompt, parent_ctx),
|
|
||||||
)
|
|
||||||
.await
|
|
||||||
{
|
|
||||||
Ok(r) => r,
|
|
||||||
Err(_) => Err(anyhow!("llm node timed out after {secs}s")),
|
|
||||||
},
|
|
||||||
None => run_with_retries(node, &prompt, parent_ctx).await,
|
|
||||||
};
|
|
||||||
parent_ctx.role = saved_role;
|
|
||||||
result
|
|
||||||
}
|
|
||||||
|
|
||||||
async fn run_with_retries(
|
|
||||||
node: &LlmNode,
|
|
||||||
prompt: &str,
|
|
||||||
ctx: &mut RequestContext,
|
|
||||||
) -> Result<String> {
|
|
||||||
let mut last_err: Option<Error> = None;
|
|
||||||
for attempt in 1..=node.max_attempts {
|
|
||||||
match run_chat_loop(node, prompt, ctx).await {
|
|
||||||
Ok(out) => return Ok(out),
|
|
||||||
Err(e) if is_transient(&e) && attempt < node.max_attempts => {
|
|
||||||
warn!("llm node attempt {attempt} failed (transient): {e}; retrying");
|
|
||||||
last_err = Some(e);
|
|
||||||
}
|
|
||||||
Err(e) => return Err(e),
|
|
||||||
}
|
|
||||||
}
|
|
||||||
Err(last_err.unwrap_or_else(|| anyhow!("llm node exhausted retries")))
|
|
||||||
}
|
|
||||||
|
|
||||||
async fn run_chat_loop(node: &LlmNode, prompt: &str, ctx: &mut RequestContext) -> Result<String> {
|
|
||||||
let abort = create_abort_signal();
|
|
||||||
let app_cfg = Arc::clone(&ctx.app.config);
|
|
||||||
let role_for_input = ctx.role.clone();
|
|
||||||
let mut input = Input::from_str(ctx, prompt, role_for_input);
|
|
||||||
let mut accumulated = String::new();
|
|
||||||
|
|
||||||
for turn in 0..node.max_iterations {
|
|
||||||
let client = input.create_client()?;
|
|
||||||
ctx.before_chat_completion(&input)?;
|
|
||||||
let (output, tool_results) =
|
|
||||||
call_chat_completions(&input, false, false, client.as_ref(), ctx, abort.clone())
|
|
||||||
.await?;
|
|
||||||
ctx.after_chat_completion(app_cfg.as_ref(), &input, &output, &tool_results)?;
|
|
||||||
|
|
||||||
if !output.is_empty() {
|
|
||||||
if !accumulated.is_empty() {
|
|
||||||
accumulated.push('\n');
|
|
||||||
}
|
|
||||||
accumulated.push_str(&output);
|
|
||||||
}
|
|
||||||
|
|
||||||
if tool_results.is_empty() {
|
|
||||||
return Ok(accumulated);
|
|
||||||
}
|
|
||||||
|
|
||||||
if turn + 1 == node.max_iterations {
|
|
||||||
bail!(
|
|
||||||
"llm node hit max_iterations ({}) before LLM concluded",
|
|
||||||
node.max_iterations
|
|
||||||
);
|
|
||||||
}
|
|
||||||
|
|
||||||
input = input.merge_tool_results(output, tool_results);
|
|
||||||
}
|
|
||||||
|
|
||||||
bail!("llm node ended without producing output")
|
|
||||||
}
|
|
||||||
|
|
||||||
fn build_inline_role(
|
|
||||||
node: &LlmNode,
|
|
||||||
instructions: Option<&str>,
|
|
||||||
regular_tools: &[String],
|
|
||||||
mcp_servers: &[String],
|
|
||||||
parent_ctx: &RequestContext,
|
|
||||||
) -> Result<Role> {
|
|
||||||
let mut role = Role::new("llm_node", instructions.unwrap_or(""));
|
|
||||||
|
|
||||||
let model = match &node.model {
|
|
||||||
Some(model_id) => {
|
|
||||||
Model::retrieve_model(parent_ctx.app.config.as_ref(), model_id, ModelType::Chat)
|
|
||||||
.with_context(|| format!("Unknown model '{model_id}' on llm node"))?
|
|
||||||
}
|
|
||||||
None => parent_ctx.current_model().clone(),
|
|
||||||
};
|
|
||||||
role.set_model(model);
|
|
||||||
|
|
||||||
if let Some(t) = node.temperature {
|
|
||||||
role.set_temperature(Some(t));
|
|
||||||
}
|
|
||||||
if let Some(p) = node.top_p {
|
|
||||||
role.set_top_p(Some(p));
|
|
||||||
}
|
|
||||||
|
|
||||||
if node.tools.as_deref().unwrap_or_default().is_empty() {
|
|
||||||
role.set_enabled_tools(Some(String::new()));
|
|
||||||
role.set_enabled_mcp_servers(Some(String::new()));
|
|
||||||
} else {
|
|
||||||
if !regular_tools.is_empty() {
|
|
||||||
role.set_enabled_tools(Some(regular_tools.join(",")));
|
|
||||||
} else {
|
|
||||||
role.set_enabled_tools(Some(String::new()));
|
|
||||||
}
|
|
||||||
if !mcp_servers.is_empty() {
|
|
||||||
role.set_enabled_mcp_servers(Some(mcp_servers.join(",")));
|
|
||||||
} else {
|
|
||||||
role.set_enabled_mcp_servers(Some(String::new()));
|
|
||||||
}
|
|
||||||
}
|
|
||||||
|
|
||||||
Ok(role)
|
|
||||||
}
|
|
||||||
|
|
||||||
fn categorize_tools(entries: Option<&[String]>) -> (Vec<String>, Vec<String>) {
|
|
||||||
let mut regular = Vec::new();
|
|
||||||
let mut mcp = Vec::new();
|
|
||||||
let Some(entries) = entries else {
|
|
||||||
return (regular, mcp);
|
|
||||||
};
|
|
||||||
|
|
||||||
for e in entries {
|
|
||||||
if let Some(server) = e.strip_prefix("mcp:") {
|
|
||||||
mcp.push(server.to_string());
|
|
||||||
} else {
|
|
||||||
regular.push(e.clone());
|
|
||||||
}
|
|
||||||
}
|
|
||||||
|
|
||||||
(regular, mcp)
|
|
||||||
}
|
|
||||||
|
|
||||||
fn validate_tools_subset(
|
|
||||||
regular: &[String],
|
|
||||||
mcp_servers: &[String],
|
|
||||||
parent_ctx: &RequestContext,
|
|
||||||
) -> Result<()> {
|
|
||||||
let agent = parent_ctx
|
|
||||||
.agent
|
|
||||||
.as_ref()
|
|
||||||
.ok_or_else(|| anyhow!("llm node requires an active agent"))?;
|
|
||||||
|
|
||||||
if !regular.is_empty() {
|
|
||||||
let known: HashSet<&str> = agent
|
|
||||||
.functions()
|
|
||||||
.declarations()
|
|
||||||
.iter()
|
|
||||||
.map(|d| d.name.as_str())
|
|
||||||
.collect();
|
|
||||||
for name in regular {
|
|
||||||
if !known.contains(name.as_str()) {
|
|
||||||
let mut avail: Vec<&str> = known.iter().copied().collect();
|
|
||||||
avail.sort();
|
|
||||||
bail!(
|
|
||||||
"llm node references unknown tool '{name}'. Agent '{}' provides: {}",
|
|
||||||
agent.name(),
|
|
||||||
avail.join(", ")
|
|
||||||
);
|
|
||||||
}
|
|
||||||
}
|
|
||||||
}
|
|
||||||
|
|
||||||
if !mcp_servers.is_empty() {
|
|
||||||
let known: HashSet<&str> = agent
|
|
||||||
.mcp_server_names()
|
|
||||||
.iter()
|
|
||||||
.map(|s| s.as_str())
|
|
||||||
.collect();
|
|
||||||
for server in mcp_servers {
|
|
||||||
if !known.contains(server.as_str()) {
|
|
||||||
let mut avail: Vec<&str> = known.iter().copied().collect();
|
|
||||||
avail.sort();
|
|
||||||
bail!(
|
|
||||||
"llm node references unknown MCP server 'mcp:{server}'. \
|
|
||||||
Agent '{}' has MCP servers: [{}]",
|
|
||||||
agent.name(),
|
|
||||||
avail.join(", ")
|
|
||||||
);
|
|
||||||
}
|
|
||||||
}
|
|
||||||
}
|
|
||||||
|
|
||||||
Ok(())
|
|
||||||
}
|
|
||||||
|
|
||||||
fn is_transient(err: &Error) -> bool {
|
|
||||||
let s = format!("{err:#}");
|
|
||||||
s.contains("timed out")
|
|
||||||
|| s.contains("rate limit")
|
|
||||||
|| s.contains("429")
|
|
||||||
|| s.contains("Connection reset")
|
|
||||||
|| s.contains("Connection refused")
|
|
||||||
|| s.contains("produced no output")
|
|
||||||
}
|
|
||||||
|
|
||||||
fn apply_state_updates_with_output(
|
|
||||||
node: &LlmNode,
|
|
||||||
state_manager: &mut StateManager,
|
|
||||||
output: &Value,
|
|
||||||
) {
|
|
||||||
if node.output_schema.is_some()
|
|
||||||
&& let Some(obj) = output.as_object()
|
|
||||||
{
|
|
||||||
for (k, v) in obj {
|
|
||||||
state_manager.state_mut().set(k.clone(), v.clone());
|
|
||||||
}
|
|
||||||
}
|
|
||||||
|
|
||||||
let Some(updates) = &node.state_updates else {
|
|
||||||
return;
|
|
||||||
};
|
|
||||||
let prev_output = state_manager.state().get(OUTPUT_KEY).cloned();
|
|
||||||
state_manager
|
|
||||||
.state_mut()
|
|
||||||
.set(OUTPUT_KEY.into(), output.clone());
|
|
||||||
|
|
||||||
for (key, template) in updates {
|
|
||||||
let value = state_manager.interpolate_lenient(template);
|
|
||||||
state_manager
|
|
||||||
.state_mut()
|
|
||||||
.set(key.clone(), Value::String(value));
|
|
||||||
}
|
|
||||||
|
|
||||||
match prev_output {
|
|
||||||
Some(v) => state_manager.state_mut().set(OUTPUT_KEY.into(), v),
|
|
||||||
None => {
|
|
||||||
state_manager
|
|
||||||
.state_mut()
|
|
||||||
.set(OUTPUT_KEY.into(), Value::Null);
|
|
||||||
}
|
|
||||||
}
|
|
||||||
}
|
|
||||||
|
|
||||||
fn format_schema_hint(schema: &Value) -> String {
|
|
||||||
let schema_json = serde_json::to_string_pretty(schema).unwrap_or_else(|_| schema.to_string());
|
|
||||||
format!(
|
|
||||||
"Respond with a JSON object that matches this schema. Output ONLY the JSON \
|
|
||||||
object with no surrounding prose or markdown fences.\n\nSchema:\n{schema_json}"
|
|
||||||
)
|
|
||||||
}
|
|
||||||
|
|
||||||
#[cfg(test)]
|
|
||||||
mod tests {
|
|
||||||
use super::super::types::*;
|
|
||||||
use super::*;
|
|
||||||
use serde_json::json;
|
|
||||||
use std::collections::HashMap;
|
|
||||||
|
|
||||||
fn manager_with(pairs: &[(&str, Value)]) -> StateManager {
|
|
||||||
let mut map = HashMap::new();
|
|
||||||
for (k, v) in pairs {
|
|
||||||
map.insert((*k).into(), v.clone());
|
|
||||||
}
|
|
||||||
StateManager::new(map)
|
|
||||||
}
|
|
||||||
|
|
||||||
fn node_with(updates: Option<HashMap<String, String>>) -> LlmNode {
|
|
||||||
LlmNode {
|
|
||||||
instructions: Some("sys".into()),
|
|
||||||
prompt: "user".into(),
|
|
||||||
tools: None,
|
|
||||||
model: None,
|
|
||||||
temperature: None,
|
|
||||||
top_p: None,
|
|
||||||
fallback: None,
|
|
||||||
max_attempts: 1,
|
|
||||||
max_iterations: 10,
|
|
||||||
state_updates: updates,
|
|
||||||
output_schema: None,
|
|
||||||
timeout: None,
|
|
||||||
}
|
|
||||||
}
|
|
||||||
|
|
||||||
#[test]
|
|
||||||
fn state_updates_expose_output_during_evaluation() {
|
|
||||||
let mut u = HashMap::new();
|
|
||||||
u.insert("response".into(), "{{output}}".into());
|
|
||||||
let node = node_with(Some(u));
|
|
||||||
let mut state = manager_with(&[]);
|
|
||||||
|
|
||||||
apply_state_updates_with_output(&node, &mut state, &json!("the answer"));
|
|
||||||
|
|
||||||
assert_eq!(state.state().get("response"), Some(&json!("the answer")));
|
|
||||||
}
|
|
||||||
|
|
||||||
#[test]
|
|
||||||
fn state_updates_can_mix_existing_keys_with_output() {
|
|
||||||
let mut u = HashMap::new();
|
|
||||||
u.insert("summary".into(), "{{topic}}: {{output}}".into());
|
|
||||||
let node = node_with(Some(u));
|
|
||||||
let mut state = manager_with(&[("topic", json!("LOINC"))]);
|
|
||||||
|
|
||||||
apply_state_updates_with_output(&node, &mut state, &json!("abc"));
|
|
||||||
|
|
||||||
assert_eq!(state.state().get("summary"), Some(&json!("LOINC: abc")));
|
|
||||||
}
|
|
||||||
|
|
||||||
#[test]
|
|
||||||
fn output_key_is_cleared_after_state_updates() {
|
|
||||||
let mut u = HashMap::new();
|
|
||||||
u.insert("k".into(), "{{output}}".into());
|
|
||||||
let node = node_with(Some(u));
|
|
||||||
let mut state = manager_with(&[]);
|
|
||||||
|
|
||||||
apply_state_updates_with_output(&node, &mut state, &json!("anything"));
|
|
||||||
|
|
||||||
assert_eq!(state.state().get(OUTPUT_KEY), Some(&json!(null)));
|
|
||||||
}
|
|
||||||
|
|
||||||
#[test]
|
|
||||||
fn pre_existing_output_value_is_restored() {
|
|
||||||
let mut u = HashMap::new();
|
|
||||||
u.insert("greeting".into(), "{{output}}".into());
|
|
||||||
let node = node_with(Some(u));
|
|
||||||
let mut state = manager_with(&[("output", json!("preserved"))]);
|
|
||||||
|
|
||||||
apply_state_updates_with_output(&node, &mut state, &json!("new"));
|
|
||||||
|
|
||||||
assert_eq!(state.state().get("greeting"), Some(&json!("new")));
|
|
||||||
assert_eq!(state.state().get(OUTPUT_KEY), Some(&json!("preserved")));
|
|
||||||
}
|
|
||||||
|
|
||||||
#[test]
|
|
||||||
fn no_state_updates_is_a_noop() {
|
|
||||||
let node = node_with(None);
|
|
||||||
let mut state = manager_with(&[("k", json!("v"))]);
|
|
||||||
|
|
||||||
apply_state_updates_with_output(&node, &mut state, &json!("x"));
|
|
||||||
|
|
||||||
assert_eq!(state.state().get("k"), Some(&json!("v")));
|
|
||||||
assert!(state.state().get(OUTPUT_KEY).is_none());
|
|
||||||
}
|
|
||||||
|
|
||||||
#[test]
|
|
||||||
fn outcome_from_success_is_continue() {
|
|
||||||
assert_eq!(
|
|
||||||
outcome_from(None, Some("fb")).unwrap(),
|
|
||||||
LlmExecutionOutcome::Continue
|
|
||||||
);
|
|
||||||
assert_eq!(
|
|
||||||
outcome_from(None, None).unwrap(),
|
|
||||||
LlmExecutionOutcome::Continue
|
|
||||||
);
|
|
||||||
}
|
|
||||||
|
|
||||||
#[test]
|
|
||||||
fn outcome_from_failure_with_fallback_is_fell_back() {
|
|
||||||
assert_eq!(
|
|
||||||
outcome_from(Some("HTTP 404"), Some("fb")).unwrap(),
|
|
||||||
LlmExecutionOutcome::FellBack("fb".to_string())
|
|
||||||
);
|
|
||||||
}
|
|
||||||
|
|
||||||
#[test]
|
|
||||||
fn outcome_from_failure_without_fallback_propagates_error() {
|
|
||||||
let err = outcome_from(Some("HTTP 404"), None).unwrap_err();
|
|
||||||
let msg = format!("{err:#}");
|
|
||||||
assert!(msg.contains("no fallback declared"), "got: {msg}");
|
|
||||||
assert!(msg.contains("HTTP 404"), "got: {msg}");
|
|
||||||
}
|
|
||||||
|
|
||||||
fn node_with_schema(updates: Option<HashMap<String, String>>, schema: Value) -> LlmNode {
|
|
||||||
let mut n = node_with(updates);
|
|
||||||
n.output_schema = Some(schema);
|
|
||||||
n
|
|
||||||
}
|
|
||||||
|
|
||||||
#[test]
|
|
||||||
fn output_schema_auto_merges_top_level_keys() {
|
|
||||||
let node = node_with_schema(None, json!({"type": "object"}));
|
|
||||||
let mut state = manager_with(&[]);
|
|
||||||
let output = json!({"goal": "do X", "summary": "details"});
|
|
||||||
|
|
||||||
apply_state_updates_with_output(&node, &mut state, &output);
|
|
||||||
|
|
||||||
assert_eq!(state.state().get("goal"), Some(&json!("do X")));
|
|
||||||
assert_eq!(state.state().get("summary"), Some(&json!("details")));
|
|
||||||
}
|
|
||||||
|
|
||||||
#[test]
|
|
||||||
fn output_schema_preserves_nested_value_types() {
|
|
||||||
let node = node_with_schema(None, json!({"type": "object"}));
|
|
||||||
let mut state = manager_with(&[]);
|
|
||||||
let output = json!({
|
|
||||||
"tags": ["a", "b"],
|
|
||||||
"config": { "key": "value" },
|
|
||||||
"count": 42
|
|
||||||
});
|
|
||||||
|
|
||||||
apply_state_updates_with_output(&node, &mut state, &output);
|
|
||||||
|
|
||||||
assert_eq!(state.state().get("tags"), Some(&json!(["a", "b"])));
|
|
||||||
assert_eq!(state.state().get("config"), Some(&json!({"key": "value"})));
|
|
||||||
assert_eq!(state.state().get("count"), Some(&json!(42)));
|
|
||||||
}
|
|
||||||
|
|
||||||
#[test]
|
|
||||||
fn output_schema_explicit_state_updates_override_auto_merge() {
|
|
||||||
let mut u = HashMap::new();
|
|
||||||
u.insert("goal".into(), "renamed-{{output.goal}}".into());
|
|
||||||
let node = node_with_schema(Some(u), json!({"type": "object"}));
|
|
||||||
let mut state = manager_with(&[]);
|
|
||||||
let output = json!({"goal": "do X"});
|
|
||||||
|
|
||||||
apply_state_updates_with_output(&node, &mut state, &output);
|
|
||||||
|
|
||||||
assert_eq!(state.state().get("goal"), Some(&json!("renamed-do X")));
|
|
||||||
}
|
|
||||||
|
|
||||||
#[test]
|
|
||||||
fn output_schema_skips_auto_merge_for_non_object() {
|
|
||||||
let node = node_with_schema(None, json!({"type": "array"}));
|
|
||||||
let mut state = manager_with(&[]);
|
|
||||||
let output = json!([1, 2, 3]);
|
|
||||||
|
|
||||||
apply_state_updates_with_output(&node, &mut state, &output);
|
|
||||||
|
|
||||||
assert!(state.state().get("0").is_none());
|
|
||||||
assert!(state.state().get(OUTPUT_KEY).is_none());
|
|
||||||
}
|
|
||||||
|
|
||||||
#[test]
|
|
||||||
fn no_schema_does_not_auto_merge() {
|
|
||||||
let node = node_with(None);
|
|
||||||
let mut state = manager_with(&[]);
|
|
||||||
let output = json!({"goal": "do X"});
|
|
||||||
|
|
||||||
apply_state_updates_with_output(&node, &mut state, &output);
|
|
||||||
|
|
||||||
assert!(state.state().get("goal").is_none());
|
|
||||||
}
|
|
||||||
|
|
||||||
#[test]
|
|
||||||
fn format_schema_hint_includes_schema_and_instruction() {
|
|
||||||
let schema = json!({"type": "object", "properties": {"goal": {"type": "string"}}});
|
|
||||||
|
|
||||||
let hint = format_schema_hint(&schema);
|
|
||||||
|
|
||||||
assert!(hint.contains("Schema:"));
|
|
||||||
assert!(hint.contains("\"goal\""));
|
|
||||||
assert!(hint.contains("JSON"));
|
|
||||||
assert!(hint.contains("ONLY"));
|
|
||||||
}
|
|
||||||
|
|
||||||
#[test]
|
|
||||||
fn categorize_tools_splits_mcp_and_regular() {
|
|
||||||
let entries = vec![
|
|
||||||
"read_query".to_string(),
|
|
||||||
"mcp:pubmed-search".to_string(),
|
|
||||||
"web_search_loki".to_string(),
|
|
||||||
"mcp:github".to_string(),
|
|
||||||
];
|
|
||||||
|
|
||||||
let (regular, mcp) = categorize_tools(Some(&entries));
|
|
||||||
|
|
||||||
assert_eq!(regular, vec!["read_query", "web_search_loki"]);
|
|
||||||
assert_eq!(mcp, vec!["pubmed-search", "github"]);
|
|
||||||
}
|
|
||||||
|
|
||||||
#[test]
|
|
||||||
fn categorize_tools_with_none_returns_empty() {
|
|
||||||
let (regular, mcp) = categorize_tools(None);
|
|
||||||
|
|
||||||
assert!(regular.is_empty());
|
|
||||||
assert!(mcp.is_empty());
|
|
||||||
}
|
|
||||||
|
|
||||||
#[test]
|
|
||||||
fn categorize_tools_with_empty_returns_empty() {
|
|
||||||
let (regular, mcp) = categorize_tools(Some(&[]));
|
|
||||||
|
|
||||||
assert!(regular.is_empty());
|
|
||||||
assert!(mcp.is_empty());
|
|
||||||
}
|
|
||||||
|
|
||||||
#[test]
|
|
||||||
fn is_transient_matches_expected_signatures() {
|
|
||||||
assert!(is_transient(&anyhow!("request timed out after 30s")));
|
|
||||||
assert!(is_transient(&anyhow!("rate limit reached")));
|
|
||||||
assert!(is_transient(&anyhow!("429 too many requests")));
|
|
||||||
assert!(is_transient(&anyhow!("Connection reset by peer")));
|
|
||||||
assert!(is_transient(&anyhow!("Connection refused")));
|
|
||||||
assert!(is_transient(&anyhow!("llm produced no output")));
|
|
||||||
}
|
|
||||||
|
|
||||||
#[test]
|
|
||||||
fn is_transient_rejects_non_transient_errors() {
|
|
||||||
assert!(!is_transient(&anyhow!("Unknown model 'foo'")));
|
|
||||||
assert!(!is_transient(&anyhow!(
|
|
||||||
"llm node references unknown tool 'bad'"
|
|
||||||
)));
|
|
||||||
assert!(!is_transient(&anyhow!("hit max_iterations")));
|
|
||||||
assert!(!is_transient(&anyhow!("authentication failed")));
|
|
||||||
}
|
|
||||||
}
|
|
||||||
@@ -1,312 +0,0 @@
|
|||||||
use super::state::StateManager;
|
|
||||||
use super::types::{Node, NodeType};
|
|
||||||
use crate::utils::dimmed_text;
|
|
||||||
use chrono::Local;
|
|
||||||
use indexmap::IndexMap;
|
|
||||||
use std::cmp::Reverse;
|
|
||||||
use std::time::Duration;
|
|
||||||
|
|
||||||
fn ts() -> String {
|
|
||||||
Local::now().format("%H:%M:%S").to_string()
|
|
||||||
}
|
|
||||||
|
|
||||||
fn fmt_secs(elapsed: Duration) -> String {
|
|
||||||
let secs = elapsed.as_secs_f64();
|
|
||||||
if secs < 1.0 {
|
|
||||||
format!("{}ms", elapsed.as_millis())
|
|
||||||
} else {
|
|
||||||
format!("{secs:.2}s")
|
|
||||||
}
|
|
||||||
}
|
|
||||||
|
|
||||||
#[derive(Debug, Clone, Default)]
|
|
||||||
struct NodeTiming {
|
|
||||||
count: usize,
|
|
||||||
total: Duration,
|
|
||||||
max: Duration,
|
|
||||||
}
|
|
||||||
|
|
||||||
impl NodeTiming {
|
|
||||||
fn record(&mut self, elapsed: Duration) {
|
|
||||||
self.count += 1;
|
|
||||||
self.total += elapsed;
|
|
||||||
if elapsed > self.max {
|
|
||||||
self.max = elapsed;
|
|
||||||
}
|
|
||||||
}
|
|
||||||
}
|
|
||||||
|
|
||||||
pub struct GraphLogger {
|
|
||||||
graph_name: String,
|
|
||||||
log_state_snapshots: bool,
|
|
||||||
silent: bool,
|
|
||||||
timings: IndexMap<String, NodeTiming>,
|
|
||||||
}
|
|
||||||
|
|
||||||
impl GraphLogger {
|
|
||||||
pub fn with_visibility(graph_name: &str, log_state_snapshots: bool, silent: bool) -> Self {
|
|
||||||
Self {
|
|
||||||
graph_name: graph_name.to_string(),
|
|
||||||
log_state_snapshots,
|
|
||||||
silent,
|
|
||||||
timings: IndexMap::new(),
|
|
||||||
}
|
|
||||||
}
|
|
||||||
|
|
||||||
pub fn graph_start(&self, start_node: &str, node_count: usize) {
|
|
||||||
info!(
|
|
||||||
"[graph:{}] start at '{}' ({} nodes)",
|
|
||||||
self.graph_name, start_node, node_count
|
|
||||||
);
|
|
||||||
if !self.silent {
|
|
||||||
eprintln!(
|
|
||||||
"{}",
|
|
||||||
dimmed_text(&format!(
|
|
||||||
"▸ graph: {} (start: {start_node})",
|
|
||||||
self.graph_name
|
|
||||||
))
|
|
||||||
);
|
|
||||||
}
|
|
||||||
}
|
|
||||||
|
|
||||||
pub fn graph_complete(&self, end_node: &str, elapsed: Duration) {
|
|
||||||
info!(
|
|
||||||
"[graph:{}] end '{}' (elapsed {:?})",
|
|
||||||
self.graph_name, end_node, elapsed
|
|
||||||
);
|
|
||||||
if !self.silent {
|
|
||||||
eprintln!(
|
|
||||||
"{}",
|
|
||||||
dimmed_text(&format!("▸ graph done in {:.2}s", elapsed.as_secs_f64()))
|
|
||||||
);
|
|
||||||
}
|
|
||||||
self.log_performance_summary();
|
|
||||||
}
|
|
||||||
|
|
||||||
pub fn graph_error(&self, error: &anyhow::Error) {
|
|
||||||
error!("[graph:{}] execution failed: {error:#}", self.graph_name);
|
|
||||||
}
|
|
||||||
|
|
||||||
pub fn node_entry(&self, node: &Node, visit: usize) {
|
|
||||||
debug!(
|
|
||||||
"[graph:{}] entering '{}' (visit {visit})",
|
|
||||||
self.graph_name, node.id
|
|
||||||
);
|
|
||||||
}
|
|
||||||
|
|
||||||
pub fn silent(&self) -> bool {
|
|
||||||
self.silent
|
|
||||||
}
|
|
||||||
|
|
||||||
pub fn node_start(&self, node: &Node, in_super_step: bool) {
|
|
||||||
narrate_node_start(self.silent, node, in_super_step);
|
|
||||||
}
|
|
||||||
|
|
||||||
pub fn super_step_start(&self, branches: &[String]) {
|
|
||||||
if self.silent {
|
|
||||||
return;
|
|
||||||
}
|
|
||||||
eprintln!(
|
|
||||||
"{}",
|
|
||||||
dimmed_text(&format!(
|
|
||||||
"▸ {} super-step start: {}",
|
|
||||||
ts(),
|
|
||||||
branches.join(", ")
|
|
||||||
))
|
|
||||||
);
|
|
||||||
}
|
|
||||||
|
|
||||||
pub fn super_step_end(&self, targets: &[String]) {
|
|
||||||
if self.silent {
|
|
||||||
return;
|
|
||||||
}
|
|
||||||
let route = if targets.is_empty() {
|
|
||||||
String::new()
|
|
||||||
} else {
|
|
||||||
format!(" -> {}", targets.join(", "))
|
|
||||||
};
|
|
||||||
eprintln!(
|
|
||||||
"{}",
|
|
||||||
dimmed_text(&format!("▸ {} super-step end{route}", ts()))
|
|
||||||
);
|
|
||||||
}
|
|
||||||
|
|
||||||
pub fn record_timing(&mut self, node_id: &str, elapsed: Duration) {
|
|
||||||
self.timings
|
|
||||||
.entry(node_id.to_string())
|
|
||||||
.or_default()
|
|
||||||
.record(elapsed);
|
|
||||||
}
|
|
||||||
|
|
||||||
pub fn routing(&self, from: &str, to: &str) {
|
|
||||||
debug!("[graph:{}] {from} -> {to}", self.graph_name);
|
|
||||||
}
|
|
||||||
|
|
||||||
pub fn validation_warning(&self, node_id: Option<&str>, message: &str) {
|
|
||||||
match node_id {
|
|
||||||
Some(id) => warn!("[graph:{}] [{id}] {message}", self.graph_name),
|
|
||||||
None => warn!("[graph:{}] {message}", self.graph_name),
|
|
||||||
}
|
|
||||||
}
|
|
||||||
|
|
||||||
pub fn state_snapshot(&self, node_id: &str, state: &StateManager) {
|
|
||||||
if !self.log_state_snapshots {
|
|
||||||
return;
|
|
||||||
}
|
|
||||||
let snapshot = state.snapshot();
|
|
||||||
let mut keys: Vec<&str> = snapshot.keys().map(String::as_str).collect();
|
|
||||||
keys.sort_unstable();
|
|
||||||
|
|
||||||
debug!(
|
|
||||||
"[graph:{}] [{node_id}] state: {} bytes, keys={:?}",
|
|
||||||
self.graph_name,
|
|
||||||
state.size_bytes(),
|
|
||||||
keys
|
|
||||||
);
|
|
||||||
trace!(
|
|
||||||
"[graph:{}] [{node_id}] full state: {:?}",
|
|
||||||
self.graph_name, snapshot
|
|
||||||
);
|
|
||||||
}
|
|
||||||
|
|
||||||
fn log_performance_summary(&self) {
|
|
||||||
if self.timings.is_empty() {
|
|
||||||
return;
|
|
||||||
}
|
|
||||||
let mut rows: Vec<(&String, &NodeTiming)> = self.timings.iter().collect();
|
|
||||||
rows.sort_by_key(|b| Reverse(b.1.total));
|
|
||||||
|
|
||||||
info!(
|
|
||||||
"[graph:{}] performance summary (slowest first):",
|
|
||||||
self.graph_name
|
|
||||||
);
|
|
||||||
|
|
||||||
for (node_id, t) in rows {
|
|
||||||
let avg = t.total / t.count.max(1) as u32;
|
|
||||||
info!(
|
|
||||||
"[graph:{}] {node_id}: {} visit(s), total {}ms, avg {}ms, max {}ms",
|
|
||||||
self.graph_name,
|
|
||||||
t.count,
|
|
||||||
t.total.as_millis(),
|
|
||||||
avg.as_millis(),
|
|
||||||
t.max.as_millis(),
|
|
||||||
);
|
|
||||||
}
|
|
||||||
}
|
|
||||||
}
|
|
||||||
|
|
||||||
pub fn narrate_node_start(silent: bool, node: &Node, in_super_step: bool) {
|
|
||||||
if silent {
|
|
||||||
return;
|
|
||||||
}
|
|
||||||
let indent = if in_super_step { " " } else { "" };
|
|
||||||
let label = node_type_label(node);
|
|
||||||
eprintln!(
|
|
||||||
"{}",
|
|
||||||
dimmed_text(&format!("▸ {} {indent}{} ({label}) start", ts(), node.id))
|
|
||||||
);
|
|
||||||
}
|
|
||||||
|
|
||||||
pub fn narrate_node_complete(
|
|
||||||
silent: bool,
|
|
||||||
node: &Node,
|
|
||||||
elapsed: Duration,
|
|
||||||
next_target: Option<&str>,
|
|
||||||
in_super_step: bool,
|
|
||||||
) {
|
|
||||||
if silent {
|
|
||||||
return;
|
|
||||||
}
|
|
||||||
let indent = if in_super_step { " " } else { "" };
|
|
||||||
let label = node_type_label(node);
|
|
||||||
let dur = fmt_secs(elapsed);
|
|
||||||
let route = next_target.map(|t| format!(" -> {t}")).unwrap_or_default();
|
|
||||||
eprintln!(
|
|
||||||
"{}",
|
|
||||||
dimmed_text(&format!(
|
|
||||||
"▸ {} {indent}{} ({label}) done in {dur}{route}",
|
|
||||||
ts(),
|
|
||||||
node.id
|
|
||||||
))
|
|
||||||
);
|
|
||||||
}
|
|
||||||
|
|
||||||
pub fn narrate_node_failed(
|
|
||||||
silent: bool,
|
|
||||||
node: &Node,
|
|
||||||
elapsed: Duration,
|
|
||||||
err: &str,
|
|
||||||
in_super_step: bool,
|
|
||||||
) {
|
|
||||||
if silent {
|
|
||||||
return;
|
|
||||||
}
|
|
||||||
let indent = if in_super_step { " " } else { "" };
|
|
||||||
let label = node_type_label(node);
|
|
||||||
let dur = fmt_secs(elapsed);
|
|
||||||
let excerpt: String = err.chars().take(120).collect();
|
|
||||||
eprintln!(
|
|
||||||
"{}",
|
|
||||||
dimmed_text(&format!(
|
|
||||||
"▸ {} {indent}{} ({label}) FAILED in {dur} -- {excerpt}",
|
|
||||||
ts(),
|
|
||||||
node.id
|
|
||||||
))
|
|
||||||
);
|
|
||||||
}
|
|
||||||
|
|
||||||
pub(super) fn node_type_label(node: &Node) -> &'static str {
|
|
||||||
match &node.node_type {
|
|
||||||
NodeType::Agent(_) => "agent",
|
|
||||||
NodeType::Script(_) => "script",
|
|
||||||
NodeType::Approval(_) => "approval",
|
|
||||||
NodeType::Input(_) => "input",
|
|
||||||
NodeType::Llm(_) => "llm",
|
|
||||||
NodeType::Rag(_) => "rag",
|
|
||||||
NodeType::End(_) => "end",
|
|
||||||
NodeType::Map(_) => "map",
|
|
||||||
}
|
|
||||||
}
|
|
||||||
|
|
||||||
#[cfg(test)]
|
|
||||||
mod tests {
|
|
||||||
use super::*;
|
|
||||||
|
|
||||||
#[test]
|
|
||||||
fn records_and_aggregates_node_timings() {
|
|
||||||
let mut logger = GraphLogger::with_visibility("g", false, false);
|
|
||||||
logger.record_timing("a", Duration::from_millis(100));
|
|
||||||
logger.record_timing("a", Duration::from_millis(300));
|
|
||||||
logger.record_timing("b", Duration::from_millis(50));
|
|
||||||
|
|
||||||
let a = logger.timings.get("a").unwrap();
|
|
||||||
assert_eq!(a.count, 2);
|
|
||||||
assert_eq!(a.total, Duration::from_millis(400));
|
|
||||||
assert_eq!(a.max, Duration::from_millis(300));
|
|
||||||
|
|
||||||
let b = logger.timings.get("b").unwrap();
|
|
||||||
assert_eq!(b.count, 1);
|
|
||||||
assert_eq!(b.total, Duration::from_millis(50));
|
|
||||||
}
|
|
||||||
|
|
||||||
#[test]
|
|
||||||
fn node_timing_max_tracks_largest() {
|
|
||||||
let mut t = NodeTiming::default();
|
|
||||||
|
|
||||||
t.record(Duration::from_millis(10));
|
|
||||||
t.record(Duration::from_millis(80));
|
|
||||||
t.record(Duration::from_millis(40));
|
|
||||||
|
|
||||||
assert_eq!(t.max, Duration::from_millis(80));
|
|
||||||
assert_eq!(t.count, 3);
|
|
||||||
assert_eq!(t.total, Duration::from_millis(130));
|
|
||||||
}
|
|
||||||
|
|
||||||
#[test]
|
|
||||||
fn new_logger_has_no_timings() {
|
|
||||||
let logger = GraphLogger::with_visibility("g", true, false);
|
|
||||||
|
|
||||||
assert!(logger.timings.is_empty());
|
|
||||||
assert!(logger.log_state_snapshots);
|
|
||||||
}
|
|
||||||
}
|
|
||||||
@@ -1,152 +0,0 @@
|
|||||||
use super::agent::AgentNodeExecutor;
|
|
||||||
use super::executor::StepContext;
|
|
||||||
use super::llm::LlmNodeExecutor;
|
|
||||||
use super::rag::RagNodeExecutor;
|
|
||||||
use super::state::StateManager;
|
|
||||||
use super::types::{MapNode, NodeType};
|
|
||||||
use crate::config::{RenderMode, RequestContext};
|
|
||||||
use crate::graph::type_name;
|
|
||||||
use anyhow::{Context, Result, anyhow};
|
|
||||||
use futures_util::future::join_all;
|
|
||||||
use serde_json::Value;
|
|
||||||
use std::collections::HashMap;
|
|
||||||
use std::sync::Arc;
|
|
||||||
use tokio::sync::Semaphore;
|
|
||||||
|
|
||||||
pub(super) struct MapNodeExecutor;
|
|
||||||
|
|
||||||
impl MapNodeExecutor {
|
|
||||||
pub(super) async fn execute(
|
|
||||||
node: &MapNode,
|
|
||||||
state: &mut StateManager,
|
|
||||||
ctx: &mut RequestContext,
|
|
||||||
step_ctx: &StepContext<'_>,
|
|
||||||
node_id: &str,
|
|
||||||
) -> Result<()> {
|
|
||||||
let over_value = state
|
|
||||||
.interpolate_raw(&node.over)
|
|
||||||
.with_context(|| format!("map node '{node_id}': evaluating `over` template"))?;
|
|
||||||
|
|
||||||
let items = over_value.as_array().ok_or_else(|| {
|
|
||||||
anyhow!(
|
|
||||||
"map node '{}': `over` template '{}' must resolve to an array, got {}",
|
|
||||||
node_id,
|
|
||||||
node.over,
|
|
||||||
type_name(&over_value)
|
|
||||||
)
|
|
||||||
})?;
|
|
||||||
let items = items.clone();
|
|
||||||
|
|
||||||
let branch_node = step_ctx
|
|
||||||
.graph
|
|
||||||
.get_node(&node.branch)
|
|
||||||
.ok_or_else(|| {
|
|
||||||
anyhow!(
|
|
||||||
"map node '{node_id}': branch '{}' not found in graph",
|
|
||||||
node.branch
|
|
||||||
)
|
|
||||||
})?
|
|
||||||
.clone();
|
|
||||||
|
|
||||||
let max_conc = node
|
|
||||||
.max_concurrency
|
|
||||||
.unwrap_or(step_ctx.max_concurrency)
|
|
||||||
.max(1);
|
|
||||||
let semaphore = Arc::new(Semaphore::new(max_conc));
|
|
||||||
let mut sub_tasks = Vec::with_capacity(items.len());
|
|
||||||
|
|
||||||
for (idx, item) in items.iter().enumerate() {
|
|
||||||
let item = item.clone();
|
|
||||||
let as_name = node.as_name.clone();
|
|
||||||
let branch_clone = branch_node.clone();
|
|
||||||
let mut sub_state = state.fork_for_branch_state();
|
|
||||||
let mut sub_ctx = ctx.fork_for_branch();
|
|
||||||
sub_ctx.render_mode = RenderMode::Silent;
|
|
||||||
let script_clone = step_ctx.script_executor.clone();
|
|
||||||
let sub_branch_id = node.branch.clone();
|
|
||||||
let sem = semaphore.clone();
|
|
||||||
let abort = step_ctx.abort_signal.clone();
|
|
||||||
|
|
||||||
sub_state.state_mut().set(as_name, item);
|
|
||||||
|
|
||||||
let task = tokio::spawn(async move {
|
|
||||||
let _permit = sem
|
|
||||||
.acquire()
|
|
||||||
.await
|
|
||||||
.expect("map semaphore should not be closed");
|
|
||||||
if abort.aborted() {
|
|
||||||
return (
|
|
||||||
idx,
|
|
||||||
sub_state,
|
|
||||||
Err(anyhow!("map sub-branch [{idx}] aborted")),
|
|
||||||
);
|
|
||||||
}
|
|
||||||
let mut state = sub_state;
|
|
||||||
let mut ctx = sub_ctx;
|
|
||||||
|
|
||||||
let exec_result: Result<()> = match &branch_clone.node_type {
|
|
||||||
NodeType::Llm(n) => LlmNodeExecutor::execute(n, &mut state, &mut ctx)
|
|
||||||
.await
|
|
||||||
.map(|_| ()),
|
|
||||||
NodeType::Agent(n) => AgentNodeExecutor::execute(n, &mut state, &mut ctx)
|
|
||||||
.await
|
|
||||||
.map(|_| ()),
|
|
||||||
NodeType::Rag(n) => {
|
|
||||||
RagNodeExecutor::execute(n, &sub_branch_id, &mut state, &mut ctx).await
|
|
||||||
}
|
|
||||||
NodeType::Script(n) => script_clone.execute(n, &mut state).await.map(|_| ()),
|
|
||||||
_ => Err(anyhow!(
|
|
||||||
"map branch '{}' has type that cannot run inside a map \
|
|
||||||
(validator should have caught this; internal error)",
|
|
||||||
branch_clone.id
|
|
||||||
)),
|
|
||||||
};
|
|
||||||
|
|
||||||
(idx, state, exec_result)
|
|
||||||
});
|
|
||||||
sub_tasks.push(task);
|
|
||||||
}
|
|
||||||
|
|
||||||
let joined = join_all(sub_tasks).await;
|
|
||||||
|
|
||||||
// Collect outputs keyed by input index so order is preserved regardless of finish order.
|
|
||||||
let mut outputs: HashMap<usize, Value> = HashMap::new();
|
|
||||||
for join_result in joined {
|
|
||||||
let (idx, sub_state, exec_result) =
|
|
||||||
join_result.map_err(|e| anyhow!("map sub-branch panicked: {e}"))?;
|
|
||||||
|
|
||||||
exec_result
|
|
||||||
.with_context(|| format!("map node '{node_id}': sub-branch [{idx}] failed"))?;
|
|
||||||
|
|
||||||
let output_value = sub_state
|
|
||||||
.state()
|
|
||||||
.get(&node.output_key)
|
|
||||||
.cloned()
|
|
||||||
.ok_or_else(|| {
|
|
||||||
anyhow!(
|
|
||||||
"map node '{node_id}': sub-branch [{idx}] did not write \
|
|
||||||
`output_key` '{}'",
|
|
||||||
node.output_key
|
|
||||||
)
|
|
||||||
})?;
|
|
||||||
|
|
||||||
outputs.insert(idx, output_value);
|
|
||||||
}
|
|
||||||
|
|
||||||
let mut collected = Vec::with_capacity(items.len());
|
|
||||||
for idx in 0..items.len() {
|
|
||||||
let value = outputs.remove(&idx).ok_or_else(|| {
|
|
||||||
anyhow!(
|
|
||||||
"map node '{node_id}': internal error: missing result for sub-branch [{idx}]"
|
|
||||||
)
|
|
||||||
})?;
|
|
||||||
collected.push(value);
|
|
||||||
}
|
|
||||||
|
|
||||||
state
|
|
||||||
.state_mut()
|
|
||||||
.set(node.collect_into.clone(), Value::Array(collected));
|
|
||||||
|
|
||||||
Ok(())
|
|
||||||
}
|
|
||||||
}
|
|
||||||
@@ -1,39 +0,0 @@
|
|||||||
pub mod agent;
|
|
||||||
pub mod dispatch;
|
|
||||||
pub mod executor;
|
|
||||||
pub mod llm;
|
|
||||||
pub mod logging;
|
|
||||||
pub mod map;
|
|
||||||
pub mod parser;
|
|
||||||
pub mod rag;
|
|
||||||
pub mod reducer;
|
|
||||||
pub mod script;
|
|
||||||
pub mod staging;
|
|
||||||
pub mod state;
|
|
||||||
pub mod structured;
|
|
||||||
pub mod types;
|
|
||||||
pub mod user_interaction;
|
|
||||||
pub mod validator;
|
|
||||||
|
|
||||||
pub use dispatch::{active_agent_graph_name, run_active_agent_graph};
|
|
||||||
pub use executor::GraphExecutor;
|
|
||||||
pub use parser::{GraphParser, agent_has_graph};
|
|
||||||
use serde_json::Value;
|
|
||||||
pub use types::{Graph, NodeType};
|
|
||||||
|
|
||||||
pub const GRAPH_SCHEMA_VERSION: &str = "1.0";
|
|
||||||
|
|
||||||
pub const DEFAULT_MAX_LOOP_ITERATIONS: usize = 100;
|
|
||||||
|
|
||||||
pub const MAX_STATE_SIZE_BYTES: usize = 32 * 1024;
|
|
||||||
|
|
||||||
pub(in crate::graph) fn type_name(value: &Value) -> &'static str {
|
|
||||||
match value {
|
|
||||||
Value::Null => "null",
|
|
||||||
Value::Bool(_) => "bool",
|
|
||||||
Value::Number(_) => "number",
|
|
||||||
Value::String(_) => "string",
|
|
||||||
Value::Array(_) => "array",
|
|
||||||
Value::Object(_) => "object",
|
|
||||||
}
|
|
||||||
}
|
|
||||||
@@ -1,462 +0,0 @@
|
|||||||
use super::types::Graph;
|
|
||||||
use crate::config::paths;
|
|
||||||
use anyhow::{Context, Error, Result, anyhow, bail};
|
|
||||||
use std::fs::read_to_string;
|
|
||||||
use std::path::{Path, PathBuf};
|
|
||||||
|
|
||||||
const SUPPORTED_VERSIONS: &[&str] = &["1.0"];
|
|
||||||
|
|
||||||
pub struct GraphParser {
|
|
||||||
base_dir: PathBuf,
|
|
||||||
}
|
|
||||||
|
|
||||||
impl GraphParser {
|
|
||||||
pub fn new(base_dir: impl Into<PathBuf>) -> Self {
|
|
||||||
Self {
|
|
||||||
base_dir: base_dir.into(),
|
|
||||||
}
|
|
||||||
}
|
|
||||||
|
|
||||||
pub fn load_from_file(&self, path: impl AsRef<Path>) -> Result<Graph> {
|
|
||||||
let path = path.as_ref();
|
|
||||||
let full_path = if path.is_absolute() {
|
|
||||||
path.to_path_buf()
|
|
||||||
} else {
|
|
||||||
self.base_dir.join(path)
|
|
||||||
};
|
|
||||||
|
|
||||||
let contents = read_to_string(&full_path)
|
|
||||||
.with_context(|| format!("Failed to read graph file at '{}'", full_path.display()))?;
|
|
||||||
|
|
||||||
self.load_from_string(&contents)
|
|
||||||
.with_context(|| format!("Failed to parse graph file at '{}'", full_path.display()))
|
|
||||||
}
|
|
||||||
|
|
||||||
pub fn load_from_string(&self, yaml: &str) -> Result<Graph> {
|
|
||||||
let mut graph: Graph = serde_yaml::from_str(yaml).map_err(enhance_yaml_error)?;
|
|
||||||
|
|
||||||
validate_schema_version(&graph.version)?;
|
|
||||||
|
|
||||||
for (key, node) in &mut graph.nodes {
|
|
||||||
if node.id.is_empty() {
|
|
||||||
node.id = key.clone();
|
|
||||||
} else if &node.id != key {
|
|
||||||
bail!(
|
|
||||||
"Node ID mismatch: key '{}' does not match node.id '{}'",
|
|
||||||
key,
|
|
||||||
node.id
|
|
||||||
);
|
|
||||||
}
|
|
||||||
}
|
|
||||||
|
|
||||||
validate_structure(&graph)?;
|
|
||||||
|
|
||||||
Ok(graph)
|
|
||||||
}
|
|
||||||
}
|
|
||||||
|
|
||||||
fn validate_schema_version(version: &str) -> Result<()> {
|
|
||||||
if !SUPPORTED_VERSIONS.contains(&version) {
|
|
||||||
bail!(
|
|
||||||
"Unsupported graph schema version '{}'. Supported versions: {}",
|
|
||||||
version,
|
|
||||||
SUPPORTED_VERSIONS.join(", ")
|
|
||||||
);
|
|
||||||
}
|
|
||||||
|
|
||||||
Ok(())
|
|
||||||
}
|
|
||||||
|
|
||||||
fn validate_structure(graph: &Graph) -> Result<()> {
|
|
||||||
if graph.name.is_empty() {
|
|
||||||
bail!("Graph must have a non-empty 'name' field");
|
|
||||||
}
|
|
||||||
|
|
||||||
if graph.nodes.is_empty() {
|
|
||||||
bail!("Graph '{}' has no nodes defined", graph.name);
|
|
||||||
}
|
|
||||||
|
|
||||||
if !graph.has_node(&graph.start) {
|
|
||||||
bail!(
|
|
||||||
"Start node '{}' not found in graph '{}'. Available nodes: {}",
|
|
||||||
graph.start,
|
|
||||||
graph.name,
|
|
||||||
graph.node_ids().join(", ")
|
|
||||||
);
|
|
||||||
}
|
|
||||||
|
|
||||||
Ok(())
|
|
||||||
}
|
|
||||||
|
|
||||||
fn enhance_yaml_error(error: serde_yaml::Error) -> Error {
|
|
||||||
let msg = error.to_string();
|
|
||||||
|
|
||||||
let hint = if msg.contains("missing field") {
|
|
||||||
"\n\nHint: Check that all required fields are present.\n\
|
|
||||||
Top-level required fields: `name`, `start`, `nodes`.\n\
|
|
||||||
Each node requires `type` plus that type's fields:\n\
|
|
||||||
- agent: `agent`, `prompt`\n\
|
|
||||||
- script: `script`\n\
|
|
||||||
- approval: `question`, `options`, `routes`, `on_other`\n\
|
|
||||||
- input: `question`\n\
|
|
||||||
- llm: `prompt`\n\
|
|
||||||
- rag: `documents`\n\
|
|
||||||
- end: (no required fields)"
|
|
||||||
} else if msg.contains("unknown field") || msg.contains("unknown variant") {
|
|
||||||
"\n\nHint: Check for typos in field names or `type:` values.\n\
|
|
||||||
Valid node types: agent, script, approval, input, llm, rag, end."
|
|
||||||
} else if msg.contains("invalid type") {
|
|
||||||
"\n\nHint: Check that field values have the correct type.\n\
|
|
||||||
- Strings should be quoted if they contain special characters\n\
|
|
||||||
- Numbers should not be quoted\n\
|
|
||||||
- Lists use YAML array syntax (- item1)\n\
|
|
||||||
- Maps use YAML object syntax (key: value)"
|
|
||||||
} else {
|
|
||||||
""
|
|
||||||
};
|
|
||||||
|
|
||||||
anyhow!("YAML parsing error: {}{}", msg, hint)
|
|
||||||
}
|
|
||||||
|
|
||||||
pub fn agent_has_graph(agent_name: &str) -> bool {
|
|
||||||
paths::agent_graph_file(agent_name).exists()
|
|
||||||
}
|
|
||||||
|
|
||||||
#[cfg(test)]
|
|
||||||
mod tests {
|
|
||||||
use super::super::GRAPH_SCHEMA_VERSION;
|
|
||||||
use super::super::types::NodeType;
|
|
||||||
use super::*;
|
|
||||||
use indoc::formatdoc;
|
|
||||||
use std::fs::File;
|
|
||||||
use std::io::Write;
|
|
||||||
use std::{env, fs, process};
|
|
||||||
|
|
||||||
fn parser() -> GraphParser {
|
|
||||||
GraphParser::new(env::current_dir().unwrap())
|
|
||||||
}
|
|
||||||
|
|
||||||
#[test]
|
|
||||||
fn parses_a_simple_graph() {
|
|
||||||
let yaml = formatdoc! {r#"
|
|
||||||
name: simple_graph
|
|
||||||
version: "1.0"
|
|
||||||
start: node1
|
|
||||||
nodes:
|
|
||||||
node1:
|
|
||||||
id: node1
|
|
||||||
type: agent
|
|
||||||
agent: test_agent
|
|
||||||
prompt: "Hello world"
|
|
||||||
next: node2
|
|
||||||
node2:
|
|
||||||
id: node2
|
|
||||||
type: end
|
|
||||||
output: done
|
|
||||||
"#};
|
|
||||||
|
|
||||||
let graph = parser().load_from_string(&yaml).unwrap();
|
|
||||||
|
|
||||||
assert_eq!(graph.name, "simple_graph");
|
|
||||||
assert_eq!(graph.start, "node1");
|
|
||||||
assert_eq!(graph.nodes.len(), 2);
|
|
||||||
assert_eq!(
|
|
||||||
graph.nodes.get("node1").unwrap().next_target(),
|
|
||||||
Some("node2")
|
|
||||||
);
|
|
||||||
}
|
|
||||||
|
|
||||||
#[test]
|
|
||||||
fn auto_fills_node_ids_from_keys() {
|
|
||||||
let yaml = formatdoc! {r#"
|
|
||||||
name: auto_id_graph
|
|
||||||
version: "1.0"
|
|
||||||
start: node1
|
|
||||||
nodes:
|
|
||||||
node1:
|
|
||||||
type: agent
|
|
||||||
agent: test_agent
|
|
||||||
prompt: Test
|
|
||||||
next: node2
|
|
||||||
node2:
|
|
||||||
type: end
|
|
||||||
output: done
|
|
||||||
"#};
|
|
||||||
|
|
||||||
let graph = parser().load_from_string(&yaml).unwrap();
|
|
||||||
|
|
||||||
assert_eq!(graph.nodes.get("node1").unwrap().id, "node1");
|
|
||||||
assert_eq!(graph.nodes.get("node2").unwrap().id, "node2");
|
|
||||||
}
|
|
||||||
|
|
||||||
#[test]
|
|
||||||
fn rejects_missing_start_node() {
|
|
||||||
let yaml = formatdoc! {r#"
|
|
||||||
name: bad_graph
|
|
||||||
version: "1.0"
|
|
||||||
start: nonexistent
|
|
||||||
nodes:
|
|
||||||
node1:
|
|
||||||
type: end
|
|
||||||
"#};
|
|
||||||
|
|
||||||
let err = parser().load_from_string(&yaml).unwrap_err().to_string();
|
|
||||||
|
|
||||||
assert!(
|
|
||||||
err.contains("Start node 'nonexistent' not found"),
|
|
||||||
"got: {err}"
|
|
||||||
);
|
|
||||||
}
|
|
||||||
|
|
||||||
#[test]
|
|
||||||
fn rejects_empty_graph_name() {
|
|
||||||
let yaml = formatdoc! {r#"
|
|
||||||
name: ""
|
|
||||||
version: "1.0"
|
|
||||||
start: node1
|
|
||||||
nodes:
|
|
||||||
node1:
|
|
||||||
type: end
|
|
||||||
"#};
|
|
||||||
|
|
||||||
let err = parser().load_from_string(&yaml).unwrap_err().to_string();
|
|
||||||
|
|
||||||
assert!(err.contains("non-empty 'name'"), "got: {err}");
|
|
||||||
}
|
|
||||||
|
|
||||||
#[test]
|
|
||||||
fn rejects_no_nodes() {
|
|
||||||
let yaml = formatdoc! {r#"
|
|
||||||
name: empty_graph
|
|
||||||
version: "1.0"
|
|
||||||
start: node1
|
|
||||||
nodes: {}
|
|
||||||
"#, "{}"};
|
|
||||||
|
|
||||||
let err = parser().load_from_string(&yaml).unwrap_err().to_string();
|
|
||||||
|
|
||||||
assert!(err.contains("no nodes defined"), "got: {err}");
|
|
||||||
}
|
|
||||||
|
|
||||||
#[test]
|
|
||||||
fn rejects_unsupported_version() {
|
|
||||||
let yaml = formatdoc! {r#"
|
|
||||||
name: future_graph
|
|
||||||
version: "2.0"
|
|
||||||
start: node1
|
|
||||||
nodes:
|
|
||||||
node1:
|
|
||||||
type: end
|
|
||||||
"#};
|
|
||||||
|
|
||||||
let err = parser().load_from_string(&yaml).unwrap_err().to_string();
|
|
||||||
|
|
||||||
assert!(
|
|
||||||
err.contains("Unsupported graph schema version"),
|
|
||||||
"got: {err}"
|
|
||||||
);
|
|
||||||
}
|
|
||||||
|
|
||||||
#[test]
|
|
||||||
fn rejects_node_id_mismatch() {
|
|
||||||
let yaml = formatdoc! {r#"
|
|
||||||
name: mismatch_graph
|
|
||||||
version: "1.0"
|
|
||||||
start: node1
|
|
||||||
nodes:
|
|
||||||
node1:
|
|
||||||
id: different_id
|
|
||||||
type: end
|
|
||||||
"#};
|
|
||||||
|
|
||||||
let err = parser().load_from_string(&yaml).unwrap_err().to_string();
|
|
||||||
|
|
||||||
assert!(err.contains("Node ID mismatch"), "got: {err}");
|
|
||||||
}
|
|
||||||
|
|
||||||
#[test]
|
|
||||||
fn parses_approval_node_with_routes() {
|
|
||||||
let yaml = formatdoc! {r#"
|
|
||||||
name: approval_graph
|
|
||||||
version: "1.0"
|
|
||||||
start: approval1
|
|
||||||
nodes:
|
|
||||||
approval1:
|
|
||||||
type: approval
|
|
||||||
question: "Proceed with deployment?"
|
|
||||||
options:
|
|
||||||
- "Yes"
|
|
||||||
- "No"
|
|
||||||
routes:
|
|
||||||
"Yes": deploy
|
|
||||||
"No": cancel
|
|
||||||
on_other: cancel
|
|
||||||
deploy:
|
|
||||||
type: end
|
|
||||||
cancel:
|
|
||||||
type: end
|
|
||||||
"#};
|
|
||||||
|
|
||||||
let graph = parser().load_from_string(&yaml).unwrap();
|
|
||||||
|
|
||||||
let approval = graph.nodes.get("approval1").unwrap();
|
|
||||||
match &approval.node_type {
|
|
||||||
NodeType::Approval(a) => {
|
|
||||||
assert_eq!(a.options.len(), 2);
|
|
||||||
assert_eq!(a.routes.len(), 2);
|
|
||||||
assert_eq!(a.routes.get("Yes").map(|s| s.as_str()), Some("deploy"));
|
|
||||||
}
|
|
||||||
_ => panic!("expected approval node"),
|
|
||||||
}
|
|
||||||
}
|
|
||||||
|
|
||||||
#[test]
|
|
||||||
fn parses_settings_overrides() {
|
|
||||||
let yaml = formatdoc! {r#"
|
|
||||||
name: settings_graph
|
|
||||||
version: "1.0"
|
|
||||||
start: node1
|
|
||||||
settings:
|
|
||||||
max_loop_iterations: 50
|
|
||||||
timeout: 300
|
|
||||||
log_state_snapshots: false
|
|
||||||
nodes:
|
|
||||||
node1:
|
|
||||||
type: end
|
|
||||||
"#};
|
|
||||||
|
|
||||||
let graph = parser().load_from_string(&yaml).unwrap();
|
|
||||||
|
|
||||||
assert_eq!(graph.settings.max_loop_iterations, 50);
|
|
||||||
assert_eq!(graph.settings.timeout, Some(300));
|
|
||||||
assert!(!graph.settings.log_state_snapshots);
|
|
||||||
assert!(graph.settings.validate_before_run);
|
|
||||||
}
|
|
||||||
|
|
||||||
#[test]
|
|
||||||
fn parses_initial_state() {
|
|
||||||
let yaml = formatdoc! {r#"
|
|
||||||
name: state_graph
|
|
||||||
version: "1.0"
|
|
||||||
start: node1
|
|
||||||
initial_state:
|
|
||||||
user_name: "Alice"
|
|
||||||
count: 42
|
|
||||||
enabled: true
|
|
||||||
nodes:
|
|
||||||
node1:
|
|
||||||
type: end
|
|
||||||
"#};
|
|
||||||
|
|
||||||
let graph = parser().load_from_string(&yaml).unwrap();
|
|
||||||
|
|
||||||
assert_eq!(graph.initial_state.len(), 3);
|
|
||||||
assert_eq!(graph.initial_state.get("user_name").unwrap(), "Alice");
|
|
||||||
assert_eq!(
|
|
||||||
graph.initial_state.get("count").unwrap(),
|
|
||||||
&serde_json::json!(42)
|
|
||||||
);
|
|
||||||
assert_eq!(
|
|
||||||
graph.initial_state.get("enabled").unwrap(),
|
|
||||||
&serde_json::json!(true)
|
|
||||||
);
|
|
||||||
}
|
|
||||||
|
|
||||||
#[test]
|
|
||||||
fn uses_default_version_when_absent() {
|
|
||||||
let yaml = formatdoc! {r#"
|
|
||||||
name: no_version
|
|
||||||
start: node1
|
|
||||||
nodes:
|
|
||||||
node1:
|
|
||||||
type: end
|
|
||||||
"#};
|
|
||||||
|
|
||||||
let graph = parser().load_from_string(&yaml).unwrap();
|
|
||||||
|
|
||||||
assert_eq!(graph.version, GRAPH_SCHEMA_VERSION);
|
|
||||||
}
|
|
||||||
|
|
||||||
#[test]
|
|
||||||
fn rejects_unknown_node_type_with_hint() {
|
|
||||||
let yaml = formatdoc! {r#"
|
|
||||||
name: bad_type
|
|
||||||
version: "1.0"
|
|
||||||
start: node1
|
|
||||||
nodes:
|
|
||||||
node1:
|
|
||||||
type: nonsense
|
|
||||||
"#};
|
|
||||||
|
|
||||||
let err = parser().load_from_string(&yaml).unwrap_err().to_string();
|
|
||||||
|
|
||||||
assert!(
|
|
||||||
err.contains("Valid node types") || err.contains("unknown variant"),
|
|
||||||
"got: {err}"
|
|
||||||
);
|
|
||||||
}
|
|
||||||
|
|
||||||
#[test]
|
|
||||||
fn rejects_malformed_yaml() {
|
|
||||||
let yaml = "name: bad\n bad: indent\nstart: a";
|
|
||||||
|
|
||||||
let result = parser().load_from_string(yaml);
|
|
||||||
|
|
||||||
assert!(result.is_err());
|
|
||||||
}
|
|
||||||
|
|
||||||
#[test]
|
|
||||||
fn missing_required_fields_have_a_hint() {
|
|
||||||
let yaml = formatdoc! {r#"
|
|
||||||
name: missing_start
|
|
||||||
version: "1.0"
|
|
||||||
nodes:
|
|
||||||
node1:
|
|
||||||
type: end
|
|
||||||
"#};
|
|
||||||
|
|
||||||
let err = parser().load_from_string(&yaml).unwrap_err().to_string();
|
|
||||||
|
|
||||||
assert!(err.contains("Hint"), "got: {err}");
|
|
||||||
}
|
|
||||||
|
|
||||||
#[test]
|
|
||||||
fn load_from_file_reads_disk() {
|
|
||||||
let dir = env::temp_dir();
|
|
||||||
let path = dir.join(format!("loki_graph_parser_test_{}.yaml", process::id()));
|
|
||||||
let yaml = formatdoc! {r#"
|
|
||||||
name: disk_graph
|
|
||||||
version: "1.0"
|
|
||||||
start: only
|
|
||||||
nodes:
|
|
||||||
only:
|
|
||||||
type: end
|
|
||||||
output: ok
|
|
||||||
"#};
|
|
||||||
{
|
|
||||||
let mut f = File::create(&path).unwrap();
|
|
||||||
f.write_all(yaml.as_bytes()).unwrap();
|
|
||||||
}
|
|
||||||
|
|
||||||
let graph = GraphParser::new(dir).load_from_file(&path).unwrap();
|
|
||||||
|
|
||||||
assert_eq!(graph.name, "disk_graph");
|
|
||||||
|
|
||||||
let _ = fs::remove_file(&path);
|
|
||||||
}
|
|
||||||
|
|
||||||
#[test]
|
|
||||||
fn load_from_file_errors_on_missing_path() {
|
|
||||||
let err = parser()
|
|
||||||
.load_from_file("/definitely/not/a/real/path/to_any_graph.yaml")
|
|
||||||
.unwrap_err()
|
|
||||||
.to_string();
|
|
||||||
|
|
||||||
assert!(err.contains("Failed to read graph file"), "got: {err}");
|
|
||||||
}
|
|
||||||
|
|
||||||
#[test]
|
|
||||||
fn agent_has_graph_false_for_unknown_agent() {
|
|
||||||
assert!(!agent_has_graph("__nonexistent_agent_for_test__"));
|
|
||||||
}
|
|
||||||
}
|
|
||||||
@@ -1,129 +0,0 @@
|
|||||||
use super::state::StateManager;
|
|
||||||
use super::types::RagNode;
|
|
||||||
use crate::config::RequestContext;
|
|
||||||
use crate::utils::create_abort_signal;
|
|
||||||
use anyhow::{Context, Result, anyhow};
|
|
||||||
use serde_json::{Map, Value};
|
|
||||||
use std::time::Duration;
|
|
||||||
use tokio::time::timeout;
|
|
||||||
|
|
||||||
const OUTPUT_KEY: &str = "output";
|
|
||||||
const DEFAULT_QUERY: &str = "{{initial_prompt}}";
|
|
||||||
const DEFAULT_RAG_TIMEOUT_SECS: u64 = 120;
|
|
||||||
|
|
||||||
pub struct RagNodeExecutor;
|
|
||||||
|
|
||||||
impl RagNodeExecutor {
|
|
||||||
pub(super) async fn execute(
|
|
||||||
node: &RagNode,
|
|
||||||
node_id: &str,
|
|
||||||
state_manager: &mut StateManager,
|
|
||||||
ctx: &mut RequestContext,
|
|
||||||
) -> Result<()> {
|
|
||||||
let query_template = node.query.as_deref().unwrap_or(DEFAULT_QUERY);
|
|
||||||
let query = state_manager
|
|
||||||
.interpolate(query_template)
|
|
||||||
.context("Failed to interpolate rag node query")?;
|
|
||||||
|
|
||||||
let rag = ctx
|
|
||||||
.agent
|
|
||||||
.as_ref()
|
|
||||||
.and_then(|a| a.graph_rag(node_id))
|
|
||||||
.ok_or_else(|| anyhow!("rag node '{node_id}' has no initialized knowledge base"))?;
|
|
||||||
|
|
||||||
let top_k = node.top_k.unwrap_or_else(|| rag.configured_top_k());
|
|
||||||
let rerank = rag.configured_reranker();
|
|
||||||
|
|
||||||
let timeout_dur = Duration::from_secs(node.timeout.unwrap_or(DEFAULT_RAG_TIMEOUT_SECS));
|
|
||||||
let abort = create_abort_signal();
|
|
||||||
let (context, sources_str, _ids) =
|
|
||||||
timeout(timeout_dur, rag.search(&query, top_k, rerank, abort))
|
|
||||||
.await
|
|
||||||
.with_context(|| {
|
|
||||||
format!(
|
|
||||||
"rag node '{node_id}' timed out after {}s",
|
|
||||||
timeout_dur.as_secs()
|
|
||||||
)
|
|
||||||
})?
|
|
||||||
.with_context(|| format!("rag node '{node_id}' retrieval failed"))?;
|
|
||||||
|
|
||||||
let output = build_rag_output(context, &sources_str);
|
|
||||||
apply_state_updates(node, state_manager, &output);
|
|
||||||
Ok(())
|
|
||||||
}
|
|
||||||
}
|
|
||||||
|
|
||||||
/// Assemble the `{{output}}` value as `{ "context": <ctx>, "sources": [...] }`.
|
|
||||||
fn build_rag_output(context: String, sources_str: &str) -> Value {
|
|
||||||
let sources: Vec<Value> = sources_str
|
|
||||||
.lines()
|
|
||||||
.map(|line| line.trim().trim_start_matches("- ").trim())
|
|
||||||
.filter(|s| !s.is_empty())
|
|
||||||
.map(|s| Value::String(s.to_string()))
|
|
||||||
.collect();
|
|
||||||
let mut obj = Map::new();
|
|
||||||
|
|
||||||
obj.insert("context".into(), Value::String(context));
|
|
||||||
obj.insert("sources".into(), Value::Array(sources));
|
|
||||||
|
|
||||||
Value::Object(obj)
|
|
||||||
}
|
|
||||||
|
|
||||||
fn apply_state_updates(node: &RagNode, state_manager: &mut StateManager, output: &Value) {
|
|
||||||
let Some(updates) = &node.state_updates else {
|
|
||||||
return;
|
|
||||||
};
|
|
||||||
let prev_output = state_manager.state().get(OUTPUT_KEY).cloned();
|
|
||||||
state_manager
|
|
||||||
.state_mut()
|
|
||||||
.set(OUTPUT_KEY.into(), output.clone());
|
|
||||||
|
|
||||||
for (key, template) in updates {
|
|
||||||
let value = state_manager.interpolate_lenient(template);
|
|
||||||
state_manager
|
|
||||||
.state_mut()
|
|
||||||
.set(key.clone(), Value::String(value));
|
|
||||||
}
|
|
||||||
|
|
||||||
match prev_output {
|
|
||||||
Some(v) => state_manager.state_mut().set(OUTPUT_KEY.into(), v),
|
|
||||||
None => state_manager
|
|
||||||
.state_mut()
|
|
||||||
.set(OUTPUT_KEY.into(), Value::Null),
|
|
||||||
}
|
|
||||||
}
|
|
||||||
|
|
||||||
#[cfg(test)]
|
|
||||||
mod tests {
|
|
||||||
use super::*;
|
|
||||||
use serde_json::json;
|
|
||||||
|
|
||||||
#[test]
|
|
||||||
fn build_rag_output_splits_bullet_sources_into_array() {
|
|
||||||
let out = build_rag_output("ctx".into(), "- a.md\n- https://x.com/spec");
|
|
||||||
|
|
||||||
assert_eq!(out["context"], json!("ctx"));
|
|
||||||
assert_eq!(out["sources"], json!(["a.md", "https://x.com/spec"]));
|
|
||||||
}
|
|
||||||
|
|
||||||
#[test]
|
|
||||||
fn build_rag_output_handles_empty_sources() {
|
|
||||||
let out = build_rag_output("ctx".into(), "");
|
|
||||||
|
|
||||||
assert_eq!(out["sources"], json!([]));
|
|
||||||
}
|
|
||||||
|
|
||||||
#[test]
|
|
||||||
fn build_rag_output_ignores_blank_lines() {
|
|
||||||
let out = build_rag_output("c".into(), "- a\n\n- b\n");
|
|
||||||
|
|
||||||
assert_eq!(out["sources"], json!(["a", "b"]));
|
|
||||||
}
|
|
||||||
|
|
||||||
#[test]
|
|
||||||
fn build_rag_output_tolerates_unprefixed_lines() {
|
|
||||||
let out = build_rag_output("c".into(), "plain/path");
|
|
||||||
|
|
||||||
assert_eq!(out["sources"], json!(["plain/path"]));
|
|
||||||
}
|
|
||||||
}
|
|
||||||
@@ -1,389 +0,0 @@
|
|||||||
use super::types::Reducer;
|
|
||||||
use crate::graph::type_name;
|
|
||||||
use anyhow::{Result, bail};
|
|
||||||
use serde_json::{Number, Value};
|
|
||||||
|
|
||||||
pub fn apply(reducer: Reducer, current: Option<&Value>, incoming: Value) -> Result<Value> {
|
|
||||||
match reducer {
|
|
||||||
Reducer::Append => apply_append(current, incoming),
|
|
||||||
Reducer::Extend => apply_extend(current, incoming),
|
|
||||||
Reducer::Concat => apply_concat(current, incoming),
|
|
||||||
Reducer::Sum => apply_sum(current, incoming),
|
|
||||||
Reducer::Max => apply_max(current, incoming),
|
|
||||||
Reducer::Min => apply_min(current, incoming),
|
|
||||||
Reducer::Merge => apply_merge(current, incoming),
|
|
||||||
Reducer::Overwrite => Ok(incoming),
|
|
||||||
}
|
|
||||||
}
|
|
||||||
|
|
||||||
fn apply_append(current: Option<&Value>, incoming: Value) -> Result<Value> {
|
|
||||||
let mut arr = match current {
|
|
||||||
None => Vec::new(),
|
|
||||||
Some(Value::Array(a)) => a.clone(),
|
|
||||||
Some(other) => bail!(
|
|
||||||
"reducer 'append' requires an array (or absent) for the current value, got {}",
|
|
||||||
type_name(other)
|
|
||||||
),
|
|
||||||
};
|
|
||||||
arr.push(incoming);
|
|
||||||
|
|
||||||
Ok(Value::Array(arr))
|
|
||||||
}
|
|
||||||
|
|
||||||
fn apply_extend(current: Option<&Value>, incoming: Value) -> Result<Value> {
|
|
||||||
let mut arr = match current {
|
|
||||||
None => Vec::new(),
|
|
||||||
Some(Value::Array(a)) => a.clone(),
|
|
||||||
Some(other) => bail!(
|
|
||||||
"reducer 'extend' requires an array (or absent) for the current value, got {}",
|
|
||||||
type_name(other)
|
|
||||||
),
|
|
||||||
};
|
|
||||||
match incoming {
|
|
||||||
Value::Array(items) => arr.extend(items),
|
|
||||||
other => bail!(
|
|
||||||
"reducer 'extend' requires an array for the incoming value, got {}",
|
|
||||||
type_name(&other)
|
|
||||||
),
|
|
||||||
}
|
|
||||||
|
|
||||||
Ok(Value::Array(arr))
|
|
||||||
}
|
|
||||||
|
|
||||||
fn apply_concat(current: Option<&Value>, incoming: Value) -> Result<Value> {
|
|
||||||
let incoming_str = match incoming {
|
|
||||||
Value::String(s) => s,
|
|
||||||
other => bail!(
|
|
||||||
"reducer 'concat' requires a string for the incoming value, got {}",
|
|
||||||
type_name(&other)
|
|
||||||
),
|
|
||||||
};
|
|
||||||
let result = match current {
|
|
||||||
None => incoming_str,
|
|
||||||
Some(Value::String(c)) => {
|
|
||||||
if c.is_empty() {
|
|
||||||
incoming_str
|
|
||||||
} else {
|
|
||||||
format!("{c}\n{incoming_str}")
|
|
||||||
}
|
|
||||||
}
|
|
||||||
Some(other) => bail!(
|
|
||||||
"reducer 'concat' requires a string (or absent) for the current value, got {}",
|
|
||||||
type_name(other)
|
|
||||||
),
|
|
||||||
};
|
|
||||||
|
|
||||||
Ok(Value::String(result))
|
|
||||||
}
|
|
||||||
|
|
||||||
fn apply_sum(current: Option<&Value>, incoming: Value) -> Result<Value> {
|
|
||||||
let i = number_or_error(&incoming, "sum", "incoming")?;
|
|
||||||
let c = match current {
|
|
||||||
None => 0.0,
|
|
||||||
Some(value) => number_or_error(value, "sum", "current")?,
|
|
||||||
};
|
|
||||||
|
|
||||||
Ok(json_number(c + i))
|
|
||||||
}
|
|
||||||
|
|
||||||
fn apply_max(current: Option<&Value>, incoming: Value) -> Result<Value> {
|
|
||||||
let i = number_or_error(&incoming, "max", "incoming")?;
|
|
||||||
match current {
|
|
||||||
None => Ok(json_number(i)),
|
|
||||||
Some(value) => {
|
|
||||||
let c = number_or_error(value, "max", "current")?;
|
|
||||||
Ok(json_number(c.max(i)))
|
|
||||||
}
|
|
||||||
}
|
|
||||||
}
|
|
||||||
|
|
||||||
fn apply_min(current: Option<&Value>, incoming: Value) -> Result<Value> {
|
|
||||||
let i = number_or_error(&incoming, "min", "incoming")?;
|
|
||||||
match current {
|
|
||||||
None => Ok(json_number(i)),
|
|
||||||
Some(value) => {
|
|
||||||
let c = number_or_error(value, "min", "current")?;
|
|
||||||
Ok(json_number(c.min(i)))
|
|
||||||
}
|
|
||||||
}
|
|
||||||
}
|
|
||||||
|
|
||||||
fn apply_merge(current: Option<&Value>, incoming: Value) -> Result<Value> {
|
|
||||||
let mut map = match current {
|
|
||||||
None => serde_json::Map::new(),
|
|
||||||
Some(Value::Object(m)) => m.clone(),
|
|
||||||
Some(other) => bail!(
|
|
||||||
"reducer 'merge' requires an object (or absent) for the current value, got {}",
|
|
||||||
type_name(other)
|
|
||||||
),
|
|
||||||
};
|
|
||||||
match incoming {
|
|
||||||
Value::Object(items) => {
|
|
||||||
for (k, v) in items {
|
|
||||||
map.insert(k, v);
|
|
||||||
}
|
|
||||||
}
|
|
||||||
other => bail!(
|
|
||||||
"reducer 'merge' requires an object for the incoming value, got {}",
|
|
||||||
type_name(&other)
|
|
||||||
),
|
|
||||||
}
|
|
||||||
|
|
||||||
Ok(Value::Object(map))
|
|
||||||
}
|
|
||||||
|
|
||||||
fn number_or_error(value: &Value, reducer_name: &str, position: &str) -> Result<f64> {
|
|
||||||
match value.as_f64() {
|
|
||||||
Some(n) => Ok(n),
|
|
||||||
None => bail!(
|
|
||||||
"reducer '{reducer_name}' requires a number for the {position} value, got {}",
|
|
||||||
type_name(value)
|
|
||||||
),
|
|
||||||
}
|
|
||||||
}
|
|
||||||
|
|
||||||
// Numeric reducers compute in f64 for simplicity. Integer typing is preserved when the result is losslessly
|
|
||||||
// representable as i64.
|
|
||||||
fn json_number(n: f64) -> Value {
|
|
||||||
if n.fract() == 0.0 && n.is_finite() && n.abs() <= (i64::MAX as f64) {
|
|
||||||
Value::Number(Number::from(n as i64))
|
|
||||||
} else {
|
|
||||||
match Number::from_f64(n) {
|
|
||||||
Some(num) => Value::Number(num),
|
|
||||||
None => Value::Null,
|
|
||||||
}
|
|
||||||
}
|
|
||||||
}
|
|
||||||
|
|
||||||
#[cfg(test)]
|
|
||||||
mod tests {
|
|
||||||
use super::*;
|
|
||||||
use serde_json::json;
|
|
||||||
|
|
||||||
#[test]
|
|
||||||
fn append_to_absent_creates_single_element_array() {
|
|
||||||
let result = apply(Reducer::Append, None, json!("a")).unwrap();
|
|
||||||
|
|
||||||
assert_eq!(result, json!(["a"]));
|
|
||||||
}
|
|
||||||
|
|
||||||
#[test]
|
|
||||||
fn append_pushes_onto_existing_array() {
|
|
||||||
let current = json!(["a", "b"]);
|
|
||||||
let result = apply(Reducer::Append, Some(¤t), json!("c")).unwrap();
|
|
||||||
|
|
||||||
assert_eq!(result, json!(["a", "b", "c"]));
|
|
||||||
}
|
|
||||||
|
|
||||||
#[test]
|
|
||||||
fn append_errors_when_current_is_not_array() {
|
|
||||||
let current = json!("not an array");
|
|
||||||
|
|
||||||
let err = apply(Reducer::Append, Some(¤t), json!("x"))
|
|
||||||
.unwrap_err()
|
|
||||||
.to_string();
|
|
||||||
|
|
||||||
assert!(err.contains("'append'"), "got: {err}");
|
|
||||||
assert!(err.contains("string"), "got: {err}");
|
|
||||||
}
|
|
||||||
|
|
||||||
#[test]
|
|
||||||
fn extend_concatenates_arrays() {
|
|
||||||
let current = json!([1, 2]);
|
|
||||||
|
|
||||||
let result = apply(Reducer::Extend, Some(¤t), json!([3, 4])).unwrap();
|
|
||||||
|
|
||||||
assert_eq!(result, json!([1, 2, 3, 4]));
|
|
||||||
}
|
|
||||||
|
|
||||||
#[test]
|
|
||||||
fn extend_from_absent_with_array() {
|
|
||||||
let result = apply(Reducer::Extend, None, json!([1, 2])).unwrap();
|
|
||||||
|
|
||||||
assert_eq!(result, json!([1, 2]));
|
|
||||||
}
|
|
||||||
|
|
||||||
#[test]
|
|
||||||
fn extend_errors_when_incoming_is_not_array() {
|
|
||||||
let err = apply(Reducer::Extend, None, json!(42))
|
|
||||||
.unwrap_err()
|
|
||||||
.to_string();
|
|
||||||
|
|
||||||
assert!(err.contains("'extend'"), "got: {err}");
|
|
||||||
assert!(err.contains("number"), "got: {err}");
|
|
||||||
assert!(err.contains("incoming"), "got: {err}");
|
|
||||||
}
|
|
||||||
|
|
||||||
#[test]
|
|
||||||
fn concat_joins_strings_with_newline() {
|
|
||||||
let current = json!("first");
|
|
||||||
|
|
||||||
let result = apply(Reducer::Concat, Some(¤t), json!("second")).unwrap();
|
|
||||||
|
|
||||||
assert_eq!(result, json!("first\nsecond"));
|
|
||||||
}
|
|
||||||
|
|
||||||
#[test]
|
|
||||||
fn concat_from_absent_yields_incoming() {
|
|
||||||
let result = apply(Reducer::Concat, None, json!("hello")).unwrap();
|
|
||||||
|
|
||||||
assert_eq!(result, json!("hello"));
|
|
||||||
}
|
|
||||||
|
|
||||||
#[test]
|
|
||||||
fn concat_skips_separator_when_current_is_empty_string() {
|
|
||||||
let current = json!("");
|
|
||||||
|
|
||||||
let result = apply(Reducer::Concat, Some(¤t), json!("first")).unwrap();
|
|
||||||
|
|
||||||
assert_eq!(result, json!("first"));
|
|
||||||
}
|
|
||||||
|
|
||||||
#[test]
|
|
||||||
fn concat_errors_when_incoming_is_not_string() {
|
|
||||||
let err = apply(Reducer::Concat, None, json!(42))
|
|
||||||
.unwrap_err()
|
|
||||||
.to_string();
|
|
||||||
|
|
||||||
assert!(err.contains("'concat'"), "got: {err}");
|
|
||||||
assert!(err.contains("number"), "got: {err}");
|
|
||||||
}
|
|
||||||
|
|
||||||
#[test]
|
|
||||||
fn sum_adds_numbers() {
|
|
||||||
let current = json!(5);
|
|
||||||
|
|
||||||
let result = apply(Reducer::Sum, Some(¤t), json!(7)).unwrap();
|
|
||||||
|
|
||||||
assert_eq!(result, json!(12));
|
|
||||||
}
|
|
||||||
|
|
||||||
#[test]
|
|
||||||
fn sum_starts_from_zero_when_current_absent() {
|
|
||||||
let result = apply(Reducer::Sum, None, json!(3.5)).unwrap();
|
|
||||||
|
|
||||||
assert_eq!(result, json!(3.5));
|
|
||||||
}
|
|
||||||
|
|
||||||
#[test]
|
|
||||||
fn sum_preserves_integer_type_for_whole_results() {
|
|
||||||
let current = json!(2);
|
|
||||||
|
|
||||||
let result = apply(Reducer::Sum, Some(¤t), json!(3)).unwrap();
|
|
||||||
|
|
||||||
assert!(result.is_i64(), "expected integer, got {result:?}");
|
|
||||||
assert_eq!(result, json!(5));
|
|
||||||
}
|
|
||||||
|
|
||||||
#[test]
|
|
||||||
fn sum_uses_float_when_result_has_fractional() {
|
|
||||||
let current = json!(1.5);
|
|
||||||
let result = apply(Reducer::Sum, Some(¤t), json!(2.25)).unwrap();
|
|
||||||
|
|
||||||
assert_eq!(result, json!(3.75));
|
|
||||||
}
|
|
||||||
|
|
||||||
#[test]
|
|
||||||
fn sum_errors_on_string_incoming() {
|
|
||||||
let err = apply(Reducer::Sum, None, json!("not a number"))
|
|
||||||
.unwrap_err()
|
|
||||||
.to_string();
|
|
||||||
|
|
||||||
assert!(err.contains("'sum'"), "got: {err}");
|
|
||||||
assert!(err.contains("string"), "got: {err}");
|
|
||||||
}
|
|
||||||
|
|
||||||
#[test]
|
|
||||||
fn max_returns_larger_of_two() {
|
|
||||||
let current = json!(5);
|
|
||||||
let result = apply(Reducer::Max, Some(¤t), json!(3)).unwrap();
|
|
||||||
assert_eq!(result, json!(5));
|
|
||||||
|
|
||||||
let result = apply(Reducer::Max, Some(¤t), json!(10)).unwrap();
|
|
||||||
assert_eq!(result, json!(10));
|
|
||||||
}
|
|
||||||
|
|
||||||
#[test]
|
|
||||||
fn max_yields_incoming_when_current_absent() {
|
|
||||||
let result = apply(Reducer::Max, None, json!(42)).unwrap();
|
|
||||||
|
|
||||||
assert_eq!(result, json!(42));
|
|
||||||
}
|
|
||||||
|
|
||||||
#[test]
|
|
||||||
fn min_returns_smaller_of_two() {
|
|
||||||
let current = json!(5);
|
|
||||||
let result = apply(Reducer::Min, Some(¤t), json!(3)).unwrap();
|
|
||||||
assert_eq!(result, json!(3));
|
|
||||||
|
|
||||||
let result = apply(Reducer::Min, Some(¤t), json!(10)).unwrap();
|
|
||||||
assert_eq!(result, json!(5));
|
|
||||||
}
|
|
||||||
|
|
||||||
#[test]
|
|
||||||
fn min_errors_on_non_numeric_current() {
|
|
||||||
let current = json!("oops");
|
|
||||||
|
|
||||||
let err = apply(Reducer::Min, Some(¤t), json!(1))
|
|
||||||
.unwrap_err()
|
|
||||||
.to_string();
|
|
||||||
|
|
||||||
assert!(err.contains("'min'"), "got: {err}");
|
|
||||||
assert!(err.contains("current"), "got: {err}");
|
|
||||||
}
|
|
||||||
|
|
||||||
#[test]
|
|
||||||
fn merge_unions_objects_with_incoming_winning_collisions() {
|
|
||||||
let current = json!({ "a": 1, "b": 2 });
|
|
||||||
let incoming = json!({ "b": 99, "c": 3 });
|
|
||||||
|
|
||||||
let result = apply(Reducer::Merge, Some(¤t), incoming).unwrap();
|
|
||||||
|
|
||||||
assert_eq!(result, json!({ "a": 1, "b": 99, "c": 3 }));
|
|
||||||
}
|
|
||||||
|
|
||||||
#[test]
|
|
||||||
fn merge_from_absent_yields_incoming_object() {
|
|
||||||
let result = apply(Reducer::Merge, None, json!({ "k": "v" })).unwrap();
|
|
||||||
|
|
||||||
assert_eq!(result, json!({ "k": "v" }));
|
|
||||||
}
|
|
||||||
|
|
||||||
#[test]
|
|
||||||
fn merge_errors_when_incoming_is_not_object() {
|
|
||||||
let err = apply(Reducer::Merge, None, json!([1, 2]))
|
|
||||||
.unwrap_err()
|
|
||||||
.to_string();
|
|
||||||
|
|
||||||
assert!(err.contains("'merge'"), "got: {err}");
|
|
||||||
assert!(err.contains("array"), "got: {err}");
|
|
||||||
}
|
|
||||||
|
|
||||||
#[test]
|
|
||||||
fn merge_errors_when_current_is_not_object() {
|
|
||||||
let current = json!("not object");
|
|
||||||
|
|
||||||
let err = apply(Reducer::Merge, Some(¤t), json!({ "k": "v" }))
|
|
||||||
.unwrap_err()
|
|
||||||
.to_string();
|
|
||||||
|
|
||||||
assert!(err.contains("'merge'"), "got: {err}");
|
|
||||||
assert!(err.contains("current"), "got: {err}");
|
|
||||||
}
|
|
||||||
|
|
||||||
#[test]
|
|
||||||
fn overwrite_ignores_current_and_returns_incoming() {
|
|
||||||
let current = json!("old");
|
|
||||||
|
|
||||||
let result = apply(Reducer::Overwrite, Some(¤t), json!("new")).unwrap();
|
|
||||||
|
|
||||||
assert_eq!(result, json!("new"));
|
|
||||||
}
|
|
||||||
|
|
||||||
#[test]
|
|
||||||
fn overwrite_works_with_absent_current() {
|
|
||||||
let result = apply(Reducer::Overwrite, None, json!(42)).unwrap();
|
|
||||||
|
|
||||||
assert_eq!(result, json!(42));
|
|
||||||
}
|
|
||||||
}
|
|
||||||
@@ -1,557 +0,0 @@
|
|||||||
use super::state::{StateManager, StateRepresentation};
|
|
||||||
use super::types::ScriptNode;
|
|
||||||
use crate::config::paths;
|
|
||||||
use crate::function::Language;
|
|
||||||
use anyhow::{Context, Result, anyhow, bail};
|
|
||||||
use serde_json::Value;
|
|
||||||
use std::collections::HashMap;
|
|
||||||
use std::env;
|
|
||||||
use std::path::{Path, PathBuf};
|
|
||||||
use std::process::Stdio;
|
|
||||||
use std::time::Duration;
|
|
||||||
use tokio::process::Command;
|
|
||||||
use tokio::time::timeout;
|
|
||||||
|
|
||||||
#[cfg(windows)]
|
|
||||||
const PATH_SEP: &str = ";";
|
|
||||||
#[cfg(not(windows))]
|
|
||||||
const PATH_SEP: &str = ":";
|
|
||||||
|
|
||||||
#[derive(Clone)]
|
|
||||||
pub struct ScriptExecutor {
|
|
||||||
base_dir: PathBuf,
|
|
||||||
extra_envs: HashMap<String, String>,
|
|
||||||
}
|
|
||||||
|
|
||||||
impl ScriptExecutor {
|
|
||||||
pub fn new(base_dir: impl Into<PathBuf>) -> Self {
|
|
||||||
let base_dir = base_dir.into();
|
|
||||||
let extra_envs = build_default_envs(&base_dir);
|
|
||||||
Self {
|
|
||||||
base_dir,
|
|
||||||
extra_envs,
|
|
||||||
}
|
|
||||||
}
|
|
||||||
|
|
||||||
pub fn with_envs(mut self, envs: HashMap<String, String>) -> Self {
|
|
||||||
self.extra_envs.extend(envs);
|
|
||||||
self
|
|
||||||
}
|
|
||||||
|
|
||||||
pub async fn execute(
|
|
||||||
&self,
|
|
||||||
node: &ScriptNode,
|
|
||||||
state_manager: &mut StateManager,
|
|
||||||
) -> Result<Option<String>> {
|
|
||||||
let script_path = self.base_dir.join(&node.script);
|
|
||||||
if !script_path.exists() {
|
|
||||||
bail!("Script file not found: '{}'", script_path.display());
|
|
||||||
}
|
|
||||||
|
|
||||||
let language = detect_language(&script_path)?;
|
|
||||||
let state_repr = state_manager.serialize_state()?;
|
|
||||||
|
|
||||||
let mut cmd = build_command(language, &script_path)?;
|
|
||||||
cmd.stdout(Stdio::piped());
|
|
||||||
cmd.stderr(Stdio::piped());
|
|
||||||
cmd.envs(&self.extra_envs);
|
|
||||||
match &state_repr {
|
|
||||||
StateRepresentation::Inline(json) => {
|
|
||||||
cmd.env("GRAPH_STATE", json);
|
|
||||||
}
|
|
||||||
StateRepresentation::File(path) => {
|
|
||||||
cmd.env("GRAPH_STATE_FILE", path);
|
|
||||||
}
|
|
||||||
}
|
|
||||||
|
|
||||||
let timeout_dur = Duration::from_secs(node.timeout);
|
|
||||||
let output = timeout(timeout_dur, cmd.output())
|
|
||||||
.await
|
|
||||||
.with_context(|| {
|
|
||||||
format!(
|
|
||||||
"Script '{}' timed out after {}s",
|
|
||||||
script_path.display(),
|
|
||||||
node.timeout
|
|
||||||
)
|
|
||||||
})?
|
|
||||||
.with_context(|| {
|
|
||||||
format!(
|
|
||||||
"Failed to spawn script process for '{}'",
|
|
||||||
script_path.display()
|
|
||||||
)
|
|
||||||
})?;
|
|
||||||
|
|
||||||
if !output.status.success() {
|
|
||||||
let stderr = String::from_utf8_lossy(&output.stderr);
|
|
||||||
bail!(
|
|
||||||
"Script '{}' failed with exit code {:?}:\n{}",
|
|
||||||
script_path.display(),
|
|
||||||
output.status.code(),
|
|
||||||
stderr.trim()
|
|
||||||
);
|
|
||||||
}
|
|
||||||
|
|
||||||
let stdout = String::from_utf8_lossy(&output.stdout);
|
|
||||||
let json_output = stdout.trim();
|
|
||||||
if json_output.is_empty() {
|
|
||||||
bail!(
|
|
||||||
"Script '{}' produced no output (scripts must emit a single JSON object on stdout)",
|
|
||||||
script_path.display()
|
|
||||||
);
|
|
||||||
}
|
|
||||||
|
|
||||||
let next = state_manager
|
|
||||||
.merge_script_output(json_output)
|
|
||||||
.with_context(|| {
|
|
||||||
format!(
|
|
||||||
"Failed to merge output from script '{}'",
|
|
||||||
script_path.display()
|
|
||||||
)
|
|
||||||
})?;
|
|
||||||
|
|
||||||
apply_state_updates(node, state_manager);
|
|
||||||
|
|
||||||
Ok(next)
|
|
||||||
}
|
|
||||||
}
|
|
||||||
|
|
||||||
fn apply_state_updates(node: &ScriptNode, state_manager: &mut StateManager) {
|
|
||||||
let Some(updates) = &node.state_updates else {
|
|
||||||
return;
|
|
||||||
};
|
|
||||||
|
|
||||||
for (key, template) in updates {
|
|
||||||
let value = state_manager.interpolate_lenient(template);
|
|
||||||
state_manager
|
|
||||||
.state_mut()
|
|
||||||
.set(key.clone(), Value::String(value));
|
|
||||||
}
|
|
||||||
}
|
|
||||||
|
|
||||||
fn build_default_envs(agent_data_dir: &Path) -> HashMap<String, String> {
|
|
||||||
let mut envs = HashMap::new();
|
|
||||||
envs.insert(
|
|
||||||
"LLM_ROOT_DIR".to_string(),
|
|
||||||
paths::config_dir().to_string_lossy().into_owned(),
|
|
||||||
);
|
|
||||||
envs.insert(
|
|
||||||
"LLM_PROMPT_UTILS_FILE".to_string(),
|
|
||||||
paths::bash_prompt_utils_file()
|
|
||||||
.to_string_lossy()
|
|
||||||
.into_owned(),
|
|
||||||
);
|
|
||||||
envs.insert(
|
|
||||||
"LLM_AGENT_DATA_DIR".to_string(),
|
|
||||||
agent_data_dir.to_string_lossy().into_owned(),
|
|
||||||
);
|
|
||||||
envs.insert("CLICOLOR_FORCE".to_string(), "1".to_string());
|
|
||||||
envs.insert("FORCE_COLOR".to_string(), "1".to_string());
|
|
||||||
|
|
||||||
if let Ok(current_path) = env::var("PATH") {
|
|
||||||
let bin_dir = paths::functions_bin_dir();
|
|
||||||
envs.insert(
|
|
||||||
"PATH".to_string(),
|
|
||||||
format!("{}{}{}", bin_dir.display(), PATH_SEP, current_path),
|
|
||||||
);
|
|
||||||
}
|
|
||||||
|
|
||||||
envs
|
|
||||||
}
|
|
||||||
|
|
||||||
fn detect_language(script_path: &Path) -> Result<Language> {
|
|
||||||
let ext = script_path
|
|
||||||
.extension()
|
|
||||||
.and_then(|e| e.to_str())
|
|
||||||
.ok_or_else(|| anyhow!("Script has no file extension: '{}'", script_path.display()))?
|
|
||||||
.to_string();
|
|
||||||
|
|
||||||
match Language::from(&ext) {
|
|
||||||
Language::Unsupported => bail!(
|
|
||||||
"Unsupported script extension '.{}' for '{}'",
|
|
||||||
ext,
|
|
||||||
script_path.display()
|
|
||||||
),
|
|
||||||
lang => Ok(lang),
|
|
||||||
}
|
|
||||||
}
|
|
||||||
|
|
||||||
fn build_command(language: Language, script_path: &Path) -> Result<Command> {
|
|
||||||
let (program, prefix_args) = language.direct_invoker().ok_or_else(|| {
|
|
||||||
anyhow!(
|
|
||||||
"No direct invoker available for script '{}'",
|
|
||||||
script_path.display()
|
|
||||||
)
|
|
||||||
})?;
|
|
||||||
let mut cmd = Command::new(program);
|
|
||||||
|
|
||||||
for arg in prefix_args {
|
|
||||||
cmd.arg(arg);
|
|
||||||
}
|
|
||||||
|
|
||||||
cmd.arg(script_path);
|
|
||||||
Ok(cmd)
|
|
||||||
}
|
|
||||||
|
|
||||||
#[cfg(test)]
|
|
||||||
mod tests {
|
|
||||||
use super::super::MAX_STATE_SIZE_BYTES;
|
|
||||||
use super::*;
|
|
||||||
use crate::utils::temp_file;
|
|
||||||
use indoc::formatdoc;
|
|
||||||
use serde_json::json;
|
|
||||||
use std::collections::HashMap;
|
|
||||||
use std::env::temp_dir;
|
|
||||||
use std::fs;
|
|
||||||
|
|
||||||
fn cmd_available(name: &str) -> bool {
|
|
||||||
which::which(name).is_ok()
|
|
||||||
}
|
|
||||||
|
|
||||||
fn write_script(contents: &str, ext: &str) -> (PathBuf, PathBuf) {
|
|
||||||
let dir = temp_file("-graph-script-test-", "");
|
|
||||||
fs::create_dir_all(&dir).unwrap();
|
|
||||||
let path = dir.join(format!("script.{ext}"));
|
|
||||||
fs::write(&path, contents).unwrap();
|
|
||||||
(dir, path)
|
|
||||||
}
|
|
||||||
|
|
||||||
fn cleanup(dir: &Path) {
|
|
||||||
let _ = fs::remove_dir_all(dir);
|
|
||||||
}
|
|
||||||
|
|
||||||
fn node_for(script_filename: &str, timeout: u64) -> ScriptNode {
|
|
||||||
ScriptNode {
|
|
||||||
script: script_filename.into(),
|
|
||||||
state_updates: None,
|
|
||||||
fallback: None,
|
|
||||||
timeout,
|
|
||||||
}
|
|
||||||
}
|
|
||||||
|
|
||||||
#[tokio::test]
|
|
||||||
async fn bash_script_merges_json_output_into_state() {
|
|
||||||
if !cmd_available("bash") {
|
|
||||||
eprintln!("skipping: bash not available");
|
|
||||||
return;
|
|
||||||
}
|
|
||||||
let (dir, path) = write_script(
|
|
||||||
r#"#!/bin/bash
|
|
||||||
echo '{"quality": 0.85, "issues": 3, "_next": "approve"}'
|
|
||||||
"#,
|
|
||||||
"sh",
|
|
||||||
);
|
|
||||||
let mut state = StateManager::new(HashMap::new());
|
|
||||||
let executor = ScriptExecutor::new(&dir);
|
|
||||||
|
|
||||||
let next = executor
|
|
||||||
.execute(
|
|
||||||
&node_for(path.file_name().unwrap().to_str().unwrap(), 5),
|
|
||||||
&mut state,
|
|
||||||
)
|
|
||||||
.await
|
|
||||||
.unwrap();
|
|
||||||
|
|
||||||
assert_eq!(next.as_deref(), Some("approve"));
|
|
||||||
assert_eq!(state.state().get("quality"), Some(&json!(0.85)));
|
|
||||||
assert_eq!(state.state().get("issues"), Some(&json!(3)));
|
|
||||||
assert!(state.state().get("_next").is_none());
|
|
||||||
cleanup(&dir);
|
|
||||||
}
|
|
||||||
|
|
||||||
#[tokio::test]
|
|
||||||
async fn bash_script_can_read_state_from_env() {
|
|
||||||
if !cmd_available("bash") || !cmd_available("python3") {
|
|
||||||
eprintln!("skipping: bash or python3 not available");
|
|
||||||
return;
|
|
||||||
}
|
|
||||||
let (dir, path) = write_script(
|
|
||||||
r#"#!/bin/bash
|
|
||||||
NAME=$(python3 -c 'import json,os; print(json.loads(os.environ["GRAPH_STATE"])["name"])')
|
|
||||||
printf '{"greeting": "hello %s"}' "$NAME"
|
|
||||||
"#,
|
|
||||||
"sh",
|
|
||||||
);
|
|
||||||
let mut initial = HashMap::new();
|
|
||||||
initial.insert("name".into(), json!("alice"));
|
|
||||||
let mut state = StateManager::new(initial);
|
|
||||||
let executor = ScriptExecutor::new(&dir);
|
|
||||||
|
|
||||||
let _ = executor
|
|
||||||
.execute(
|
|
||||||
&node_for(path.file_name().unwrap().to_str().unwrap(), 5),
|
|
||||||
&mut state,
|
|
||||||
)
|
|
||||||
.await
|
|
||||||
.unwrap();
|
|
||||||
|
|
||||||
assert_eq!(state.state().get("greeting"), Some(&json!("hello alice")));
|
|
||||||
cleanup(&dir);
|
|
||||||
}
|
|
||||||
|
|
||||||
#[tokio::test]
|
|
||||||
async fn script_without_next_returns_none() {
|
|
||||||
if !cmd_available("bash") {
|
|
||||||
return;
|
|
||||||
}
|
|
||||||
let (dir, path) = write_script(
|
|
||||||
r#"#!/bin/bash
|
|
||||||
echo '{"ok": true}'
|
|
||||||
"#,
|
|
||||||
"sh",
|
|
||||||
);
|
|
||||||
let mut state = StateManager::new(HashMap::new());
|
|
||||||
let executor = ScriptExecutor::new(&dir);
|
|
||||||
|
|
||||||
let next = executor
|
|
||||||
.execute(
|
|
||||||
&node_for(path.file_name().unwrap().to_str().unwrap(), 5),
|
|
||||||
&mut state,
|
|
||||||
)
|
|
||||||
.await
|
|
||||||
.unwrap();
|
|
||||||
|
|
||||||
assert!(next.is_none());
|
|
||||||
assert_eq!(state.state().get("ok"), Some(&json!(true)));
|
|
||||||
cleanup(&dir);
|
|
||||||
}
|
|
||||||
|
|
||||||
#[tokio::test]
|
|
||||||
async fn state_updates_apply_after_json_merge() {
|
|
||||||
if !cmd_available("bash") {
|
|
||||||
return;
|
|
||||||
}
|
|
||||||
let (dir, path) = write_script(
|
|
||||||
r#"#!/bin/bash
|
|
||||||
echo '{"raw": "hello"}'
|
|
||||||
"#,
|
|
||||||
"sh",
|
|
||||||
);
|
|
||||||
let mut node = node_for(path.file_name().unwrap().to_str().unwrap(), 5);
|
|
||||||
let mut updates = HashMap::new();
|
|
||||||
updates.insert("decorated".into(), "[{{raw}}]".into());
|
|
||||||
node.state_updates = Some(updates);
|
|
||||||
|
|
||||||
let mut state = StateManager::new(HashMap::new());
|
|
||||||
let executor = ScriptExecutor::new(&dir);
|
|
||||||
executor.execute(&node, &mut state).await.unwrap();
|
|
||||||
|
|
||||||
assert_eq!(state.state().get("raw"), Some(&json!("hello")));
|
|
||||||
assert_eq!(state.state().get("decorated"), Some(&json!("[hello]")));
|
|
||||||
cleanup(&dir);
|
|
||||||
}
|
|
||||||
|
|
||||||
#[tokio::test]
|
|
||||||
async fn missing_script_file_errors_before_spawning() {
|
|
||||||
let mut state = StateManager::new(HashMap::new());
|
|
||||||
let executor = ScriptExecutor::new(temp_dir());
|
|
||||||
|
|
||||||
let err = executor
|
|
||||||
.execute(&node_for("__does_not_exist__.sh", 5), &mut state)
|
|
||||||
.await
|
|
||||||
.unwrap_err()
|
|
||||||
.to_string();
|
|
||||||
|
|
||||||
assert!(err.contains("Script file not found"), "got: {err}");
|
|
||||||
}
|
|
||||||
|
|
||||||
#[tokio::test]
|
|
||||||
async fn empty_stdout_errors() {
|
|
||||||
if !cmd_available("bash") {
|
|
||||||
return;
|
|
||||||
}
|
|
||||||
let (dir, path) = write_script("#!/bin/bash\n", "sh");
|
|
||||||
let mut state = StateManager::new(HashMap::new());
|
|
||||||
let executor = ScriptExecutor::new(&dir);
|
|
||||||
|
|
||||||
let err = executor
|
|
||||||
.execute(
|
|
||||||
&node_for(path.file_name().unwrap().to_str().unwrap(), 5),
|
|
||||||
&mut state,
|
|
||||||
)
|
|
||||||
.await
|
|
||||||
.unwrap_err()
|
|
||||||
.to_string();
|
|
||||||
|
|
||||||
assert!(err.contains("produced no output"), "got: {err}");
|
|
||||||
cleanup(&dir);
|
|
||||||
}
|
|
||||||
|
|
||||||
#[tokio::test]
|
|
||||||
async fn non_json_output_errors() {
|
|
||||||
if !cmd_available("bash") {
|
|
||||||
return;
|
|
||||||
}
|
|
||||||
let (dir, path) = write_script(
|
|
||||||
&formatdoc! {r#"
|
|
||||||
#!/bin/bash
|
|
||||||
echo "not json at all"
|
|
||||||
"#},
|
|
||||||
"sh",
|
|
||||||
);
|
|
||||||
let mut state = StateManager::new(HashMap::new());
|
|
||||||
let executor = ScriptExecutor::new(&dir);
|
|
||||||
|
|
||||||
let err = executor
|
|
||||||
.execute(
|
|
||||||
&node_for(path.file_name().unwrap().to_str().unwrap(), 5),
|
|
||||||
&mut state,
|
|
||||||
)
|
|
||||||
.await
|
|
||||||
.unwrap_err()
|
|
||||||
.to_string();
|
|
||||||
|
|
||||||
assert!(err.contains("merge output"), "got: {err}");
|
|
||||||
cleanup(&dir);
|
|
||||||
}
|
|
||||||
|
|
||||||
#[tokio::test]
|
|
||||||
async fn non_zero_exit_errors_and_includes_stderr() {
|
|
||||||
if !cmd_available("bash") {
|
|
||||||
return;
|
|
||||||
}
|
|
||||||
let (dir, path) = write_script(
|
|
||||||
&formatdoc! {r#"
|
|
||||||
#!/bin/bash
|
|
||||||
echo "bad happened" >&2
|
|
||||||
exit 7
|
|
||||||
"#},
|
|
||||||
"sh",
|
|
||||||
);
|
|
||||||
let mut state = StateManager::new(HashMap::new());
|
|
||||||
let executor = ScriptExecutor::new(&dir);
|
|
||||||
|
|
||||||
let err = executor
|
|
||||||
.execute(
|
|
||||||
&node_for(path.file_name().unwrap().to_str().unwrap(), 5),
|
|
||||||
&mut state,
|
|
||||||
)
|
|
||||||
.await
|
|
||||||
.unwrap_err()
|
|
||||||
.to_string();
|
|
||||||
|
|
||||||
assert!(err.contains("exit code"), "got: {err}");
|
|
||||||
assert!(err.contains("bad happened"), "got: {err}");
|
|
||||||
cleanup(&dir);
|
|
||||||
}
|
|
||||||
|
|
||||||
#[tokio::test]
|
|
||||||
async fn execution_timeout_is_enforced() {
|
|
||||||
if !cmd_available("bash") {
|
|
||||||
return;
|
|
||||||
}
|
|
||||||
let (dir, path) = write_script(
|
|
||||||
r#"#!/bin/bash
|
|
||||||
sleep 5
|
|
||||||
echo '{"ok":true}'
|
|
||||||
"#,
|
|
||||||
"sh",
|
|
||||||
);
|
|
||||||
let mut state = StateManager::new(HashMap::new());
|
|
||||||
let executor = ScriptExecutor::new(&dir);
|
|
||||||
|
|
||||||
let err = executor
|
|
||||||
.execute(
|
|
||||||
&node_for(path.file_name().unwrap().to_str().unwrap(), 1),
|
|
||||||
&mut state,
|
|
||||||
)
|
|
||||||
.await
|
|
||||||
.unwrap_err()
|
|
||||||
.to_string();
|
|
||||||
|
|
||||||
assert!(err.contains("timed out"), "got: {err}");
|
|
||||||
cleanup(&dir);
|
|
||||||
}
|
|
||||||
|
|
||||||
#[tokio::test]
|
|
||||||
async fn large_state_is_delivered_via_file_env_var() {
|
|
||||||
if !cmd_available("bash") || !cmd_available("python3") {
|
|
||||||
return;
|
|
||||||
}
|
|
||||||
let big = "x".repeat(MAX_STATE_SIZE_BYTES + 1024);
|
|
||||||
let mut initial = HashMap::new();
|
|
||||||
initial.insert("blob".into(), json!(big));
|
|
||||||
|
|
||||||
let (dir, path) = write_script(
|
|
||||||
r#"#!/bin/bash
|
|
||||||
if [ -n "$GRAPH_STATE_FILE" ]; then
|
|
||||||
LEN=$(python3 -c 'import json,os; print(len(json.load(open(os.environ["GRAPH_STATE_FILE"]))["blob"]))')
|
|
||||||
printf '{"blob_len": %s, "via_file": true}' "$LEN"
|
|
||||||
elif [ -n "$GRAPH_STATE" ]; then
|
|
||||||
echo '{"via_file": false}'
|
|
||||||
fi
|
|
||||||
"#,
|
|
||||||
"sh",
|
|
||||||
);
|
|
||||||
|
|
||||||
let mut state = StateManager::new(initial);
|
|
||||||
let executor = ScriptExecutor::new(&dir);
|
|
||||||
executor
|
|
||||||
.execute(
|
|
||||||
&node_for(path.file_name().unwrap().to_str().unwrap(), 10),
|
|
||||||
&mut state,
|
|
||||||
)
|
|
||||||
.await
|
|
||||||
.unwrap();
|
|
||||||
|
|
||||||
assert_eq!(state.state().get("via_file"), Some(&json!(true)));
|
|
||||||
let len = state.state().get("blob_len").unwrap().as_i64().unwrap();
|
|
||||||
assert_eq!(len as usize, big.len());
|
|
||||||
cleanup(&dir);
|
|
||||||
}
|
|
||||||
|
|
||||||
#[tokio::test]
|
|
||||||
async fn python_script_can_emit_routing_and_state() {
|
|
||||||
if !cmd_available("python3") {
|
|
||||||
eprintln!("skipping: python3 not available");
|
|
||||||
return;
|
|
||||||
}
|
|
||||||
let (dir, path) = write_script(
|
|
||||||
r#"import os, json
|
|
||||||
state = json.loads(os.environ["GRAPH_STATE"])
|
|
||||||
print(json.dumps({
|
|
||||||
"_next": "next_node",
|
|
||||||
"doubled": state.get("n", 0) * 2,
|
|
||||||
}))
|
|
||||||
"#,
|
|
||||||
"py",
|
|
||||||
);
|
|
||||||
let mut initial = HashMap::new();
|
|
||||||
initial.insert("n".into(), json!(21));
|
|
||||||
let mut state = StateManager::new(initial);
|
|
||||||
|
|
||||||
let executor = ScriptExecutor::new(&dir);
|
|
||||||
let next = executor
|
|
||||||
.execute(
|
|
||||||
&node_for(path.file_name().unwrap().to_str().unwrap(), 5),
|
|
||||||
&mut state,
|
|
||||||
)
|
|
||||||
.await
|
|
||||||
.unwrap();
|
|
||||||
|
|
||||||
assert_eq!(next.as_deref(), Some("next_node"));
|
|
||||||
assert_eq!(state.state().get("doubled"), Some(&json!(42)));
|
|
||||||
cleanup(&dir);
|
|
||||||
}
|
|
||||||
|
|
||||||
#[tokio::test]
|
|
||||||
async fn unknown_extension_is_rejected() {
|
|
||||||
let (dir, path) = write_script("echo hi", "xyz");
|
|
||||||
let mut state = StateManager::new(HashMap::new());
|
|
||||||
let executor = ScriptExecutor::new(&dir);
|
|
||||||
|
|
||||||
let err = executor
|
|
||||||
.execute(
|
|
||||||
&node_for(path.file_name().unwrap().to_str().unwrap(), 5),
|
|
||||||
&mut state,
|
|
||||||
)
|
|
||||||
.await
|
|
||||||
.unwrap_err()
|
|
||||||
.to_string();
|
|
||||||
|
|
||||||
assert!(
|
|
||||||
err.contains("Unsupported script extension '.xyz'"),
|
|
||||||
"got: {err}"
|
|
||||||
);
|
|
||||||
cleanup(&dir);
|
|
||||||
}
|
|
||||||
}
|
|
||||||
@@ -1,9 +0,0 @@
|
|||||||
use serde_json::Value;
|
|
||||||
use std::collections::HashMap;
|
|
||||||
|
|
||||||
#[derive(Debug, Clone)]
|
|
||||||
pub struct BranchWrites {
|
|
||||||
pub node_id: String,
|
|
||||||
pub invocation_index: usize,
|
|
||||||
pub writes: HashMap<String, Value>,
|
|
||||||
}
|
|
||||||
-1021
File diff suppressed because it is too large
Load Diff
@@ -1,189 +0,0 @@
|
|||||||
use crate::client::call_chat_completions;
|
|
||||||
use crate::config::{Input, RequestContext, Role, RoleLike};
|
|
||||||
use crate::utils::create_abort_signal;
|
|
||||||
use anyhow::{Context, Result, bail};
|
|
||||||
use serde_json::Value;
|
|
||||||
use std::sync::Arc;
|
|
||||||
|
|
||||||
const EXTRACTOR_ROLE_NAME: &str = "__structured_output__";
|
|
||||||
|
|
||||||
const EXTRACTOR_ROLE_PROMPT: &str = "\
|
|
||||||
Extract a JSON object from the user's input that strictly conforms to the provided JSON Schema.
|
|
||||||
|
|
||||||
Rules:
|
|
||||||
- Output ONLY the JSON object. No prose, no explanation, no markdown fences, no <think> tokens.
|
|
||||||
- The first character of your response must be `{` and the last must be `}`.
|
|
||||||
- Every key marked `required` in the schema MUST appear in the output.
|
|
||||||
- All values MUST match the types specified in the schema.
|
|
||||||
- If the input is already a valid JSON object matching the schema, return it unchanged.
|
|
||||||
- If a field cannot be determined from the input, use `null` (when allowed) or your best inferred value.
|
|
||||||
- Do NOT invent fields not present in the schema.";
|
|
||||||
|
|
||||||
pub async fn extract(raw: &str, schema: &Value, parent_ctx: &mut RequestContext) -> Result<Value> {
|
|
||||||
if let Some(parsed) = try_parse_json(raw) {
|
|
||||||
return Ok(parsed);
|
|
||||||
}
|
|
||||||
|
|
||||||
extract_via_extractor(raw, schema, parent_ctx, false).await
|
|
||||||
}
|
|
||||||
|
|
||||||
async fn extract_via_extractor(
|
|
||||||
raw: &str,
|
|
||||||
schema: &Value,
|
|
||||||
parent_ctx: &mut RequestContext,
|
|
||||||
is_repair: bool,
|
|
||||||
) -> Result<Value> {
|
|
||||||
let role = build_extractor_role()?;
|
|
||||||
let prompt = build_extractor_prompt(raw, schema, is_repair);
|
|
||||||
|
|
||||||
let saved_role = parent_ctx.role.clone();
|
|
||||||
parent_ctx.role = Some(role);
|
|
||||||
let result = run_one_shot(&prompt, parent_ctx).await;
|
|
||||||
parent_ctx.role = saved_role;
|
|
||||||
|
|
||||||
let output = result.context("Structured-output extractor LLM call failed")?;
|
|
||||||
|
|
||||||
match try_parse_json(&output) {
|
|
||||||
Some(value) => Ok(value),
|
|
||||||
None if is_repair => bail!(
|
|
||||||
"Structured-output extractor failed to produce valid JSON after repair retry. \
|
|
||||||
Last response:\n{output}"
|
|
||||||
),
|
|
||||||
None => Box::pin(extract_via_extractor(&output, schema, parent_ctx, true)).await,
|
|
||||||
}
|
|
||||||
}
|
|
||||||
|
|
||||||
fn build_extractor_role() -> Result<Role> {
|
|
||||||
let mut role = Role::new(EXTRACTOR_ROLE_NAME, EXTRACTOR_ROLE_PROMPT);
|
|
||||||
role.set_enabled_tools(Some(String::new()));
|
|
||||||
role.set_enabled_mcp_servers(Some(String::new()));
|
|
||||||
Ok(role)
|
|
||||||
}
|
|
||||||
|
|
||||||
fn build_extractor_prompt(raw: &str, schema: &Value, is_repair: bool) -> String {
|
|
||||||
let schema_json = serde_json::to_string_pretty(schema).unwrap_or_else(|_| schema.to_string());
|
|
||||||
if is_repair {
|
|
||||||
format!(
|
|
||||||
"Your previous response was not valid JSON. Output ONLY a JSON object \
|
|
||||||
matching this schema. No prose, no fences.\n\nSchema:\n{schema_json}\n\nInput:\n{raw}"
|
|
||||||
)
|
|
||||||
} else {
|
|
||||||
format!("Schema:\n{schema_json}\n\nInput:\n{raw}")
|
|
||||||
}
|
|
||||||
}
|
|
||||||
|
|
||||||
async fn run_one_shot(prompt: &str, ctx: &mut RequestContext) -> Result<String> {
|
|
||||||
let abort = create_abort_signal();
|
|
||||||
let app_cfg = Arc::clone(&ctx.app.config);
|
|
||||||
let role_for_input = ctx.role.clone();
|
|
||||||
let input = Input::from_str(ctx, prompt, role_for_input);
|
|
||||||
let client = input.create_client()?;
|
|
||||||
ctx.before_chat_completion(&input)?;
|
|
||||||
let (output, tool_results) =
|
|
||||||
call_chat_completions(&input, false, false, client.as_ref(), ctx, abort).await?;
|
|
||||||
ctx.after_chat_completion(app_cfg.as_ref(), &input, &output, &tool_results)?;
|
|
||||||
|
|
||||||
Ok(output)
|
|
||||||
}
|
|
||||||
|
|
||||||
fn try_parse_json(raw: &str) -> Option<Value> {
|
|
||||||
let cleaned = strip_code_fences(raw.trim());
|
|
||||||
|
|
||||||
serde_json::from_str(cleaned).ok()
|
|
||||||
}
|
|
||||||
|
|
||||||
fn strip_code_fences(s: &str) -> &str {
|
|
||||||
let after_open = s
|
|
||||||
.strip_prefix("```json")
|
|
||||||
.or_else(|| s.strip_prefix("```"))
|
|
||||||
.map(str::trim_start)
|
|
||||||
.unwrap_or(s);
|
|
||||||
after_open
|
|
||||||
.strip_suffix("```")
|
|
||||||
.map(str::trim_end)
|
|
||||||
.unwrap_or(after_open)
|
|
||||||
}
|
|
||||||
|
|
||||||
#[cfg(test)]
|
|
||||||
mod tests {
|
|
||||||
use super::*;
|
|
||||||
use serde_json::json;
|
|
||||||
|
|
||||||
#[test]
|
|
||||||
fn try_parse_json_accepts_plain_object() {
|
|
||||||
let v = try_parse_json(r#"{"a": 1}"#).unwrap();
|
|
||||||
|
|
||||||
assert_eq!(v, json!({"a": 1}));
|
|
||||||
}
|
|
||||||
|
|
||||||
#[test]
|
|
||||||
fn try_parse_json_strips_json_fences() {
|
|
||||||
let raw = "```json\n{\"a\": 1}\n```";
|
|
||||||
|
|
||||||
let v = try_parse_json(raw).unwrap();
|
|
||||||
|
|
||||||
assert_eq!(v, json!({"a": 1}));
|
|
||||||
}
|
|
||||||
|
|
||||||
#[test]
|
|
||||||
fn try_parse_json_strips_bare_fences() {
|
|
||||||
let raw = "```\n{\"a\": 1}\n```";
|
|
||||||
|
|
||||||
let v = try_parse_json(raw).unwrap();
|
|
||||||
|
|
||||||
assert_eq!(v, json!({"a": 1}));
|
|
||||||
}
|
|
||||||
|
|
||||||
#[test]
|
|
||||||
fn try_parse_json_tolerates_whitespace() {
|
|
||||||
let v = try_parse_json(" \n {\"x\": true}\n\n").unwrap();
|
|
||||||
|
|
||||||
assert_eq!(v, json!({"x": true}));
|
|
||||||
}
|
|
||||||
|
|
||||||
#[test]
|
|
||||||
fn try_parse_json_returns_none_on_prose() {
|
|
||||||
assert!(try_parse_json("Here is the result: it's good").is_none());
|
|
||||||
}
|
|
||||||
|
|
||||||
#[test]
|
|
||||||
fn try_parse_json_returns_none_on_partial_json() {
|
|
||||||
assert!(try_parse_json("{\"a\": ").is_none());
|
|
||||||
}
|
|
||||||
|
|
||||||
#[test]
|
|
||||||
fn try_parse_json_accepts_arrays() {
|
|
||||||
let v = try_parse_json("[1, 2, 3]").unwrap();
|
|
||||||
|
|
||||||
assert_eq!(v, json!([1, 2, 3]));
|
|
||||||
}
|
|
||||||
|
|
||||||
#[test]
|
|
||||||
fn build_extractor_prompt_includes_schema_and_input() {
|
|
||||||
let schema = json!({"type": "object"});
|
|
||||||
|
|
||||||
let prompt = build_extractor_prompt("hello", &schema, false);
|
|
||||||
|
|
||||||
assert!(prompt.contains("Schema:"));
|
|
||||||
assert!(prompt.contains("Input:"));
|
|
||||||
assert!(prompt.contains("hello"));
|
|
||||||
}
|
|
||||||
|
|
||||||
#[test]
|
|
||||||
fn build_extractor_prompt_repair_includes_repair_instruction() {
|
|
||||||
let schema = json!({"type": "object"});
|
|
||||||
|
|
||||||
let prompt = build_extractor_prompt("oops", &schema, true);
|
|
||||||
|
|
||||||
assert!(prompt.contains("previous response"));
|
|
||||||
assert!(prompt.contains("oops"));
|
|
||||||
}
|
|
||||||
|
|
||||||
#[test]
|
|
||||||
fn build_extractor_role_disables_tools_and_mcp() {
|
|
||||||
let role = build_extractor_role().expect("builtin role must exist");
|
|
||||||
|
|
||||||
assert_eq!(role.enabled_tools().as_deref(), Some(""));
|
|
||||||
assert_eq!(role.enabled_mcp_servers().as_deref(), Some(""));
|
|
||||||
}
|
|
||||||
}
|
|
||||||
-1116
File diff suppressed because it is too large
Load Diff
@@ -1,369 +0,0 @@
|
|||||||
use super::state::StateManager;
|
|
||||||
use super::types::{ApprovalNode, InputNode};
|
|
||||||
use crate::config::RequestContext;
|
|
||||||
use crate::function::user_interaction::{USER_FUNCTION_PREFIX, handle_user_tool};
|
|
||||||
use anyhow::{Context, Result, anyhow, bail};
|
|
||||||
use serde_json::{Value, json};
|
|
||||||
use std::collections::HashMap;
|
|
||||||
|
|
||||||
const CHOICE_KEY: &str = "choice";
|
|
||||||
const INPUT_KEY: &str = "input";
|
|
||||||
|
|
||||||
pub struct ApprovalNodeExecutor;
|
|
||||||
|
|
||||||
impl ApprovalNodeExecutor {
|
|
||||||
pub async fn execute(
|
|
||||||
node: &ApprovalNode,
|
|
||||||
state_manager: &mut StateManager,
|
|
||||||
ctx: &mut RequestContext,
|
|
||||||
) -> Result<String> {
|
|
||||||
let question = state_manager
|
|
||||||
.interpolate(&node.question)
|
|
||||||
.context("Failed to interpolate approval question")?;
|
|
||||||
|
|
||||||
let response = handle_user_tool(
|
|
||||||
ctx,
|
|
||||||
&format!("{USER_FUNCTION_PREFIX}ask"),
|
|
||||||
&json!({ "question": question, "options": node.options }),
|
|
||||||
)
|
|
||||||
.await
|
|
||||||
.context("user__ask failed")?;
|
|
||||||
|
|
||||||
if let Some(err) = response.get("error").and_then(Value::as_str) {
|
|
||||||
bail!("Approval interaction failed: {err}");
|
|
||||||
}
|
|
||||||
|
|
||||||
let choice = response
|
|
||||||
.get("answer")
|
|
||||||
.and_then(Value::as_str)
|
|
||||||
.context("Approval response missing 'answer' field")?
|
|
||||||
.to_string();
|
|
||||||
|
|
||||||
apply_state_updates_with_var(&node.state_updates, state_manager, CHOICE_KEY, &choice);
|
|
||||||
|
|
||||||
resolve_approval_route(node, &choice)
|
|
||||||
}
|
|
||||||
}
|
|
||||||
|
|
||||||
pub struct InputNodeExecutor;
|
|
||||||
|
|
||||||
impl InputNodeExecutor {
|
|
||||||
pub async fn execute(
|
|
||||||
node: &InputNode,
|
|
||||||
node_next: Option<&str>,
|
|
||||||
state_manager: &mut StateManager,
|
|
||||||
ctx: &mut RequestContext,
|
|
||||||
) -> Result<String> {
|
|
||||||
let question = build_input_question(node, state_manager)?;
|
|
||||||
|
|
||||||
let response = handle_user_tool(
|
|
||||||
ctx,
|
|
||||||
&format!("{USER_FUNCTION_PREFIX}input"),
|
|
||||||
&json!({ "question": question }),
|
|
||||||
)
|
|
||||||
.await
|
|
||||||
.context("user__input failed")?;
|
|
||||||
|
|
||||||
if let Some(err) = response.get("error").and_then(Value::as_str) {
|
|
||||||
bail!("Input interaction failed: {err}");
|
|
||||||
}
|
|
||||||
|
|
||||||
let raw = response
|
|
||||||
.get("answer")
|
|
||||||
.and_then(Value::as_str)
|
|
||||||
.context("Input response missing 'answer' field")?
|
|
||||||
.to_string();
|
|
||||||
|
|
||||||
let input_text = if raw.is_empty() {
|
|
||||||
node.default
|
|
||||||
.as_ref()
|
|
||||||
.map(|t| state_manager.interpolate_lenient(t))
|
|
||||||
.unwrap_or_default()
|
|
||||||
} else {
|
|
||||||
raw
|
|
||||||
};
|
|
||||||
|
|
||||||
if let Some(expr) = &node.validation
|
|
||||||
&& !validate_length(&input_text, expr)?
|
|
||||||
{
|
|
||||||
bail!(
|
|
||||||
"Input failed validation '{}' (got {} chars)",
|
|
||||||
expr,
|
|
||||||
input_text.chars().count()
|
|
||||||
);
|
|
||||||
}
|
|
||||||
|
|
||||||
apply_state_updates_with_var(&node.state_updates, state_manager, INPUT_KEY, &input_text);
|
|
||||||
|
|
||||||
node_next
|
|
||||||
.map(String::from)
|
|
||||||
.ok_or_else(|| anyhow!("Input node has no `next` set"))
|
|
||||||
}
|
|
||||||
}
|
|
||||||
|
|
||||||
fn build_input_question(node: &InputNode, state_manager: &StateManager) -> Result<String> {
|
|
||||||
let mut question = state_manager
|
|
||||||
.interpolate(&node.question)
|
|
||||||
.context("Failed to interpolate input question")?;
|
|
||||||
|
|
||||||
if let Some(default_template) = &node.default {
|
|
||||||
let default = state_manager.interpolate_lenient(default_template);
|
|
||||||
if !default.is_empty() {
|
|
||||||
question = format!("{question} [default: {default}]");
|
|
||||||
}
|
|
||||||
}
|
|
||||||
|
|
||||||
Ok(question)
|
|
||||||
}
|
|
||||||
|
|
||||||
fn resolve_approval_route(node: &ApprovalNode, choice: &str) -> Result<String> {
|
|
||||||
if let Some(target) = node.routes.get(choice) {
|
|
||||||
return Ok(target.clone());
|
|
||||||
}
|
|
||||||
|
|
||||||
Ok(node.on_other.clone())
|
|
||||||
}
|
|
||||||
|
|
||||||
fn apply_state_updates_with_var(
|
|
||||||
updates: &Option<HashMap<String, String>>,
|
|
||||||
state_manager: &mut StateManager,
|
|
||||||
var_name: &str,
|
|
||||||
var_value: &str,
|
|
||||||
) {
|
|
||||||
let Some(updates) = updates else {
|
|
||||||
return;
|
|
||||||
};
|
|
||||||
let prev = state_manager.state().get(var_name).cloned();
|
|
||||||
state_manager
|
|
||||||
.state_mut()
|
|
||||||
.set(var_name.into(), Value::String(var_value.to_string()));
|
|
||||||
|
|
||||||
for (key, template) in updates {
|
|
||||||
let value = state_manager.interpolate_lenient(template);
|
|
||||||
state_manager
|
|
||||||
.state_mut()
|
|
||||||
.set(key.clone(), Value::String(value));
|
|
||||||
}
|
|
||||||
|
|
||||||
match prev {
|
|
||||||
Some(v) => state_manager.state_mut().set(var_name.into(), v),
|
|
||||||
None => {
|
|
||||||
state_manager.state_mut().set(var_name.into(), Value::Null);
|
|
||||||
}
|
|
||||||
}
|
|
||||||
}
|
|
||||||
|
|
||||||
/// Evaluate a `len(input) OP N` expression where OP is one of `>`, `>=`,
|
|
||||||
/// `<`, `<=`, `==`. Lengths are byte counts (matches Rust's `str::len`).
|
|
||||||
/// Other expressions are rejected at runtime.
|
|
||||||
fn validate_length(input: &str, expr: &str) -> Result<bool> {
|
|
||||||
let trimmed = expr.trim();
|
|
||||||
let after_len = trimmed
|
|
||||||
.strip_prefix("len(input)")
|
|
||||||
.map(str::trim)
|
|
||||||
.ok_or_else(|| {
|
|
||||||
anyhow!(
|
|
||||||
"Unsupported validation expression '{expr}'; only `len(input) OP N` is supported"
|
|
||||||
)
|
|
||||||
})?;
|
|
||||||
|
|
||||||
let (op, rhs_str) = if let Some(rest) = after_len.strip_prefix(">=") {
|
|
||||||
(">=", rest)
|
|
||||||
} else if let Some(rest) = after_len.strip_prefix("<=") {
|
|
||||||
("<=", rest)
|
|
||||||
} else if let Some(rest) = after_len.strip_prefix("==") {
|
|
||||||
("==", rest)
|
|
||||||
} else if let Some(rest) = after_len.strip_prefix('>') {
|
|
||||||
(">", rest)
|
|
||||||
} else if let Some(rest) = after_len.strip_prefix('<') {
|
|
||||||
("<", rest)
|
|
||||||
} else {
|
|
||||||
bail!("No comparison operator in validation expression '{expr}'");
|
|
||||||
};
|
|
||||||
|
|
||||||
let rhs: usize = rhs_str
|
|
||||||
.trim()
|
|
||||||
.parse()
|
|
||||||
.with_context(|| format!("Invalid right-hand side in validation '{expr}'"))?;
|
|
||||||
|
|
||||||
let len = input.len();
|
|
||||||
Ok(match op {
|
|
||||||
">=" => len >= rhs,
|
|
||||||
"<=" => len <= rhs,
|
|
||||||
"==" => len == rhs,
|
|
||||||
">" => len > rhs,
|
|
||||||
"<" => len < rhs,
|
|
||||||
_ => unreachable!(),
|
|
||||||
})
|
|
||||||
}
|
|
||||||
|
|
||||||
#[cfg(test)]
|
|
||||||
mod tests {
|
|
||||||
use super::super::types::*;
|
|
||||||
use super::*;
|
|
||||||
use serde_json::json;
|
|
||||||
|
|
||||||
fn manager_with(pairs: &[(&str, Value)]) -> StateManager {
|
|
||||||
let mut map = HashMap::new();
|
|
||||||
for (k, v) in pairs {
|
|
||||||
map.insert((*k).into(), v.clone());
|
|
||||||
}
|
|
||||||
|
|
||||||
StateManager::new(map)
|
|
||||||
}
|
|
||||||
|
|
||||||
fn approval(options: &[&str], routes: &[(&str, &str)], on_other: &str) -> ApprovalNode {
|
|
||||||
let mut r = HashMap::new();
|
|
||||||
for (k, v) in routes {
|
|
||||||
r.insert((*k).into(), (*v).into());
|
|
||||||
}
|
|
||||||
|
|
||||||
ApprovalNode {
|
|
||||||
question: "?".into(),
|
|
||||||
options: options.iter().map(|s| (*s).into()).collect(),
|
|
||||||
routes: r,
|
|
||||||
on_other: on_other.into(),
|
|
||||||
state_updates: None,
|
|
||||||
}
|
|
||||||
}
|
|
||||||
|
|
||||||
fn input(question: &str) -> InputNode {
|
|
||||||
InputNode {
|
|
||||||
question: question.into(),
|
|
||||||
default: None,
|
|
||||||
validation: None,
|
|
||||||
state_updates: None,
|
|
||||||
}
|
|
||||||
}
|
|
||||||
|
|
||||||
#[test]
|
|
||||||
fn validate_length_supports_all_comparison_operators() {
|
|
||||||
assert!(validate_length("hello", "len(input) > 0").unwrap());
|
|
||||||
assert!(!validate_length("", "len(input) > 0").unwrap());
|
|
||||||
assert!(validate_length("hello", "len(input) >= 5").unwrap());
|
|
||||||
assert!(!validate_length("hi", "len(input) >= 5").unwrap());
|
|
||||||
assert!(validate_length("hello", "len(input) < 10").unwrap());
|
|
||||||
assert!(!validate_length("hello world!", "len(input) < 10").unwrap());
|
|
||||||
assert!(validate_length("hi", "len(input) <= 2").unwrap());
|
|
||||||
assert!(!validate_length("hello", "len(input) <= 2").unwrap());
|
|
||||||
assert!(validate_length("hello", "len(input) == 5").unwrap());
|
|
||||||
assert!(!validate_length("hello", "len(input) == 3").unwrap());
|
|
||||||
}
|
|
||||||
|
|
||||||
#[test]
|
|
||||||
fn validate_length_handles_whitespace() {
|
|
||||||
assert!(validate_length("hi", " len(input) >= 1 ").unwrap());
|
|
||||||
}
|
|
||||||
|
|
||||||
#[test]
|
|
||||||
fn validate_length_rejects_unsupported_expressions() {
|
|
||||||
assert!(validate_length("x", "matches /[a-z]+/").is_err());
|
|
||||||
assert!(validate_length("x", "len(input)").is_err());
|
|
||||||
assert!(validate_length("x", "len(input) >").is_err());
|
|
||||||
assert!(validate_length("x", "len(input) >= abc").is_err());
|
|
||||||
}
|
|
||||||
|
|
||||||
#[test]
|
|
||||||
fn approval_route_lookup_returns_target_on_match() {
|
|
||||||
let node = approval(
|
|
||||||
&["yes", "no"],
|
|
||||||
&[("yes", "deploy"), ("no", "cancel")],
|
|
||||||
"clarify",
|
|
||||||
);
|
|
||||||
|
|
||||||
assert_eq!(resolve_approval_route(&node, "yes").unwrap(), "deploy");
|
|
||||||
assert_eq!(resolve_approval_route(&node, "no").unwrap(), "cancel");
|
|
||||||
}
|
|
||||||
|
|
||||||
#[test]
|
|
||||||
fn approval_route_lookup_falls_back_to_on_other_for_unknown_choice() {
|
|
||||||
let node = approval(
|
|
||||||
&["yes", "no"],
|
|
||||||
&[("yes", "deploy"), ("no", "cancel")],
|
|
||||||
"clarify",
|
|
||||||
);
|
|
||||||
|
|
||||||
assert_eq!(resolve_approval_route(&node, "maybe").unwrap(), "clarify");
|
|
||||||
assert_eq!(
|
|
||||||
resolve_approval_route(&node, "free-form text").unwrap(),
|
|
||||||
"clarify"
|
|
||||||
);
|
|
||||||
}
|
|
||||||
|
|
||||||
#[test]
|
|
||||||
fn state_updates_expose_choice_during_evaluation_only() {
|
|
||||||
let mut updates = HashMap::new();
|
|
||||||
updates.insert("decision".into(), "{{choice}}".into());
|
|
||||||
let mut state = manager_with(&[]);
|
|
||||||
|
|
||||||
apply_state_updates_with_var(&Some(updates), &mut state, CHOICE_KEY, "approve");
|
|
||||||
|
|
||||||
assert_eq!(state.state().get("decision"), Some(&json!("approve")));
|
|
||||||
assert_eq!(state.state().get(CHOICE_KEY), Some(&Value::Null));
|
|
||||||
}
|
|
||||||
|
|
||||||
#[test]
|
|
||||||
fn state_updates_preserve_pre_existing_var_value() {
|
|
||||||
let mut updates = HashMap::new();
|
|
||||||
updates.insert("decision".into(), "{{choice}}".into());
|
|
||||||
let mut state = manager_with(&[("choice", json!("preserved"))]);
|
|
||||||
|
|
||||||
apply_state_updates_with_var(&Some(updates), &mut state, CHOICE_KEY, "approve");
|
|
||||||
|
|
||||||
assert_eq!(state.state().get("decision"), Some(&json!("approve")));
|
|
||||||
assert_eq!(state.state().get(CHOICE_KEY), Some(&json!("preserved")));
|
|
||||||
}
|
|
||||||
|
|
||||||
#[test]
|
|
||||||
fn state_updates_for_input_use_input_key() {
|
|
||||||
let mut updates = HashMap::new();
|
|
||||||
updates.insert("api_key".into(), "{{input}}".into());
|
|
||||||
let mut state = manager_with(&[]);
|
|
||||||
|
|
||||||
apply_state_updates_with_var(&Some(updates), &mut state, INPUT_KEY, "sk-12345");
|
|
||||||
|
|
||||||
assert_eq!(state.state().get("api_key"), Some(&json!("sk-12345")));
|
|
||||||
assert_eq!(state.state().get(INPUT_KEY), Some(&Value::Null));
|
|
||||||
}
|
|
||||||
|
|
||||||
#[test]
|
|
||||||
fn input_question_appends_default_when_present() {
|
|
||||||
let state = manager_with(&[("name", json!("alice"))]);
|
|
||||||
let mut node = input("Hi, what's your name?");
|
|
||||||
node.default = Some("{{name}}".into());
|
|
||||||
|
|
||||||
let q = build_input_question(&node, &state).unwrap();
|
|
||||||
|
|
||||||
assert_eq!(q, "Hi, what's your name? [default: alice]");
|
|
||||||
}
|
|
||||||
|
|
||||||
#[test]
|
|
||||||
fn input_question_omits_default_when_blank_after_interpolation() {
|
|
||||||
let state = manager_with(&[]);
|
|
||||||
let mut node = input("Enter value:");
|
|
||||||
node.default = Some("{{missing}}".into());
|
|
||||||
|
|
||||||
let q = build_input_question(&node, &state).unwrap();
|
|
||||||
|
|
||||||
assert_eq!(q, "Enter value:");
|
|
||||||
}
|
|
||||||
|
|
||||||
#[test]
|
|
||||||
fn input_question_uses_no_default_when_field_absent() {
|
|
||||||
let state = manager_with(&[]);
|
|
||||||
let node = input("Enter value:");
|
|
||||||
|
|
||||||
let q = build_input_question(&node, &state).unwrap();
|
|
||||||
|
|
||||||
assert_eq!(q, "Enter value:");
|
|
||||||
}
|
|
||||||
|
|
||||||
#[test]
|
|
||||||
fn no_state_updates_means_var_never_appears_in_state() {
|
|
||||||
let mut state = manager_with(&[]);
|
|
||||||
|
|
||||||
apply_state_updates_with_var(&None, &mut state, CHOICE_KEY, "approve");
|
|
||||||
|
|
||||||
assert!(state.state().get(CHOICE_KEY).is_none());
|
|
||||||
}
|
|
||||||
}
|
|
||||||
File diff suppressed because it is too large
Load Diff
-27
@@ -2,7 +2,6 @@ mod cli;
|
|||||||
mod client;
|
mod client;
|
||||||
mod config;
|
mod config;
|
||||||
mod function;
|
mod function;
|
||||||
mod graph;
|
|
||||||
mod rag;
|
mod rag;
|
||||||
mod render;
|
mod render;
|
||||||
mod repl;
|
mod repl;
|
||||||
@@ -83,23 +82,8 @@ async fn main() -> Result<()> {
|
|||||||
|
|
||||||
let log_path = setup_logger()?;
|
let log_path = setup_logger()?;
|
||||||
|
|
||||||
if let Some(version) = &cli.update {
|
|
||||||
let version = version.clone();
|
|
||||||
let force = cli.force;
|
|
||||||
return tokio::task::spawn_blocking(move || config::run_self_update(version, force))
|
|
||||||
.await?;
|
|
||||||
}
|
|
||||||
|
|
||||||
install_builtins()?;
|
install_builtins()?;
|
||||||
|
|
||||||
if let Some(category) = cli.install {
|
|
||||||
return config::install_assets(category);
|
|
||||||
}
|
|
||||||
|
|
||||||
if let Some(url) = cli.install_from.as_deref() {
|
|
||||||
return config::install_remote(url, cli.filter, cli.install_force);
|
|
||||||
}
|
|
||||||
|
|
||||||
if let Some(client_arg) = &cli.authenticate {
|
if let Some(client_arg) = &cli.authenticate {
|
||||||
let cfg = Config::load_with_interpolation(true).await?;
|
let cfg = Config::load_with_interpolation(true).await?;
|
||||||
let app_config = AppConfig::from_config(cfg)?;
|
let app_config = AppConfig::from_config(cfg)?;
|
||||||
@@ -327,17 +311,6 @@ async fn start_directive(
|
|||||||
abort_signal: AbortSignal,
|
abort_signal: AbortSignal,
|
||||||
) -> Result<()> {
|
) -> Result<()> {
|
||||||
let app: Arc<AppConfig> = Arc::clone(&ctx.app.config);
|
let app: Arc<AppConfig> = Arc::clone(&ctx.app.config);
|
||||||
|
|
||||||
if graph::active_agent_graph_name(ctx).is_some() {
|
|
||||||
ctx.before_chat_completion(&input)?;
|
|
||||||
let output =
|
|
||||||
graph::run_active_agent_graph(ctx, &input.text(), abort_signal.clone()).await?;
|
|
||||||
app.print_markdown(&output)?;
|
|
||||||
ctx.after_chat_completion(app.as_ref(), &input, &output, &[])?;
|
|
||||||
ctx.exit_session()?;
|
|
||||||
return Ok(());
|
|
||||||
}
|
|
||||||
|
|
||||||
let client = input.create_client()?;
|
let client = input.create_client()?;
|
||||||
let extract_code = !*IS_STDOUT_TERMINAL && code_mode;
|
let extract_code = !*IS_STDOUT_TERMINAL && code_mode;
|
||||||
ctx.before_chat_completion(&input)?;
|
ctx.before_chat_completion(&input)?;
|
||||||
|
|||||||
+11
-18
@@ -8,7 +8,6 @@ use crate::vault::interpolate_secrets;
|
|||||||
use anyhow::{Context, Result, anyhow};
|
use anyhow::{Context, Result, anyhow};
|
||||||
use futures_util::{StreamExt, TryStreamExt, stream};
|
use futures_util::{StreamExt, TryStreamExt, stream};
|
||||||
use http::{HeaderName, HeaderValue};
|
use http::{HeaderName, HeaderValue};
|
||||||
use indexmap::IndexMap;
|
|
||||||
use indoc::formatdoc;
|
use indoc::formatdoc;
|
||||||
use rmcp::service::RunningService;
|
use rmcp::service::RunningService;
|
||||||
use rmcp::transport::StreamableHttpClientTransport;
|
use rmcp::transport::StreamableHttpClientTransport;
|
||||||
@@ -50,29 +49,23 @@ impl Clone for ServerCatalog {
|
|||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
#[derive(Debug, Clone, Deserialize, Serialize)]
|
#[derive(Debug, Clone, Deserialize)]
|
||||||
pub(crate) struct McpServersConfig {
|
pub(crate) struct McpServersConfig {
|
||||||
#[serde(rename = "mcpServers")]
|
#[serde(rename = "mcpServers")]
|
||||||
pub mcp_servers: IndexMap<String, McpServer>,
|
pub mcp_servers: HashMap<String, McpServer>,
|
||||||
}
|
}
|
||||||
|
|
||||||
#[derive(Debug, Clone, Deserialize, Serialize, PartialEq, Eq)]
|
#[derive(Debug, Clone, Deserialize)]
|
||||||
#[serde(deny_unknown_fields)]
|
#[serde(deny_unknown_fields)]
|
||||||
pub(crate) struct McpServer {
|
pub(crate) struct McpServer {
|
||||||
#[serde(rename = "type")]
|
#[serde(rename = "type")]
|
||||||
pub transport_type: McpTransportType,
|
pub transport_type: McpTransportType,
|
||||||
#[serde(skip_serializing_if = "Option::is_none")]
|
|
||||||
pub command: Option<String>,
|
pub command: Option<String>,
|
||||||
#[serde(skip_serializing_if = "Option::is_none")]
|
|
||||||
pub args: Option<Vec<String>>,
|
pub args: Option<Vec<String>>,
|
||||||
#[serde(skip_serializing_if = "Option::is_none")]
|
pub env: Option<HashMap<String, JsonField>>,
|
||||||
pub env: Option<IndexMap<String, JsonField>>,
|
|
||||||
#[serde(skip_serializing_if = "Option::is_none")]
|
|
||||||
pub cwd: Option<String>,
|
pub cwd: Option<String>,
|
||||||
#[serde(skip_serializing_if = "Option::is_none")]
|
|
||||||
pub url: Option<String>,
|
pub url: Option<String>,
|
||||||
#[serde(skip_serializing_if = "Option::is_none")]
|
pub headers: Option<HashMap<String, String>>,
|
||||||
pub headers: Option<IndexMap<String, String>>,
|
|
||||||
}
|
}
|
||||||
|
|
||||||
impl McpServer {
|
impl McpServer {
|
||||||
@@ -118,7 +111,7 @@ impl McpServer {
|
|||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
#[derive(Debug, Clone, Deserialize, Serialize, PartialEq, Eq, Hash)]
|
#[derive(Debug, Clone, Deserialize, PartialEq, Eq, Hash)]
|
||||||
#[serde(rename_all = "lowercase")]
|
#[serde(rename_all = "lowercase")]
|
||||||
pub(crate) enum McpTransportType {
|
pub(crate) enum McpTransportType {
|
||||||
Stdio,
|
Stdio,
|
||||||
@@ -126,7 +119,7 @@ pub(crate) enum McpTransportType {
|
|||||||
Sse,
|
Sse,
|
||||||
}
|
}
|
||||||
|
|
||||||
#[derive(Debug, Clone, Deserialize, Serialize, PartialEq, Eq)]
|
#[derive(Debug, Clone, Deserialize)]
|
||||||
#[serde(untagged)]
|
#[serde(untagged)]
|
||||||
pub(crate) enum JsonField {
|
pub(crate) enum JsonField {
|
||||||
Str(String),
|
Str(String),
|
||||||
@@ -359,7 +352,7 @@ pub(crate) async fn spawn_mcp_server(
|
|||||||
|
|
||||||
async fn spawn_http_mcp_server(
|
async fn spawn_http_mcp_server(
|
||||||
url: &str,
|
url: &str,
|
||||||
headers: Option<&IndexMap<String, String>>,
|
headers: Option<&HashMap<String, String>>,
|
||||||
) -> Result<Arc<ConnectedServer>> {
|
) -> Result<Arc<ConnectedServer>> {
|
||||||
let transport = if let Some(hdrs) = headers
|
let transport = if let Some(hdrs) = headers
|
||||||
&& !hdrs.is_empty()
|
&& !hdrs.is_empty()
|
||||||
@@ -389,7 +382,7 @@ async fn spawn_http_mcp_server(
|
|||||||
|
|
||||||
async fn spawn_sse_mcp_server(
|
async fn spawn_sse_mcp_server(
|
||||||
url: &str,
|
url: &str,
|
||||||
headers: Option<&IndexMap<String, String>>,
|
headers: Option<&HashMap<String, String>>,
|
||||||
) -> Result<Arc<ConnectedServer>> {
|
) -> Result<Arc<ConnectedServer>> {
|
||||||
let sse = LegacySseTransport::connect(url, headers)
|
let sse = LegacySseTransport::connect(url, headers)
|
||||||
.await
|
.await
|
||||||
@@ -489,7 +482,7 @@ mod tests {
|
|||||||
}
|
}
|
||||||
|
|
||||||
fn make_registry_with_config(server_names: &[&str]) -> McpRegistry {
|
fn make_registry_with_config(server_names: &[&str]) -> McpRegistry {
|
||||||
let mut mcp_servers = IndexMap::new();
|
let mut mcp_servers = HashMap::new();
|
||||||
for name in server_names {
|
for name in server_names {
|
||||||
mcp_servers.insert(name.to_string(), stdio_server("echo"));
|
mcp_servers.insert(name.to_string(), stdio_server("echo"));
|
||||||
}
|
}
|
||||||
@@ -537,7 +530,7 @@ mod tests {
|
|||||||
|
|
||||||
#[test]
|
#[test]
|
||||||
fn validate_stdio_with_headers_fails() {
|
fn validate_stdio_with_headers_fails() {
|
||||||
let mut headers = IndexMap::new();
|
let mut headers = HashMap::new();
|
||||||
headers.insert("Auth".into(), "Bearer tok".into());
|
headers.insert("Auth".into(), "Bearer tok".into());
|
||||||
let spec = McpServer {
|
let spec = McpServer {
|
||||||
transport_type: McpTransportType::Stdio,
|
transport_type: McpTransportType::Stdio,
|
||||||
|
|||||||
+14
-24
@@ -1,14 +1,13 @@
|
|||||||
use anyhow::{Context, Result, anyhow};
|
use anyhow::{Context, Result, anyhow};
|
||||||
use eventsource_stream::{EventStream, Eventsource};
|
|
||||||
use fmt::{Display, Formatter};
|
use fmt::{Display, Formatter};
|
||||||
use futures_util::StreamExt;
|
use futures_util::StreamExt;
|
||||||
use futures_util::stream::BoxStream;
|
|
||||||
use indexmap::IndexMap;
|
|
||||||
use mpsc::error::SendError;
|
use mpsc::error::SendError;
|
||||||
use mpsc::{OwnedPermit, Receiver, Sender, channel};
|
use mpsc::{OwnedPermit, Receiver, Sender, channel};
|
||||||
|
use reqwest::Client;
|
||||||
use reqwest::header::{HeaderMap, HeaderName, HeaderValue};
|
use reqwest::header::{HeaderMap, HeaderName, HeaderValue};
|
||||||
use reqwest::{Client, header};
|
use reqwest_eventsource::{Event, EventSource};
|
||||||
use rmcp::model::{ClientJsonRpcMessage, ServerJsonRpcMessage};
|
use rmcp::model::{ClientJsonRpcMessage, ServerJsonRpcMessage};
|
||||||
|
use std::collections::HashMap;
|
||||||
use std::error::Error;
|
use std::error::Error;
|
||||||
use std::fmt;
|
use std::fmt;
|
||||||
use std::future::Future;
|
use std::future::Future;
|
||||||
@@ -18,8 +17,6 @@ use tokio::sync::mpsc;
|
|||||||
use tokio::time::Duration;
|
use tokio::time::Duration;
|
||||||
use url::Url;
|
use url::Url;
|
||||||
|
|
||||||
type SseEventStream = EventStream<BoxStream<'static, reqwest::Result<bytes::Bytes>>>;
|
|
||||||
|
|
||||||
const CHANNEL_BUF: usize = 64;
|
const CHANNEL_BUF: usize = 64;
|
||||||
|
|
||||||
pub struct LegacySseTransport {
|
pub struct LegacySseTransport {
|
||||||
@@ -28,10 +25,7 @@ pub struct LegacySseTransport {
|
|||||||
}
|
}
|
||||||
|
|
||||||
impl LegacySseTransport {
|
impl LegacySseTransport {
|
||||||
pub async fn connect(
|
pub async fn connect(sse_url: &str, headers: Option<&HashMap<String, String>>) -> Result<Self> {
|
||||||
sse_url: &str,
|
|
||||||
headers: Option<&IndexMap<String, String>>,
|
|
||||||
) -> Result<Self> {
|
|
||||||
let base_url =
|
let base_url =
|
||||||
Url::parse(sse_url).with_context(|| format!("Invalid SSE URL: {sse_url}"))?;
|
Url::parse(sse_url).with_context(|| format!("Invalid SSE URL: {sse_url}"))?;
|
||||||
|
|
||||||
@@ -53,15 +47,8 @@ impl LegacySseTransport {
|
|||||||
.build()
|
.build()
|
||||||
.context("Failed to build HTTP client")?;
|
.context("Failed to build HTTP client")?;
|
||||||
|
|
||||||
let response = client
|
let request = client.get(sse_url);
|
||||||
.get(sse_url)
|
let mut es = EventSource::new(request).context("Failed to open SSE connection")?;
|
||||||
.header(header::ACCEPT, "text/event-stream")
|
|
||||||
.send()
|
|
||||||
.await
|
|
||||||
.context("Failed to open SSE connection")?
|
|
||||||
.error_for_status()
|
|
||||||
.context("SSE server returned an error status")?;
|
|
||||||
let mut es: SseEventStream = response.bytes_stream().boxed().eventsource();
|
|
||||||
|
|
||||||
let post_endpoint = wait_for_endpoint_event(&mut es, &base_url).await?;
|
let post_endpoint = wait_for_endpoint_event(&mut es, &base_url).await?;
|
||||||
|
|
||||||
@@ -96,17 +83,18 @@ impl LegacySseTransport {
|
|||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
async fn wait_for_endpoint_event(es: &mut SseEventStream, base_url: &Url) -> Result<String> {
|
async fn wait_for_endpoint_event(es: &mut EventSource, base_url: &Url) -> Result<String> {
|
||||||
let timeout = Duration::from_secs(30);
|
let timeout = Duration::from_secs(30);
|
||||||
tokio::time::timeout(timeout, async {
|
tokio::time::timeout(timeout, async {
|
||||||
while let Some(event) = es.next().await {
|
while let Some(event) = es.next().await {
|
||||||
match event {
|
match event {
|
||||||
Ok(msg) if msg.event == "endpoint" => {
|
Ok(Event::Open) => {}
|
||||||
|
Ok(Event::Message(msg)) if msg.event == "endpoint" => {
|
||||||
let endpoint = msg.data.trim().to_string();
|
let endpoint = msg.data.trim().to_string();
|
||||||
let resolved = resolve_endpoint(&endpoint, base_url)?;
|
let resolved = resolve_endpoint(&endpoint, base_url)?;
|
||||||
return Ok(resolved);
|
return Ok(resolved);
|
||||||
}
|
}
|
||||||
Ok(_) => {}
|
Ok(Event::Message(_)) => {}
|
||||||
Err(e) => {
|
Err(e) => {
|
||||||
return Err(anyhow!(
|
return Err(anyhow!(
|
||||||
"SSE connection error while waiting for endpoint event: {e}"
|
"SSE connection error while waiting for endpoint event: {e}"
|
||||||
@@ -132,10 +120,10 @@ fn resolve_endpoint(endpoint: &str, base_url: &Url) -> Result<String> {
|
|||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
async fn sse_reader_task(mut es: SseEventStream, tx: Sender<ServerJsonRpcMessage>) {
|
async fn sse_reader_task(mut es: EventSource, tx: Sender<ServerJsonRpcMessage>) {
|
||||||
while let Some(event) = es.next().await {
|
while let Some(event) = es.next().await {
|
||||||
match event {
|
match event {
|
||||||
Ok(msg) if msg.event == "message" => {
|
Ok(Event::Message(msg)) if msg.event == "message" => {
|
||||||
match serde_json::from_str::<ServerJsonRpcMessage>(&msg.data) {
|
match serde_json::from_str::<ServerJsonRpcMessage>(&msg.data) {
|
||||||
Ok(rpc_msg) => {
|
Ok(rpc_msg) => {
|
||||||
if tx.send(rpc_msg).await.is_err() {
|
if tx.send(rpc_msg).await.is_err() {
|
||||||
@@ -148,12 +136,14 @@ async fn sse_reader_task(mut es: SseEventStream, tx: Sender<ServerJsonRpcMessage
|
|||||||
}
|
}
|
||||||
}
|
}
|
||||||
Ok(_) => {}
|
Ok(_) => {}
|
||||||
|
Err(reqwest_eventsource::Error::StreamEnded) => break,
|
||||||
Err(e) => {
|
Err(e) => {
|
||||||
error!("SSE stream error: {e}");
|
error!("SSE stream error: {e}");
|
||||||
break;
|
break;
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
es.close();
|
||||||
}
|
}
|
||||||
|
|
||||||
async fn post_writer_task(
|
async fn post_writer_task(
|
||||||
|
|||||||
+11
-135
@@ -16,8 +16,7 @@ use parking_lot::RwLock;
|
|||||||
use serde::{Deserialize, Serialize};
|
use serde::{Deserialize, Serialize};
|
||||||
use serde_json::json;
|
use serde_json::json;
|
||||||
use std::{
|
use std::{
|
||||||
collections::HashMap, env, fmt, fmt::Debug, fs, hash::Hash, path::Path, sync::Arc,
|
collections::HashMap, env, fmt::Debug, fs, hash::Hash, path::Path, sync::Arc, time::Duration,
|
||||||
time::Duration,
|
|
||||||
};
|
};
|
||||||
use tokio::time::sleep;
|
use tokio::time::sleep;
|
||||||
|
|
||||||
@@ -57,7 +56,7 @@ pub struct Rag {
|
|||||||
}
|
}
|
||||||
|
|
||||||
impl Debug for Rag {
|
impl Debug for Rag {
|
||||||
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
|
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
|
||||||
f.debug_struct("Rag")
|
f.debug_struct("Rag")
|
||||||
.field("name", &self.name)
|
.field("name", &self.name)
|
||||||
.field("path", &self.path)
|
.field("path", &self.path)
|
||||||
@@ -82,126 +81,11 @@ impl Clone for Rag {
|
|||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
#[derive(Debug, Clone, Default)]
|
|
||||||
pub struct RagInitConfig {
|
|
||||||
pub embedding_model: Option<String>,
|
|
||||||
pub chunk_size: Option<usize>,
|
|
||||||
pub chunk_overlap: Option<usize>,
|
|
||||||
pub reranker_model: Option<String>,
|
|
||||||
pub top_k: Option<usize>,
|
|
||||||
pub batch_size: Option<usize>,
|
|
||||||
}
|
|
||||||
|
|
||||||
impl Rag {
|
impl Rag {
|
||||||
fn create_embeddings_client(&self, model: Model) -> Result<Box<dyn Client>> {
|
fn create_embeddings_client(&self, model: Model) -> Result<Box<dyn Client>> {
|
||||||
init_client(&self.app_config, model)
|
init_client(&self.app_config, model)
|
||||||
}
|
}
|
||||||
|
|
||||||
pub async fn init_with_config(
|
|
||||||
app: &AppConfig,
|
|
||||||
name: &str,
|
|
||||||
save_path: &Path,
|
|
||||||
doc_paths: &[String],
|
|
||||||
config: &RagInitConfig,
|
|
||||||
abort_signal: AbortSignal,
|
|
||||||
) -> Result<Self> {
|
|
||||||
if doc_paths.is_empty() {
|
|
||||||
bail!("Cannot build RAG knowledge base '{name}' with no documents");
|
|
||||||
}
|
|
||||||
println!("⚙ Initializing RAG...");
|
|
||||||
let data = Self::resolve_init_data(app, config)?;
|
|
||||||
let mut rag = Self::create(app, name, save_path, data)?;
|
|
||||||
let loaders = app.document_loaders.clone();
|
|
||||||
let (spinner, spinner_rx) = Spinner::create("");
|
|
||||||
abortable_run_with_spinner_rx(
|
|
||||||
rag.sync_documents(doc_paths, true, loaders, Some(spinner)),
|
|
||||||
spinner_rx,
|
|
||||||
abort_signal,
|
|
||||||
)
|
|
||||||
.await?;
|
|
||||||
if rag.save()? {
|
|
||||||
println!("✓ Saved RAG to '{}'.", save_path.display());
|
|
||||||
}
|
|
||||||
Ok(rag)
|
|
||||||
}
|
|
||||||
|
|
||||||
fn resolve_init_data(app: &AppConfig, config: &RagInitConfig) -> Result<RagData> {
|
|
||||||
let embedding_model_id = config
|
|
||||||
.embedding_model
|
|
||||||
.clone()
|
|
||||||
.or_else(|| app.rag_embedding_model.clone());
|
|
||||||
let embedding_model_id = match embedding_model_id {
|
|
||||||
Some(value) => {
|
|
||||||
println!("Embedding model: {value}");
|
|
||||||
value
|
|
||||||
}
|
|
||||||
None => {
|
|
||||||
if !*IS_STDOUT_TERMINAL {
|
|
||||||
bail!(
|
|
||||||
"RAG knowledge base needs an embedding model. Set `embedding_model` \
|
|
||||||
on the rag node, or run the agent interactively once."
|
|
||||||
);
|
|
||||||
}
|
|
||||||
let models = list_models(app, ModelType::Embedding);
|
|
||||||
if models.is_empty() {
|
|
||||||
bail!("No available embedding model");
|
|
||||||
}
|
|
||||||
select_embedding_model(&models)?
|
|
||||||
}
|
|
||||||
};
|
|
||||||
let embedding_model =
|
|
||||||
Model::retrieve_model(app, &embedding_model_id, ModelType::Embedding)?;
|
|
||||||
|
|
||||||
let chunk_size = match config.chunk_size.or(app.rag_chunk_size) {
|
|
||||||
Some(value) => {
|
|
||||||
println!("Chunk size: {value}");
|
|
||||||
value
|
|
||||||
}
|
|
||||||
None => {
|
|
||||||
if !*IS_STDOUT_TERMINAL {
|
|
||||||
bail!(
|
|
||||||
"RAG knowledge base needs a chunk_size. Set `chunk_size` on the \
|
|
||||||
rag node, or run the agent interactively once."
|
|
||||||
);
|
|
||||||
}
|
|
||||||
set_chunk_size(&embedding_model)?
|
|
||||||
}
|
|
||||||
};
|
|
||||||
let chunk_overlap = match config.chunk_overlap.or(app.rag_chunk_overlap) {
|
|
||||||
Some(value) => {
|
|
||||||
println!("Chunk overlap: {value}");
|
|
||||||
value
|
|
||||||
}
|
|
||||||
None => {
|
|
||||||
if !*IS_STDOUT_TERMINAL {
|
|
||||||
bail!(
|
|
||||||
"RAG knowledge base needs a chunk_overlap. Set `chunk_overlap` on \
|
|
||||||
the rag node, or run the agent interactively once."
|
|
||||||
);
|
|
||||||
}
|
|
||||||
set_chunk_overlay(chunk_size / 20)?
|
|
||||||
}
|
|
||||||
};
|
|
||||||
|
|
||||||
let reranker_model = config
|
|
||||||
.reranker_model
|
|
||||||
.clone()
|
|
||||||
.or_else(|| app.rag_reranker_model.clone());
|
|
||||||
let top_k = config.top_k.unwrap_or(app.rag_top_k);
|
|
||||||
let batch_size = config
|
|
||||||
.batch_size
|
|
||||||
.or_else(|| embedding_model.max_batch_size());
|
|
||||||
|
|
||||||
Ok(RagData::new(
|
|
||||||
embedding_model.id(),
|
|
||||||
chunk_size,
|
|
||||||
chunk_overlap,
|
|
||||||
reranker_model,
|
|
||||||
top_k,
|
|
||||||
batch_size,
|
|
||||||
))
|
|
||||||
}
|
|
||||||
|
|
||||||
pub async fn init(
|
pub async fn init(
|
||||||
app: &AppConfig,
|
app: &AppConfig,
|
||||||
name: &str,
|
name: &str,
|
||||||
@@ -431,14 +315,6 @@ impl Rag {
|
|||||||
self.name == TEMP_RAG_NAME
|
self.name == TEMP_RAG_NAME
|
||||||
}
|
}
|
||||||
|
|
||||||
pub fn configured_top_k(&self) -> usize {
|
|
||||||
self.data.top_k
|
|
||||||
}
|
|
||||||
|
|
||||||
pub fn configured_reranker(&self) -> Option<&str> {
|
|
||||||
self.data.reranker_model.as_deref()
|
|
||||||
}
|
|
||||||
|
|
||||||
pub async fn search(
|
pub async fn search(
|
||||||
&self,
|
&self,
|
||||||
text: &str,
|
text: &str,
|
||||||
@@ -447,7 +323,7 @@ impl Rag {
|
|||||||
abort_signal: AbortSignal,
|
abort_signal: AbortSignal,
|
||||||
) -> Result<(String, String, Vec<DocumentId>)> {
|
) -> Result<(String, String, Vec<DocumentId>)> {
|
||||||
let ret = abortable_run_with_spinner(
|
let ret = abortable_run_with_spinner(
|
||||||
self.hybrid_search(text, top_k, rerank_model),
|
self.hybird_search(text, top_k, rerank_model),
|
||||||
"Searching",
|
"Searching",
|
||||||
abort_signal,
|
abort_signal,
|
||||||
)
|
)
|
||||||
@@ -707,7 +583,7 @@ impl Rag {
|
|||||||
Ok(())
|
Ok(())
|
||||||
}
|
}
|
||||||
|
|
||||||
async fn hybrid_search(
|
async fn hybird_search(
|
||||||
&self,
|
&self,
|
||||||
query: &str,
|
query: &str,
|
||||||
top_k: usize,
|
top_k: usize,
|
||||||
@@ -905,7 +781,7 @@ pub struct RagData {
|
|||||||
}
|
}
|
||||||
|
|
||||||
impl Debug for RagData {
|
impl Debug for RagData {
|
||||||
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
|
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
|
||||||
f.debug_struct("RagData")
|
f.debug_struct("RagData")
|
||||||
.field("embedding_model", &self.embedding_model)
|
.field("embedding_model", &self.embedding_model)
|
||||||
.field("chunk_size", &self.chunk_size)
|
.field("chunk_size", &self.chunk_size)
|
||||||
@@ -1033,7 +909,7 @@ pub type FileId = usize;
|
|||||||
pub struct DocumentId(usize);
|
pub struct DocumentId(usize);
|
||||||
|
|
||||||
impl Debug for DocumentId {
|
impl Debug for DocumentId {
|
||||||
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
|
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
|
||||||
let (file_index, document_index) = self.split();
|
let (file_index, document_index) = self.split();
|
||||||
f.write_fmt(format_args!("{file_index}-{document_index}"))
|
f.write_fmt(format_args!("{file_index}-{document_index}"))
|
||||||
}
|
}
|
||||||
@@ -1075,8 +951,8 @@ impl SelectOption {
|
|||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
impl fmt::Display for SelectOption {
|
impl std::fmt::Display for SelectOption {
|
||||||
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
|
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
|
||||||
write!(f, "{} ({})", self.value, self.description)
|
write!(f, "{} ({})", self.value, self.description)
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
@@ -1380,13 +1256,13 @@ mod tests {
|
|||||||
|
|
||||||
#[test]
|
#[test]
|
||||||
fn get_separators_returns_language_specific() {
|
fn get_separators_returns_language_specific() {
|
||||||
let rs_seps = get_separators("rs");
|
let rs_seps = splitter::get_separators("rs");
|
||||||
assert!(rs_seps.iter().any(|s| s.contains("fn ")));
|
assert!(rs_seps.iter().any(|s| s.contains("fn ")));
|
||||||
|
|
||||||
let py_seps = get_separators("py");
|
let py_seps = splitter::get_separators("py");
|
||||||
assert!(py_seps.iter().any(|s| s.contains("def ")));
|
assert!(py_seps.iter().any(|s| s.contains("def ")));
|
||||||
|
|
||||||
let md_seps = get_separators("md");
|
let md_seps = splitter::get_separators("md");
|
||||||
assert!(md_seps.iter().any(|s| s.contains("# ")));
|
assert!(md_seps.iter().any(|s| s.contains("# ")));
|
||||||
}
|
}
|
||||||
|
|
||||||
|
|||||||
@@ -17,11 +17,7 @@ pub async fn render_stream(
|
|||||||
rx: UnboundedReceiver<SseEvent>,
|
rx: UnboundedReceiver<SseEvent>,
|
||||||
app: &AppConfig,
|
app: &AppConfig,
|
||||||
abort_signal: AbortSignal,
|
abort_signal: AbortSignal,
|
||||||
silent: bool,
|
|
||||||
) -> Result<()> {
|
) -> Result<()> {
|
||||||
if silent {
|
|
||||||
return drain_silently(rx, &abort_signal).await;
|
|
||||||
}
|
|
||||||
let ret = if *IS_STDOUT_TERMINAL && app.highlight {
|
let ret = if *IS_STDOUT_TERMINAL && app.highlight {
|
||||||
let render_options = app.render_options()?;
|
let render_options = app.render_options()?;
|
||||||
let mut render = MarkdownRender::init(render_options)?;
|
let mut render = MarkdownRender::init(render_options)?;
|
||||||
@@ -32,22 +28,6 @@ pub async fn render_stream(
|
|||||||
ret.map_err(|err| err.context("Failed to reader stream"))
|
ret.map_err(|err| err.context("Failed to reader stream"))
|
||||||
}
|
}
|
||||||
|
|
||||||
async fn drain_silently(
|
|
||||||
mut rx: UnboundedReceiver<SseEvent>,
|
|
||||||
abort_signal: &AbortSignal,
|
|
||||||
) -> Result<()> {
|
|
||||||
loop {
|
|
||||||
if abort_signal.aborted() {
|
|
||||||
break;
|
|
||||||
}
|
|
||||||
match rx.recv().await {
|
|
||||||
Some(SseEvent::Done) | None => break,
|
|
||||||
Some(SseEvent::Text(_)) => {}
|
|
||||||
}
|
|
||||||
}
|
|
||||||
Ok(())
|
|
||||||
}
|
|
||||||
|
|
||||||
pub fn render_error(err: anyhow::Error) {
|
pub fn render_error(err: anyhow::Error) {
|
||||||
eprintln!("{}", error_text(&pretty_error(&err)));
|
eprintln!("{}", error_text(&pretty_error(&err)));
|
||||||
}
|
}
|
||||||
|
|||||||
+45
-124
@@ -7,22 +7,20 @@ use self::highlighter::ReplHighlighter;
|
|||||||
use self::prompt::ReplPrompt;
|
use self::prompt::ReplPrompt;
|
||||||
|
|
||||||
use crate::client::{call_chat_completions, call_chat_completions_streaming, init_client, oauth};
|
use crate::client::{call_chat_completions, call_chat_completions_streaming, init_client, oauth};
|
||||||
|
use crate::config::paths;
|
||||||
use crate::config::{
|
use crate::config::{
|
||||||
AgentVariables, AppConfig, AssertState, Input, LastMessage, RequestContext, StateFlags,
|
AgentVariables, AppConfig, AssertState, Input, LastMessage, RequestContext, StateFlags,
|
||||||
macro_execute,
|
macro_execute,
|
||||||
};
|
};
|
||||||
use crate::config::{AssetCategory, paths};
|
|
||||||
use crate::render::render_error;
|
use crate::render::render_error;
|
||||||
use crate::utils::{
|
use crate::utils::{
|
||||||
AbortSignal, abortable_run_with_spinner, create_abort_signal, dimmed_text, set_text, temp_file,
|
AbortSignal, abortable_run_with_spinner, create_abort_signal, dimmed_text, set_text, temp_file,
|
||||||
};
|
};
|
||||||
|
|
||||||
use crate::{config, graph, resolve_oauth_client};
|
use crate::resolve_oauth_client;
|
||||||
use anyhow::{Context, Result, bail};
|
use anyhow::{Context, Result, bail};
|
||||||
use crossterm::cursor::SetCursorStyle;
|
use crossterm::cursor::SetCursorStyle;
|
||||||
use fancy_regex::Regex;
|
use fancy_regex::Regex;
|
||||||
use indoc::indoc;
|
|
||||||
use log::warn;
|
|
||||||
use parking_lot::RwLock;
|
use parking_lot::RwLock;
|
||||||
use reedline::CursorConfig;
|
use reedline::CursorConfig;
|
||||||
use reedline::{
|
use reedline::{
|
||||||
@@ -33,20 +31,10 @@ use reedline::{
|
|||||||
use reedline::{MenuBuilder, Signal};
|
use reedline::{MenuBuilder, Signal};
|
||||||
use std::sync::LazyLock;
|
use std::sync::LazyLock;
|
||||||
use std::{env, process, sync::Arc};
|
use std::{env, process, sync::Arc};
|
||||||
use tokio::task;
|
|
||||||
|
|
||||||
const MENU_NAME: &str = "completion_menu";
|
const MENU_NAME: &str = "completion_menu";
|
||||||
|
|
||||||
pub const DEFAULT_CONTINUATION_PROMPT: &str = indoc! {"
|
static REPL_COMMANDS: LazyLock<[ReplCommand; 39]> = LazyLock::new(|| {
|
||||||
[SYSTEM REMINDER - TODO CONTINUATION]
|
|
||||||
You have incomplete tasks. Rules:
|
|
||||||
1. BEFORE marking a todo done: verify the work compiles/works. No premature completion.
|
|
||||||
2. If a todo is broad (e.g. \"implement X and implement Y\"): break it into specific subtasks FIRST using todo__add, then work on those.\n\
|
|
||||||
3. Each todo should be atomic and be \"single responsibility\" - completable in one focused action.
|
|
||||||
4. Continue with the next pending item now. Call tools immediately."
|
|
||||||
};
|
|
||||||
|
|
||||||
static REPL_COMMANDS: LazyLock<[ReplCommand; 42]> = LazyLock::new(|| {
|
|
||||||
[
|
[
|
||||||
ReplCommand::new(".help", "Show this help guide", AssertState::pass()),
|
ReplCommand::new(".help", "Show this help guide", AssertState::pass()),
|
||||||
ReplCommand::new(".info", "Show system info", AssertState::pass()),
|
ReplCommand::new(".info", "Show system info", AssertState::pass()),
|
||||||
@@ -60,11 +48,6 @@ static REPL_COMMANDS: LazyLock<[ReplCommand; 42]> = LazyLock::new(|| {
|
|||||||
"Modify configuration file",
|
"Modify configuration file",
|
||||||
AssertState::False(StateFlags::AGENT),
|
AssertState::False(StateFlags::AGENT),
|
||||||
),
|
),
|
||||||
ReplCommand::new(
|
|
||||||
".edit mcp-config",
|
|
||||||
"Modify the MCP servers configuration file",
|
|
||||||
AssertState::False(StateFlags::AGENT),
|
|
||||||
),
|
|
||||||
ReplCommand::new(".model", "Switch LLM model", AssertState::pass()),
|
ReplCommand::new(".model", "Switch LLM model", AssertState::pass()),
|
||||||
ReplCommand::new(
|
ReplCommand::new(
|
||||||
".prompt",
|
".prompt",
|
||||||
@@ -158,7 +141,7 @@ static REPL_COMMANDS: LazyLock<[ReplCommand; 42]> = LazyLock::new(|| {
|
|||||||
ReplCommand::new(
|
ReplCommand::new(
|
||||||
".clear todo",
|
".clear todo",
|
||||||
"Clear the todo list and stop auto-continuation",
|
"Clear the todo list and stop auto-continuation",
|
||||||
AssertState::pass(),
|
AssertState::True(StateFlags::AGENT),
|
||||||
),
|
),
|
||||||
ReplCommand::new(
|
ReplCommand::new(
|
||||||
".rag",
|
".rag",
|
||||||
@@ -218,16 +201,6 @@ static REPL_COMMANDS: LazyLock<[ReplCommand; 42]> = LazyLock::new(|| {
|
|||||||
"View or modify the Loki vault",
|
"View or modify the Loki vault",
|
||||||
AssertState::pass(),
|
AssertState::pass(),
|
||||||
),
|
),
|
||||||
ReplCommand::new(
|
|
||||||
".install",
|
|
||||||
"Reinstall bundled assets, or install assets from a remote git repo (.install remote <url>)",
|
|
||||||
AssertState::pass(),
|
|
||||||
),
|
|
||||||
ReplCommand::new(
|
|
||||||
".update",
|
|
||||||
"Update Loki to the latest release (or a specified version)",
|
|
||||||
AssertState::pass(),
|
|
||||||
),
|
|
||||||
ReplCommand::new(".exit", "Exit REPL", AssertState::pass()),
|
ReplCommand::new(".exit", "Exit REPL", AssertState::pass()),
|
||||||
]
|
]
|
||||||
});
|
});
|
||||||
@@ -514,12 +487,6 @@ pub async fn run_repl_command(
|
|||||||
),
|
),
|
||||||
},
|
},
|
||||||
".session" => {
|
".session" => {
|
||||||
if let Some(name) = graph::active_agent_graph_name(ctx) {
|
|
||||||
bail!(
|
|
||||||
"Graph-based agent '{name}' does not support sessions. \
|
|
||||||
The graph manages its own state."
|
|
||||||
);
|
|
||||||
}
|
|
||||||
let app = Arc::clone(&ctx.app.config);
|
let app = Arc::clone(&ctx.app.config);
|
||||||
ctx.use_session(app.as_ref(), args, abort_signal.clone())
|
ctx.use_session(app.as_ref(), args, abort_signal.clone())
|
||||||
.await?;
|
.await?;
|
||||||
@@ -531,41 +498,13 @@ pub async fn run_repl_command(
|
|||||||
};
|
};
|
||||||
eprintln!("\n📢 {}", color.italic().paint("Autonaming the session."),);
|
eprintln!("\n📢 {}", color.italic().paint("Autonaming the session."),);
|
||||||
if let Err(err) = ctx.autoname_session(app.as_ref()).await {
|
if let Err(err) = ctx.autoname_session(app.as_ref()).await {
|
||||||
warn!("Failed to autonaming the session: {err}");
|
log::warn!("Failed to autonaming the session: {err}");
|
||||||
}
|
}
|
||||||
if let Some(session) = ctx.session.as_mut() {
|
if let Some(session) = ctx.session.as_mut() {
|
||||||
session.set_autonaming(false);
|
session.set_autonaming(false);
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
".install" => {
|
|
||||||
let trimmed = args.map(str::trim).unwrap_or("");
|
|
||||||
let mut parts = trimmed.splitn(2, char::is_whitespace);
|
|
||||||
match parts.next() {
|
|
||||||
Some("remote") => {
|
|
||||||
let rest = parts.next().unwrap_or("").trim();
|
|
||||||
config::install_remote_from_repl_args(rest)?;
|
|
||||||
}
|
|
||||||
Some(name) if !name.is_empty() => match AssetCategory::parse(name) {
|
|
||||||
Some(category) => config::install_assets(category)?,
|
|
||||||
None => println!(
|
|
||||||
"Unknown asset category '{name}'. Valid categories: {}",
|
|
||||||
AssetCategory::NAMES.join(", ")
|
|
||||||
),
|
|
||||||
},
|
|
||||||
_ => println!(
|
|
||||||
"Usage: .install <{}> | .install remote <git-url>",
|
|
||||||
AssetCategory::NAMES.join("|")
|
|
||||||
),
|
|
||||||
}
|
|
||||||
}
|
|
||||||
".update" => {
|
|
||||||
if ctx.macro_flag {
|
|
||||||
bail!("Cannot perform this operation because you are in a macro")
|
|
||||||
}
|
|
||||||
let version = args.map(|s| s.trim().to_string());
|
|
||||||
task::spawn_blocking(move || config::run_self_update(version, false)).await??;
|
|
||||||
}
|
|
||||||
".rag" => {
|
".rag" => {
|
||||||
ctx.use_rag(args, abort_signal.clone()).await?;
|
ctx.use_rag(args, abort_signal.clone()).await?;
|
||||||
}
|
}
|
||||||
@@ -656,13 +595,8 @@ pub async fn run_repl_command(
|
|||||||
let app = Arc::clone(&ctx.app.config);
|
let app = Arc::clone(&ctx.app.config);
|
||||||
ctx.edit_agent_config(app.as_ref())?;
|
ctx.edit_agent_config(app.as_ref())?;
|
||||||
}
|
}
|
||||||
Some("mcp-config") => {
|
|
||||||
ctx.edit_mcp_config()?;
|
|
||||||
}
|
|
||||||
_ => {
|
_ => {
|
||||||
println!(
|
println!(r#"Usage: .edit <config|role|session|rag-docs|agent-config>"#)
|
||||||
r#"Usage: .edit <config|mcp-config|role|session|rag-docs|agent-config>"#
|
|
||||||
)
|
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
@@ -830,18 +764,25 @@ pub async fn run_repl_command(
|
|||||||
bail!("Use '.empty session' instead");
|
bail!("Use '.empty session' instead");
|
||||||
}
|
}
|
||||||
Some("todo") => {
|
Some("todo") => {
|
||||||
let config = ctx.auto_continue_config();
|
let cleared = match ctx.agent.as_mut() {
|
||||||
if !config.enabled {
|
Some(agent) => {
|
||||||
bail!(
|
if !agent.auto_continue_enabled() {
|
||||||
"Auto-continue is not enabled. Set 'auto_continue: true' in your config to enable it."
|
bail!(
|
||||||
);
|
"The todo system is not enabled for this agent. Set 'auto_continue: true' in the agent's config.yaml to enable it."
|
||||||
}
|
);
|
||||||
if ctx.todo_list.is_empty() {
|
}
|
||||||
println!("Todo list is already empty.");
|
if ctx.todo_list.is_empty() {
|
||||||
} else {
|
println!("Todo list is already empty.");
|
||||||
ctx.clear_todo_list();
|
false
|
||||||
println!("Todo list cleared.");
|
} else {
|
||||||
}
|
ctx.clear_todo_list();
|
||||||
|
println!("Todo list cleared.");
|
||||||
|
true
|
||||||
|
}
|
||||||
|
}
|
||||||
|
None => bail!("No active agent"),
|
||||||
|
};
|
||||||
|
let _ = cleared;
|
||||||
}
|
}
|
||||||
_ => unknown_command()?,
|
_ => unknown_command()?,
|
||||||
},
|
},
|
||||||
@@ -914,18 +855,8 @@ async fn ask(
|
|||||||
tokio::time::sleep(std::time::Duration::from_millis(100)).await;
|
tokio::time::sleep(std::time::Duration::from_millis(100)).await;
|
||||||
}
|
}
|
||||||
|
|
||||||
let app = Arc::clone(&ctx.app.config);
|
|
||||||
|
|
||||||
if graph::active_agent_graph_name(ctx).is_some() {
|
|
||||||
ctx.before_chat_completion(&input)?;
|
|
||||||
let output =
|
|
||||||
graph::run_active_agent_graph(ctx, &input.text(), abort_signal.clone()).await?;
|
|
||||||
app.print_markdown(&output)?;
|
|
||||||
ctx.after_chat_completion(app.as_ref(), &input, &output, &[])?;
|
|
||||||
return Ok(());
|
|
||||||
}
|
|
||||||
|
|
||||||
let client = input.create_client()?;
|
let client = input.create_client()?;
|
||||||
|
let app = Arc::clone(&ctx.app.config);
|
||||||
ctx.before_chat_completion(&input)?;
|
ctx.before_chat_completion(&input)?;
|
||||||
let (output, tool_results) = if input.stream() {
|
let (output, tool_results) = if input.stream() {
|
||||||
call_chat_completions_streaming(&input, client.as_ref(), ctx, abort_signal.clone()).await?
|
call_chat_completions_streaming(&input, client.as_ref(), ctx, abort_signal.clone()).await?
|
||||||
@@ -950,22 +881,19 @@ async fn ask(
|
|||||||
)
|
)
|
||||||
.await
|
.await
|
||||||
} else {
|
} else {
|
||||||
let do_continue = should_continue(ctx);
|
let should_continue = agent_should_continue(ctx);
|
||||||
|
|
||||||
if do_continue {
|
if should_continue {
|
||||||
let full_prompt = {
|
let full_prompt = {
|
||||||
let config = ctx.auto_continue_config();
|
|
||||||
let todo_state = ctx.todo_list.render_for_model();
|
let todo_state = ctx.todo_list.render_for_model();
|
||||||
let remaining = ctx.todo_list.incomplete_count();
|
let remaining = ctx.todo_list.incomplete_count();
|
||||||
ctx.set_last_continuation_response(output.clone());
|
ctx.set_last_continuation_response(output.clone());
|
||||||
ctx.increment_auto_continue_count();
|
ctx.increment_auto_continue_count();
|
||||||
|
let agent = ctx.agent.as_mut().expect("agent checked above");
|
||||||
let count = ctx.auto_continue_count;
|
let count = ctx.auto_continue_count;
|
||||||
let max = config.max_continues;
|
let max = agent.max_auto_continues();
|
||||||
|
|
||||||
let prompt = config
|
let prompt = agent.continuation_prompt();
|
||||||
.continuation_prompt
|
|
||||||
.as_deref()
|
|
||||||
.unwrap_or(DEFAULT_CONTINUATION_PROMPT);
|
|
||||||
|
|
||||||
let color = if app.light_theme() {
|
let color = if app.light_theme() {
|
||||||
nu_ansi_term::Color::LightGray
|
nu_ansi_term::Color::LightGray
|
||||||
@@ -993,7 +921,7 @@ async fn ask(
|
|||||||
};
|
};
|
||||||
eprintln!("\n📢 {}", color.italic().paint("Autonaming the session."),);
|
eprintln!("\n📢 {}", color.italic().paint("Autonaming the session."),);
|
||||||
if let Err(err) = ctx.autoname_session(app.as_ref()).await {
|
if let Err(err) = ctx.autoname_session(app.as_ref()).await {
|
||||||
warn!("Failed to autonaming the session: {err}");
|
log::warn!("Failed to autonaming the session: {err}");
|
||||||
}
|
}
|
||||||
if let Some(session) = ctx.session.as_mut() {
|
if let Some(session) = ctx.session.as_mut() {
|
||||||
session.set_autonaming(false);
|
session.set_autonaming(false);
|
||||||
@@ -1006,7 +934,7 @@ async fn ask(
|
|||||||
.is_some_and(|s| s.needs_compression(app.compression_threshold));
|
.is_some_and(|s| s.needs_compression(app.compression_threshold));
|
||||||
|
|
||||||
if needs_compression {
|
if needs_compression {
|
||||||
let agent_can_continue_after_compress = should_continue(ctx);
|
let agent_can_continue_after_compress = agent_should_continue(ctx);
|
||||||
|
|
||||||
if let Some(session) = ctx.session.as_mut() {
|
if let Some(session) = ctx.session.as_mut() {
|
||||||
session.set_compressing(true);
|
session.set_compressing(true);
|
||||||
@@ -1020,7 +948,7 @@ async fn ask(
|
|||||||
eprintln!("\n📢 {}", color.italic().paint("Compressing the session."),);
|
eprintln!("\n📢 {}", color.italic().paint("Compressing the session."),);
|
||||||
|
|
||||||
if let Err(err) = ctx.compress_session().await {
|
if let Err(err) = ctx.compress_session().await {
|
||||||
warn!("Failed to compress the session: {err}");
|
log::warn!("Failed to compress the session: {err}");
|
||||||
}
|
}
|
||||||
if let Some(session) = ctx.session.as_mut() {
|
if let Some(session) = ctx.session.as_mut() {
|
||||||
session.set_compressing(false);
|
session.set_compressing(false);
|
||||||
@@ -1028,17 +956,14 @@ async fn ask(
|
|||||||
|
|
||||||
if agent_can_continue_after_compress {
|
if agent_can_continue_after_compress {
|
||||||
let full_prompt = {
|
let full_prompt = {
|
||||||
let config = ctx.auto_continue_config();
|
|
||||||
let todo_state = ctx.todo_list.render_for_model();
|
let todo_state = ctx.todo_list.render_for_model();
|
||||||
let remaining = ctx.todo_list.incomplete_count();
|
let remaining = ctx.todo_list.incomplete_count();
|
||||||
ctx.increment_auto_continue_count();
|
ctx.increment_auto_continue_count();
|
||||||
|
let agent = ctx.agent.as_mut().expect("agent checked above");
|
||||||
let count = ctx.auto_continue_count;
|
let count = ctx.auto_continue_count;
|
||||||
let max = config.max_continues;
|
let max = agent.max_auto_continues();
|
||||||
|
|
||||||
let prompt = config
|
let prompt = agent.continuation_prompt();
|
||||||
.continuation_prompt
|
|
||||||
.as_deref()
|
|
||||||
.unwrap_or(DEFAULT_CONTINUATION_PROMPT);
|
|
||||||
|
|
||||||
let color = if app.light_theme() {
|
let color = if app.light_theme() {
|
||||||
nu_ansi_term::Color::LightGray
|
nu_ansi_term::Color::LightGray
|
||||||
@@ -1064,12 +989,10 @@ async fn ask(
|
|||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
fn should_continue(ctx: &RequestContext) -> bool {
|
fn agent_should_continue(ctx: &RequestContext) -> bool {
|
||||||
let config = ctx.auto_continue_config();
|
ctx.agent.as_ref().is_some_and(|agent| {
|
||||||
ctx.app.config.function_calling_support
|
agent.auto_continue_enabled() && ctx.auto_continue_count < agent.max_auto_continues()
|
||||||
&& config.enabled
|
}) && ctx.todo_list.has_incomplete()
|
||||||
&& ctx.auto_continue_count < config.max_continues
|
|
||||||
&& ctx.todo_list.has_incomplete()
|
|
||||||
}
|
}
|
||||||
|
|
||||||
fn reset_continuation(ctx: &mut RequestContext) {
|
fn reset_continuation(ctx: &mut RequestContext) {
|
||||||
@@ -1265,8 +1188,8 @@ mod tests {
|
|||||||
}
|
}
|
||||||
|
|
||||||
#[test]
|
#[test]
|
||||||
fn repl_commands_has_42_entries() {
|
fn repl_commands_has_39_entries() {
|
||||||
assert_eq!(REPL_COMMANDS.len(), 42);
|
assert_eq!(REPL_COMMANDS.len(), 39);
|
||||||
}
|
}
|
||||||
|
|
||||||
#[test]
|
#[test]
|
||||||
@@ -1388,15 +1311,13 @@ mod tests {
|
|||||||
}
|
}
|
||||||
|
|
||||||
#[test]
|
#[test]
|
||||||
fn repl_commands_clear_todo_always_available() {
|
fn repl_commands_clear_todo_requires_agent() {
|
||||||
let cmd = REPL_COMMANDS
|
let cmd = REPL_COMMANDS
|
||||||
.iter()
|
.iter()
|
||||||
.find(|c| c.name == ".clear todo")
|
.find(|c| c.name == ".clear todo")
|
||||||
.unwrap();
|
.unwrap();
|
||||||
assert!(cmd.is_valid(StateFlags::AGENT));
|
assert!(cmd.is_valid(StateFlags::AGENT));
|
||||||
assert!(cmd.is_valid(StateFlags::empty()));
|
assert!(!cmd.is_valid(StateFlags::empty()));
|
||||||
assert!(cmd.is_valid(StateFlags::SESSION));
|
|
||||||
assert!(cmd.is_valid(StateFlags::ROLE));
|
|
||||||
}
|
}
|
||||||
|
|
||||||
#[test]
|
#[test]
|
||||||
|
|||||||
+16
-2
@@ -34,6 +34,7 @@ use is_terminal::IsTerminal;
|
|||||||
use std::borrow::Cow;
|
use std::borrow::Cow;
|
||||||
use std::sync::LazyLock;
|
use std::sync::LazyLock;
|
||||||
use std::{cmp, env, path::PathBuf, process};
|
use std::{cmp, env, path::PathBuf, process};
|
||||||
|
use unicode_segmentation::UnicodeSegmentation;
|
||||||
|
|
||||||
pub static CODE_BLOCK_RE: LazyLock<Regex> =
|
pub static CODE_BLOCK_RE: LazyLock<Regex> =
|
||||||
LazyLock::new(|| Regex::new(r"(?ms)```\w*(.*)```").unwrap());
|
LazyLock::new(|| Regex::new(r"(?ms)```\w*(.*)```").unwrap());
|
||||||
@@ -73,8 +74,21 @@ pub fn parse_bool(value: &str) -> Option<bool> {
|
|||||||
}
|
}
|
||||||
|
|
||||||
pub fn estimate_token_length(text: &str) -> usize {
|
pub fn estimate_token_length(text: &str) -> usize {
|
||||||
let weighted: usize = text.chars().map(|c| if c.is_ascii() { 1 } else { 2 }).sum();
|
let words: Vec<&str> = text.unicode_words().collect();
|
||||||
weighted.div_ceil(4)
|
let mut output: f32 = 0.0;
|
||||||
|
for word in words {
|
||||||
|
if word.is_ascii() {
|
||||||
|
output += 1.3;
|
||||||
|
} else {
|
||||||
|
let count = word.chars().count();
|
||||||
|
if count == 1 {
|
||||||
|
output += 1.0
|
||||||
|
} else {
|
||||||
|
output += (count as f32) * 0.5;
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
output.ceil() as usize
|
||||||
}
|
}
|
||||||
|
|
||||||
pub fn strip_think_tag(text: &str) -> Cow<'_, str> {
|
pub fn strip_think_tag(text: &str) -> Cow<'_, str> {
|
||||||
|
|||||||
+16
-55
@@ -241,23 +241,23 @@ fn add_file(files: &mut IndexSet<String>, suffixes: Option<&Vec<String>>, path:
|
|||||||
}
|
}
|
||||||
|
|
||||||
fn is_valid_extension(suffixes: Option<&Vec<String>>, path: &Path) -> bool {
|
fn is_valid_extension(suffixes: Option<&Vec<String>>, path: &Path) -> bool {
|
||||||
let Some(suffixes) = suffixes else {
|
let filename_regex = Regex::new(r"^.+\.*").unwrap();
|
||||||
return true;
|
if let Some(suffixes) = suffixes
|
||||||
};
|
&& !suffixes.is_empty()
|
||||||
if suffixes.is_empty() {
|
{
|
||||||
return true;
|
if let Ok(Some(_)) = filename_regex.find(&suffixes.join(",")) {
|
||||||
}
|
let file_name = path
|
||||||
|
.file_name()
|
||||||
let file_name = path.file_name().and_then(|v| v.to_str());
|
.and_then(|v| v.to_str())
|
||||||
let extension = path.extension().and_then(|v| v.to_str());
|
.expect("invalid filename")
|
||||||
|
.to_string();
|
||||||
suffixes.iter().any(|suffix| {
|
return suffixes.contains(&file_name);
|
||||||
if suffix.contains('.') {
|
} else if let Some(extension) = path.extension().map(|v| v.to_string_lossy().to_string()) {
|
||||||
Some(suffix.as_str()) == file_name
|
return suffixes.contains(&extension);
|
||||||
} else {
|
|
||||||
Some(suffix.as_str()) == extension
|
|
||||||
}
|
}
|
||||||
})
|
return false;
|
||||||
|
}
|
||||||
|
true
|
||||||
}
|
}
|
||||||
|
|
||||||
#[cfg(test)]
|
#[cfg(test)]
|
||||||
@@ -352,43 +352,4 @@ mod tests {
|
|||||||
)
|
)
|
||||||
);
|
);
|
||||||
}
|
}
|
||||||
|
|
||||||
#[test]
|
|
||||||
fn test_is_valid_extension() {
|
|
||||||
let md_ext = vec!["md".to_string()];
|
|
||||||
let md_txt_ext = vec!["md".to_string(), "txt".to_string()];
|
|
||||||
let test_md_filename = vec!["test.md".to_string()];
|
|
||||||
let mixed = vec!["md".to_string(), "test.txt".to_string()];
|
|
||||||
|
|
||||||
assert!(is_valid_extension(None, Path::new("Agents.md")));
|
|
||||||
assert!(is_valid_extension(Some(&vec![]), Path::new("Agents.md")));
|
|
||||||
|
|
||||||
assert!(is_valid_extension(Some(&md_ext), Path::new("Agents.md")));
|
|
||||||
assert!(is_valid_extension(
|
|
||||||
Some(&md_ext),
|
|
||||||
Path::new("/home/atusa/code/loki.wiki/Agents.md")
|
|
||||||
));
|
|
||||||
assert!(!is_valid_extension(Some(&md_ext), Path::new("notes.txt")));
|
|
||||||
assert!(!is_valid_extension(Some(&md_ext), Path::new("README")));
|
|
||||||
|
|
||||||
assert!(is_valid_extension(Some(&md_txt_ext), Path::new("a.md")));
|
|
||||||
assert!(is_valid_extension(Some(&md_txt_ext), Path::new("a.txt")));
|
|
||||||
assert!(!is_valid_extension(Some(&md_txt_ext), Path::new("a.rs")));
|
|
||||||
|
|
||||||
assert!(is_valid_extension(
|
|
||||||
Some(&test_md_filename),
|
|
||||||
Path::new("dir/test.md")
|
|
||||||
));
|
|
||||||
assert!(!is_valid_extension(
|
|
||||||
Some(&test_md_filename),
|
|
||||||
Path::new("dir/Agents.md")
|
|
||||||
));
|
|
||||||
|
|
||||||
assert!(is_valid_extension(Some(&mixed), Path::new("Agents.md")));
|
|
||||||
assert!(is_valid_extension(Some(&mixed), Path::new("dir/test.txt")));
|
|
||||||
assert!(!is_valid_extension(
|
|
||||||
Some(&mixed),
|
|
||||||
Path::new("dir/other.txt")
|
|
||||||
));
|
|
||||||
}
|
|
||||||
}
|
}
|
||||||
|
|||||||
Reference in New Issue
Block a user