Skip to content

Context Window Stats Utility

Track how much of the model's context window a MemAgent has consumed at any point during a conversation. The agent records usage every time it calls the underlying LLM and makes the latest snapshot accessible through logs and code.

What the Agent Records

Each snapshot includes:

  • timestamp: ISO-8601 timestamp for the measurement
  • stage: the agent stage that triggered the measurement (e.g., iteration_1, memory_compression)
  • prompt_tokens: number of tokens sent to the model
  • completion_tokens: tokens generated by the model
  • total_tokens: sum of prompt + completion tokens
  • context_window_tokens: configured or inferred window size
  • percentage_used: total usage / context window × 100

If the provider does not return usage information, MemoRizz falls back to None so downstream code can handle missing values gracefully.

Logging

By default the agent logs each measurement at INFO level:

Context window usage (iteration_2): 2,350/128,000 tokens (1.84%) | prompt=2,100 completion=250

Monitor your existing log stream (e.g., tail -f app.log) to watch the token budget drain in real time.

Programmatic Access

Use memagent.get_context_window_stats() to retrieve the most recent snapshot after an interaction:

response = agent.run("Summarize the workshop agenda we discussed yesterday.")

stats = agent.get_context_window_stats()
if stats:
    print(
        f"Total tokens: {stats['total_tokens']}"
        f" ({stats['percentage_used']:.2f}% of {stats['context_window_tokens']})"
    )
else:
    print("Provider did not return usage information.")

Snapshots are ordinary dictionaries, so you can emit them to observability pipelines, dashboards, or audits.

Configuring the Context Window

MemoRizz tries to detect the context window automatically:

  1. Use the explicit context_window_tokens argument passed to MemAgent or MemAgentBuilder.
  2. If not provided, look for context_window_tokens / max_context_tokens / context_window inside llm_config.
  3. Fall back to the provider's built-in knowledge (OpenAI & Azure expose known limits; Hugging Face derives the tokenizer limit).

You can override the inferred value at any time:

agent = (MemAgentBuilder()
    .with_llm_config({"provider": "openai", "model": "gpt-4o-mini"})
    .with_memory_provider(provider)
    .build()
)

# Later, adjust for a custom fine-tuned model
agent._context_window_tokens = 32_000

Tip: When you know the exact budget (e.g., for a fine-tuned or local model) always pass it explicitly so percentage calculations remain accurate.

Provider Support

  • OpenAI / Azure OpenAI: Token usage comes directly from the API response (response.usage).
  • Hugging Face: The provider counts tokens via the active tokenizer (falling back to whitespace splitting when needed).
  • Custom Providers: Implement the LLMProvider protocol’s get_last_usage() and get_context_window_tokens() methods to plug into the same reporting pipeline.

With these hooks in place, every MemAgent—single or multi-step—can report how close it is to the model’s context limit, helping you catch runaway prompts before they overflow the window.