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(crate) tool_registry: Arc<RwLock<ToolRegistry>>,
20    #[cfg(feature = "skill")]
21    pub(crate) 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    /// Set a custom sandbox policy.
82    #[cfg(feature = "sandbox")]
83    pub fn with_sandbox_policy(mut self, policy: SandboxPolicy) -> Self {
84        self.sandbox_policy = policy;
85        self
86    }
87
88    #[cfg(feature = "tool")]
89    pub fn with_tool_registry(self, tool_registry: ToolRegistry) -> Self {
90        let snapshot = tool_registry.available_tools_json();
91        let _ = self.tool_tx.send(snapshot);
92        Self {
93            tool_registry: Arc::new(RwLock::new(tool_registry)),
94            ..self
95        }
96    }
97
98    #[cfg(feature = "skill")]
99    pub fn with_skill_registry(self, skill_registry: SkillRegistry) -> Self {
100        let prompt = skill_registry.get_active_skills_prompt();
101        let _ = self.skill_tx.send(prompt);
102        Self {
103            skill_registry: Arc::new(RwLock::new(skill_registry)),
104            ..self
105        }
106    }
107
108    #[cfg(feature = "tool")]
109    pub(crate) async fn add_tool(&mut self, tool: Box<dyn Tool>) {
110        let mut registry = self.tool_registry.write().await;
111        registry.add_tool(tool);
112        let _ = self.tool_tx.send(registry.available_tools_json());
113    }
114
115    #[cfg(feature = "tool")]
116    pub(crate) async fn remove_tool(&mut self, name: &str) {
117        let mut registry = self.tool_registry.write().await;
118        registry.remove_tool(name);
119        let _ = self.tool_tx.send(registry.available_tools_json());
120    }
121
122    #[cfg(feature = "tool")]
123    pub(crate) async fn set_tool_availability(&mut self, _name: &str, _available: bool) {
124        let registry = self.tool_registry.read().await;
125        let _ = self.tool_tx.send(registry.available_tools_json());
126    }
127
128    pub(crate) fn set_client(&mut self, client: async_openai::Client<OpenAIConfig>) {
129        self.llm_client = client.clone();
130        let _ = self.client_tx.send(client);
131    }
132
133    pub(crate) fn set_model(&mut self, model: impl Into<String>) {
134        let model = model.into();
135        self.model = model.clone();
136        let _ = self.model_tx.send(model);
137    }
138
139    #[cfg(feature = "skill")]
140    pub(crate) async fn add_skill(&mut self, skill: Skill) {
141        let mut registry = self.skill_registry.write().await;
142        registry.add(skill);
143        let _ = self.skill_tx.send(registry.get_active_skills_prompt());
144    }
145
146    #[cfg(feature = "skill")]
147    pub(crate) async fn remove_skill(&mut self, name: &str) {
148        let mut registry = self.skill_registry.write().await;
149        registry.remove(name);
150        let _ = self.skill_tx.send(registry.get_active_skills_prompt());
151    }
152
153    #[cfg(feature = "skill")]
154    pub(crate) async fn activate_skill(&mut self, name: &str) -> bool {
155        let mut registry = self.skill_registry.write().await;
156        let ok = registry.activate(name);
157        if ok {
158            let _ = self.skill_tx.send(registry.get_active_skills_prompt());
159        }
160        ok
161    }
162
163    #[cfg(feature = "skill")]
164    pub(crate) async fn deactivate_skill(&mut self, name: &str) -> bool {
165        let mut registry = self.skill_registry.write().await;
166        let ok = registry.deactivate(name);
167        if ok {
168            let _ = self.skill_tx.send(registry.get_active_skills_prompt());
169        }
170        ok
171    }
172
173    #[cfg(feature = "skill")]
174    pub(crate) fn skill_prompt_now(&self) -> String {
175        self.skill_tx.borrow().clone()
176    }
177
178    #[cfg(feature = "skill")]
179    pub(crate) fn set_skill_prompt(&mut self, prompt: String) {
180        let _ = self.skill_tx.send(prompt);
181    }
182
183    pub(crate) fn model(&self) -> &str {
184        &self.model
185    }
186}
187
188#[derive(Debug, Clone)]
189pub struct FuneraEnvWatcher {
190    #[cfg(feature = "tool")]
191    tool_rx: watch::Receiver<JsonValue>,
192    client_rx: watch::Receiver<async_openai::Client<OpenAIConfig>>,
193    model_rx: watch::Receiver<String>,
194    #[cfg(feature = "skill")]
195    skill_rx: watch::Receiver<String>,
196}
197
198impl FuneraEnvWatcher {
199    #[cfg(feature = "tool")]
200    pub fn watch_tool(&mut self) -> JsonValue {
201        self.tool_rx.borrow_and_update().clone()
202    }
203
204    pub fn watch_client(&mut self) -> async_openai::Client<OpenAIConfig> {
205        self.client_rx.borrow_and_update().clone()
206    }
207
208    pub fn watch_model(&mut self) -> String {
209        self.model_rx.borrow_and_update().clone()
210    }
211
212    #[cfg(feature = "skill")]
213    pub fn watch_skill(&mut self) -> String {
214        self.skill_rx.borrow_and_update().clone()
215    }
216
217    #[cfg(feature = "tool")]
218    pub fn has_tool_changed(&self) -> bool {
219        self.tool_rx.has_changed().unwrap_or(false)
220    }
221
222    pub fn has_client_changed(&self) -> bool {
223        self.client_rx.has_changed().unwrap_or(false)
224    }
225
226    pub fn has_model_changed(&self) -> bool {
227        self.model_rx.has_changed().unwrap_or(false)
228    }
229
230    #[cfg(feature = "skill")]
231    pub fn has_skill_changed(&self) -> bool {
232        self.skill_rx.has_changed().unwrap_or(false)
233    }
234
235    pub fn use_client(&mut self) -> async_openai::Client<OpenAIConfig> {
236        self.watch_client()
237    }
238
239    #[cfg(feature = "tool")]
240    pub async fn tool_changed(&mut self) -> Result<(), RecvError> {
241        self.tool_rx.changed().await
242    }
243
244    pub async fn client_changed(&mut self) -> Result<(), RecvError> {
245        self.client_rx.changed().await
246    }
247
248    pub async fn model_changed(&mut self) -> Result<(), RecvError> {
249        self.model_rx.changed().await
250    }
251
252    #[cfg(feature = "skill")]
253    pub async fn skill_changed(&mut self) -> Result<(), RecvError> {
254        self.skill_rx.changed().await
255    }
256}