Add support for customizing OpenAI endpoint (#34)

* Bump `orch` to `0.0.16`

* Ignore `.DS_Store` files

* Add ability to customize OpenAI base endpoint

* Fix subcommand search tests

* Add test

.github/workflows/ci.yml

@@ -66,7 +66,7 @@ jobs:
       - name: Run E2E tests
         run: ./scripts/e2e-tests.sh
         env:
-          OPENAI_API_KEY_E2E: ${{ secrets.OPENAI_API_KEY_E2E }}
+          OPENAI_API_KEY: ${{ secrets.OPENAI_API_KEY_E2E }}
 
   linting:
     name: Linting and formatting

.gitignore

@@ -11,4 +11,7 @@ __pycache__
 .pytest_cache
 
 .env*.local
-temp
+temp
+
+# macOS
+.DS_Store

Cargo.toml

@@ -7,8 +7,8 @@ repository = "https://github.com/guywaldman/magic-cli"
 edition = "2021"
 
 [dependencies]
-orch = { version = "0.0.15" }
-orch_response = { version = "0.0.15" }                 # Will be bundled inside `orch` (#10)
+orch = { version = "0.0.16" }
+orch_response = { version = "0.0.16" }                 # Will be bundled inside `orch` (#10)
 chrono = "0.4.38"
 clap = { version = "4.5.7", features = ["derive"] }
 clipboard = "0.5.0"

src/cli/config/keys.rs

@@ -192,6 +192,20 @@ impl ConfigKeys {
                     }),
                 ).secret()
             );
+            keys.insert(
+                "openai.api_endpoint".to_string(),
+                ConfigurationKey::new(
+                    "openai.api_endpoint".to_string(),
+                    "Custom API endpoint for the OpenAI API.".to_string(),
+                    Box::new(|config: &mut MagicCliConfig, value: &str| {
+                        if config.openai_config.is_none() {
+                            config.openai_config = Some(OpenAiConfig::default());
+                        }
+                        config.openai_config.as_mut().unwrap().api_endpoint = Some(value.to_string());
+                        Ok(())
+                    }),
+                ).secret()
+            );
             keys.insert(
                 "openai.model".to_string(),
                 ConfigurationKey::new(

src/cli/config/mcli_config.rs

@@ -237,14 +237,19 @@ impl MagicCliConfigManager {
                 let Some(api_key) = openai_config.api_key.clone() else {
                     return Err(MagicCliConfigError::MissingConfigKey("api_key".to_owned()));
                 };
-                let openai = OpenAiBuilder::new()
+                let mut openai_builder = OpenAiBuilder::new()
                     .with_model(model)
                     .with_embeddings_model(embedding_model)
-                    .with_api_key(api_key)
+                    .with_api_key(api_key);
+                if let Some(api_endpoint) = openai_config.api_endpoint.clone() {
+                    openai_builder = openai_builder.with_api_endpoint(api_endpoint);
+                }
+                let openai = openai_builder
                     .try_build()
                     .map_err(|e| MagicCliConfigError::Configuration(e.to_string()))?;
                 Ok(Box::new(openai))
             }
+            _ => Err(MagicCliConfigError::Configuration("Invalid LLM provider".to_string())),
         }
     }
 

src/cli/subcommand_search.rs

@@ -78,7 +78,7 @@ impl MagicCliSubcommand for SearchSubcommand {
             };
             if !allow_remote_llm {
                 return Err(Box::new(MagicCliConfigError::Configuration(
-                    "Using remote LLM but `allow_remote_llm` is set to false. Set it to `true` if you are willing for remote LLM providers such as OpenAI to embed your shell history which may contains sensitive information.".to_string(),
+                    "Using remote LLM but `search.allow_remote_llm` is set to false. Set it to `true` if you are willing for remote LLM providers such as OpenAI to embed your shell history which may contains sensitive information.".to_string(),
                 )));
             }
         }

src/lm/openai_config.rs

@@ -6,6 +6,7 @@ use crate::cli::config::{ConfigOptions, MagicCliConfigError};
 #[derive(Debug, Clone, Serialize, Deserialize)]
 pub struct OpenAiConfig {
     pub api_key: Option<String>,
+    pub api_endpoint: Option<String>,
     pub model: Option<String>,
     pub embedding_model: Option<String>,
 }
@@ -14,6 +15,7 @@ impl Default for OpenAiConfig {
     fn default() -> Self {
         Self {
             api_key: None,
+            api_endpoint: None,
             model: Some(openai_model::GPT_4O_MINI.to_string()),
             embedding_model: Some(openai_embedding_model::TEXT_EMBEDDING_ADA_002.to_string()),
         }
@@ -29,6 +31,10 @@ impl ConfigOptions for OpenAiConfig {
             populated = true;
             self.api_key = defaults.api_key;
         }
+        if self.api_endpoint.is_none() {
+            populated = true;
+            self.api_endpoint = defaults.api_endpoint;
+        }
         if self.model.is_none() {
             populated = true;
             self.model = defaults.model;

tests/test_config.py

@@ -68,3 +68,21 @@ def test_custom_config_set(self):
                 results = run_subcommand("config", ["get", "general.llm"], mcli_args=["--config", f.name])
                 assert results.status == 0
                 assert "openai" in results.stdout
+
+    def test_config_openai_api_endpoint_misconfigured(self):
+        """Tests that the config subcommand fails if the user tries to set the OpenAI API endpoint but the configuration is not set correctly."""
+        with tempfile.NamedTemporaryFile(mode="w", delete=False) as tmp:
+            with open(tmp.name, "w") as f:
+                f.write(json.dumps({"general": {"llm": "openai"}}))
+
+            results = run_subcommand(
+                "config", ["set", "--key", "openai.api_endpoint", "--value", "http://example.org"], mcli_args=["--config", f.name]
+            )
+            assert results.status == 0
+            results = run_subcommand(
+                "config", ["set", "--key", "openai.api_key", "--value", "non-existent-api-key"], mcli_args=["--config", f.name]
+            )
+            assert results.status == 0
+
+            results = run_subcommand("suggest", ["'Print the current directory using `ls`. Use only `ls`'"], mcli_args=["--config", f.name])
+            assert results.status != 0

tests/test_search.py

@@ -8,13 +8,21 @@ class TestSearchSubcommand:
     """Tests for the `search` subcommand, which allows users to semantically search for commands across their shell history."""
 
     def test_basic_search_openai(self):
+        """Tests basic search with OpenAI."""
         env = Env.from_env()
 
         with tempfile.TemporaryDirectory() as index_dir, tempfile.NamedTemporaryFile(
             mode="w", delete=False
         ) as config, tempfile.NamedTemporaryFile(mode="w", delete=False) as shell_history:
             with open(config.name, "w") as f:
-                f.write(json.dumps({"general": {"llm": "openai"}, "search": {"shell_history": shell_history.name, "index_dir": index_dir}}))
+                f.write(
+                    json.dumps(
+                        {
+                            "general": {"llm": "openai"},
+                            "search": {"shell_history": shell_history.name, "index_dir": index_dir, "allow_remote_llm": True},
+                        }
+                    )
+                )
 
             with open(shell_history.name, "w") as f:
                 f.write("echo foobar\n")
@@ -28,14 +36,22 @@ def test_basic_search_openai(self):
             assert results.status == 0
             assert "echo foobar" in results.stdout
 
-    def test_search_with_remote_llm(self):
+    def test_search_with_remote_llm_misconfigured(self):
+        """Tests that the search subcommand fails if the user tries to use a remote LLM but the configuration is not set correctly."""
         env = Env.from_env()
 
         with tempfile.TemporaryDirectory() as index_dir, tempfile.NamedTemporaryFile(
             mode="w", delete=False
         ) as config, tempfile.NamedTemporaryFile(mode="w", delete=False) as shell_history:
             with open(config.name, "w") as f:
-                f.write(json.dumps({"general": {"llm": "openai"}, "search": {"shell_history": shell_history.name, "index_dir": index_dir}}))
+                f.write(
+                    json.dumps(
+                        {
+                            "general": {"llm": "openai"},
+                            "search": {"shell_history": shell_history.name, "index_dir": index_dir},
+                        }
+                    )
+                )
 
             with open(shell_history.name, "w") as f:
                 f.write("echo foobar\n")
@@ -46,4 +62,5 @@ def test_search_with_remote_llm(self):
             set_config("openai.api_key", env.openai_api_key, config_path=config.name)
 
             results = run_subcommand("search", ["foobar", "--output-only"], mcli_args=["--config", config.name])
-            print(results)
+            assert results.status != 0
+            assert "search.allow_remote_llm" in results.stderr

tests/utils.py

@@ -61,8 +61,8 @@ def from_env(cls):
             print(f"Loading environment variables from '{env_path}'")
             dotenv.load_dotenv(dotenv_path=env_path)
 
-        openai_api_key = os.environ.get("OPENAI_API_KEY_E2E")
+        openai_api_key = os.environ.get("OPENAI_API_KEY")
         if openai_api_key is None:
-            raise Exception("OPENAI_API_KEY_E2E environment variable not set")
+            raise Exception("OPENAI_API_KEY environment variable not set")
 
         return cls(openai_api_key=openai_api_key)