lib.rs (13821B)
1 use async_openai::{ 2 config::OpenAIConfig, 3 types::{ChatCompletionRequestMessage, CreateChatCompletionRequest}, 4 Client, 5 }; 6 use chrono::{Duration, Local}; 7 use egui_wgpu::RenderState; 8 use enostr::KeypairUnowned; 9 use futures::StreamExt; 10 use nostrdb::Transaction; 11 use notedeck::{AppAction, AppContext, JobsCache}; 12 use std::collections::HashMap; 13 use std::string::ToString; 14 use std::sync::mpsc::{self, Receiver}; 15 use std::sync::Arc; 16 17 pub use avatar::DaveAvatar; 18 pub use config::ModelConfig; 19 pub use messages::{DaveApiResponse, Message}; 20 pub use quaternion::Quaternion; 21 pub use tools::{ 22 PartialToolCall, QueryCall, QueryResponse, Tool, ToolCall, ToolCalls, ToolResponse, 23 ToolResponses, 24 }; 25 pub use ui::{DaveAction, DaveResponse, DaveUi}; 26 pub use vec3::Vec3; 27 28 mod avatar; 29 mod config; 30 mod messages; 31 mod quaternion; 32 mod tools; 33 mod ui; 34 mod vec3; 35 36 pub struct Dave { 37 chat: Vec<Message>, 38 /// A 3d representation of dave. 39 avatar: Option<DaveAvatar>, 40 input: String, 41 tools: Arc<HashMap<String, Tool>>, 42 client: async_openai::Client<OpenAIConfig>, 43 incoming_tokens: Option<Receiver<DaveApiResponse>>, 44 model_config: ModelConfig, 45 jobs: JobsCache, 46 } 47 48 /// Calculate an anonymous user_id from a keypair 49 fn calculate_user_id(keypair: KeypairUnowned) -> String { 50 use sha2::{Digest, Sha256}; 51 // pubkeys have degraded privacy, don't do that 52 let key_input = keypair 53 .secret_key 54 .map(|sk| sk.as_secret_bytes()) 55 .unwrap_or(keypair.pubkey.bytes()); 56 let hex_key = hex::encode(key_input); 57 let input = format!("{hex_key}notedeck_dave_user_id"); 58 hex::encode(Sha256::digest(input)) 59 } 60 61 impl Dave { 62 pub fn avatar_mut(&mut self) -> Option<&mut DaveAvatar> { 63 self.avatar.as_mut() 64 } 65 66 fn system_prompt() -> Message { 67 let now = Local::now(); 68 let yesterday = now - Duration::hours(24); 69 let date = now.format("%Y-%m-%d %H:%M:%S"); 70 let timestamp = now.timestamp(); 71 let yesterday_timestamp = yesterday.timestamp(); 72 73 Message::System(format!( 74 r#" 75 You are an AI agent for the nostr protocol called Dave, created by Damus. nostr is a decentralized social media and internet communications protocol. You are embedded in a nostr browser called 'Damus Notedeck'. 76 77 - The current date is {date} ({timestamp} unix timestamp if needed for queries). 78 79 - Yesterday (-24hrs) was {yesterday_timestamp}. You can use this in combination with `since` queries for pulling notes for summarizing notes the user might have missed while they were away. 80 81 # Response Guidelines 82 83 - You *MUST* call the present_notes tool with a list of comma-separated note id references when referring to notes so that the UI can display them. Do *NOT* include note id references in the text response, but you *SHOULD* use ^1, ^2, etc to reference note indices passed to present_notes. 84 - When a user asks for a digest instead of specific query terms, make sure to include both since and until to pull notes for the correct range. 85 - When tasked with open-ended queries such as looking for interesting notes or summarizing the day, make sure to add enough notes to the context (limit: 100-200) so that it returns enough data for summarization. 86 "# 87 )) 88 } 89 90 pub fn new(render_state: Option<&RenderState>) -> Self { 91 let model_config = ModelConfig::default(); 92 //let model_config = ModelConfig::ollama(); 93 let client = Client::with_config(model_config.to_api()); 94 95 let input = "".to_string(); 96 let avatar = render_state.map(DaveAvatar::new); 97 let mut tools: HashMap<String, Tool> = HashMap::new(); 98 for tool in tools::dave_tools() { 99 tools.insert(tool.name().to_string(), tool); 100 } 101 102 Dave { 103 client, 104 avatar, 105 incoming_tokens: None, 106 tools: Arc::new(tools), 107 input, 108 model_config, 109 chat: vec![], 110 jobs: JobsCache::default(), 111 } 112 } 113 114 /// Process incoming tokens from the ai backend 115 fn process_events(&mut self, app_ctx: &AppContext) -> bool { 116 // Should we continue sending requests? Set this to true if 117 // we have tool responses to send back to the ai 118 let mut should_send = false; 119 120 let Some(recvr) = &self.incoming_tokens else { 121 return should_send; 122 }; 123 124 while let Ok(res) = recvr.try_recv() { 125 if let Some(avatar) = &mut self.avatar { 126 avatar.random_nudge(); 127 } 128 match res { 129 DaveApiResponse::Failed(err) => self.chat.push(Message::Error(err)), 130 131 DaveApiResponse::Token(token) => match self.chat.last_mut() { 132 Some(Message::Assistant(msg)) => *msg = msg.clone() + &token, 133 Some(_) => self.chat.push(Message::Assistant(token)), 134 None => {} 135 }, 136 137 DaveApiResponse::ToolCalls(toolcalls) => { 138 tracing::info!("got tool calls: {:?}", toolcalls); 139 self.chat.push(Message::ToolCalls(toolcalls.clone())); 140 141 let txn = Transaction::new(app_ctx.ndb).unwrap(); 142 for call in &toolcalls { 143 // execute toolcall 144 match call.calls() { 145 ToolCalls::PresentNotes(present) => { 146 self.chat.push(Message::ToolResponse(ToolResponse::new( 147 call.id().to_owned(), 148 ToolResponses::PresentNotes(present.note_ids.len() as i32), 149 ))); 150 151 should_send = true; 152 } 153 154 ToolCalls::Invalid(invalid) => { 155 should_send = true; 156 157 self.chat.push(Message::tool_error( 158 call.id().to_string(), 159 invalid.error.clone(), 160 )); 161 } 162 163 ToolCalls::Query(search_call) => { 164 should_send = true; 165 166 let resp = search_call.execute(&txn, app_ctx.ndb); 167 self.chat.push(Message::ToolResponse(ToolResponse::new( 168 call.id().to_owned(), 169 ToolResponses::Query(resp), 170 ))) 171 } 172 } 173 } 174 } 175 } 176 } 177 178 should_send 179 } 180 181 fn ui(&mut self, app_ctx: &mut AppContext, ui: &mut egui::Ui) -> DaveResponse { 182 DaveUi::new(self.model_config.trial, &self.chat, &mut self.input).ui( 183 app_ctx, 184 &mut self.jobs, 185 ui, 186 ) 187 } 188 189 fn handle_new_chat(&mut self) { 190 self.chat = vec![]; 191 self.input.clear(); 192 } 193 194 /// Handle a user send action triggered by the ui 195 fn handle_user_send(&mut self, app_ctx: &AppContext, ui: &egui::Ui) { 196 self.chat.push(Message::User(self.input.clone())); 197 self.send_user_message(app_ctx, ui.ctx()); 198 self.input.clear(); 199 } 200 201 fn send_user_message(&mut self, app_ctx: &AppContext, ctx: &egui::Context) { 202 let messages: Vec<ChatCompletionRequestMessage> = { 203 let txn = Transaction::new(app_ctx.ndb).expect("txn"); 204 self.chat 205 .iter() 206 .filter_map(|c| c.to_api_msg(&txn, app_ctx.ndb)) 207 .collect() 208 }; 209 tracing::debug!("sending messages, latest: {:?}", messages.last().unwrap()); 210 211 let user_id = calculate_user_id(app_ctx.accounts.get_selected_account().keypair()); 212 213 let ctx = ctx.clone(); 214 let client = self.client.clone(); 215 let tools = self.tools.clone(); 216 let model_name = self.model_config.model().to_owned(); 217 218 let (tx, rx) = mpsc::channel(); 219 self.incoming_tokens = Some(rx); 220 221 tokio::spawn(async move { 222 let mut token_stream = match client 223 .chat() 224 .create_stream(CreateChatCompletionRequest { 225 model: model_name, 226 stream: Some(true), 227 messages, 228 tools: Some(tools::dave_tools().iter().map(|t| t.to_api()).collect()), 229 user: Some(user_id), 230 ..Default::default() 231 }) 232 .await 233 { 234 Err(err) => { 235 tracing::error!("openai chat error: {err}"); 236 return; 237 } 238 239 Ok(stream) => stream, 240 }; 241 242 let mut all_tool_calls: HashMap<u32, PartialToolCall> = HashMap::new(); 243 244 while let Some(token) = token_stream.next().await { 245 let token = match token { 246 Ok(token) => token, 247 Err(err) => { 248 tracing::error!("failed to get token: {err}"); 249 let _ = tx.send(DaveApiResponse::Failed(err.to_string())); 250 return; 251 } 252 }; 253 254 for choice in &token.choices { 255 let resp = &choice.delta; 256 257 // if we have tool call arg chunks, collect them here 258 if let Some(tool_calls) = &resp.tool_calls { 259 for tool in tool_calls { 260 let entry = all_tool_calls.entry(tool.index).or_default(); 261 262 if let Some(id) = &tool.id { 263 entry.id_mut().get_or_insert(id.clone()); 264 } 265 266 if let Some(name) = tool.function.as_ref().and_then(|f| f.name.as_ref()) 267 { 268 entry.name_mut().get_or_insert(name.to_string()); 269 } 270 271 if let Some(argchunk) = 272 tool.function.as_ref().and_then(|f| f.arguments.as_ref()) 273 { 274 entry 275 .arguments_mut() 276 .get_or_insert_with(String::new) 277 .push_str(argchunk); 278 } 279 } 280 } 281 282 if let Some(content) = &resp.content { 283 if let Err(err) = tx.send(DaveApiResponse::Token(content.to_owned())) { 284 tracing::error!("failed to send dave response token to ui: {err}"); 285 } 286 ctx.request_repaint(); 287 } 288 } 289 } 290 291 let mut parsed_tool_calls = vec![]; 292 for (_index, partial) in all_tool_calls { 293 let Some(unknown_tool_call) = partial.complete() else { 294 tracing::error!("could not complete partial tool call: {:?}", partial); 295 continue; 296 }; 297 298 match unknown_tool_call.parse(&tools) { 299 Ok(tool_call) => { 300 parsed_tool_calls.push(tool_call); 301 } 302 Err(err) => { 303 // TODO: we should be 304 tracing::error!( 305 "failed to parse tool call {:?}: {}", 306 unknown_tool_call, 307 err, 308 ); 309 310 if let Some(id) = partial.id() { 311 // we have an id, so we can communicate the error 312 // back to the ai 313 parsed_tool_calls.push(ToolCall::invalid( 314 id.to_string(), 315 partial.name, 316 partial.arguments, 317 err.to_string(), 318 )); 319 } 320 } 321 }; 322 } 323 324 if !parsed_tool_calls.is_empty() { 325 tx.send(DaveApiResponse::ToolCalls(parsed_tool_calls)) 326 .unwrap(); 327 ctx.request_repaint(); 328 } 329 330 tracing::debug!("stream closed"); 331 }); 332 } 333 } 334 335 impl notedeck::App for Dave { 336 fn update(&mut self, ctx: &mut AppContext<'_>, ui: &mut egui::Ui) -> Option<AppAction> { 337 /* 338 self.app 339 .frame_history 340 .on_new_frame(ctx.input(|i| i.time), frame.info().cpu_usage); 341 */ 342 let mut app_action: Option<AppAction> = None; 343 344 // always insert system prompt if we have no context 345 if self.chat.is_empty() { 346 self.chat.push(Dave::system_prompt()); 347 } 348 349 //update_dave(self, ctx, ui.ctx()); 350 let should_send = self.process_events(ctx); 351 if let Some(action) = self.ui(ctx, ui).action { 352 match action { 353 DaveAction::ToggleChrome => { 354 app_action = Some(AppAction::ToggleChrome); 355 } 356 DaveAction::Note(n) => { 357 app_action = Some(AppAction::Note(n)); 358 } 359 DaveAction::NewChat => { 360 self.handle_new_chat(); 361 } 362 DaveAction::Send => { 363 self.handle_user_send(ctx, ui); 364 } 365 } 366 } 367 368 if should_send { 369 self.send_user_message(ctx, ui.ctx()); 370 } 371 372 app_action 373 } 374 }