feat: support multi bots and custom url (#150)

1 files changed, 19 insertions, 9 deletions

diff --git a/src/main.rs b/src/main.rs
index b201fd6..ec8abee 100644
--- a/src/main.rs
+++ b/src/main.rs
@@ -8,11 +8,12 @@ mod term;
 mod utils;
 
 use crate::cli::Cli;
-use crate::client::ChatGptClient;
+use crate::client::Client;
 use crate::config::{Config, SharedConfig};
 
-use anyhow::{anyhow, Result};
+use anyhow::{anyhow, Context, Result};
 use clap::Parser;
+use client::{init_client, list_models};
 use crossbeam::sync::WaitGroup;
 use is_terminal::IsTerminal;
 use parking_lot::RwLock;
@@ -21,6 +22,7 @@ use repl::{AbortSignal, Repl};
 use std::io::{stdin, Read};
 use std::sync::Arc;
 use std::{io::stdout, process::exit};
+use tokio::runtime::Runtime;
 use utils::cl100k_base_singleton;
 
 fn main() -> Result<()> {
@@ -36,8 +38,8 @@ fn main() -> Result<()> {
         exit(0);
     }
     if cli.list_models {
-        for (name, _) in &config::MODELS {
-            println!("{name}");
+        for model in list_models(&config.read()) {
+            println!("{}", model.stringify());
         }
         exit(0);
     }
@@ -69,24 +71,25 @@ fn main() -> Result<()> {
         exit(0);
     }
     let no_stream = cli.no_stream;
-    let client = ChatGptClient::init(config.clone())?;
+    let runtime = init_runtime()?;
+    let client = init_client(config.clone(), runtime)?;
     if atty::isnt(atty::Stream::Stdin) {
         let mut input = String::new();
         stdin().read_to_string(&mut input)?;
         if let Some(text) = text {
             input = format!("{text}\n{input}");
         }
-        start_directive(&client, &config, &input, no_stream)
+        start_directive(client.as_ref(), &config, &input, no_stream)
     } else {
         match text {
-            Some(text) => start_directive(&client, &config, &text, no_stream),
+            Some(text) => start_directive(client.as_ref(), &config, &text, no_stream),
             None => start_interactive(client, config),
         }
     }
 }
 
 fn start_directive(
-    client: &ChatGptClient,
+    client: &dyn Client,
     config: &SharedConfig,
     input: &str,
     no_stream: bool,
@@ -120,9 +123,16 @@ fn start_directive(
     config.read().save_message(input, &output)
 }
 
-fn start_interactive(client: ChatGptClient, config: SharedConfig) -> Result<()> {
+fn start_interactive(client: Box<dyn Client>, config: SharedConfig) -> Result<()> {
     cl100k_base_singleton();
     config.write().on_repl()?;
     let mut repl = Repl::init(config.clone())?;
     repl.run(client, config)
 }
+
+fn init_runtime() -> Result<Runtime> {
+    tokio::runtime::Builder::new_current_thread()
+        .enable_all()
+        .build()
+        .with_context(|| "Failed to init tokio")
+}