refactor: support future llms

pythops · pythops · commit dab1b2cdeebd · 2023-04-25T23:13:16.000+02:00
diff --git a/src/app.rs b/src/app.rs
@@ -1,10 +1,12 @@
 use std;
 use std::collections::HashMap;
 
-use crate::config::AppConfig;
+use crate::config::Config;
 use crate::notification::Notification;
 use crossterm::event::KeyCode;
 
+use std::sync::Arc;
+
 pub type AppResult<T> = std::result::Result<T, Box<dyn std::error::Error>>;
 
 #[derive(Debug)]
@@ -31,16 +33,16 @@ pub struct App {
     pub previous_key: KeyCode,
     pub focused_block: FocusedBlock,
     pub show_help_popup: bool,
-    pub gpt_messages: Vec<HashMap<String, String>>,
+    pub llm_messages: Vec<HashMap<String, String>>,
     pub history: Vec<Vec<String>>,
     pub show_history_popup: bool,
     pub history_thread_index: usize,
-    pub config: AppConfig,
+    pub config: Arc<Config>,
     pub notifications: Vec<Notification>,
 }
 
-impl Default for App {
-    fn default() -> Self {
+impl App {
+    pub fn new(config: Arc<Config>) -> Self {
         Self {
             running: true,
             prompt: String::from(">_ "),
@@ -50,20 +52,14 @@ impl Default for App {
             previous_key: KeyCode::Null,
             focused_block: FocusedBlock::Prompt,
             show_help_popup: false,
-            gpt_messages: Vec::new(),
+            llm_messages: Vec::new(),
             history: Vec::new(),
             show_history_popup: false,
             history_thread_index: 0,
-            config: AppConfig::load(),
+            config,
             notifications: Vec::new(),
         }
     }
-}
-
-impl App {
-    pub fn new() -> Self {
-        Self::default()
-    }
 
     pub fn tick(&mut self) {
         self.notifications.retain(|n| n.ttl > 0);
diff --git a/src/config.rs b/src/config.rs
@@ -4,42 +4,50 @@ use dirs;
 use serde::Deserialize;
 
 #[derive(Deserialize, Debug)]
-pub struct AppConfig {
+pub struct Config {
     #[serde(default = "default_archive_file_name")]
     pub archive_file_name: String,
 
     #[serde(default)]
     pub key_bindings: KeyBindings,
 
     #[serde(default)]
-    pub gpt: GPT,
+    pub gpt: GPTConfig,
 }
 
 pub fn default_archive_file_name() -> String {
     String::from("tenere.archive")
 }
 
 #[derive(Deserialize, Debug, Clone)]
-pub struct GPT {
+pub struct GPTConfig {
     pub openai_api_key: Option<String>,
 
-    #[serde(default = "GPT::default_model")]
+    #[serde(default = "GPTConfig::default_model")]
     pub model: String,
+
+    #[serde(default = "GPTConfig::default_url")]
+    pub url: String,
 }
 
-impl Default for GPT {
+impl Default for GPTConfig {
     fn default() -> Self {
         Self {
             openai_api_key: None,
             model: String::from("gpt-3.5-turbo"),
+            url: String::from("https://api.openai.com/v1/chat/completions"),
         }
     }
 }
 
-impl GPT {
+impl GPTConfig {
     pub fn default_model() -> String {
         String::from("gpt-3.5-turbo")
     }
+
+    pub fn default_url() -> String {
+        String::from("https://api.openai.com/v1/chat/completions")
+    }
 }
 
 #[derive(Deserialize, Debug)]
@@ -86,15 +94,15 @@ impl KeyBindings {
     }
 }
 
-impl AppConfig {
+impl Config {
     pub fn load() -> Self {
         let conf_path = dirs::config_dir()
             .unwrap()
             .join("tenere")
             .join("config.toml");
 
         let config = std::fs::read_to_string(conf_path).unwrap_or(String::new());
-        let app_config: AppConfig = toml::from_str(&config).unwrap();
+        let app_config: Config = toml::from_str(&config).unwrap();
         app_config
     }
 }
diff --git a/src/event.rs b/src/event.rs
@@ -11,7 +11,7 @@ pub enum Event {
     Key(KeyEvent),
     Mouse(MouseEvent),
     Resize(u16, u16),
-    GPTResponse(String),
+    LLMAnswer(String),
     Notification(Notification),
 }
 
diff --git a/src/gpt.rs b/src/gpt.rs
@@ -1,3 +1,5 @@
+use crate::config::GPTConfig;
+use crate::llm::LLM;
 use reqwest::header::HeaderMap;
 use serde_json::{json, Value};
 use std;
@@ -7,13 +9,15 @@ use std::collections::HashMap;
 pub struct GPT {
     client: reqwest::blocking::Client,
     openai_api_key: String,
+    url: String,
 }
 
 impl GPT {
-    pub fn new(api_key: Option<String>) -> Self {
+    pub fn new(config: GPTConfig) -> Self {
         let openai_api_key = match std::env::var("OPENAI_API_KEY") {
             Ok(key) => key,
-            Err(_) => api_key
+            Err(_) => config
+                .openai_api_key
                 .ok_or_else(|| {
                     eprintln!(
                         r#"Can not find the openai api key
@@ -28,15 +32,16 @@ You need to define one wether in the configuration file or as an environment var
         Self {
             client: reqwest::blocking::Client::new(),
             openai_api_key,
+            url: config.url,
         }
     }
+}
 
-    pub fn ask(
+impl LLM for GPT {
+    fn ask(
         &self,
         chat_messages: Vec<HashMap<String, String>>,
     ) -> Result<String, Box<dyn std::error::Error>> {
-        let url = "https://api.openai.com/v1/chat/completions";
-
         let mut headers = HeaderMap::new();
         headers.insert("Content-Type", "application/json".parse().unwrap());
         headers.insert(
@@ -61,7 +66,12 @@ You need to define one wether in the configuration file or as an environment var
             "messages": messages
         });
 
-        let response = self.client.post(url).headers(headers).json(&body).send()?;
+        let response = self
+            .client
+            .post(&self.url)
+            .headers(headers)
+            .json(&body)
+            .send()?;
 
         match response.error_for_status() {
             Ok(res) => {
diff --git a/src/handler.rs b/src/handler.rs
@@ -1,8 +1,9 @@
 use crate::{
     app::{App, AppResult, FocusedBlock, Mode},
     event::Event,
-    gpt::GPT,
 };
+
+use crate::llm::LLM;
 use crossterm::event::{KeyCode, KeyEvent, KeyModifiers};
 use std::sync::mpsc::Sender;
 use std::{collections::HashMap, thread};
@@ -13,7 +14,7 @@ use std::sync::Arc;
 pub fn handle_key_events(
     key_event: KeyEvent,
     app: &mut App,
-    gpt: Arc<GPT>,
+    llm: Arc<impl LLM + 'static>,
     sender: Sender<Event>,
 ) -> AppResult<()> {
     match app.mode {
@@ -30,25 +31,25 @@ pub fn handle_key_events(
             }
 
             KeyCode::Enter => {
-                let mut conv: HashMap<String, String> = HashMap::new();
-
                 let user_input: String = app.prompt.drain(3..).collect();
                 let user_input = user_input.trim();
                 if user_input.is_empty() {
                     return Ok(());
                 }
                 app.chat.push(format!(" : {}\n", user_input));
 
-                conv.insert("role".to_string(), "user".to_string());
-                conv.insert("content".to_string(), user_input.to_string());
-                app.gpt_messages.push(conv.clone());
+                let conv = HashMap::from([
+                    ("role".into(), "user".into()),
+                    ("content".into(), user_input.into()),
+                ]);
+                app.llm_messages.push(conv);
 
-                let gpt_messages = app.gpt_messages.clone();
+                let llm_messages = app.llm_messages.clone();
 
                 thread::spawn(move || {
-                    let response = gpt.ask(gpt_messages.to_vec());
+                    let response = llm.ask(llm_messages.to_vec());
                     sender
-                        .send(Event::GPTResponse(match response {
+                        .send(Event::LLMAnswer(match response {
                             Ok(answer) => answer,
                             Err(e) => e.to_string(),
                         }))
@@ -89,7 +90,7 @@ pub fn handle_key_events(
                 app.prompt = String::from(">_ ");
                 app.history.push(app.chat.clone());
                 app.chat = Vec::new();
-                app.gpt_messages = Vec::new();
+                app.llm_messages = Vec::new();
                 app.scroll = 0;
             }
 
diff --git a/src/lib.rs b/src/lib.rs
@@ -15,3 +15,5 @@ pub mod config;
 pub mod ui;
 
 pub mod notification;
+
+pub mod llm;
diff --git a/src/llm.rs b/src/llm.rs
@@ -0,0 +1,25 @@
+use crate::config::Config;
+use crate::gpt::GPT;
+use std::collections::HashMap;
+
+use std::sync::Arc;
+pub trait LLM: Send + Sync {
+    fn ask(
+        &self,
+        chat_messages: Vec<HashMap<String, String>>,
+    ) -> Result<String, Box<dyn std::error::Error>>;
+}
+
+pub enum LLMBackend {
+    ChatGPT,
+}
+
+pub struct LLMModel {}
+
+impl LLMModel {
+    pub fn init(model: LLMBackend, config: Arc<Config>) -> impl LLM {
+        match model {
+            LLMBackend::ChatGPT => GPT::new(config.gpt.clone()),
+        }
+    }
+}
diff --git a/src/main.rs b/src/main.rs
@@ -2,22 +2,25 @@ use std::collections::HashMap;
 use std::{env, io};
 use tenere::app::{App, AppResult};
 use tenere::cli;
+use tenere::config::Config;
 use tenere::event::{Event, EventHandler};
-use tenere::gpt::GPT;
 use tenere::handler::handle_key_events;
 use tenere::tui::Tui;
 use tui::backend::CrosstermBackend;
 use tui::Terminal;
 
+use tenere::llm::{LLMBackend, LLMModel};
+
 use std::sync::Arc;
 
 use clap::crate_version;
 
 fn main() -> AppResult<()> {
     cli::cli().version(crate_version!()).get_matches();
 
-    let mut app = App::new();
-    let gpt = Arc::new(GPT::new(app.config.gpt.openai_api_key.clone()));
+    let config = Arc::new(Config::load());
+    let mut app = App::new(config.clone());
+    let llm = Arc::new(LLMModel::init(LLMBackend::ChatGPT, config));
 
     let backend = CrosstermBackend::new(io::stderr());
     let terminal = Terminal::new(backend)?;
@@ -30,18 +33,18 @@ fn main() -> AppResult<()> {
         match tui.events.next()? {
             Event::Tick => app.tick(),
             Event::Key(key_event) => {
-                handle_key_events(key_event, &mut app, gpt.clone(), tui.events.sender.clone())?
+                handle_key_events(key_event, &mut app, llm.clone(), tui.events.sender.clone())?
             }
             Event::Mouse(_) => {}
             Event::Resize(_, _) => {}
-            Event::GPTResponse(response) => {
+            Event::LLMAnswer(answer) => {
                 app.chat.pop();
-                app.chat.push(format!("🤖: {}\n", response));
+                app.chat.push(format!("🤖: {}\n", answer));
                 app.chat.push("\n".to_string());
                 let mut conv: HashMap<String, String> = HashMap::new();
                 conv.insert("role".to_string(), "user".to_string());
-                conv.insert("content".to_string(), response.clone());
-                app.gpt_messages.push(conv);
+                conv.insert("content".to_string(), answer);
+                app.llm_messages.push(conv);
             }
             Event::Notification(notification) => {
                 app.notifications.push(notification);

Original file line number	Diff line number	Diff line change
`@@ -4,42 +4,50 @@ use dirs;`
`4`	`4`	`use serde::Deserialize;`
`5`	`5`
`6`	`6`	`#[derive(Deserialize, Debug)]`
`7`		`-pub struct AppConfig {`
	`7`	`+pub struct Config {`
`8`	`8`	`#[serde(default = "default_archive_file_name")]`
`9`	`9`	`pub archive_file_name: String,`
`10`	`10`
`11`	`11`	`#[serde(default)]`
`12`	`12`	`pub key_bindings: KeyBindings,`
`13`	`13`
`14`	`14`	`#[serde(default)]`
`15`		`- pub gpt: GPT,`
	`15`	`+ pub gpt: GPTConfig,`
`16`	`16`	`}`
`17`	`17`
`18`	`18`	`pub fn default_archive_file_name() -> String {`
`19`	`19`	`String::from("tenere.archive")`
`20`	`20`	`}`
`21`	`21`
`22`	`22`	`#[derive(Deserialize, Debug, Clone)]`
`23`		`-pub struct GPT {`
	`23`	`+pub struct GPTConfig {`
`24`	`24`	`pub openai_api_key: Option<String>,`
`25`	`25`
`26`		`- #[serde(default = "GPT::default_model")]`
	`26`	`+ #[serde(default = "GPTConfig::default_model")]`
`27`	`27`	`pub model: String,`
	`28`	`+`
	`29`	`+ #[serde(default = "GPTConfig::default_url")]`
	`30`	`+ pub url: String,`
`28`	`31`	`}`
`29`	`32`
`30`		`-impl Default for GPT {`
	`33`	`+impl Default for GPTConfig {`
`31`	`34`	`fn default() -> Self {`
`32`	`35`	`Self {`
`33`	`36`	`openai_api_key: None,`
`34`	`37`	`model: String::from("gpt-3.5-turbo"),`
	`38`	`+ url: String::from("https://api.openai.com/v1/chat/completions"),`
`35`	`39`	`}`
`36`	`40`	`}`
`37`	`41`	`}`
`38`	`42`
`39`		`-impl GPT {`
	`43`	`+impl GPTConfig {`
`40`	`44`	`pub fn default_model() -> String {`
`41`	`45`	`String::from("gpt-3.5-turbo")`
`42`	`46`	`}`
	`47`	`+`
	`48`	`+ pub fn default_url() -> String {`
	`49`	`+ String::from("https://api.openai.com/v1/chat/completions")`
	`50`	`+ }`
`43`	`51`	`}`
`44`	`52`
`45`	`53`	`#[derive(Deserialize, Debug)]`
`@@ -86,15 +94,15 @@ impl KeyBindings {`
`86`	`94`	`}`
`87`	`95`	`}`
`88`	`96`
`89`		`-impl AppConfig {`
	`97`	`+impl Config {`
`90`	`98`	`pub fn load() -> Self {`
`91`	`99`	`let conf_path = dirs::config_dir()`
`92`	`100`	`.unwrap()`
`93`	`101`	`.join("tenere")`
`94`	`102`	`.join("config.toml");`
`95`	`103`
`96`	`104`	`let config = std::fs::read_to_string(conf_path).unwrap_or(String::new());`
`97`		`- let app_config: AppConfig = toml::from_str(&config).unwrap();`
	`105`	`+ let app_config: Config = toml::from_str(&config).unwrap();`
`98`	`106`	`app_config`
`99`	`107`	`}`
`100`	`108`	`}`
Original file line number	Diff line number	Diff line change
`@@ -11,7 +11,7 @@ pub enum Event {`
`11`	`11`	`Key(KeyEvent),`
`12`	`12`	`Mouse(MouseEvent),`
`13`	`13`	`Resize(u16, u16),`
`14`		`- GPTResponse(String),`
	`14`	`+ LLMAnswer(String),`
`15`	`15`	`Notification(Notification),`
`16`	`16`	`}`
`17`	`17`