diff --git a/edsl/language_models/key_management/KeyLookupBuilder.py b/edsl/language_models/key_management/KeyLookupBuilder.py index 7eb62ea83..a0aa52fe9 100644 --- a/edsl/language_models/key_management/KeyLookupBuilder.py +++ b/edsl/language_models/key_management/KeyLookupBuilder.py @@ -36,7 +36,7 @@ class KeyLookupBuilder: >>> builder = KeyLookupBuilder(fetch_order=("config", "env")) >>> builder.DEFAULT_RPM - 10 + 100 >>> builder.DEFAULT_TPM 2000000 >>> builder.fetch_order @@ -54,8 +54,12 @@ class KeyLookupBuilder: ('openai', 'rpm') """ - DEFAULT_RPM = 10 - DEFAULT_TPM = 2000000 + # DEFAULT_RPM = 10 + # DEFAULT_TPM = 2000000 + from edsl.config import CONFIG + + DEFAULT_RPM = int(CONFIG.get("EDSL_SERVICE_RPM_BASELINE")) + DEFAULT_TPM = int(CONFIG.get("EDSL_SERVICE_TPM_BASELINE")) def __init__(self, fetch_order: Optional[tuple[str]] = None): if fetch_order is None: diff --git a/tests/language_models/test_KeyLookupBuilder.py b/tests/language_models/test_KeyLookupBuilder.py index dca52d40e..7dd67d3b5 100644 --- a/tests/language_models/test_KeyLookupBuilder.py +++ b/tests/language_models/test_KeyLookupBuilder.py @@ -137,8 +137,8 @@ def test_default_limits(): result = builder.get_language_model_input("test") - assert result.rpm == builder.DEFAULT_RPM - assert result.tpm == builder.DEFAULT_TPM + assert float(result.rpm) == float(builder.DEFAULT_RPM) + assert float(result.tpm) == float(builder.DEFAULT_TPM) def test_build_method():