notedeck

One damus client to rule them all
git clone git://jb55.com/notedeck
Log | Files | Refs | README | LICENSE

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 }