From f548922490cbbd5d9d519671e858a8ddb127ed04 Mon Sep 17 00:00:00 2001 From: Tanishq <74068132+tanishqbaweja@users.noreply.github.com> Date: Sun, 25 Oct 2026 23:44:09 +0530 Subject: [PATCH] fix(ollama): honor explicit API bases without losing global defaults --- .../get_llm_provider_logic.py | 2 +- litellm/main.py | 6 +- tests/unit/test_main.py | 64 +++++++++++++++++++ 3 files changed, 68 insertions(+), 4 deletions(-) diff --git a/litellm/litellm_core_utils/get_llm_provider_logic.py b/litellm/litellm_core_utils/get_llm_provider_logic.py index de51ecd0b21..9ca67b66eb1 100644 --- a/litellm/litellm_core_utils/get_llm_provider_logic.py +++ b/litellm/litellm_core_utils/get_llm_provider_logic.py @@ -697,7 +697,7 @@ def _get_openai_compatible_provider_info( api_base = api_base or get_secret("NEBIUS_API_BASE") or "https://api.studio.nebius.ai/v1" dynamic_api_key = api_key or get_secret_str("NEBIUS_API_KEY") elif custom_llm_provider == "ollama": - api_base = api_base or get_secret("OLLAMA_API_BASE") or "http://localhost:11434" + api_base = api_base or litellm.api_base or get_secret("OLLAMA_API_BASE") or "http://localhost:11434" dynamic_api_key = api_key or get_secret_str("OLLAMA_API_KEY") elif (custom_llm_provider == "ai21_chat") or (custom_llm_provider == "ai21" and model in litellm.ai21_chat_models): api_base = api_base or get_secret("AI21_API_BASE") or "https://api.ai21.com/studio/v1" diff --git a/litellm/main.py b/litellm/main.py index f30d93361e5..76d2ea8641d 100644 --- a/litellm/main.py +++ b/litellm/main.py @@ -4597,7 +4597,7 @@ def _complete_ollama(ctx: CompletionDispatchContext) -> _CompletionDispatchResul stream: Final = ctx.stream timeout: Final = ctx.timeout - api_base = litellm.api_base or api_base or get_secret("OLLAMA_API_BASE") or "http://localhost:11434" + api_base = api_base or litellm.api_base or get_secret("OLLAMA_API_BASE") or "http://localhost:11434" if api_key is not None and "Authorization" not in headers: headers["Authorization"] = f"Bearer {api_key}" @@ -4637,7 +4637,7 @@ def _complete_ollama_chat(ctx: CompletionDispatchContext) -> _CompletionDispatch stream: Final = ctx.stream timeout: Final = ctx.timeout - api_base = litellm.api_base or api_base or get_secret("OLLAMA_API_BASE") or "http://localhost:11434" + api_base = api_base or litellm.api_base or get_secret("OLLAMA_API_BASE") or "http://localhost:11434" api_key = api_key or litellm.ollama_key or os.environ.get("OLLAMA_API_KEY") or litellm.api_key if api_key is not None and "Authorization" not in headers: @@ -7280,7 +7280,7 @@ def embedding( api_key=api_key, ) elif custom_llm_provider == "ollama": - api_base = litellm.api_base or api_base or get_secret_str("OLLAMA_API_BASE") or "http://localhost:11434" + api_base = api_base or litellm.api_base or get_secret_str("OLLAMA_API_BASE") or "http://localhost:11434" if isinstance(input, str): input = [input] diff --git a/tests/unit/test_main.py b/tests/unit/test_main.py index 70b268bf2a2..8b2a02513fe 100644 --- a/tests/unit/test_main.py +++ b/tests/unit/test_main.py @@ -95,6 +95,70 @@ def preserve_litellm_completion_state(monkeypatch: pytest.MonkeyPatch) -> None: WHITE_PNG: Final = (Path(__file__).parents[1] / "white_100x100.png").read_bytes() +@pytest.mark.parametrize("asynchronous", [False, True]) +@pytest.mark.parametrize( + ("provider", "endpoint"), + [("ollama", "generate"), ("ollama_chat", "chat"), ("ollama", "embed")], +) +@pytest.mark.parametrize( + ("explicit_base", "global_base", "env_base", "expected_base"), + [ + ("http://explicit.example", "http://global.example", "http://env.example", "http://explicit.example"), + (None, "http://global.example", "http://env.example", "http://global.example"), + ("", "http://global.example", "http://env.example", "http://global.example"), + (None, None, "http://env.example", "http://env.example"), + (None, None, None, "http://localhost:11434"), + ], +) +@pytest.mark.asyncio +async def test_ollama_requests_respect_api_base_precedence( + asynchronous: bool, + provider: str, + endpoint: str, + explicit_base: str | None, + global_base: str | None, + env_base: str | None, + expected_base: str, + monkeypatch: pytest.MonkeyPatch, + respx_mock: respx.MockRouter, +) -> None: + monkeypatch.setattr(litellm, "api_base", global_base) + monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) + monkeypatch.setenv("OLLAMA_API_BASE", env_base or "") + respx_mock.post(url__regex=r".*/api/show").respond(json={"model_info": {}, "capabilities": ["completion"]}) + route: Final = respx_mock.post(url__regex=rf".*/api/{endpoint}").respond( + json={ + "model": "routing-test", + "created_at": "2026-01-01T00:00:00Z", + "response": "OK", + "message": {"role": "assistant", "content": "OK"}, + "embeddings": [[0.1, 0.2, 0.3]], + "done": True, + "done_reason": "stop", + "prompt_eval_count": 3, + "eval_count": 1, + } + ) + args: Final = {"model": f"{provider}/routing-test", "api_base": explicit_base} + if endpoint == "embed": + response: Final = ( + await litellm.aembedding(**args, input="hello") + if asynchronous + else litellm.embedding(**args, input="hello") + ) + assert response.data == [{"object": "embedding", "index": 0, "embedding": [0.1, 0.2, 0.3]}] + else: + messages: Final = [{"role": "user", "content": "hello"}] + completion_response: Final = ( + await litellm.acompletion(**args, messages=messages) + if asynchronous + else litellm.completion(**args, messages=messages) + ) + assert completion_response.choices[0].message.content == "OK" + assert route.call_count == 1 + assert str(route.calls[0].request.url) == f"{expected_base}/api/{endpoint}" + + @pytest.mark.parametrize("asynchronous", [False, True]) @pytest.mark.parametrize("stream", [False, True]) @pytest.mark.parametrize(