1use 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 let url_parsed = Url::parse(&url)?;
28
29 let (ws_stream, _) = connect_async_with_config(url_parsed, None, false).await?;
31
32 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 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 }
74 Message::Pong(_) => {
75 }
77 _ => {
78 }
80 }
81 }
82 });
83
84 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 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 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 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 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 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 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 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 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 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 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 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 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 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 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 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;