rs_agent/
agent_orchestrators.rs

1//! CodeMode orchestration and tool integration
2//!
3//! This module handles the integration of CodeMode with the agent system,
4//! matching the structure from go-agent's agent_orchestrators.go.
5
6use std::collections::HashMap;
7use std::sync::Arc;
8
9use anyhow::anyhow;
10use async_trait::async_trait;
11use rs_utcp::plugins::codemode::{
12    CodeModeArgs, CodeModeResult, CodeModeUtcp, CodemodeOrchestrator, LlmModel,
13};
14use serde_json::Value;
15
16use crate::error::AgentError;
17use crate::models::LLM;
18use crate::tools::Tool;
19use crate::types::{Message, Role, ToolRequest, ToolResponse, ToolSpec};
20
21/// Adapter that exposes the UTCP CodeMode runtime as a tool in the agent catalog.
22///
23/// This allows agents to execute code snippets via the `codemode.run_code` tool.
24pub struct CodeModeTool {
25    engine: Arc<CodeModeUtcp>,
26}
27
28impl CodeModeTool {
29    pub fn new(engine: Arc<CodeModeUtcp>) -> Self {
30        Self { engine }
31    }
32
33    fn spec_from_engine(&self) -> ToolSpec {
34        let schema = self.engine.tool();
35        let input_schema = serde_json::to_value(&schema.inputs)
36            .unwrap_or_else(|_| serde_json::json!({"type": "object"}));
37
38        ToolSpec {
39            name: schema.name,
40            description: schema.description,
41            input_schema,
42            examples: None,
43        }
44    }
45}
46
47#[async_trait]
48impl Tool for CodeModeTool {
49    fn spec(&self) -> ToolSpec {
50        self.spec_from_engine()
51    }
52
53    async fn invoke(&self, req: ToolRequest) -> crate::Result<ToolResponse> {
54        let code = req
55            .arguments
56            .get("code")
57            .and_then(|v| v.as_str())
58            .ok_or_else(|| AgentError::ToolError("codemode.run_code requires `code`".into()))?;
59
60        let timeout = req.arguments.get("timeout").and_then(|v| v.as_u64());
61
62        let result = self
63            .engine
64            .execute(CodeModeArgs {
65                code: code.to_string(),
66                timeout,
67            })
68            .await
69            .map_err(|e| AgentError::ToolError(e.to_string()))?;
70
71        let content = serialize_result(&result);
72        Ok(ToolResponse {
73            content,
74            metadata: Some(HashMap::from([(
75                "provider".to_string(),
76                "codemode".to_string(),
77            )])),
78        })
79    }
80}
81
82/// Bridge that lets the CodeMode orchestrator reuse an `rs-agent` LLM.
83///
84/// This adapter allows the CodeMode orchestrator to call into any LLM provider
85/// that implements the rs-agent LLM trait.
86pub struct CodemodeLlmAdapter {
87    llm: Arc<dyn LLM>,
88}
89
90impl CodemodeLlmAdapter {
91    pub fn new(llm: Arc<dyn LLM>) -> Self {
92        Self { llm }
93    }
94}
95
96#[async_trait]
97impl LlmModel for CodemodeLlmAdapter {
98    async fn complete(&self, prompt: &str) -> anyhow::Result<Value> {
99        let messages = vec![Message {
100            role: Role::User,
101            content: prompt.to_string(),
102            metadata: None,
103        }];
104
105        let result = self
106            .llm
107            .generate(messages, None)
108            .await
109            .map_err(|e| anyhow!(e.to_string()))?;
110
111        let cleaned = strip_code_fence(&result.content);
112        Ok(Value::String(cleaned))
113    }
114}
115
116/// Builds a CodeMode orchestrator with the given engine and LLM.
117///
118/// The orchestrator can automatically route natural language queries to tool chains
119/// or executable code snippets.
120pub fn build_orchestrator(engine: Arc<CodeModeUtcp>, llm: Arc<dyn LLM>) -> CodemodeOrchestrator {
121    let adapter = CodemodeLlmAdapter::new(llm);
122    CodemodeOrchestrator::new(engine, Arc::new(adapter))
123}
124
125/// Convenience function to format orchestrator output for agent responses.
126pub fn format_codemode_value(value: &Value) -> String {
127    if let Some(s) = value.as_str() {
128        return s.to_string();
129    }
130
131    serde_json::to_string(value).unwrap_or_else(|_| format!("{value:?}"))
132}
133
134fn serialize_result(result: &CodeModeResult) -> String {
135    serde_json::to_string(result).unwrap_or_else(|_| format!("{result:?}"))
136}
137
138fn strip_code_fence(s: &str) -> String {
139    let trimmed = s.trim();
140    if !trimmed.starts_with("```") {
141        return trimmed.to_string();
142    }
143
144    // Remove opening ```lang if present
145    let mut inner = trimmed.trim_start_matches("```");
146    if let Some(pos) = inner.find('\n') {
147        inner = &inner[pos + 1..];
148    }
149
150    // Remove closing ```
151    if let Some(end) = inner.rfind("```") {
152        inner = &inner[..end];
153    }
154
155    inner.trim().to_string()
156}
157
158#[cfg(test)]
159mod tests {
160    use super::*;
161
162    #[test]
163    fn strip_code_fence_removes_markdown_fences() {
164        assert_eq!(strip_code_fence("```rust\nlet x = 5;\n```"), "let x = 5;");
165        assert_eq!(strip_code_fence("```\ncode\n```"), "code");
166        assert_eq!(strip_code_fence("plain text"), "plain text");
167    }
168
169    #[test]
170    fn format_codemode_value_handles_strings_and_json() {
171        assert_eq!(format_codemode_value(&Value::String("test".into())), "test");
172        assert_eq!(
173            format_codemode_value(&Value::Number(42.into())),
174            "42"
175        );
176    }
177}
rs_agent/agent_orchestrators.rs

rs_agent/
agent_orchestrators.rs
