diff --git a/gcode/agent.py b/gcode/agent.py index 606bd0f..05e260a 100644 --- a/gcode/agent.py +++ b/gcode/agent.py @@ -12,6 +12,55 @@ MAX_HISTORY = 30 +def _usage_value(usage: dict, *names): + """Return the first non-None value among the given keys, so zero counts win.""" + for name in names: + if usage.get(name) is not None: + return usage[name] + return None + + +def _print_usage(response, ui) -> None: + """Print a compact footer with token usage when the provider supplies it. + + Prefers LangChain's ``usage_metadata`` and falls back to + ``response_metadata['token_usage' | 'usage']``. Prints nothing when no + usage is available and never raises. + """ + try: + usage = getattr(response, "usage_metadata", None) + if not isinstance(usage, dict) or all(v is None for v in usage.values()): + meta = getattr(response, "response_metadata", {}) or {} + cand: dict = {} + if isinstance(meta, dict): + cand = meta.get("token_usage") or meta.get("usage") or {} + usage = cand if isinstance(cand, dict) else None + if not usage or all(v is None for v in usage.values()): + return + inp = _usage_value(usage, "input_tokens", "prompt_tokens", "promptTokens") + out = _usage_value(usage, "output_tokens", "completion_tokens", "completionTokens") + total = _usage_value(usage, "total_tokens", "totalTokens") + parts = [] + if inp is not None and out is not None: + parts.append(f"{inp} in / {out} out") + elif inp is not None: + parts.append(f"{inp} in") + elif out is not None: + parts.append(f"{out} out") + if total is not None: + parts.append(f"total {total}") + cost = _usage_value(usage, "cost", "total_cost") + if cost is not None: + try: + parts.append(f"~${float(cost):.4f}") + except (TypeError, ValueError): + parts.append(f"cost {cost}") + if parts: + ui.info(f"[dim]Usage: {' · '.join(parts)}[/dim]") + except Exception: + return + + def build_model(model_id: str, api_key: str): """Build a ChatOpenAI model bound to all GCode tools. @@ -82,10 +131,14 @@ def _stream(messages: list, model, ui) -> AIMessage: if interrupted: ui.info("(streaming stopped by user)") # Store the canonical AIMessage (not the chunk) for clean history + reloads. + # usage_metadata/response_metadata carry the provider's token counts, which + # _print_usage reads; dropping them would make the usage footer always empty. return AIMessage( content=accumulated.content, tool_calls=[] if interrupted else accumulated.tool_calls, additional_kwargs=accumulated.additional_kwargs, + response_metadata=accumulated.response_metadata, + usage_metadata=accumulated.usage_metadata, id=accumulated.id, ) @@ -128,6 +181,7 @@ def run_turn(user_input: str, messages: list, model, ui) -> None: return messages.append(response) + _print_usage(response, ui) errored = False while getattr(response, "tool_calls", None): @@ -149,6 +203,7 @@ def run_turn(user_input: str, messages: list, model, ui) -> None: break messages.append(response) + _print_usage(response, ui) if errored: return diff --git a/tests/test_agent.py b/tests/test_agent.py index ace84fd..a14f71a 100644 --- a/tests/test_agent.py +++ b/tests/test_agent.py @@ -1,8 +1,8 @@ -"""Unit tests for the agent loop, focused on Ctrl+C interrupt handling.""" +"""Unit tests for the agent loop: Ctrl+C interrupts and the usage footer.""" from unittest.mock import Mock, patch -from gcode.agent import _run_tool, _stream +from gcode.agent import _print_usage, _run_tool, _stream from langchain_core.messages import AIMessage, AIMessageChunk @@ -71,3 +71,87 @@ def test_run_tool_returns_cancelled_on_keyboard_interrupt(): assert result == "Command execution cancelled by user." ui.tool_result.assert_called_once_with("failing_tool", "Command execution cancelled by user.") + + +class _UsageModel: + """A model whose final stream chunk carries the provider's usage metadata.""" + + def stream(self, messages): + yield AIMessageChunk(content="hi") + yield AIMessageChunk( + content="", + usage_metadata={"input_tokens": 10, "output_tokens": 5, "total_tokens": 15}, + ) + + +def test_stream_preserves_usage_metadata(): + ui = _FakeUI() + msg = _stream([], _UsageModel(), ui) + + assert msg.usage_metadata == {"input_tokens": 10, "output_tokens": 5, "total_tokens": 15} + + +def test_print_usage_from_usage_metadata(): + ui = _FakeUI() + msg = AIMessage( + content="", usage_metadata={"input_tokens": 842, "output_tokens": 128, "total_tokens": 970} + ) + + _print_usage(msg, ui) + + assert ("info", "[dim]Usage: 842 in / 128 out · total 970[/dim]") in ui.calls + + +def test_print_usage_falls_back_to_response_metadata(): + ui = _FakeUI() + msg = AIMessage( + content="", + response_metadata={ + "token_usage": {"prompt_tokens": 3, "completion_tokens": 4, "total_tokens": 7} + }, + ) + + _print_usage(msg, ui) + + assert ("info", "[dim]Usage: 3 in / 4 out · total 7[/dim]") in ui.calls + + +def test_print_usage_keeps_zero_counts(): + ui = _FakeUI() + msg = AIMessage( + content="", + response_metadata={ + "token_usage": {"prompt_tokens": 0, "completion_tokens": 7, "total_tokens": 7} + }, + ) + + _print_usage(msg, ui) + + assert ("info", "[dim]Usage: 0 in / 7 out · total 7[/dim]") in ui.calls + + +def test_print_usage_shows_cost_when_present(): + ui = _FakeUI() + msg = AIMessage( + content="", + response_metadata={ + "token_usage": { + "prompt_tokens": 3, + "completion_tokens": 4, + "total_tokens": 7, + "total_cost": 0.0123456, + } + }, + ) + + _print_usage(msg, ui) + + assert ("info", "[dim]Usage: 3 in / 4 out · total 7 · ~$0.0123[/dim]") in ui.calls + + +def test_print_usage_silent_without_usage(): + ui = _FakeUI() + + _print_usage(AIMessage(content="no usage here"), ui) + + assert not any(call[0] == "info" for call in ui.calls if isinstance(call, tuple))