Skip to main content

funera_core/
env.rs

1use std::sync::Arc;
2
3use async_openai::config::OpenAIConfig;
4
5#[cfg(feature = "skill")]
6use crate::re_act::skills::{Skill, SkillRegistry};
7#[cfg(feature = "tool")]
8use crate::re_act::tool::{Tool, ToolRegistry};
9#[cfg(feature = "sandbox")]
10use crate::security::sandbox::SandboxPolicy;
11use serde_json::Value as JsonValue;
12use tokio::sync::{
13    RwLock,
14    watch::{self, error::RecvError},
15};
16
17pub struct FuneraEnv {
18    #[cfg(feature = "tool")]
19    pub tool_registry: Arc<RwLock<ToolRegistry>>,
20    #[cfg(feature = "skill")]
21    pub skill_registry: Arc<RwLock<SkillRegistry>>,
22    llm_client: async_openai::Client<OpenAIConfig>,
23    model: String,
24    #[cfg(feature = "tool")]
25    tool_tx: watch::Sender<JsonValue>,
26    client_tx: watch::Sender<async_openai::Client<OpenAIConfig>>,
27    model_tx: watch::Sender<String>,
28    #[cfg(feature = "skill")]
29    skill_tx: watch::Sender<String>,
30    #[cfg(feature = "sandbox")]
31    sandbox_policy: SandboxPolicy,
32}
33
34impl FuneraEnv {
35    pub fn new(
36        llm_client: async_openai::Client<OpenAIConfig>,
37        model: impl Into<String>,
38    ) -> (Self, FuneraEnvWatcher) {
39        let model = model.into();
40        let (client_tx, client_rx) = watch::channel(llm_client.clone());
41        let (model_tx, model_rx) = watch::channel(model.clone());
42
43        #[cfg(feature = "tool")]
44        let tool_registry = Arc::new(RwLock::new(ToolRegistry::new()));
45        #[cfg(feature = "tool")]
46        let (tool_tx, tool_rx) = watch::channel(JsonValue::Array(Vec::new()));
47
48        #[cfg(feature = "skill")]
49        let skill_registry = Arc::new(RwLock::new(SkillRegistry::new()));
50        #[cfg(feature = "skill")]
51        let (skill_tx, skill_rx) = watch::channel(String::new());
52
53        (
54            Self {
55                #[cfg(feature = "tool")]
56                tool_registry,
57                #[cfg(feature = "skill")]
58                skill_registry,
59                llm_client,
60                model,
61                #[cfg(feature = "tool")]
62                tool_tx,
63                client_tx,
64                model_tx,
65                #[cfg(feature = "skill")]
66                skill_tx,
67                #[cfg(feature = "sandbox")]
68                sandbox_policy: SandboxPolicy::default(),
69            },
70            FuneraEnvWatcher {
71                #[cfg(feature = "tool")]
72                tool_rx,
73                client_rx,
74                model_rx,
75                #[cfg(feature = "skill")]
76                skill_rx,
77            },
78        )
79    }
80
81    /// Access the current sandbox policy.
82    #[cfg(feature = "sandbox")]
83    pub fn sandbox_policy(&self) -> &SandboxPolicy {
84        &self.sandbox_policy
85    }
86
87    /// Set a custom sandbox policy.
88    #[cfg(feature = "sandbox")]
89    pub fn with_sandbox_policy(mut self, policy: SandboxPolicy) -> Self {
90        self.sandbox_policy = policy;
91        self
92    }
93
94    #[cfg(feature = "tool")]
95    pub fn with_tool_registry(self, tool_registry: ToolRegistry) -> Self {
96        let snapshot = tool_registry.available_tools_json();
97        let _ = self.tool_tx.send(snapshot);
98        Self {
99            tool_registry: Arc::new(RwLock::new(tool_registry)),
100            ..self
101        }
102    }
103
104    #[cfg(feature = "skill")]
105    pub fn with_skill_registry(self, skill_registry: SkillRegistry) -> Self {
106        let prompt = skill_registry.get_active_skills_prompt();
107        let _ = self.skill_tx.send(prompt);
108        Self {
109            skill_registry: Arc::new(RwLock::new(skill_registry)),
110            ..self
111        }
112    }
113
114    #[cfg(feature = "tool")]
115    pub async fn add_tool(&mut self, tool: Box<dyn Tool>) {
116        let mut registry = self.tool_registry.write().await;
117        registry.add_tool(tool);
118        let _ = self.tool_tx.send(registry.available_tools_json());
119    }
120
121    #[cfg(feature = "tool")]
122    pub async fn remove_tool(&mut self, name: &str) {
123        let mut registry = self.tool_registry.write().await;
124        registry.remove_tool(name);
125        let _ = self.tool_tx.send(registry.available_tools_json());
126    }
127
128    #[cfg(feature = "tool")]
129    pub async fn set_tool_availability(&mut self, _name: &str, _available: bool) {
130        let registry = self.tool_registry.read().await;
131        let _ = self.tool_tx.send(registry.available_tools_json());
132    }
133
134    pub fn set_client(&mut self, client: async_openai::Client<OpenAIConfig>) {
135        self.llm_client = client.clone();
136        let _ = self.client_tx.send(client);
137    }
138
139    pub fn set_model(&mut self, model: impl Into<String>) {
140        let model = model.into();
141        self.model = model.clone();
142        let _ = self.model_tx.send(model);
143    }
144
145    #[cfg(feature = "skill")]
146    pub async fn add_skill(&mut self, skill: Skill) {
147        let mut registry = self.skill_registry.write().await;
148        registry.add(skill);
149        let _ = self.skill_tx.send(registry.get_active_skills_prompt());
150    }
151
152    #[cfg(feature = "skill")]
153    pub async fn remove_skill(&mut self, name: &str) {
154        let mut registry = self.skill_registry.write().await;
155        registry.remove(name);
156        let _ = self.skill_tx.send(registry.get_active_skills_prompt());
157    }
158
159    #[cfg(feature = "skill")]
160    pub async fn activate_skill(&mut self, name: &str) -> bool {
161        let mut registry = self.skill_registry.write().await;
162        let ok = registry.activate(name);
163        if ok {
164            let _ = self.skill_tx.send(registry.get_active_skills_prompt());
165        }
166        ok
167    }
168
169    #[cfg(feature = "skill")]
170    pub async fn deactivate_skill(&mut self, name: &str) -> bool {
171        let mut registry = self.skill_registry.write().await;
172        let ok = registry.deactivate(name);
173        if ok {
174            let _ = self.skill_tx.send(registry.get_active_skills_prompt());
175        }
176        ok
177    }
178
179    #[cfg(feature = "skill")]
180    pub fn skill_prompt_now(&self) -> String {
181        self.skill_tx.borrow().clone()
182    }
183
184    #[cfg(feature = "skill")]
185    pub fn set_skill_prompt(&mut self, prompt: String) {
186        let _ = self.skill_tx.send(prompt);
187    }
188}
189
190#[derive(Debug, Clone)]
191pub struct FuneraEnvWatcher {
192    #[cfg(feature = "tool")]
193    tool_rx: watch::Receiver<JsonValue>,
194    client_rx: watch::Receiver<async_openai::Client<OpenAIConfig>>,
195    model_rx: watch::Receiver<String>,
196    #[cfg(feature = "skill")]
197    skill_rx: watch::Receiver<String>,
198}
199
200impl FuneraEnvWatcher {
201    #[cfg(feature = "tool")]
202    pub fn watch_tool(&mut self) -> JsonValue {
203        self.tool_rx.borrow_and_update().clone()
204    }
205
206    pub fn watch_client(&mut self) -> async_openai::Client<OpenAIConfig> {
207        self.client_rx.borrow_and_update().clone()
208    }
209
210    pub fn watch_model(&mut self) -> String {
211        self.model_rx.borrow_and_update().clone()
212    }
213
214    #[cfg(feature = "skill")]
215    pub fn watch_skill(&mut self) -> String {
216        self.skill_rx.borrow_and_update().clone()
217    }
218
219    #[cfg(feature = "tool")]
220    pub fn has_tool_changed(&self) -> bool {
221        self.tool_rx.has_changed().unwrap_or(false)
222    }
223
224    pub fn has_client_changed(&self) -> bool {
225        self.client_rx.has_changed().unwrap_or(false)
226    }
227
228    pub fn has_model_changed(&self) -> bool {
229        self.model_rx.has_changed().unwrap_or(false)
230    }
231
232    #[cfg(feature = "skill")]
233    pub fn has_skill_changed(&self) -> bool {
234        self.skill_rx.has_changed().unwrap_or(false)
235    }
236
237    pub fn use_client(&mut self) -> async_openai::Client<OpenAIConfig> {
238        self.watch_client()
239    }
240
241    #[cfg(feature = "tool")]
242    pub async fn tool_changed(&mut self) -> Result<(), RecvError> {
243        self.tool_rx.changed().await
244    }
245
246    pub async fn client_changed(&mut self) -> Result<(), RecvError> {
247        self.client_rx.changed().await
248    }
249
250    pub async fn model_changed(&mut self) -> Result<(), RecvError> {
251        self.model_rx.changed().await
252    }
253
254    #[cfg(feature = "skill")]
255    pub async fn skill_changed(&mut self) -> Result<(), RecvError> {
256        self.skill_rx.changed().await
257    }
258}