feat: implement RoundCostAwarenessHook for budget visibility #781#1020
Open
jhawpetoss6-collab wants to merge 1 commit intomassgen:mainfrom
Open
feat: implement RoundCostAwarenessHook for budget visibility #781#1020jhawpetoss6-collab wants to merge 1 commit intomassgen:mainfrom
jhawpetoss6-collab wants to merge 1 commit intomassgen:mainfrom
Conversation
📝 WalkthroughWalkthroughThis PR adds a new Changes
Estimated code review effort🎯 1 (Trivial) | ⏱️ ~3 minutes Possibly related issues
🚥 Pre-merge checks | ✅ 3 | ❌ 3❌ Failed checks (3 warnings)
✅ Passed checks (3 passed)
✏️ Tip: You can configure your own custom pre-merge checks in the settings. ✨ Finishing Touches🧪 Generate unit tests (beta)
Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out. Comment |
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 3
🧹 Nitpick comments (2)
massgen/hooks/RoundCostAwarenessHook.py (2)
7-19: Add Google-style docstrings to new methods.
__init__,update_cost, andget_status_stringare changed/new functions and should include Google-style docstrings for maintainability and consistency.As per coding guidelines, "
**/*.py: For new or changed functions, include Google-style docstrings".🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@massgen/hooks/RoundCostAwarenessHook.py` around lines 7 - 19, Add Google-style docstrings to the three changed methods: __init__, update_cost, and get_status_string in RoundCostAwarenessHook; each docstring should summarize the method, describe arguments (e.g., budget: float | None for __init__, cost: float for update_cost), mention return values (if any, e.g., str for get_status_string), and note side effects (updates cumulative_cost and prints warning in update_cost). Place the docstrings immediately below each def using the Google style sections (Args, Returns, and Raises/Side Effects as appropriate) to match project guidelines.
14-14: Use structured logging instead ofHook infrastructure code should use
loggingrather thanhook_middleware.pyhandles hook operations. This ensures warnings are properly routed through the logging system and can be controlled via log levels and configuration.🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@massgen/hooks/RoundCostAwarenessHook.py` at line 14, Replace the print warning in RoundCostAwarenessHook with structured logging: import logging (or reuse the module/logger used by hook_middleware.py), get a logger (e.g., logger = logging.getLogger(__name__)), and call logger.warning with the formatted message including self.cumulative_cost and self.budget (the current print is in the RoundCostAwarenessHook method that references those attributes). Ensure the message mirrors the current text but uses logger.warning so it flows through the existing logging/configuration.
🤖 Prompt for all review comments with AI agents
Verify each finding against the current code and only fix it if needed.
Inline comments:
In `@massgen/hooks/RoundCostAwarenessHook.py`:
- Around line 7-9: In RoundCostAwarenessHook.__init__, validate the incoming
budget explicitly: if budget is None leave as-is, but if budget is negative
raise a ValueError to reject invalid budgets; set self.budget = budget (allow
0). Then update any places in this class (e.g., methods that currently do checks
like "if self.budget") to use an explicit None check ("if self.budget is not
None") so a budget of 0 is handled correctly when formatting warnings/status or
comparing cumulative_cost. Ensure all references to budget use the same
None-check logic.
- Around line 11-13: In RoundCostAwarenessHook.update_cost validate the incoming
cost before mutating self.cumulative_cost: ensure cost is a numeric (e.g., int
or float), not NaN, and non-negative; if the check fails raise a ValueError (or
TypeError for wrong type) and do not change cumulative_cost or trigger budget
logic; add this guard at the top of update_cost to reject negative/non-numeric
inputs and only then increment self.cumulative_cost and evaluate the budget
threshold.
- Line 1: RoundCostAwarenessHook must inherit from PatternHook and implement the
required async execute(function_name, arguments, context, **kwargs) ->
HookResult so the framework can call it; update the class signature to subclass
PatternHook and add an async def execute(...) that returns a HookResult and
delegates to existing methods. In __init__, update_cost, and get_status_string
add Google-style docstrings, ensure update_cost validates that cost is numeric
(int/float) before accumulating, change the budget check from "if self.budget
and" to "if self.budget is not None and ..." to allow budget=0, and replace
print(...) with structured logging via logger.warning (or logger.error as
appropriate). Ensure execute uses the class methods
(update_cost/get_status_string) and imports/returns the correct HookResult type
expected by the framework.
---
Nitpick comments:
In `@massgen/hooks/RoundCostAwarenessHook.py`:
- Around line 7-19: Add Google-style docstrings to the three changed methods:
__init__, update_cost, and get_status_string in RoundCostAwarenessHook; each
docstring should summarize the method, describe arguments (e.g., budget: float |
None for __init__, cost: float for update_cost), mention return values (if any,
e.g., str for get_status_string), and note side effects (updates cumulative_cost
and prints warning in update_cost). Place the docstrings immediately below each
def using the Google style sections (Args, Returns, and Raises/Side Effects as
appropriate) to match project guidelines.
- Line 14: Replace the print warning in RoundCostAwarenessHook with structured
logging: import logging (or reuse the module/logger used by hook_middleware.py),
get a logger (e.g., logger = logging.getLogger(__name__)), and call
logger.warning with the formatted message including self.cumulative_cost and
self.budget (the current print is in the RoundCostAwarenessHook method that
references those attributes). Ensure the message mirrors the current text but
uses logger.warning so it flows through the existing logging/configuration.
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Pro
Run ID: ab2d26c3-478a-4e51-b026-bfe568cd844c
📒 Files selected for processing (1)
massgen/hooks/RoundCostAwarenessHook.py
| @@ -0,0 +1,19 @@ | |||
| class RoundCostAwarenessHook: | |||
There was a problem hiding this comment.
🧩 Analysis chain
🏁 Script executed:
#!/bin/bash
# Verify whether RoundCostAwarenessHook is registered and follows existing hook base patterns.
set -euo pipefail
echo "== Hook class declarations =="
rg -nP --type=py -C2 'class\s+\w+Hook\s*(\(|:)' massgen
echo
echo "== RoundCostAwarenessHook references =="
rg -nP --type=py -C3 '\bRoundCostAwarenessHook\b|hook_type\s*==\s*"RoundCostAwareness"|register_hooks_from_config\s*\(' massgen
echo
echo "== Existing hook inheritance patterns =="
rg -nP --type=py -C2 'class\s+\w+\(PatternHook\)|class\s+\w+\(FunctionHook\)' massgenRepository: massgen/MassGen
Length of output: 30750
🏁 Script executed:
cat -n massgen/hooks/RoundCostAwarenessHook.pyRepository: massgen/MassGen
Length of output: 908
🏁 Script executed:
ls -la massgen/hooks/ && echo "---" && cat massgen/hooks/__init__.py 2>/dev/null || echo "No __init__.py in hooks/"Repository: massgen/MassGen
Length of output: 276
🏁 Script executed:
echo "=== Searching for any imports of RoundCostAwarenessHook ===" && \
rg -nP "from.*RoundCostAwarenessHook|import.*RoundCostAwarenessHook|massgen\.hooks" --type=py massgen && \
echo "=== Checking if register_hooks_from_config dynamically loads hooks ===" && \
rg -nP -A20 "def register_hooks_from_config" massgen/mcp_tools/hooks.py | head -60Repository: massgen/MassGen
Length of output: 1090
🏁 Script executed:
rg -nP -A60 "def register_hooks_from_config" massgen/mcp_tools/hooks.py | head -120Repository: massgen/MassGen
Length of output: 2617
🏁 Script executed:
rg -nP -A100 "def _create_hook_from_config" massgen/mcp_tools/hooks.py | head -150Repository: massgen/MassGen
Length of output: 3930
Inherit from PatternHook and implement required execute() method—hook will not execute in framework otherwise.
RoundCostAwarenessHook cannot function as a hook in the MassGen framework. The framework expects all hooks to inherit from PatternHook and implement an async def execute(function_name, arguments, context, **kwargs) -> HookResult method. This class has neither. If someone tries to register it via configuration, it will fail at runtime.
Additionally:
- Line 7-19: Add Google-style docstrings to new methods (
__init__,update_cost,get_status_string) - Line 13:
if self.budget andfails for validbudget=0. Use explicitis not Nonecheck instead. - Line 14: Use structured logging (
logger.warning) instead ofprint() - Line 12: Validate that
costis a valid numeric type before accumulating
🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed.
In `@massgen/hooks/RoundCostAwarenessHook.py` at line 1, RoundCostAwarenessHook
must inherit from PatternHook and implement the required async
execute(function_name, arguments, context, **kwargs) -> HookResult so the
framework can call it; update the class signature to subclass PatternHook and
add an async def execute(...) that returns a HookResult and delegates to
existing methods. In __init__, update_cost, and get_status_string add
Google-style docstrings, ensure update_cost validates that cost is numeric
(int/float) before accumulating, change the budget check from "if self.budget
and" to "if self.budget is not None and ..." to allow budget=0, and replace
print(...) with structured logging via logger.warning (or logger.error as
appropriate). Ensure execute uses the class methods
(update_cost/get_status_string) and imports/returns the correct HookResult type
expected by the framework.
Comment on lines
+7
to
+9
| def __init__(self, budget=None): | ||
| self.cumulative_cost = 0.0 | ||
| self.budget = budget |
There was a problem hiding this comment.
Budget boundary handling is incorrect for 0 and invalid for negatives.
Using if self.budget treats 0 as “no budget,” disabling warnings/status budget formatting. Also, negative budgets are accepted silently. Use explicit None checks and validate non-negative budget in __init__.
Proposed fix
class RoundCostAwarenessHook:
@@
def __init__(self, budget=None):
+ if budget is not None and budget < 0:
+ raise ValueError("budget must be >= 0")
self.cumulative_cost = 0.0
self.budget = budget
@@
- if self.budget and self.cumulative_cost > self.budget * 0.75:
+ if self.budget is not None and self.cumulative_cost > self.budget * 0.75:
print(f"Warning: Nearing budget limit [${self.cumulative_cost:.2f} / ${self.budget:.2f}]")
@@
- if self.budget:
+ if self.budget is not None:
return f"Cost: ${self.cumulative_cost:.2f} / ${self.budget:.2f} budget"Also applies to: 13-18
🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed.
In `@massgen/hooks/RoundCostAwarenessHook.py` around lines 7 - 9, In
RoundCostAwarenessHook.__init__, validate the incoming budget explicitly: if
budget is None leave as-is, but if budget is negative raise a ValueError to
reject invalid budgets; set self.budget = budget (allow 0). Then update any
places in this class (e.g., methods that currently do checks like "if
self.budget") to use an explicit None check ("if self.budget is not None") so a
budget of 0 is handled correctly when formatting warnings/status or comparing
cumulative_cost. Ensure all references to budget use the same None-check logic.
Comment on lines
+11
to
+13
| def update_cost(self, cost): | ||
| self.cumulative_cost += cost | ||
| if self.budget and self.cumulative_cost > self.budget * 0.75: |
There was a problem hiding this comment.
update_cost should reject negative/non-numeric costs.
Without validation, cumulative_cost can be decreased or crash on invalid input types.
Proposed fix
def update_cost(self, cost):
+ if not isinstance(cost, (int, float)):
+ raise TypeError("cost must be numeric")
+ if cost < 0:
+ raise ValueError("cost must be >= 0")
self.cumulative_cost += cost🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed.
In `@massgen/hooks/RoundCostAwarenessHook.py` around lines 11 - 13, In
RoundCostAwarenessHook.update_cost validate the incoming cost before mutating
self.cumulative_cost: ensure cost is a numeric (e.g., int or float), not NaN,
and non-negative; if the check fails raise a ValueError (or TypeError for wrong
type) and do not change cumulative_cost or trigger budget logic; add this guard
at the top of update_cost to reject negative/non-numeric inputs and only then
increment self.cumulative_cost and evaluate the budget threshold.
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
1 participant
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
This PR implements the
RoundCostAwarenessHookfor MassGen, providing agents and users with real-time visibility into per-round expenditure (#781).Changes:
RoundCostAwarenessHookfor tracking cumulative token and operation costs./claim #781
Summary by CodeRabbit
New Features