diff --git a/server-rs/crates/module-editor-agent/Cargo.toml b/server-rs/crates/module-editor-agent/Cargo.toml index 7a0593a2e..188b483f0 100644 --- a/server-rs/crates/module-editor-agent/Cargo.toml +++ b/server-rs/crates/module-editor-agent/Cargo.toml @@ -9,12 +9,7 @@ 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 } -tokio = { workspace = true, features = ["macros", "rt", "rt-multi-thread"] } - -[dev-dependencies] -dotenvy = { workspace = true } 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 deleted file mode 100644 index b5c85d330..000000000 --- a/server-rs/crates/module-editor-agent/examples/llm_chat_agent.rs +++ /dev/null @@ -1,386 +0,0 @@ -use std::env; -use std::path::PathBuf; - -use module_editor_agent::agent::agent::{Agent, LlmApiAdaptor}; -use module_editor_agent::agent::agent_builder::AgentBuilder; -use module_editor_agent::agent::error::PromptError; -use module_editor_agent::agent::hook::Hook; -use module_editor_agent::agent::memory::AgentMemory; -use module_editor_agent::agent::run::Flow; -use module_editor_agent::agent::tool::{Tool, ToolCall, ToolDyn}; -use platform_llm::{LlmClient, LlmConfig, LlmError, LlmMessage, LlmProvider}; -use serde::{Deserialize, Serialize}; -use serde_json::Value; -// --------------------------------------------------------------------------- -// 1. LLM completion model adapter -// --------------------------------------------------------------------------- - -struct LlmCompletionModel { - client: LlmClient, -} - -impl LlmApiAdaptor for LlmCompletionModel { - async fn complete(&self, messages: &[LlmMessage]) -> Result { - use platform_llm::LlmTextRequest; - let request = LlmTextRequest::new(messages.to_vec()).with_request_timeout_ms(30_000); - let response = self - .client - .request_text(request) - .await - .map_err(|e| PromptError::CompletionError(e.to_string()))?; - Ok(response.content) - } - - fn tool_result_message(&self, tool_name: &str, output: &str) -> LlmMessage { - LlmMessage::system(format!("Tool '{tool_name}' returned: {output}")) - } - - fn build_assistant_message(&self, text: &str) -> LlmMessage { - LlmMessage::assistant(text) - } -} - -struct ToolValidationHook { - valid_names: Vec, -} - -impl ToolValidationHook { - fn new(names: Vec) -> Self { - Self { valid_names: names } - } -} - -impl Hook for ToolValidationHook { - fn before_tool_call(&self, tool_call: &ToolCall) -> Flow { - if self.valid_names.iter().any(|n| n == &tool_call.name) { - Flow::Continue - } else { - Flow::Skip - } - } -} - -// --------------------------------------------------------------------------- -// 3. Concrete agent builder -// --------------------------------------------------------------------------- - -struct LlmChatAgentBuilder { - client: Option, - system_prompt_parts: Vec, - tools: Vec>, - hooks: Vec>, - max_turns: usize, - memory_data: Option>>, - context: Option, -} - -impl AgentBuilder for LlmChatAgentBuilder { - type Client = LlmClient; - - fn new() -> Self { - Self { - client: None, - system_prompt_parts: Vec::new(), - tools: Vec::new(), - hooks: Vec::new(), - max_turns: 10, - memory_data: None, - context: None, - } - } - - fn with_client(mut self, client: LlmClient) -> Self { - self.client = Some(client); - self - } - - fn system_prompt(mut self, system_prompt: impl Into) -> Self { - self.system_prompt_parts.push(system_prompt.into()); - self - } - - fn tool(mut self, tool: impl Tool + Send + Sync + 'static) -> Self { - self.tools.push(Box::new(tool)); - self - } - - fn add_hook(mut self, hook: impl Hook + 'static) -> Self { - self.hooks.push(Box::new(hook)); - self - } - - fn max_turns(mut self, n: usize) -> Self { - self.max_turns = n; - self - } - - fn memory(mut self, memory: impl AgentMemory + 'static) -> Self { - self.memory_data = Some(Box::new(memory)); - self - } - - fn context(mut self, context: Value) -> Self { - self.context = Some(context); - self - } - - fn build(self) -> Agent { - let model = LlmCompletionModel { - client: self.client.expect("call .with_client() first"), - }; - let mut agent = Agent::new(model); - let tool_specs = self - .tools - .iter() - .map(|tool| ToolPromptSpec { - name: tool.tool_name().to_string(), - description: tool.description(), - parameters: tool.parameters(), - }) - .collect::>(); - let base_prompt = self.system_prompt_parts.join("\n\n"); - let system_prompt = build_tools_system_prompt(&base_prompt, &tool_specs); - agent.tools = self.tools; - agent.hooks = self.hooks; - agent.default_max_turns = self.max_turns; - agent.memory = self.memory_data; - agent.context = self.context; - agent.system_prompt = Some(LlmMessage::system(&system_prompt)); - agent - } -} - -// --------------------------------------------------------------------------- -// 4. Echo tool -// --------------------------------------------------------------------------- - -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, - _context: serde_json::Value, - ) -> Result { - Ok(EchoOutput { result: args.input }) - } -} - -// --------------------------------------------------------------------------- -// 6. System prompt builder -// --------------------------------------------------------------------------- - -struct ToolPromptSpec { - name: String, - description: String, - parameters: serde_json::Value, -} - -fn tool_names(agent: &Agent) -> Vec { - agent - .tools() - .iter() - .map(|t| t.tool_name().to_string()) - .collect() -} - -fn build_tools_system_prompt(base_prompt: &str, tool_specs: &[ToolPromptSpec]) -> 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 tool_specs.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\": { \"argument_name\": \"argument_value\" }\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 tool_specs { - prompt.push_str(&format!("- {}\n", tool.name)); - prompt.push_str(&format!(" Description: {}\n", tool.description)); - prompt.push_str(" Arguments JSON Schema:\n"); - let parameters = serde_json::to_string_pretty(&tool.parameters) - .unwrap_or_else(|_| tool.parameters.to_string()); - prompt.push_str(&indent_multiline(¶meters, " ")); - prompt.push_str("\n"); - } - - 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 -} - -fn indent_multiline(text: &str, indent: &str) -> String { - text.lines() - .map(|line| format!("{indent}{line}")) - .collect::>() - .join("\n") -} - -// --------------------------------------------------------------------------- -// 7. Config loading via dotenvy -// --------------------------------------------------------------------------- - -/// Reads LLM config from environment variables (after dotenvy has loaded the -/// `.env` file into the process environment). -fn load_config_from_env() -> Result> { - let provider_str = env::var("GENARRATIVE_LLM_PROVIDER").unwrap_or_else(|_| "ark".to_string()); - let provider = match provider_str.as_str() { - "ark" => Ok(LlmProvider::Ark), - "dashscope" | "dash_scope" => Ok(LlmProvider::DashScope), - "openai_compatible" | "openai-compatible" => Ok(LlmProvider::OpenAiCompatible), - other => Err(format!("unsupported provider: {other}")), - }?; - - let base_url = env::var("GENARRATIVE_LLM_BASE_URL") - .or_else(|_| env::var("VITE_LLM_BASE_URL")) - .map_err(|_| "missing GENARRATIVE_LLM_BASE_URL or VITE_LLM_BASE_URL".to_string())?; - - let api_key = env::var("GENARRATIVE_LLM_API_KEY") - .or_else(|_| env::var("LLM_API_KEY")) - .or_else(|_| env::var("ARK_API_KEY")) - .map_err(|_| "missing GENARRATIVE_LLM_API_KEY, LLM_API_KEY, or ARK_API_KEY".to_string())?; - - let model = env::var("GENARRATIVE_LLM_MODEL") - .or_else(|_| env::var("VITE_LLM_MODEL")) - .map_err(|_| "missing GENARRATIVE_LLM_MODEL or VITE_LLM_MODEL".to_string())?; - - let request_timeout_ms = env::var("GENARRATIVE_LLM_REQUEST_TIMEOUT_MS") - .ok() - .and_then(|v| v.parse::().ok()) - .unwrap_or(30_000); - - let max_retries = env::var("GENARRATIVE_LLM_MAX_RETRIES") - .ok() - .and_then(|v| v.parse::().ok()) - .unwrap_or(2); - - let retry_backoff_ms = env::var("GENARRATIVE_LLM_RETRY_BACKOFF_MS") - .ok() - .and_then(|v| v.parse::().ok()) - .unwrap_or(1_000); - - Ok(LlmConfig::new( - provider, - base_url, - api_key, - model, - request_timeout_ms, - max_retries, - retry_backoff_ms, - )?) -} - -// --------------------------------------------------------------------------- -// 8. Agent runner -// --------------------------------------------------------------------------- - -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 names = tool_names(&agent); - let tool_hook = ToolValidationHook::new(names.clone()); - - let output = agent - .prompt(LlmMessage::user( - "Use the echo tool to echo 'Hello from the JSON harness!', then tell me what it said.", - )) - .max_turns(5) - .add_hook(tool_hook) - .await - .expect("agent should succeed"); - - println!("--- Final Agent Response ---"); - println!("{}", output.text); - println!("-----------------------------"); - - Ok(()) -} - -fn main() -> Result<(), Box> { - let env_path = std::env::args() - .nth(1) - .map(PathBuf::from) - .unwrap_or_else(|| PathBuf::from(".env.local")); - - if env_path.exists() { - println!("Loading environment from {} …", env_path.display()); - dotenvy::from_path(&env_path) - .map_err(|e| format!("failed to load env file '{}': {e}", env_path.display()))?; - } else { - println!( - "Env file '{}' not found — falling back to current environment.", - env_path.display() - ); - } - - let config = load_config_from_env()?; - let rt = tokio::runtime::Runtime::new()?; - rt.block_on(run_chat_agent(config))?; - - Ok(()) -} diff --git a/server-rs/crates/module-editor-agent/src/agent/run.rs b/server-rs/crates/module-editor-agent/src/agent/run.rs index dd7b90fb6..6397ef09e 100644 --- a/server-rs/crates/module-editor-agent/src/agent/run.rs +++ b/server-rs/crates/module-editor-agent/src/agent/run.rs @@ -9,10 +9,17 @@ use serde::Deserialize; use serde_json::Value; use std::pin::Pin; +pub type TextOutput = String; + +#[derive(Debug, Clone)] +pub struct ToolCallOutput { + pub tool_call: ToolCall, + pub message: String, +} #[derive(Debug, Clone)] pub enum PromptOutput { - Text(String), - Tool(ToolCall), + Text(TextOutput), + Tool(ToolCallOutput), } #[derive(Debug, Clone, Copy, PartialEq, Eq)] @@ -218,13 +225,16 @@ where let output_json = serde_json::to_string(&json_output) .map_err(|e| PromptError::InternalError(e.to_string()))?; - let output_str = format!( + let overall_message = format!( "[tool_call:{tc_id}] args: {arg_json} output: {output_json}" ); let msg = - agent.model.tool_result_message(&tc.name, &output_str); + agent.model.tool_result_message(&tc.name, &overall_message); memory.append_message(msg); - prompt_result.push(Tool(tc.clone())) + prompt_result.push(Tool(ToolCallOutput { + tool_call: tc.clone(), + message: overall_message.clone(), + })) } ToolOutcome::InternalError(failure) if failure.fatal => { return Err(PromptError::ToolError(failure.message));