diff --git a/README.md b/README.md index 988b43d..b613d89 100644 --- a/README.md +++ b/README.md @@ -2,8 +2,8 @@ A simple common-line interface for [ChatGPT API](https://platform.openai.com/docs/api-reference/chat/create). -- 🌟 Streaming output! -- 💡 One-shot mode to get a quick answer +- 🌟 Streaming output! +- 💡 One-shot mode to get a quick answer - 🤖 Interactive mode to have a conversation Screen Shot 2023-03-07 at 09 30 53 @@ -20,15 +20,17 @@ cargo install --path . You'll need a OpenAI API key (you can get one [here](https://platform.openai.com/account/api-keys)), and you'll need to export your API Key as an environment variable: +You can also set a OpenAI API base environment variable, just like [openai-python](https://github.com/openai/openai-python/blob/main/openai/__init__.py#L37) ```bash export OPENAI_API_KEY= +# export OPENAI_API_BASE="https://api.openai.com/v1" ``` Then you can start an interactive conversation with ChatGPT: ```bash -heygpt +heygpt ``` OR use the one-shot mode by providing a prompt: diff --git a/src/main.rs b/src/main.rs index 02c84e1..6d8419b 100644 --- a/src/main.rs +++ b/src/main.rs @@ -51,6 +51,7 @@ We generally recommend altering this or temperature but not both."# const READLINE_HISTORY: &str = ".heygpt_history"; const OPENAI_API_KEY: &str = "OPENAI_API_KEY"; +const OPENAI_API_BASE: &str = "OPENAI_API_BASE"; #[tokio::main(flavor = "current_thread")] async fn main() -> Result<()> { @@ -62,10 +63,12 @@ async fn main() -> Result<()> { let api_key = std::env::var(OPENAI_API_KEY).map_err(|_| anyhow!("{} not set", OPENAI_API_KEY))?; + let api_base = std::env::var(OPENAI_API_BASE).unwrap_or("https://api.openai.com/v1".into()); + // Enter interactive mode if prompt is empty let interactive = options.prompt.is_empty(); - let mut session = Session::new(options, api_key); + let mut session = Session::new(options, api_key, api_base); if !interactive { session.run_one_shot().await?; } else { @@ -82,15 +85,19 @@ struct Session { /// OpenAI API key api_key: String, + /// OpenAI API base URL + api_base: String, + /// Messages history messages: Vec, } impl Session { - pub fn new(options: Options, api_key: String) -> Self { + pub fn new(options: Options, api_key: String, api_base: String) -> Self { Self { options, api_key, + api_base, messages: Vec::new(), } } @@ -178,7 +185,7 @@ impl Session { let client = Client::new(); let req = client - .post("https://api.openai.com/v1/chat/completions".to_string()) + .post(format!("{}/chat/completions", &self.api_base)) .headers(headers) .json(&data);