opensource/dify
2
0
Fork 0
mirror of https://github.com/langgenius/dify.git synced 2026-08-15 04:59:46 +08:00
Code Issues Packages Projects Releases Wiki Activity

refactor: generalize hosted model quota lifecycle (#40722)

Browse Source 
Co-authored-by: autofix-ci[bot] <114827586+autofix-ci[bot]@users.noreply.github.com>
This commit is contained in:
林玮 (Jade Lin) 2026-08-13 08:10:02 +00:00 committed by GitHub
parent 36dfc6c216
commit b9248690b0
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
5 changed files with 434 additions and 44 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

2
api/core/app/llm/__init__.py
View File

@ -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",
]

110
api/core/app/llm/quota.py
View File

@ -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,

95
api/core/model_manager.py
View File

@ -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

83
api/tests/unit_tests/core/app/test_llm_quota.py
View File

@ -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,

188
api/tests/unit_tests/core/test_model_manager.py
View File

@ -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())