Skip to main content

satori_client/
lib.rs

1// src/lib.rs
2
3use serde;
4use std::collections::HashMap;
5use std::sync::{Arc, Mutex};
6use tokio::sync::mpsc::{self, Sender};
7use tokio::task;
8use tokio::time::Duration;
9use tokio_tungstenite::{connect_async_with_config, tungstenite::protocol::Message};
10use url::Url;
11use uuid::Uuid;
12
13pub type Callback = Box<dyn Fn(serde_json::Value) + Send + Sync + 'static>;
14
15#[derive(Clone)]
16pub struct Satori {
17    username: String,
18    password: String,
19    url: String,
20    sender: Sender<Message>,
21    subscriptions: Arc<Mutex<HashMap<String, Callback>>>,
22}
23
24impl Satori {
25    pub async fn connect(username: String, password: String, url: String) -> anyhow::Result<Self> {
26        // Configure WebSocket connection with no timeout
27        let url_parsed = Url::parse(&url)?;
28
29        // Create a custom connector with no timeout
30        let (ws_stream, _) = connect_async_with_config(url_parsed, None, false).await?;
31
32        // Configure the WebSocket to have no timeout by using keepalive settings
33        // The WebSocket connection will maintain the connection without timing out
34
35        let (mut write, read) = ws_stream.split();
36        let (sender, mut receiver) = mpsc::channel(100);
37        let subscriptions: Arc<Mutex<HashMap<String, Callback>>> =
38            Arc::new(Mutex::new(HashMap::new()));
39
40        let _subscriptions_clone = subscriptions.clone();
41
42        task::spawn(async move {
43            while let Some(msg) = receiver.recv().await {
44                let _ = write.send(msg).await;
45            }
46        });
47
48        let subscriptions_clone2 = subscriptions.clone();
49
50        // Spawn a task to handle incoming messages and ping responses
51        task::spawn(async move {
52            let mut read = read;
53            while let Some(Ok(msg)) = read.next().await {
54                match msg {
55                    Message::Text(txt) => {
56                        if let Ok(json) = serde_json::from_str::<serde_json::Value>(&txt) {
57                            if json.get("type")
58                                == Some(&serde_json::Value::String("notification".to_string()))
59                            {
60                                if let Some(key) = json.get("key").and_then(|k| k.as_str()) {
61                                    if let Some(cb) = subscriptions_clone2.lock().unwrap().get(key)
62                                    {
63                                        cb(json["data"].clone());
64                                    }
65                                }
66                            }
67                        }
68                    }
69                    Message::Ping(_data) => {
70                        // Respond to ping with pong to keep connection alive
71                        // Note: We can't send pong here since write is moved to another task
72                        // The WebSocket library should handle ping/pong automatically
73                    }
74                    Message::Pong(_) => {
75                        // Handle pong messages (optional)
76                    }
77                    _ => {
78                        // Handle other message types if needed
79                    }
80                }
81            }
82        });
83
84        // Spawn a task to send periodic pings to keep the connection alive
85        let sender_clone = sender.clone();
86        task::spawn(async move {
87            let mut interval = tokio::time::interval(Duration::from_secs(30));
88            loop {
89                interval.tick().await;
90                if let Err(e) = sender_clone.send(Message::Ping(vec![])).await {
91                    eprintln!("Failed to send ping: {}", e);
92                    break;
93                }
94            }
95        });
96
97        Ok(Self {
98            username,
99            password,
100            url,
101            sender,
102            subscriptions,
103        })
104    }
105
106    pub async fn send(
107        &self,
108        mut payload: serde_json::Map<String, serde_json::Value>,
109    ) -> anyhow::Result<serde_json::Value> {
110        let id = Uuid::new_v4().to_string();
111        payload.insert("id".into(), id.clone().into());
112        payload.insert("username".into(), self.username.clone().into());
113        payload.insert("password".into(), self.password.clone().into());
114
115        let msg = Message::Text(serde_json::Value::Object(payload.clone()).to_string());
116        self.sender.send(msg).await?;
117        Ok(serde_json::json!({"status": "sent", "id": id}))
118    }
119
120    pub async fn command(
121        &self,
122        command: &str,
123        args: serde_json::Value,
124    ) -> anyhow::Result<serde_json::Value> {
125        let mut payload = args.as_object().cloned().unwrap_or_default();
126        payload.insert("command".into(), command.into());
127        self.send(payload).await
128    }
129
130    pub async fn set_notify<F>(&self, key: &str, callback: F) -> anyhow::Result<()>
131    where
132        F: Fn(serde_json::Value) + Send + Sync + 'static,
133    {
134        self.subscriptions
135            .lock()
136            .unwrap()
137            .insert(key.to_string(), Box::new(callback));
138        let mut payload = serde_json::Map::new();
139        payload.insert("command".into(), "NOTIFY".into());
140        payload.insert("key".into(), key.into());
141        self.send(payload).await?;
142        Ok(())
143    }
144
145    // Convenience methods
146    pub async fn set(&self, args: serde_json::Value) -> anyhow::Result<serde_json::Value> {
147        self.command("SET", args).await
148    }
149
150    pub async fn query(&self, args: serde_json::Value) -> anyhow::Result<serde_json::Value> {
151        self.command("QUERY", args).await
152    }
153
154    pub async fn ann(&self, args: serde_json::Value) -> anyhow::Result<serde_json::Value> {
155        self.command("ANN", args).await
156    }
157
158    pub async fn set_middleware(
159        &self,
160        args: serde_json::Value,
161    ) -> anyhow::Result<serde_json::Value> {
162        self.command("SET_MIDDLEWARE", args).await
163    }
164
165    pub async fn get_access_frequency(
166        &self,
167        args: serde_json::Value,
168    ) -> anyhow::Result<serde_json::Value> {
169        self.command("GET_ACCESS_FREQUENCY", args).await
170    }
171
172    pub async fn get_operations(&self) -> anyhow::Result<serde_json::Value> {
173        self.command("GET_OPERATIONS", serde_json::Value::Null)
174            .await
175    }
176
177    pub async fn ask(&self, args: serde_json::Value) -> anyhow::Result<serde_json::Value> {
178        self.command("ASK", args).await
179    }
180
181    pub async fn get(&self, args: serde_json::Value) -> anyhow::Result<serde_json::Value> {
182        self.command("GET", args).await
183    }
184    pub async fn put(&self, args: serde_json::Value) -> anyhow::Result<serde_json::Value> {
185        self.command("PUT", args).await
186    }
187    pub async fn delete(&self, args: serde_json::Value) -> anyhow::Result<serde_json::Value> {
188        self.command("DELETE", args).await
189    }
190
191    pub async fn encrypt(&self, args: serde_json::Value) -> anyhow::Result<serde_json::Value> {
192        self.command("ENCRYPT", args).await
193    }
194    pub async fn decrypt(&self, args: serde_json::Value) -> anyhow::Result<serde_json::Value> {
195        self.command("DECRYPT", args).await
196    }
197    pub async fn push(&self, args: serde_json::Value) -> anyhow::Result<serde_json::Value> {
198        self.command("PUSH", args).await
199    }
200    pub async fn pop(&self, args: serde_json::Value) -> anyhow::Result<serde_json::Value> {
201        self.command("POP", args).await
202    }
203
204    pub async fn splice(&self, args: serde_json::Value) -> anyhow::Result<serde_json::Value> {
205        self.command("SPLICE", args).await
206    }
207    pub async fn remove(&self, args: serde_json::Value) -> anyhow::Result<serde_json::Value> {
208        self.command("REMOVE", args).await
209    }
210    pub async fn dfs(&self, args: serde_json::Value) -> anyhow::Result<serde_json::Value> {
211        self.command("DFS", args).await
212    }
213    pub async fn set_vertex(&self, args: serde_json::Value) -> anyhow::Result<serde_json::Value> {
214        self.command("SET_VERTEX", args).await
215    }
216    pub async fn get_vertex(&self, args: serde_json::Value) -> anyhow::Result<serde_json::Value> {
217        self.command("GET_VERTEX", args).await
218    }
219    pub async fn delete_vertex(
220        &self,
221        args: serde_json::Value,
222    ) -> anyhow::Result<serde_json::Value> {
223        self.command("DELETE_VERTEX", args).await
224    }
225
226    // ========== Graph Operations ==========
227
228    /// Performs breadth-first search on the graph.
229    /// Returns all nodes reachable from a starting node.
230    pub async fn graph_bfs(&self, args: serde_json::Value) -> anyhow::Result<serde_json::Value> {
231        self.command("GRAPH_BFS", args).await
232    }
233
234    /// Performs depth-first search on the graph.
235    /// Returns all nodes reachable from a starting node.
236    pub async fn graph_dfs(&self, args: serde_json::Value) -> anyhow::Result<serde_json::Value> {
237        self.command("GRAPH_DFS", args).await
238    }
239
240    /// Finds the shortest path between a start node and end node using Dijkstra's algorithm.
241    pub async fn graph_shortest_path(&self, args: serde_json::Value) -> anyhow::Result<serde_json::Value> {
242        self.command("GRAPH_SHORTEST_PATH", args).await
243    }
244
245    /// Identifies all connected components in the graph.
246    pub async fn graph_connected_components(&self, args: serde_json::Value) -> anyhow::Result<serde_json::Value> {
247        self.command("GRAPH_CONNECTED_COMPONENTS", args).await
248    }
249
250    /// Finds strongly connected components using Tarjan's algorithm.
251    pub async fn graph_scc(&self, args: serde_json::Value) -> anyhow::Result<serde_json::Value> {
252        self.command("GRAPH_SCC", args).await
253    }
254
255    /// Calculates degree centrality (number of connections) for each node in the graph.
256    pub async fn graph_degree_centrality(&self, args: serde_json::Value) -> anyhow::Result<serde_json::Value> {
257        self.command("GRAPH_DEGREE_CENTRALITY", args).await
258    }
259
260    /// Calculates closeness centrality for each node.
261    pub async fn graph_closeness_centrality(&self, args: serde_json::Value) -> anyhow::Result<serde_json::Value> {
262        self.command("GRAPH_CLOSENESS_CENTRALITY", args).await
263    }
264
265    /// Finds the centroid node (node with highest closeness centrality).
266    pub async fn graph_centroid(&self, args: serde_json::Value) -> anyhow::Result<serde_json::Value> {
267        self.command("GRAPH_CENTROID", args).await
268    }
269
270    // ========== AI Operations ==========
271
272    /// Alias for ANN - performs approximate nearest neighbor search.
273    pub async fn get_similar(&self, args: serde_json::Value) -> anyhow::Result<serde_json::Value> {
274        self.command("GET_SIMILAR", args).await
275    }
276
277    // ========== Mindspace Operations ==========
278
279    /// Creates a new mindspace (cognitive context) for AI-powered conversations.
280    pub async fn set_mindspace(&self, args: serde_json::Value) -> anyhow::Result<serde_json::Value> {
281        self.command("SET_MINDSPACE", args).await
282    }
283
284    /// Alias for SET_MINDSPACE.
285    pub async fn create_mindspace(&self, args: serde_json::Value) -> anyhow::Result<serde_json::Value> {
286        self.command("SET_MINDSPACE", args).await
287    }
288
289    /// Deletes a mindspace and all its associated context.
290    pub async fn delete_mindspace(&self, args: serde_json::Value) -> anyhow::Result<serde_json::Value> {
291        self.command("DELETE_MINDSPACE", args).await
292    }
293
294    /// Sends a message to a mindspace and receives an AI-generated response.
295    pub async fn chat_mindspace(&self, args: serde_json::Value) -> anyhow::Result<serde_json::Value> {
296        self.command("CHAT_MINDSPACE", args).await
297    }
298
299    /// Imports text corpus into a mindspace for semantic search and context retrieval.
300    pub async fn lecture_mindspace(&self, args: serde_json::Value) -> anyhow::Result<serde_json::Value> {
301        self.command("LECTURE_MINDSPACE", args).await
302    }
303}
304
305use futures_util::SinkExt;
306use futures_util::StreamExt;