« prev ^ index » next coverage.py v7.15.4, created at 2026-08-25 07:19 +0800
1"""Shared token accounting helpers for agent strategies.
3Every strategy funnels LLM completions through these helpers so token
4usage is counted consistently. Semantics mirror the original
5``function_calling`` implementation: prefer the prompt/completion split
6when either side is reported, fall back to ``usage.total_tokens``, and
7return ``0`` when usage is missing entirely.
8"""
10from __future__ import annotations
12from lexigram.contracts.ai.llm import Completion, CompletionProtocol
14CompletionInput = Completion | CompletionProtocol
17def token_split(completion: CompletionInput) -> tuple[int, int]:
18 """Extract the prompt/completion token split from a completion.
20 Args:
21 completion: LLM completion result.
23 Returns:
24 Tuple of ``(prompt_tokens, completion_tokens)``. Both are
25 ``0`` when usage is missing or not reported.
26 """
27 usage = getattr(completion, "usage", None)
28 if not usage:
29 return 0, 0
30 if isinstance(usage, dict):
31 return (
32 int(usage.get("prompt_tokens", 0) or 0),
33 int(usage.get("completion_tokens", 0) or 0),
34 )
35 return (
36 int(getattr(usage, "prompt_tokens", 0) or 0),
37 int(getattr(usage, "completion_tokens", 0) or 0),
38 )
41def count_tokens(completion: CompletionInput) -> int:
42 """Count total tokens for a completion.
44 Prefers the prompt/completion split when either side is reported;
45 otherwise falls back to ``usage.total_tokens``. Returns ``0`` when
46 usage is missing entirely.
48 Args:
49 completion: LLM completion result.
51 Returns:
52 Total token count consumed by the completion.
53 """
54 prompt, completion_tokens = token_split(completion)
55 if prompt or completion_tokens:
56 return prompt + completion_tokens
57 usage = getattr(completion, "usage", None)
58 if isinstance(usage, dict):
59 return int(usage.get("total_tokens", 0) or 0)
60 return int(getattr(usage, "total_tokens", 0) or 0)
63class TokenAccumulator:
64 """Mutable token totals accumulated across multiple LLM calls."""
66 def __init__(self) -> None:
67 self.prompt_tokens = 0
68 self.completion_tokens = 0
69 self.total_tokens = 0
71 def add(self, completion: CompletionInput) -> None:
72 """Accumulate usage from one completion.
74 Args:
75 completion: LLM completion result.
76 """
77 prompt, completion_tokens = token_split(completion)
78 self.prompt_tokens += prompt
79 self.completion_tokens += completion_tokens
80 self.total_tokens += count_tokens(completion)
83__all__ = ["TokenAccumulator", "count_tokens", "token_split"]