diff --git a/cortex/ask.py b/cortex/ask.py index 33c06351..c66971d9 100644 --- a/cortex/ask.py +++ b/cortex/ask.py @@ -12,6 +12,8 @@ import subprocess from typing import Any +from cortex.config_utils import get_ollama_model + class SystemInfoGatherer: """Gathers local system information for context-aware responses.""" @@ -169,11 +171,18 @@ def _default_model(self) -> str: elif self.provider == "claude": return "claude-sonnet-4-20250514" elif self.provider == "ollama": - return "llama3.2" + return self._get_ollama_model() elif self.provider == "fake": return "fake" return "gpt-4" + def _get_ollama_model(self) -> str: + """Determine which Ollama model to use. + + Delegates to the shared ``get_ollama_model()`` utility function. + """ + return get_ollama_model() + def _initialize_client(self): if self.provider == "openai": try: diff --git a/cortex/config_utils.py b/cortex/config_utils.py new file mode 100644 index 00000000..6f523816 --- /dev/null +++ b/cortex/config_utils.py @@ -0,0 +1,48 @@ +"""Configuration utilities for Cortex. + +This module provides shared configuration helpers used across the codebase. +""" + +import json +import os +from pathlib import Path + + +def get_ollama_model() -> str: + """Determine which Ollama model to use. + + The model name is resolved using the following precedence: + + 1. If the ``OLLAMA_MODEL`` environment variable is set, its value is + returned. + 2. Otherwise, if ``~/.cortex/config.json`` exists and contains an + ``"ollama_model"`` key, that value is returned. + 3. If neither of the above sources provides a model name, the + hard-coded default ``"llama3.2"`` is used. + + Any errors encountered while reading or parsing the configuration + file are silently ignored, and the resolution continues to the next + step in the precedence chain. + + Returns: + The Ollama model name to use. + """ + # Try environment variable first + env_model = os.environ.get("OLLAMA_MODEL") + if env_model: + return env_model + + # Try config file + try: + config_file = Path.home() / ".cortex" / "config.json" + if config_file.exists(): + with open(config_file, encoding="utf-8") as f: + config = json.load(f) + model = config.get("ollama_model") + if model: + return model + except (OSError, json.JSONDecodeError): + pass # Ignore file/parse errors + + # Default to llama3.2 + return "llama3.2" diff --git a/cortex/dependency_check.py b/cortex/dependency_check.py index d42e610f..1c070076 100644 --- a/cortex/dependency_check.py +++ b/cortex/dependency_check.py @@ -43,7 +43,7 @@ def format_installation_instructions(missing: list[str]) -> str: ╰─────────────────────────────────────────────────────────────────╯ Cortex requires the following packages that are not installed: - {', '.join(missing)} + {", ".join(missing)} To fix this, run ONE of the following: diff --git a/cortex/first_run_wizard.py b/cortex/first_run_wizard.py index c31f9fb0..6c7d22f9 100644 --- a/cortex/first_run_wizard.py +++ b/cortex/first_run_wizard.py @@ -371,16 +371,43 @@ def _setup_ollama(self) -> StepResult: print("\n✗ Failed to install Ollama") return StepResult(success=True, data={"api_provider": "none"}) - # Pull a small model - print("\nPulling llama3.2 model (this may take a few minutes)...") + # Let user choose model or use default + print("\nWhich Ollama model would you like to use?") + print(" 1. llama3.2 (2GB) - Recommended for most users") + print(" 2. llama3.2:1b (1.3GB) - Faster, less RAM") + print(" 3. mistral (4GB) - Alternative quality model") + print(" 4. phi3 (2.3GB) - Microsoft's efficient model") + print(" 5. Custom (enter your own)") + + model_choices = { + "1": "llama3.2", + "2": "llama3.2:1b", + "3": "mistral", + "4": "phi3", + } + + choice = self._prompt("\nEnter choice [1]: ", default="1") + + if choice == "5": + model_name = self._prompt("Enter model name: ", default="llama3.2") + elif choice in model_choices: + model_name = model_choices[choice] + else: + print(f"Invalid choice '{choice}', using default model llama3.2") + model_name = "llama3.2" + + # Pull the selected model + print(f"\nPulling {model_name} model (this may take a few minutes)...") try: - subprocess.run(["ollama", "pull", "llama3.2"], check=True) + subprocess.run(["ollama", "pull", model_name], check=True) print("\n✓ Model ready!") except subprocess.CalledProcessError: - print("\n⚠ Could not pull model - you can do this later with: ollama pull llama3.2") + print( + f"\n⚠ Could not pull model - you can do this later with: ollama pull {model_name}" + ) self.config["api_provider"] = "ollama" - self.config["ollama_model"] = "llama3.2" + self.config["ollama_model"] = model_name return StepResult(success=True, data={"api_provider": "ollama"}) diff --git a/cortex/llm/interpreter.py b/cortex/llm/interpreter.py index 74870d75..069771b8 100644 --- a/cortex/llm/interpreter.py +++ b/cortex/llm/interpreter.py @@ -4,6 +4,8 @@ from enum import Enum from typing import TYPE_CHECKING, Any, Optional +from cortex.config_utils import get_ollama_model + if TYPE_CHECKING: from cortex.semantic_cache import SemanticCache @@ -67,28 +69,11 @@ def __init__( self._initialize_client() def _get_ollama_model(self) -> str: - """Get Ollama model from config file or environment.""" - # Try environment variable first - env_model = os.environ.get("OLLAMA_MODEL") - if env_model: - return env_model - - # Try config file - try: - from pathlib import Path + """Get Ollama model from config file or environment. - config_file = Path.home() / ".cortex" / "config.json" - if config_file.exists(): - with open(config_file) as f: - config = json.load(f) - model = config.get("ollama_model") - if model: - return model - except Exception: - pass # Ignore errors reading config - - # Default to llama3.2 - return "llama3.2" + Delegates to the shared ``get_ollama_model()`` utility function. + """ + return get_ollama_model() def _initialize_client(self): if self.provider == APIProvider.OPENAI: @@ -112,7 +97,8 @@ def _initialize_client(self): ollama_base_url = os.environ.get("OLLAMA_BASE_URL", "http://localhost:11434") self.client = OpenAI( - api_key="ollama", base_url=f"{ollama_base_url}/v1" # Dummy key, not used + api_key="ollama", + base_url=f"{ollama_base_url}/v1", # Dummy key, not used ) except ImportError: raise ImportError("OpenAI package not installed. Run: pip install openai") diff --git a/cortex/sandbox/docker_sandbox.py b/cortex/sandbox/docker_sandbox.py index 71e57fc8..ca0073fc 100644 --- a/cortex/sandbox/docker_sandbox.py +++ b/cortex/sandbox/docker_sandbox.py @@ -250,8 +250,7 @@ def require_docker(self) -> str: ) if result.returncode != 0: raise DockerNotFoundError( - "Docker daemon is not running.\n" - "Start Docker with: sudo systemctl start docker" + "Docker daemon is not running.\nStart Docker with: sudo systemctl start docker" ) except subprocess.TimeoutExpired: raise DockerNotFoundError("Docker daemon is not responding.") diff --git a/tests/test_ask.py b/tests/test_ask.py index 0fe53176..88133493 100644 --- a/tests/test_ask.py +++ b/tests/test_ask.py @@ -4,6 +4,7 @@ import sys import tempfile import unittest +from pathlib import Path from unittest.mock import MagicMock, patch sys.path.insert(0, os.path.join(os.path.dirname(__file__), "..")) @@ -240,8 +241,35 @@ def test_default_model_claude(self): def test_default_model_ollama(self): """Test default model for Ollama.""" + # Test with environment variable + + # Save and clear any existing OLLAMA_MODEL + original_model = os.environ.get("OLLAMA_MODEL") + + # Test with custom env variable + os.environ["OLLAMA_MODEL"] = "test-model" handler = AskHandler(api_key="test", provider="ollama") - self.assertEqual(handler.model, "llama3.2") + self.assertEqual(handler.model, "test-model") + + # Clean up + if original_model is not None: + os.environ["OLLAMA_MODEL"] = original_model + else: + os.environ.pop("OLLAMA_MODEL", None) + + # Test deterministic default behavior when no env var or config file exists. + # Point the home directory to a temporary location without ~/.cortex/config.json + # Also ensure OLLAMA_MODEL is not set in the environment so get_ollama_model() + # exercises the built-in default model lookup. + env_without_ollama = {k: v for k, v in os.environ.items() if k != "OLLAMA_MODEL"} + with ( + tempfile.TemporaryDirectory() as tmpdir, + patch("cortex.config_utils.Path.home", return_value=Path(tmpdir)), + patch.dict(os.environ, env_without_ollama, clear=True), + ): + handler2 = AskHandler(api_key="test", provider="ollama") + # When no env var and no config file exist, AskHandler should use its built-in default. + self.assertEqual(handler2.model, "llama3.2") def test_default_model_fake(self): """Test default model for fake provider.""" diff --git a/tests/test_cli.py b/tests/test_cli.py index 093a8e50..bed29ab4 100644 --- a/tests/test_cli.py +++ b/tests/test_cli.py @@ -81,14 +81,16 @@ def test_print_success(self, mock_stdout): self.assertTrue(True) @patch.dict(os.environ, {}, clear=True) - def test_install_no_api_key(self): - # When no API key is set, the CLI falls back to Ollama. - # If Ollama is running, this should succeed. If not, it should fail. - # We'll mock Ollama to be unavailable to test the failure case. - with patch("cortex.llm.interpreter.CommandInterpreter.parse") as mock_parse: - mock_parse.side_effect = RuntimeError("Ollama not available") - result = self.cli.install("docker") - self.assertEqual(result, 1) + @patch("cortex.cli.CommandInterpreter") + def test_install_no_api_key(self, mock_interpreter_class): + # Should work with Ollama (no API key needed) + mock_interpreter = Mock() + mock_interpreter.parse.return_value = ["apt update", "apt install docker"] + mock_interpreter_class.return_value = mock_interpreter + + result = self.cli.install("docker") + # Should succeed with Ollama as fallback provider + self.assertEqual(result, 0) @patch.dict(os.environ, {"OPENAI_API_KEY": "sk-test-openai-key-123"}, clear=True) @patch("cortex.cli.CommandInterpreter") diff --git a/tests/test_ollama_integration.py b/tests/test_ollama_integration.py index fb688d86..f5b0a1ef 100755 --- a/tests/test_ollama_integration.py +++ b/tests/test_ollama_integration.py @@ -174,6 +174,11 @@ def test_routing_decision(): """Test routing logic with Ollama.""" print("4. Testing routing decision...") + # Get available model + test_model = os.environ.get("OLLAMA_MODEL") or get_available_ollama_model() + if not test_model: + pytest.skip("No Ollama models available") + try: router = LLMRouter( ollama_base_url="http://localhost:11434", @@ -204,6 +209,11 @@ def test_stats_tracking(): """Test that stats tracking works with Ollama.""" print("5. Testing stats tracking...") + # Get available model + test_model = os.environ.get("OLLAMA_MODEL") or get_available_ollama_model() + if not test_model: + pytest.skip("No Ollama models available") + try: router = LLMRouter( ollama_base_url="http://localhost:11434",