Skip to main content

vv_agent/sdk/client/
agents.rs

1use std::collections::BTreeMap;
2use std::sync::Arc;
3
4use crate::llm::LlmClient;
5use crate::runtime::AgentRuntime;
6
7use super::super::resources::AgentResourceLoader;
8use super::super::types::{AgentDefinition, AgentSDKOptions};
9use super::runtime::{configure_runtime_from_options, SettingsRunAgent};
10use super::AgentSDKClient;
11
12impl AgentSDKClient {
13    pub fn new(options: AgentSDKOptions) -> Self {
14        let mut agents = BTreeMap::new();
15        let mut prompt_templates = BTreeMap::new();
16        let mut resource_skill_directories = Vec::new();
17        let mut resource_diagnostics = Vec::new();
18
19        if options.auto_discover_resources {
20            let mut loader = options
21                .resource_loader
22                .clone()
23                .unwrap_or_else(|| AgentResourceLoader::new(options.workspace.clone()));
24            let discovered = loader.discover();
25            agents = discovered.agents;
26            prompt_templates = discovered.prompts;
27            resource_skill_directories = discovered.skill_directories;
28            resource_diagnostics = discovered.diagnostics;
29        }
30
31        let runtime_options = options.clone();
32
33        Self {
34            options,
35            default_agent: None,
36            agents,
37            prompt_templates,
38            resource_skill_directories,
39            resource_diagnostics,
40            runtime: Arc::new(SettingsRunAgent {
41                options: runtime_options,
42            }),
43        }
44    }
45
46    pub fn new_with_agent(options: AgentSDKOptions, agent: AgentDefinition) -> Self {
47        let mut client = Self::new(options);
48        client.default_agent = Some(agent);
49        client
50    }
51
52    pub fn new_with_agents(
53        options: AgentSDKOptions,
54        agents: BTreeMap<String, AgentDefinition>,
55    ) -> Result<Self, String> {
56        let mut client = Self::new(options);
57        client.register_agents(agents)?;
58        Ok(client)
59    }
60
61    pub fn with_runtime<C: LlmClient + Clone + 'static>(
62        mut self,
63        mut runtime: AgentRuntime<C>,
64    ) -> Self {
65        configure_runtime_from_options(&mut runtime, &self.options);
66        self.runtime = Arc::new(runtime);
67        self
68    }
69
70    pub fn set_default_agent(&mut self, definition: AgentDefinition) {
71        self.default_agent = Some(definition);
72    }
73
74    pub fn register_agent(
75        &mut self,
76        name: impl Into<String>,
77        definition: AgentDefinition,
78    ) -> Result<(), String> {
79        let name = name.into().trim().to_string();
80        if name.is_empty() {
81            return Err("Agent name cannot be empty".to_string());
82        }
83        self.agents.insert(name, definition);
84        Ok(())
85    }
86
87    pub fn register_agents(
88        &mut self,
89        agents: BTreeMap<String, AgentDefinition>,
90    ) -> Result<(), String> {
91        for (name, definition) in agents {
92            self.register_agent(name, definition)?;
93        }
94        Ok(())
95    }
96
97    pub fn list_agents(&self) -> Vec<String> {
98        self.agents.keys().cloned().collect()
99    }
100
101    pub fn resource_diagnostics(&self) -> Vec<String> {
102        self.resource_diagnostics.clone()
103    }
104
105    pub(super) fn get_agent(&self, agent_name: &str) -> Result<&AgentDefinition, String> {
106        if agent_name.is_empty() {
107            return Err("Agent name cannot be empty".to_string());
108        }
109        self.agents.get(agent_name).ok_or_else(|| {
110            let available = self.list_agents().join(", ");
111            format!("Unknown agent: {agent_name}. Available: {available}")
112        })
113    }
114
115    pub(super) fn default_or_only_agent(
116        &self,
117        empty_message: &str,
118        multiple_prefix: &str,
119    ) -> Result<(String, AgentDefinition), String> {
120        if let Some(agent) = self.default_agent.clone() {
121            return Ok(("default".to_string(), agent));
122        }
123        if self.agents.len() == 1 {
124            let (name, definition) = self.agents.iter().next().expect("single agent");
125            return Ok((name.clone(), definition.clone()));
126        }
127        if self.agents.is_empty() {
128            return Err(empty_message.to_string());
129        }
130        let available = self.list_agents().join(", ");
131        Err(format!("{multiple_prefix} {available}"))
132    }
133
134    pub(super) fn effective_definition(&self, mut definition: AgentDefinition) -> AgentDefinition {
135        if self.options.bash_shell.is_some() && definition.bash_shell.is_none() {
136            definition.bash_shell = self.options.bash_shell.clone();
137        }
138        if !self.options.windows_shell_priority.is_empty()
139            && definition.windows_shell_priority.is_empty()
140        {
141            definition.windows_shell_priority = self.options.windows_shell_priority.clone();
142        }
143        if !self.options.bash_env.is_empty() {
144            let mut bash_env = self.options.bash_env.clone();
145            bash_env.extend(definition.bash_env.clone());
146            definition.bash_env = bash_env;
147        }
148        if definition.system_prompt.is_none() {
149            if let Some(template_name) = definition.system_prompt_template.as_deref() {
150                if let Some(template) = self.prompt_templates.get(template_name) {
151                    if !template.trim().is_empty() {
152                        definition.description = template.clone();
153                    }
154                }
155            }
156        }
157        if definition.skill_directories.is_empty() && !self.resource_skill_directories.is_empty() {
158            definition.skill_directories = self.resource_skill_directories.clone();
159        }
160        definition
161    }
162}