Skip to content

Commit 58e33c3

Browse files
le-big-macclaude
andcommitted
Track fallback LM cost in _subcall at max_depth
When _subcall hits max_depth and does a plain LM completion instead of spawning a child RLM, the cost was not added to _cumulative_cost. This meant fallback subcall spending was invisible to budget tracking. Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
1 parent b5bfb11 commit 58e33c3

2 files changed

Lines changed: 32 additions & 0 deletions

File tree

‎rlm/core/rlm.py‎

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -696,6 +696,9 @@ def _subcall(self, prompt: str, model: str | None = None) -> RLMChatCompletion:
696696
end_time = time.perf_counter()
697697
model_usage = client.get_last_usage()
698698
usage_summary = UsageSummary(model_usage_summaries={root_model: model_usage})
699+
# Track fallback LM cost in parent's cumulative total
700+
if usage_summary.total_cost:
701+
self._cumulative_cost += usage_summary.total_cost
699702
return RLMChatCompletion(
700703
root_model=root_model,
701704
prompt=prompt,

‎tests/test_depth_metadata.py‎

Lines changed: 29 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -345,6 +345,35 @@ def execute_with_child_cost(code_str):
345345
assert exc_info.value.spent == 10.0
346346
assert exc_info.value.budget == 5.0
347347

348+
def test_fallback_lm_cost_tracked(self):
349+
"""At max_depth, _subcall does a plain LM completion. Its cost must
350+
still be added to _cumulative_cost for budget tracking."""
351+
352+
rlm_inst = RLM(
353+
backend="openai",
354+
backend_kwargs={"model_name": "test"},
355+
max_budget=5.0,
356+
max_depth=1, # next_depth=1 >= 1 → fallback path
357+
)
358+
359+
assert rlm_inst._cumulative_cost == 0.0
360+
361+
# Mock get_client to return a client that reports $3 cost
362+
mock_client = Mock()
363+
mock_client.model_name = "test"
364+
mock_client.completion.return_value = "fallback answer"
365+
mock_client.get_last_usage.return_value = ModelUsageSummary(
366+
total_calls=1, total_input_tokens=100, total_output_tokens=50, total_cost=3.0
367+
)
368+
369+
with patch("rlm.core.rlm.get_client", return_value=mock_client):
370+
result = rlm_inst._subcall("test prompt")
371+
372+
assert result.response == "fallback answer"
373+
assert result.usage_summary.total_cost == 3.0
374+
# Fallback LM cost must be tracked in parent's _cumulative_cost
375+
assert rlm_inst._cumulative_cost == 3.0
376+
348377
def test_token_limit_check_raises(self):
349378
"""_check_iteration_limits should raise TokenLimitExceededError when tokens exceeded."""
350379
from rlm.core.types import RLMIteration

0 commit comments

Comments
 (0)