v1.3.0: Full Rig integration - Multi-agent AI framework

rig-service/src/agent.rs

@@ -0,0 +1,217 @@
+//! Agent management and multi-agent orchestration
+
+use anyhow::Result;
+use rig::{
+    agent::Agent,
+    client::{CompletionClient, ProviderClient},
+    completion::{Completion, Message},
+    providers::openai,
+};
+use serde::{Deserialize, Serialize};
+use std::sync::Arc;
+use tokio::sync::RwLock;
+use uuid::Uuid;
+
+use crate::tools::{Tool, ToolRegistry, ToolResult};
+
+/// Agent configuration
+#[derive(Debug, Clone, Serialize, Deserialize)]
+pub struct AgentConfig {
+    pub name: String,
+    pub preamble: String,
+    pub model: String,
+    pub provider: String,
+    pub temperature: f32,
+    pub max_turns: u32,
+}
+
+/// Agent session
+#[derive(Debug, Clone)]
+pub struct AgentSession {
+    pub id: String,
+    pub config: AgentConfig,
+    pub messages: Vec<Message>,
+}
+
+/// Agent Council for multi-agent orchestration
+#[derive(Debug, Clone)]
+pub struct AgentCouncil {
+    pub id: String,
+    pub name: String,
+    pub agents: Vec<AgentSession>,
+}
+
+/// Agent manager
+#[derive(Debug, Clone)]
+pub struct AgentManager {
+    sessions: Arc<RwLock<Vec<AgentSession>>>,
+    councils: Arc<RwLock<Vec<AgentCouncil>>>,
+    tool_registry: ToolRegistry,
+}
+
+impl AgentManager {
+    pub fn new(tool_registry: ToolRegistry) -> Self {
+        Self {
+            sessions: Arc::new(RwLock::new(Vec::new())),
+            councils: Arc::new(RwLock::new(Vec::new())),
+            tool_registry,
+        }
+    }
+    
+    /// Create a new agent session
+    pub async fn create_session(&self, config: AgentConfig) -> Result<String> {
+        let session = AgentSession {
+            id: Uuid::new_v4().to_string(),
+            config,
+            messages: Vec::new(),
+        };
+        
+        let id = session.id.clone();
+        let mut sessions = self.sessions.write().await;
+        sessions.push(session);
+        
+        Ok(id)
+    }
+    
+    /// Get a session by ID
+    pub async fn get_session(&self, id: &str) -> Option<AgentSession> {
+        let sessions = self.sessions.read().await;
+        sessions.iter().find(|s| s.id == id).cloned()
+    }
+    
+    /// Execute agent prompt
+    pub async fn execute_prompt(
+        &self,
+        session_id: &str,
+        prompt: &str,
+    ) -> Result<String> {
+        let session = self.get_session(session_id)
+            .await
+            .ok_or_else(|| anyhow::anyhow!("Session not found"))?;
+        
+        // Create Rig client based on provider
+        let client = self.create_client(&session.config.provider)?;
+        
+        // Build agent with Rig
+        let agent = client
+            .agent(&session.config.model)
+            .preamble(&session.config.preamble)
+            .temperature(session.config.temperature)
+            .build();
+        
+        // Execute prompt
+        let response = agent.prompt(prompt).await?;
+        
+        // Store message
+        let mut sessions = self.sessions.write().await;
+        if let Some(session) = sessions.iter_mut().find(|s| s.id == session_id) {
+            session.messages.push(Message {
+                role: "user".to_string(),
+                content: prompt.to_string(),
+            });
+            session.messages.push(Message {
+                role: "assistant".to_string(),
+                content: response.clone(),
+            });
+        }
+        
+        Ok(response)
+    }
+    
+    /// Create agent council
+    pub async fn create_council(
+        &self,
+        name: &str,
+        agent_configs: Vec<AgentConfig>,
+    ) -> Result<String> {
+        let mut agents = Vec::new();
+        
+        for config in agent_configs {
+            let session = AgentSession {
+                id: Uuid::new_v4().to_string(),
+                config,
+                messages: Vec::new(),
+            };
+            agents.push(session);
+        }
+        
+        let council = AgentCouncil {
+            id: Uuid::new_v4().to_string(),
+            name: name.to_string(),
+            agents,
+        };
+        
+        let council_id = council.id.clone();
+        let mut councils = self.councils.write().await;
+        councils.push(council);
+        
+        Ok(council_id)
+    }
+    
+    /// Execute council orchestration
+    pub async fn execute_council(
+        &self,
+        council_id: &str,
+        task: &str,
+    ) -> Result<String> {
+        let council = self.councils.read()
+            .await
+            .iter()
+            .find(|c| c.id == council_id)
+            .cloned()
+            .ok_or_else(|| anyhow::anyhow!("Council not found"))?;
+        
+        let mut results = Vec::new();
+        
+        // Execute task with each agent
+        for agent in &council.agents {
+            match self.execute_prompt(&agent.id, task).await {
+                Ok(result) => {
+                    results.push(format!("{}: {}", agent.config.name, result));
+                }
+                Err(e) => {
+                    results.push(format!("{}: Error - {}", agent.config.name, e));
+                }
+            }
+        }
+        
+        // Synthesize results
+        Ok(results.join("\n\n"))
+    }
+    
+    /// Create Rig client for provider
+    fn create_client(&self, provider: &str) -> Result<openai::Client> {
+        match provider.to_lowercase().as_str() {
+            "openai" => {
+                let api_key = std::env::var("OPENAI_API_KEY")
+                    .map_err(|_| anyhow::anyhow!("OPENAI_API_KEY not set"))?;
+                Ok(openai::Client::new(&api_key))
+            }
+            _ => {
+                // Default to OpenAI for now
+                let api_key = std::env::var("OPENAI_API_KEY")
+                    .unwrap_or_else(|_| "dummy".to_string());
+                Ok(openai::Client::new(&api_key))
+            }
+        }
+    }
+    
+    /// List all sessions
+    pub async fn list_sessions(&self) -> Vec<AgentSession> {
+        let sessions = self.sessions.read().await;
+        sessions.clone()
+    }
+    
+    /// List all councils
+    pub async fn list_councils(&self) -> Vec<AgentCouncil> {
+        let councils = self.councils.read().await;
+        councils.clone()
+    }
+    
+    /// Delete a session
+    pub async fn delete_session(&self, id: &str) -> Result<()> {
+        let mut sessions = self.sessions.write().await;
+        sessions.retain(|s| s.id != id);
+        Ok(())
+    }
+}