diff --git a/backend/ee/onyx/server/query_and_chat/token_limit.py b/backend/ee/onyx/server/query_and_chat/token_limit.py index 1d8e383d9f..d5e7618c11 100644 --- a/backend/ee/onyx/server/query_and_chat/token_limit.py +++ b/backend/ee/onyx/server/query_and_chat/token_limit.py @@ -6,6 +6,7 @@ from sqlalchemy.orm import Session from onyx.configs.constants import TokenRateLimitScope from onyx.db.api_key import is_api_key_email_address from onyx.db.engine.sql_engine import get_session_with_current_tenant +from onyx.db.enums import AccountType from onyx.db.models import User from onyx.db.token_limit import ( fetch_all_user_token_rate_limits, @@ -37,8 +38,9 @@ def _check_token_rate_limits(user: User) -> None: if user.is_anonymous: _user_is_rate_limited_by_global() - elif is_api_key_email_address(user.email): - # API keys are only rate limited by global settings + elif user.account_type == AccountType.SERVICE_ACCOUNT or is_api_key_email_address( + user.email + ): _user_is_rate_limited_by_global() else: diff --git a/backend/onyx/onyxbot/slack/handlers/handle_regular_answer.py b/backend/onyx/onyxbot/slack/handlers/handle_regular_answer.py index ad118fda30..19de7182c3 100644 --- a/backend/onyx/onyxbot/slack/handlers/handle_regular_answer.py +++ b/backend/onyx/onyxbot/slack/handlers/handle_regular_answer.py @@ -346,6 +346,16 @@ def handle_regular_answer( send_as_ephemeral=send_as_ephemeral, ) finally: + if feedback_reminder_id and message_info.sender_id: + try: + client.chat_deleteScheduledMessage( + channel=message_info.sender_id, + scheduled_message_id=feedback_reminder_id, + ) + except Exception: + logger.exception( + "Unable to delete scheduled feedback reminder after Slack budget reply" + ) if not is_slash_command: update_emote_react( emoji=ONYX_BOT_REACT_EMOJI, @@ -354,7 +364,7 @@ def handle_regular_answer( remove=True, client=client, ) - return True + return False try: filters = BaseFilters( diff --git a/backend/tests/external_dependency_unit/ee/onyx/server/query_and_chat/test_service_account_budget_scope.py b/backend/tests/external_dependency_unit/ee/onyx/server/query_and_chat/test_service_account_budget_scope.py new file mode 100644 index 0000000000..a2b7098033 --- /dev/null +++ b/backend/tests/external_dependency_unit/ee/onyx/server/query_and_chat/test_service_account_budget_scope.py @@ -0,0 +1,121 @@ +from collections.abc import Generator +from datetime import datetime, timezone + +import pytest +from sqlalchemy import delete +from sqlalchemy.orm import Session + +import ee.onyx.server.query_and_chat.token_limit as ee_token_limit +import onyx.server.query_and_chat.token_limit as token_limit +from onyx.configs.constants import TokenRateLimitScope +from onyx.db.enums import AccountType +from onyx.db.llm_usage import LLMUsageRecord +from onyx.db.models import TokenRateLimit, User, UserUsage +from onyx.db.user_usage import get_token_window_start, record_user_usage +from onyx.error_handling.error_codes import OnyxErrorCode +from onyx.error_handling.exceptions import OnyxError +from onyx.tracing.flows import LLMFlow +from tests.external_dependency_unit.conftest import create_test_user, delete_test_user + +pytestmark = pytest.mark.usefixtures("tenant_context") + + +def _token_limit(scope: TokenRateLimitScope) -> TokenRateLimit: + return TokenRateLimit( + enabled=True, + token_budget=1, + cost_budget_cents=None, + period_hours=24, + scope=scope, + ) + + +def _record_over_budget_usage(db_session: Session, user: User) -> None: + record_user_usage( + db_session=db_session, + user_id=str(user.id), + usage=LLMUsageRecord( + model="service-account-budget-test-model", + flow=LLMFlow.CHAT_RESPONSE.value, + provider=None, + input_tokens=1_500, + output_tokens=500, + cache_read_tokens=0, + cost_cents=0.0, + window_start=get_token_window_start(datetime.now(timezone.utc), 24), + ), + ) + db_session.commit() + + +@pytest.fixture +def users_over_a_per_user_budget( + db_session: Session, monkeypatch: pytest.MonkeyPatch +) -> Generator[tuple[User, User], None, None]: + standard = create_test_user(db_session, "budget-scope-standard") + service = create_test_user( + db_session, "budget-scope-service", account_type=AccountType.SERVICE_ACCOUNT + ) + for user in (standard, service): + _record_over_budget_usage(db_session, user) + + monkeypatch.setattr( + ee_token_limit, + "fetch_all_user_token_rate_limits", + lambda **_: [_token_limit(TokenRateLimitScope.USER)], + ) + monkeypatch.setattr( + ee_token_limit, "_user_is_rate_limited_by_group", lambda _: None + ) + monkeypatch.setattr( + token_limit, "fetch_all_global_token_rate_limits", lambda **_: [] + ) + + yield standard, service + + db_session.rollback() + db_session.execute( + delete(UserUsage).where(UserUsage.user_id.in_([standard.id, service.id])) + ) + delete_test_user(db_session, standard, service) + db_session.commit() + + +def test_standard_user_over_per_user_budget_is_limited( + users_over_a_per_user_budget: tuple[User, User], +) -> None: + standard, _ = users_over_a_per_user_budget + + with pytest.raises(OnyxError) as exc_info: + ee_token_limit._check_token_rate_limits(standard) + + assert exc_info.value.error_code is OnyxErrorCode.RATE_LIMITED + assert exc_info.value.extra is not None + assert exc_info.value.extra["scope"] == TokenRateLimitScope.USER.value + + +def test_service_account_is_not_held_to_per_user_budget( + users_over_a_per_user_budget: tuple[User, User], +) -> None: + _, service = users_over_a_per_user_budget + + ee_token_limit._check_token_rate_limits(service) + + +def test_service_account_is_still_held_to_global_budget( + users_over_a_per_user_budget: tuple[User, User], + monkeypatch: pytest.MonkeyPatch, +) -> None: + _, service = users_over_a_per_user_budget + monkeypatch.setattr( + token_limit, + "fetch_all_global_token_rate_limits", + lambda **_: [_token_limit(TokenRateLimitScope.GLOBAL)], + ) + + with pytest.raises(OnyxError) as exc_info: + ee_token_limit._check_token_rate_limits(service) + + assert exc_info.value.error_code is OnyxErrorCode.RATE_LIMITED + assert exc_info.value.extra is not None + assert exc_info.value.extra["scope"] == TokenRateLimitScope.GLOBAL.value diff --git a/backend/tests/external_dependency_unit/server/test_llm_entrypoint_budgets.py b/backend/tests/external_dependency_unit/server/test_llm_entrypoint_budgets.py index 0930e401a1..d034be8378 100644 --- a/backend/tests/external_dependency_unit/server/test_llm_entrypoint_budgets.py +++ b/backend/tests/external_dependency_unit/server/test_llm_entrypoint_budgets.py @@ -198,6 +198,7 @@ def test_slack_bot_replies_with_budget_message_instead_of_answering( ), ) + client = MagicMock() with ( patch( f"{_HANDLE_REGULAR_ANSWER}.get_user_by_email", @@ -227,11 +228,11 @@ def test_slack_bot_replies_with_budget_message_instead_of_answering( message_info=message_info, slack_channel_config=slack_channel_config, receiver_ids=None, - client=MagicMock(), + client=client, channel="C123", logger=MagicMock(), db_session=db_session, - feedback_reminder_id=None, + feedback_reminder_id="scheduled-reminder", ) mock_update_react.assert_called_once() @@ -239,5 +240,8 @@ def test_slack_bot_replies_with_budget_message_instead_of_answering( mock_stream.assert_not_called() mock_respond.assert_called_once() assert mock_respond.call_args.kwargs["text"] == _BUDGET_MESSAGE + client.chat_deleteScheduledMessage.assert_called_once_with( + channel="U123", scheduled_message_id="scheduled-reminder" + ) if not reply_fails: - assert result is True + assert result is False diff --git a/backend/tests/unit/ee/onyx/server/gateway/test_gateway_cost_cap.py b/backend/tests/unit/ee/onyx/server/gateway/test_gateway_cost_cap.py index 4fed227102..1b528dc09c 100644 --- a/backend/tests/unit/ee/onyx/server/gateway/test_gateway_cost_cap.py +++ b/backend/tests/unit/ee/onyx/server/gateway/test_gateway_cost_cap.py @@ -3,21 +3,24 @@ The chat route runs `check_llm_cost_limit_for_provider` for every provider it calls. The gateway's generating routes did not, so a tenant over its weekly cap could keep spending Onyx's managed keys through `/v1/*`. The cap itself is -real here; only the usage lookup is stubbed to report "over the cap". +real here. """ from __future__ import annotations -from collections.abc import Callable +from collections.abc import Callable, Generator from typing import Any, cast from unittest.mock import MagicMock, patch import pytest -from fastapi import HTTPException, Request +from fastapi import Request from sqlalchemy.orm import Session from ee.onyx.server.gateway import api as gateway_api from onyx.db.models import User +from onyx.db.usage import UsageLimitExceededError, UsageType +from onyx.error_handling.error_codes import OnyxErrorCode +from onyx.error_handling.exceptions import OnyxError from onyx.server import usage_limits from onyx.server.gateway.models import ( AnthropicMessagesRequest, @@ -72,7 +75,16 @@ def _provider(api_key: str) -> tuple[LLMProviderView, ModelConfigurationView]: def _over_cap(**_kwargs: object) -> None: - raise HTTPException(status_code=429, detail="LLM usage limit exceeded") + raise UsageLimitExceededError(usage_type=UsageType.LLM_COST, current=150, limit=100) + + +@pytest.fixture(autouse=True) +def _paid_tenant_with_cap() -> Generator[None, None, None]: + with ( + patch.object(usage_limits, "is_tenant_on_trial_fn", return_value=False), + patch.object(usage_limits, "get_limit_for_usage_type", return_value=100), + ): + yield @pytest.mark.parametrize( @@ -99,10 +111,10 @@ def test_managed_key_over_cap_never_reaches_the_provider( patch.object(usage_limits, "USAGE_LIMITS_ENABLED", True), patch.object(usage_limits, "_ONYX_MANAGED_API_KEYS", {_MANAGED_KEY}), patch.object( - usage_limits, "check_usage_and_raise", side_effect=_over_cap + usage_limits, "check_usage_limit", side_effect=_over_cap ) as usage_check, patch.multiple(gateway_api, **handlers), - pytest.raises(HTTPException) as exc_info, + pytest.raises(OnyxError) as exc_info, ): route( request=make_request(), @@ -111,6 +123,7 @@ def test_managed_key_over_cap_never_reaches_the_provider( db_session=cast(Session, MagicMock(spec=Session)), ) + assert exc_info.value.error_code is OnyxErrorCode.RATE_LIMITED assert exc_info.value.status_code == 429 usage_check.assert_called_once() for handler in handlers.values(): @@ -133,7 +146,7 @@ def test_own_provider_key_is_not_subject_to_the_cap() -> None: patch.object(usage_limits, "USAGE_LIMITS_ENABLED", True), patch.object(usage_limits, "_ONYX_MANAGED_API_KEYS", {_MANAGED_KEY}), patch.object( - usage_limits, "check_usage_and_raise", side_effect=_over_cap + usage_limits, "check_usage_limit", side_effect=_over_cap ) as usage_check, patch.object(gateway_api, "handle_chat_completion") as handle, ):