|
13 | 13 | import inspect |
14 | 14 | import json |
15 | 15 | import logging |
| 16 | +import math |
16 | 17 | import os |
17 | 18 | import re |
18 | 19 | import subprocess |
@@ -300,6 +301,113 @@ def _find_matching_close_paren(text: str, open_pos: int) -> Optional[int]: |
300 | 301 | return None |
301 | 302 |
|
302 | 303 |
|
| 304 | +def _safe_number(value: Any) -> int: |
| 305 | + """Coerce a usage-stat value to a non-negative int; anything else |
| 306 | + (string, None, nested structure, bool) is untrusted input and yields 0 |
| 307 | + rather than raising — a malformed stat must never break the run that |
| 308 | + carries it.""" |
| 309 | + if isinstance(value, bool): |
| 310 | + return 0 |
| 311 | + if isinstance(value, (int, float)) and value >= 0: |
| 312 | + return int(value) |
| 313 | + return 0 |
| 314 | + |
| 315 | + |
| 316 | +def _sum_conversation_tokens( |
| 317 | + conversation: List[Dict[str, Any]], |
| 318 | + tool_usage_entries: Optional[List[Dict[str, Any]]] = None, |
| 319 | +) -> Tuple[int, int]: |
| 320 | + """Sum input/output tokens from per-step 'stats' entries already appended |
| 321 | + to conversation, plus any tool-reported usage folded in separately (see |
| 322 | + ``_extract_tool_usage``). Returns (total_input, total_output).""" |
| 323 | + total_input = 0 |
| 324 | + total_output = 0 |
| 325 | + for entry in conversation: |
| 326 | + if entry.get("role") == "system" and isinstance(entry.get("content"), dict): |
| 327 | + content = entry["content"] |
| 328 | + if content.get("type") == "stats" and "performance_stats" in content: |
| 329 | + stats = content["performance_stats"] |
| 330 | + total_input += _safe_number(stats.get("input_tokens")) |
| 331 | + total_output += _safe_number(stats.get("output_tokens")) |
| 332 | + for usage in tool_usage_entries or []: |
| 333 | + total_input += _safe_number( |
| 334 | + usage.get("prompt_tokens") or usage.get("input_tokens") |
| 335 | + ) |
| 336 | + total_output += _safe_number( |
| 337 | + usage.get("completion_tokens") or usage.get("output_tokens") |
| 338 | + ) |
| 339 | + return total_input, total_output |
| 340 | + |
| 341 | + |
| 342 | +def _query_ttft_seconds(conversation: List[Dict[str, Any]]) -> Optional[float]: |
| 343 | + """Turn's ttft = the FIRST step's own time_to_first_token; a later step's |
| 344 | + value would drop all earlier tool-decision latency. None when step 1 has |
| 345 | + no positive value — never a fabricated 0.0.""" |
| 346 | + for entry in conversation: |
| 347 | + if entry.get("role") == "system" and isinstance(entry.get("content"), dict): |
| 348 | + content = entry["content"] |
| 349 | + if content.get("type") == "stats" and "performance_stats" in content: |
| 350 | + if content.get("step") != 1: |
| 351 | + # Step 1's own poll failed/was skipped — never misattribute |
| 352 | + # a later step's latency as the turn's ttft. |
| 353 | + return None |
| 354 | + stats = content["performance_stats"] |
| 355 | + ttft = ( |
| 356 | + stats.get("time_to_first_token") |
| 357 | + if isinstance(stats, dict) |
| 358 | + else None |
| 359 | + ) |
| 360 | + if ( |
| 361 | + isinstance(ttft, (int, float)) |
| 362 | + and not isinstance(ttft, bool) |
| 363 | + and math.isfinite(ttft) |
| 364 | + and ttft > 0 |
| 365 | + ): |
| 366 | + return float(ttft) |
| 367 | + return None |
| 368 | + return None |
| 369 | + |
| 370 | + |
| 371 | +# Only these field names are ever accepted from a tool's self-reported |
| 372 | +# ``usage`` dict — deliberately narrower than "any dict under a `usage` key", |
| 373 | +# so a tool with an unrelated `usage` value (rate-limit/quota/disk usage, not |
| 374 | +# LLM tokens) is never misread as token accounting (#2899). |
| 375 | +_TOOL_USAGE_TOKEN_FIELDS = ( |
| 376 | + "prompt_tokens", |
| 377 | + "completion_tokens", |
| 378 | + "input_tokens", |
| 379 | + "output_tokens", |
| 380 | +) |
| 381 | + |
| 382 | + |
| 383 | +def _extract_tool_usage(tool_result: Any) -> Optional[Dict[str, Any]]: |
| 384 | + """Pull a tool-reported usage dict off a tool's own return payload, if |
| 385 | + present and shaped like real token accounting. Some tools make their own |
| 386 | + internal LLM calls outside the normal per-step chat-completion accounting |
| 387 | + (e.g. a triage tool that classifies many items with its own client calls) |
| 388 | + and report the aggregate on their own return value instead of through the |
| 389 | + per-step stats path. Never raises — a malformed payload (bad JSON, wrong |
| 390 | + shape, non-numeric fields) yields ``None``, the same as "no usage to |
| 391 | + report".""" |
| 392 | + try: |
| 393 | + payload = tool_result |
| 394 | + if isinstance(payload, str): |
| 395 | + payload = json.loads(payload) |
| 396 | + if not isinstance(payload, dict): |
| 397 | + return None |
| 398 | + usage = payload.get("usage") |
| 399 | + if not isinstance(usage, dict): |
| 400 | + return None |
| 401 | + has_real_token_field = any( |
| 402 | + isinstance(usage.get(f), (int, float)) |
| 403 | + and not isinstance(usage.get(f), bool) |
| 404 | + for f in _TOOL_USAGE_TOKEN_FIELDS |
| 405 | + ) |
| 406 | + return usage if has_real_token_field else None |
| 407 | + except (ValueError, TypeError): |
| 408 | + return None |
| 409 | + |
| 410 | + |
303 | 411 | # Suffix appended to the last tool-result message when ``single_tool_per_turn`` |
304 | 412 | # agents have completed their one tool call. The model sees this and emits a |
305 | 413 | # short final reply instead of calling another tool. Greppable for fixtures |
@@ -533,6 +641,10 @@ def __init__( |
533 | 641 | # post-registration skill-set load both see the explicit request. |
534 | 642 | self._requested_skill_set = skill_set |
535 | 643 | self.error_history = [] # Store error history for learning |
| 644 | + # Safe default so _execute_tool -> _fold_tool_usage never AttributeErrors |
| 645 | + # if called outside the normal process_query loop (e.g. directly in a |
| 646 | + # test); _process_query_impl resets this per-turn (#2899). |
| 647 | + self._tool_reported_usage: List[Dict[str, Any]] = [] |
536 | 648 | self.conversation_history = ( |
537 | 649 | [] |
538 | 650 | ) # Store conversation history for session persistence |
@@ -2653,6 +2765,22 @@ def _tool_requires_confirmation(self, tool_name: str) -> bool: |
2653 | 2765 | return bool(flag) |
2654 | 2766 | return tool_name.startswith("mcp_") |
2655 | 2767 |
|
| 2768 | + def _fold_tool_usage(self, tool_name: str, tool_result: Any) -> None: |
| 2769 | + """Record a tool's self-reported LLM usage (see ``_extract_tool_usage``) |
| 2770 | + against this turn's running total. Called from the single success path |
| 2771 | + inside ``_execute_tool`` so every caller is covered uniformly. Never |
| 2772 | + raises — extraction failures are already swallowed by |
| 2773 | + ``_extract_tool_usage``; this method only appends. |
| 2774 | +
|
| 2775 | + A tool whose internal LLM calls already route through ``self.chat`` |
| 2776 | + would double-count here — a constraint on future tools, not a live one. |
| 2777 | + """ |
| 2778 | + usage = _extract_tool_usage(tool_result) |
| 2779 | + if usage is None: |
| 2780 | + return |
| 2781 | + logger.debug("Tool '%s' reported its own LLM usage: %s", tool_name, usage) |
| 2782 | + self._tool_reported_usage.append(usage) |
| 2783 | + |
2656 | 2784 | def _execute_tool(self, tool_name: str, tool_args: Dict[str, Any]) -> Any: |
2657 | 2785 | """ |
2658 | 2786 | Execute a tool by name with the provided arguments. |
@@ -2796,6 +2924,7 @@ def _execute_tool(self, tool_name: str, tool_args: Dict[str, Any]) -> Any: |
2796 | 2924 | try: |
2797 | 2925 | result = self._call_tool_bounded(tool, tool_args, tool_name) |
2798 | 2926 | logger.debug(f"Tool execution result: {result}") |
| 2927 | + self._fold_tool_usage(tool_name, result) |
2799 | 2928 | return result |
2800 | 2929 | except ToolExecutionTimeout as e: |
2801 | 2930 | # Bounded-execution guard fired: the tool body blocked past its |
@@ -3666,6 +3795,10 @@ def _process_query_impl( |
3666 | 3795 | self.current_step = 0 |
3667 | 3796 | self.total_plan_steps = 0 |
3668 | 3797 | self.plan_iterations = 0 # Reset plan iteration counter |
| 3798 | + # Tool-reported LLM usage this turn (see _fold_tool_usage / #2899) — |
| 3799 | + # reset per-turn since an Agent instance persists across queries in |
| 3800 | + # an interactive session. |
| 3801 | + self._tool_reported_usage: List[Dict[str, Any]] = [] |
3669 | 3802 |
|
3670 | 3803 | # Add user query to the conversation history |
3671 | 3804 | conversation.append({"role": "user", "content": user_input}) |
@@ -5538,7 +5671,20 @@ def _process_query_impl( |
5538 | 5671 |
|
5539 | 5672 | final_answer = self.finalize_answer(answer_candidate, conversation) |
5540 | 5673 | self.execution_state = self.STATE_COMPLETION |
5541 | | - self.console.print_final_answer(final_answer, streaming=self.streaming) |
| 5674 | + # Compute the real token total BEFORE printing the answer so it |
| 5675 | + # can ride the same event, instead of the post-loop aggregation |
| 5676 | + # below which runs after print_final_answer already fired |
| 5677 | + # (#2899). Output tokens only — "tokens actually generated", |
| 5678 | + # matching the tok/s calc downstream which is also output-only. |
| 5679 | + _pre_input_tokens, pre_output_tokens = _sum_conversation_tokens( |
| 5680 | + conversation, self._tool_reported_usage |
| 5681 | + ) |
| 5682 | + self.console.print_final_answer( |
| 5683 | + final_answer, |
| 5684 | + streaming=self.streaming, |
| 5685 | + total_tokens=pre_output_tokens, |
| 5686 | + ttft_seconds=_query_ttft_seconds(conversation), |
| 5687 | + ) |
5542 | 5688 | break |
5543 | 5689 |
|
5544 | 5690 | # Check if we're at the limit and ask user if they want to continue |
@@ -5603,18 +5749,14 @@ def _process_query_impl( |
5603 | 5749 | # Calculate total duration |
5604 | 5750 | total_duration = time.time() - start_time |
5605 | 5751 |
|
5606 | | - # Aggregate token counts from conversation stats |
5607 | | - total_input_tokens = 0 |
5608 | | - total_output_tokens = 0 |
5609 | | - for entry in conversation: |
5610 | | - if entry.get("role") == "system" and isinstance(entry.get("content"), dict): |
5611 | | - content = entry["content"] |
5612 | | - if content.get("type") == "stats" and "performance_stats" in content: |
5613 | | - stats = content["performance_stats"] |
5614 | | - if stats.get("input_tokens") is not None: |
5615 | | - total_input_tokens += stats["input_tokens"] |
5616 | | - if stats.get("output_tokens") is not None: |
5617 | | - total_output_tokens += stats["output_tokens"] |
| 5752 | + # Aggregate token counts from conversation stats, plus any usage a |
| 5753 | + # tool self-reported (e.g. a triage tool's internal per-message LLM |
| 5754 | + # calls, #2899) — same helper the pre-answer computation above uses, |
| 5755 | + # reading identical inputs since nothing mutates conversation or |
| 5756 | + # self._tool_reported_usage between the two calls. |
| 5757 | + total_input_tokens, total_output_tokens = _sum_conversation_tokens( |
| 5758 | + conversation, self._tool_reported_usage |
| 5759 | + ) |
5618 | 5760 |
|
5619 | 5761 | # Return the result |
5620 | 5762 | has_errors = len(self.error_history) > 0 |
|
0 commit comments