| 1 | import hashlib |
| 2 | import math |
| 3 | from typing import Any |
| 4 | |
| 5 | from helpers import files, history, skills, tokens |
| 6 | from helpers.llm_result import result_from_metadata |
| 7 | |
| 8 | |
| 9 | PARTS_KEY = "context_window_usage" |
| 10 | CACHE_KEY = "_context_window_usage_cache" |
| 11 | PROVIDER_USAGE_KEY = "context_window_provider_usage" |
| 12 | USAGE_KEYS = ( |
| 13 | "messages", |
| 14 | "system_tools", |
| 15 | "skills", |
| 16 | "mcp_tools", |
| 17 | "system_prompt", |
| 18 | "extras", |
| 19 | ) |
| 20 | MEASURED_KEYS = tuple(key for key in USAGE_KEYS if key != "system_prompt") |
| 21 | |
| 22 | |
| 23 | def reset(agent: Any) -> None: |
| 24 | params = _temporary_params(agent) |
| 25 | if params is not None: |
| 26 | params[PARTS_KEY] = {} |
| 27 | |
| 28 | |
| 29 | def discard(agent: Any) -> None: |
| 30 | params = _temporary_params(agent) |
| 31 | if params is not None: |
| 32 | params.pop(PARTS_KEY, None) |
| 33 | |
| 34 | |
| 35 | def record_prompt(agent: Any, key: str, prompt: Any) -> None: |
| 36 | parts = _parts(agent) |
| 37 | if parts is None or key not in MEASURED_KEYS: |
| 38 | return |
| 39 | text = files.remove_code_fences(str(prompt or ""), language="json") |
| 40 | parts[key] = _cached_tokens(agent, f"prompt:{key}", text) |
| 41 | |
| 42 | |
| 43 | def capture_context(agent: Any, loop_data: Any) -> None: |
| 44 | parts = _parts(agent) |
| 45 | if parts is None or loop_data is None: |
| 46 | return |
| 47 | |
| 48 | output = list(getattr(loop_data, "history_output", None) or []) |
| 49 | parts["_history_output"] = output |
| 50 | skill_output = [message for message in output if skills.skill_instruction_name(message)] |
| 51 | skill_tokens = _output_tokens(agent, "history_skills", skill_output) |
| 52 | parts["messages"] = max(_history_tokens(agent, output) - skill_tokens, 0) |
| 53 | parts["skills"] = parts.get("skills", 0) + skill_tokens |
| 54 | |
| 55 | protocol_values = { |
| 56 | **getattr(loop_data, "protocol_persistent", {}), |
| 57 | **getattr(loop_data, "protocol_temporary", {}), |
| 58 | } |
| 59 | extras_values = { |
| 60 | **getattr(loop_data, "extras_persistent", {}), |
| 61 | **getattr(loop_data, "extras_temporary", {}), |
| 62 | } |
| 63 | protocol = agent._build_context_message( |
| 64 | "agent.context.protocol.md", |
| 65 | "protocol", |
| 66 | protocol_values, |
| 67 | include_empty=False, |
| 68 | ) |
| 69 | extras = agent._build_context_message( |
| 70 | "agent.context.extras.md", |
| 71 | "extras", |
| 72 | extras_values, |
| 73 | include_empty=True, |
| 74 | ) |
| 75 | parts["extras"] = _output_tokens(agent, "extras", protocol + extras) |
| 76 | |
| 77 | |
| 78 | def finalize(agent: Any) -> None: |
| 79 | params = _temporary_params(agent) |
| 80 | parts = params.pop(PARTS_KEY, None) if params is not None else None |
| 81 | window = agent.get_data(agent.DATA_NAME_CTX_WINDOW) if agent else None |
| 82 | if not isinstance(parts, dict) or not isinstance(window, dict): |
| 83 | return |
| 84 | |
| 85 | history_output = parts.pop("_history_output", None) |
| 86 | total = _non_negative_int(window.get("tokens")) |
| 87 | usage = {key: _non_negative_int(parts.get(key)) for key in MEASURED_KEYS} |
| 88 | measured_total = sum(usage.values()) |
| 89 | if total and measured_total >= total and isinstance(history_output, list): |
| 90 | message_output = [ |
| 91 | message |
| 92 | for message in history_output |
| 93 | if not skills.skill_instruction_name(message) |
| 94 | ] |
| 95 | usage["messages"] = _output_tokens( |
| 96 | agent, "history_messages", message_output |
| 97 | ) |
| 98 | measured_total = sum(usage.values()) |
| 99 | if measured_total > total and measured_total: |
| 100 | usage = _scale_to_total(usage, total, measured_total) |
| 101 | measured_total = total |
| 102 | usage["system_prompt"] = total - measured_total |
| 103 | usage = {key: usage.get(key, 0) for key in USAGE_KEYS} |
| 104 | |
| 105 | updated = dict(window) |
| 106 | updated["usage"] = usage |
| 107 | agent.set_data(agent.DATA_NAME_CTX_WINDOW, updated) |
| 108 | |
| 109 | |
| 110 | def usage_snapshot(value: Any) -> dict[str, int]: |
| 111 | if not isinstance(value, dict): |
| 112 | return {} |
| 113 | return {key: _non_negative_int(value.get(key)) for key in USAGE_KEYS} |
| 114 | |
| 115 | |
| 116 | def capture_provider_usage(agent: Any, result: Any) -> None: |
| 117 | if agent is None or result is None or not hasattr(result, "usage"): |
| 118 | return |
| 119 | |
| 120 | snapshot = provider_usage_snapshot(getattr(result, "usage", None)) |
| 121 | agent.set_data( |
| 122 | PROVIDER_USAGE_KEY, |
| 123 | snapshot if snapshot else {"available": False}, |
| 124 | ) |
| 125 | |
| 126 | |
| 127 | def latest_provider_usage(agent: Any) -> dict[str, int | float]: |
| 128 | data = getattr(agent, "data", None) |
| 129 | if isinstance(data, dict) and PROVIDER_USAGE_KEY in data: |
| 130 | stored = data.get(PROVIDER_USAGE_KEY) |
| 131 | if isinstance(stored, dict) and stored.get("available") is False: |
| 132 | return {} |
| 133 | return provider_usage_snapshot(stored) |
| 134 | |
| 135 | all_messages = getattr(getattr(agent, "history", None), "all_messages", None) |
| 136 | if not callable(all_messages): |
| 137 | return {} |
| 138 | for message in reversed(all_messages()): |
| 139 | if not getattr(message, "ai", False): |
| 140 | continue |
| 141 | result = result_from_metadata(getattr(message, "metadata", None)) |
| 142 | if result: |
| 143 | return provider_usage_snapshot(result.usage) |
| 144 | return {} |
| 145 | |
| 146 | |
| 147 | def provider_usage_snapshot(value: Any) -> dict[str, int | float]: |
| 148 | if not isinstance(value, dict): |
| 149 | return {} |
| 150 | |
| 151 | input_details = { |
| 152 | **_mapping(value.get("prompt_tokens_details")), |
| 153 | **_mapping(value.get("input_tokens_details")), |
| 154 | } |
| 155 | result: dict[str, int | float] = {} |
| 156 | fields = { |
| 157 | "input_tokens": (value.get("input_tokens"), value.get("prompt_tokens")), |
| 158 | "cached_tokens": ( |
| 159 | input_details.get("cached_tokens"), |
| 160 | input_details.get("cache_read_tokens"), |
| 161 | value.get("cache_read_input_tokens"), |
| 162 | value.get("cached_tokens"), |
| 163 | ), |
| 164 | "output_tokens": ( |
| 165 | value.get("output_tokens"), |
| 166 | value.get("completion_tokens"), |
| 167 | ), |
| 168 | } |
| 169 | for key, values in fields.items(): |
| 170 | number = _optional_non_negative_int(*values) |
| 171 | if number is not None: |
| 172 | result[key] = number |
| 173 | |
| 174 | cost = _optional_non_negative_float( |
| 175 | value.get("cost"), value.get("response_cost") |
| 176 | ) |
| 177 | if cost is not None: |
| 178 | result["cost"] = cost |
| 179 | return result |
| 180 | |
| 181 | |
| 182 | def _parts(agent: Any) -> dict[str, Any] | None: |
| 183 | params = _temporary_params(agent) |
| 184 | value = params.get(PARTS_KEY) if params is not None else None |
| 185 | return value if isinstance(value, dict) else None |
| 186 | |
| 187 | |
| 188 | def _temporary_params(agent: Any) -> dict[str, Any] | None: |
| 189 | loop_data = getattr(agent, "loop_data", None) |
| 190 | params = getattr(loop_data, "params_temporary", None) |
| 191 | return params if isinstance(params, dict) else None |
| 192 | |
| 193 | |
| 194 | def _history_tokens(agent: Any, output: list[history.OutputMessage]) -> int: |
| 195 | get_tokens = getattr(getattr(agent, "history", None), "get_tokens", None) |
| 196 | if callable(get_tokens): |
| 197 | return _non_negative_int(get_tokens()) |
| 198 | return _count_output_tokens(output) |
| 199 | |
| 200 | |
| 201 | def _output_tokens( |
| 202 | agent: Any, cache_key: str, output: list[history.OutputMessage] |
| 203 | ) -> int: |
| 204 | text = history.output_text(output, ai_label="assistant", human_label="user") |
| 205 | return _cached_tokens(agent, cache_key, text) |
| 206 | |
| 207 | |
| 208 | def _count_output_tokens(output: list[history.OutputMessage]) -> int: |
| 209 | text = history.output_text(output, ai_label="assistant", human_label="user") |
| 210 | return tokens.approximate_prompt_tokens(text) |
| 211 | |
| 212 | |
| 213 | def _cached_tokens(agent: Any, key: str, text: str) -> int: |
| 214 | cache = _cache(agent) |
| 215 | digest = hashlib.sha256(text.encode("utf-8")).hexdigest() |
| 216 | cached = cache.get(key) if cache is not None else None |
| 217 | if isinstance(cached, tuple) and len(cached) == 2 and cached[0] == digest: |
| 218 | return _non_negative_int(cached[1]) |
| 219 | |
| 220 | count = tokens.approximate_prompt_tokens(text) |
| 221 | if cache is not None: |
| 222 | cache[key] = (digest, count) |
| 223 | return count |
| 224 | |
| 225 | |
| 226 | def _cache(agent: Any) -> dict[str, tuple[str, int]] | None: |
| 227 | data = getattr(agent, "data", None) |
| 228 | if not isinstance(data, dict): |
| 229 | return None |
| 230 | cache = data.get(CACHE_KEY) |
| 231 | if not isinstance(cache, dict): |
| 232 | cache = {} |
| 233 | data[CACHE_KEY] = cache |
| 234 | return cache |
| 235 | |
| 236 | |
| 237 | def _non_negative_int(value: Any) -> int: |
| 238 | try: |
| 239 | return max(int(value or 0), 0) |
| 240 | except (TypeError, ValueError): |
| 241 | return 0 |
| 242 | |
| 243 | |
| 244 | def _mapping(value: Any) -> dict[str, Any]: |
| 245 | return value if isinstance(value, dict) else {} |
| 246 | |
| 247 | |
| 248 | def _optional_non_negative_int(*values: Any) -> int | None: |
| 249 | for value in values: |
| 250 | if value is None: |
| 251 | continue |
| 252 | try: |
| 253 | return max(int(value), 0) |
| 254 | except (TypeError, ValueError): |
| 255 | continue |
| 256 | return None |
| 257 | |
| 258 | |
| 259 | def _optional_non_negative_float(*values: Any) -> float | None: |
| 260 | for value in values: |
| 261 | if value is None: |
| 262 | continue |
| 263 | try: |
| 264 | number = float(value) |
| 265 | except (TypeError, ValueError): |
| 266 | continue |
| 267 | if math.isfinite(number): |
| 268 | return max(number, 0) |
| 269 | return None |
| 270 | |
| 271 | |
| 272 | def _scale_to_total(values: dict[str, int], total: int, current: int) -> dict[str, int]: |
| 273 | scaled = {key: value * total // current for key, value in values.items()} |
| 274 | remainder = total - sum(scaled.values()) |
| 275 | order = sorted(values, key=lambda key: values[key] * total % current, reverse=True) |
| 276 | for key in order[:remainder]: |
| 277 | scaled[key] += 1 |
| 278 | return scaled |