netra-zen 1.0.0__py3-none-any.whl
This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
- agent_interface/__init__.py +26 -0
- agent_interface/base_agent.py +351 -0
- netra_zen-1.0.0.dist-info/METADATA +576 -0
- netra_zen-1.0.0.dist-info/RECORD +15 -0
- netra_zen-1.0.0.dist-info/WHEEL +5 -0
- netra_zen-1.0.0.dist-info/entry_points.txt +2 -0
- netra_zen-1.0.0.dist-info/licenses/LICENSE.md +1 -0
- netra_zen-1.0.0.dist-info/top_level.txt +4 -0
- token_budget/__init__.py +1 -0
- token_budget/budget_manager.py +200 -0
- token_budget/models.py +74 -0
- token_budget/visualization.py +22 -0
- token_transparency/__init__.py +20 -0
- token_transparency/claude_pricing_engine.py +327 -0
- zen_orchestrator.py +2884 -0
@@ -0,0 +1,200 @@
|
|
1
|
+
"""Token budget manager - enhanced implementation with cost support."""
|
2
|
+
|
3
|
+
from .models import CommandBudgetInfo, BudgetType
|
4
|
+
from typing import Dict, Optional, List, Union
|
5
|
+
import sys
|
6
|
+
from pathlib import Path
|
7
|
+
|
8
|
+
# Import ClaudePricingEngine for cost calculations
|
9
|
+
sys.path.insert(0, str(Path(__file__).parent.parent))
|
10
|
+
try:
|
11
|
+
from token_transparency import ClaudePricingEngine, TokenUsageData
|
12
|
+
except ImportError:
|
13
|
+
ClaudePricingEngine = None
|
14
|
+
TokenUsageData = None
|
15
|
+
|
16
|
+
class TokenBudgetManager:
|
17
|
+
"""Manages budgets for overall session and individual commands with support for both tokens and cost."""
|
18
|
+
|
19
|
+
def __init__(self, overall_budget: Optional[Union[int, float]] = None,
|
20
|
+
enforcement_mode: str = "warn",
|
21
|
+
budget_type: Union[str, BudgetType] = BudgetType.TOKENS,
|
22
|
+
overall_cost_budget: Optional[float] = None):
|
23
|
+
"""
|
24
|
+
Initialize the budget manager.
|
25
|
+
|
26
|
+
Args:
|
27
|
+
overall_budget: Overall token budget (backward compatibility)
|
28
|
+
enforcement_mode: Action when budget exceeded ("warn" or "block")
|
29
|
+
budget_type: Type of budget ("tokens", "cost", or "mixed")
|
30
|
+
overall_cost_budget: Overall cost budget in USD
|
31
|
+
"""
|
32
|
+
# Handle backward compatibility
|
33
|
+
if overall_cost_budget is not None:
|
34
|
+
self.overall_budget = overall_cost_budget
|
35
|
+
self.budget_type = BudgetType.COST
|
36
|
+
else:
|
37
|
+
self.overall_budget = overall_budget
|
38
|
+
if isinstance(budget_type, str):
|
39
|
+
self.budget_type = BudgetType(budget_type.lower())
|
40
|
+
else:
|
41
|
+
self.budget_type = budget_type
|
42
|
+
|
43
|
+
self.enforcement_mode = enforcement_mode
|
44
|
+
self.command_budgets: Dict[str, CommandBudgetInfo] = {}
|
45
|
+
self.total_usage: Union[int, float] = 0.0 if self.is_cost_budget else 0
|
46
|
+
|
47
|
+
# Initialize pricing engine if needed for cost calculations
|
48
|
+
self.pricing_engine = None
|
49
|
+
if ClaudePricingEngine and (self.is_cost_budget or self.is_mixed_budget):
|
50
|
+
self.pricing_engine = ClaudePricingEngine()
|
51
|
+
|
52
|
+
@property
|
53
|
+
def is_token_budget(self) -> bool:
|
54
|
+
"""Check if this is a token-based budget."""
|
55
|
+
return self.budget_type == BudgetType.TOKENS
|
56
|
+
|
57
|
+
@property
|
58
|
+
def is_cost_budget(self) -> bool:
|
59
|
+
"""Check if this is a cost-based budget."""
|
60
|
+
return self.budget_type == BudgetType.COST
|
61
|
+
|
62
|
+
@property
|
63
|
+
def is_mixed_budget(self) -> bool:
|
64
|
+
"""Check if this is a mixed budget."""
|
65
|
+
return self.budget_type == BudgetType.MIXED
|
66
|
+
|
67
|
+
def set_command_budget(self, command_name: str, limit: int):
|
68
|
+
"""Sets the token budget for a specific command (backward compatibility)."""
|
69
|
+
if command_name in self.command_budgets:
|
70
|
+
# Preserve existing usage when updating budget limit
|
71
|
+
existing_usage = self.command_budgets[command_name].used
|
72
|
+
self.command_budgets[command_name] = CommandBudgetInfo(
|
73
|
+
limit=limit, used=existing_usage, budget_type=BudgetType.TOKENS)
|
74
|
+
else:
|
75
|
+
self.command_budgets[command_name] = CommandBudgetInfo(
|
76
|
+
limit=limit, budget_type=BudgetType.TOKENS)
|
77
|
+
|
78
|
+
def set_command_cost_budget(self, command_name: str, limit: float):
|
79
|
+
"""Sets the cost budget for a specific command in USD."""
|
80
|
+
if command_name in self.command_budgets:
|
81
|
+
# Preserve existing usage when updating budget limit
|
82
|
+
existing_usage = self.command_budgets[command_name].used
|
83
|
+
self.command_budgets[command_name] = CommandBudgetInfo(
|
84
|
+
limit=limit, used=existing_usage, budget_type=BudgetType.COST)
|
85
|
+
else:
|
86
|
+
self.command_budgets[command_name] = CommandBudgetInfo(
|
87
|
+
limit=limit, budget_type=BudgetType.COST)
|
88
|
+
|
89
|
+
def record_usage(self, command_name: str, tokens: int):
|
90
|
+
"""Records token usage for a command and updates the overall total (backward compatibility)."""
|
91
|
+
if self.is_cost_budget and self.pricing_engine:
|
92
|
+
# Convert tokens to cost for cost-based budgets
|
93
|
+
cost = self.convert_tokens_to_cost(tokens)
|
94
|
+
self.record_cost_usage(command_name, cost)
|
95
|
+
else:
|
96
|
+
# Traditional token usage
|
97
|
+
self.total_usage += tokens
|
98
|
+
if command_name in self.command_budgets:
|
99
|
+
self.command_budgets[command_name].used += tokens
|
100
|
+
|
101
|
+
def record_cost_usage(self, command_name: str, cost: float):
|
102
|
+
"""Records cost usage for a command and updates the overall total."""
|
103
|
+
self.total_usage += cost
|
104
|
+
if command_name in self.command_budgets:
|
105
|
+
self.command_budgets[command_name].used += cost
|
106
|
+
|
107
|
+
def check_budget(self, command_name: str, estimated_tokens: int) -> tuple[bool, str]:
|
108
|
+
"""Checks if a command can run based on its budget and the overall budget (backward compatibility).
|
109
|
+
|
110
|
+
Returns:
|
111
|
+
tuple: (can_run: bool, reason: str) - reason explains which budget would be exceeded
|
112
|
+
"""
|
113
|
+
if self.is_cost_budget:
|
114
|
+
# Convert tokens to cost for cost budget checking
|
115
|
+
estimated_cost = self.convert_tokens_to_cost(estimated_tokens)
|
116
|
+
return self.check_cost_budget(command_name, estimated_cost)
|
117
|
+
else:
|
118
|
+
return self._check_token_budget(command_name, estimated_tokens)
|
119
|
+
|
120
|
+
def check_cost_budget(self, command_name: str, estimated_cost: float) -> tuple[bool, str]:
|
121
|
+
"""Checks if a command can run based on cost budgets.
|
122
|
+
|
123
|
+
Returns:
|
124
|
+
tuple: (can_run: bool, reason: str) - reason explains which budget would be exceeded
|
125
|
+
"""
|
126
|
+
# Check overall cost budget FIRST (takes precedence)
|
127
|
+
if self.overall_budget is not None and (self.total_usage + estimated_cost) > self.overall_budget:
|
128
|
+
projected_total = self.total_usage + estimated_cost
|
129
|
+
return False, f"Overall cost budget exceeded: ${projected_total:.4f}/${self.overall_budget:.4f}"
|
130
|
+
|
131
|
+
# Check per-command cost budget
|
132
|
+
if command_name in self.command_budgets:
|
133
|
+
command_budget = self.command_budgets[command_name]
|
134
|
+
if command_budget.is_cost_budget and (command_budget.used + estimated_cost) > command_budget.limit:
|
135
|
+
projected_command = command_budget.used + estimated_cost
|
136
|
+
return False, f"Command '{command_name}' cost budget exceeded: ${projected_command:.4f}/${command_budget.limit:.4f}"
|
137
|
+
|
138
|
+
return True, "Within cost budget limits"
|
139
|
+
|
140
|
+
def _check_token_budget(self, command_name: str, estimated_tokens: int) -> tuple[bool, str]:
|
141
|
+
"""Internal method for checking token budgets."""
|
142
|
+
# Check overall budget FIRST (takes precedence)
|
143
|
+
if self.overall_budget is not None and (self.total_usage + estimated_tokens) > self.overall_budget:
|
144
|
+
projected_total = self.total_usage + estimated_tokens
|
145
|
+
return False, f"Overall budget exceeded: {projected_total}/{self.overall_budget} tokens"
|
146
|
+
|
147
|
+
# Check per-command budget
|
148
|
+
if command_name in self.command_budgets:
|
149
|
+
command_budget = self.command_budgets[command_name]
|
150
|
+
if command_budget.is_token_budget and (command_budget.used + estimated_tokens) > command_budget.limit:
|
151
|
+
projected_command = command_budget.used + estimated_tokens
|
152
|
+
return False, f"Command '{command_name}' budget exceeded: {projected_command}/{command_budget.limit} tokens"
|
153
|
+
|
154
|
+
return True, "Within budget limits"
|
155
|
+
|
156
|
+
def convert_tokens_to_cost(self, tokens: int, model: str = "claude-3-5-sonnet") -> float:
|
157
|
+
"""Convert tokens to cost using the pricing engine."""
|
158
|
+
if not self.pricing_engine:
|
159
|
+
raise AttributeError("Pricing engine not available for cost conversion")
|
160
|
+
|
161
|
+
# Create usage data for cost calculation
|
162
|
+
# Assume equal split between input and output tokens for estimation
|
163
|
+
input_tokens = int(tokens * 0.6) # Estimate 60% input
|
164
|
+
output_tokens = tokens - input_tokens # Remaining 40% output
|
165
|
+
|
166
|
+
usage_data = TokenUsageData(
|
167
|
+
input_tokens=input_tokens,
|
168
|
+
output_tokens=output_tokens,
|
169
|
+
model=model
|
170
|
+
)
|
171
|
+
|
172
|
+
cost_breakdown = self.pricing_engine.calculate_cost(usage_data)
|
173
|
+
return cost_breakdown.total_cost
|
174
|
+
|
175
|
+
def convert_cost_to_tokens(self, cost: float, model: str = "claude-3-5-sonnet") -> int:
|
176
|
+
"""Convert cost to approximate token count using the pricing engine."""
|
177
|
+
if not self.pricing_engine:
|
178
|
+
raise AttributeError("Pricing engine not available for cost conversion")
|
179
|
+
|
180
|
+
# Get model pricing
|
181
|
+
model_pricing = self.pricing_engine.pricing_config.MODEL_PRICING.get(
|
182
|
+
model, self.pricing_engine.pricing_config.MODEL_PRICING["claude-3-5-sonnet"]
|
183
|
+
)
|
184
|
+
|
185
|
+
# Use average of input and output pricing for estimation
|
186
|
+
avg_price_per_million = (model_pricing["input"] + model_pricing["output"]) / 2
|
187
|
+
tokens = int((cost / avg_price_per_million) * 1_000_000)
|
188
|
+
|
189
|
+
return tokens
|
190
|
+
|
191
|
+
def set_budget_parameter_type(self, budget_type: Union[str, BudgetType]):
|
192
|
+
"""Set the budget parameter type (tokens, cost, or mixed)."""
|
193
|
+
if isinstance(budget_type, str):
|
194
|
+
self.budget_type = BudgetType(budget_type.lower())
|
195
|
+
else:
|
196
|
+
self.budget_type = budget_type
|
197
|
+
|
198
|
+
# Initialize pricing engine if switching to cost or mixed mode
|
199
|
+
if ClaudePricingEngine and (self.is_cost_budget or self.is_mixed_budget) and not self.pricing_engine:
|
200
|
+
self.pricing_engine = ClaudePricingEngine()
|
token_budget/models.py
ADDED
@@ -0,0 +1,74 @@
|
|
1
|
+
"""Token budget data models - enhanced implementation with cost support."""
|
2
|
+
|
3
|
+
from typing import Dict, Optional, Union
|
4
|
+
from enum import Enum
|
5
|
+
|
6
|
+
class BudgetType(Enum):
|
7
|
+
"""Types of budget tracking supported."""
|
8
|
+
TOKENS = "tokens"
|
9
|
+
COST = "cost"
|
10
|
+
MIXED = "mixed"
|
11
|
+
|
12
|
+
class CommandBudgetInfo:
|
13
|
+
"""Tracks the budget status for a single command with support for both tokens and cost."""
|
14
|
+
|
15
|
+
def __init__(self, limit: Union[int, float], used: Union[int, float] = 0,
|
16
|
+
budget_type: BudgetType = BudgetType.TOKENS):
|
17
|
+
"""
|
18
|
+
Initialize command budget info.
|
19
|
+
|
20
|
+
Args:
|
21
|
+
limit: Budget limit (tokens as int, cost as float)
|
22
|
+
used: Current usage (tokens as int, cost as float)
|
23
|
+
budget_type: Type of budget (tokens, cost, or mixed)
|
24
|
+
"""
|
25
|
+
self.limit = limit
|
26
|
+
self.used = used
|
27
|
+
self.budget_type = budget_type
|
28
|
+
|
29
|
+
@property
|
30
|
+
def remaining(self) -> Union[int, float]:
|
31
|
+
"""Get remaining budget amount."""
|
32
|
+
return self.limit - self.used
|
33
|
+
|
34
|
+
@property
|
35
|
+
def percentage(self) -> float:
|
36
|
+
"""Get percentage of budget used."""
|
37
|
+
return (self.used / self.limit * 100) if self.limit > 0 else 0
|
38
|
+
|
39
|
+
@property
|
40
|
+
def is_token_budget(self) -> bool:
|
41
|
+
"""Check if this is a token-based budget."""
|
42
|
+
return self.budget_type == BudgetType.TOKENS
|
43
|
+
|
44
|
+
@property
|
45
|
+
def is_cost_budget(self) -> bool:
|
46
|
+
"""Check if this is a cost-based budget."""
|
47
|
+
return self.budget_type == BudgetType.COST
|
48
|
+
|
49
|
+
@property
|
50
|
+
def is_mixed_budget(self) -> bool:
|
51
|
+
"""Check if this is a mixed budget (both tokens and cost)."""
|
52
|
+
return self.budget_type == BudgetType.MIXED
|
53
|
+
|
54
|
+
def format_limit(self) -> str:
|
55
|
+
"""Format the limit for display."""
|
56
|
+
if self.is_cost_budget:
|
57
|
+
return f"${self.limit:.4f}"
|
58
|
+
else:
|
59
|
+
return f"{int(self.limit)} tokens"
|
60
|
+
|
61
|
+
def format_used(self) -> str:
|
62
|
+
"""Format the used amount for display."""
|
63
|
+
if self.is_cost_budget:
|
64
|
+
return f"${self.used:.4f}"
|
65
|
+
else:
|
66
|
+
return f"{int(self.used)} tokens"
|
67
|
+
|
68
|
+
def format_remaining(self) -> str:
|
69
|
+
"""Format the remaining amount for display."""
|
70
|
+
remaining = self.remaining
|
71
|
+
if self.is_cost_budget:
|
72
|
+
return f"${remaining:.4f}"
|
73
|
+
else:
|
74
|
+
return f"{int(remaining)} tokens"
|
@@ -0,0 +1,22 @@
|
|
1
|
+
"""Token budget visualization utilities."""
|
2
|
+
|
3
|
+
def render_progress_bar(used: int, total: int, width: int = 20) -> str:
|
4
|
+
"""Renders an ASCII progress bar."""
|
5
|
+
if total == 0:
|
6
|
+
return "[NO BUDGET SET]"
|
7
|
+
|
8
|
+
percentage = min(used / total, 1.0)
|
9
|
+
filled_width = int(percentage * width)
|
10
|
+
|
11
|
+
# Use ASCII characters for Windows compatibility
|
12
|
+
bar = '#' * filled_width + '-' * (width - filled_width)
|
13
|
+
|
14
|
+
# Color coding (ANSI escape codes)
|
15
|
+
color_start = '\033[92m' # Green
|
16
|
+
if percentage > 0.9:
|
17
|
+
color_start = '\033[91m' # Red
|
18
|
+
elif percentage > 0.7:
|
19
|
+
color_start = '\033[93m' # Yellow
|
20
|
+
color_end = '\033[0m' # Reset
|
21
|
+
|
22
|
+
return f"[{color_start}{bar}{color_end}] {percentage:.0%}"
|
@@ -0,0 +1,20 @@
|
|
1
|
+
"""
|
2
|
+
Token Transparency Module for Zen Claude Orchestrator
|
3
|
+
|
4
|
+
This module provides transparent token usage tracking and cost calculation
|
5
|
+
for Claude Code instances, ensuring compliance with official Claude pricing.
|
6
|
+
"""
|
7
|
+
|
8
|
+
from .claude_pricing_engine import (
|
9
|
+
ClaudePricingEngine,
|
10
|
+
ClaudePricingConfig,
|
11
|
+
TokenUsageData,
|
12
|
+
CostBreakdown
|
13
|
+
)
|
14
|
+
|
15
|
+
__all__ = [
|
16
|
+
'ClaudePricingEngine',
|
17
|
+
'ClaudePricingConfig',
|
18
|
+
'TokenUsageData',
|
19
|
+
'CostBreakdown'
|
20
|
+
]
|
@@ -0,0 +1,327 @@
|
|
1
|
+
#!/usr/bin/env python3
|
2
|
+
"""
|
3
|
+
Claude Code Pricing Compliance Engine
|
4
|
+
|
5
|
+
Provides accurate token counting and cost calculation based on official Claude pricing.
|
6
|
+
Designed to be the SSOT for all Claude Code pricing calculations within zen.
|
7
|
+
|
8
|
+
Key Features:
|
9
|
+
- Model detection from API responses
|
10
|
+
- Accurate cache pricing based on duration
|
11
|
+
- Tool cost calculation
|
12
|
+
- Compliance with Claude pricing documentation
|
13
|
+
- Extensible for future Claude Code agent support
|
14
|
+
"""
|
15
|
+
|
16
|
+
from dataclasses import dataclass
|
17
|
+
from typing import Dict, Optional, Tuple, Any
|
18
|
+
import re
|
19
|
+
import json
|
20
|
+
import logging
|
21
|
+
|
22
|
+
logger = logging.getLogger(__name__)
|
23
|
+
|
24
|
+
@dataclass
|
25
|
+
class ClaudePricingConfig:
|
26
|
+
"""Current Claude pricing rates as of 2024-2025"""
|
27
|
+
|
28
|
+
# Model pricing per million tokens (input, output)
|
29
|
+
MODEL_PRICING = {
|
30
|
+
"claude-opus-4": {"input": 15.0, "output": 75.0},
|
31
|
+
"claude-opus-4.1": {"input": 15.0, "output": 75.0},
|
32
|
+
"claude-sonnet-4": {"input": 3.0, "output": 15.0},
|
33
|
+
"claude-sonnet-3.7": {"input": 3.0, "output": 15.0},
|
34
|
+
"claude-3-5-sonnet": {"input": 3.0, "output": 15.0},
|
35
|
+
"claude-haiku-3.5": {"input": 0.8, "output": 4.0},
|
36
|
+
}
|
37
|
+
|
38
|
+
# Cache pricing multipliers
|
39
|
+
CACHE_READ_MULTIPLIER = 0.1 # 10% of base input price
|
40
|
+
CACHE_5MIN_WRITE_MULTIPLIER = 1.25 # 25% premium
|
41
|
+
CACHE_1HOUR_WRITE_MULTIPLIER = 2.0 # 100% premium
|
42
|
+
|
43
|
+
# Tool pricing (per 1000 calls)
|
44
|
+
TOOL_PRICING = {
|
45
|
+
"web_search": 10.0, # $10 per 1000 searches
|
46
|
+
"web_fetch": 0.0, # No additional charge
|
47
|
+
"default": 0.0 # Most tools have no additional charge
|
48
|
+
}
|
49
|
+
|
50
|
+
@dataclass
|
51
|
+
class TokenUsageData:
|
52
|
+
"""Token usage data with detailed breakdown"""
|
53
|
+
input_tokens: int = 0
|
54
|
+
output_tokens: int = 0
|
55
|
+
cache_read_tokens: int = 0
|
56
|
+
cache_creation_tokens: int = 0
|
57
|
+
cache_type: str = "5min" # "5min" or "1hour"
|
58
|
+
total_tokens: int = 0
|
59
|
+
tool_calls: int = 0
|
60
|
+
model: str = "claude-3-5-sonnet"
|
61
|
+
|
62
|
+
def __post_init__(self):
|
63
|
+
"""Calculate total if not provided"""
|
64
|
+
if self.total_tokens == 0:
|
65
|
+
self.total_tokens = (self.input_tokens + self.output_tokens +
|
66
|
+
self.cache_read_tokens + self.cache_creation_tokens)
|
67
|
+
|
68
|
+
@dataclass
|
69
|
+
class CostBreakdown:
|
70
|
+
"""Detailed cost breakdown for transparency"""
|
71
|
+
input_cost: float = 0.0
|
72
|
+
output_cost: float = 0.0
|
73
|
+
cache_read_cost: float = 0.0
|
74
|
+
cache_creation_cost: float = 0.0
|
75
|
+
tool_cost: float = 0.0
|
76
|
+
total_cost: float = 0.0
|
77
|
+
model_used: str = ""
|
78
|
+
cache_type: str = ""
|
79
|
+
|
80
|
+
def __post_init__(self):
|
81
|
+
"""Calculate total cost"""
|
82
|
+
self.total_cost = (self.input_cost + self.output_cost +
|
83
|
+
self.cache_read_cost + self.cache_creation_cost + self.tool_cost)
|
84
|
+
|
85
|
+
class ClaudePricingEngine:
|
86
|
+
"""
|
87
|
+
Claude Code pricing compliance engine for accurate cost calculation.
|
88
|
+
|
89
|
+
Ensures compliance with official Claude pricing documentation and provides
|
90
|
+
detailed transparency for token usage costs.
|
91
|
+
"""
|
92
|
+
|
93
|
+
def __init__(self):
|
94
|
+
self.pricing_config = ClaudePricingConfig()
|
95
|
+
|
96
|
+
def detect_model_from_response(self, response_data: Dict[str, Any]) -> str:
|
97
|
+
"""
|
98
|
+
Detect Claude model from API response or usage data.
|
99
|
+
|
100
|
+
Args:
|
101
|
+
response_data: API response or usage data containing model information
|
102
|
+
|
103
|
+
Returns:
|
104
|
+
Model name string, defaults to claude-3-5-sonnet if not detected
|
105
|
+
"""
|
106
|
+
# Try multiple locations where model might be specified
|
107
|
+
model_locations = [
|
108
|
+
response_data.get('model'),
|
109
|
+
response_data.get('model_name'),
|
110
|
+
response_data.get('usage', {}).get('model'),
|
111
|
+
response_data.get('message', {}).get('model'),
|
112
|
+
response_data.get('metadata', {}).get('model')
|
113
|
+
]
|
114
|
+
|
115
|
+
for model in model_locations:
|
116
|
+
if model and isinstance(model, str):
|
117
|
+
# Normalize model name
|
118
|
+
normalized = self._normalize_model_name(model)
|
119
|
+
if normalized in self.pricing_config.MODEL_PRICING:
|
120
|
+
return normalized
|
121
|
+
|
122
|
+
# Default fallback
|
123
|
+
logger.debug("Model not detected in response, defaulting to claude-3-5-sonnet")
|
124
|
+
return "claude-3-5-sonnet"
|
125
|
+
|
126
|
+
def _normalize_model_name(self, model_name: str) -> str:
|
127
|
+
"""Normalize model name to match pricing config keys"""
|
128
|
+
model_name = model_name.lower().strip()
|
129
|
+
|
130
|
+
# Handle various model name formats
|
131
|
+
if "opus" in model_name:
|
132
|
+
if "4.1" in model_name:
|
133
|
+
return "claude-opus-4.1"
|
134
|
+
elif "4" in model_name:
|
135
|
+
return "claude-opus-4"
|
136
|
+
elif "sonnet" in model_name:
|
137
|
+
if "4" in model_name:
|
138
|
+
return "claude-sonnet-4"
|
139
|
+
elif "3.7" in model_name:
|
140
|
+
return "claude-sonnet-3.7"
|
141
|
+
elif "3.5" in model_name or "3-5" in model_name:
|
142
|
+
return "claude-3-5-sonnet"
|
143
|
+
elif "haiku" in model_name:
|
144
|
+
if "3.5" in model_name:
|
145
|
+
return "claude-haiku-3.5"
|
146
|
+
|
147
|
+
return model_name
|
148
|
+
|
149
|
+
def detect_cache_type(self, response_data: Dict[str, Any]) -> str:
|
150
|
+
"""
|
151
|
+
Detect cache type (5min vs 1hour) from response data.
|
152
|
+
|
153
|
+
Args:
|
154
|
+
response_data: API response data
|
155
|
+
|
156
|
+
Returns:
|
157
|
+
"5min" or "1hour", defaults to "5min"
|
158
|
+
"""
|
159
|
+
# Look for cache type indicators in response
|
160
|
+
cache_indicators = [
|
161
|
+
response_data.get('cache_type'),
|
162
|
+
response_data.get('usage', {}).get('cache_type'),
|
163
|
+
response_data.get('metadata', {}).get('cache_type')
|
164
|
+
]
|
165
|
+
|
166
|
+
for indicator in cache_indicators:
|
167
|
+
if indicator:
|
168
|
+
if "1hour" in str(indicator).lower() or "60min" in str(indicator).lower():
|
169
|
+
return "1hour"
|
170
|
+
elif "5min" in str(indicator).lower():
|
171
|
+
return "5min"
|
172
|
+
|
173
|
+
# Default to 5min cache
|
174
|
+
return "5min"
|
175
|
+
|
176
|
+
def calculate_cost(self, usage_data: TokenUsageData,
|
177
|
+
authoritative_cost: Optional[float] = None,
|
178
|
+
tool_tokens: Optional[Dict[str, int]] = None) -> CostBreakdown:
|
179
|
+
"""
|
180
|
+
Calculate detailed cost breakdown with Claude pricing compliance.
|
181
|
+
|
182
|
+
Args:
|
183
|
+
usage_data: Token usage information
|
184
|
+
authoritative_cost: SDK-provided cost (preferred when available)
|
185
|
+
tool_tokens: Dictionary of tool names to token counts for tool cost calculation
|
186
|
+
|
187
|
+
Returns:
|
188
|
+
Detailed cost breakdown for transparency
|
189
|
+
"""
|
190
|
+
# Use authoritative cost if provided (most accurate)
|
191
|
+
if authoritative_cost is not None:
|
192
|
+
breakdown = CostBreakdown(
|
193
|
+
model_used=usage_data.model,
|
194
|
+
cache_type=usage_data.cache_type
|
195
|
+
)
|
196
|
+
breakdown.total_cost = authoritative_cost
|
197
|
+
return breakdown
|
198
|
+
|
199
|
+
# Get model pricing
|
200
|
+
model_pricing = self.pricing_config.MODEL_PRICING.get(
|
201
|
+
usage_data.model,
|
202
|
+
self.pricing_config.MODEL_PRICING["claude-3-5-sonnet"]
|
203
|
+
)
|
204
|
+
|
205
|
+
# Calculate base costs
|
206
|
+
input_cost = (usage_data.input_tokens / 1_000_000) * model_pricing["input"]
|
207
|
+
output_cost = (usage_data.output_tokens / 1_000_000) * model_pricing["output"]
|
208
|
+
|
209
|
+
# Calculate cache costs with correct multipliers
|
210
|
+
cache_read_cost = (usage_data.cache_read_tokens / 1_000_000) * \
|
211
|
+
(model_pricing["input"] * self.pricing_config.CACHE_READ_MULTIPLIER)
|
212
|
+
|
213
|
+
# Cache creation cost depends on cache type
|
214
|
+
cache_multiplier = (self.pricing_config.CACHE_1HOUR_WRITE_MULTIPLIER
|
215
|
+
if usage_data.cache_type == "1hour"
|
216
|
+
else self.pricing_config.CACHE_5MIN_WRITE_MULTIPLIER)
|
217
|
+
|
218
|
+
cache_creation_cost = (usage_data.cache_creation_tokens / 1_000_000) * \
|
219
|
+
(model_pricing["input"] * cache_multiplier)
|
220
|
+
|
221
|
+
# Calculate tool costs based on token usage
|
222
|
+
tool_cost = 0.0
|
223
|
+
if tool_tokens:
|
224
|
+
for tool_name, tokens in tool_tokens.items():
|
225
|
+
# Tool tokens are charged at the same rate as input tokens for the model
|
226
|
+
tool_cost += (tokens / 1_000_000) * model_pricing["input"]
|
227
|
+
|
228
|
+
return CostBreakdown(
|
229
|
+
input_cost=input_cost,
|
230
|
+
output_cost=output_cost,
|
231
|
+
cache_read_cost=cache_read_cost,
|
232
|
+
cache_creation_cost=cache_creation_cost,
|
233
|
+
tool_cost=tool_cost,
|
234
|
+
model_used=usage_data.model,
|
235
|
+
cache_type=usage_data.cache_type
|
236
|
+
)
|
237
|
+
|
238
|
+
def parse_claude_response(self, response_line: str) -> Optional[TokenUsageData]:
|
239
|
+
"""
|
240
|
+
Parse token usage from Claude Code response line with model detection.
|
241
|
+
|
242
|
+
Args:
|
243
|
+
response_line: Single line from Claude Code output
|
244
|
+
|
245
|
+
Returns:
|
246
|
+
TokenUsageData if parsing successful, None otherwise
|
247
|
+
"""
|
248
|
+
line = response_line.strip()
|
249
|
+
if not line.startswith('{'):
|
250
|
+
return None
|
251
|
+
|
252
|
+
try:
|
253
|
+
json_data = json.loads(line)
|
254
|
+
|
255
|
+
# Detect model and cache type
|
256
|
+
model = self.detect_model_from_response(json_data)
|
257
|
+
cache_type = self.detect_cache_type(json_data)
|
258
|
+
|
259
|
+
# Extract usage data
|
260
|
+
usage_data = None
|
261
|
+
if 'usage' in json_data:
|
262
|
+
usage_data = json_data['usage']
|
263
|
+
elif 'message' in json_data and isinstance(json_data['message'], dict):
|
264
|
+
usage_data = json_data['message'].get('usage')
|
265
|
+
|
266
|
+
if usage_data and isinstance(usage_data, dict):
|
267
|
+
return TokenUsageData(
|
268
|
+
input_tokens=int(usage_data.get('input_tokens', 0)),
|
269
|
+
output_tokens=int(usage_data.get('output_tokens', 0)),
|
270
|
+
cache_read_tokens=int(usage_data.get('cache_read_input_tokens', 0)),
|
271
|
+
cache_creation_tokens=int(usage_data.get('cache_creation_input_tokens', 0)),
|
272
|
+
total_tokens=int(usage_data.get('total_tokens', 0)),
|
273
|
+
model=model,
|
274
|
+
cache_type=cache_type
|
275
|
+
)
|
276
|
+
|
277
|
+
except (json.JSONDecodeError, ValueError, KeyError) as e:
|
278
|
+
logger.debug(f"Failed to parse Claude response: {e}")
|
279
|
+
|
280
|
+
return None
|
281
|
+
|
282
|
+
def get_transparency_report(self, usage_data: TokenUsageData,
|
283
|
+
cost_breakdown: CostBreakdown,
|
284
|
+
tool_tokens: Optional[Dict[str, int]] = None) -> Dict[str, Any]:
|
285
|
+
"""
|
286
|
+
Generate transparency report for token usage and costs.
|
287
|
+
|
288
|
+
Args:
|
289
|
+
usage_data: Token usage information
|
290
|
+
cost_breakdown: Detailed cost breakdown
|
291
|
+
tool_tokens: Tool-specific token usage
|
292
|
+
|
293
|
+
Returns:
|
294
|
+
Comprehensive transparency report
|
295
|
+
"""
|
296
|
+
return {
|
297
|
+
"model_used": usage_data.model,
|
298
|
+
"cache_type": usage_data.cache_type,
|
299
|
+
"token_breakdown": {
|
300
|
+
"input_tokens": usage_data.input_tokens,
|
301
|
+
"output_tokens": usage_data.output_tokens,
|
302
|
+
"cache_read_tokens": usage_data.cache_read_tokens,
|
303
|
+
"cache_creation_tokens": usage_data.cache_creation_tokens,
|
304
|
+
"total_tokens": usage_data.total_tokens,
|
305
|
+
"tool_tokens": tool_tokens or {}
|
306
|
+
},
|
307
|
+
"cost_breakdown": {
|
308
|
+
"input_cost_usd": round(cost_breakdown.input_cost, 6),
|
309
|
+
"output_cost_usd": round(cost_breakdown.output_cost, 6),
|
310
|
+
"cache_read_cost_usd": round(cost_breakdown.cache_read_cost, 6),
|
311
|
+
"cache_creation_cost_usd": round(cost_breakdown.cache_creation_cost, 6),
|
312
|
+
"tool_cost_usd": round(cost_breakdown.tool_cost, 6),
|
313
|
+
"total_cost_usd": round(cost_breakdown.total_cost, 6)
|
314
|
+
},
|
315
|
+
"pricing_rates": {
|
316
|
+
"model_rates": self.pricing_config.MODEL_PRICING[usage_data.model],
|
317
|
+
"cache_read_multiplier": self.pricing_config.CACHE_READ_MULTIPLIER,
|
318
|
+
"cache_write_multiplier": (self.pricing_config.CACHE_1HOUR_WRITE_MULTIPLIER
|
319
|
+
if usage_data.cache_type == "1hour"
|
320
|
+
else self.pricing_config.CACHE_5MIN_WRITE_MULTIPLIER)
|
321
|
+
},
|
322
|
+
"compliance_info": {
|
323
|
+
"pricing_source": "https://docs.claude.com/en/docs/about-claude/pricing",
|
324
|
+
"last_updated": "2024-2025",
|
325
|
+
"model_detected": usage_data.model != "claude-3-5-sonnet"
|
326
|
+
}
|
327
|
+
}
|