Merge pull request #104 from getmetal/task/sp/allow-openai-env-variable

Add `OPENAI_API_BASE` config
getmetal · Dec 19, 2023 · 2631be7 · 2631be7
2 parents dc4faeb + 4e52a01
commit 2631be7
Show file tree

Hide file tree

Showing 2 changed files with 19 additions and 4 deletions.
diff --git a/README.md b/README.md
@@ -98,6 +98,7 @@ Searches are segmented (filtered) by the session id provided automatically.
 - `PORT` (default:8000) - Motorhead Server Port
 - `OPENAI_API_KEY`- [Your api key](https://platform.openai.com/account/api-keys) to connect to OpenAI.
 - `REDIS_URL` (required)- URL used to connect to `redis`.
+- `OPENAI_API_BASE` (default:https://api.openai.com/v1) - OpenAI API Base URL
 
 ### Azure deployment
 

diff --git a/src/models.rs b/src/models.rs
@@ -55,10 +55,24 @@ impl Manager for OpenAIClientManager {
                     completion_client: Client::with_config(config),
                 }
             }
-            _ => AnyOpenAIClient::OpenAI {
-                embedding_client: Client::new(),
-                completion_client: Client::new(),
-            },
+            _ => {
+                let openai_api_base = env::var("OPENAI_API_BASE");
+
+                if let Ok(openai_api_base) = openai_api_base {
+                    let embedding_config = OpenAIConfig::default().with_api_base(&openai_api_base);
+                    let completion_config = OpenAIConfig::default().with_api_base(&openai_api_base);
+
+                    AnyOpenAIClient::OpenAI {
+                        embedding_client: Client::with_config(embedding_config),
+                        completion_client: Client::with_config(completion_config),
+                    }
+                } else {
+                    AnyOpenAIClient::OpenAI {
+                        embedding_client: Client::new(),
+                        completion_client: Client::new(),
+                    }
+                }
+            }
         };
         Ok(openai_client)
     }