mirror of
https://github.com/anthropics/claude-code-sdk-python.git
synced 2025-12-23 09:19:52 +00:00
feat: add max_budget_usd option to Python SDK (#293)
Add support for limiting API costs using the max_budget_usd option, mirroring the TypeScript SDK functionality. When the budget is exceeded, query execution stops and returns a result with subtype 'error_max_budget_usd'. - Add max_budget_usd field to ClaudeAgentOptions - Pass --max-budget-usd flag to Claude Code CLI - Add test coverage for budget limit behavior 🤖 Generated with [Claude Code](https://claude.com/claude-code) Co-authored-by: Claude <noreply@anthropic.com>
This commit is contained in:
parent
7be296f12e
commit
ae800c5ec8
4 changed files with 169 additions and 0 deletions
95
examples/max_budget_usd.py
Normal file
95
examples/max_budget_usd.py
Normal file
|
|
@ -0,0 +1,95 @@
|
|||
#!/usr/bin/env python3
|
||||
"""Example demonstrating max_budget_usd option for cost control."""
|
||||
|
||||
import anyio
|
||||
|
||||
from claude_agent_sdk import (
|
||||
AssistantMessage,
|
||||
ClaudeAgentOptions,
|
||||
ResultMessage,
|
||||
TextBlock,
|
||||
query,
|
||||
)
|
||||
|
||||
|
||||
async def without_budget():
|
||||
"""Example without budget limit."""
|
||||
print("=== Without Budget Limit ===")
|
||||
|
||||
async for message in query(prompt="What is 2 + 2?"):
|
||||
if isinstance(message, AssistantMessage):
|
||||
for block in message.content:
|
||||
if isinstance(block, TextBlock):
|
||||
print(f"Claude: {block.text}")
|
||||
elif isinstance(message, ResultMessage):
|
||||
if message.total_cost_usd:
|
||||
print(f"Total cost: ${message.total_cost_usd:.4f}")
|
||||
print(f"Status: {message.subtype}")
|
||||
print()
|
||||
|
||||
|
||||
async def with_reasonable_budget():
|
||||
"""Example with budget that won't be exceeded."""
|
||||
print("=== With Reasonable Budget ($0.10) ===")
|
||||
|
||||
options = ClaudeAgentOptions(
|
||||
max_budget_usd=0.10, # 10 cents - plenty for a simple query
|
||||
)
|
||||
|
||||
async for message in query(prompt="What is 2 + 2?", options=options):
|
||||
if isinstance(message, AssistantMessage):
|
||||
for block in message.content:
|
||||
if isinstance(block, TextBlock):
|
||||
print(f"Claude: {block.text}")
|
||||
elif isinstance(message, ResultMessage):
|
||||
if message.total_cost_usd:
|
||||
print(f"Total cost: ${message.total_cost_usd:.4f}")
|
||||
print(f"Status: {message.subtype}")
|
||||
print()
|
||||
|
||||
|
||||
async def with_tight_budget():
|
||||
"""Example with very tight budget that will likely be exceeded."""
|
||||
print("=== With Tight Budget ($0.0001) ===")
|
||||
|
||||
options = ClaudeAgentOptions(
|
||||
max_budget_usd=0.0001, # Very small budget - will be exceeded quickly
|
||||
)
|
||||
|
||||
async for message in query(
|
||||
prompt="Read the README.md file and summarize it", options=options
|
||||
):
|
||||
if isinstance(message, AssistantMessage):
|
||||
for block in message.content:
|
||||
if isinstance(block, TextBlock):
|
||||
print(f"Claude: {block.text}")
|
||||
elif isinstance(message, ResultMessage):
|
||||
if message.total_cost_usd:
|
||||
print(f"Total cost: ${message.total_cost_usd:.4f}")
|
||||
print(f"Status: {message.subtype}")
|
||||
|
||||
# Check if budget was exceeded
|
||||
if message.subtype == "error_max_budget_usd":
|
||||
print("⚠️ Budget limit exceeded!")
|
||||
print(
|
||||
"Note: The cost may exceed the budget by up to one API call's worth"
|
||||
)
|
||||
print()
|
||||
|
||||
|
||||
async def main():
|
||||
"""Run all examples."""
|
||||
print("This example demonstrates using max_budget_usd to control API costs.\n")
|
||||
|
||||
await without_budget()
|
||||
await with_reasonable_budget()
|
||||
await with_tight_budget()
|
||||
|
||||
print(
|
||||
"\nNote: Budget checking happens after each API call completes,\n"
|
||||
"so the final cost may slightly exceed the specified budget.\n"
|
||||
)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
anyio.run(main)
|
||||
|
|
@ -115,6 +115,9 @@ class SubprocessCLITransport(Transport):
|
|||
if self._options.max_turns:
|
||||
cmd.extend(["--max-turns", str(self._options.max_turns)])
|
||||
|
||||
if self._options.max_budget_usd is not None:
|
||||
cmd.extend(["--max-budget-usd", str(self._options.max_budget_usd)])
|
||||
|
||||
if self._options.disallowed_tools:
|
||||
cmd.extend(["--disallowedTools", ",".join(self._options.disallowed_tools)])
|
||||
|
||||
|
|
|
|||
|
|
@ -518,6 +518,7 @@ class ClaudeAgentOptions:
|
|||
continue_conversation: bool = False
|
||||
resume: str | None = None
|
||||
max_turns: int | None = None
|
||||
max_budget_usd: float | None = None
|
||||
disallowed_tools: list[str] = field(default_factory=list)
|
||||
model: str | None = None
|
||||
permission_prompt_tool_name: str | None = None
|
||||
|
|
|
|||
|
|
@ -212,3 +212,73 @@ class TestIntegration:
|
|||
assert call_kwargs["options"].continue_conversation is True
|
||||
|
||||
anyio.run(_test)
|
||||
|
||||
def test_max_budget_usd_option(self):
|
||||
"""Test query with max_budget_usd option."""
|
||||
|
||||
async def _test():
|
||||
with patch(
|
||||
"claude_agent_sdk._internal.client.SubprocessCLITransport"
|
||||
) as mock_transport_class:
|
||||
mock_transport = AsyncMock()
|
||||
mock_transport_class.return_value = mock_transport
|
||||
|
||||
# Mock the message stream that exceeds budget
|
||||
async def mock_receive():
|
||||
yield {
|
||||
"type": "assistant",
|
||||
"message": {
|
||||
"role": "assistant",
|
||||
"content": [
|
||||
{"type": "text", "text": "Starting to read..."}
|
||||
],
|
||||
"model": "claude-opus-4-1-20250805",
|
||||
},
|
||||
}
|
||||
yield {
|
||||
"type": "result",
|
||||
"subtype": "error_max_budget_usd",
|
||||
"duration_ms": 500,
|
||||
"duration_api_ms": 400,
|
||||
"is_error": False,
|
||||
"num_turns": 1,
|
||||
"session_id": "test-session-budget",
|
||||
"total_cost_usd": 0.0002,
|
||||
"usage": {
|
||||
"input_tokens": 100,
|
||||
"output_tokens": 50,
|
||||
},
|
||||
}
|
||||
|
||||
mock_transport.read_messages = mock_receive
|
||||
mock_transport.connect = AsyncMock()
|
||||
mock_transport.close = AsyncMock()
|
||||
mock_transport.end_input = AsyncMock()
|
||||
mock_transport.write = AsyncMock()
|
||||
mock_transport.is_ready = Mock(return_value=True)
|
||||
|
||||
# Run query with very small budget
|
||||
messages = []
|
||||
async for msg in query(
|
||||
prompt="Read the readme",
|
||||
options=ClaudeAgentOptions(max_budget_usd=0.0001),
|
||||
):
|
||||
messages.append(msg)
|
||||
|
||||
# Verify results
|
||||
assert len(messages) == 2
|
||||
|
||||
# Check result message
|
||||
assert isinstance(messages[1], ResultMessage)
|
||||
assert messages[1].subtype == "error_max_budget_usd"
|
||||
assert messages[1].is_error is False
|
||||
assert messages[1].total_cost_usd == 0.0002
|
||||
assert messages[1].total_cost_usd is not None
|
||||
assert messages[1].total_cost_usd > 0
|
||||
|
||||
# Verify transport was created with max_budget_usd option
|
||||
mock_transport_class.assert_called_once()
|
||||
call_kwargs = mock_transport_class.call_args.kwargs
|
||||
assert call_kwargs["options"].max_budget_usd == 0.0001
|
||||
|
||||
anyio.run(_test)
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue