mirror of
https://github.com/All-Hands-AI/OpenHands.git
synced 2026-01-09 14:57:59 -05:00
fix(backend): preserve users custom llm settings during settings migrations (#12134)
Co-authored-by: Xingyao Wang <xingyao@all-hands.dev>
This commit is contained in:
@@ -19,6 +19,7 @@ from server.constants import (
|
|||||||
LITE_LLM_API_URL,
|
LITE_LLM_API_URL,
|
||||||
LITE_LLM_TEAM_ID,
|
LITE_LLM_TEAM_ID,
|
||||||
REQUIRE_PAYMENT,
|
REQUIRE_PAYMENT,
|
||||||
|
USER_SETTINGS_VERSION_TO_MODEL,
|
||||||
get_default_litellm_model,
|
get_default_litellm_model,
|
||||||
)
|
)
|
||||||
from server.logger import logger
|
from server.logger import logger
|
||||||
@@ -202,6 +203,53 @@ class SaasSettingsStore(SettingsStore):
|
|||||||
)
|
)
|
||||||
return None
|
return None
|
||||||
|
|
||||||
|
def _has_custom_settings(
|
||||||
|
self, settings: Settings, old_user_version: int | None
|
||||||
|
) -> bool:
|
||||||
|
"""
|
||||||
|
Check if user has custom LLM settings that should be preserved.
|
||||||
|
Returns True if user customized either model or base_url.
|
||||||
|
|
||||||
|
Args:
|
||||||
|
settings: The user's current settings
|
||||||
|
old_user_version: The user's old settings version, if any
|
||||||
|
|
||||||
|
Returns:
|
||||||
|
True if user has custom settings, False if using old defaults
|
||||||
|
"""
|
||||||
|
# Normalize values
|
||||||
|
user_model = (
|
||||||
|
settings.llm_model.strip()
|
||||||
|
if settings.llm_model and settings.llm_model.strip()
|
||||||
|
else None
|
||||||
|
)
|
||||||
|
user_base_url = (
|
||||||
|
settings.llm_base_url.strip()
|
||||||
|
if settings.llm_base_url and settings.llm_base_url.strip()
|
||||||
|
else None
|
||||||
|
)
|
||||||
|
|
||||||
|
# Custom base_url = definitely custom settings (BYOK)
|
||||||
|
if user_base_url and user_base_url != LITE_LLM_API_URL:
|
||||||
|
return True
|
||||||
|
|
||||||
|
# No model set = using defaults
|
||||||
|
if not user_model:
|
||||||
|
return False
|
||||||
|
|
||||||
|
# Check if model matches old version's default
|
||||||
|
if (
|
||||||
|
old_user_version
|
||||||
|
and old_user_version < CURRENT_USER_SETTINGS_VERSION
|
||||||
|
and old_user_version in USER_SETTINGS_VERSION_TO_MODEL
|
||||||
|
):
|
||||||
|
old_default_base = USER_SETTINGS_VERSION_TO_MODEL[old_user_version]
|
||||||
|
user_model_base = user_model.split('/')[-1]
|
||||||
|
if user_model_base == old_default_base:
|
||||||
|
return False # Matches old default
|
||||||
|
|
||||||
|
return True # Custom model
|
||||||
|
|
||||||
async def update_settings_with_litellm_default(
|
async def update_settings_with_litellm_default(
|
||||||
self, settings: Settings
|
self, settings: Settings
|
||||||
) -> Settings | None:
|
) -> Settings | None:
|
||||||
@@ -213,6 +261,17 @@ class SaasSettingsStore(SettingsStore):
|
|||||||
return None
|
return None
|
||||||
local_deploy = os.environ.get('LOCAL_DEPLOYMENT', None)
|
local_deploy = os.environ.get('LOCAL_DEPLOYMENT', None)
|
||||||
key = LITE_LLM_API_KEY
|
key = LITE_LLM_API_KEY
|
||||||
|
|
||||||
|
# Check if user has custom settings
|
||||||
|
has_custom = self._has_custom_settings(settings, settings.user_version)
|
||||||
|
|
||||||
|
# Determine model to use (needed before LiteLLM user creation)
|
||||||
|
llm_model_to_use = (
|
||||||
|
settings.llm_model
|
||||||
|
if has_custom and settings.llm_model
|
||||||
|
else get_default_litellm_model()
|
||||||
|
)
|
||||||
|
|
||||||
if not local_deploy:
|
if not local_deploy:
|
||||||
# Get user info to add to litellm
|
# Get user info to add to litellm
|
||||||
token_manager = TokenManager()
|
token_manager = TokenManager()
|
||||||
@@ -276,7 +335,7 @@ class SaasSettingsStore(SettingsStore):
|
|||||||
|
|
||||||
# Create the new litellm user
|
# Create the new litellm user
|
||||||
response = await self._create_user_in_lite_llm(
|
response = await self._create_user_in_lite_llm(
|
||||||
client, email, max_budget, spend
|
client, email, max_budget, spend, llm_model_to_use
|
||||||
)
|
)
|
||||||
if not response.is_success:
|
if not response.is_success:
|
||||||
logger.warning(
|
logger.warning(
|
||||||
@@ -285,7 +344,7 @@ class SaasSettingsStore(SettingsStore):
|
|||||||
)
|
)
|
||||||
# Litellm insists on unique email addresses - it is possible the email address was registered with a different user.
|
# Litellm insists on unique email addresses - it is possible the email address was registered with a different user.
|
||||||
response = await self._create_user_in_lite_llm(
|
response = await self._create_user_in_lite_llm(
|
||||||
client, None, max_budget, spend
|
client, None, max_budget, spend, llm_model_to_use
|
||||||
)
|
)
|
||||||
|
|
||||||
# User failed to create in litellm - this is an unforseen error state...
|
# User failed to create in litellm - this is an unforseen error state...
|
||||||
@@ -311,11 +370,17 @@ class SaasSettingsStore(SettingsStore):
|
|||||||
extra={'user_id': self.user_id},
|
extra={'user_id': self.user_id},
|
||||||
)
|
)
|
||||||
|
|
||||||
|
if has_custom:
|
||||||
|
settings.llm_model = settings.llm_model or get_default_litellm_model()
|
||||||
|
settings.llm_base_url = settings.llm_base_url or LITE_LLM_API_URL
|
||||||
|
settings.llm_api_key = settings.llm_api_key or SecretStr(key)
|
||||||
|
else:
|
||||||
|
settings.llm_model = get_default_litellm_model()
|
||||||
|
settings.llm_base_url = LITE_LLM_API_URL
|
||||||
|
settings.llm_api_key = SecretStr(key)
|
||||||
|
|
||||||
settings.agent = 'CodeActAgent'
|
settings.agent = 'CodeActAgent'
|
||||||
# Use the model corresponding to the current user settings version
|
|
||||||
settings.llm_model = get_default_litellm_model()
|
|
||||||
settings.llm_api_key = SecretStr(key)
|
|
||||||
settings.llm_base_url = LITE_LLM_API_URL
|
|
||||||
return settings
|
return settings
|
||||||
|
|
||||||
@classmethod
|
@classmethod
|
||||||
@@ -398,7 +463,12 @@ class SaasSettingsStore(SettingsStore):
|
|||||||
)
|
)
|
||||||
|
|
||||||
async def _create_user_in_lite_llm(
|
async def _create_user_in_lite_llm(
|
||||||
self, client: httpx.AsyncClient, email: str | None, max_budget: int, spend: int
|
self,
|
||||||
|
client: httpx.AsyncClient,
|
||||||
|
email: str | None,
|
||||||
|
max_budget: int,
|
||||||
|
spend: int,
|
||||||
|
llm_model: str,
|
||||||
):
|
):
|
||||||
response = await client.post(
|
response = await client.post(
|
||||||
f'{LITE_LLM_API_URL}/user/new',
|
f'{LITE_LLM_API_URL}/user/new',
|
||||||
@@ -413,7 +483,7 @@ class SaasSettingsStore(SettingsStore):
|
|||||||
'send_invite_email': False,
|
'send_invite_email': False,
|
||||||
'metadata': {
|
'metadata': {
|
||||||
'version': CURRENT_USER_SETTINGS_VERSION,
|
'version': CURRENT_USER_SETTINGS_VERSION,
|
||||||
'model': get_default_litellm_model(),
|
'model': llm_model,
|
||||||
},
|
},
|
||||||
'key_alias': f'OpenHands Cloud - user {self.user_id}',
|
'key_alias': f'OpenHands Cloud - user {self.user_id}',
|
||||||
},
|
},
|
||||||
|
|||||||
@@ -6,6 +6,7 @@ from server.constants import (
|
|||||||
CURRENT_USER_SETTINGS_VERSION,
|
CURRENT_USER_SETTINGS_VERSION,
|
||||||
LITE_LLM_API_URL,
|
LITE_LLM_API_URL,
|
||||||
LITE_LLM_TEAM_ID,
|
LITE_LLM_TEAM_ID,
|
||||||
|
get_default_litellm_model,
|
||||||
)
|
)
|
||||||
from storage.saas_settings_store import SaasSettingsStore
|
from storage.saas_settings_store import SaasSettingsStore
|
||||||
from storage.user_settings import UserSettings
|
from storage.user_settings import UserSettings
|
||||||
@@ -393,10 +394,11 @@ async def test_create_user_in_lite_llm(settings_store):
|
|||||||
mock_response = AsyncMock()
|
mock_response = AsyncMock()
|
||||||
mock_response.is_success = True
|
mock_response.is_success = True
|
||||||
mock_client.post.return_value = mock_response
|
mock_client.post.return_value = mock_response
|
||||||
|
test_model = 'custom-model/test-model'
|
||||||
|
|
||||||
# Test with email
|
# Test with email
|
||||||
await settings_store._create_user_in_lite_llm(
|
await settings_store._create_user_in_lite_llm(
|
||||||
mock_client, 'test@example.com', 50, 10
|
mock_client, 'test@example.com', 50, 10, test_model
|
||||||
)
|
)
|
||||||
|
|
||||||
# Get the actual call arguments
|
# Get the actual call arguments
|
||||||
@@ -412,11 +414,11 @@ async def test_create_user_in_lite_llm(settings_store):
|
|||||||
assert call_args['json']['auto_create_key'] is True
|
assert call_args['json']['auto_create_key'] is True
|
||||||
assert call_args['json']['send_invite_email'] is False
|
assert call_args['json']['send_invite_email'] is False
|
||||||
assert call_args['json']['metadata']['version'] == CURRENT_USER_SETTINGS_VERSION
|
assert call_args['json']['metadata']['version'] == CURRENT_USER_SETTINGS_VERSION
|
||||||
assert 'model' in call_args['json']['metadata']
|
assert call_args['json']['metadata']['model'] == test_model
|
||||||
|
|
||||||
# Test with None email
|
# Test with None email
|
||||||
mock_client.post.reset_mock()
|
mock_client.post.reset_mock()
|
||||||
await settings_store._create_user_in_lite_llm(mock_client, None, 25, 15)
|
await settings_store._create_user_in_lite_llm(mock_client, None, 25, 15, test_model)
|
||||||
|
|
||||||
# Get the actual call arguments
|
# Get the actual call arguments
|
||||||
call_args = mock_client.post.call_args[1]
|
call_args = mock_client.post.call_args[1]
|
||||||
@@ -431,12 +433,12 @@ async def test_create_user_in_lite_llm(settings_store):
|
|||||||
assert call_args['json']['auto_create_key'] is True
|
assert call_args['json']['auto_create_key'] is True
|
||||||
assert call_args['json']['send_invite_email'] is False
|
assert call_args['json']['send_invite_email'] is False
|
||||||
assert call_args['json']['metadata']['version'] == CURRENT_USER_SETTINGS_VERSION
|
assert call_args['json']['metadata']['version'] == CURRENT_USER_SETTINGS_VERSION
|
||||||
assert 'model' in call_args['json']['metadata']
|
assert call_args['json']['metadata']['model'] == test_model
|
||||||
|
|
||||||
# Verify response is returned correctly
|
# Verify response is returned correctly
|
||||||
assert (
|
assert (
|
||||||
await settings_store._create_user_in_lite_llm(
|
await settings_store._create_user_in_lite_llm(
|
||||||
mock_client, 'email@test.com', 30, 7
|
mock_client, 'email@test.com', 30, 7, test_model
|
||||||
)
|
)
|
||||||
== mock_response
|
== mock_response
|
||||||
)
|
)
|
||||||
@@ -464,3 +466,808 @@ async def test_encryption(settings_store):
|
|||||||
# But we should be able to decrypt it when loading
|
# But we should be able to decrypt it when loading
|
||||||
loaded_settings = await settings_store.load()
|
loaded_settings = await settings_store.load()
|
||||||
assert loaded_settings.llm_api_key.get_secret_value() == 'secret_key'
|
assert loaded_settings.llm_api_key.get_secret_value() == 'secret_key'
|
||||||
|
|
||||||
|
|
||||||
|
@pytest.mark.asyncio
|
||||||
|
async def test_update_settings_with_litellm_default_preserves_custom_model(
|
||||||
|
settings_store, mock_litellm_api, session_maker
|
||||||
|
):
|
||||||
|
# Arrange: User has a custom LLM model set
|
||||||
|
custom_model = 'anthropic/claude-3-5-sonnet-20241022'
|
||||||
|
settings = Settings(llm_model=custom_model)
|
||||||
|
|
||||||
|
with (
|
||||||
|
patch(
|
||||||
|
'server.auth.token_manager.TokenManager.get_user_info_from_user_id',
|
||||||
|
AsyncMock(return_value={'email': 'user@example.com'}),
|
||||||
|
),
|
||||||
|
patch('storage.saas_settings_store.session_maker', session_maker),
|
||||||
|
):
|
||||||
|
# Act: Update settings with LiteLLM defaults
|
||||||
|
updated_settings = await settings_store.update_settings_with_litellm_default(
|
||||||
|
settings
|
||||||
|
)
|
||||||
|
|
||||||
|
# Assert: Custom model is preserved
|
||||||
|
assert updated_settings is not None
|
||||||
|
assert updated_settings.llm_model == custom_model
|
||||||
|
assert updated_settings.agent == 'CodeActAgent'
|
||||||
|
assert updated_settings.llm_api_key is not None
|
||||||
|
|
||||||
|
# Assert: LiteLLM metadata contains user's custom model
|
||||||
|
call_args = mock_litellm_api.return_value.__aenter__.return_value.post.call_args[1]
|
||||||
|
assert call_args['json']['metadata']['model'] == custom_model
|
||||||
|
|
||||||
|
|
||||||
|
@pytest.mark.asyncio
|
||||||
|
async def test_update_settings_with_litellm_default_uses_default_when_no_model(
|
||||||
|
settings_store, mock_litellm_api, session_maker
|
||||||
|
):
|
||||||
|
# Arrange: User has no model set (new user scenario)
|
||||||
|
settings = Settings()
|
||||||
|
|
||||||
|
with (
|
||||||
|
patch(
|
||||||
|
'server.auth.token_manager.TokenManager.get_user_info_from_user_id',
|
||||||
|
AsyncMock(return_value={'email': 'newuser@example.com'}),
|
||||||
|
),
|
||||||
|
patch('storage.saas_settings_store.session_maker', session_maker),
|
||||||
|
):
|
||||||
|
# Act: Update settings with LiteLLM defaults
|
||||||
|
updated_settings = await settings_store.update_settings_with_litellm_default(
|
||||||
|
settings
|
||||||
|
)
|
||||||
|
|
||||||
|
# Assert: Default model is assigned
|
||||||
|
assert updated_settings is not None
|
||||||
|
expected_default = get_default_litellm_model()
|
||||||
|
assert updated_settings.llm_model == expected_default
|
||||||
|
assert updated_settings.agent == 'CodeActAgent'
|
||||||
|
|
||||||
|
# Assert: LiteLLM metadata contains default model
|
||||||
|
call_args = mock_litellm_api.return_value.__aenter__.return_value.post.call_args[1]
|
||||||
|
assert call_args['json']['metadata']['model'] == expected_default
|
||||||
|
|
||||||
|
|
||||||
|
@pytest.mark.asyncio
|
||||||
|
async def test_update_settings_with_litellm_default_handles_empty_string_model(
|
||||||
|
settings_store, mock_litellm_api, session_maker
|
||||||
|
):
|
||||||
|
# Arrange: User has empty string as model (edge case)
|
||||||
|
settings = Settings(llm_model='')
|
||||||
|
|
||||||
|
with (
|
||||||
|
patch(
|
||||||
|
'server.auth.token_manager.TokenManager.get_user_info_from_user_id',
|
||||||
|
AsyncMock(return_value={'email': 'user@example.com'}),
|
||||||
|
),
|
||||||
|
patch('storage.saas_settings_store.session_maker', session_maker),
|
||||||
|
):
|
||||||
|
# Act: Update settings with LiteLLM defaults
|
||||||
|
updated_settings = await settings_store.update_settings_with_litellm_default(
|
||||||
|
settings
|
||||||
|
)
|
||||||
|
|
||||||
|
# Assert: Default model is used (empty string treated as no model)
|
||||||
|
assert updated_settings is not None
|
||||||
|
expected_default = get_default_litellm_model()
|
||||||
|
assert updated_settings.llm_model == expected_default
|
||||||
|
|
||||||
|
|
||||||
|
@pytest.mark.asyncio
|
||||||
|
async def test_update_settings_with_litellm_default_handles_whitespace_model(
|
||||||
|
settings_store, mock_litellm_api, session_maker
|
||||||
|
):
|
||||||
|
# Arrange: User has whitespace-only model (edge case)
|
||||||
|
settings = Settings(llm_model=' ')
|
||||||
|
|
||||||
|
with (
|
||||||
|
patch(
|
||||||
|
'server.auth.token_manager.TokenManager.get_user_info_from_user_id',
|
||||||
|
AsyncMock(return_value={'email': 'user@example.com'}),
|
||||||
|
),
|
||||||
|
patch('storage.saas_settings_store.session_maker', session_maker),
|
||||||
|
):
|
||||||
|
# Act: Update settings with LiteLLM defaults
|
||||||
|
updated_settings = await settings_store.update_settings_with_litellm_default(
|
||||||
|
settings
|
||||||
|
)
|
||||||
|
|
||||||
|
# Assert: Default model is used (whitespace treated as no model)
|
||||||
|
assert updated_settings is not None
|
||||||
|
expected_default = get_default_litellm_model()
|
||||||
|
assert updated_settings.llm_model == expected_default
|
||||||
|
|
||||||
|
|
||||||
|
@pytest.mark.asyncio
|
||||||
|
async def test_update_settings_with_litellm_default_preserves_custom_api_key(
|
||||||
|
settings_store, mock_litellm_api, session_maker
|
||||||
|
):
|
||||||
|
# Arrange: User has a custom API key and custom model (so has_custom=True)
|
||||||
|
custom_api_key = 'sk-custom-user-api-key-12345'
|
||||||
|
custom_model = 'gpt-4'
|
||||||
|
settings = Settings(llm_model=custom_model, llm_api_key=SecretStr(custom_api_key))
|
||||||
|
|
||||||
|
with (
|
||||||
|
patch(
|
||||||
|
'server.auth.token_manager.TokenManager.get_user_info_from_user_id',
|
||||||
|
AsyncMock(return_value={'email': 'user@example.com'}),
|
||||||
|
),
|
||||||
|
patch('storage.saas_settings_store.session_maker', session_maker),
|
||||||
|
):
|
||||||
|
# Act: Update settings with LiteLLM defaults
|
||||||
|
updated_settings = await settings_store.update_settings_with_litellm_default(
|
||||||
|
settings
|
||||||
|
)
|
||||||
|
|
||||||
|
# Assert: Custom API key is preserved when user has custom settings
|
||||||
|
assert updated_settings is not None
|
||||||
|
assert updated_settings.llm_api_key.get_secret_value() == custom_api_key
|
||||||
|
assert updated_settings.llm_api_key.get_secret_value() != 'test_api_key'
|
||||||
|
|
||||||
|
|
||||||
|
@pytest.mark.asyncio
|
||||||
|
async def test_update_settings_with_litellm_default_preserves_custom_base_url(
|
||||||
|
settings_store, mock_litellm_api, session_maker
|
||||||
|
):
|
||||||
|
# Arrange: User has a custom base URL
|
||||||
|
custom_base_url = 'https://api.custom-llm-provider.com/v1'
|
||||||
|
settings = Settings(llm_base_url=custom_base_url)
|
||||||
|
|
||||||
|
with (
|
||||||
|
patch(
|
||||||
|
'server.auth.token_manager.TokenManager.get_user_info_from_user_id',
|
||||||
|
AsyncMock(return_value={'email': 'user@example.com'}),
|
||||||
|
),
|
||||||
|
patch('storage.saas_settings_store.session_maker', session_maker),
|
||||||
|
):
|
||||||
|
# Act: Update settings with LiteLLM defaults
|
||||||
|
updated_settings = await settings_store.update_settings_with_litellm_default(
|
||||||
|
settings
|
||||||
|
)
|
||||||
|
|
||||||
|
# Assert: Custom base URL is preserved
|
||||||
|
assert updated_settings is not None
|
||||||
|
assert updated_settings.llm_base_url == custom_base_url
|
||||||
|
assert updated_settings.llm_base_url != LITE_LLM_API_URL
|
||||||
|
|
||||||
|
|
||||||
|
@pytest.mark.asyncio
|
||||||
|
async def test_update_settings_with_litellm_default_preserves_custom_api_key_and_base_url(
|
||||||
|
settings_store, mock_litellm_api, session_maker
|
||||||
|
):
|
||||||
|
# Arrange: User has both custom API key and base URL
|
||||||
|
custom_api_key = 'sk-custom-user-api-key-67890'
|
||||||
|
custom_base_url = 'https://api.another-llm-provider.com/v1'
|
||||||
|
custom_model = 'openai/gpt-4'
|
||||||
|
settings = Settings(
|
||||||
|
llm_model=custom_model,
|
||||||
|
llm_api_key=SecretStr(custom_api_key),
|
||||||
|
llm_base_url=custom_base_url,
|
||||||
|
)
|
||||||
|
|
||||||
|
with (
|
||||||
|
patch(
|
||||||
|
'server.auth.token_manager.TokenManager.get_user_info_from_user_id',
|
||||||
|
AsyncMock(return_value={'email': 'user@example.com'}),
|
||||||
|
),
|
||||||
|
patch('storage.saas_settings_store.session_maker', session_maker),
|
||||||
|
):
|
||||||
|
# Act: Update settings with LiteLLM defaults
|
||||||
|
updated_settings = await settings_store.update_settings_with_litellm_default(
|
||||||
|
settings
|
||||||
|
)
|
||||||
|
|
||||||
|
# Assert: All custom settings are preserved
|
||||||
|
assert updated_settings is not None
|
||||||
|
assert updated_settings.llm_model == custom_model
|
||||||
|
assert updated_settings.llm_api_key.get_secret_value() == custom_api_key
|
||||||
|
assert updated_settings.llm_base_url == custom_base_url
|
||||||
|
|
||||||
|
|
||||||
|
@pytest.mark.asyncio
|
||||||
|
async def test_update_settings_with_litellm_default_uses_default_api_key_when_none(
|
||||||
|
settings_store, mock_litellm_api, session_maker
|
||||||
|
):
|
||||||
|
# Arrange: User has no API key set
|
||||||
|
settings = Settings(llm_api_key=None)
|
||||||
|
|
||||||
|
with (
|
||||||
|
patch(
|
||||||
|
'server.auth.token_manager.TokenManager.get_user_info_from_user_id',
|
||||||
|
AsyncMock(return_value={'email': 'user@example.com'}),
|
||||||
|
),
|
||||||
|
patch('storage.saas_settings_store.session_maker', session_maker),
|
||||||
|
):
|
||||||
|
# Act: Update settings with LiteLLM defaults
|
||||||
|
updated_settings = await settings_store.update_settings_with_litellm_default(
|
||||||
|
settings
|
||||||
|
)
|
||||||
|
|
||||||
|
# Assert: Default LiteLLM API key is assigned
|
||||||
|
assert updated_settings is not None
|
||||||
|
assert updated_settings.llm_api_key is not None
|
||||||
|
assert updated_settings.llm_api_key.get_secret_value() == 'test_api_key'
|
||||||
|
|
||||||
|
|
||||||
|
@pytest.mark.asyncio
|
||||||
|
async def test_update_settings_with_litellm_default_uses_default_base_url_when_none(
|
||||||
|
settings_store, mock_litellm_api, session_maker
|
||||||
|
):
|
||||||
|
# Arrange: User has no base URL set
|
||||||
|
settings = Settings(llm_base_url=None)
|
||||||
|
|
||||||
|
with (
|
||||||
|
patch(
|
||||||
|
'server.auth.token_manager.TokenManager.get_user_info_from_user_id',
|
||||||
|
AsyncMock(return_value={'email': 'user@example.com'}),
|
||||||
|
),
|
||||||
|
patch('storage.saas_settings_store.session_maker', session_maker),
|
||||||
|
patch('storage.saas_settings_store.LITE_LLM_API_URL', 'http://test.url'),
|
||||||
|
):
|
||||||
|
# Act: Update settings with LiteLLM defaults
|
||||||
|
updated_settings = await settings_store.update_settings_with_litellm_default(
|
||||||
|
settings
|
||||||
|
)
|
||||||
|
|
||||||
|
# Assert: Default LiteLLM base URL is assigned (using mocked value)
|
||||||
|
assert updated_settings is not None
|
||||||
|
assert updated_settings.llm_base_url == 'http://test.url'
|
||||||
|
|
||||||
|
|
||||||
|
@pytest.mark.asyncio
|
||||||
|
async def test_update_settings_with_litellm_default_handles_empty_api_key(
|
||||||
|
settings_store, mock_litellm_api, session_maker
|
||||||
|
):
|
||||||
|
# Arrange: User has empty string as API key (edge case)
|
||||||
|
settings = Settings(llm_api_key=SecretStr(''))
|
||||||
|
|
||||||
|
with (
|
||||||
|
patch(
|
||||||
|
'server.auth.token_manager.TokenManager.get_user_info_from_user_id',
|
||||||
|
AsyncMock(return_value={'email': 'user@example.com'}),
|
||||||
|
),
|
||||||
|
patch('storage.saas_settings_store.session_maker', session_maker),
|
||||||
|
):
|
||||||
|
# Act: Update settings with LiteLLM defaults
|
||||||
|
updated_settings = await settings_store.update_settings_with_litellm_default(
|
||||||
|
settings
|
||||||
|
)
|
||||||
|
|
||||||
|
# Assert: Default API key is used (empty string treated as no key)
|
||||||
|
assert updated_settings is not None
|
||||||
|
assert updated_settings.llm_api_key.get_secret_value() == 'test_api_key'
|
||||||
|
|
||||||
|
|
||||||
|
@pytest.mark.asyncio
|
||||||
|
async def test_update_settings_with_litellm_default_handles_empty_base_url(
|
||||||
|
settings_store, mock_litellm_api, session_maker
|
||||||
|
):
|
||||||
|
# Arrange: User has empty string as base URL (edge case)
|
||||||
|
settings = Settings(llm_base_url='')
|
||||||
|
|
||||||
|
with (
|
||||||
|
patch(
|
||||||
|
'server.auth.token_manager.TokenManager.get_user_info_from_user_id',
|
||||||
|
AsyncMock(return_value={'email': 'user@example.com'}),
|
||||||
|
),
|
||||||
|
patch('storage.saas_settings_store.session_maker', session_maker),
|
||||||
|
patch('storage.saas_settings_store.LITE_LLM_API_URL', 'http://test.url'),
|
||||||
|
):
|
||||||
|
# Act: Update settings with LiteLLM defaults
|
||||||
|
updated_settings = await settings_store.update_settings_with_litellm_default(
|
||||||
|
settings
|
||||||
|
)
|
||||||
|
|
||||||
|
# Assert: Default base URL is used (empty string treated as no URL)
|
||||||
|
assert updated_settings is not None
|
||||||
|
assert updated_settings.llm_base_url == 'http://test.url'
|
||||||
|
|
||||||
|
|
||||||
|
@pytest.mark.asyncio
|
||||||
|
async def test_update_settings_with_litellm_default_handles_whitespace_api_key(
|
||||||
|
settings_store, mock_litellm_api, session_maker
|
||||||
|
):
|
||||||
|
# Arrange: User has whitespace-only API key (edge case)
|
||||||
|
settings = Settings(llm_api_key=SecretStr(' '))
|
||||||
|
|
||||||
|
with (
|
||||||
|
patch(
|
||||||
|
'server.auth.token_manager.TokenManager.get_user_info_from_user_id',
|
||||||
|
AsyncMock(return_value={'email': 'user@example.com'}),
|
||||||
|
),
|
||||||
|
patch('storage.saas_settings_store.session_maker', session_maker),
|
||||||
|
):
|
||||||
|
# Act: Update settings with LiteLLM defaults
|
||||||
|
updated_settings = await settings_store.update_settings_with_litellm_default(
|
||||||
|
settings
|
||||||
|
)
|
||||||
|
|
||||||
|
# Assert: Default API key is used (whitespace treated as no key)
|
||||||
|
assert updated_settings is not None
|
||||||
|
assert updated_settings.llm_api_key.get_secret_value() == 'test_api_key'
|
||||||
|
|
||||||
|
|
||||||
|
@pytest.mark.asyncio
|
||||||
|
async def test_update_settings_with_litellm_default_handles_whitespace_base_url(
|
||||||
|
settings_store, mock_litellm_api, session_maker
|
||||||
|
):
|
||||||
|
# Arrange: User has whitespace-only base URL (edge case)
|
||||||
|
settings = Settings(llm_base_url=' ')
|
||||||
|
|
||||||
|
with (
|
||||||
|
patch(
|
||||||
|
'server.auth.token_manager.TokenManager.get_user_info_from_user_id',
|
||||||
|
AsyncMock(return_value={'email': 'user@example.com'}),
|
||||||
|
),
|
||||||
|
patch('storage.saas_settings_store.session_maker', session_maker),
|
||||||
|
patch('storage.saas_settings_store.LITE_LLM_API_URL', 'http://test.url'),
|
||||||
|
):
|
||||||
|
# Act: Update settings with LiteLLM defaults
|
||||||
|
updated_settings = await settings_store.update_settings_with_litellm_default(
|
||||||
|
settings
|
||||||
|
)
|
||||||
|
|
||||||
|
# Assert: Default base URL is used (whitespace treated as no URL)
|
||||||
|
assert updated_settings is not None
|
||||||
|
assert updated_settings.llm_base_url == 'http://test.url'
|
||||||
|
|
||||||
|
|
||||||
|
# Tests for version migration and helper methods
|
||||||
|
|
||||||
|
|
||||||
|
@pytest.mark.asyncio
|
||||||
|
async def test_has_custom_settings_with_custom_base_url(settings_store):
|
||||||
|
# Arrange: User with custom base URL (BYOR)
|
||||||
|
with patch('storage.saas_settings_store.LITE_LLM_API_URL', 'http://default.url'):
|
||||||
|
settings = Settings(llm_base_url='http://custom.url')
|
||||||
|
|
||||||
|
# Act: Check if has custom settings
|
||||||
|
has_custom = settings_store._has_custom_settings(settings, None)
|
||||||
|
|
||||||
|
# Assert: Custom base URL detected
|
||||||
|
assert has_custom is True
|
||||||
|
|
||||||
|
|
||||||
|
@pytest.mark.asyncio
|
||||||
|
async def test_has_custom_settings_with_default_base_url(settings_store):
|
||||||
|
# Arrange: User with default base URL
|
||||||
|
with patch('storage.saas_settings_store.LITE_LLM_API_URL', 'http://default.url'):
|
||||||
|
settings = Settings(llm_base_url='http://default.url')
|
||||||
|
|
||||||
|
# Act: Check if has custom settings
|
||||||
|
has_custom = settings_store._has_custom_settings(settings, None)
|
||||||
|
|
||||||
|
# Assert: No custom settings (no model set)
|
||||||
|
assert has_custom is False
|
||||||
|
|
||||||
|
|
||||||
|
@pytest.mark.asyncio
|
||||||
|
async def test_has_custom_settings_with_no_model(settings_store):
|
||||||
|
# Arrange: User with no model set
|
||||||
|
with patch('storage.saas_settings_store.LITE_LLM_API_URL', 'http://default.url'):
|
||||||
|
settings = Settings(llm_model=None, llm_base_url='http://default.url')
|
||||||
|
|
||||||
|
# Act: Check if has custom settings
|
||||||
|
has_custom = settings_store._has_custom_settings(settings, None)
|
||||||
|
|
||||||
|
# Assert: No custom settings (using defaults)
|
||||||
|
assert has_custom is False
|
||||||
|
|
||||||
|
|
||||||
|
@pytest.mark.asyncio
|
||||||
|
async def test_has_custom_settings_with_empty_model(settings_store):
|
||||||
|
# Arrange: User with empty model
|
||||||
|
with patch('storage.saas_settings_store.LITE_LLM_API_URL', 'http://default.url'):
|
||||||
|
settings = Settings(llm_model='', llm_base_url='http://default.url')
|
||||||
|
|
||||||
|
# Act: Check if has custom settings
|
||||||
|
has_custom = settings_store._has_custom_settings(settings, None)
|
||||||
|
|
||||||
|
# Assert: No custom settings (empty treated as no model)
|
||||||
|
assert has_custom is False
|
||||||
|
|
||||||
|
|
||||||
|
@pytest.mark.asyncio
|
||||||
|
async def test_has_custom_settings_with_whitespace_model(settings_store):
|
||||||
|
# Arrange: User with whitespace-only model
|
||||||
|
with patch('storage.saas_settings_store.LITE_LLM_API_URL', 'http://default.url'):
|
||||||
|
settings = Settings(llm_model=' ', llm_base_url='http://default.url')
|
||||||
|
|
||||||
|
# Act: Check if has custom settings
|
||||||
|
has_custom = settings_store._has_custom_settings(settings, None)
|
||||||
|
|
||||||
|
# Assert: No custom settings (whitespace treated as no model)
|
||||||
|
assert has_custom is False
|
||||||
|
|
||||||
|
|
||||||
|
@pytest.mark.asyncio
|
||||||
|
async def test_has_custom_settings_with_custom_model(settings_store):
|
||||||
|
# Arrange: User with custom model
|
||||||
|
with patch('storage.saas_settings_store.LITE_LLM_API_URL', 'http://default.url'):
|
||||||
|
settings = Settings(llm_model='gpt-4', llm_base_url='http://default.url')
|
||||||
|
|
||||||
|
# Act: Check if has custom settings
|
||||||
|
has_custom = settings_store._has_custom_settings(settings, None)
|
||||||
|
|
||||||
|
# Assert: Custom model detected
|
||||||
|
assert has_custom is True
|
||||||
|
|
||||||
|
|
||||||
|
@pytest.mark.asyncio
|
||||||
|
async def test_has_custom_settings_matches_old_default_model(settings_store):
|
||||||
|
# Arrange: User with old version and model matching old default
|
||||||
|
with (
|
||||||
|
patch('storage.saas_settings_store.LITE_LLM_API_URL', 'http://default.url'),
|
||||||
|
patch('server.constants.CURRENT_USER_SETTINGS_VERSION', 5),
|
||||||
|
patch(
|
||||||
|
'server.constants.USER_SETTINGS_VERSION_TO_MODEL',
|
||||||
|
{1: 'claude-3-5-sonnet-20241022'},
|
||||||
|
),
|
||||||
|
):
|
||||||
|
settings = Settings(
|
||||||
|
llm_model='litellm_proxy/prod/claude-3-5-sonnet-20241022',
|
||||||
|
llm_base_url='http://default.url',
|
||||||
|
)
|
||||||
|
|
||||||
|
# Act: Check if has custom settings
|
||||||
|
has_custom = settings_store._has_custom_settings(settings, 1)
|
||||||
|
|
||||||
|
# Assert: Matches old default, so not custom
|
||||||
|
assert has_custom is False
|
||||||
|
|
||||||
|
|
||||||
|
@pytest.mark.asyncio
|
||||||
|
async def test_has_custom_settings_matches_old_default_by_base_name(settings_store):
|
||||||
|
# Arrange: User with old version and model matching old default by base name
|
||||||
|
with (
|
||||||
|
patch('storage.saas_settings_store.LITE_LLM_API_URL', 'http://default.url'),
|
||||||
|
patch('server.constants.CURRENT_USER_SETTINGS_VERSION', 5),
|
||||||
|
patch(
|
||||||
|
'server.constants.USER_SETTINGS_VERSION_TO_MODEL',
|
||||||
|
{1: 'claude-3-5-sonnet-20241022'},
|
||||||
|
),
|
||||||
|
):
|
||||||
|
settings = Settings(
|
||||||
|
llm_model='anthropic/claude-3-5-sonnet-20241022',
|
||||||
|
llm_base_url='http://default.url',
|
||||||
|
)
|
||||||
|
|
||||||
|
# Act: Check if has custom settings
|
||||||
|
has_custom = settings_store._has_custom_settings(settings, 1)
|
||||||
|
|
||||||
|
# Assert: Matches old default by base name, so not custom
|
||||||
|
assert has_custom is False
|
||||||
|
|
||||||
|
|
||||||
|
@pytest.mark.asyncio
|
||||||
|
async def test_has_custom_settings_with_old_version_but_custom_model(settings_store):
|
||||||
|
# Arrange: User with old version but custom model
|
||||||
|
with (
|
||||||
|
patch('storage.saas_settings_store.LITE_LLM_API_URL', 'http://default.url'),
|
||||||
|
patch('server.constants.CURRENT_USER_SETTINGS_VERSION', 5),
|
||||||
|
patch(
|
||||||
|
'server.constants.USER_SETTINGS_VERSION_TO_MODEL',
|
||||||
|
{1: 'claude-3-5-sonnet-20241022'},
|
||||||
|
),
|
||||||
|
):
|
||||||
|
settings = Settings(llm_model='gpt-4', llm_base_url='http://default.url')
|
||||||
|
|
||||||
|
# Act: Check if has custom settings
|
||||||
|
has_custom = settings_store._has_custom_settings(settings, 1)
|
||||||
|
|
||||||
|
# Assert: Custom model detected
|
||||||
|
assert has_custom is True
|
||||||
|
|
||||||
|
|
||||||
|
@pytest.mark.asyncio
|
||||||
|
async def test_has_custom_settings_with_current_version(settings_store):
|
||||||
|
# Arrange: User with current version
|
||||||
|
with (
|
||||||
|
patch('storage.saas_settings_store.LITE_LLM_API_URL', 'http://default.url'),
|
||||||
|
patch('server.constants.CURRENT_USER_SETTINGS_VERSION', 5),
|
||||||
|
patch(
|
||||||
|
'server.constants.USER_SETTINGS_VERSION_TO_MODEL',
|
||||||
|
{1: 'claude-3-5-sonnet-20241022', 5: 'claude-opus-4-5-20251101'},
|
||||||
|
),
|
||||||
|
):
|
||||||
|
settings = Settings(
|
||||||
|
llm_model='claude-3-5-sonnet-20241022', llm_base_url='http://default.url'
|
||||||
|
)
|
||||||
|
|
||||||
|
# Act: Check if has custom settings
|
||||||
|
has_custom = settings_store._has_custom_settings(settings, 5)
|
||||||
|
|
||||||
|
# Assert: Current version, so model is custom (not old default)
|
||||||
|
assert has_custom is True
|
||||||
|
|
||||||
|
|
||||||
|
@pytest.mark.asyncio
|
||||||
|
async def test_has_custom_settings_with_none_version(settings_store):
|
||||||
|
# Arrange: User with no version
|
||||||
|
with patch('storage.saas_settings_store.LITE_LLM_API_URL', 'http://default.url'):
|
||||||
|
settings = Settings(
|
||||||
|
llm_model='claude-3-5-sonnet-20241022', llm_base_url='http://default.url'
|
||||||
|
)
|
||||||
|
|
||||||
|
# Act: Check if has custom settings
|
||||||
|
has_custom = settings_store._has_custom_settings(settings, None)
|
||||||
|
|
||||||
|
# Assert: No version, so model is custom
|
||||||
|
assert has_custom is True
|
||||||
|
|
||||||
|
|
||||||
|
@pytest.mark.asyncio
|
||||||
|
async def test_has_custom_settings_with_invalid_version(settings_store):
|
||||||
|
# Arrange: User with invalid version
|
||||||
|
with (
|
||||||
|
patch('storage.saas_settings_store.LITE_LLM_API_URL', 'http://default.url'),
|
||||||
|
patch('server.constants.CURRENT_USER_SETTINGS_VERSION', 5),
|
||||||
|
patch(
|
||||||
|
'server.constants.USER_SETTINGS_VERSION_TO_MODEL',
|
||||||
|
{1: 'claude-3-5-sonnet-20241022'},
|
||||||
|
),
|
||||||
|
):
|
||||||
|
settings = Settings(
|
||||||
|
llm_model='claude-3-5-sonnet-20241022', llm_base_url='http://default.url'
|
||||||
|
)
|
||||||
|
|
||||||
|
# Act: Check if has custom settings
|
||||||
|
has_custom = settings_store._has_custom_settings(settings, 99)
|
||||||
|
|
||||||
|
# Assert: Invalid version, so model is custom
|
||||||
|
assert has_custom is True
|
||||||
|
|
||||||
|
|
||||||
|
@pytest.mark.asyncio
|
||||||
|
async def test_has_custom_settings_normalizes_whitespace(settings_store):
|
||||||
|
# Arrange: Settings with whitespace in values
|
||||||
|
with patch('storage.saas_settings_store.LITE_LLM_API_URL', 'http://default.url'):
|
||||||
|
settings = Settings(
|
||||||
|
llm_model=' claude-3-5-sonnet-20241022 ',
|
||||||
|
llm_base_url=' http://default.url ',
|
||||||
|
)
|
||||||
|
|
||||||
|
# Act: Check if has custom settings
|
||||||
|
has_custom = settings_store._has_custom_settings(settings, None)
|
||||||
|
|
||||||
|
# Assert: Whitespace is normalized, custom model detected
|
||||||
|
assert has_custom is True
|
||||||
|
|
||||||
|
|
||||||
|
@pytest.mark.asyncio
|
||||||
|
async def test_update_settings_upgrades_user_from_old_defaults(
|
||||||
|
settings_store, mock_litellm_api, session_maker
|
||||||
|
):
|
||||||
|
# Arrange: User with old version using old defaults
|
||||||
|
old_version = 1
|
||||||
|
old_model = 'litellm_proxy/prod/claude-3-5-sonnet-20241022'
|
||||||
|
settings = Settings(llm_model=old_model, llm_base_url=LITE_LLM_API_URL)
|
||||||
|
|
||||||
|
# Use a consistent test URL
|
||||||
|
test_base_url = 'http://test.url'
|
||||||
|
|
||||||
|
with (
|
||||||
|
patch('storage.saas_settings_store.session_maker', session_maker),
|
||||||
|
patch(
|
||||||
|
'server.constants.USER_SETTINGS_VERSION_TO_MODEL',
|
||||||
|
{1: 'claude-3-5-sonnet-20241022', 5: 'claude-opus-4-5-20251101'},
|
||||||
|
),
|
||||||
|
patch(
|
||||||
|
'storage.saas_settings_store.USER_SETTINGS_VERSION_TO_MODEL',
|
||||||
|
{1: 'claude-3-5-sonnet-20241022', 5: 'claude-opus-4-5-20251101'},
|
||||||
|
),
|
||||||
|
patch('server.constants.CURRENT_USER_SETTINGS_VERSION', 5),
|
||||||
|
patch('storage.saas_settings_store.CURRENT_USER_SETTINGS_VERSION', 5),
|
||||||
|
patch('storage.saas_settings_store.LITE_LLM_API_URL', test_base_url),
|
||||||
|
patch(
|
||||||
|
'storage.saas_settings_store.get_default_litellm_model',
|
||||||
|
return_value='litellm_proxy/prod/claude-opus-4-5-20251101',
|
||||||
|
),
|
||||||
|
patch(
|
||||||
|
'server.auth.token_manager.TokenManager.get_user_info_from_user_id',
|
||||||
|
AsyncMock(return_value={'email': 'user@example.com'}),
|
||||||
|
),
|
||||||
|
):
|
||||||
|
# Create existing user settings with old version
|
||||||
|
with session_maker() as session:
|
||||||
|
existing_settings = UserSettings(
|
||||||
|
keycloak_user_id=settings_store.user_id,
|
||||||
|
user_version=old_version,
|
||||||
|
llm_model=old_model,
|
||||||
|
llm_base_url=test_base_url,
|
||||||
|
)
|
||||||
|
session.add(existing_settings)
|
||||||
|
session.commit()
|
||||||
|
|
||||||
|
# Update settings to use test_base_url
|
||||||
|
# Set user_version to match the database so _has_custom_settings can detect old defaults
|
||||||
|
settings = Settings(
|
||||||
|
llm_model=old_model, llm_base_url=test_base_url, user_version=old_version
|
||||||
|
)
|
||||||
|
|
||||||
|
# Act: Update settings
|
||||||
|
updated_settings = await settings_store.update_settings_with_litellm_default(
|
||||||
|
settings
|
||||||
|
)
|
||||||
|
|
||||||
|
# Assert: Settings upgraded to new defaults
|
||||||
|
assert updated_settings is not None
|
||||||
|
assert (
|
||||||
|
updated_settings.llm_model == 'litellm_proxy/prod/claude-opus-4-5-20251101'
|
||||||
|
)
|
||||||
|
assert updated_settings.llm_base_url == test_base_url
|
||||||
|
|
||||||
|
|
||||||
|
@pytest.mark.asyncio
|
||||||
|
async def test_update_settings_preserves_custom_settings_during_upgrade(
|
||||||
|
settings_store, mock_litellm_api, session_maker
|
||||||
|
):
|
||||||
|
# Arrange: User with old version but custom settings
|
||||||
|
old_version = 1
|
||||||
|
custom_model = 'gpt-4'
|
||||||
|
custom_base_url = 'http://custom.url'
|
||||||
|
settings = Settings(llm_model=custom_model, llm_base_url=custom_base_url)
|
||||||
|
|
||||||
|
with (
|
||||||
|
patch('storage.saas_settings_store.session_maker', session_maker),
|
||||||
|
patch(
|
||||||
|
'server.constants.USER_SETTINGS_VERSION_TO_MODEL',
|
||||||
|
{1: 'claude-3-5-sonnet-20241022'},
|
||||||
|
),
|
||||||
|
patch(
|
||||||
|
'server.auth.token_manager.TokenManager.get_user_info_from_user_id',
|
||||||
|
AsyncMock(return_value={'email': 'user@example.com'}),
|
||||||
|
),
|
||||||
|
):
|
||||||
|
# Create existing user settings with old version
|
||||||
|
with session_maker() as session:
|
||||||
|
existing_settings = UserSettings(
|
||||||
|
keycloak_user_id=settings_store.user_id,
|
||||||
|
user_version=old_version,
|
||||||
|
llm_model=custom_model,
|
||||||
|
llm_base_url=custom_base_url,
|
||||||
|
)
|
||||||
|
session.add(existing_settings)
|
||||||
|
session.commit()
|
||||||
|
|
||||||
|
# Act: Update settings
|
||||||
|
updated_settings = await settings_store.update_settings_with_litellm_default(
|
||||||
|
settings
|
||||||
|
)
|
||||||
|
|
||||||
|
# Assert: Custom settings preserved
|
||||||
|
assert updated_settings is not None
|
||||||
|
assert updated_settings.llm_model == custom_model
|
||||||
|
assert updated_settings.llm_base_url == custom_base_url
|
||||||
|
|
||||||
|
|
||||||
|
@pytest.mark.asyncio
|
||||||
|
async def test_update_settings_migrates_billing_margin_v3_to_v4(
|
||||||
|
settings_store, mock_litellm_api, session_maker
|
||||||
|
):
|
||||||
|
# Arrange: User with version 3 and billing margin
|
||||||
|
old_version = 3
|
||||||
|
billing_margin = 2.0
|
||||||
|
max_budget = 10.0
|
||||||
|
spend = 5.0
|
||||||
|
|
||||||
|
settings = Settings()
|
||||||
|
|
||||||
|
mock_get_response = AsyncMock()
|
||||||
|
mock_get_response.is_success = True
|
||||||
|
mock_get_response.json = MagicMock(
|
||||||
|
return_value={'user_info': {'max_budget': max_budget, 'spend': spend}}
|
||||||
|
)
|
||||||
|
|
||||||
|
mock_post_response = AsyncMock()
|
||||||
|
mock_post_response.is_success = True
|
||||||
|
mock_post_response.json = MagicMock(return_value={'key': 'test_api_key'})
|
||||||
|
|
||||||
|
with (
|
||||||
|
patch('storage.saas_settings_store.session_maker', session_maker),
|
||||||
|
patch(
|
||||||
|
'server.auth.token_manager.TokenManager.get_user_info_from_user_id',
|
||||||
|
AsyncMock(return_value={'email': 'user@example.com'}),
|
||||||
|
),
|
||||||
|
patch('httpx.AsyncClient') as mock_client,
|
||||||
|
):
|
||||||
|
mock_client.return_value.__aenter__.return_value.get.return_value = (
|
||||||
|
mock_get_response
|
||||||
|
)
|
||||||
|
mock_client.return_value.__aenter__.return_value.post.return_value = (
|
||||||
|
mock_post_response
|
||||||
|
)
|
||||||
|
|
||||||
|
# Create existing user settings with version 3 and billing margin
|
||||||
|
with session_maker() as session:
|
||||||
|
existing_settings = UserSettings(
|
||||||
|
keycloak_user_id=settings_store.user_id,
|
||||||
|
user_version=old_version,
|
||||||
|
billing_margin=billing_margin,
|
||||||
|
)
|
||||||
|
session.add(existing_settings)
|
||||||
|
session.commit()
|
||||||
|
|
||||||
|
# Act: Update settings
|
||||||
|
updated_settings = await settings_store.update_settings_with_litellm_default(
|
||||||
|
settings
|
||||||
|
)
|
||||||
|
|
||||||
|
# Assert: Settings updated
|
||||||
|
assert updated_settings is not None
|
||||||
|
|
||||||
|
# Assert: Billing margin applied to budget
|
||||||
|
call_args = mock_client.return_value.__aenter__.return_value.post.call_args[1]
|
||||||
|
assert call_args['json']['max_budget'] == max_budget * billing_margin
|
||||||
|
assert call_args['json']['spend'] == spend * billing_margin
|
||||||
|
|
||||||
|
# Assert: Billing margin reset to 1.0
|
||||||
|
with session_maker() as session:
|
||||||
|
updated_user_settings = (
|
||||||
|
session.query(UserSettings)
|
||||||
|
.filter(UserSettings.keycloak_user_id == settings_store.user_id)
|
||||||
|
.first()
|
||||||
|
)
|
||||||
|
assert updated_user_settings.billing_margin == 1.0
|
||||||
|
|
||||||
|
|
||||||
|
@pytest.mark.asyncio
|
||||||
|
async def test_update_settings_skips_billing_margin_migration_when_already_v4(
|
||||||
|
settings_store, mock_litellm_api, session_maker
|
||||||
|
):
|
||||||
|
# Arrange: User with version 4
|
||||||
|
version = 4
|
||||||
|
billing_margin = 2.0
|
||||||
|
max_budget = 10.0
|
||||||
|
spend = 5.0
|
||||||
|
|
||||||
|
settings = Settings()
|
||||||
|
|
||||||
|
mock_get_response = AsyncMock()
|
||||||
|
mock_get_response.is_success = True
|
||||||
|
mock_get_response.json = MagicMock(
|
||||||
|
return_value={'user_info': {'max_budget': max_budget, 'spend': spend}}
|
||||||
|
)
|
||||||
|
|
||||||
|
mock_post_response = AsyncMock()
|
||||||
|
mock_post_response.is_success = True
|
||||||
|
mock_post_response.json = MagicMock(return_value={'key': 'test_api_key'})
|
||||||
|
|
||||||
|
with (
|
||||||
|
patch('storage.saas_settings_store.session_maker', session_maker),
|
||||||
|
patch(
|
||||||
|
'server.auth.token_manager.TokenManager.get_user_info_from_user_id',
|
||||||
|
AsyncMock(return_value={'email': 'user@example.com'}),
|
||||||
|
),
|
||||||
|
patch('httpx.AsyncClient') as mock_client,
|
||||||
|
):
|
||||||
|
mock_client.return_value.__aenter__.return_value.get.return_value = (
|
||||||
|
mock_get_response
|
||||||
|
)
|
||||||
|
mock_client.return_value.__aenter__.return_value.post.return_value = (
|
||||||
|
mock_post_response
|
||||||
|
)
|
||||||
|
|
||||||
|
# Create existing user settings with version 4
|
||||||
|
with session_maker() as session:
|
||||||
|
existing_settings = UserSettings(
|
||||||
|
keycloak_user_id=settings_store.user_id,
|
||||||
|
user_version=version,
|
||||||
|
billing_margin=billing_margin,
|
||||||
|
)
|
||||||
|
session.add(existing_settings)
|
||||||
|
session.commit()
|
||||||
|
|
||||||
|
# Act: Update settings
|
||||||
|
updated_settings = await settings_store.update_settings_with_litellm_default(
|
||||||
|
settings
|
||||||
|
)
|
||||||
|
|
||||||
|
# Assert: Settings updated
|
||||||
|
assert updated_settings is not None
|
||||||
|
|
||||||
|
# Assert: Billing margin NOT applied (version >= 4)
|
||||||
|
call_args = mock_client.return_value.__aenter__.return_value.post.call_args[1]
|
||||||
|
assert call_args['json']['max_budget'] == max_budget
|
||||||
|
assert call_args['json']['spend'] == spend
|
||||||
|
|||||||
@@ -30,6 +30,7 @@ class Settings(BaseModel):
|
|||||||
llm_model: str | None = None
|
llm_model: str | None = None
|
||||||
llm_api_key: SecretStr | None = None
|
llm_api_key: SecretStr | None = None
|
||||||
llm_base_url: str | None = None
|
llm_base_url: str | None = None
|
||||||
|
user_version: int | None = None
|
||||||
remote_runtime_resource_factor: int | None = None
|
remote_runtime_resource_factor: int | None = None
|
||||||
# Planned to be removed from settings
|
# Planned to be removed from settings
|
||||||
secrets_store: Secrets = Field(default_factory=Secrets, frozen=True)
|
secrets_store: Secrets = Field(default_factory=Secrets, frozen=True)
|
||||||
|
|||||||
Reference in New Issue
Block a user