Merge pull request #679 from tisnik/lcore-741-quota-limiters-configuration

LCORE-741: Proper quota limiters configuration

Author: tisnik

docs/config.png

@@ -45,7 +45,7 @@ class "Configuration" as src.models.config.Configuration {
   llama_stack
   mcp_servers : Optional[list[ModelContextProtocolServer]]
   name : str
-  quota_handlers : Optional[QuotaHandlersConfig]
+  quota_handlers : Optional[QuotaHandlersConfiguration]
   service
   user_data_collection
   dump(filename: str) -> None
@@ -135,11 +135,23 @@ class "PostgreSQLDatabaseConfiguration" as src.models.config.PostgreSQLDatabaseC
   user : str
   check_postgres_configuration() -> Self
 }
-class "QuotaHandlersConfig" as src.models.config.QuotaHandlersConfig {
+class "QuotaHandlersConfiguration" as src.models.config.QuotaHandlersConfiguration {
   enable_token_history : bool
+  limiters : Optional[list[QuotaLimiterConfiguration]]
   postgres : Optional[PostgreSQLDatabaseConfiguration]
+  scheduler : Optional[QuotaSchedulerConfiguration]
   sqlite : Optional[SQLiteDatabaseConfiguration]
 }
+class "QuotaLimiterConfiguration" as src.models.config.QuotaLimiterConfiguration {
+  initial_quota : Annotated
+  name : str
+  period : str
+  quota_increase : Annotated
+  type : Literal['user_limiter', 'cluster_limiter']
+}
+class "QuotaSchedulerConfiguration" as src.models.config.QuotaSchedulerConfiguration {
+  period : Annotated
+}
 class "SQLiteDatabaseConfiguration" as src.models.config.SQLiteDatabaseConfiguration {
   db_path : str
 }
@@ -184,7 +196,8 @@ src.models.config.JwtRoleRule --|> src.models.config.ConfigurationBase
 src.models.config.LlamaStackConfiguration --|> src.models.config.ConfigurationBase
 src.models.config.ModelContextProtocolServer --|> src.models.config.ConfigurationBase
 src.models.config.PostgreSQLDatabaseConfiguration --|> src.models.config.ConfigurationBase
-src.models.config.QuotaHandlersConfig --|> src.models.config.ConfigurationBase
+src.models.config.QuotaHandlersConfiguration --|> src.models.config.ConfigurationBase
+src.models.config.QuotaLimiterConfiguration --|> src.models.config.ConfigurationBase
 src.models.config.SQLiteDatabaseConfiguration --|> src.models.config.ConfigurationBase
 src.models.config.ServiceConfiguration --|> src.models.config.ConfigurationBase
 src.models.config.TLSConfiguration --|> src.models.config.ConfigurationBase

docs/config.puml

@@ -17,6 +17,7 @@
     FilePath,
     AnyHttpUrl,
     PositiveInt,
+    NonNegativeInt,
     SecretStr,
 )
 
@@ -564,11 +565,31 @@ class ByokRag(ConfigurationBase):
     db_path: FilePath
 
 
-class QuotaHandlersConfig(ConfigurationBase):
+class QuotaLimiterConfiguration(ConfigurationBase):
+    """Configuration for one quota limiter."""
+
+    type: Literal["user_limiter", "cluster_limiter"]
+    name: str
+    initial_quota: NonNegativeInt
+    quota_increase: NonNegativeInt
+    period: str
+
+
+class QuotaSchedulerConfiguration(BaseModel):
+    """Quota scheduler configuration."""
+
+    period: PositiveInt = 1
+
+
+class QuotaHandlersConfiguration(ConfigurationBase):
     """Quota limiter configuration."""
 
     sqlite: Optional[SQLiteDatabaseConfiguration] = None
     postgres: Optional[PostgreSQLDatabaseConfiguration] = None
+    limiters: list[QuotaLimiterConfiguration] = Field(default_factory=list)
+    scheduler: QuotaSchedulerConfiguration = Field(
+        default_factory=QuotaSchedulerConfiguration
+    )
     enable_token_history: bool = False
 
 
@@ -591,7 +612,9 @@ class Configuration(ConfigurationBase):
         default_factory=ConversationCacheConfiguration
     )
     byok_rag: list[ByokRag] = Field(default_factory=list)
-    quota_handlers: QuotaHandlersConfig = Field(default_factory=QuotaHandlersConfig)
+    quota_handlers: QuotaHandlersConfiguration = Field(
+        default_factory=QuotaHandlersConfiguration
+    )
 
     def dump(self, filename: str = "configuration.json") -> None:
         """Dump actual configuration into JSON file."""

docs/config.svg

@@ -13,6 +13,9 @@
     PostgreSQLDatabaseConfiguration,
     CORSConfiguration,
     Configuration,
+    QuotaHandlersConfiguration,
+    QuotaLimiterConfiguration,
+    QuotaSchedulerConfiguration,
     ServiceConfiguration,
     InferenceConfiguration,
     TLSConfiguration,
@@ -175,6 +178,8 @@ def test_dump_configuration(tmp_path) -> None:
             "quota_handlers": {
                 "sqlite": None,
                 "postgres": None,
+                "limiters": [],
+                "scheduler": {"period": 1},
                 "enable_token_history": False,
             },
         }
@@ -293,3 +298,201 @@ def test_dump_configuration_with_more_mcp_servers(tmp_path) -> None:
                 "url": "http://localhost:8083",
             },
         ]
+
+
+def test_dump_configuration_with_quota_limiters(tmp_path) -> None:
+    """
+    Test that the Configuration object can be serialized to a JSON file and
+    that the resulting file contains all expected sections and values.
+
+    Please note that redaction process is not in place.
+    """
+    cfg = Configuration(
+        name="test_name",
+        service=ServiceConfiguration(
+            tls_config=TLSConfiguration(
+                tls_certificate_path=Path("tests/configuration/server.crt"),
+                tls_key_path=Path("tests/configuration/server.key"),
+                tls_key_password=Path("tests/configuration/password"),
+            ),
+            cors=CORSConfiguration(
+                allow_origins=["foo_origin", "bar_origin", "baz_origin"],
+                allow_credentials=False,
+                allow_methods=["foo_method", "bar_method", "baz_method"],
+                allow_headers=["foo_header", "bar_header", "baz_header"],
+            ),
+        ),
+        llama_stack=LlamaStackConfiguration(
+            use_as_library_client=True,
+            library_client_config_path="tests/configuration/run.yaml",
+            api_key="whatever",
+        ),
+        user_data_collection=UserDataCollection(
+            feedback_enabled=False, feedback_storage=None
+        ),
+        database=DatabaseConfiguration(
+            sqlite=None,
+            postgres=PostgreSQLDatabaseConfiguration(
+                db="lightspeed_stack",
+                user="ls_user",
+                password="ls_password",
+                port=5432,
+                ca_cert_path=None,
+                ssl_mode="require",
+                gss_encmode="disable",
+            ),
+        ),
+        mcp_servers=[],
+        customization=None,
+        inference=InferenceConfiguration(
+            default_provider="default_provider",
+            default_model="default_model",
+        ),
+        quota_handlers=QuotaHandlersConfiguration(
+            limiters=[
+                QuotaLimiterConfiguration(
+                    type="user_limiter",
+                    name="user_monthly_limits",
+                    initial_quota=1,
+                    quota_increase=10,
+                    period="2 seconds",
+                ),
+                QuotaLimiterConfiguration(
+                    type="cluster_limiter",
+                    name="cluster_monthly_limits",
+                    initial_quota=2,
+                    quota_increase=20,
+                    period="1 month",
+                ),
+            ],
+            scheduler=QuotaSchedulerConfiguration(period=10),
+            enable_token_history=True,
+        ),
+    )
+    assert cfg is not None
+    dump_file = tmp_path / "test.json"
+    cfg.dump(dump_file)
+
+    with open(dump_file, "r", encoding="utf-8") as fin:
+        content = json.load(fin)
+        # content should be loaded
+        assert content is not None
+
+        # all sections must exists
+        assert "name" in content
+        assert "service" in content
+        assert "llama_stack" in content
+        assert "user_data_collection" in content
+        assert "mcp_servers" in content
+        assert "authentication" in content
+        assert "authorization" in content
+        assert "customization" in content
+        assert "inference" in content
+        assert "database" in content
+        assert "byok_rag" in content
+        assert "quota_handlers" in content
+
+        # check the whole deserialized JSON file content
+        assert content == {
+            "name": "test_name",
+            "service": {
+                "host": "localhost",
+                "port": 8080,
+                "auth_enabled": False,
+                "workers": 1,
+                "color_log": True,
+                "access_log": True,
+                "tls_config": {
+                    "tls_certificate_path": "tests/configuration/server.crt",
+                    "tls_key_password": "tests/configuration/password",
+                    "tls_key_path": "tests/configuration/server.key",
+                },
+                "cors": {
+                    "allow_credentials": False,
+                    "allow_headers": [
+                        "foo_header",
+                        "bar_header",
+                        "baz_header",
+                    ],
+                    "allow_methods": [
+                        "foo_method",
+                        "bar_method",
+                        "baz_method",
+                    ],
+                    "allow_origins": [
+                        "foo_origin",
+                        "bar_origin",
+                        "baz_origin",
+                    ],
+                },
+            },
+            "llama_stack": {
+                "url": None,
+                "use_as_library_client": True,
+                "api_key": "**********",
+                "library_client_config_path": "tests/configuration/run.yaml",
+            },
+            "user_data_collection": {
+                "feedback_enabled": False,
+                "feedback_storage": None,
+                "transcripts_enabled": False,
+                "transcripts_storage": None,
+            },
+            "mcp_servers": [],
+            "authentication": {
+                "module": "noop",
+                "skip_tls_verification": False,
+                "k8s_ca_cert_path": None,
+                "k8s_cluster_api": None,
+                "jwk_config": None,
+            },
+            "customization": None,
+            "inference": {
+                "default_provider": "default_provider",
+                "default_model": "default_model",
+            },
+            "database": {
+                "sqlite": None,
+                "postgres": {
+                    "host": "localhost",
+                    "port": 5432,
+                    "db": "lightspeed_stack",
+                    "user": "ls_user",
+                    "password": "**********",
+                    "ssl_mode": "require",
+                    "gss_encmode": "disable",
+                    "namespace": "lightspeed-stack",
+                    "ca_cert_path": None,
+                },
+            },
+            "authorization": None,
+            "conversation_cache": {
+                "memory": None,
+                "postgres": None,
+                "sqlite": None,
+                "type": None,
+            },
+            "byok_rag": [],
+            "quota_handlers": {
+                "sqlite": None,
+                "postgres": None,
+                "limiters": [
+                    {
+                        "initial_quota": 1,
+                        "name": "user_monthly_limits",
+                        "period": "2 seconds",
+                        "quota_increase": 10,
+                        "type": "user_limiter",
+                    },
+                    {
+                        "initial_quota": 2,
+                        "name": "cluster_monthly_limits",
+                        "period": "1 month",
+                        "quota_increase": 20,
+                        "type": "cluster_limiter",
+                    },
+                ],
+                "scheduler": {"period": 10},
+                "enable_token_history": True,
+            },
+        }

src/models/config.py

@@ -0,0 +1,20 @@
+"""Unit tests for QuotaHandlersConfiguration model."""
+
+from models.config import QuotaHandlersConfiguration, QuotaSchedulerConfiguration
+
+
+def test_quota_handlers_configuration() -> None:
+    """Test the quota handlers configuration."""
+    cfg = QuotaHandlersConfiguration(
+        sqlite=None,
+        postgres=None,
+        limiters=[],
+        scheduler=QuotaSchedulerConfiguration(period=10),
+        enable_token_history=False,
+    )
+    assert cfg is not None
+    assert cfg.sqlite is None
+    assert cfg.postgres is None
+    assert cfg.limiters == []
+    assert cfg.scheduler is not None
+    assert not cfg.enable_token_history

tests/unit/models/config/test_dump_configuration.py

@@ -0,0 +1,60 @@
+"""Unit tests for QuotaLimiterConfig model."""
+
+import pytest
+
+from models.config import QuotaLimiterConfiguration
+
+
+def test_quota_limiter_configuration() -> None:
+    """Test the default configuration."""
+    cfg = QuotaLimiterConfiguration(
+        type="cluster_limiter",
+        name="cluster_monthly_limits",
+        initial_quota=0,
+        quota_increase=10,
+        period="3 seconds",
+    )
+    assert cfg is not None
+    assert cfg.type == "cluster_limiter"
+    assert cfg.name == "cluster_monthly_limits"
+    assert cfg.initial_quota == 0
+    assert cfg.quota_increase == 10
+    assert cfg.period == "3 seconds"
+
+
+def test_quota_limiter_configuration_improper_value_1() -> None:
+    """Test the default configuration."""
+    with pytest.raises(ValueError, match="Input should be greater than or equal to 0"):
+        _ = QuotaLimiterConfiguration(
+            type="cluster_limiter",
+            name="cluster_monthly_limits",
+            initial_quota=-1,
+            quota_increase=10,
+            period="3 seconds",
+        )
+
+
+def test_quota_limiter_configuration_improper_value_2() -> None:
+    """Test the default configuration."""
+    with pytest.raises(ValueError, match="Input should be greater than or equal to 0"):
+        _ = QuotaLimiterConfiguration(
+            type="cluster_limiter",
+            name="cluster_monthly_limits",
+            initial_quota=1,
+            quota_increase=-10,
+            period="3 seconds",
+        )
+
+
+def test_quota_limiter_configuration_improper_value_3() -> None:
+    """Test the default configuration."""
+    with pytest.raises(
+        ValueError, match="Input should be 'user_limiter' or 'cluster_limiter'"
+    ):
+        _ = QuotaLimiterConfiguration(
+            type="unknown_limiter",
+            name="cluster_monthly_limits",
+            initial_quota=1,
+            quota_increase=10,
+            period="3 seconds",
+        )