Skip to content

feat: add context_tokens to internal Agents State - #12102

Merged
sjrl merged 7 commits into
mainfrom
add-context-tokens-agent
Jul 23, 2026
Merged

feat: add context_tokens to internal Agents State#12102
sjrl merged 7 commits into
mainfrom
add-context-tokens-agent

Conversation

@sjrl

@sjrl sjrl commented Jul 21, 2026

Copy link
Copy Markdown
Contributor

Related Issues

Proposed Changes:

The Agent now tracks an approximate current context-window size in its internal State under context_tokens, refreshed after every LLM call with that reply's prompt-plus-completion tokens (normalized across the prompt_tokens/completion_tokens and input_tokens/output_tokens key conventions). Unlike token_usage, which accumulates across the whole run, context_tokens is replaced each call. Hooks can read it via state.get("context_tokens") — for example, a before_llm hook that triggers context compaction once the value crosses a threshold. It is a best-effort snapshot: it is 0 when the generator does not report usage, and does not count messages appended after the latest call until the next call refreshes it.

How did you test it?

New tests

Notes for the reviewer

I decided to add a new utils file agent/utils.py to start collecting these util functions that are crowding the top of the agent.py file. After this PR I will open a new one to move more util functionality into that file. This is not a breaking change since all util functions being moved are private.

I also ran live runs against different model providers APIs using our chat generators to see what usage looks like and confirm that the agent's util methods works for our most popular providers. I've collected them in test_normalizes_provider_shapes

Checklist

  • I have read the contributors guidelines and the code of conduct.
  • I have updated the related issue with new insights and changes.
  • I have added unit tests and updated the docstrings.
  • I've used one of the conventional commit types for my PR title: fix:, feat:, build:, chore:, ci:, docs:, style:, refactor:, perf:, test: and added ! in case the PR includes breaking changes.
  • I have documented my code.
  • I have added a release note file, following the contributors guidelines.
  • I have run pre-commit hooks and fixed any issue.

@vercel

vercel Bot commented Jul 21, 2026

Copy link
Copy Markdown

The latest updates on your projects. Learn more about Vercel for GitHub.

1 Skipped Deployment
Project Deployment Actions Updated (UTC)
haystack-docs Ignored Ignored Preview Jul 22, 2026 1:43pm

Request Review

@sjrl sjrl self-assigned this Jul 21, 2026
@github-actions github-actions Bot added the type:documentation Improvements on the docs label Jul 21, 2026
@github-actions

github-actions Bot commented Jul 21, 2026

Copy link
Copy Markdown
Contributor

Coverage report

Click to see where and how coverage changed

FileStatementsMissingCoverageCoverage
(new stmts)
Lines missing
  haystack/components/agents
  agent.py
  utils.py
Project Total  

This report was generated by python-coverage-comment-action

@sjrl
sjrl marked this pull request as ready for review July 22, 2026 13:45
@sjrl
sjrl requested a review from a team as a code owner July 22, 2026 13:45
@sjrl
sjrl requested review from julian-risch and removed request for a team July 22, 2026 13:45

@julian-risch julian-risch left a comment

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Looks very good to me! 🚀

@sjrl
sjrl merged commit 478471a into main Jul 23, 2026
29 checks passed
@sjrl
sjrl deleted the add-context-tokens-agent branch July 23, 2026 15:12
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

topic:tests type:documentation Improvements on the docs

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants