Merge branch 'Bowser1704-feat/openai-host-base'

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
 
 <img width="1022" alt="Screen Shot 2023-03-07 at 09 30 53" src="https://user-images.githubusercontent.com/10192522/223295925-00eed881-cdfc-4f46-9510-1e0bd1c99e60.png">
@@ -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=<your 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:

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<Message>,
 }
 
 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);