This commit is contained in:
Tuan Pham 2026-11-06 00:15:37 +00:00 • committed by GitHub
commit 349012fcd1
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
5 changed files with 134 additions and 1 deletions

View file

@ -6,6 +6,7 @@ import tempfile
from collections.abc import Callable, Mapping, Sequence
from dataclasses import dataclass
from datetime import datetime, timedelta, timezone
from functools import lru_cache
from pathlib import Path
from threading import Lock
from typing import Final, Protocol
@ -468,6 +469,32 @@ def get_token_creator(
key_file_path,
)
token_getter: Final = _cached_token_getter(
auth_url,
client_id,
client_secret,
cert_str,
key_str,
cert_file_path,
key_file_path,
timeout,
expiry_buffer_minutes,
)
return token_getter, credentials["base_url"], credentials["resource_group"]
@lru_cache(maxsize=64)
def _cached_token_getter(
auth_url: str | None,
client_id: str | None,
client_secret: str | None,
cert_str: str | None,
key_str: str | None,
cert_file_path: str | None,
key_file_path: str | None,
timeout: float,
expiry_buffer_minutes: int,
) -> Callable[[], str]:
lock: Final = Lock()
token: str | None = None
token_expiry: datetime | None = None
@ -517,4 +544,4 @@ def get_token_creator(
token, token_expiry = _fetch_token()
return token
return get_token, credentials["base_url"], credentials["resource_group"]
return get_token

View file

@ -204,3 +204,21 @@ async def test_sap_chat_required_headers(
f"Header '{header_name}' has incorrect value. "
f"Expected: '{expected_value}', Got: '{request.headers[header_name]}'"
)
@pytest.mark.parametrize("sync_mode", [True, False])
@pytest.mark.asyncio
async def test_sap_chat_reuses_oauth_token_across_calls(respx_mock, aicore, sap_api_response, sync_mode):
import litellm
aicore.token.respond(json={"access_token": "mock-token", "expires_in": 43200})
completion = respx_mock.post(f"{aicore.deployment_url}/v2/completion").respond(json=sap_api_response)
messages = [{"role": "user", "content": "Hello"}]
for _ in range(3):
if sync_mode:
litellm.completion(model="sap/gpt-4o", messages=messages)
else:
await litellm.acompletion(model="sap/gpt-4o", messages=messages)
assert (aicore.token.call_count, completion.call_count) == (1, 3)

View file

@ -0,0 +1,57 @@
import json
import uuid
from dataclasses import dataclass
import pytest
import respx
import litellm
@dataclass(frozen=True, slots=True)
class MockAICore:
deployment_url: str
token: respx.Route
@pytest.fixture
def aicore(respx_mock, monkeypatch) -> MockAICore:
for var in (
"AICORE_CLIENT_ID",
"AICORE_CLIENT_SECRET",
"AICORE_AUTH_URL",
"AICORE_RESOURCE_GROUP",
"AICORE_BASE_URL",
"AICORE_CERT_URL",
"AICORE_SERVICE_KEY",
"VCAP_SERVICES",
):
monkeypatch.delenv(var, raising=False)
monkeypatch.setenv("AICORE_HOME", "notexist")
monkeypatch.setattr(litellm, "disable_aiohttp_transport", True)
host = f"https://{uuid.uuid4().hex}.mock-sap.com"
service_key = {
"clientid": "mock-client",
"clientsecret": "mock-secret",
"url": f"{host}/auth",
"serviceurls": {"AI_API_URL": f"{host}/api"},
}
monkeypatch.setattr(litellm, "sap_service_key", json.dumps(service_key))
base_url = f"{host}/api/v2"
deployment_url = f"{base_url}/inference/deployments/orchestration"
respx_mock.get(f"{base_url}/lm/deployments").respond(
json={
"resources": [
{
"configurationId": "orchestration-config",
"scenarioId": "orchestration",
"deploymentUrl": deployment_url,
"createdAt": "2026-01-01T00:00:00Z",
}
]
}
)
respx_mock.get(f"{base_url}/lm/configurations/orchestration-config").respond(json={"executableId": "orchestration"})
return MockAICore(deployment_url=deployment_url, token=respx_mock.post(f"{host}/auth/oauth/token"))

View file

@ -1667,3 +1667,20 @@ async def test_sap_embedding_required_headers(
f"Header '{header_name}' has incorrect value. "
f"Expected: '{expected_value}', Got: '{request.headers[header_name]}'"
)
@pytest.mark.parametrize("sync_mode", [True, False])
@pytest.mark.asyncio
async def test_sap_embedding_reuses_oauth_token_across_calls(respx_mock, aicore, sap_api_response, sync_mode):
import litellm
aicore.token.respond(json={"access_token": "mock-token", "expires_in": 43200})
embeddings = respx_mock.post(f"{aicore.deployment_url}/v2/embeddings").respond(json=sap_api_response)
for _ in range(3):
if sync_mode:
litellm.embedding(model="sap/text-embedding-3-small", input="Hi")
else:
await litellm.aembedding(model="sap/text-embedding-3-small", input="Hi")
assert (aicore.token.call_count, embeddings.call_count) == (1, 3)

View file

@ -160,3 +160,17 @@ def test_credentials_without_authentication_mode(monkeypatch):
# validate_credentials should raise because no authentication mode is provided
with pytest.raises(ValueError, match="SAP AI Core credentials are incomplete"):
sap_credentials.validate_credentials(**creds)
@pytest.mark.parametrize(
"expires_in, expected_token_requests",
[(43200, 1), (3000, 3)],
ids=["reused-while-valid", "refreshed-inside-60-minute-buffer"],
)
def test_token_is_shared_across_token_creators_until_refresh_buffer(aicore, expires_in, expected_token_requests):
aicore.token.respond(json={"access_token": "mock-token", "expires_in": expires_in})
tokens = tuple(sap_credentials.get_token_creator()[0]() for _ in range(3))
assert tokens == ("Bearer mock-token",) * 3
assert aicore.token.call_count == expected_token_requests