From b9248690b0a9077e0bb5c0aa645533ae173bf235 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=E6=9E=97=E7=8E=AE=20=28Jade=20Lin=29?= Date: Thu, 13 Aug 2026 08:10:02 +0000 Subject: [PATCH] refactor: generalize hosted model quota lifecycle (#40722) Co-authored-by: autofix-ci[bot] <114827586+autofix-ci[bot]@users.noreply.github.com> --- api/core/app/llm/__init__.py | 2 + api/core/app/llm/quota.py | 110 ++++++---- api/core/model_manager.py | 95 ++++++++- .../unit_tests/core/app/test_llm_quota.py | 83 ++++++++ .../unit_tests/core/test_model_manager.py | 188 +++++++++++++++++- 5 files changed, 434 insertions(+), 44 deletions(-) diff --git a/api/core/app/llm/__init__.py b/api/core/app/llm/__init__.py index 6f4e6909e18..817c3aa713a 100644 --- a/api/core/app/llm/__init__.py +++ b/api/core/app/llm/__init__.py @@ -6,6 +6,7 @@ from .quota import ( ensure_llm_quota_available, ensure_llm_quota_available_for_model, reserve_llm_quota_for_model, + reserve_model_quota_for_model, ) __all__ = [ @@ -14,4 +15,5 @@ __all__ = [ "ensure_llm_quota_available", "ensure_llm_quota_available_for_model", "reserve_llm_quota_for_model", + "reserve_model_quota_for_model", ] diff --git a/api/core/app/llm/quota.py b/api/core/app/llm/quota.py index e4e84502cfa..ad28f13a6af 100644 --- a/api/core/app/llm/quota.py +++ b/api/core/app/llm/quota.py @@ -1,9 +1,7 @@ -"""Tenant-scoped helpers for checking and deducting LLM provider quota. +"""Tenant-scoped helpers for checking and deducting hosted model quota. -System-hosted quota accounting is currently defined only for LLM models. Keep -the public helpers LLM-specific so callers do not carry unused model-type -plumbing, and fail loudly if the deprecated ``ModelInstance`` wrappers are used -with a non-LLM model. +The reservation entry point covers every model invocation type. Legacy quota +helpers remain LLM-specific because token-based settlement requires LLM usage. """ import warnings @@ -30,27 +28,28 @@ from models.provider_ids import ModelProviderID from services.credit_pool_service import CreditPoolReservation, CreditPoolService -class LLMQuotaReservationState(StrEnum): +class ModelQuotaReservationState(StrEnum): RESERVED = auto() COMMITTED = auto() RELEASED = auto() @dataclass -class LLMQuotaReservation: - """Quota reserved for one system-hosted LLM invocation.""" +class ModelQuotaReservation: + """Quota reserved for one system-hosted model invocation.""" tenant_id: str provider: str + model_type: ModelType model: str provider_configuration: Any quota_unit: QuotaUnit | None = None credit_pool_reservation: CreditPoolReservation | None = None - requires_usage: bool = False - _state: LLMQuotaReservationState = field(default=LLMQuotaReservationState.RESERVED, init=False, repr=False) + requires_settlement: bool = False + _state: ModelQuotaReservationState = field(default=ModelQuotaReservationState.RESERVED, init=False, repr=False) @property - def state(self) -> LLMQuotaReservationState: + def state(self) -> ModelQuotaReservationState: return self._state @property @@ -58,37 +57,41 @@ class LLMQuotaReservation: return self.credit_pool_reservation is not None def commit(self, usage: LLMUsage | None = None) -> None: - if self._state == LLMQuotaReservationState.COMMITTED: + if self._state == ModelQuotaReservationState.COMMITTED: return - if self._state == LLMQuotaReservationState.RELEASED: - raise RuntimeError("Cannot commit a released LLM quota reservation.") + if self._state == ModelQuotaReservationState.RELEASED: + raise RuntimeError("Cannot commit a released model quota reservation.") if self.credit_pool_reservation is not None: self.credit_pool_reservation.commit() - elif self.requires_usage: - if usage is None: - raise ValueError("Accurate terminal usage is required for token-based LLM quota settlement.") - used_quota = _resolve_llm_used_quota( + elif self.requires_settlement: + used_quota = _resolve_model_used_quota( system_configuration=self.provider_configuration.system_configuration, + model_type=self.model_type, model=self.model, usage=usage, ) - _deduct_used_llm_quota( + _deduct_used_model_quota( tenant_id=self.tenant_id, provider=self.provider, provider_configuration=self.provider_configuration, used_quota=used_quota, ) - self._state = LLMQuotaReservationState.COMMITTED + self._state = ModelQuotaReservationState.COMMITTED def release(self) -> None: - if self._state in {LLMQuotaReservationState.COMMITTED, LLMQuotaReservationState.RELEASED}: + if self._state in {ModelQuotaReservationState.COMMITTED, ModelQuotaReservationState.RELEASED}: return if self.credit_pool_reservation is not None: self.credit_pool_reservation.release() - self._state = LLMQuotaReservationState.RELEASED + self._state = ModelQuotaReservationState.RELEASED + + +# Compatibility aliases for callers that still import the LLM-specific names. +LLMQuotaReservationState = ModelQuotaReservationState +LLMQuotaReservation = ModelQuotaReservation def _get_provider_configuration(*, tenant_id: str, provider: str): @@ -111,19 +114,22 @@ def _get_current_quota_configuration(system_configuration): ) -def reserve_llm_quota_for_model(*, tenant_id: str, provider: str, model: str) -> LLMQuotaReservation: - """Reserve system-hosted LLM quota before invoking the provider.""" +def reserve_model_quota_for_model( + *, tenant_id: str, provider: str, model_type: ModelType, model: str +) -> ModelQuotaReservation: + """Reserve system-hosted model quota before invoking the provider.""" provider_configuration = _get_provider_configuration(tenant_id=tenant_id, provider=provider) - reservation = LLMQuotaReservation( + reservation = ModelQuotaReservation( tenant_id=tenant_id, provider=provider, + model_type=model_type, model=model, provider_configuration=provider_configuration, ) if provider_configuration.using_provider_type != ProviderType.SYSTEM: return reservation - provider_model = provider_configuration.get_provider_model(model_type=ModelType.LLM, model=model) + provider_model = provider_configuration.get_provider_model(model_type=model_type, model=model) if provider_model and provider_model.status == ModelStatus.QUOTA_EXCEEDED: raise QuotaExceededError(f"Model provider {provider} quota exceeded.") @@ -147,20 +153,40 @@ def reserve_llm_quota_for_model(*, tenant_id: str, provider: str, model: str) -> case _: raise ValueError(f"Unsupported hosted credit pool quota unit: {quota_configuration.quota_unit}") + reservation_meta = {"source": "llm.invoke", "provider": provider, "model": model} + if model_type != ModelType.LLM: + reservation_meta = { + "source": "model.invoke", + "provider": provider, + "model_type": model_type.value, + "model": model, + } reservation.credit_pool_reservation = CreditPoolService.reserve_credits( tenant_id=tenant_id, credits_required=amount, pool_type="paid" if quota_type == ProviderQuotaType.PAID else "trial", request_id=str(uuid4()), session_factory=db.session, - meta={"source": "llm.invoke", "provider": provider, "model": model}, + meta=reservation_meta, ) elif quota_type == ProviderQuotaType.FREE: - reservation.requires_usage = True + if quota_configuration.quota_unit == QuotaUnit.TOKENS and model_type != ModelType.LLM: + raise ValueError("Token-based quota settlement only supports LLM invocations.") + reservation.requires_settlement = True return reservation +def reserve_llm_quota_for_model(*, tenant_id: str, provider: str, model: str) -> ModelQuotaReservation: + """Reserve system-hosted LLM quota before invoking the provider.""" + return reserve_model_quota_for_model( + tenant_id=tenant_id, + provider=provider, + model_type=ModelType.LLM, + model=model, + ) + + def ensure_llm_quota_available_for_model(*, tenant_id: str, provider: str, model: str) -> None: """Raise when a tenant-bound LLM model is already out of quota.""" provider_configuration = _get_provider_configuration(tenant_id=tenant_id, provider=provider) @@ -175,8 +201,10 @@ def ensure_llm_quota_available_for_model(*, tenant_id: str, provider: str, model raise QuotaExceededError(f"Model provider {provider} quota exceeded.") -def _resolve_llm_used_quota(*, system_configuration, model: str, usage: LLMUsage) -> int | None: - """Compute the quota impact for an LLM invocation under the current quota mode.""" +def _resolve_model_used_quota( + *, system_configuration, model_type: ModelType, model: str, usage: LLMUsage | None +) -> int | None: + """Compute the quota impact for a model invocation under the current quota mode.""" quota_unit = None for quota_configuration in system_configuration.quota_configurations: if quota_configuration.quota_type == system_configuration.current_quota_type: @@ -190,6 +218,8 @@ def _resolve_llm_used_quota(*, system_configuration, model: str, usage: LLMUsage used_quota = None if quota_unit: if quota_unit == QuotaUnit.TOKENS: + if model_type != ModelType.LLM or usage is None: + raise ValueError("Accurate terminal usage is required for token-based LLM quota settlement.") used_quota = usage.total_tokens elif quota_unit == QuotaUnit.CREDITS: used_quota = dify_config.get_model_credits(model) @@ -199,7 +229,17 @@ def _resolve_llm_used_quota(*, system_configuration, model: str, usage: LLMUsage return used_quota -def _deduct_free_llm_quota( +def _resolve_llm_used_quota(*, system_configuration, model: str, usage: LLMUsage) -> int | None: + """Compute the quota impact for an LLM invocation under the current quota mode.""" + return _resolve_model_used_quota( + system_configuration=system_configuration, + model_type=ModelType.LLM, + model=model, + usage=usage, + ) + + +def _deduct_free_model_quota( *, tenant_id: str, provider: str, @@ -238,8 +278,8 @@ def _deduct_free_llm_quota( raise QuotaExceededError(f"Model provider {provider} quota exceeded.") -def _deduct_used_llm_quota(*, tenant_id: str, provider: str, provider_configuration, used_quota: int | None) -> None: - """Apply a resolved LLM quota charge against the current provider quota bucket.""" +def _deduct_used_model_quota(*, tenant_id: str, provider: str, provider_configuration, used_quota: int | None) -> None: + """Apply a resolved model quota charge against the current provider quota bucket.""" if provider_configuration.using_provider_type != ProviderType.SYSTEM: return @@ -264,7 +304,7 @@ def _deduct_used_llm_quota(*, tenant_id: str, provider: str, provider_configurat session=db.session(), ) case ProviderQuotaType.FREE: - _deduct_free_llm_quota( + _deduct_free_model_quota( tenant_id=tenant_id, provider=provider, quota_type=system_configuration.current_quota_type, @@ -282,7 +322,7 @@ def deduct_llm_quota_for_model(*, tenant_id: str, provider: str, model: str, usa model=model, usage=usage, ) - _deduct_used_llm_quota( + _deduct_used_model_quota( tenant_id=tenant_id, provider=provider, provider_configuration=provider_configuration, diff --git a/api/core/model_manager.py b/api/core/model_manager.py index c07cc74583d..00059d2a35b 100644 --- a/api/core/model_manager.py +++ b/api/core/model_manager.py @@ -443,14 +443,15 @@ class ModelInstance: class QuotaManagedModelInstance(ModelInstance): - """A system-hosted LLM instance that owns quota settlement per invocation.""" + """A system-hosted model instance that owns quota settlement per invocation.""" def reserve_quota(self): - from core.app.llm.quota import reserve_llm_quota_for_model + from core.app.llm.quota import reserve_model_quota_for_model - return reserve_llm_quota_for_model( + return reserve_model_quota_for_model( tenant_id=self.provider_model_bundle.configuration.tenant_id, provider=self.provider, + model_type=self.model_type_instance.model_type, model=self.model_name, ) @@ -459,7 +460,16 @@ class QuotaManagedModelInstance(ModelInstance): try: reservation.release() except Exception: - logger.exception("Failed to release LLM quota reservation") + logger.exception("Failed to release model quota reservation") + + def _invoke_with_quota(self, function: Callable[P, R], *args: P.args, **kwargs: P.kwargs) -> R: + reservation = self.reserve_quota() + try: + response = function(*args, **kwargs) + reservation.commit() + return response + finally: + self.release_quota_safely(reservation) @overload def invoke_llm( @@ -584,6 +594,78 @@ class QuotaManagedModelInstance(ModelInstance): finally: self.release_quota_safely(reservation) + @override + def invoke_text_embedding( + self, texts: list[str], input_type: EmbeddingInputType = EmbeddingInputType.DOCUMENT + ) -> EmbeddingResult: + return self._invoke_with_quota(super().invoke_text_embedding, texts=texts, input_type=input_type) + + @override + def invoke_multimodal_embedding( + self, + multimodel_documents: list[dict], + input_type: EmbeddingInputType = EmbeddingInputType.DOCUMENT, + ) -> EmbeddingResult: + return self._invoke_with_quota( + super().invoke_multimodal_embedding, + multimodel_documents=multimodel_documents, + input_type=input_type, + ) + + @override + def invoke_rerank( + self, + query: str, + docs: list[str], + score_threshold: float | None = None, + top_n: int | None = None, + ) -> RerankResult: + return self._invoke_with_quota( + super().invoke_rerank, + query=query, + docs=docs, + score_threshold=score_threshold, + top_n=top_n, + ) + + @override + def invoke_multimodal_rerank( + self, + query: MultimodalRerankInput, + docs: list[MultimodalRerankInput], + score_threshold: float | None = None, + top_n: int | None = None, + ) -> RerankResult: + return self._invoke_with_quota( + super().invoke_multimodal_rerank, + query=query, + docs=docs, + score_threshold=score_threshold, + top_n=top_n, + ) + + @override + def invoke_moderation(self, text: str) -> bool: + return self._invoke_with_quota(super().invoke_moderation, text=text) + + @override + def invoke_speech2text(self, file: IO[bytes]) -> str: + return self._invoke_with_quota(super().invoke_speech2text, file=file) + + @override + def invoke_tts(self, content_text: str, voice: str = "") -> Iterable[bytes]: + return self._invoke_tts_stream(content_text=content_text, voice=voice) + + def _invoke_tts_stream(self, *, content_text: str, voice: str) -> Generator[bytes, None, None]: + reservation = self.reserve_quota() + try: + response = super().invoke_tts(content_text=content_text, voice=voice) + for chunk in response: + reservation.commit() + yield chunk + finally: + self.release_quota_safely(reservation) + class ModelManager: """Resolves :class:`ModelInstance` objects for a tenant and provider. @@ -645,10 +727,7 @@ class ModelManager: @staticmethod def _model_instance_class(provider_model_bundle: ProviderModelBundle, model_type: ModelType) -> type[ModelInstance]: - if ( - model_type == ModelType.LLM - and provider_model_bundle.configuration.using_provider_type == ProviderType.SYSTEM - ): + if provider_model_bundle.configuration.using_provider_type == ProviderType.SYSTEM: return QuotaManagedModelInstance return ModelInstance diff --git a/api/tests/unit_tests/core/app/test_llm_quota.py b/api/tests/unit_tests/core/app/test_llm_quota.py index e7297639610..64372ae91d2 100644 --- a/api/tests/unit_tests/core/app/test_llm_quota.py +++ b/api/tests/unit_tests/core/app/test_llm_quota.py @@ -16,6 +16,7 @@ from core.app.llm.quota import ( ensure_llm_quota_available, ensure_llm_quota_available_for_model, reserve_llm_quota_for_model, + reserve_model_quota_for_model, ) from core.entities.model_entities import ModelStatus from core.entities.provider_entities import ProviderQuotaType, QuotaUnit @@ -148,6 +149,88 @@ def test_reserve_llm_quota_uses_exact_credit_pool_reservation() -> None: credit_reservation.release.assert_not_called() +def test_reserve_non_llm_quota_uses_model_type_and_credit_pool_reservation() -> None: + credit_reservation = MagicMock() + provider_configuration = SimpleNamespace( + using_provider_type=ProviderType.SYSTEM, + get_provider_model=MagicMock(return_value=SimpleNamespace(status=ModelStatus.ACTIVE)), + system_configuration=SimpleNamespace( + current_quota_type=ProviderQuotaType.TRIAL, + quota_configurations=[ + SimpleNamespace( + quota_type=ProviderQuotaType.TRIAL, + quota_unit=QuotaUnit.CREDITS, + quota_limit=100, + ) + ], + ), + ) + provider_manager = MagicMock() + provider_manager.get_configurations.return_value.get.return_value = provider_configuration + + with ( + patch("core.app.llm.quota.create_plugin_provider_manager", return_value=provider_manager), + patch.object(type(dify_config), "get_model_credits", return_value=3), + patch("core.app.llm.quota.CreditPoolService.reserve_credits", return_value=credit_reservation) as reserve, + ): + reservation = reserve_model_quota_for_model( + tenant_id="tenant-id", + provider="openai", + model_type=ModelType.TEXT_EMBEDDING, + model="text-embedding-3-small", + ) + reservation.commit() + + provider_configuration.get_provider_model.assert_called_once_with( + model_type=ModelType.TEXT_EMBEDDING, + model="text-embedding-3-small", + ) + reserve.assert_called_once_with( + tenant_id="tenant-id", + credits_required=3, + pool_type="trial", + request_id=ANY, + session_factory=ANY, + meta={ + "source": "model.invoke", + "provider": "openai", + "model_type": "text-embedding", + "model": "text-embedding-3-small", + }, + ) + credit_reservation.commit.assert_called_once_with() + + +def test_reserve_non_llm_quota_rejects_free_token_settlement() -> None: + provider_configuration = SimpleNamespace( + using_provider_type=ProviderType.SYSTEM, + get_provider_model=MagicMock(return_value=SimpleNamespace(status=ModelStatus.ACTIVE)), + system_configuration=SimpleNamespace( + current_quota_type=ProviderQuotaType.FREE, + quota_configurations=[ + SimpleNamespace( + quota_type=ProviderQuotaType.FREE, + quota_unit=QuotaUnit.TOKENS, + quota_limit=100, + ) + ], + ), + ) + provider_manager = MagicMock() + provider_manager.get_configurations.return_value.get.return_value = provider_configuration + + with ( + patch("core.app.llm.quota.create_plugin_provider_manager", return_value=provider_manager), + pytest.raises(ValueError, match="only supports LLM invocations"), + ): + reserve_model_quota_for_model( + tenant_id="tenant-id", + provider="openai", + model_type=ModelType.TEXT_EMBEDDING, + model="text-embedding-3-small", + ) + + def test_reserve_llm_quota_requires_accurate_usage_for_free_tokens() -> None: provider_configuration = SimpleNamespace( using_provider_type=ProviderType.SYSTEM, diff --git a/api/tests/unit_tests/core/test_model_manager.py b/api/tests/unit_tests/core/test_model_manager.py index fa873e8d865..398d6c4f158 100644 --- a/api/tests/unit_tests/core/test_model_manager.py +++ b/api/tests/unit_tests/core/test_model_manager.py @@ -1,3 +1,5 @@ +from collections.abc import Callable +from io import BytesIO from unittest.mock import MagicMock, patch import pytest @@ -77,6 +79,7 @@ def _build_model_manager_bundle( *, provider_type: ProviderType, restrict_models: list[RestrictModel], + model_type: ModelType = ModelType.LLM, ) -> tuple[ModelManager, MagicMock]: provider_manager = MagicMock() bundle = MagicMock() @@ -96,7 +99,7 @@ def _build_model_manager_bundle( ) ] bundle.configuration.get_current_credentials.return_value = {"api_key": "hosted"} - bundle.model_type_instance.model_type = ModelType.LLM + bundle.model_type_instance.model_type = model_type provider_manager.get_provider_model_bundle.return_value = bundle return ModelManager(provider_manager), bundle @@ -112,6 +115,33 @@ def test_model_manager_wraps_allowlisted_system_llm() -> None: assert isinstance(model_instance, QuotaManagedModelInstance) +@pytest.mark.parametrize("model_type", list(ModelType)) +def test_model_manager_wraps_every_system_model_type(model_type: ModelType) -> None: + manager, _ = _build_model_manager_bundle( + provider_type=ProviderType.SYSTEM, + restrict_models=[RestrictModel(model="hosted-model", model_type=model_type)], + model_type=model_type, + ) + + model_instance = manager.get_model_instance("tenant-1", "openai", model_type, "hosted-model") + + assert isinstance(model_instance, QuotaManagedModelInstance) + + +def test_model_manager_does_not_wrap_custom_non_llm_model() -> None: + manager, _ = _build_model_manager_bundle( + provider_type=ProviderType.CUSTOM, + restrict_models=[], + model_type=ModelType.TEXT_EMBEDDING, + ) + + model_instance = manager.get_model_instance( + "tenant-1", "openai", ModelType.TEXT_EMBEDDING, "text-embedding-3-small" + ) + + assert type(model_instance) is ModelInstance + + def test_model_manager_rejects_system_model_by_exact_name() -> None: manager, bundle = _build_model_manager_bundle( provider_type=ProviderType.SYSTEM, @@ -275,6 +305,162 @@ def test_quota_managed_usage_stream_does_not_deliver_when_settlement_fails() -> reservation.release.assert_called_once_with() +@pytest.mark.parametrize( + ("model_type", "method_name", "invoke_model"), + [ + ( + ModelType.TEXT_EMBEDDING, + "invoke_text_embedding", + lambda model_instance: model_instance.invoke_text_embedding(texts=["hello"]), + ), + ( + ModelType.TEXT_EMBEDDING, + "invoke_multimodal_embedding", + lambda model_instance: model_instance.invoke_multimodal_embedding( + multimodel_documents=[{"content": "image"}] + ), + ), + ( + ModelType.RERANK, + "invoke_rerank", + lambda model_instance: model_instance.invoke_rerank(query="hello", docs=["document"]), + ), + ( + ModelType.RERANK, + "invoke_multimodal_rerank", + lambda model_instance: model_instance.invoke_multimodal_rerank(query=MagicMock(), docs=[MagicMock()]), + ), + ( + ModelType.MODERATION, + "invoke_moderation", + lambda model_instance: model_instance.invoke_moderation(text="hello"), + ), + ( + ModelType.SPEECH2TEXT, + "invoke_speech2text", + lambda model_instance: model_instance.invoke_speech2text(file=BytesIO(b"audio")), + ), + ], +) +def test_quota_managed_non_llm_invocation_finalizes_reservation( + model_type: ModelType, + method_name: str, + invoke_model: Callable[[ModelInstance], object], +) -> None: + manager, _ = _build_model_manager_bundle( + provider_type=ProviderType.SYSTEM, + restrict_models=[RestrictModel(model="hosted-model", model_type=model_type)], + model_type=model_type, + ) + model_instance = manager.get_model_instance("tenant-1", "openai", model_type, "hosted-model") + result = MagicMock() + reservation = MagicMock() + + with ( + patch.object(model_instance, "reserve_quota", return_value=reservation), + patch.object(ModelInstance, method_name, return_value=result) as invoke, + ): + response = invoke_model(model_instance) + + assert response is result + invoke.assert_called_once() + reservation.commit.assert_called_once_with() + reservation.release.assert_called_once_with() + + +def test_quota_managed_non_llm_invocation_releases_when_provider_fails() -> None: + manager, _ = _build_model_manager_bundle( + provider_type=ProviderType.SYSTEM, + restrict_models=[RestrictModel(model="embedding-model", model_type=ModelType.TEXT_EMBEDDING)], + model_type=ModelType.TEXT_EMBEDDING, + ) + model_instance = manager.get_model_instance("tenant-1", "openai", ModelType.TEXT_EMBEDDING, "embedding-model") + reservation = MagicMock() + + with ( + patch.object(model_instance, "reserve_quota", return_value=reservation), + patch.object(ModelInstance, "invoke_text_embedding", side_effect=RuntimeError("provider failed")), + pytest.raises(RuntimeError, match="provider failed"), + ): + model_instance.invoke_text_embedding(texts=["hello"]) + + reservation.commit.assert_not_called() + reservation.release.assert_called_once_with() + + +def test_quota_managed_tts_commits_before_first_chunk() -> None: + manager, _ = _build_model_manager_bundle( + provider_type=ProviderType.SYSTEM, + restrict_models=[RestrictModel(model="tts-model", model_type=ModelType.TTS)], + model_type=ModelType.TTS, + ) + model_instance = manager.get_model_instance("tenant-1", "openai", ModelType.TTS, "tts-model") + reservation = MagicMock() + events: list[str] = [] + reservation.commit.side_effect = lambda: events.append("commit") + + def provider_stream(): + events.append("provider") + yield b"audio" + + with ( + patch.object(model_instance, "reserve_quota", return_value=reservation), + patch.object(ModelInstance, "invoke_tts", return_value=provider_stream()), + ): + response = iter(model_instance.invoke_tts(content_text="hello", voice="voice")) + assert next(response) == b"audio" + events.append("delivered") + with pytest.raises(StopIteration): + next(response) + + assert events == ["provider", "commit", "delivered"] + reservation.commit.assert_called_once_with() + reservation.release.assert_called_once_with() + + +def test_quota_managed_tts_releases_when_provider_fails_before_first_chunk() -> None: + manager, _ = _build_model_manager_bundle( + provider_type=ProviderType.SYSTEM, + restrict_models=[RestrictModel(model="tts-model", model_type=ModelType.TTS)], + model_type=ModelType.TTS, + ) + model_instance = manager.get_model_instance("tenant-1", "openai", ModelType.TTS, "tts-model") + reservation = MagicMock() + + def failing_stream(): + raise RuntimeError("provider failed") + yield b"" + + with ( + patch.object(model_instance, "reserve_quota", return_value=reservation), + patch.object(ModelInstance, "invoke_tts", return_value=failing_stream()), + pytest.raises(RuntimeError, match="provider failed"), + ): + list(model_instance.invoke_tts(content_text="hello")) + + reservation.commit.assert_not_called() + reservation.release.assert_called_once_with() + + +def test_quota_managed_non_inference_helper_does_not_reserve_quota() -> None: + manager, _ = _build_model_manager_bundle( + provider_type=ProviderType.SYSTEM, + restrict_models=[RestrictModel(model="embedding-model", model_type=ModelType.TEXT_EMBEDDING)], + model_type=ModelType.TEXT_EMBEDDING, + ) + model_instance = manager.get_model_instance("tenant-1", "openai", ModelType.TEXT_EMBEDDING, "embedding-model") + + with ( + patch.object(model_instance, "reserve_quota") as reserve, + patch.object(ModelInstance, "get_text_embedding_num_tokens", return_value=[1]) as count_tokens, + ): + result = model_instance.get_text_embedding_num_tokens(["hello"]) + + assert result == [1] + count_tokens.assert_called_once_with(["hello"]) + reserve.assert_not_called() + + def test_lb_model_manager_fetch_next(mocker: MockerFixture, lb_model_manager: LBModelManager): # initialize redis client redis_client.initialize(redis.Redis())