From ddac3d5386ab01f290e782d53fc52d9051c6c960 Mon Sep 17 00:00:00 2001 From: Andrei Darashenka <6311180+darashenka@users.noreply.github.com> Date: Fri, 25 Jul 2025 18:14:35 +0200 Subject: [PATCH] honor OLLAMA_API_KEY for ollama_chat --- litellm/llms/ollama/chat/transformation.py | 2 + litellm/llms/ollama/common_utils.py | 21 ++++++++-- .../llms/ollama/completion/transformation.py | 18 +++++++++ .../llms/ollama/test_ollama_model_info.py | 38 ++++++++++++++++--- 4 files changed, 70 insertions(+), 9 deletions(-) diff --git a/litellm/llms/ollama/chat/transformation.py b/litellm/llms/ollama/chat/transformation.py index d4ce4052a7..2006ddf0e2 100644 --- a/litellm/llms/ollama/chat/transformation.py +++ b/litellm/llms/ollama/chat/transformation.py @@ -229,6 +229,8 @@ class OllamaChatConfig(BaseConfig): api_key: Optional[str] = None, api_base: Optional[str] = None, ) -> dict: + if api_key is not None and "Authorization" not in headers: + headers["Authorization"] = f"Bearer {api_key}" return headers def get_complete_url( diff --git a/litellm/llms/ollama/common_utils.py b/litellm/llms/ollama/common_utils.py index daff7a1206..166ceee27f 100644 --- a/litellm/llms/ollama/common_utils.py +++ b/litellm/llms/ollama/common_utils.py @@ -57,8 +57,20 @@ class OllamaModelInfo(BaseLLMModelInfo): """ @staticmethod - def get_api_key(api_key=None) -> None: - return None # Ollama does not use an API key by default + def get_api_key(api_key=None) -> Optional[str]: + """Get API key from environment variables or litellm configuration""" + import os + + import litellm + from litellm.secret_managers.main import get_secret_str + + return ( + os.environ.get("OLLAMA_API_KEY") + or litellm.api_key + or litellm.openai_key + or get_secret_str("OLLAMA_API_KEY") + ) + @staticmethod def get_api_base(api_base: Optional[str] = None) -> str: @@ -73,9 +85,12 @@ class OllamaModelInfo(BaseLLMModelInfo): """ base = self.get_api_base(api_base) + api_key = self.get_api_key() + headers = { "Authorization": f"Bearer {api_key}" } if api_key else {} + names: set[str] = set() try: - resp = httpx.get(f"{base}/api/tags") + resp = httpx.get(f"{base}/api/tags", headers=headers) resp.raise_for_status() data = resp.json() # Expecting a dict with a 'models' list diff --git a/litellm/llms/ollama/completion/transformation.py b/litellm/llms/ollama/completion/transformation.py index aa1da616d8..917e26878a 100644 --- a/litellm/llms/ollama/completion/transformation.py +++ b/litellm/llms/ollama/completion/transformation.py @@ -199,6 +199,21 @@ class OllamaConfig(BaseConfig): return v return None + @staticmethod + def get_api_key() -> Optional[str]: + """Get API key from environment variables or litellm configuration""" + import os + + import litellm + from litellm.secret_managers.main import get_secret_str + + return ( + os.environ.get("OLLAMA_API_KEY") + or litellm.api_key + or litellm.openai_key + or get_secret_str("OLLAMA_API_KEY") + ) + def get_model_info(self, model: str) -> ModelInfoBase: """ curl http://localhost:11434/api/show -d '{ @@ -208,11 +223,14 @@ class OllamaConfig(BaseConfig): if model.startswith("ollama/") or model.startswith("ollama_chat/"): model = model.split("/", 1)[1] api_base = get_secret_str("OLLAMA_API_BASE") or "http://localhost:11434" + api_key = self.get_api_key() + headers = { "Authorization": f"Bearer {api_key}" } if api_key else {} try: response = litellm.module_level_client.post( url=f"{api_base}/api/show", json={"name": model}, + headers=headers, ) except Exception as e: raise Exception( diff --git a/tests/test_litellm/llms/ollama/test_ollama_model_info.py b/tests/test_litellm/llms/ollama/test_ollama_model_info.py index adb079763d..b0aa464bf9 100644 --- a/tests/test_litellm/llms/ollama/test_ollama_model_info.py +++ b/tests/test_litellm/llms/ollama/test_ollama_model_info.py @@ -1,8 +1,6 @@ -import json import os import sys -import uuid -from unittest.mock import MagicMock, patch +from unittest.mock import patch import pytest @@ -56,6 +54,7 @@ class TestOllamaModelInfo: get_models should extract and return sorted unique model names. """ calls = [] + call_headers = [] sample = { "models": [ {"name": "zeta"}, @@ -65,8 +64,9 @@ class TestOllamaModelInfo: ] } - def mock_get(url): + def mock_get(url, headers): calls.append(url) + call_headers.append(headers) return DummyResponse(sample, status_code=200) monkeypatch.setattr(httpx, "get", mock_get) @@ -76,6 +76,32 @@ class TestOllamaModelInfo: assert models == ["alpha", "zeta"] # Ensure correct endpoint was called assert calls and calls[0].endswith("/api/tags") + assert call_headers and call_headers[0] == {} + + def test_get_models_from_dict_response_api_key(self, monkeypatch): + """ + When the /api/tags endpoint returns a dict with a 'models' list, + get_models should extract and return sorted unique model names. + """ + calls = [] + call_headers = [] + + def mock_get(url, headers): + calls.append(url) + call_headers.append(headers) + return DummyResponse({}, status_code=200) + + old_environ = dict(os.environ) + os.environ.update({"OLLAMA_API_KEY": "test_api_key"}) + monkeypatch.setattr(httpx, "get", mock_get) + info = OllamaModelInfo() + models = info.get_models() + os.environ.clear() + os.environ.update(old_environ) + assert models == [] + # Ensure correct endpoint was called + assert calls and calls[0].endswith("/api/tags") + assert call_headers and call_headers[0] == {'Authorization': 'Bearer test_api_key'} def test_get_models_from_list_response(self, monkeypatch): """ @@ -88,7 +114,7 @@ class TestOllamaModelInfo: {}, # no name/model key should be ignored ] - def mock_get(url): + def mock_get(url, headers): return DummyResponse(sample, status_code=200) monkeypatch.setattr(httpx, "get", mock_get) @@ -102,7 +128,7 @@ class TestOllamaModelInfo: fall back to the static models_by_provider list prefixed by 'ollama/'. """ - def mock_get(url): + def mock_get(url, headers): raise Exception("connection failure") monkeypatch.setattr(httpx, "get", mock_get)