Wraps an :class:`LLMBackend` and sums token usage across calls. ``Planner`` / ``Executor`` call ``generate()`` and only need a string; this wrapper transparently pulls usage from the inner backend's ``generate_with_usage()`` and accumulates it in-place. Totals are reset at the star
| 26 | |
| 27 | |
| 28 | class _TokenMeter(LLMBackend): |
| 29 | """Wraps an :class:`LLMBackend` and sums token usage across calls. |
| 30 | |
| 31 | ``Planner`` / ``Executor`` call ``generate()`` and only need a string; |
| 32 | this wrapper transparently pulls usage from the inner backend's |
| 33 | ``generate_with_usage()`` and accumulates it in-place. Totals are |
| 34 | reset at the start of each :meth:`PlanExecuteRunner.run` call so |
| 35 | per-run span attributes reflect that run alone. |
| 36 | """ |
| 37 | |
| 38 | def __init__(self, inner: LLMBackend) -> None: |
| 39 | self._inner = inner |
| 40 | self.input_tokens = 0 |
| 41 | self.output_tokens = 0 |
| 42 | |
| 43 | def reset(self) -> None: |
| 44 | self.input_tokens = 0 |
| 45 | self.output_tokens = 0 |
| 46 | |
| 47 | def generate(self, prompt: str, temperature: float = 0.0) -> str: |
| 48 | result = self._inner.generate_with_usage(prompt, temperature) |
| 49 | self.input_tokens += result.input_tokens |
| 50 | self.output_tokens += result.output_tokens |
| 51 | return result.text |
| 52 | |
| 53 | def generate_with_usage(self, prompt: str, temperature: float = 0.0) -> LLMResult: |
| 54 | result = self._inner.generate_with_usage(prompt, temperature) |
| 55 | self.input_tokens += result.input_tokens |
| 56 | self.output_tokens += result.output_tokens |
| 57 | return result |
| 58 | |
| 59 | @property |
| 60 | def model_id(self) -> str: |
| 61 | return self._inner.model_id |
| 62 | |
| 63 | |
| 64 | _log = logging.getLogger(__name__) |