diff --git a/docs/config.png b/docs/config.png index 0c709cfd..97d7fb01 100644 Binary files a/docs/config.png and b/docs/config.png differ diff --git a/docs/config.puml b/docs/config.puml index a3202893..fcad48b1 100644 --- a/docs/config.puml +++ b/docs/config.puml @@ -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 diff --git a/docs/config.svg b/docs/config.svg index 86e2926e..ee97368d 100644 --- a/docs/config.svg +++ b/docs/config.svg @@ -1,5 +1,5 @@ - + @@ -15,13 +15,13 @@ - - - - Action - - name - + + + + Action + + name + @@ -95,7 +95,7 @@ llama_stack mcp_servers : Optional[list[ModelContextProtocolServer]] name : str - quota_handlers : Optional[QuotaHandlersConfig] + quota_handlers : Optional[QuotaHandlersConfiguration] service user_data_collection @@ -103,13 +103,13 @@ - - - - ConfigurationBase - - model_config - + + + + ConfigurationBase + + model_config + @@ -154,17 +154,17 @@ - - - - DatabaseConfiguration - - config - db_type - postgres : Optional[PostgreSQLDatabaseConfiguration] - sqlite : Optional[SQLiteDatabaseConfiguration] - - check_database_configuration() -> Self + + + + DatabaseConfiguration + + config + db_type + postgres : Optional[PostgreSQLDatabaseConfiguration] + sqlite : Optional[SQLiteDatabaseConfiguration] + + check_database_configuration() -> Self @@ -190,13 +190,13 @@ - - - - JsonPathOperator - - name - + + + + JsonPathOperator + + name + @@ -223,21 +223,21 @@ - - - - JwtRoleRule - - compiled_regex - jsonpath : str - negate : bool - operator - roles : list[str] - value : Any - - check_jsonpath() -> Self - check_regex_pattern() -> Self - check_roles() -> Self + + + + JwtRoleRule + + compiled_regex + jsonpath : str + negate : bool + operator + roles : list[str] + value : Any + + check_jsonpath() -> Self + check_regex_pattern() -> Self + check_roles() -> Self @@ -284,31 +284,57 @@ check_postgres_configuration() -> Self - - - - - - QuotaHandlersConfig - - enable_token_history : bool - postgres : Optional[PostgreSQLDatabaseConfiguration] - sqlite : Optional[SQLiteDatabaseConfiguration] - + + + + + + QuotaHandlersConfiguration + + enable_token_history : bool + limiters : Optional[list[QuotaLimiterConfiguration]] + postgres : Optional[PostgreSQLDatabaseConfiguration] + scheduler : Optional[QuotaSchedulerConfiguration] + sqlite : Optional[SQLiteDatabaseConfiguration] + + + + + + + + QuotaLimiterConfiguration + + initial_quota : Annotated + name : str + period : str + quota_increase : Annotated + type : Literal['user_limiter', 'cluster_limiter'] + + + + + + + + QuotaSchedulerConfiguration + + period : Annotated + - - - - SQLiteDatabaseConfiguration - - db_path : str - + + + + SQLiteDatabaseConfiguration + + db_path : str + - + ServiceConfiguration @@ -326,20 +352,20 @@ - - - - TLSConfiguration - - tls_certificate_path : Optional[FilePath] - tls_key_password : Optional[FilePath] - tls_key_path : Optional[FilePath] - - check_tls_configuration() -> Self + + + + TLSConfiguration + + tls_certificate_path : Optional[FilePath] + tls_key_password : Optional[FilePath] + tls_key_path : Optional[FilePath] + + check_tls_configuration() -> Self - + UserDataCollection @@ -353,150 +379,155 @@ - - + + - - + + - - + + - - + + - - + + - - + + - - + + - - + + - - + + - - + + - - + + - - + + - - + + - - + + - - + + - - + + - - + + - - - - + + + + + + + + + - - + + - - + + - - + + - - + + - + custom_profile - - - operator + + + operator - + llama_stack - - - sqlite + + + sqlite - + service - + user_data_collection - + diff --git a/src/models/config.py b/src/models/config.py index cd753b02..c2542dad 100644 --- a/src/models/config.py +++ b/src/models/config.py @@ -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.""" diff --git a/tests/unit/models/config/test_dump_configuration.py b/tests/unit/models/config/test_dump_configuration.py index 8fc18c12..1709b8ea 100644 --- a/tests/unit/models/config/test_dump_configuration.py +++ b/tests/unit/models/config/test_dump_configuration.py @@ -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, + }, + } diff --git a/tests/unit/models/config/test_quota_handlers_config.py b/tests/unit/models/config/test_quota_handlers_config.py new file mode 100644 index 00000000..7d4f5a98 --- /dev/null +++ b/tests/unit/models/config/test_quota_handlers_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 diff --git a/tests/unit/models/config/test_quota_limiter_config.py b/tests/unit/models/config/test_quota_limiter_config.py new file mode 100644 index 00000000..b704ebf0 --- /dev/null +++ b/tests/unit/models/config/test_quota_limiter_config.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", + ) diff --git a/tests/unit/models/config/test_quota_scheduler_config.py b/tests/unit/models/config/test_quota_scheduler_config.py new file mode 100644 index 00000000..8dd62377 --- /dev/null +++ b/tests/unit/models/config/test_quota_scheduler_config.py @@ -0,0 +1,18 @@ +"""Unit tests for QuotaSchedulerConfig model.""" + +from models.config import QuotaSchedulerConfiguration + + +def test_quota_scheduler_default_configuration() -> None: + """Test the default configuration.""" + cfg = QuotaSchedulerConfiguration() + assert cfg is not None + # default value + assert cfg.period == 1 + + +def test_quota_scheduler_custom_configuration() -> None: + """Test the custom configuration.""" + cfg = QuotaSchedulerConfiguration(period=10) + assert cfg is not None + assert cfg.period == 10