From 5ce9d6625853a8234c6633698c2e66139e08ecd9 Mon Sep 17 00:00:00 2001 From: jay Date: Sun, 6 Sep 2026 06:44:33 -0700 Subject: [PATCH] add token usage --- AGENTS.md | 11 +- README.md | 1 + oli_bot/agent.py | 23 ++ oli_bot/backends/base.py | 13 +- oli_bot/backends/huggingface.py | 48 +++- oli_bot/backends/ollama.py | 49 +++- oli_bot/backends/openai.py | 73 +++++- oli_bot/backends/transformers.py | 16 +- oli_bot/chat.py | 62 +++-- oli_bot/models.py | 30 +++ oli_bot/screens/mcp_setup.py | 4 +- oli_bot/screens/taglines.py | 2 +- oli_bot/sessions.py | 9 + oli_bot/tools/manager.py | 4 +- tests/test_agent_scaffolding.py | 18 +- tests/test_backend_stream_flush_extra.py | 2 +- tests/test_backend_tool_call_flush.py | 3 +- tests/test_token_usage.py | 285 +++++++++++++++++++++++ 18 files changed, 610 insertions(+), 43 deletions(-) create mode 100644 tests/test_token_usage.py diff --git a/AGENTS.md b/AGENTS.md index aed6ce5..d80dd35 100644 --- a/AGENTS.md +++ b/AGENTS.md @@ -11,12 +11,12 @@ This file is the top-level orientation map. Depth lives under [docs/](docs/). | Path | Purpose | | --- | --- | -| [oli_bot/chat.py](oli_bot/chat.py) | Textual TUI app (`OliBot`). Owns command handling, session UI, the `#command-suggestions` autocomplete `ListView`, and (when pooling is enabled) the "Active Sub-Agents" `Tree`. Slash-command names come from the module-level `COMMANDS` tuple. | +| [oli_bot/chat.py](oli_bot/chat.py) | Textual TUI app (`OliBot`). Owns command handling, session UI, the `#command-suggestions` autocomplete `ListView`, and (when pooling is enabled) the "Active Sub-Agents" `Tree`. Tracks cumulative session token usage (persisted per session) and renders it in the `#status-bar`. Slash-command names come from the module-level `COMMANDS` tuple. | | [oli_bot/api_server.py](oli_bot/api_server.py) | FastAPI app exposing the harness over an OpenAI-compatible REST API (`GET /v1/models`, `POST /v1/chat/completions` streaming + non-streaming, `GET /health`). Stateless from the caller's POV; a single process-private `Agent` is shared across requests and serialised with a `threading.RLock`. Auto-approves permissions (no human), but offline/dry-run still apply. | -| [oli_bot/agent.py](oli_bot/agent.py) | `Agent` — mode + system prompt owner; orchestrates the tool-calling loop and streams typed events (`TextChunk`, `ThinkingChunk`, `ToolCallChunk`, `ToolCallExecuting`, `ToolCallResult`, `StreamChunk`, `Error`, `Done`). Also hosts `sanitize_tool_history`, `stream_sub_agent_run`, and `AgentPool` (built from [oli_bot/agents.yaml](oli_bot/agents.yaml) when `--use-pool` is set). | -| [oli_bot/backends/](oli_bot/backends/) | Backend package — `ModelBackend` ABC, `OllamaBackend`, `OpenAIBackend`, `HuggingFaceBackend`, `TransformersBackend`, and the `create_model_backend()` factory. Also hosts the shared `_StreamingThinkParser` and per-backend message formatting (Ollama native `images`, OpenAI `image_url` or Bedrock-native blocks via `openai_vision_style`, textual placeholder for text-only backends). See [docs/BACKENDS.md](docs/BACKENDS.md). | +| [oli_bot/agent.py](oli_bot/agent.py) | `Agent` — mode + system prompt owner; orchestrates the tool-calling loop and streams typed events (`TextChunk`, `ThinkingChunk`, `ToolCallChunk`, `ToolCallExecuting`, `ToolCallResult`, `StreamChunk`, `UsageEvent`, `Error`, `Done`). Aggregates per-call `UsageChunk`s from each backend round into a single per-run `UsageEvent`. Also hosts `sanitize_tool_history`, `_merge_usage`, `stream_sub_agent_run`, and `AgentPool` (built from [oli_bot/agents.yaml](oli_bot/agents.yaml) when `--use-pool` is set). | +| [oli_bot/backends/](oli_bot/backends/) | Backend package — `ModelBackend` ABC, `OllamaBackend`, `OpenAIBackend`, `HuggingFaceBackend`, `TransformersBackend`, and the `create_model_backend()` factory. Also hosts the shared `_StreamingThinkParser` and per-backend message formatting (Ollama native `images`, OpenAI `image_url` or Bedrock-native blocks via `openai_vision_style`, textual placeholder for text-only backends). Every backend surfaces a trailing `UsageChunk`: exact counts from provider usage where available (OpenAI `usage`/`stream_options`, Ollama `prompt_eval_count`/`eval_count`, HF `usage`), else a `~chars/4` estimate via `estimate_tokens`. See [docs/BACKENDS.md](docs/BACKENDS.md). | | [oli_bot/screens/](oli_bot/screens/) | All `ModalScreen` subclasses: `PermissionScreen`, `ConfirmScreen`, `ModelPickerScreen`, `ServerListScreen`, `MCPSetupScreen`, `SessionListScreen`, `WorkspaceListScreen`, `SubAgentViewScreen`, `ConfigScreen`, `InputPromptScreen`, plus `taglines.py` / `todo_widget.py`. | -| [oli_bot/models.py](oli_bot/models.py) | Shared dataclasses: `Message`, `ToolCall`, `ModelResponse`, `HostConfig`, `MCPServerConfig`, `ProfileData`, `SubAgentRun`, `ImageAttachment`, `TodoItem` / `TodoListState`, plus `AgentEvent` variants and the `AgentRole` enum. `Message.images` is in-memory only (dropped on session save). | +| [oli_bot/models.py](oli_bot/models.py) | Shared dataclasses: `Message`, `ToolCall`, `ModelResponse`, `HostConfig`, `MCPServerConfig`, `ProfileData`, `SubAgentRun`, `ImageAttachment`, `TodoItem` / `TodoListState`, plus `AgentEvent` variants and the `AgentRole` enum. Token accounting lives here too: `Usage` (prompt/completion/`estimated` flag), the per-call `UsageChunk` stream event, and the per-run `UsageEvent`. `Message.images` is in-memory only (dropped on session save); `ModelResponse.usage` is optionally set by backends. | | [oli_bot/config.py](oli_bot/config.py) | `AppConfig` — `pydantic_settings.BaseSettings`. Env vars prefixed `OLI_`, plus `.env` support and `OLI_TRUNCATION_SMALL` / `_LARGE` aliases via `AliasChoices`. Module-level `configs = AppConfig()` singleton. See [docs/CONFIGURE.md](docs/CONFIGURE.md). | | [oli_bot/settings.py](oli_bot/settings.py) | `SettingsManager` — load/save/merge `~/.config/oli/settings.json`; precedence `settings.json` > `OLI_*` env > SDK-standard env (`OPENAI_API_KEY`, `OPENAI_BASE_URL`, `HUGGINGFACE_API_KEY`, `HF_TOKEN`) > declared defaults. Empty API-key strings in JSON fall through to env. | | [oli_bot/profiles/](oli_bot/profiles/) | `ProfileManifest` / `PermissionsManifest` (Pydantic) in `schema.py`; `ProfilePermissionEnforcer` (layered allow/deny glob patterns, base-profile inheritance, deny-overrides-allow) in `permissions.py`. Built-in profiles ship as sibling directories. | @@ -69,6 +69,7 @@ See [docs/AGENT-POOLING.md](docs/AGENT-POOLING.md) for the `agents.yaml` schema, - **Session lifecycle** — new UUID by default; `-s`/`--load-session ` loads a specific one, `--resume-last` loads the newest for the active server (mutually exclusive). Auto-saved after every response; `/sessions` opens `SessionListScreen` for browse/switch/rename/delete. On exit `main()` prints the resume hint. - **Tool loop** — `Agent._tool_loop` calls `backend.stream_generate()`, streaming text and thinking to the TUI while accumulating tool calls; after each round it calls `MCPClientManager.drain_builtin_attachments()` so `view_image` results ride into the next turn as a synthetic user-role `Message(images=...)` (kept **after** all `role=tool` messages so the assistant→tool block stays contiguous for `sanitize_tool_history`). If iterations are exhausted, `_final_stream` does one text-only pass to guarantee a reply. +- **Token accounting** — each backend yields a trailing `UsageChunk` (exact where the provider reports counts, else `~chars/4` via `estimate_tokens`). The agent sums per-call chunks across a run and emits a single `UsageEvent` before `Done`; the TUI adds it to the session total in the `#status-bar` (`~` prefix when any value was estimated). The running total persists in session JSON (`total_tokens` / `total_tokens_estimated`) and is reset on `/clear`, `/home`, and new-session flows. - **Error paths** — a failed round yields `Error` for the UI then `Done(full_text="")`; the empty `full_text` deliberately trips the TUI's "skip empty assistant append" branch so error text never enters `self.messages` (prevents the Bedrock/Anthropic 400 poisoning loop). All `stream_generate()` implementations re-raise instead of swallowing. - **Permission gating** — `BuiltinToolManager.call_tool()` runs the profile enforcer, then `Session.needs_permission()`, then (if needed) `confirm_callback(description)`; the TUI shows `PermissionScreen` and the callback returns `"once"`, `"session"`, or `"deny"`. Session grants persist per scope for the process lifetime. `glob`/`grep` targeting patterns like `.env*`, `*.pem`, `*secret*` trigger the `workspace_sensitive` scope even inside the workspace. - **Built-in tool naming** — registered as `builtin__` and dispatched by `MCPClientManager.call_tool`. @@ -147,7 +148,7 @@ pip install -e '.[dev]' pytest ``` -Tests under `tests/` cover: `AgentPool` scaffolding (lookup, `${VAR}` expansion, `root-agent` exclusion, per-agent backend overrides, concurrent dispatch), config env-var precedence, session round-trip + `save_session` ID propagation, permission matrix (workspace scoping, sensitive files, `glob`/`grep` sensitive-pattern gating), truncation boundary preservation, security regressions (`git` restricted to read-only subcommands in `run_command`, `find -exec/-delete` blocked, SSRF for loopback/private/link-local/non-`http(s)`), `OpenAIBackend.stream_generate` tool-call flushing on all finish reasons, and the `api_server.py` OpenAI-compatible endpoints. `tests/conftest.py` scrubs `OLI_*` env vars so runs are hermetic. +Tests under `tests/` cover: `AgentPool` scaffolding (lookup, `${VAR}` expansion, `root-agent` exclusion, per-agent backend overrides, concurrent dispatch), config env-var precedence, session round-trip + `save_session` ID propagation, permission matrix (workspace scoping, sensitive files, `glob`/`grep` sensitive-pattern gating), truncation boundary preservation, security regressions (`git` restricted to read-only subcommands in `run_command`, `find -exec/-delete` blocked, SSRF for loopback/private/link-local/non-`http(s)`), `OpenAIBackend.stream_generate` tool-call flushing on all finish reasons, token-usage accounting (backend `UsageChunk` exact/estimate paths, agent `UsageEvent` aggregation, session `total_tokens` round-trip), and the `api_server.py` OpenAI-compatible endpoints. `tests/conftest.py` scrubs `OLI_*` env vars so runs are hermetic. ## Code style diff --git a/README.md b/README.md index d9b7642..cc9cca0 100644 --- a/README.md +++ b/README.md @@ -44,6 +44,7 @@ Concretely, that means: - **Agent pooling (optional)** — with `--use-pool`, the root agent can fan tasks out concurrently to vendor-agnostic sub-agents defined in [agents.yaml](agents.yaml) via a `dispatch` tool. Each pool entry binds a model _and_ a backend, so dispatch decisions are also compute-location decisions — a frontier model can plan while sensitive work stays on a local model, or a local root can fan out to faster remote SLMs for latency-sensitive tool calls. - **Permission system** — write operations and sensitive reads require user approval. Session grants, workspace scoping, and profile-level allow/deny lists. - **Streaming Markdown** responses in a Textual TUI, with in-app model/server/session/profile management. +- **Session token counter** — the TUI status bar tracks cumulative session token usage (exact where the provider reports counts, `~`-prefixed estimates otherwise), persisted per session and summed from OpenAI/`stream_options`, Ollama eval counts, HuggingFace usage, and Transformers runs. Reset on `/clear`, `/home`, and new-session flows. - **OpenAI-compatible API server** — run the same agent harness behind `/v1/models` and `/v1/chat/completions` (streaming + non-streaming) so any workflow that speaks the OpenAI wire protocol (the `openai` Python SDK, curl, or plain REST) can drive the agent. ## Roadmap / areas of active exploration diff --git a/oli_bot/agent.py b/oli_bot/agent.py index 95d2804..7260346 100644 --- a/oli_bot/agent.py +++ b/oli_bot/agent.py @@ -30,6 +30,9 @@ StreamChunk, SubAgentRun, ThinkingChunk, + Usage, + UsageChunk, + UsageEvent, ) logger = logging.getLogger(__name__) @@ -41,6 +44,7 @@ | AssistantResponse | StreamChunk | ThinkingChunk + | UsageEvent | Error | Done ) @@ -113,6 +117,15 @@ def sanitize_tool_history(messages: List[Message]) -> List[Message]: return kept +def _merge_usage(acc: Usage, u: Usage) -> Usage: + """Sum per-call usage into a cumulative run total (flag estimated if any).""" + return Usage( + prompt_tokens=acc.prompt_tokens + u.prompt_tokens, + completion_tokens=acc.completion_tokens + u.completion_tokens, + estimated=acc.estimated or u.estimated, + ) + + class Agent: def __init__( self, @@ -287,6 +300,7 @@ async def _tool_loop( tools: Optional[List[Dict[str, Any]]], confirm_callback: Optional[Callable[[str], Any]], ) -> AsyncIterator[AgentEvent]: + run_usage = Usage() for _ in range(self.config.max_tool_iterations): stream = self.backend.stream_generate( messages, tools=tools if tools else [] @@ -310,6 +324,8 @@ async def _tool_loop( yield StreamChunk(event.text) elif isinstance(event, ThinkingChunk): yield event + elif isinstance(event, UsageChunk): + run_usage = _merge_usage(run_usage, event.usage) elif isinstance(event, ToolCallChunk): tool_calls = event.tool_calls for tc in tool_calls: @@ -344,6 +360,8 @@ async def _tool_loop( return if not tool_calls: + if run_usage.total_tokens: + yield UsageEvent(usage=run_usage) yield Done(full_text=full_response) return @@ -410,6 +428,7 @@ async def _final_stream( ) -> AsyncIterator[AgentEvent]: full_response = "" errored = False + usage = Usage() try: # Forward tools so backends that require a matching tool schema # for any toolUse/toolResult blocks in history (e.g. Bedrock via @@ -428,6 +447,8 @@ async def _final_stream( yield StreamChunk(event.text) elif isinstance(event, ThinkingChunk): yield event + elif isinstance(event, UsageChunk): + usage = _merge_usage(usage, event.usage) except asyncio.TimeoutError: err_text = ( f"Response timed out (no data for " @@ -448,6 +469,8 @@ async def _final_stream( else: if not full_response: full_response = "The model completed all tool calls but did not produce a final summary." + if usage.total_tokens: + yield UsageEvent(usage=usage) yield Done(full_text=full_response) diff --git a/oli_bot/backends/base.py b/oli_bot/backends/base.py index 63137bd..3e83a92 100644 --- a/oli_bot/backends/base.py +++ b/oli_bot/backends/base.py @@ -11,6 +11,7 @@ TextChunk, ThinkingChunk, ToolCallChunk, + UsageChunk, ) logger = logging.getLogger(__name__) @@ -20,7 +21,16 @@ MAX_TOKENS = configs.max_tokens TEMPERATURE = configs.temperature -StreamEvent = TextChunk | ToolCallChunk | ThinkingChunk +StreamEvent = TextChunk | ToolCallChunk | ThinkingChunk | UsageChunk + + +def estimate_tokens(text: str) -> int: + """Rough token estimate for providers that do not report usage. + + Mirrors the heuristic used by the OpenAI-compatible API layer (~4 chars per + token), so numbers stay consistent everywhere estimation is needed. + """ + return max(1, (len(text) + 3) // 4) class ModelBackend(ABC): @@ -51,6 +61,7 @@ async def stream_generate( "MAX_TOKENS", "TEMPERATURE", "StreamEvent", + "estimate_tokens", "logger", "Any", "AsyncIterator", diff --git a/oli_bot/backends/huggingface.py b/oli_bot/backends/huggingface.py index 28b76f4..1d59521 100644 --- a/oli_bot/backends/huggingface.py +++ b/oli_bot/backends/huggingface.py @@ -13,9 +13,17 @@ ThinkingChunk, ToolCall, ToolCallChunk, + Usage, + UsageChunk, ) -from .base import MAX_TOKENS, TEMPERATURE, ModelBackend, StreamEvent +from .base import ( + MAX_TOKENS, + TEMPERATURE, + ModelBackend, + StreamEvent, + estimate_tokens, +) from .messages import _format_messages, _format_tools from .streaming import _StreamingThinkParser @@ -86,10 +94,29 @@ async def generate( ) for i, tc in enumerate(message.tool_calls) ] + usage = None + usage_field = getattr(response, "usage", None) + if ( + usage_field is not None + and getattr(usage_field, "prompt_tokens", None) is not None + ): + usage = Usage( + prompt_tokens=usage_field.prompt_tokens, + completion_tokens=getattr(usage_field, "completion_tokens", 0) or 0, + ) + else: + usage = Usage( + prompt_tokens=estimate_tokens( + json.dumps(_format_messages(messages, stringify_arguments=True)) + ), + completion_tokens=estimate_tokens(message.content or ""), + estimated=True, + ) return ModelResponse( content=message.content or "", tool_calls=tool_calls or None, finish_reason=response.choices[0].finish_reason or "stop", + usage=usage, ) except Exception as e: logger.exception( @@ -113,13 +140,17 @@ async def stream_generate( tool_calls_acc: Dict[int, dict] = {} flushed = False parser = _StreamingThinkParser() + streamed_text = "" + chunk_usage = None async for chunk in stream: + chunk_usage = getattr(chunk, "usage", None) or chunk_usage if chunk.choices and chunk.choices[0].delta: delta = chunk.choices[0].delta reasoning = getattr(delta, "reasoning_content", None) if reasoning: yield ThinkingChunk(reasoning) if delta.content: + streamed_text += delta.content for kind, text in parser.feed(delta.content): if text: yield ( @@ -192,6 +223,21 @@ async def stream_generate( for data in tool_calls_acc.values() ] yield ToolCallChunk(tool_calls) + + if chunk_usage is not None: + usage = Usage( + prompt_tokens=getattr(chunk_usage, "prompt_tokens", 0) or 0, + completion_tokens=getattr(chunk_usage, "completion_tokens", 0) or 0, + ) + else: + usage = Usage( + prompt_tokens=estimate_tokens( + json.dumps(_format_messages(messages, stringify_arguments=True)) + ), + completion_tokens=estimate_tokens(streamed_text), + estimated=True, + ) + yield UsageChunk(usage) except Exception as e: logger.exception( "Error in %s.stream_generate: %s", self.__class__.__name__, str(e) diff --git a/oli_bot/backends/ollama.py b/oli_bot/backends/ollama.py index f647580..8c30b3e 100644 --- a/oli_bot/backends/ollama.py +++ b/oli_bot/backends/ollama.py @@ -1,5 +1,6 @@ from __future__ import annotations +import json import logging from typing import Any, AsyncIterator, Dict, List, Optional @@ -12,9 +13,17 @@ ThinkingChunk, ToolCall, ToolCallChunk, + Usage, + UsageChunk, ) -from .base import MAX_TOKENS, TEMPERATURE, ModelBackend, StreamEvent +from .base import ( + MAX_TOKENS, + TEMPERATURE, + ModelBackend, + StreamEvent, + estimate_tokens, +) from .messages import _format_messages, _format_tools from .streaming import _StreamingThinkParser @@ -73,10 +82,29 @@ async def generate( for i, tc in enumerate(response.message.tool_calls) ] + usage = None + if ( + getattr(response, "prompt_eval_count", None) is not None + and getattr(response, "eval_count", None) is not None + ): + usage = Usage( + prompt_tokens=response.prompt_eval_count or 0, + completion_tokens=response.eval_count or 0, + ) + else: + usage = Usage( + prompt_tokens=estimate_tokens( + json.dumps(_format_messages(messages, image_style="ollama")) + ), + completion_tokens=estimate_tokens(response.message.content or ""), + estimated=True, + ) + return ModelResponse( content=response.message.content or "", tool_calls=tool_calls, finish_reason="stop", + usage=usage, ) except Exception as e: logger.exception( @@ -103,12 +131,20 @@ async def stream_generate( stream = await self.client.chat(**kwargs) yielded = False parser = _StreamingThinkParser() + streamed_text = "" + prompt_eval = None + eval_count = None async for chunk in stream: + if getattr(chunk, "prompt_eval_count", None) is not None: + prompt_eval = chunk.prompt_eval_count + if getattr(chunk, "eval_count", None) is not None: + eval_count = chunk.eval_count thinking_field = getattr(chunk.message, "thinking", None) if thinking_field: yield ThinkingChunk(thinking_field) yielded = True if chunk.message.content: + streamed_text += chunk.message.content for kind, text in parser.feed(chunk.message.content): if text: yield ( @@ -139,6 +175,17 @@ async def stream_generate( if text: yield ThinkingChunk(text) if kind == "thinking" else TextChunk(text) yielded = True + if prompt_eval is not None and eval_count is not None: + usage = Usage(prompt_tokens=prompt_eval, completion_tokens=eval_count) + else: + usage = Usage( + prompt_tokens=estimate_tokens( + json.dumps(_format_messages(messages, image_style="ollama")) + ), + completion_tokens=estimate_tokens(streamed_text), + estimated=True, + ) + yield UsageChunk(usage) if not yielded: logger.warning( "Ollama stream completed with no content or tool calls (model=%s)", diff --git a/oli_bot/backends/openai.py b/oli_bot/backends/openai.py index e498080..58ca658 100644 --- a/oli_bot/backends/openai.py +++ b/oli_bot/backends/openai.py @@ -14,9 +14,17 @@ ThinkingChunk, ToolCall, ToolCallChunk, + Usage, + UsageChunk, ) -from .base import MAX_TOKENS, TEMPERATURE, ModelBackend, StreamEvent +from .base import ( + MAX_TOKENS, + TEMPERATURE, + ModelBackend, + StreamEvent, + estimate_tokens, +) from .messages import _format_messages, _format_tools, _validate_message_content_blocks from .streaming import _StreamingThinkParser @@ -85,10 +93,25 @@ async def generate( parameters=json.loads(call.function.arguments or "{}"), ) ) + usage = None + if response.usage is not None: + usage = Usage( + prompt_tokens=response.usage.prompt_tokens or 0, + completion_tokens=response.usage.completion_tokens or 0, + ) + else: + usage = Usage( + prompt_tokens=estimate_tokens(json.dumps(formatted)), + completion_tokens=estimate_tokens( + response.choices[0].message.content or "" + ), + estimated=True, + ) return ModelResponse( content=response.choices[0].message.content, tool_calls=tool_calls, finish_reason="stop", + usage=usage, ) except Exception as e: logger.exception( @@ -115,23 +138,44 @@ async def stream_generate( ) _validate_message_content_blocks(formatted_messages) - response = await self.client.chat.completions.create( - model=self.model, - messages=formatted_messages, - tools=_format_tools(tools) if tools else [], - tool_choice="auto" if tools else "none", - stream=True, - ) + try: + response = await self.client.chat.completions.create( + model=self.model, + messages=formatted_messages, + tools=_format_tools(tools) if tools else [], + tool_choice="auto" if tools else "none", + stream=True, + stream_options={"include_usage": True}, + ) + except Exception as e: + # Some OpenAI-compatible providers (Bedrock proxies, LM Studio, + # older vLLM, ...) reject the stream_options param. Retry once + # without it; usage then falls back to estimation. + logger.debug( + "Provider rejected stream_options, retrying without usage: %s", e + ) + response = await self.client.chat.completions.create( + model=self.model, + messages=formatted_messages, + tools=_format_tools(tools) if tools else [], + tool_choice="auto" if tools else "none", + stream=True, + ) tool_calls_acc: Dict[int, dict] = {} flushed = False parser = _StreamingThinkParser() + streamed_text = "" + chunk_usage = None async for chunk in response: delta = chunk.choices[0].delta if chunk.choices else None + if getattr(chunk, "usage", None) is not None: + chunk_usage = chunk.usage if delta: reasoning = getattr(delta, "reasoning_content", None) if reasoning: yield ThinkingChunk(reasoning) if delta.content: + streamed_text += delta.content for kind, text in parser.feed(delta.content): if text: yield ( @@ -206,6 +250,19 @@ async def stream_generate( ] yield ToolCallChunk(tool_calls) + if chunk_usage is not None: + usage = Usage( + prompt_tokens=chunk_usage.prompt_tokens or 0, + completion_tokens=chunk_usage.completion_tokens or 0, + ) + else: + usage = Usage( + prompt_tokens=estimate_tokens(json.dumps(formatted_messages)), + completion_tokens=estimate_tokens(streamed_text), + estimated=True, + ) + yield UsageChunk(usage) + except Exception as e: logger.exception( "Error in %s.stream_generate: %s", self.__class__.__name__, str(e) diff --git a/oli_bot/backends/transformers.py b/oli_bot/backends/transformers.py index 690e923..158cec6 100644 --- a/oli_bot/backends/transformers.py +++ b/oli_bot/backends/transformers.py @@ -13,9 +13,11 @@ ThinkingChunk, ToolCall, ToolCallChunk, + Usage, + UsageChunk, ) -from .base import MAX_TOKENS, TEMPERATURE, ModelBackend, StreamEvent +from .base import MAX_TOKENS, TEMPERATURE, ModelBackend, StreamEvent, estimate_tokens from .messages import _append_image_placeholder_text from .streaming import _StreamingThinkParser @@ -262,6 +264,10 @@ async def generate( content=clean_text, tool_calls=tool_calls or None, finish_reason="stop", + usage=Usage( + prompt_tokens=int(encoded["input_ids"].shape[1]), + completion_tokens=int(new_tokens.shape[0]), + ), ) except Exception as e: logger.exception( @@ -334,6 +340,14 @@ async def stream_generate( if tool_calls: yield ToolCallChunk(tool_calls) + yield UsageChunk( + Usage( + prompt_tokens=int(encoded["input_ids"].shape[1]), + completion_tokens=estimate_tokens(accumulated), + estimated=True, + ) + ) + except Exception as e: logger.exception( "Error in %s.stream_generate: %s", self.__class__.__name__, str(e) diff --git a/oli_bot/chat.py b/oli_bot/chat.py index 450e947..0885df6 100644 --- a/oli_bot/chat.py +++ b/oli_bot/chat.py @@ -53,7 +53,7 @@ ) from .config import configs -from .models import SubAgentRun +from .models import SubAgentRun, UsageEvent from .tools.manager import BuiltinToolManager from .tools.memory import _current_sub_run from .mcp_client import MCPClientManager @@ -293,6 +293,8 @@ def __init__( if self.backend.model is None or not self.backend.model: self.backend.model = self._get_large_model() or None self.messages: list[Message] = [] + self.total_tokens = 0 + self.tokens_estimated = False cwd = Path.cwd() if is_sensitive_path(cwd): self.session = Session(workspace=None) @@ -351,6 +353,8 @@ def __init__( [_message_from_dict(m) for m in data.get("messages", [])] ) self.current_session_id = data["id"] + self.total_tokens = data.get("total_tokens", 0) or 0 + self.tokens_estimated = bool(data.get("total_tokens_estimated", False)) if data.get("model"): self.backend.model = data["model"] loaded = True @@ -365,6 +369,10 @@ def __init__( [_message_from_dict(m) for m in data.get("messages", [])] ) self.current_session_id = data["id"] + self.total_tokens = data.get("total_tokens", 0) or 0 + self.tokens_estimated = bool( + data.get("total_tokens_estimated", False) + ) if data.get("model"): self.backend.model = data["model"] loaded = True @@ -560,6 +568,8 @@ def _save_session(self) -> None: messages=self.messages, model=self.backend.model or "", profile=self.agent.profile_name or "", + total_tokens=self.total_tokens, + tokens_estimated=self.tokens_estimated, ) if new_id != self.current_session_id: logger.info( @@ -585,6 +595,9 @@ def update_header(self) -> None: f"{self._status_badges()} {model_part} " f"[dim]:: {self.agent.profile_name}{server_part}[/dim]" ) + if self.total_tokens > 0: + est = "~" if self.tokens_estimated else "" + info += f" [dim]·[/dim] {est}{self.total_tokens:,} tok" status.add_row(hints, info) self.query_one("#status-bar", Static).update(status) @@ -825,9 +838,7 @@ async def _mcp_add(self) -> None: async def _mcp_edit(self, name: str) -> None: servers = {s.name: s for s in self.mcp_manager.list_servers()} if name not in servers: - self._add_message( - "System", f"[red]MCP server '{name}' not found.[/red]" - ) + self._add_message("System", f"[red]MCP server '{name}' not found.[/red]") return cfg = servers[name] existing = { @@ -850,9 +861,7 @@ async def _mcp_edit(self, name: str) -> None: transport=result.get("transport", "stdio"), url=result.get("url", ""), ) - self._add_message( - "System", f"MCP server [bold]{name}[/bold] updated." - ) + self._add_message("System", f"MCP server [bold]{name}[/bold] updated.") except ValueError as e: self._add_message("System", f"[red]{e}[/red]") @@ -919,6 +928,9 @@ def _sessions_new(self) -> None: self.messages.append( Message(role="system", content=self.agent.system_prompt) ) + self.total_tokens = 0 + self.tokens_estimated = False + self.update_header() self.query_one("#chat-log").remove_children() self._add_message("System", "New session created.") @@ -1025,6 +1037,9 @@ def _sessions_switch(self, target: str) -> None: [_message_from_dict(m) for m in data.get("messages", [])] ) self.current_session_id = session_id + self.total_tokens = data.get("total_tokens", 0) or 0 + self.tokens_estimated = bool(data.get("total_tokens_estimated", False)) + self.update_header() self.query_one("#chat-log").remove_children() self._clear_todo_widget() self._add_message("System", f"Switched to session [bold]{data['name']}[/bold].") @@ -1058,6 +1073,9 @@ def _sessions_delete(self, target: str) -> None: self.messages.append( Message(role="system", content=self.agent.system_prompt) ) + self.total_tokens = 0 + self.tokens_estimated = False + self.update_header() self.query_one("#chat-log").remove_children() self._add_message( "System", "Current session was deleted. New session created." @@ -1113,6 +1131,9 @@ async def _sessions_purge(self, server: str | None = None) -> None: profile=self.agent.profile_name or "", system_prompt=self.agent.system_prompt or "", ) + self.total_tokens = 0 + self.tokens_estimated = False + self.update_header() self.query_one("#chat-log").remove_children() self._add_message( "System", @@ -1999,6 +2020,9 @@ def action_clear_chat(self) -> None: self.messages.append( Message(role="system", content=self.agent.system_prompt) ) + self.total_tokens = 0 + self.tokens_estimated = False + self.update_header() self.query_one("#chat-log").remove_children() self._clear_todo_widget() self._add_message("System", "Conversation cleared.") @@ -2009,6 +2033,9 @@ def action_go_home(self) -> None: return self._save_session() self.messages.clear() + self.total_tokens = 0 + self.tokens_estimated = False + self.update_header() chat_log = self.query_one("#chat-log") try: welcome = chat_log.query_one("#welcome") @@ -2288,6 +2315,10 @@ def stop_spinner() -> None: ) except Exception: logger.warning("Failed to update error panel") + case UsageEvent(usage): + self.total_tokens += usage.total_tokens + self.tokens_estimated = self.tokens_estimated or usage.estimated + self.update_header() case AgentDone(full_text): stop_spinner() if think_widget is not None: @@ -2478,13 +2509,10 @@ def _register_dispatch_tool(self) -> None: ) ) - agent_description = ( - "Name of the sub-agent to run this task. " - + ( - f"Each agent belongs to a specific pool — {pool_summary}." - if has_multiple_pools - else f"Available: {', '.join(all_agent_names)}." - ) + agent_description = "Name of the sub-agent to run this task. " + ( + f"Each agent belongs to a specific pool — {pool_summary}." + if has_multiple_pools + else f"Available: {', '.join(all_agent_names)}." ) task_item_properties: dict = { @@ -2622,10 +2650,10 @@ def _todo_node_label(todo: dict) -> str: }.get(priority, "·") status_icon, style_open, style_close = { - "pending": ("○", "[#6b7d74]", "[/#6b7d74]"), + "pending": ("○", "[#6b7d74]", "[/#6b7d74]"), "in_progress": ("▶", "[bold #2ecc71]", "[/bold #2ecc71]"), - "completed": ("✓", "[dim #a9dfbf]", "[/dim #a9dfbf]"), - "cancelled": ("✗", "[dim]", "[/dim]"), + "completed": ("✓", "[dim #a9dfbf]", "[/dim #a9dfbf]"), + "cancelled": ("✗", "[dim]", "[/dim]"), }.get(status, ("·", "[dim]", "[/dim]")) return f"{priority_dot} {status_icon} {style_open}{content}{style_close}" diff --git a/oli_bot/models.py b/oli_bot/models.py index 1448536..ffa6b51 100644 --- a/oli_bot/models.py +++ b/oli_bot/models.py @@ -39,12 +39,30 @@ class ToolCall: parameters: Dict[str, Any] = field(default_factory=dict) +@dataclass +class Usage: + """Token counts for a single model call. + + ``estimated`` is True when any count was derived heuristically (e.g. the + transformers/stream fallback) rather than reported by the provider. + """ + + prompt_tokens: int = 0 + completion_tokens: int = 0 + estimated: bool = False + + @property + def total_tokens(self) -> int: + return self.prompt_tokens + self.completion_tokens + + @dataclass class ModelResponse: content: str tool_calls: Optional[List[ToolCall]] = None finish_reason: str = "stop" error: str = "" + usage: Optional[Usage] = None @dataclass @@ -72,6 +90,11 @@ class ToolCallChunk: tool_calls: List[ToolCall] +@dataclass +class UsageChunk: + usage: Usage + + @dataclass class ToolCallExecuting: name: str @@ -104,6 +127,13 @@ class Done: full_text: str +@dataclass +class UsageEvent: + """Cumulative token usage for a single agent run (one user turn).""" + + usage: Usage + + @dataclass class ProfileData: system_prompt: str diff --git a/oli_bot/screens/mcp_setup.py b/oli_bot/screens/mcp_setup.py index 50b288e..d256a6d 100644 --- a/oli_bot/screens/mcp_setup.py +++ b/oli_bot/screens/mcp_setup.py @@ -61,7 +61,9 @@ class MCPSetupScreen(ModalScreen[Optional[Dict[str, Any]]]): BINDINGS = [("escape", "cancel", "Cancel")] - def __init__(self, existing: Optional[Dict[str, Any]] = None, **kwargs: Any) -> None: + def __init__( + self, existing: Optional[Dict[str, Any]] = None, **kwargs: Any + ) -> None: """Create an MCP setup screen. Parameters diff --git a/oli_bot/screens/taglines.py b/oli_bot/screens/taglines.py index 1a546a2..f9768d6 100644 --- a/oli_bot/screens/taglines.py +++ b/oli_bot/screens/taglines.py @@ -80,7 +80,7 @@ "Loading personality... 100% (citation needed)", "Achievement unlocked: opened the app", "85% fat free!", - "I'm back, baby!" + "I'm back, baby!", ] diff --git a/oli_bot/sessions.py b/oli_bot/sessions.py index 88be03a..21d9d30 100644 --- a/oli_bot/sessions.py +++ b/oli_bot/sessions.py @@ -318,6 +318,8 @@ def create_session( "server": server, "model": model or "", "profile": profile or "", + "total_tokens": 0, + "total_tokens_estimated": False, "messages": [], } if system_prompt: @@ -336,6 +338,8 @@ def save_session( messages: List[Message], model: str, profile: str, + total_tokens: int = 0, + tokens_estimated: bool = False, ) -> str: """Persist a session; return the (possibly new) session id. @@ -380,6 +384,11 @@ def save_session( data["model"] = model or data.get("model", "") data["profile"] = profile or data.get("profile", "") data["messages"] = [_message_to_dict(m) for m in messages] + if created_new: + total_tokens = 0 + tokens_estimated = False + data["total_tokens"] = int(total_tokens or 0) + data["total_tokens_estimated"] = bool(tokens_estimated) path.write_text(json.dumps(data, indent=2), encoding="utf-8") if created_new: logger.info("Session recreated under new id: %s", session_id) diff --git a/oli_bot/tools/manager.py b/oli_bot/tools/manager.py index e869007..21822b2 100644 --- a/oli_bot/tools/manager.py +++ b/oli_bot/tools/manager.py @@ -233,7 +233,9 @@ def get_todos(self): cancelled_count=by_status.get("cancelled", 0), ) - def set_todo_callback(self, callback: Optional[Callable[[list[dict]], None]]) -> None: + def set_todo_callback( + self, callback: Optional[Callable[[list[dict]], None]] + ) -> None: """Register a callback that fires whenever the root agent's todo list changes.""" self._todo_change_callback = callback diff --git a/tests/test_agent_scaffolding.py b/tests/test_agent_scaffolding.py index 4b95e52..8819029 100644 --- a/tests/test_agent_scaffolding.py +++ b/tests/test_agent_scaffolding.py @@ -290,7 +290,9 @@ def test_select_agent_wrong_pool_raises(pool): pool.agent_pool["coding"] = {"code-writer": _make_agent("code-writer")} # code-writer lives in coding, not default - with pytest.raises(ValueError, match="code-writer not found in agent pool 'default'"): + with pytest.raises( + ValueError, match="code-writer not found in agent pool 'default'" + ): pool.select_agent("default", "code-writer") @@ -353,8 +355,14 @@ def test_multi_pool_build_from_yaml(monkeypatch): assert set(built_pool.agent_pool.keys()) == {"default", "coding"} assert built_pool.list_agents("default") == ["search-agent"] assert built_pool.list_agents("coding") == ["code-writer"] - assert built_pool.agent_pool["default"]["search-agent"].backend.base_url == "http://default-host:11434" - assert built_pool.agent_pool["coding"]["code-writer"].backend.base_url == "http://coding-host:11434" + assert ( + built_pool.agent_pool["default"]["search-agent"].backend.base_url + == "http://default-host:11434" + ) + assert ( + built_pool.agent_pool["coding"]["code-writer"].backend.base_url + == "http://coding-host:11434" + ) def test_sub_agent_run_pool_name_defaults_to_default(): @@ -369,5 +377,7 @@ def test_sub_agent_run_pool_name_can_be_set(): """SubAgentRun.pool_name stores the pool when explicitly provided.""" from oli_bot.models import SubAgentRun - run = SubAgentRun(task_id="t", agent_name="code-writer", task="write code", pool_name="coding") + run = SubAgentRun( + task_id="t", agent_name="code-writer", task="write code", pool_name="coding" + ) assert run.pool_name == "coding" diff --git a/tests/test_backend_stream_flush_extra.py b/tests/test_backend_stream_flush_extra.py index 38e9082..f69bbb7 100644 --- a/tests/test_backend_stream_flush_extra.py +++ b/tests/test_backend_stream_flush_extra.py @@ -106,7 +106,7 @@ async def test_ollama_yields_text_then_tool_call_in_the_same_chunk_order(): b = _make_ollama(chunks) events = [ev async for ev in b.stream_generate([], tools=[])] kinds = [type(e).__name__ for e in events] - assert kinds == ["TextChunk", "ToolCallChunk"] + assert kinds == ["TextChunk", "ToolCallChunk", "UsageChunk"] @pytest.mark.asyncio diff --git a/tests/test_backend_tool_call_flush.py b/tests/test_backend_tool_call_flush.py index 5dd344d..a7bc990 100644 --- a/tests/test_backend_tool_call_flush.py +++ b/tests/test_backend_tool_call_flush.py @@ -246,7 +246,8 @@ async def test_text_only_response_yields_no_tool_call_chunk(): events = [ev async for ev in _make_backend(chunks).stream_generate([], tools=[])] kinds = [type(e).__name__ for e in events] - assert kinds == ["TextChunk", "TextChunk"] + assert kinds == ["TextChunk", "TextChunk", "UsageChunk"] + assert not any(isinstance(e, ToolCallChunk) for e in events) assert "".join(e.text for e in events if isinstance(e, TextChunk)) == "hello world" diff --git a/tests/test_token_usage.py b/tests/test_token_usage.py new file mode 100644 index 0000000..6884748 --- /dev/null +++ b/tests/test_token_usage.py @@ -0,0 +1,285 @@ +"""Token-usage plumbing: backend UsageChunk emission, agent UsageEvent +aggregation, and session persistence of the running total.""" + +import pytest +from types import SimpleNamespace + +from oli_bot.agent import Agent +from oli_bot.config import AppConfig +from oli_bot.models import ( + Done, + Message, + TextChunk, + ToolCall, + ToolCallChunk, + Usage, + UsageChunk, + UsageEvent, +) +from oli_bot.sessions import ConversationStore +from oli_bot.backends import OllamaBackend, OpenAIBackend + +# --------------------------------------------------------------------------- # +# OpenAI streaming # +# --------------------------------------------------------------------------- # + + +class _AsyncIter: + def __init__(self, chunks): + self._chunks = chunks + + def __aiter__(self): + return self._agen() + + async def _agen(self): + for c in self._chunks: + yield c + + +class _DummyCompletions: + def __init__(self, chunks): + self._chunks = chunks + + async def create(self, **_kwargs): + return _AsyncIter(self._chunks) + + +def _openai_chunk(*, content=None, finish=None, usage=None): + delta = SimpleNamespace(content=content, tool_calls=None) + choice = SimpleNamespace(delta=delta, finish_reason=finish) + return SimpleNamespace(choices=[choice], usage=usage) + + +def _make_openai(chunks, *, create=None): + b = OpenAIBackend.__new__(OpenAIBackend) + b.model = "gpt-fake" + b.vision_style = "openai" + if create is None: + create = _DummyCompletions(chunks).create + b.client = SimpleNamespace( + chat=SimpleNamespace(completions=SimpleNamespace(create=create)) + ) + return b + + +@pytest.mark.asyncio +async def test_openai_stream_exact_usage_from_final_chunk(): + chunks = [ + _openai_chunk(content="hi ", finish=None), + _openai_chunk(content="there", finish="stop"), + _openai_chunk( + finish=None, + usage=SimpleNamespace(prompt_tokens=12, completion_tokens=9), + ), + ] + b = _make_openai(chunks) + events = [ev async for ev in b.stream_generate([], tools=[])] + + usage_chunks = [e for e in events if isinstance(e, UsageChunk)] + assert len(usage_chunks) == 1 + u = usage_chunks[0].usage + assert u.prompt_tokens == 12 + assert u.completion_tokens == 9 + assert u.total_tokens == 21 + assert u.estimated is False + + +@pytest.mark.asyncio +async def test_openai_stream_estimates_when_no_usage_reported(): + chunks = [ + _openai_chunk(content="hello", finish="stop"), + ] + b = _make_openai(chunks) + events = [ev async for ev in b.stream_generate([], tools=[])] + + usage_chunks = [e for e in events if isinstance(e, UsageChunk)] + assert len(usage_chunks) == 1 + u = usage_chunks[0].usage + assert u.estimated is True + assert u.prompt_tokens >= 1 + assert u.completion_tokens >= 1 + + +@pytest.mark.asyncio +async def test_openai_retries_without_stream_options_when_rejected(): + calls = [] + calls_history = [] + + async def flaky_create(**kwargs): + calls.append(1) + calls_history.append(kwargs) + if kwargs.get("stream_options") is not None: + raise RuntimeError("unknown parameter: stream_options") + return _AsyncIter([_openai_chunk(content="ok", finish="stop")]) + + b = _make_openai([], create=flaky_create) + events = [ev async for ev in b.stream_generate([], tools=[])] + + assert len(calls) == 2 + assert "stream_options" not in calls_history[1] + texts = [e.text for e in events if isinstance(e, TextChunk)] + assert "".join(texts) == "ok" + assert any(isinstance(e, UsageChunk) for e in events) + + +# --------------------------------------------------------------------------- # +# Ollama streaming # +# --------------------------------------------------------------------------- # + + +class _OllamaAsyncClient: + def __init__(self, chunks): + self._chunks = chunks + + async def chat(self, **_kwargs): + return _AsyncIter(self._chunks) + + +def _ollama_chunk(*, content=None, prompt_eval=None, eval_count=None): + return SimpleNamespace( + message=SimpleNamespace(content=content, tool_calls=None, thinking=None), + prompt_eval_count=prompt_eval, + eval_count=eval_count, + ) + + +def _make_ollama(chunks): + b = OllamaBackend.__new__(OllamaBackend) + b.model = "ollama-fake" + b.base_url = "http://x" + b.client = _OllamaAsyncClient(chunks) + return b + + +@pytest.mark.asyncio +async def test_ollama_stream_exact_usage_from_eval_counts(): + chunks = [ + _ollama_chunk(content="hello", prompt_eval=10, eval_count=3), + _ollama_chunk(content=" world", prompt_eval=10, eval_count=5), + ] + b = _make_ollama(chunks) + events = [ev async for ev in b.stream_generate([], tools=[])] + + usage_chunks = [e for e in events if isinstance(e, UsageChunk)] + assert len(usage_chunks) == 1 + u = usage_chunks[0].usage + assert u.prompt_tokens == 10 + assert u.completion_tokens == 5 + assert u.estimated is False + + +@pytest.mark.asyncio +async def test_ollama_stream_estimates_when_counts_absent(): + chunks = [_ollama_chunk(content="hi", prompt_eval=None, eval_count=None)] + b = _make_ollama(chunks) + events = [ev async for ev in b.stream_generate([], tools=[])] + + usage_chunks = [e for e in events if isinstance(e, UsageChunk)] + assert len(usage_chunks) == 1 + assert usage_chunks[0].usage.estimated is True + + +# --------------------------------------------------------------------------- # +# Agent aggregation # +# --------------------------------------------------------------------------- # + + +class _UsageToolStub: + """First call: tool call + usage. Second call: final text + usage.""" + + model = "stub" + + def __init__(self): + self._calls = 0 + + async def stream_generate(self, messages, tools=None): + self._calls += 1 + if self._calls == 1: + yield ToolCallChunk( + [ToolCall(id="c1", name="builtin__echo", parameters={"x": 1})] + ) + yield UsageChunk(Usage(prompt_tokens=5, completion_tokens=3)) + else: + yield TextChunk("done") + yield UsageChunk(Usage(prompt_tokens=10, completion_tokens=7)) + + +class _StubMCP: + async def call_tool(self, name, params, confirm_callback=None): + return "ok" + + +class _NoUsageStub: + model = "stub" + + async def stream_generate(self, messages, tools=None): + yield TextChunk("plain text") + + +def _agent(backend) -> Agent: + return Agent( + role="default", + backend=backend, + mcp_manager=_StubMCP(), + profile_name="default", + mode="agent", + config=AppConfig(_env_file=None, max_tool_iterations=3), + ) + + +@pytest.mark.asyncio +async def test_agent_emits_single_usage_event_summed_across_tool_iterations(): + a = _agent(_UsageToolStub()) + msgs = [Message(role="system", content="sys"), Message(role="user", content="hi")] + events = [ev async for ev in a.process(msgs)] + + usage_events = [e for e in events if isinstance(e, UsageEvent)] + assert len(usage_events) == 1 + u = usage_events[0].usage + assert u.prompt_tokens == 15 + assert u.completion_tokens == 10 + assert u.estimated is False + + # UsageChunk is consumed internally, never forwarded to the UI. + assert not any(isinstance(e, UsageChunk) for e in events) + assert any(isinstance(e, Done) for e in events) + + +@pytest.mark.asyncio +async def test_agent_emits_no_usage_event_when_backend_reports_none(): + a = _agent(_NoUsageStub()) + msgs = [Message(role="system", content="sys"), Message(role="user", content="hi")] + events = [ev async for ev in a.process(msgs)] + + assert not any(isinstance(e, UsageEvent) for e in events) + assert any(isinstance(e, Done) for e in events) + + +# --------------------------------------------------------------------------- # +# Session persistence # +# --------------------------------------------------------------------------- # + + +def test_session_round_trips_token_totals(tmp_path): + store = ConversationStore(sessions_dir=tmp_path) + sid = store.create_session("srv", "m", "p", "system") + store.save_session( + "srv", + sid, + [Message(role="user", content="hi")], + "m", + "p", + total_tokens=1234, + tokens_estimated=True, + ) + data = store.load_session("srv", sid) + assert data["total_tokens"] == 1234 + assert data["total_tokens_estimated"] is True + + +def test_new_session_defaults_to_zero_tokens(tmp_path): + store = ConversationStore(sessions_dir=tmp_path) + sid = store.create_session("srv", "m", "p", "") + data = store.load_session("srv", sid) + assert data["total_tokens"] == 0 + assert data["total_tokens_estimated"] is False