diff --git a/graphify/llm.py b/graphify/llm.py index 48546b865..e4743c36e 100644 --- a/graphify/llm.py +++ b/graphify/llm.py @@ -56,6 +56,22 @@ def _get_tokenizer(): # Cached at import time. None if tiktoken is unavailable; consumers must handle. _TOKENIZER = _get_tokenizer() + +def _resolve_ollama_base_url(default: str) -> str: + ollama_base_url = os.environ.get("OLLAMA_BASE_URL") + if ollama_base_url is not None: + return ollama_base_url + ollama_host = os.environ.get("OLLAMA_HOST") + if ollama_host is None: + return default + if not ollama_host.startswith(("http://", "https://")): + ollama_host = f"http://{ollama_host}" + ollama_host = ollama_host.rstrip("/") + if not ollama_host.endswith("/v1"): + ollama_host = f"{ollama_host}/v1" + return ollama_host + + BACKENDS: dict[str, dict] = { "claude": { # ANTHROPIC_BASE_URL points the backend at any Anthropic-compatible @@ -83,7 +99,7 @@ def _get_tokenizer(): "max_tokens": 16384, }, "ollama": { - "base_url": os.environ.get("OLLAMA_BASE_URL", "http://localhost:11434/v1"), + "base_url": _resolve_ollama_base_url("http://localhost:11434/v1"), "default_model": os.environ.get("OLLAMA_MODEL", "qwen2.5-coder:7b"), "env_key": "OLLAMA_API_KEY", "pricing": {"input": 0.0, "output": 0.0}, @@ -1572,7 +1588,7 @@ def extract_files_direct( # Ollama ignores auth but the OpenAI client library requires a non-empty # string. Use a placeholder and surface a visible warning so this never # silently routes traffic without the user realising — see F-029. - ollama_url = os.environ.get("OLLAMA_BASE_URL", cfg.get("base_url", "")) + ollama_url = _resolve_ollama_base_url(cfg.get("base_url", "")) _validate_ollama_base_url(ollama_url) print( "[graphify] WARNING: ollama backend selected with no OLLAMA_API_KEY set; " @@ -2373,7 +2389,7 @@ def _call_llm( cfg = BACKENDS[backend] key = _get_backend_api_key(backend) if not key and backend == "ollama": - ollama_url = os.environ.get("OLLAMA_BASE_URL", cfg.get("base_url", "")) + ollama_url = _resolve_ollama_base_url(cfg.get("base_url", "")) _validate_ollama_base_url(ollama_url) key = "ollama" if not key and backend not in ("bedrock", "claude-cli"): diff --git a/tests/test_llm_backends.py b/tests/test_llm_backends.py index e6556ba4a..c1392ded5 100644 --- a/tests/test_llm_backends.py +++ b/tests/test_llm_backends.py @@ -22,6 +22,34 @@ def _clear_backend_env(monkeypatch): monkeypatch.delenv(env_key, raising=False) +def test_resolve_ollama_base_url_prefers_base_url(monkeypatch): + monkeypatch.setenv("OLLAMA_BASE_URL", "custom-base-url") + monkeypatch.setenv("OLLAMA_HOST", "ignored-host:11434") + + assert llm._resolve_ollama_base_url("default-url") == "custom-base-url" + + +def test_resolve_ollama_base_url_normalizes_host_without_scheme(monkeypatch): + monkeypatch.delenv("OLLAMA_BASE_URL", raising=False) + monkeypatch.setenv("OLLAMA_HOST", "myhost:11434") + + assert llm._resolve_ollama_base_url("default-url") == "http://myhost:11434/v1" + + +def test_resolve_ollama_base_url_preserves_normalized_host(monkeypatch): + monkeypatch.delenv("OLLAMA_BASE_URL", raising=False) + monkeypatch.setenv("OLLAMA_HOST", "https://myhost:11434/v1") + + assert llm._resolve_ollama_base_url("default-url") == "https://myhost:11434/v1" + + +def test_resolve_ollama_base_url_returns_default_without_env(monkeypatch): + monkeypatch.delenv("OLLAMA_BASE_URL", raising=False) + monkeypatch.delenv("OLLAMA_HOST", raising=False) + + assert llm._resolve_ollama_base_url("default-url") == "default-url" + + def test_gemini_accepts_gemini_api_key(monkeypatch): _clear_backend_env(monkeypatch) monkeypatch.setenv("GEMINI_API_KEY", "gemini-key")