From acddece5205c34338ad8519ab98f2420a9ddb0dc Mon Sep 17 00:00:00 2001
From: sigoden <sigoden@gmail.com>
Date: Fri, 17 Mar 2023 15:09:06 +0800
Subject: [PATCH] feat: add `config.organization_id` (#77)

---
 src/client.rs     | 10 ++++++++--
 src/config/mod.rs | 24 ++++++++++++------------
 2 files changed, 20 insertions(+), 14 deletions(-)

diff --git a/src/client.rs b/src/client.rs
index 14d65d0..d366509 100644
--- a/src/client.rs
+++ b/src/client.rs
@@ -153,12 +153,18 @@ impl ChatGptClient {
                 .and_then(|m| m.insert("stream".into(), json!(true)));
         }
 
-        let builder = self
+        let (api_key, organization_id) = self.config.read().get_api_key();
+
+        let mut builder = self
             .build_client()?
             .post(API_URL)
-            .bearer_auth(self.config.read().get_api_key())
+            .bearer_auth(api_key)
             .json(&body);
 
+        if let Some(organization_id) = organization_id {
+            builder = builder.header("OpenAI-Organization", organization_id);
+        }
+
         Ok(builder)
     }
 }
diff --git a/src/config/mod.rs b/src/config/mod.rs
index 1607711..147a817 100644
--- a/src/config/mod.rs
+++ b/src/config/mod.rs
@@ -32,8 +32,7 @@ const CONFIG_FILE_NAME: &str = "config.yaml";
 const ROLES_FILE_NAME: &str = "roles.yaml";
 const HISTORY_FILE_NAME: &str = "history.txt";
 const MESSAGE_FILE_NAME: &str = "messages.md";
-const SET_COMPLETIONS: [&str; 9] = [
-    ".set api_key",
+const SET_COMPLETIONS: [&str; 8] = [
     ".set temperature",
     ".set save true",
     ".set save false",
@@ -49,6 +48,8 @@ const SET_COMPLETIONS: [&str; 9] = [
 pub struct Config {
     /// Openai api key
     pub api_key: Option<String>,
+    /// Openai organization id
+    pub organization_id: Option<String>,
     /// Openai model
     #[serde(rename(serialize = "model", deserialize = "model"))]
     pub model_name: Option<String>,
@@ -85,6 +86,7 @@ impl Default for Config {
     fn default() -> Self {
         Self {
             api_key: None,
+            organization_id: None,
             model_name: None,
             temperature: None,
             save: false,
@@ -201,8 +203,10 @@ impl Config {
         Self::local_file(CONFIG_FILE_NAME)
     }
 
-    pub fn get_api_key(&self) -> &String {
-        self.api_key.as_ref().expect("api_key not set")
+    pub fn get_api_key(&self) -> (String, Option<String>) {
+        let api_key = self.api_key.as_ref().expect("api_key not set");
+        let organization_id = self.organization_id.as_ref();
+        (api_key.into(), organization_id.cloned())
     }
 
     pub fn roles_file() -> Result<PathBuf> {
@@ -324,11 +328,14 @@ impl Config {
             .temperature
             .map(|v| v.to_string())
             .unwrap_or("-".into());
+        let (api_key, organization_id) = self.get_api_key();
+        let organization_id = organization_id.unwrap_or("-".into());
         let items = vec![
             ("config_file", file_info(&Config::config_file()?)),
             ("roles_file", file_info(&Config::roles_file()?)),
             ("messages_file", file_info(&Config::messages_file()?)),
-            ("api_key", self.get_api_key().to_string()),
+            ("api_key", api_key),
+            ("organization_id", organization_id),
             ("model", self.model.0.to_string()),
             ("temperature", temperature),
             ("save", self.save.to_string()),
@@ -367,13 +374,6 @@ impl Config {
         let value = parts[1];
         let unset = value == "null";
         match key {
-            "api_key" => {
-                if unset {
-                    bail!("Error: Not allowed");
-                } else {
-                    self.api_key = Some(value.to_string());
-                }
-            }
             "temperature" => {
                 if unset {
                     self.temperature = None;