diff --git a/agent/.gitignore b/agent/.gitignore new file mode 100644 index 000000000..2c96eb1b6 --- /dev/null +++ b/agent/.gitignore @@ -0,0 +1,2 @@ +target/ +Cargo.lock diff --git a/agent/Cargo.toml b/agent/Cargo.toml new file mode 100644 index 000000000..39f4a172f --- /dev/null +++ b/agent/Cargo.toml @@ -0,0 +1,26 @@ +[workspace] + +[package] +name = "iii-agent" +version = "0.1.0" +edition = "2021" +publish = false + +[[bin]] +name = "iii-agent" +path = "src/main.rs" + +[dependencies] +iii-sdk = "=0.11.3" +tokio = { version = "1", features = ["rt-multi-thread", "macros", "sync", "signal"] } +serde = { version = "1", features = ["derive"] } +serde_json = "1" +serde_yaml = "0.9" +anyhow = "1" +tracing = "0.1" +tracing-subscriber = { version = "0.3", features = ["fmt", "env-filter"] } +clap = { version = "4", features = ["derive", "env"] } +chrono = { version = "0.4", features = ["serde"] } +reqwest = { version = "0.12", default-features = false, features = ["json", "rustls-tls", "stream"] } +futures-util = "0.3" +uuid = { version = "1", features = ["v4"] } diff --git a/agent/README.md b/agent/README.md new file mode 100644 index 000000000..685551224 --- /dev/null +++ b/agent/README.md @@ -0,0 +1,70 @@ +# iii-agent + +Linear, PostHog, Attio — they all shipped the same thing: a chat bar as the primary interface. iii-agent brings this to the iii console. It dynamically discovers every function registered by every connected worker, lets users ask questions in natural language, and the LLM decides which functions to call. "What's slow in my system?" triggers `eval::analyze_traces`. "Show me the topology" triggers `introspect::diagram`. The agent composes the answer from real data, not hallucinations. + +**Plug and play:** Build with `cargo build --release`, set `ANTHROPIC_API_KEY` in your environment, then run `./target/release/iii-agent --url ws://your-engine:49134`. It registers 7 functions, discovers all available tools from other workers, and starts accepting chat via `agent::chat`. Connect more workers and they're automatically available — no restart needed. + +## Functions + +| Function ID | Description | +|---|---| +| `agent::chat` | Send a message and get a structured JSON-UI response | +| `agent::chat_stream` | Send a message with streaming response via iii Streams | +| `agent::discover` | List all available functions the agent can orchestrate | +| `agent::plan` | Generate an execution plan DAG without executing | +| `agent::session_create` | Create a new chat session | +| `agent::session_history` | Retrieve conversation history for a session | +| `agent::session_cleanup` | Clean up expired sessions (cron-triggered) | + +## iii Primitives Used + +- **State** -- session history, cached tool definitions +- **Streams** -- streaming chat responses via `agent:events:{session_id}` group +- **Cron** -- hourly session cleanup +- **HTTP** -- chat, discovery, planning, and session management endpoints + +## Prerequisites + +- Rust 1.75+ +- Running iii engine on `ws://127.0.0.1:49134` +- `ANTHROPIC_API_KEY` environment variable set + +## Build + +```bash +cargo build --release +``` + +## Usage + +```bash +# Load the key from your secret manager (keychain, 1password, doppler, etc.) +# into the environment before launching the worker — never paste the literal +# key on the command line, since it lands in shell history and `ps` output. +export ANTHROPIC_API_KEY="$(security find-generic-password -s anthropic-api-key -w)" +./target/release/iii-agent --url ws://127.0.0.1:49134 --config ./config.yaml +``` + +``` +Options: + --config Path to config.yaml [default: ./config.yaml] + --url WebSocket URL of the iii engine [default: ws://127.0.0.1:49134] + --manifest Output module manifest as JSON and exit + -h, --help Print help +``` + +## Configuration + +```yaml +anthropic_model: "claude-sonnet-4-20250514" # model to use for chat +max_tokens: 4096 # max tokens per LLM response +max_iterations: 10 # max tool-use loops per message +session_ttl_hours: 24 # session expiry +cron_session_cleanup: "0 0 * * * *" # hourly cleanup schedule +``` + +## Tests + +```bash +cargo test +``` diff --git a/agent/build.rs b/agent/build.rs new file mode 100644 index 000000000..81caa36d6 --- /dev/null +++ b/agent/build.rs @@ -0,0 +1,6 @@ +fn main() { + println!( + "cargo:rustc-env=TARGET={}", + std::env::var("TARGET").unwrap() + ); +} diff --git a/agent/config.yaml b/agent/config.yaml new file mode 100644 index 000000000..5c88c130a --- /dev/null +++ b/agent/config.yaml @@ -0,0 +1,5 @@ +anthropic_model: "claude-sonnet-4-20250514" +max_tokens: 4096 +max_iterations: 10 +session_ttl_hours: 24 +cron_session_cleanup: "0 0 * * * *" diff --git a/agent/src/config.rs b/agent/src/config.rs new file mode 100644 index 000000000..dd0a38ada --- /dev/null +++ b/agent/src/config.rs @@ -0,0 +1,109 @@ +use anyhow::Result; +use serde::Deserialize; + +#[derive(Deserialize, Debug, Clone)] +pub struct AgentConfig { + #[serde(default = "default_model")] + pub anthropic_model: String, + #[serde(default = "default_max_tokens")] + pub max_tokens: u32, + #[serde(default = "default_max_iterations")] + pub max_iterations: u32, + #[serde(default = "default_session_ttl_hours")] + pub session_ttl_hours: u64, + #[serde(default = "default_cron_session_cleanup")] + pub cron_session_cleanup: String, +} + +fn default_model() -> String { + "claude-haiku-4-5-20251001".to_string() +} + +fn default_max_tokens() -> u32 { + 4096 +} + +fn default_max_iterations() -> u32 { + 10 +} + +fn default_session_ttl_hours() -> u64 { + 24 +} + +fn default_cron_session_cleanup() -> String { + "0 0 * * * *".to_string() +} + +impl Default for AgentConfig { + fn default() -> Self { + AgentConfig { + anthropic_model: default_model(), + max_tokens: default_max_tokens(), + max_iterations: default_max_iterations(), + session_ttl_hours: default_session_ttl_hours(), + cron_session_cleanup: default_cron_session_cleanup(), + } + } +} + +pub fn load_config(path: &str) -> Result { + let contents = std::fs::read_to_string(path)?; + let config: AgentConfig = serde_yaml::from_str(&contents)?; + validate(&config)?; + Ok(config) +} + +fn validate(cfg: &AgentConfig) -> Result<()> { + if cfg.anthropic_model.trim().is_empty() { + anyhow::bail!("config: anthropic_model must be non-empty"); + } + if cfg.max_tokens == 0 { + anyhow::bail!("config: max_tokens must be >= 1"); + } + if cfg.max_iterations == 0 { + anyhow::bail!("config: max_iterations must be >= 1"); + } + if cfg.session_ttl_hours == 0 { + anyhow::bail!("config: session_ttl_hours must be >= 1"); + } + if cfg.cron_session_cleanup.trim().is_empty() { + anyhow::bail!("config: cron_session_cleanup must be non-empty"); + } + Ok(()) +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn test_config_defaults() { + let config = AgentConfig::default(); + assert_eq!(config.anthropic_model, "claude-haiku-4-5-20251001"); + assert_eq!(config.max_tokens, 4096); + assert_eq!(config.max_iterations, 10); + assert_eq!(config.session_ttl_hours, 24); + } + + #[test] + fn test_config_from_yaml() { + let yaml = r#" +anthropic_model: "claude-sonnet-4-20250514" +max_tokens: 8192 +max_iterations: 5 +"#; + let config: AgentConfig = serde_yaml::from_str(yaml).unwrap(); + assert_eq!(config.anthropic_model, "claude-sonnet-4-20250514"); + assert_eq!(config.max_tokens, 8192); + assert_eq!(config.max_iterations, 5); + assert_eq!(config.session_ttl_hours, 24); + } + + #[test] + fn test_config_empty_yaml() { + let config: AgentConfig = serde_yaml::from_str("{}").unwrap(); + assert_eq!(config.anthropic_model, "claude-haiku-4-5-20251001"); + assert_eq!(config.max_tokens, 4096); + } +} diff --git a/agent/src/discovery.rs b/agent/src/discovery.rs new file mode 100644 index 000000000..f3e43249c --- /dev/null +++ b/agent/src/discovery.rs @@ -0,0 +1,284 @@ +use iii_sdk::{FunctionInfo, III}; +use serde_json::json; + +use crate::llm::ToolDef; + +// Infrastructure namespaces the agent should never call as tools (it would +// recurse into itself, leak engine internals, or invoke routing primitives +// directly). Everything else is exposed — new worker namespaces are +// auto-discovered without a code change. Override with +// `discovery_excluded_prefixes` in config when a deployment needs a tighter +// boundary. +pub const DEFAULT_EXCLUDED_PREFIXES: &[&str] = &[ + "agent::", + "engine::", + "state::", + "stream::", + "iii.", +]; + +pub async fn discover_tools(iii: &III) -> Vec { + discover_tools_with(iii, DEFAULT_EXCLUDED_PREFIXES).await +} + +pub async fn discover_tools_with(iii: &III, excluded: &[&str]) -> Vec { + let functions = match iii.list_functions().await { + Ok(fns) => fns, + Err(e) => { + tracing::warn!(error = %e, "failed to discover functions"); + return Vec::new(); + } + }; + + let tools: Vec = functions + .into_iter() + .filter(|f| !f.function_id.is_empty()) + .filter(|f| !is_excluded(&f.function_id, excluded)) + .filter(|f| has_valid_schema(f)) + .map(|f| function_to_tool(&f)) + .filter(|t| !t.name.is_empty()) + .collect(); + + tracing::info!(count = tools.len(), "discovered tools"); + tools +} + +pub fn function_to_tool(f: &FunctionInfo) -> ToolDef { + ToolDef { + name: sanitize_tool_name(&f.function_id), + description: f.description.clone().unwrap_or_default(), + input_schema: f + .request_format + .clone() + .unwrap_or(json!({"type": "object", "properties": {}})), + } +} + +pub fn tool_name_to_function_id(tool_name: &str) -> String { + tool_name.replace("__", "::") +} + +pub fn sanitize_tool_name(function_id: &str) -> String { + let sanitized: String = function_id + .replace("::", "__") + .chars() + .map(|c| if c.is_ascii_alphanumeric() || c == '_' || c == '-' { c } else { '_' }) + .collect(); + if sanitized.len() > 128 { + sanitized[..128].to_string() + } else { + sanitized + } +} + +pub fn functions_to_tools(functions: &[FunctionInfo]) -> Vec { + functions + .iter() + .filter(|f| !is_excluded(&f.function_id, DEFAULT_EXCLUDED_PREFIXES)) + .map(|f| function_to_tool(f)) + .collect() +} + +pub fn build_capabilities_summary(tools: &[ToolDef]) -> String { + if tools.is_empty() { + return "No external functions are currently available.".to_string(); + } + + let mut summary = String::from("Available functions:\n"); + for tool in tools { + summary.push_str(&format!("- {}: {}\n", tool.name, tool.description)); + } + summary +} + +// Build the capabilities summary keyed by engine `function_id` (eval::metrics) +// rather than the Anthropic-sanitized tool name (eval__metrics). Use this +// when the model is asked to echo a function_id back — e.g., the planner +// fills `steps[*].function_id`, which downstream executors invoke as-is. +pub async fn build_planner_capabilities(iii: &III) -> String { + let functions = match iii.list_functions().await { + Ok(fns) => fns, + Err(_) => return "No external functions are currently available.".to_string(), + }; + + let eligible: Vec = functions + .into_iter() + .filter(|f| !f.function_id.is_empty()) + .filter(|f| !is_excluded(&f.function_id, DEFAULT_EXCLUDED_PREFIXES)) + .filter(has_valid_schema) + .collect(); + + if eligible.is_empty() { + return "No external functions are currently available.".to_string(); + } + + let mut summary = String::from("Available functions (call these exact ids):\n"); + for f in eligible { + let desc = f.description.as_deref().unwrap_or(""); + summary.push_str(&format!("- {}: {}\n", f.function_id, desc)); + } + summary +} + +fn is_excluded(function_id: &str, excluded: &[&str]) -> bool { + excluded.iter().any(|prefix| function_id.starts_with(prefix)) +} + +fn has_valid_schema(f: &FunctionInfo) -> bool { + match &f.request_format { + Some(schema) => { + schema.get("type").and_then(|t| t.as_str()) == Some("object") + } + None => true, + } +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn test_sanitize_colons() { + assert_eq!(sanitize_tool_name("eval::metrics"), "eval__metrics"); + } + + #[test] + fn test_sanitize_dots() { + assert_eq!(sanitize_tool_name("iii.on_functions.abc"), "iii_on_functions_abc"); + } + + #[test] + fn test_sanitize_uuid() { + let result = sanitize_tool_name("iii.callback.a1b2c3d4-e5f6"); + assert!(result.chars().all(|c| c.is_ascii_alphanumeric() || c == '_' || c == '-')); + } + + #[test] + fn test_sanitize_truncate() { + let long = "a".repeat(200); + assert_eq!(sanitize_tool_name(&long).len(), 128); + } + + #[test] + fn test_tool_name_to_function_id_roundtrip() { + assert_eq!(tool_name_to_function_id("eval__metrics"), "eval::metrics"); + } + + #[test] + fn test_worker_prefixes_not_excluded() { + assert!(!is_excluded("eval::metrics", DEFAULT_EXCLUDED_PREFIXES)); + assert!(!is_excluded("introspect::topology", DEFAULT_EXCLUDED_PREFIXES)); + assert!(!is_excluded("sensor::scan", DEFAULT_EXCLUDED_PREFIXES)); + assert!(!is_excluded("guardrails::check_input", DEFAULT_EXCLUDED_PREFIXES)); + assert!(!is_excluded("coding::scaffold", DEFAULT_EXCLUDED_PREFIXES)); + assert!(!is_excluded("experiment::create", DEFAULT_EXCLUDED_PREFIXES)); + assert!(!is_excluded("publish", DEFAULT_EXCLUDED_PREFIXES)); + } + + #[test] + fn test_infrastructure_prefixes_excluded() { + assert!(is_excluded("state::get", DEFAULT_EXCLUDED_PREFIXES)); + assert!(is_excluded("engine::health", DEFAULT_EXCLUDED_PREFIXES)); + assert!(is_excluded("stream::set", DEFAULT_EXCLUDED_PREFIXES)); + assert!(is_excluded("agent::chat", DEFAULT_EXCLUDED_PREFIXES)); + assert!(is_excluded("iii.on_functions_available.abc", DEFAULT_EXCLUDED_PREFIXES)); + } + + #[test] + fn test_has_valid_schema_with_object() { + let f = FunctionInfo { + function_id: "test".into(), + description: None, + request_format: Some(json!({"type": "object", "properties": {}})), + response_format: None, + metadata: None, + }; + assert!(has_valid_schema(&f)); + } + + #[test] + fn test_has_valid_schema_none() { + let f = FunctionInfo { + function_id: "test".into(), + description: None, + request_format: None, + response_format: None, + metadata: None, + }; + assert!(has_valid_schema(&f)); + } + + #[test] + fn test_has_invalid_schema() { + let f = FunctionInfo { + function_id: "test".into(), + description: None, + request_format: Some(json!({"type": "string"})), + response_format: None, + metadata: None, + }; + assert!(!has_valid_schema(&f)); + } + + #[test] + fn test_capabilities_summary_empty() { + let result = build_capabilities_summary(&[]); + assert!(result.contains("No external functions")); + } + + #[test] + fn test_capabilities_summary_with_tools() { + let tools = vec![ToolDef { + name: "eval__metrics".into(), + description: "Calculate metrics".into(), + input_schema: json!({}), + }]; + let result = build_capabilities_summary(&tools); + assert!(result.contains("eval__metrics")); + assert!(result.contains("Calculate metrics")); + } + + #[test] + fn test_system_prompt_contains_rules() { + let tools = vec![]; + let prompt = build_system_prompt(&tools); + assert!(prompt.contains("plain text")); + assert!(prompt.contains("markdown")); + assert!(prompt.contains("Do NOT wrap")); + } + + #[test] + fn test_function_to_tool() { + let f = FunctionInfo { + function_id: "eval::metrics".into(), + description: Some("Compute P50/P95/P99".into()), + request_format: Some(json!({"type": "object", "properties": {"function_id": {"type": "string"}}})), + response_format: None, + metadata: None, + }; + let tool = function_to_tool(&f); + assert_eq!(tool.name, "eval__metrics"); + assert_eq!(tool.description, "Compute P50/P95/P99"); + assert!(tool.input_schema.get("properties").is_some()); + } +} + +pub fn build_system_prompt(tools: &[ToolDef]) -> String { + let capabilities = build_capabilities_summary(tools); + + format!( + "You are the iii agent, an intelligent assistant for the iii engine.\n\ + \n\ + You have access to functions registered by connected workers. Use them to answer \ + questions about the system, analyze performance, and manage the engine.\n\ + \n\ + Rules:\n\ + - Call the available functions to gather real data before answering.\n\ + - Respond with plain text. Use markdown for formatting (tables, lists, code blocks).\n\ + - Be concise and data-driven.\n\ + - When showing data, use markdown tables.\n\ + - Do NOT wrap your response in JSON objects.\n\ + \n\ + {capabilities}" + ) +} diff --git a/agent/src/functions/chat.rs b/agent/src/functions/chat.rs new file mode 100644 index 000000000..3bb1dc908 --- /dev/null +++ b/agent/src/functions/chat.rs @@ -0,0 +1,237 @@ +use std::future::Future; +use std::pin::Pin; +use std::sync::Arc; + +use iii_sdk::{IIIError, TriggerRequest, III}; +use serde_json::{Value, json}; + +use crate::config::AgentConfig; +use crate::discovery; +use crate::llm::{ + ContentBlock, LlmClient, LlmRequest, Message, MessageContent, +}; +use crate::state; + +pub fn build_handler( + iii: III, + config: Arc, + llm: Arc, +) -> impl Fn(Value) -> Pin> + Send>> + + Send + + Sync + + 'static { + move |payload: Value| { + let iii = iii.clone(); + let config = config.clone(); + let llm = llm.clone(); + + Box::pin(async move { handle_chat(iii, config, llm, payload).await }) + } +} + +async fn handle_chat( + iii: III, + config: Arc, + llm: Arc, + payload: Value, +) -> Result { + let session_id = payload + .get("session_id") + .and_then(|v| v.as_str()) + .unwrap_or("") + .to_string(); + + let user_message = payload + .get("message") + .and_then(|v| v.as_str()) + .ok_or_else(|| IIIError::Handler("missing 'message' field".to_string()))? + .to_string(); + + let tools = discovery::discover_tools(&iii).await; + let system_prompt = discovery::build_system_prompt(&tools); + + let mut messages = load_history(&iii, &session_id).await; + + messages.push(Message { + role: "user".to_string(), + content: MessageContent::Text(user_message), + }); + + let mut iterations = 0u32; + let max_iterations = config.max_iterations; + + loop { + if iterations >= max_iterations { + break; + } + iterations += 1; + + let request = LlmRequest { + model: config.anthropic_model.clone(), + max_tokens: config.max_tokens, + system: system_prompt.clone(), + messages: messages.clone(), + tools: if tools.is_empty() { + None + } else { + Some(tools.clone()) + }, + }; + + let response = llm + .send(&request) + .await + .map_err(|e| IIIError::Handler(format!("LLM request failed: {}", e)))?; + + let tool_uses = LlmClient::extract_tool_uses(&response); + + if tool_uses.is_empty() { + let text = LlmClient::extract_text(&response); + + messages.push(Message { + role: "assistant".to_string(), + content: MessageContent::Text(text.clone()), + }); + + save_history(&iii, &session_id, &messages).await; + + return Ok(build_response(&text, &response)); + } + + let mut assistant_blocks: Vec = Vec::new(); + for block in &response.content { + assistant_blocks.push(block.clone()); + } + + messages.push(Message { + role: "assistant".to_string(), + content: MessageContent::Blocks(assistant_blocks), + }); + + let mut tool_result_blocks: Vec = Vec::new(); + + for tool_use in &tool_uses { + let function_id = discovery::tool_name_to_function_id(&tool_use.name); + + let result = execute_tool(&iii, &function_id, &tool_use.input).await; + + let (content, is_error) = match result { + Ok(val) => (serde_json::to_string(&val).unwrap_or_default(), None), + Err(e) => (format!("Error: {}", e), Some(true)), + }; + + tool_result_blocks.push(ContentBlock::ToolResult { + tool_use_id: tool_use.id.clone(), + content, + is_error, + }); + } + + messages.push(Message { + role: "user".to_string(), + content: MessageContent::Blocks(tool_result_blocks), + }); + } + + let text = "Reached maximum iterations without a final response.".to_string(); + save_history(&iii, &session_id, &messages).await; + + Ok(json!({ + "elements": [{"type": "text", "content": text}], + "session_id": session_id, + "iterations": iterations + })) +} + +async fn execute_tool(iii: &III, function_id: &str, input: &Value) -> Result { + iii.trigger(TriggerRequest { + function_id: function_id.to_string(), + payload: input.clone(), + action: None, + timeout_ms: Some(30000), + }) + .await +} + +fn build_response(text: &str, response: &crate::llm::LlmResponse) -> Value { + let elements = parse_ui_elements(text); + + let mut result = json!({ + "elements": elements, + }); + + if let Some(usage) = &response.usage { + result["usage"] = json!({ + "input_tokens": usage.input_tokens, + "output_tokens": usage.output_tokens + }); + } + + result +} + +fn parse_ui_elements(text: &str) -> Vec { + if let Ok(parsed) = serde_json::from_str::(text) { + if parsed.is_array() { + if let Some(arr) = parsed.as_array() { + return arr.clone(); + } + } + if parsed.get("type").is_some() { + return vec![parsed]; + } + } + + vec![json!({"type": "text", "content": text})] +} + +// Session records are stored as { "created_at": , "messages": [...] }. +// Read the messages array out of that envelope rather than treating the +// whole record as Vec — preserves `created_at` so cleanup's +// age-based expiry and history tools keep working across turns. +async fn load_history(iii: &III, session_id: &str) -> Vec { + if session_id.is_empty() { + return Vec::new(); + } + + match state::state_get(iii, "agent:sessions", session_id).await { + Ok(val) => { + let inner = val.get("value").unwrap_or(&val); + if let Some(arr) = inner.get("messages") { + serde_json::from_value::>(arr.clone()).unwrap_or_default() + } else if inner.is_array() { + // Legacy shape: the value was saved as a bare messages array + // by an earlier version. Load it so old sessions still open. + serde_json::from_value::>(inner.clone()).unwrap_or_default() + } else { + Vec::new() + } + } + Err(_) => Vec::new(), + } +} + +// Preserve (or mint) created_at so session_cleanup's TTL still fires and +// session_history's response shape stays stable after every turn. +async fn save_history(iii: &III, session_id: &str, messages: &[Message]) { + if session_id.is_empty() { + return; + } + + let created_at = match state::state_get(iii, "agent:sessions", session_id).await { + Ok(val) => val + .get("value") + .and_then(|v| v.get("created_at")) + .and_then(|v| v.as_str()) + .map(|s| s.to_string()) + .unwrap_or_else(|| chrono::Utc::now().to_rfc3339()), + Err(_) => chrono::Utc::now().to_rfc3339(), + }; + + let value = json!({ + "created_at": created_at, + "messages": serde_json::to_value(messages).unwrap_or(json!([])), + }); + let _ = state::state_set(iii, "agent:sessions", session_id, &value).await; +} + diff --git a/agent/src/functions/chat_stream.rs b/agent/src/functions/chat_stream.rs new file mode 100644 index 000000000..8d3a1b45e --- /dev/null +++ b/agent/src/functions/chat_stream.rs @@ -0,0 +1,349 @@ +use std::future::Future; +use std::pin::Pin; +use std::sync::Arc; + +use futures_util::StreamExt; +use iii_sdk::{IIIError, TriggerRequest, III}; +use serde_json::{Value, json}; +use uuid::Uuid; + +use crate::config::AgentConfig; +use crate::discovery; +use crate::llm::{ + ContentBlock, LlmClient, LlmRequest, Message, MessageContent, StreamEvent, +}; +use crate::state; + +pub fn build_handler( + iii: III, + config: Arc, + llm: Arc, +) -> impl Fn(Value) -> Pin> + Send>> + + Send + + Sync + + 'static { + move |payload: Value| { + let iii = iii.clone(); + let config = config.clone(); + let llm = llm.clone(); + + Box::pin(async move { handle_chat_stream(iii, config, llm, payload).await }) + } +} + +async fn handle_chat_stream( + iii: III, + config: Arc, + llm: Arc, + payload: Value, +) -> Result { + // Mint a fresh session_id when the caller omits one. Without this, + // every session-less request writes events into the shared + // `agent:events:` group, so concurrent callers interleave each + // other's streamed output. + let session_id = payload + .get("session_id") + .and_then(|v| v.as_str()) + .filter(|s| !s.is_empty()) + .map(|s| s.to_string()) + .unwrap_or_else(|| Uuid::new_v4().to_string()); + + let user_message = payload + .get("message") + .and_then(|v| v.as_str()) + .ok_or_else(|| IIIError::Handler("missing 'message' field".to_string()))? + .to_string(); + + let stream_group = format!("agent:events:{}", session_id); + + let tools = discovery::discover_tools(&iii).await; + let system_prompt = discovery::build_system_prompt(&tools); + + let mut messages = load_history(&iii, &session_id).await; + + messages.push(Message { + role: "user".to_string(), + content: MessageContent::Text(user_message), + }); + + let mut iterations = 0u32; + let max_iterations = config.max_iterations; + let mut full_text = String::new(); + + loop { + if iterations >= max_iterations { + break; + } + iterations += 1; + + let request = LlmRequest { + model: config.anthropic_model.clone(), + max_tokens: config.max_tokens, + system: system_prompt.clone(), + messages: messages.clone(), + tools: if tools.is_empty() { + None + } else { + Some(tools.clone()) + }, + }; + + let stream_result = llm.send_stream(&request).await; + + let mut event_stream = match stream_result { + Ok(s) => s, + Err(e) => { + emit_event(&iii, &stream_group, &json!({ + "type": "error", + "message": format!("LLM stream failed: {}", e) + })) + .await; + return Err(IIIError::Handler(format!("LLM stream failed: {}", e))); + } + }; + + let mut current_tool_name = String::new(); + let mut current_tool_id = String::new(); + let mut current_tool_input_json = String::new(); + let mut collected_tool_uses: Vec<(String, String, Value)> = Vec::new(); + let mut has_tool_use = false; + + while let Some(event_result) = event_stream.next().await { + match event_result { + Ok(event) => { + process_stream_event( + &iii, + &stream_group, + &event, + &mut full_text, + &mut current_tool_name, + &mut current_tool_id, + &mut current_tool_input_json, + &mut collected_tool_uses, + &mut has_tool_use, + ) + .await; + } + Err(e) => { + tracing::warn!(error = %e, "stream event error"); + } + } + } + + if !current_tool_id.is_empty() { + let input: Value = + serde_json::from_str(¤t_tool_input_json).unwrap_or(json!({})); + collected_tool_uses.push(( + current_tool_id.clone(), + current_tool_name.clone(), + input, + )); + } + + if !has_tool_use { + messages.push(Message { + role: "assistant".to_string(), + content: MessageContent::Text(full_text.clone()), + }); + + save_history(&iii, &session_id, &messages).await; + + emit_event(&iii, &stream_group, &json!({"type": "done"})).await; + + return Ok(json!({ + "stream_group": stream_group, + "session_id": session_id, + "iterations": iterations + })); + } + + let mut assistant_blocks: Vec = Vec::new(); + if !full_text.is_empty() { + assistant_blocks.push(ContentBlock::Text { + text: full_text.clone(), + }); + } + for (tool_id, tool_name, tool_input) in &collected_tool_uses { + assistant_blocks.push(ContentBlock::ToolUse { + id: tool_id.clone(), + name: tool_name.clone(), + input: tool_input.clone(), + }); + } + + messages.push(Message { + role: "assistant".to_string(), + content: MessageContent::Blocks(assistant_blocks), + }); + + let mut tool_result_blocks: Vec = Vec::new(); + + for (tool_id, tool_name, tool_input) in &collected_tool_uses { + let function_id = discovery::tool_name_to_function_id(tool_name); + + emit_event(&iii, &stream_group, &json!({ + "type": "tool_use", + "name": function_id, + "input": tool_input + })) + .await; + + let result = iii + .trigger(TriggerRequest { + function_id: function_id.clone(), + payload: tool_input.clone(), + action: None, + timeout_ms: Some(30000), + }) + .await; + + let (content, is_error) = match &result { + Ok(val) => (serde_json::to_string(val).unwrap_or_default(), None), + Err(e) => (format!("Error: {}", e), Some(true)), + }; + + emit_event(&iii, &stream_group, &json!({ + "type": "tool_result", + "name": function_id, + "result": match &result { + Ok(v) => v.clone(), + Err(e) => json!({"error": e.to_string()}) + } + })) + .await; + + tool_result_blocks.push(ContentBlock::ToolResult { + tool_use_id: tool_id.clone(), + content, + is_error, + }); + } + + messages.push(Message { + role: "user".to_string(), + content: MessageContent::Blocks(tool_result_blocks), + }); + + full_text.clear(); + } + + emit_event(&iii, &stream_group, &json!({"type": "done"})).await; + save_history(&iii, &session_id, &messages).await; + + Ok(json!({ + "stream_group": stream_group, + "session_id": session_id, + "iterations": iterations + })) +} + +async fn process_stream_event( + iii: &III, + stream_group: &str, + event: &StreamEvent, + full_text: &mut String, + current_tool_name: &mut String, + current_tool_id: &mut String, + current_tool_input_json: &mut String, + collected_tool_uses: &mut Vec<(String, String, Value)>, + has_tool_use: &mut bool, +) { + match event.event_type.as_str() { + "content_block_start" => { + if let Some(ContentBlock::ToolUse { id, name, .. }) = &event.content_block { + *current_tool_id = id.clone(); + *current_tool_name = name.clone(); + current_tool_input_json.clear(); + *has_tool_use = true; + } + } + "content_block_delta" => { + if let Some(delta) = &event.delta { + if let Some(text) = &delta.text { + full_text.push_str(text); + emit_event( + iii, + stream_group, + &json!({"type": "text_delta", "text": text}), + ) + .await; + } + if let Some(partial_json) = &delta.partial_json { + current_tool_input_json.push_str(partial_json); + } + } + } + "content_block_stop" => { + if !current_tool_id.is_empty() { + let input: Value = + serde_json::from_str(current_tool_input_json).unwrap_or(json!({})); + collected_tool_uses.push(( + current_tool_id.clone(), + current_tool_name.clone(), + input, + )); + current_tool_id.clear(); + current_tool_name.clear(); + current_tool_input_json.clear(); + } + } + _ => {} + } +} + +async fn emit_event(iii: &III, stream_group: &str, event: &Value) { + let _ = iii + .trigger(TriggerRequest { + function_id: "stream::set".to_string(), + payload: json!({ + "scope": stream_group, + "key": uuid::Uuid::new_v4().to_string(), + "value": event + }), + action: None, + timeout_ms: Some(5000), + }) + .await; +} + +async fn load_history(iii: &III, session_id: &str) -> Vec { + if session_id.is_empty() { + return Vec::new(); + } + + match state::state_get(iii, "agent:sessions", session_id).await { + Ok(val) => { + let inner = val.get("value").unwrap_or(&val); + if let Some(arr) = inner.get("messages") { + serde_json::from_value::>(arr.clone()).unwrap_or_default() + } else if inner.is_array() { + serde_json::from_value::>(inner.clone()).unwrap_or_default() + } else { + Vec::new() + } + } + Err(_) => Vec::new(), + } +} + +async fn save_history(iii: &III, session_id: &str, messages: &[Message]) { + if session_id.is_empty() { + return; + } + + let created_at = match state::state_get(iii, "agent:sessions", session_id).await { + Ok(val) => val + .get("value") + .and_then(|v| v.get("created_at")) + .and_then(|v| v.as_str()) + .map(|s| s.to_string()) + .unwrap_or_else(|| chrono::Utc::now().to_rfc3339()), + Err(_) => chrono::Utc::now().to_rfc3339(), + }; + + let value = json!({ + "created_at": created_at, + "messages": serde_json::to_value(messages).unwrap_or(json!([])), + }); + let _ = state::state_set(iii, "agent:sessions", session_id, &value).await; +} diff --git a/agent/src/functions/discover.rs b/agent/src/functions/discover.rs new file mode 100644 index 000000000..052942708 --- /dev/null +++ b/agent/src/functions/discover.rs @@ -0,0 +1,38 @@ +use std::future::Future; +use std::pin::Pin; + +use iii_sdk::{IIIError, III}; +use serde_json::{Value, json}; + +use crate::discovery; + +pub fn build_handler( + iii: III, +) -> impl Fn(Value) -> Pin> + Send>> + + Send + + Sync + + 'static { + move |_payload: Value| { + let iii = iii.clone(); + + Box::pin(async move { + let tools = discovery::discover_tools(&iii).await; + + let functions: Vec = tools + .iter() + .map(|t| { + json!({ + "name": t.name, + "description": t.description, + "input_schema": t.input_schema + }) + }) + .collect(); + + Ok(json!({ + "functions": functions, + "count": functions.len() + })) + }) + } +} diff --git a/agent/src/functions/mod.rs b/agent/src/functions/mod.rs new file mode 100644 index 000000000..15970044b --- /dev/null +++ b/agent/src/functions/mod.rs @@ -0,0 +1,5 @@ +pub mod chat; +pub mod chat_stream; +pub mod discover; +pub mod plan; +pub mod session; diff --git a/agent/src/functions/plan.rs b/agent/src/functions/plan.rs new file mode 100644 index 000000000..9c4e22a1f --- /dev/null +++ b/agent/src/functions/plan.rs @@ -0,0 +1,90 @@ +use std::future::Future; +use std::pin::Pin; +use std::sync::Arc; + +use iii_sdk::{IIIError, III}; +use serde_json::{Value, json}; + +use crate::config::AgentConfig; +use crate::discovery; +use crate::llm::{LlmClient, LlmRequest, Message, MessageContent}; + +pub fn build_handler( + iii: III, + config: Arc, + llm: Arc, +) -> impl Fn(Value) -> Pin> + Send>> + + Send + + Sync + + 'static { + move |payload: Value| { + let iii = iii.clone(); + let config = config.clone(); + let llm = llm.clone(); + + Box::pin(async move { handle_plan(iii, config, llm, payload).await }) + } +} + +async fn handle_plan( + iii: III, + config: Arc, + llm: Arc, + payload: Value, +) -> Result { + let query = payload + .get("query") + .and_then(|v| v.as_str()) + .ok_or_else(|| IIIError::Handler("missing 'query' field".to_string()))? + .to_string(); + + // Planner output is fed to downstream executors via function_id, so the + // capabilities block must name engine ids (eval::metrics) not the + // sanitized tool names (eval__metrics) the chat handler uses. + let capabilities = discovery::build_planner_capabilities(&iii).await; + + let system = format!( + "You are a planning agent for the iii engine. Given a user query, generate an execution \ + plan as a DAG of iii.trigger() calls. Do NOT execute anything.\n\ + \n\ + Return a JSON object with:\n\ + - \"steps\": an array of step objects, each with:\n\ + - \"id\": unique step identifier (e.g. \"step_1\")\n\ + - \"function_id\": the function to call\n\ + - \"payload\": the payload object\n\ + - \"depends_on\": array of step IDs this step depends on (empty if root)\n\ + - \"description\": human-readable description of what this step does\n\ + - \"summary\": a brief description of the overall plan\n\ + \n\ + {capabilities}" + ); + + let messages = vec![Message { + role: "user".to_string(), + content: MessageContent::Text(query), + }]; + + let request = LlmRequest { + model: config.anthropic_model.clone(), + max_tokens: config.max_tokens, + system, + messages, + tools: None, + }; + + let response = llm + .send(&request) + .await + .map_err(|e| IIIError::Handler(format!("LLM request failed: {}", e)))?; + + let text = LlmClient::extract_text(&response); + + let plan: Value = serde_json::from_str(&text).unwrap_or_else(|_| { + json!({ + "steps": [], + "summary": text + }) + }); + + Ok(plan) +} diff --git a/agent/src/functions/session.rs b/agent/src/functions/session.rs new file mode 100644 index 000000000..96d88cdb0 --- /dev/null +++ b/agent/src/functions/session.rs @@ -0,0 +1,132 @@ +use std::future::Future; +use std::pin::Pin; + +use iii_sdk::{IIIError, III}; +use serde_json::{Value, json}; +use uuid::Uuid; + +use crate::state; + +pub fn build_create_handler( + iii: III, +) -> impl Fn(Value) -> Pin> + Send>> + + Send + + Sync + + 'static { + move |_payload: Value| { + let iii = iii.clone(); + + Box::pin(async move { + let session_id = Uuid::new_v4().to_string(); + let now = chrono::Utc::now().to_rfc3339(); + + let session_data = json!({ + "created_at": now, + "messages": [] + }); + + state::state_set(&iii, "agent:sessions", &session_id, &session_data) + .await + .map_err(|e| IIIError::Handler(format!("failed to create session: {}", e)))?; + + Ok(json!({ + "session_id": session_id, + "created_at": now + })) + }) + } +} + +pub fn build_history_handler( + iii: III, +) -> impl Fn(Value) -> Pin> + Send>> + + Send + + Sync + + 'static { + move |payload: Value| { + let iii = iii.clone(); + + Box::pin(async move { + let session_id = payload + .get("session_id") + .and_then(|v| v.as_str()) + .ok_or_else(|| IIIError::Handler("missing 'session_id' field".to_string()))? + .to_string(); + + let result = state::state_get(&iii, "agent:sessions", &session_id).await; + + match result { + Ok(val) => Ok(json!({ + "session_id": session_id, + "history": val.get("value").cloned().unwrap_or(json!(null)) + })), + Err(_) => Ok(json!({ + "session_id": session_id, + "history": null, + "error": "session not found" + })), + } + }) + } +} + +pub fn build_cleanup_handler( + iii: III, + ttl_hours: u64, +) -> impl Fn(Value) -> Pin> + Send>> + + Send + + Sync + + 'static { + move |_payload: Value| { + let iii = iii.clone(); + let ttl_hours = ttl_hours as i64; + + Box::pin(async move { + let sessions = state::state_list(&iii, "agent:sessions").await; + + let mut cleaned = 0u64; + + if let Ok(list) = sessions { + if let Some(keys) = list.get("keys").and_then(|v| v.as_array()) { + let now = chrono::Utc::now(); + + for key_val in keys { + if let Some(key) = key_val.as_str() { + if let Ok(session) = + state::state_get(&iii, "agent:sessions", key).await + { + let should_delete = session + .get("value") + .and_then(|v| v.get("created_at")) + .and_then(|v| v.as_str()) + .and_then(|ts| { + chrono::DateTime::parse_from_rfc3339(ts).ok() + }) + .map(|created| { + let age = now + .signed_duration_since(created.with_timezone(&chrono::Utc)); + age.num_hours() > ttl_hours + }) + .unwrap_or(false); + + if should_delete { + let _ = state::state_delete( + &iii, + "agent:sessions", + key, + ) + .await; + cleaned += 1; + } + } + } + } + } + } + + Ok(json!({ + "cleaned_sessions": cleaned + })) + }) + } +} diff --git a/agent/src/llm.rs b/agent/src/llm.rs new file mode 100644 index 000000000..ea9963d82 --- /dev/null +++ b/agent/src/llm.rs @@ -0,0 +1,219 @@ +use anyhow::{Result, anyhow}; +use futures_util::StreamExt; +use reqwest::Client; +use serde::{Deserialize, Serialize}; +use serde_json::Value; + +const ANTHROPIC_API_URL: &str = "https://api.anthropic.com/v1/messages"; +const ANTHROPIC_VERSION: &str = "2023-06-01"; + +#[derive(Debug, Clone, Serialize)] +pub struct LlmRequest { + pub model: String, + pub max_tokens: u32, + pub system: String, + pub messages: Vec, + #[serde(skip_serializing_if = "Option::is_none")] + pub tools: Option>, +} + +#[derive(Debug, Clone, Serialize, Deserialize)] +pub struct Message { + pub role: String, + pub content: MessageContent, +} + +#[derive(Debug, Clone, Serialize, Deserialize)] +#[serde(untagged)] +pub enum MessageContent { + Text(String), + Blocks(Vec), +} + +#[derive(Debug, Clone, Serialize, Deserialize)] +#[serde(tag = "type")] +pub enum ContentBlock { + #[serde(rename = "text")] + Text { text: String }, + #[serde(rename = "tool_use")] + ToolUse { + id: String, + name: String, + input: Value, + }, + #[serde(rename = "tool_result")] + ToolResult { + tool_use_id: String, + content: String, + #[serde(skip_serializing_if = "Option::is_none")] + is_error: Option, + }, +} + +#[derive(Debug, Clone, Serialize, Deserialize)] +pub struct ToolDef { + pub name: String, + pub description: String, + pub input_schema: Value, +} + +#[derive(Debug, Clone, Deserialize)] +pub struct LlmResponse { + pub content: Vec, + #[allow(dead_code)] + pub stop_reason: Option, + pub usage: Option, +} + +#[derive(Debug, Clone, Deserialize)] +pub struct Usage { + pub input_tokens: Option, + pub output_tokens: Option, +} + +#[derive(Debug, Clone)] +pub struct ToolUse { + pub id: String, + pub name: String, + pub input: Value, +} + +#[derive(Debug, Clone, Deserialize)] +pub struct StreamEvent { + #[serde(rename = "type")] + pub event_type: String, + #[serde(default)] + pub delta: Option, + #[serde(default)] + pub content_block: Option, + #[serde(default)] + #[allow(dead_code)] + pub index: Option, +} + +#[derive(Debug, Clone, Deserialize)] +pub struct StreamDelta { + #[serde(rename = "type")] + #[allow(dead_code)] + pub delta_type: Option, + pub text: Option, + pub partial_json: Option, +} + +pub struct LlmClient { + client: Client, + api_key: String, +} + +impl LlmClient { + pub fn new(api_key: String) -> Self { + Self { + client: Client::new(), + api_key, + } + } + + pub fn from_env() -> Result { + let api_key = std::env::var("ANTHROPIC_API_KEY") + .map_err(|_| anyhow!("ANTHROPIC_API_KEY environment variable not set"))?; + Ok(Self::new(api_key)) + } + + pub async fn send(&self, request: &LlmRequest) -> Result { + let response = self + .client + .post(ANTHROPIC_API_URL) + .header("x-api-key", &self.api_key) + .header("anthropic-version", ANTHROPIC_VERSION) + .header("content-type", "application/json") + .json(request) + .send() + .await?; + + let status = response.status(); + if !status.is_success() { + let body = response.text().await.unwrap_or_default(); + return Err(anyhow!("Anthropic API error {}: {}", status, body)); + } + + let llm_response: LlmResponse = response.json().await?; + Ok(llm_response) + } + + pub async fn send_stream( + &self, + request: &LlmRequest, + ) -> Result>> { + let mut stream_request = serde_json::to_value(request)?; + stream_request["stream"] = serde_json::Value::Bool(true); + + let response = self + .client + .post(ANTHROPIC_API_URL) + .header("x-api-key", &self.api_key) + .header("anthropic-version", ANTHROPIC_VERSION) + .header("content-type", "application/json") + .json(&stream_request) + .send() + .await?; + + let status = response.status(); + if !status.is_success() { + let body = response.text().await.unwrap_or_default(); + return Err(anyhow!("Anthropic API error {}: {}", status, body)); + } + + let byte_stream = response.bytes_stream(); + + let event_stream = byte_stream.map(|chunk_result| { + let chunk = chunk_result.map_err(|e| anyhow!("stream read error: {}", e))?; + let text = String::from_utf8_lossy(&chunk); + let mut events = Vec::new(); + + for line in text.lines() { + if let Some(data) = line.strip_prefix("data: ") { + if data == "[DONE]" { + continue; + } + if let Ok(event) = serde_json::from_str::(data) { + events.push(event); + } + } + } + + Ok(events) + }); + + Ok(event_stream.flat_map(|result| { + let items: Vec> = match result { + Ok(events) => events.into_iter().map(Ok).collect(), + Err(e) => vec![Err(e)], + }; + futures_util::stream::iter(items) + })) + } + + pub fn extract_text(response: &LlmResponse) -> String { + let mut text = String::new(); + for block in &response.content { + if let ContentBlock::Text { text: t } = block { + text.push_str(t); + } + } + text + } + + pub fn extract_tool_uses(response: &LlmResponse) -> Vec { + let mut tool_uses = Vec::new(); + for block in &response.content { + if let ContentBlock::ToolUse { id, name, input } = block { + tool_uses.push(ToolUse { + id: id.clone(), + name: name.clone(), + input: input.clone(), + }); + } + } + tool_uses + } +} diff --git a/agent/src/main.rs b/agent/src/main.rs new file mode 100644 index 000000000..66c36f327 --- /dev/null +++ b/agent/src/main.rs @@ -0,0 +1,383 @@ +use std::sync::Arc; + +use anyhow::Result; +use clap::Parser; +use iii_sdk::{ + register_worker, InitOptions, OtelConfig, RegisterFunctionMessage, RegisterTriggerInput, +}; +use serde_json::json; + +mod config; +mod discovery; +mod functions; +mod llm; +mod manifest; +mod state; + +#[derive(Parser, Debug)] +#[command(name = "iii-agent", about = "III engine AI agent — chat orchestrator")] +struct Cli { + #[arg(long, default_value = "./config.yaml")] + config: String, + + #[arg(long, default_value = "ws://127.0.0.1:49134")] + url: String, + + #[arg(long)] + manifest: bool, + + #[arg(long, env = "ANTHROPIC_API_KEY")] + api_key: Option, +} + +#[tokio::main] +async fn main() -> Result<()> { + tracing_subscriber::fmt() + .with_env_filter( + tracing_subscriber::EnvFilter::try_from_default_env() + .unwrap_or_else(|_| tracing_subscriber::EnvFilter::new("info")), + ) + .init(); + + let cli = Cli::parse(); + + if cli.manifest { + let manifest = manifest::build_manifest(); + println!("{}", serde_json::to_string_pretty(&manifest).unwrap()); + return Ok(()); + } + + let agent_config = match config::load_config(&cli.config) { + Ok(c) => { + tracing::info!( + model = %c.anthropic_model, + max_iterations = c.max_iterations, + "loaded config from {}", + cli.config + ); + c + } + Err(e) => { + tracing::warn!(error = %e, path = %cli.config, "failed to load config, using defaults"); + config::AgentConfig::default() + } + }; + + let config = Arc::new(agent_config); + + let llm_client = if let Some(ref key) = cli.api_key { + llm::LlmClient::new(key.clone()) + } else { + llm::LlmClient::from_env()? + }; + let llm = Arc::new(llm_client); + + tracing::info!(url = %cli.url, "connecting to III engine"); + + let iii = register_worker( + &cli.url, + InitOptions { + otel: Some(OtelConfig::default()), + ..Default::default() + }, + ); + + let chat_handler = functions::chat::build_handler(iii.clone(), config.clone(), llm.clone()); + let _chat_fn = iii.register_function_with( + RegisterFunctionMessage { + id: "agent::chat".to_string(), + description: Some( + "Send a message to the AI agent and get a structured response".to_string(), + ), + request_format: Some(json!({ + "type": "object", + "properties": { + "session_id": { + "type": "string", + "description": "Session ID for conversation continuity" + }, + "message": { + "type": "string", + "description": "User message to the agent" + } + }, + "required": ["message"] + })), + response_format: Some(json!({ + "type": "object", + "properties": { + "elements": { + "type": "array", + "description": "JSON-UI elements for console rendering" + }, + "usage": { + "type": "object", + "properties": { + "input_tokens": { "type": "integer" }, + "output_tokens": { "type": "integer" } + } + } + } + })), + metadata: None, + invocation: None, + }, + chat_handler, + ); + + let chat_stream_handler = + functions::chat_stream::build_handler(iii.clone(), config.clone(), llm.clone()); + let _chat_stream_fn = iii.register_function_with( + RegisterFunctionMessage { + id: "agent::chat_stream".to_string(), + description: Some( + "Send a message to the AI agent with streaming response via iii Streams" + .to_string(), + ), + request_format: Some(json!({ + "type": "object", + "properties": { + "session_id": { + "type": "string", + "description": "Session ID for conversation continuity" + }, + "message": { + "type": "string", + "description": "User message to the agent" + } + }, + "required": ["message"] + })), + response_format: Some(json!({ + "type": "object", + "properties": { + "stream_group": { "type": "string" }, + "session_id": { "type": "string" }, + "iterations": { "type": "integer" } + } + })), + metadata: None, + invocation: None, + }, + chat_stream_handler, + ); + + let discover_handler = functions::discover::build_handler(iii.clone()); + let _discover_fn = iii.register_function_with( + RegisterFunctionMessage { + id: "agent::discover".to_string(), + description: Some( + "List all available functions that the agent can orchestrate".to_string(), + ), + request_format: Some(json!({ + "type": "object", + "properties": {} + })), + response_format: Some(json!({ + "type": "object", + "properties": { + "functions": { "type": "array" }, + "count": { "type": "integer" } + } + })), + metadata: None, + invocation: None, + }, + discover_handler, + ); + + let plan_handler = functions::plan::build_handler(iii.clone(), config.clone(), llm.clone()); + let _plan_fn = iii.register_function_with( + RegisterFunctionMessage { + id: "agent::plan".to_string(), + description: Some( + "Generate an execution plan DAG from a query without executing".to_string(), + ), + request_format: Some(json!({ + "type": "object", + "properties": { + "query": { + "type": "string", + "description": "The query to generate a plan for" + } + }, + "required": ["query"] + })), + response_format: Some(json!({ + "type": "object", + "properties": { + "steps": { "type": "array" }, + "summary": { "type": "string" } + } + })), + metadata: None, + invocation: None, + }, + plan_handler, + ); + + let session_create_handler = functions::session::build_create_handler(iii.clone()); + let _session_create_fn = iii.register_function_with( + RegisterFunctionMessage { + id: "agent::session_create".to_string(), + description: Some("Create a new chat session".to_string()), + request_format: Some(json!({ + "type": "object", + "properties": {} + })), + response_format: Some(json!({ + "type": "object", + "properties": { + "session_id": { "type": "string" }, + "created_at": { "type": "string" } + } + })), + metadata: None, + invocation: None, + }, + session_create_handler, + ); + + let session_history_handler = functions::session::build_history_handler(iii.clone()); + let _session_history_fn = iii.register_function_with( + RegisterFunctionMessage { + id: "agent::session_history".to_string(), + description: Some("Retrieve conversation history for a session".to_string()), + request_format: Some(json!({ + "type": "object", + "properties": { + "session_id": { + "type": "string", + "description": "Session ID to retrieve history for" + } + }, + "required": ["session_id"] + })), + response_format: Some(json!({ + "type": "object", + "properties": { + "session_id": { "type": "string" }, + "history": {} + } + })), + metadata: None, + invocation: None, + }, + session_history_handler, + ); + + let _http_chat = iii.register_trigger(RegisterTriggerInput { + trigger_type: "http".to_string(), + function_id: "agent::chat".to_string(), + config: json!({ + "api_path": "agent/chat", + "http_method": "POST" + }), + metadata: None, + }); + + let _http_discover = iii.register_trigger(RegisterTriggerInput { + trigger_type: "http".to_string(), + function_id: "agent::discover".to_string(), + config: json!({ + "api_path": "agent/discover", + "http_method": "GET" + }), + metadata: None, + }); + + let _http_chat_stream = iii.register_trigger(RegisterTriggerInput { + trigger_type: "http".to_string(), + function_id: "agent::chat_stream".to_string(), + config: json!({ + "api_path": "agent/chat/stream", + "http_method": "POST" + }), + metadata: None, + }); + + let _http_plan = iii.register_trigger(RegisterTriggerInput { + trigger_type: "http".to_string(), + function_id: "agent::plan".to_string(), + config: json!({ + "api_path": "agent/plan", + "http_method": "POST" + }), + metadata: None, + }); + + let _http_session_create = iii.register_trigger(RegisterTriggerInput { + trigger_type: "http".to_string(), + function_id: "agent::session_create".to_string(), + config: json!({ + "api_path": "agent/session", + "http_method": "POST" + }), + metadata: None, + }); + + let _http_session_history = iii.register_trigger(RegisterTriggerInput { + trigger_type: "http".to_string(), + function_id: "agent::session_history".to_string(), + config: json!({ + "api_path": "agent/session/history", + "http_method": "POST" + }), + metadata: None, + }); + + let iii_for_refresh = iii.clone(); + let _functions_guard = iii.on_functions_available(move |functions| { + let tools = discovery::functions_to_tools(&functions); + let tools_json = serde_json::to_value(&tools).unwrap_or(json!([])); + + let iii_inner = iii_for_refresh.clone(); + tokio::spawn(async move { + let _ = state::state_set( + &iii_inner, + "agent:tools", + "cached", + &tools_json, + ) + .await; + tracing::info!(count = tools_json.as_array().map(|a| a.len()).unwrap_or(0), "tool cache refreshed"); + }); + }); + + let session_cleanup_handler = + functions::session::build_cleanup_handler(iii.clone(), config.session_ttl_hours); + let _cleanup_fn = iii.register_function_with( + RegisterFunctionMessage { + id: "agent::session_cleanup".to_string(), + description: Some("Clean up expired sessions".to_string()), + request_format: Some(json!({"type": "object", "properties": {}})), + response_format: Some(json!({ + "type": "object", + "properties": { + "cleaned_sessions": { "type": "integer" } + } + })), + metadata: None, + invocation: None, + }, + session_cleanup_handler, + ); + + let _cron_cleanup = iii.register_trigger(RegisterTriggerInput { + trigger_type: "cron".to_string(), + function_id: "agent::session_cleanup".to_string(), + config: json!({ + "cron_expression": config.cron_session_cleanup + }), + metadata: None, + }); + + tracing::info!("iii-agent registered 7 functions, 6 HTTP triggers, 1 cron trigger, 1 subscribe trigger"); + + tokio::signal::ctrl_c().await?; + + tracing::info!("iii-agent shutting down"); + iii.shutdown_async().await; + + Ok(()) +} diff --git a/agent/src/manifest.rs b/agent/src/manifest.rs new file mode 100644 index 000000000..4cdee521e --- /dev/null +++ b/agent/src/manifest.rs @@ -0,0 +1,50 @@ +use serde::Serialize; + +#[derive(Serialize)] +pub struct ModuleManifest { + pub name: String, + pub version: String, + pub description: String, + pub default_config: serde_json::Value, + pub supported_targets: Vec, +} + +pub fn build_manifest() -> ModuleManifest { + ModuleManifest { + name: env!("CARGO_PKG_NAME").to_string(), + version: env!("CARGO_PKG_VERSION").to_string(), + description: "III engine AI agent — chat orchestrator with dynamic function discovery" + .to_string(), + default_config: serde_json::json!({ + "anthropic_model": "claude-sonnet-4-20250514", + "max_tokens": 4096, + "max_iterations": 10, + "session_ttl_hours": 24, + "cron_session_cleanup": "0 0 * * * *" + }), + supported_targets: vec![env!("TARGET").to_string()], + } +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn test_manifest_json_output() { + let manifest = build_manifest(); + let json = serde_json::to_string_pretty(&manifest).unwrap(); + let parsed: serde_json::Value = serde_json::from_str(&json).unwrap(); + assert!(parsed.is_object()); + assert_eq!(parsed["name"], "iii-agent"); + } + + #[test] + fn test_manifest_has_required_fields() { + let manifest = build_manifest(); + assert!(!manifest.name.is_empty()); + assert!(!manifest.version.is_empty()); + assert!(!manifest.description.is_empty()); + assert!(!manifest.supported_targets.is_empty()); + } +} diff --git a/agent/src/state.rs b/agent/src/state.rs new file mode 100644 index 000000000..c8d01d1ce --- /dev/null +++ b/agent/src/state.rs @@ -0,0 +1,42 @@ +use iii_sdk::{IIIError, TriggerRequest, III}; +use serde_json::{Value, json}; + +pub async fn state_get(iii: &III, scope: &str, key: &str) -> Result { + iii.trigger(TriggerRequest { + function_id: "state::get".to_string(), + payload: json!({ "scope": scope, "key": key }), + action: None, + timeout_ms: Some(5000), + }) + .await +} + +pub async fn state_set(iii: &III, scope: &str, key: &str, value: &Value) -> Result { + iii.trigger(TriggerRequest { + function_id: "state::set".to_string(), + payload: json!({ "scope": scope, "key": key, "value": value }), + action: None, + timeout_ms: Some(5000), + }) + .await +} + +pub async fn state_delete(iii: &III, scope: &str, key: &str) -> Result { + iii.trigger(TriggerRequest { + function_id: "state::delete".to_string(), + payload: json!({ "scope": scope, "key": key }), + action: None, + timeout_ms: Some(5000), + }) + .await +} + +pub async fn state_list(iii: &III, scope: &str) -> Result { + iii.trigger(TriggerRequest { + function_id: "state::list".to_string(), + payload: json!({ "scope": scope }), + action: None, + timeout_ms: Some(5000), + }) + .await +}