diff --git a/portkey_ai/api_resources/global_constants.py b/portkey_ai/api_resources/global_constants.py index c9e95f60..9948b62e 100644 --- a/portkey_ai/api_resources/global_constants.py +++ b/portkey_ai/api_resources/global_constants.py @@ -44,6 +44,7 @@ PORTKEY_GATEWAY_URL = PORTKEY_BASE_URL LOCAL_BASE_URL = "http://localhost:8787/v1" PORTKEY_API_KEY_ENV = "PORTKEY_API_KEY" +PORTKEY_BASE_URL_ENV = "PORTKEY_BASE_URL" PORTKEY_PROXY_ENV = "PORTKEY_PROXY" OPEN_AI_API_KEY = "OPENAI_API_KEY" DEFAULT_CONNECTION_LIMITS = httpx.Limits( diff --git a/portkey_ai/api_resources/utils.py b/portkey_ai/api_resources/utils.py index bd124b8a..4f9a52bc 100644 --- a/portkey_ai/api_resources/utils.py +++ b/portkey_ai/api_resources/utils.py @@ -42,6 +42,7 @@ MISSING_MODE_MESSAGE, PORTKEY_BASE_URL, PORTKEY_API_KEY_ENV, + PORTKEY_BASE_URL_ENV, PORTKEY_HEADER_PREFIX, PORTKEY_PROXY_ENV, ) @@ -509,7 +510,9 @@ def set_base_url(base_url, api_key): if base_url: return base_url - env_base_url = os.environ.get(PORTKEY_BASE_URL) or os.environ.get(PORTKEY_PROXY_ENV) + env_base_url = os.environ.get(PORTKEY_BASE_URL_ENV) or os.environ.get( + PORTKEY_PROXY_ENV + ) if env_base_url: return env_base_url diff --git a/tests/test_set_base_url.py b/tests/test_set_base_url.py new file mode 100644 index 00000000..b7c7f6c4 --- /dev/null +++ b/tests/test_set_base_url.py @@ -0,0 +1,48 @@ +import pytest + +from portkey_ai.api_resources.global_constants import ( + LOCAL_BASE_URL, + PORTKEY_BASE_URL, +) +from portkey_ai.api_resources.utils import set_base_url + + +@pytest.fixture(autouse=True) +def clear_base_url_env(monkeypatch): + # set_base_url reads process env vars; isolate the test from the ambient + # environment so PORTKEY_BASE_URL / PORTKEY_PROXY / PORTKEY_API_KEY do not leak in. + monkeypatch.delenv("PORTKEY_BASE_URL", raising=False) + monkeypatch.delenv("PORTKEY_PROXY", raising=False) + monkeypatch.delenv("PORTKEY_API_KEY", raising=False) + + +class TestSetBaseUrl: + def test_explicit_base_url_takes_precedence(self, monkeypatch): + monkeypatch.setenv("PORTKEY_BASE_URL", "https://env-gateway.example/v1") + assert ( + set_base_url("https://explicit.example/v1", None) + == "https://explicit.example/v1" + ) + + def test_portkey_base_url_env_is_honored_without_api_key(self, monkeypatch): + monkeypatch.setenv("PORTKEY_BASE_URL", "https://env-gateway.example/v1") + assert set_base_url(None, None) == "https://env-gateway.example/v1" + + def test_portkey_base_url_env_is_honored_with_api_key(self, monkeypatch): + monkeypatch.setenv("PORTKEY_BASE_URL", "https://env-gateway.example/v1") + assert set_base_url(None, "pk-fakekey") == "https://env-gateway.example/v1" + + def test_portkey_base_url_env_wins_over_portkey_proxy_env(self, monkeypatch): + monkeypatch.setenv("PORTKEY_BASE_URL", "https://env-gateway.example/v1") + monkeypatch.setenv("PORTKEY_PROXY", "https://proxy.example/v1") + assert set_base_url(None, None) == "https://env-gateway.example/v1" + + def test_portkey_proxy_env_used_when_base_url_env_absent(self, monkeypatch): + monkeypatch.setenv("PORTKEY_PROXY", "https://proxy.example/v1") + assert set_base_url(None, None) == "https://proxy.example/v1" + + def test_defaults_to_localhost_without_api_key(self): + assert set_base_url(None, None) == LOCAL_BASE_URL + + def test_defaults_to_portkey_base_url_with_api_key(self): + assert set_base_url(None, "pk-fakekey") == PORTKEY_BASE_URL