Add OpenAI model setup and configuration updates

- Implemented methods to retrieve model definitions and infer embedding dimensions for OpenAI-compatible models.
- Enhanced the `_upsert_openai_models` method to update or create model entries in the configuration, ensuring defaults are set correctly.
- Updated the setup process to prompt users for custom API settings, including base URL and model names, improving flexibility in configuration.
- Added unit tests to validate the OpenAI setup flow, ensuring proper handling of custom API initialization and configuration updates.

Author: 0xrushi

backends/advanced/init.py

@@ -164,6 +164,92 @@ def prompt_with_existing_masked(self, prompt_text: str, env_key: str, placeholde
             default=default
         )
 
+    def _get_model_def(self, config: Dict[str, Any], model_name: str) -> Dict[str, Any]:
+        """Get a model definition by name from config.yml."""
+        models = config.get("models", [])
+        if not isinstance(models, list):
+            return {}
+        return next((m for m in models if m.get("name") == model_name), {})
+
+    def _infer_embedding_dimensions(self, model_name: str, fallback: int = 1536) -> int:
+        """Infer embedding dimensions for common models."""
+        known_dimensions = {
+            "text-embedding-3-small": 1536,
+            "text-embedding-3-large": 3072,
+            "text-embedding-ada-002": 1536,
+            "nomic-embed-text-v1.5": 768,
+            "nomic-embed-text:latest": 768,
+        }
+        return known_dimensions.get(model_name, fallback)
+
+    def _upsert_openai_models(
+        self,
+        api_key: str,
+        base_url: str,
+        llm_model_name: str,
+        embedding_model_name: str,
+    ) -> None:
+        """Update or create openai-llm/openai-embed in config.yml and set defaults."""
+        config = self.config_manager.get_full_config()
+        models = config.get("models", [])
+        if not isinstance(models, list):
+            models = []
+
+        openai_llm = self._get_model_def(config, "openai-llm")
+        openai_embed = self._get_model_def(config, "openai-embed")
+
+        llm_params = openai_llm.get("model_params", {})
+        if not isinstance(llm_params, dict):
+            llm_params = {}
+        llm_params.setdefault("temperature", 0.2)
+        llm_params.setdefault("max_tokens", 2000)
+
+        embedding_dimensions = openai_embed.get("embedding_dimensions")
+        if not isinstance(embedding_dimensions, int) or embedding_dimensions <= 0:
+            embedding_dimensions = self._infer_embedding_dimensions(embedding_model_name)
+
+        llm_payload = {
+            "name": "openai-llm",
+            "description": "OpenAI/OpenAI-compatible LLM",
+            "model_type": "llm",
+            "model_provider": "openai",
+            "api_family": "openai",
+            "model_name": llm_model_name,
+            "model_url": base_url,
+            "api_key": api_key,
+            "model_params": llm_params,
+            "model_output": "json",
+        }
+        embed_payload = {
+            "name": "openai-embed",
+            "description": "OpenAI/OpenAI-compatible embeddings",
+            "model_type": "embedding",
+            "model_provider": "openai",
+            "api_family": "openai",
+            "model_name": embedding_model_name,
+            "model_url": base_url,
+            "api_key": api_key,
+            "embedding_dimensions": embedding_dimensions,
+            "model_output": "vector",
+        }
+
+        def upsert_model(payload: Dict[str, Any]):
+            for idx, model in enumerate(models):
+                if model.get("name") == payload["name"]:
+                    models[idx] = {**model, **payload}
+                    return
+            models.append(payload)
+
+        upsert_model(llm_payload)
+        upsert_model(embed_payload)
+
+        config["models"] = models
+        if "defaults" not in config or not isinstance(config["defaults"], dict):
+            config["defaults"] = {}
+        config["defaults"]["llm"] = "openai-llm"
+        config["defaults"]["embedding"] = "openai-embed"
+
+        self.config_manager.save_full_config(config)
 
     def setup_authentication(self):
         """Configure authentication settings"""
@@ -307,17 +393,29 @@ def setup_llm(self):
         self.console.print()
 
         choices = {
-            "1": "OpenAI (GPT-4, GPT-3.5 - requires API key)",
+            "1": "OpenAI / OpenAI-compatible (custom base URL, API key, model names)",
             "2": "Ollama (local models - runs locally)",
             "3": "Skip (no memory extraction)"
         }
 
         choice = self.prompt_choice("Which LLM provider will you use?", choices, "1")
 
         if choice == "1":
-            self.console.print("[blue][INFO][/blue] OpenAI selected")
+            self.console.print("[blue][INFO][/blue] OpenAI/OpenAI-compatible selected")
             self.console.print("Get your API key from: https://platform.openai.com/api-keys")
 
+            existing_cfg = self.config_manager.get_full_config()
+            openai_llm = self._get_model_def(existing_cfg, "openai-llm")
+            openai_embed = self._get_model_def(existing_cfg, "openai-embed")
+
+            default_base_url = openai_llm.get("model_url") or openai_embed.get("model_url") or "https://api.openai.com/v1"
+            default_llm_model = openai_llm.get("model_name") or "gpt-4o-mini"
+            default_embedding_model = openai_embed.get("model_name") or "text-embedding-3-small"
+
+            base_url = self.prompt_value("OpenAI-compatible base URL", default_base_url)
+            llm_model_name = self.prompt_value("LLM model name", default_llm_model)
+            embedding_model_name = self.prompt_value("Embedding model name", default_embedding_model)
+
             # Use the new masked prompt function
             api_key = self.prompt_with_existing_masked(
                 prompt_text="OpenAI API key (leave empty to skip)",
@@ -329,11 +427,19 @@ def setup_llm(self):
 
             if api_key:
                 self.config["OPENAI_API_KEY"] = api_key
-                # Update config.yml to use OpenAI models
-                self.config_manager.update_config_defaults({"llm": "openai-llm", "embedding": "openai-embed"})
+                # Update config.yml openai model definitions and defaults
+                self._upsert_openai_models(
+                    api_key=api_key,
+                    base_url=base_url,
+                    llm_model_name=llm_model_name,
+                    embedding_model_name=embedding_model_name,
+                )
                 self.console.print("[green][SUCCESS][/green] OpenAI configured in config.yml")
                 self.console.print("[blue][INFO][/blue] Set defaults.llm: openai-llm")
                 self.console.print("[blue][INFO][/blue] Set defaults.embedding: openai-embed")
+                self.console.print(f"[blue][INFO][/blue] Set openai-llm.model_url: {base_url}")
+                self.console.print(f"[blue][INFO][/blue] Set openai-llm.model_name: {llm_model_name}")
+                self.console.print(f"[blue][INFO][/blue] Set openai-embed.model_name: {embedding_model_name}")
             else:
                 self.console.print("[yellow][WARNING][/yellow] No API key provided - memory extraction will not work")
 

backends/advanced/tests/test_init_llm_setup.py

@@ -0,0 +1,136 @@
+"""Unit tests for OpenAI custom API setup/initialization flow in init.py.
+
+These tests verify that wizard setup can initialize OpenAI-compatible providers
+with custom API endpoints (base URL), API keys, and model names.
+"""
+
+import importlib.util
+from pathlib import Path
+from unittest.mock import MagicMock, call
+
+
+_INIT_PATH = Path(__file__).resolve().parents[1] / "init.py"
+_SPEC = importlib.util.spec_from_file_location("advanced_init", _INIT_PATH)
+_MODULE = importlib.util.module_from_spec(_SPEC)
+assert _SPEC and _SPEC.loader
+_SPEC.loader.exec_module(_MODULE)
+ChronicleSetup = _MODULE.ChronicleSetup
+
+
+def _build_setup_with_mocks() -> ChronicleSetup:
+    """Create ChronicleSetup instance without running __init__ side effects."""
+    setup = ChronicleSetup.__new__(ChronicleSetup)
+    setup.console = MagicMock()
+    setup.config = {}
+    setup.config_manager = MagicMock()
+    setup.print_section = MagicMock()
+    return setup
+
+
+def test_upsert_openai_models_updates_existing_defs_and_defaults():
+    """Checks OpenAI custom API config upsert in init flow.
+
+    Verifies that init setup updates OpenAI model definitions with custom API
+    settings and switches defaults to those OpenAI entries.
+    """
+    setup = _build_setup_with_mocks()
+    setup.config_manager.get_full_config.return_value = {
+        "defaults": {"llm": "local-llm", "embedding": "local-embed"},
+        "models": [
+            {
+                "name": "openai-llm",
+                "model_type": "llm",
+                "model_provider": "openai",
+                "model_name": "gpt-4o-mini",
+                "model_url": "https://api.openai.com/v1",
+                "api_key": "old-key",
+                "model_params": {"temperature": 0.3},
+            },
+            {
+                "name": "openai-embed",
+                "model_type": "embedding",
+                "model_provider": "openai",
+                "model_name": "text-embedding-3-small",
+                "model_url": "https://api.openai.com/v1",
+                "api_key": "old-key",
+                "embedding_dimensions": 1536,
+            },
+        ],
+    }
+
+    setup._upsert_openai_models(
+        api_key="new-key",
+        base_url="http://custom.example/v1",
+        llm_model_name="gpt-oss-20b",
+        embedding_model_name="text-embedding-3-large",
+    )
+
+    saved_config = setup.config_manager.save_full_config.call_args[0][0]
+    saved_models = {m["name"]: m for m in saved_config["models"]}
+
+    assert saved_config["defaults"]["llm"] == "openai-llm"
+    assert saved_config["defaults"]["embedding"] == "openai-embed"
+
+    assert saved_models["openai-llm"]["model_url"] == "http://custom.example/v1"
+    assert saved_models["openai-llm"]["api_key"] == "new-key"
+    assert saved_models["openai-llm"]["model_name"] == "gpt-oss-20b"
+    # Existing params are preserved and missing defaults are filled.
+    assert saved_models["openai-llm"]["model_params"]["temperature"] == 0.3
+    assert saved_models["openai-llm"]["model_params"]["max_tokens"] == 2000
+
+    assert saved_models["openai-embed"]["model_url"] == "http://custom.example/v1"
+    assert saved_models["openai-embed"]["api_key"] == "new-key"
+    assert saved_models["openai-embed"]["model_name"] == "text-embedding-3-large"
+    # Existing embedding dimensions are preserved.
+    assert saved_models["openai-embed"]["embedding_dimensions"] == 1536
+
+
+def test_setup_llm_openai_prompts_for_custom_values_and_updates_models():
+    """Checks init OpenAI setup prompts for custom API initialization values."""
+    setup = _build_setup_with_mocks()
+    setup.prompt_choice = MagicMock(return_value="1")
+    setup.prompt_value = MagicMock(
+        side_effect=["http://my-openai-compatible/v1", "my-chat-model", "my-embed-model"]
+    )
+    setup.prompt_with_existing_masked = MagicMock(return_value="test-api-key")
+    setup._upsert_openai_models = MagicMock()
+    setup.config_manager.get_full_config.return_value = {
+        "models": [
+            {"name": "openai-llm", "model_url": "https://api.openai.com/v1", "model_name": "gpt-4o-mini"},
+            {"name": "openai-embed", "model_url": "https://api.openai.com/v1", "model_name": "text-embedding-3-small"},
+        ]
+    }
+
+    setup.setup_llm()
+
+    setup.prompt_value.assert_has_calls(
+        [
+            call("OpenAI-compatible base URL", "https://api.openai.com/v1"),
+            call("LLM model name", "gpt-4o-mini"),
+            call("Embedding model name", "text-embedding-3-small"),
+        ]
+    )
+    setup._upsert_openai_models.assert_called_once_with(
+        api_key="test-api-key",
+        base_url="http://my-openai-compatible/v1",
+        llm_model_name="my-chat-model",
+        embedding_model_name="my-embed-model",
+    )
+    assert setup.config["OPENAI_API_KEY"] == "test-api-key"
+
+
+def test_setup_llm_openai_skips_upsert_when_api_key_missing():
+    """Checks init OpenAI custom API setup guards against missing API key."""
+    setup = _build_setup_with_mocks()
+    setup.prompt_choice = MagicMock(return_value="1")
+    setup.prompt_value = MagicMock(
+        side_effect=["https://api.openai.com/v1", "gpt-4o-mini", "text-embedding-3-small"]
+    )
+    setup.prompt_with_existing_masked = MagicMock(return_value="")
+    setup._upsert_openai_models = MagicMock()
+    setup.config_manager.get_full_config.return_value = {"models": []}
+
+    setup.setup_llm()
+
+    setup._upsert_openai_models.assert_not_called()
+    assert "OPENAI_API_KEY" not in setup.config