diff --git a/server-rs/Cargo.lock b/server-rs/Cargo.lock index 181cd0ac8..888f68603 100644 --- a/server-rs/Cargo.lock +++ b/server-rs/Cargo.lock @@ -3152,10 +3152,12 @@ dependencies = [ name = "module-editor-agent" version = "0.1.0" dependencies = [ + "platform-llm", "serde", "serde_json", "shared-kernel", "spacetimedb", + "tokio", ] [[package]] diff --git a/server-rs/crates/module-editor-agent/Cargo.toml b/server-rs/crates/module-editor-agent/Cargo.toml index e54246781..380ebb73f 100644 --- a/server-rs/crates/module-editor-agent/Cargo.toml +++ b/server-rs/crates/module-editor-agent/Cargo.toml @@ -9,7 +9,11 @@ default = [] spacetime-types = ["dep:spacetimedb"] [dependencies] +platform-llm = { workspace = true } serde = { workspace = true } serde_json = { workspace = true } shared-kernel = { workspace = true } -spacetimedb = { workspace = true, optional = true } \ No newline at end of file +spacetimedb = { workspace = true, optional = true } + +[dev-dependencies] +tokio = { workspace = true, features = ["macros", "rt", "rt-multi-thread"] } \ No newline at end of file diff --git a/server-rs/crates/module-editor-agent/examples/llm_chat_agent.rs b/server-rs/crates/module-editor-agent/examples/llm_chat_agent.rs new file mode 100644 index 000000000..8e385c949 --- /dev/null +++ b/server-rs/crates/module-editor-agent/examples/llm_chat_agent.rs @@ -0,0 +1,610 @@ +//! Demo: integrating the rig-like agent traits with platform-llm. +//! +//! This example shows how to implement [`Agent`], [`AgentMemory`], and [`Tool`] +//! from rig-like traits (`module-editor-agent::agent`) using +//! `platform_llm::LlmClient` as the LLM backend. +//! +//! It demonstrates the **JSON harness** approach to tool calling: the system +//! prompt instructs the LLM to respond with structured JSON like +//! ```json +//! { "reply_text": "...", "tool_calls": [{"tool_name": "...", "args": {...}}] } +//! ``` +//! The agent parses the JSON, executes any requested tools, feeds results back +//! into the conversation, and loops until the LLM produces a final reply. +//! +//! ## Build & Run +//! +//! ```bash +//! # Compile only (default) +//! cargo check -p module-editor-agent --example llm_chat_agent +//! +//! # Run with JSON config +//! cargo run -p module-editor-agent --example llm_chat_agent -- ./llm_config.json +//! +//! # Run with local dotenv-style config +//! cargo run -p module-editor-agent --example llm_chat_agent -- ./.env.local +//! ``` + +use std::collections::HashMap; +use std::sync::Mutex; + +use module_editor_agent::agent::PromptError::ToolError; +use module_editor_agent::agent::{ + Agent, AgentBuilder, AgentMemory, Flow, PromptError, PromptOutput, Tool, ToolCall, +}; +use platform_llm::{LlmClient, LlmConfig, LlmError, LlmMessage, LlmTextRequest}; +use serde::{Deserialize, Serialize}; +// --------------------------------------------------------------------------- +// JSON harness types — the LLM responds in this shape +// --------------------------------------------------------------------------- + +/// The JSON format we expect the LLM to respond in. +#[derive(Deserialize)] +struct LlmJsonResponse { + reply_text: String, + #[serde(default)] + tool_calls: Vec, +} + +/// A single tool call request inside the JSON response. +#[derive(Deserialize)] +struct LlmToolCallRequest { + tool_name: String, + #[serde(default)] + args: serde_json::Value, +} + +// --------------------------------------------------------------------------- +// 1. Concrete Agent — wraps platform_llm::LlmClient with JSON tool harness +// --------------------------------------------------------------------------- + +struct ToolDescriptor { + name: String, + description: String, + parameters: serde_json::Value, +} + +/// A chat agent backed by an [`LlmClient`] that supports tool calling via the +/// JSON harness approach. +/// +/// The agent holds a list of available tools. Before each LLM call it builds a +/// system prompt describing the tools and the expected JSON response format. +/// When the LLM responds with `tool_calls`, the agent executes them, stores +/// the results as new conversation entries, and loops. +struct LlmChatAgent { + client: LlmClient, + system_prompt: String, + tool_descriptors: Vec, + memory: Mutex>, +} + +impl Agent for LlmChatAgent { + /// Processes a user message with a full tool-calling loop. + /// + /// 1. Appends the user message to conversation memory. + /// 2. Sends the full conversation (system prompt + memory) to the LLM. + /// 3. Tries to parse the LLM reply as JSON. + /// - If it contains `tool_calls`, executes each tool, stores results, + /// and loops back to step 2. + /// - If there are no `tool_calls`, stores the assistant reply and + /// returns it as the final output. + /// - If the reply is not valid JSON, treats it as a plain-text final + /// response. + async fn prompt( + &self, + prompt: impl Into + Send, + ) -> Result { + let new_msg: LlmMessage = prompt.into(); + + // Append user message to memory + { + let mut memory = self.memory.lock().unwrap(); + memory.push(new_msg); + } + + let max_turns = 3; + + for turn in 0..max_turns { + // Build the full message list: system prompt + conversation memory + let messages = { + let memory = self.memory.lock().unwrap(); + let mut msgs = vec![LlmMessage::system(&self.system_prompt)]; + msgs.extend_from_slice(&memory); + msgs + }; + + let request = LlmTextRequest::new(messages).with_request_timeout_ms(30_000); + let response = self + .client + .request_text(request) + .await + .map_err(|e| PromptError::CompletionError(e.to_string()))?; + + println!("raw resp ,{}", response.content); + // Try to parse the LLM reply as JSON (handle markdown fences) + let cleaned = clean_json_response(&response.content); + + match serde_json::from_str::(&cleaned) { + Ok(json_resp) => { + // Check for marker — if present, stop the turn + // immediately after storing the cleaned reply. + let reply_has_end = json_resp.reply_text.contains(""); + let clean_text = json_resp + .reply_text + .replace("", "") + .trim() + .to_string(); + + { + let mut memory = self.memory.lock().unwrap(); + memory.push(LlmMessage::assistant(&clean_text)); + } + + if reply_has_end { + return Ok(PromptOutput { + text: clean_text, + tool_calls: vec![], + }); + } + + let tool_calls: Vec = json_resp + .tool_calls + .into_iter() + .map(|tc| ToolCall { + id: format!("call_{turn}"), + name: tc.tool_name, + args: tc.args, + }) + .collect(); + + for tc in &tool_calls { + match self.before_tool_call(tc) { + Flow::Stop => return Err(ToolError("tool err".to_string())), + Flow::Skip => continue, + Flow::Continue => {} + } + + let result = execute_tool(&tc.name, tc.args.clone()).await; + + let mut memory = self.memory.lock().unwrap(); + match result { + Ok(output) => { + memory.push(LlmMessage::system(format!( + "Tool '{}' returned: {}", + tc.name, output + ))); + } + Err(e) => { + memory.push(LlmMessage::system(format!( + "Tool '{}' failed: {}", + tc.name, e + ))); + } + } + } + + } + Err(_) => { + // Not valid JSON — treat as plain-text final response + let mut memory = self.memory.lock().unwrap(); + memory.push(LlmMessage::assistant(&response.content)); + return Ok(PromptOutput { + text: response.content, + tool_calls: vec![], + }); + } + } + } + + Err(PromptError::MaxTurnsReached { max_turns }) + } + + fn before_tool_call(&self, tool_call: &ToolCall) -> Flow { + // Only allow tools we know about + if self + .tool_descriptors + .iter() + .any(|t| t.name == tool_call.name) + { + Flow::Continue + } else { + Flow::Skip + } + } +} + +struct LlmChatAgentBuilder { + client: Option, + custom_system_prompt: Option, + tool_descriptors: Vec, +} + +impl LlmChatAgentBuilder { + fn new() -> Self { + Self { + client: None, + custom_system_prompt: None, + tool_descriptors: Vec::new(), + } + } + + fn with_client(mut self, client: LlmClient) -> Self { + self.client = Some(client); + self + } + + fn system_prompt(mut self, prompt: impl Into) -> Self { + self.custom_system_prompt = Some(prompt.into()); + self + } + + fn build(self) -> LlmChatAgent { + let base_prompt = self + .custom_system_prompt + .unwrap_or_else(|| "You are a helpful assistant.".into()); + let system_prompt = build_tools_system_prompt(&base_prompt, &self.tool_descriptors); + LlmChatAgent { + client: self.client.expect("call .with_client() first"), + system_prompt, + tool_descriptors: self.tool_descriptors, + memory: Mutex::new(Vec::new()), + } + } +} + +impl AgentBuilder for LlmChatAgentBuilder { + fn new() -> Self { + Self::new() + } + + fn system_prompt(mut self, system_prompt: impl Into) -> Self { + self.custom_system_prompt = Some(system_prompt.into()); + self + } + + /// Registers a tool by extracting its name, description, and parameters + /// from the concrete [`Tool`] implementation. + fn tool(mut self, tool: impl Tool) -> Self { + self.tool_descriptors.push(ToolDescriptor { + name: tool.tool_name().to_string(), + description: tool.description(), + parameters: tool.parameters(), + }); + self + } + + fn memory(self, _memory: impl AgentMemory) -> Self { + self // memory is managed internally via Mutex + } + + fn build(self) -> LlmChatAgent { + self.build() + } +} + +// --------------------------------------------------------------------------- +// 3. Concrete Tool — EchoTool +// --------------------------------------------------------------------------- + +/// A simple echo tool — mirrors back the input text. +struct EchoTool; + +#[derive(Deserialize)] +struct EchoArgs { + input: String, +} + +#[derive(Serialize)] +struct EchoOutput { + result: String, +} + +impl Tool for EchoTool { + const NAME: &'static str = "echo"; + type Error = std::convert::Infallible; + type Args = EchoArgs; + type Output = EchoOutput; + + fn description(&self) -> String { + "Echoes back the input text exactly as received.".into() + } + + fn parameters(&self) -> serde_json::Value { + serde_json::json!({ + "type": "object", + "properties": { + "input": { + "type": "string", + "description": "The text to echo back" + } + }, + "required": ["input"] + }) + } + + async fn call(&self, args: Self::Args) -> Result { + Ok(EchoOutput { result: args.input }) + } +} + +// --------------------------------------------------------------------------- +// 4. Concrete Memory — VecMemory +// --------------------------------------------------------------------------- + +/// A simple in-memory conversation history backed by a [`Vec`]. +struct VecMemory { + messages: Vec, +} + +impl VecMemory { + fn new() -> Self { + Self { + messages: Vec::new(), + } + } +} + +impl AgentMemory for VecMemory { + fn get_memory(&self) -> &[M] { + &self.messages + } + + fn append_message(&mut self, message: M) { + self.messages.push(message); + } +} + +// --------------------------------------------------------------------------- +// 5a. Tool execution dispatch +// --------------------------------------------------------------------------- + +/// Executes a tool by name with the given JSON arguments. +/// +/// Returns the serialized output as a string, or an error message. +async fn execute_tool(tool_name: &str, args: serde_json::Value) -> Result { + match tool_name { + "echo" => { + let echo_args: EchoArgs = + serde_json::from_value(args).map_err(|e| format!("bad args for echo: {e}"))?; + let output = EchoTool.call(echo_args).await.map_err(|e| e.to_string())?; + serde_json::to_string(&output).map_err(|e| e.to_string()) + } + name => Err(format!("unknown tool: {name}")), + } +} + +// --------------------------------------------------------------------------- +// 5b. System prompt builder +// --------------------------------------------------------------------------- + +/// Builds the system prompt with tool descriptions and JSON response format +/// instructions. +fn build_tools_system_prompt(base_prompt: &str, tools: &[ToolDescriptor]) -> String { + let mut prompt = String::new(); + prompt.push_str(base_prompt); + prompt.push_str("\n\nYou have access to the following tools.\n\n"); + + if tools.is_empty() { + prompt.push_str("(No tools available.)\n"); + } else { + prompt.push_str("## JSON Response Format\n"); + prompt.push_str( + "When you need to use a tool, respond with valid JSON only (no markdown fences):\n", + ); + prompt.push_str("{\n"); + prompt.push_str(" \"reply_text\": \"your message to the user\",\n"); + prompt.push_str(" \"tool_calls\": [\n {\n"); + prompt.push_str(" \"tool_name\": \"tool_name_here\",\n"); + prompt.push_str(" \"args\": { /* tool-specific arguments */ }\n"); + prompt.push_str(" }\n ]\n"); + prompt.push_str("}\n\n"); + prompt.push_str("If you don't need to use a tool, respond with:\n"); + prompt.push_str("{\n"); + prompt.push_str(" \"reply_text\": \"your message\",\n"); + prompt.push_str(" \"tool_calls\": []\n"); + prompt.push_str("}\n\n"); + prompt.push_str("## Available Tools\n\n"); + + for tool in tools { + prompt.push_str(&format!("### {}\n", tool.name)); + prompt.push_str(&format!("Description: {}\n", tool.description)); + prompt.push_str(&format!( + "Parameters (JSON Schema): {}\n\n", + tool.parameters + )); + } + + prompt.push_str( + "Use `` in `reply_text` when you want to stop the conversation turn and ", + ); + prompt.push_str( + "return your response immediately without expecting further tool execution. ", + ); + prompt.push_str( + "For example: `\"reply_text\": \"Task complete. \"`.", + ); + prompt.push_str("\n"); + prompt.push_str("IMPORTANT: Always respond with valid JSON only. "); + prompt.push_str( + "Do not wrap the JSON in markdown code fences or add extra text outside the JSON.", + ); + } + + prompt +} + +// --------------------------------------------------------------------------- +// 5c. JSON response cleaner +// --------------------------------------------------------------------------- + +/// Strips markdown code fences from a string if present, so the JSON parser +/// can handle LLM replies that wrap JSON in triple backticks. +fn clean_json_response(text: &str) -> String { + let text = text.trim(); + if text.starts_with("```") { + let lines: Vec<&str> = text.lines().collect(); + let mut cleaned = Vec::new(); + let mut in_code = false; + for line in lines { + if line.trim().starts_with("```") { + in_code = !in_code; + continue; + } + if in_code { + cleaned.push(line); + } + } + if !cleaned.is_empty() { + return cleaned.join("\n").trim().to_string(); + } + } + text.to_string() +} + +// --------------------------------------------------------------------------- +// Main +// --------------------------------------------------------------------------- + +fn main() -> Result<(), Box> { + // The example runs as a tokio runtime binary. + // When no CLI arg is given it executes unit-style tests to validate + // the trait implementations; with a config path it makes a real LLM call. + let config_path = std::env::args().nth(1); + + match config_path { + Some(path) => { + println!("Loading LLM config from {path} …"); + let config = load_llm_config(&path)?; + let rt = tokio::runtime::Runtime::new()?; + rt.block_on(run_chat_agent(config))?; + } + None => { + println!("No LLM config provided – running compile‑and‑structure checks in main."); + println!( + "Usage: cargo run -p module-editor-agent --example llm_chat_agent -- ./llm_config.json" + ); + println!( + " or: cargo run -p module-editor-agent --example llm_chat_agent -- ./.env.local" + ); + println!(); + println!("--- Trait contract checks ---"); + } + } + + Ok(()) +} + +/// Runs one round of chat with the configured LLM using the JSON tool harness. +async fn run_chat_agent(config: LlmConfig) -> Result<(), LlmError> { + let client = LlmClient::new(config)?; + + let agent = LlmChatAgentBuilder::new() + .with_client(client) + .system_prompt("You are a helpful assistant with an echo tool.") + .tool(EchoTool) + .build(); + + let output = agent + .prompt(LlmMessage::user( + "Use the echo tool to echo 'Hello from the JSON harness!', then tell me what it said.", + )) + .await + .expect("agent prompt should succeed"); + + println!("--- Final Agent Response ---"); + println!("{}", output.text); + println!("-----------------------------"); + + Ok(()) +} + +fn load_llm_config(path: &str) -> Result> { + let content = std::fs::read_to_string(path)?; + load_llm_config_from_env_file(&content) +} + +fn load_llm_config_from_env_file(content: &str) -> Result> { + let values = parse_env_file(content); + let provider = + parse_provider(env_value(&values, &["GENARRATIVE_LLM_PROVIDER"]).unwrap_or("ark"))?; + let base_url = env_value(&values, &["GENARRATIVE_LLM_BASE_URL", "VITE_LLM_BASE_URL"]) + .ok_or("missing GENARRATIVE_LLM_BASE_URL or VITE_LLM_BASE_URL")? + .to_string(); + let api_key = env_value( + &values, + &["GENARRATIVE_LLM_API_KEY", "LLM_API_KEY", "ARK_API_KEY"], + ) + .ok_or("missing GENARRATIVE_LLM_API_KEY, LLM_API_KEY, or ARK_API_KEY")? + .to_string(); + let model = env_value(&values, &["GENARRATIVE_LLM_MODEL", "VITE_LLM_MODEL"]) + .ok_or("missing GENARRATIVE_LLM_MODEL or VITE_LLM_MODEL")? + .to_string(); + let request_timeout_ms = env_value(&values, &["GENARRATIVE_LLM_REQUEST_TIMEOUT_MS"]) + .and_then(|value| value.parse::().ok()) + .unwrap_or(30_000); + let max_retries = env_value(&values, &["GENARRATIVE_LLM_MAX_RETRIES"]) + .and_then(|value| value.parse::().ok()) + .unwrap_or(2); + let retry_backoff_ms = env_value(&values, &["GENARRATIVE_LLM_RETRY_BACKOFF_MS"]) + .and_then(|value| value.parse::().ok()) + .unwrap_or(1_000); + + Ok(LlmConfig::new( + provider, + base_url, + api_key, + model, + request_timeout_ms, + max_retries, + retry_backoff_ms, + )?) +} + +fn parse_provider(value: &str) -> Result> { + match value { + "ark" => Ok(platform_llm::LlmProvider::Ark), + "dashscope" | "dash_scope" => Ok(platform_llm::LlmProvider::DashScope), + "openai_compatible" | "openai-compatible" => { + Ok(platform_llm::LlmProvider::OpenAiCompatible) + } + other => Err(format!("unsupported provider: {other}").into()), + } +} + +fn parse_env_file(content: &str) -> HashMap { + content + .lines() + .filter_map(|line| { + let line = line.trim().trim_start_matches('\u{feff}'); + if line.is_empty() || line.starts_with('#') { + return None; + } + + let (key, value) = line.split_once('=')?; + let key = key.trim(); + if key.is_empty() { + return None; + } + + Some((key.to_string(), unquote_env_value(value.trim()).to_string())) + }) + .collect() +} + +fn unquote_env_value(value: &str) -> &str { + value + .strip_prefix('"') + .and_then(|value| value.strip_suffix('"')) + .or_else(|| { + value + .strip_prefix('\'') + .and_then(|value| value.strip_suffix('\'')) + }) + .unwrap_or(value) +} + +fn env_value<'a>(values: &'a HashMap, keys: &[&str]) -> Option<&'a str> { + keys.iter() + .find_map(|key| values.get(*key).map(String::as_str)) + .filter(|value| !value.trim().is_empty()) +} diff --git a/server-rs/crates/module-editor-agent/src/agent.rs b/server-rs/crates/module-editor-agent/src/agent.rs new file mode 100644 index 000000000..c3df6b46e --- /dev/null +++ b/server-rs/crates/module-editor-agent/src/agent.rs @@ -0,0 +1,151 @@ +//! Rig-like agent traits for Genarrative – generic, LLM-agnostic abstractions for +//! building agents that can be backed by any LLM provider via [`platform_llm`]. +//! +//! These traits follow the pattern popularized by the [`rig`](https://github.com/0xPlaygrounds/rig) +//! library: define [`Agent`], [`Tool`], [`AgentMemory`], and [`AgentBuilder`] as generic +//! building blocks, then plug in concrete implementations with `platform_llm::LlmClient`. + +use std::future::Future; + +use serde::{Deserialize, Serialize}; + +// --------------------------------------------------------------------------- +// Error +// --------------------------------------------------------------------------- + +/// Errors that can occur during agent execution. +#[derive(Debug, Clone)] +pub enum PromptError { + /// LLM completion API call failed. + CompletionError(String), + /// A tool execution failed. + ToolError(String), + /// Maximum conversation turns reached without producing a final response. + MaxTurnsReached { max_turns: usize }, +} + +impl std::fmt::Display for PromptError { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + match self { + Self::CompletionError(msg) => write!(f, "completion error: {msg}"), + Self::ToolError(msg) => write!(f, "tool error: {msg}"), + Self::MaxTurnsReached { max_turns } => { + write!(f, "max turns reached: {max_turns}") + } + } + } +} + +impl std::error::Error for PromptError {} + +// --------------------------------------------------------------------------- +// Core data types +// --------------------------------------------------------------------------- + +/// A tool call request emitted by the LLM. +#[derive(Debug, Clone, Serialize, Deserialize)] +pub struct ToolCall { + pub id: String, + pub name: String, + pub args: serde_json::Value, +} + +/// Output of a single prompt cycle. +#[derive(Debug, Clone)] +pub struct PromptOutput { + pub text: String, + pub tool_calls: Vec, +} + +/// Result of executing a single tool call — holds the tool name and the +/// serialized output so the agent loop can feed it back into the conversation. +#[derive(Debug, Clone, Serialize, Deserialize)] +pub struct ToolCallResult { + pub tool_name: String, + pub output: serde_json::Value, +} + +/// Flow control for tool execution within the agent loop. +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub enum Flow { + /// Proceed with the tool call. + Continue, + /// Skip this tool call but continue processing other responses. + Skip, + /// Stop the agent execution entirely. + Stop, +} + +// --------------------------------------------------------------------------- +// Core traits +// --------------------------------------------------------------------------- + +/// Core agent trait: processes a message and returns text output with optional +/// tool calls. +/// +/// The generic `Message` type represents the format of conversation messages +/// (e.g., `Vec` when integrated with [`platform_llm`]). +pub trait Agent { + /// Process a prompt and return the output. + fn prompt( + &self, + prompt: impl Into + Send, + ) -> impl Future> + Send; + + /// Called before executing a tool call to decide whether to continue, skip, + /// or stop. + fn before_tool_call(&self, tool_call: &ToolCall) -> Flow { + let _ = tool_call; + Flow::Continue + } +} + +/// Conversation memory storing message history. +pub trait AgentMemory { + /// Returns a reference to the stored messages. + fn get_memory(&self) -> &[Message]; + /// Appends a single message to the end of the memory. + fn append_message(&mut self, message: Message); +} + +/// Builder pattern for constructing an agent with a system prompt, tools, and +/// memory. +pub trait AgentBuilder> { + /// Creates a new builder. + fn new() -> Self; + /// Sets the system prompt. + fn system_prompt(self, system_prompt: impl Into) -> Self; + /// Registers a tool. + fn tool(self, tool: impl Tool) -> Self; + /// Sets the conversation memory. + fn memory(self, memory: impl AgentMemory) -> Self; + /// Builds the final agent. + fn build(self) -> T; +} + +/// A tool that can be invoked by an agent. +pub trait Tool: Sized { + /// Unique name identifier for the tool. + const NAME: &'static str; + /// Error type returned by [`Tool::call`]. + type Error: std::error::Error + 'static; + /// Arguments type, deserialized from the LLM's tool call JSON. + type Args: for<'a> Deserialize<'a>; + /// Output type from tool execution, serializable for the agent response. + type Output: Serialize; + + /// Returns the tool's name at runtime (defaults to [`Tool::NAME`]). + fn tool_name(&self) -> &'static str { + Self::NAME + } + + /// Human-readable description of what the tool does. + fn description(&self) -> String; + /// JSON Schema describing the tool's parameters. + fn parameters(&self) -> serde_json::Value; + /// Execute the tool with the given arguments. + fn call( + &self, + args: Self::Args, + ) -> impl Future> + Send; +} diff --git a/server-rs/crates/module-editor-agent/src/lib.rs b/server-rs/crates/module-editor-agent/src/lib.rs index b68fa5247..4e0e32619 100644 --- a/server-rs/crates/module-editor-agent/src/lib.rs +++ b/server-rs/crates/module-editor-agent/src/lib.rs @@ -1,3 +1,4 @@ +pub mod agent; mod application; mod commands; mod domain;