diff --git a/features/steps/acms_budget_enforcement_steps.py b/features/steps/acms_budget_enforcement_steps.py index 46a061c74..2113e749f 100644 --- a/features/steps/acms_budget_enforcement_steps.py +++ b/features/steps/acms_budget_enforcement_steps.py @@ -8,18 +8,14 @@ from cleveragents.acms.budget_enforcement import BudgetEnforcer @given("a budget enforcer with max_file_size of {size:d} bytes") -def step_create_budget_enforcer_with_max_file_size( - context: object, size: int -) -> None: +def step_create_budget_enforcer_with_max_file_size(context: object, size: int) -> None: """Create a fresh budget enforcer with specified max_file_size.""" context.budget_enforcer = BudgetEnforcer(max_file_size=size, max_total_size=10000) context.file_counter = 0 @given("a budget enforcer with max_total_size of {size:d} bytes") -def step_create_budget_enforcer_with_max_total_size( - context: object, size: int -) -> None: +def step_create_budget_enforcer_with_max_total_size(context: object, size: int) -> None: """Update the budget enforcer's max_total_size. Always reconstructs a fresh BudgetEnforcer unconditionally -- @@ -58,9 +54,7 @@ def step_add_file_of_exact_size(context: object, size: int) -> None: @when("I add a file named {filename} of {size:d} bytes to the context") -def step_add_named_file( - context: object, filename: str, size: int -) -> None: +def step_add_named_file(context: object, filename: str, size: int) -> None: """Add a named file of specified size to the context.""" filename = filename.strip('"') content = "x" * size @@ -106,18 +100,14 @@ def step_add_multibyte_utf8_file(context: object, char_count: int) -> None: context.budget_enforcer.add_file(filename, content) -@then( - "the file should be included in the assembled context" -) +@then("the file should be included in the assembled context") def step_file_should_be_included(context: object) -> None: """Verify that the last added file was included.""" files = context.budget_enforcer.get_included_files() assert len(files) > 0, "No files were included in the assembled context" -@then( - "the file should be excluded from the assembled context" -) +@then("the file should be excluded from the assembled context") def step_file_should_be_excluded(context: object) -> None: """Verify that the last added file was excluded.""" violations = context.budget_enforcer.get_violations() @@ -125,23 +115,17 @@ def step_file_should_be_excluded(context: object) -> None: @then("the total context size should be {size:d} bytes") -def step_total_size_should_be( - context: object, size: int -) -> None: +def step_total_size_should_be(context: object, size: int) -> None: """Verify the total context size.""" actual_size = context.budget_enforcer.get_total_size() assert actual_size == size, f"Expected total size {size}, got {actual_size}" @then("the total context size should not exceed {size:d} bytes") -def step_total_size_should_not_exceed( - context: object, size: int -) -> None: +def step_total_size_should_not_exceed(context: object, size: int) -> None: """Verify the total context size does not exceed limit.""" actual_size = context.budget_enforcer.get_total_size() - assert ( - actual_size <= size - ), f"Total size {actual_size} exceeds limit {size}" + assert actual_size <= size, f"Total size {actual_size} exceeds limit {size}" @then("a budget violation warning should be generated for the file") @@ -153,18 +137,14 @@ def step_budget_violation_warning_generated( assert len(violations) > 0, "No budget violations were recorded" -@then( - "a budget violation warning should be generated for exceeding max_total_size" -) +@then("a budget violation warning should be generated for exceeding max_total_size") def step_budget_violation_for_total_size(context: object) -> None: """Verify that a budget violation for max_total_size was generated.""" violations = context.budget_enforcer.get_violations() total_size_violations = [ v for v in violations if v.violation_type == "max_total_size" ] - assert ( - len(total_size_violations) > 0 - ), "No max_total_size violations were recorded" + assert len(total_size_violations) > 0, "No max_total_size violations were recorded" @then("a budget violation error should be generated") @@ -176,9 +156,7 @@ def step_budget_violation_error_generated( assert len(violations) > 0, "No budget violations were recorded" -@then( - "the error message should indicate the file size exceeds max_file_size" -) +@then("the error message should indicate the file size exceeds max_file_size") def step_error_message_indicates_file_size_exceeded( context: object, ) -> None: @@ -199,12 +177,8 @@ def step_error_message_includes_size_and_limit( violations = context.budget_enforcer.get_violations() assert len(violations) > 0, "No violations found" violation = violations[-1] - assert violation.file_size is not None, ( - "File size not included in violation" - ) - assert violation.limit is not None, ( - "Limit not included in violation" - ) + assert violation.file_size is not None, "File size not included in violation" + assert violation.limit is not None, "Limit not included in violation" @then("no budget violation should be generated") @@ -236,9 +210,7 @@ def step_no_additional_files_beyond_budget(context: object) -> None: @then("the budget violation should include the filename {filename}") -def step_violation_includes_filename( - context: object, filename: str -) -> None: +def step_violation_includes_filename(context: object, filename: str) -> None: """Verify the violation includes the filename.""" filename = filename.strip('"') violations = context.budget_enforcer.get_violations() @@ -250,9 +222,7 @@ def step_violation_includes_filename( @then("the budget violation should include the file size {size:d}") -def step_violation_includes_file_size( - context: object, size: int -) -> None: +def step_violation_includes_file_size(context: object, size: int) -> None: """Verify the violation includes the file size.""" violations = context.budget_enforcer.get_violations() assert len(violations) > 0, "No violations found" @@ -263,9 +233,7 @@ def step_violation_includes_file_size( @then("the budget violation should include the limit {limit:d}") -def step_violation_includes_limit( - context: object, limit: int -) -> None: +def step_violation_includes_limit(context: object, limit: int) -> None: """Verify the violation includes the limit.""" violations = context.budget_enforcer.get_violations() assert len(violations) > 0, "No violations found" diff --git a/robot/helper_acms_budget_enforcement.py b/robot/helper_acms_budget_enforcement.py index e4b869968..45eb32138 100644 --- a/robot/helper_acms_budget_enforcement.py +++ b/robot/helper_acms_budget_enforcement.py @@ -63,9 +63,7 @@ def _test_total_size_cutoff() -> None: f"Expected total size 2000, got {enforcer.get_total_size()}" ) violations = enforcer.get_violations() - total_violations = [ - v for v in violations if v.violation_type == "max_total_size" - ] + total_violations = [v for v in violations if v.violation_type == "max_total_size"] assert len(total_violations) == 1, ( f"Expected 1 max_total_size violation, got {len(total_violations)}" ) @@ -118,17 +116,13 @@ def _test_constructor_validation() -> None: BudgetEnforcer(max_file_size=-1, max_total_size=5000) raise AssertionError("Expected ValueError for negative max_file_size") except ValueError as e: - assert "max_file_size" in str(e), ( - f"Error should mention max_file_size: {e}" - ) + assert "max_file_size" in str(e), f"Error should mention max_file_size: {e}" try: BudgetEnforcer(max_file_size=100, max_total_size=0) raise AssertionError("Expected ValueError for zero max_total_size") except ValueError as e: - assert "max_total_size" in str(e), ( - f"Error should mention max_total_size: {e}" - ) + assert "max_total_size" in str(e), f"Error should mention max_total_size: {e}" try: BudgetEnforcer(max_file_size=5000, max_total_size=1000) @@ -188,9 +182,7 @@ def _test_multibyte_utf8() -> None: content_boundary = emoji_char * 250 assert len(content_boundary.encode("utf-8")) == 1000 included2 = enforcer2.add_file("boundary.txt", content_boundary) - assert included2 is True, ( - "File with exactly 1000 bytes should be included" - ) + assert included2 is True, "File with exactly 1000 bytes should be included" print("multibyte-utf8-ok") diff --git a/src/cleveragents/acms/budget_enforcement.py b/src/cleveragents/acms/budget_enforcement.py index 0eac29234..8eb6464ad 100644 --- a/src/cleveragents/acms/budget_enforcement.py +++ b/src/cleveragents/acms/budget_enforcement.py @@ -13,7 +13,7 @@ from dataclasses import dataclass, field _VALID_VIOLATION_TYPES: frozenset[str] = frozenset({"max_file_size", "max_total_size"}) -@dataclass +@dataclass(slots=True) class BudgetViolation: """Represents a budget constraint violation.""" @@ -43,7 +43,7 @@ class BudgetViolation: ) -@dataclass +@dataclass(slots=True) class ContextFile: """Represents a file in the assembled context.""" @@ -60,7 +60,7 @@ class ContextFile: self.size = len(self.content.encode("utf-8")) -@dataclass +@dataclass(slots=True) class BudgetEnforcer: """Enforces budget constraints on assembled context.