agentstate-reducer

agentstate-reducer

Framework-agnostic message reducer for AI agent state management. Works with LangGraph, CrewAI, and plain dicts.

What It Does

Automatically prunes message history when it exceeds a threshold, keeping conversations manageable:

• Windowed pruning: Trigger at max_messages, retain min_messages
• Token-budget pruning: Trigger at max_tokens, prune down to target_tokens — whole messages only, never truncated
• System message preservation: Index 0 (system prompt) is never pruned (configurable)
• ToolMessage cascade: When an AI message is pruned, linked ToolMessages are pruned too
• Optional summarization: Callback with pruned messages to generate an LLM summary
• Long-term memory hooks (on_prune) — new in 0.4.0: pruned messages are handed to any callable the moment they leave the window, with the app's memory namespace, exactly once. Wire them to a LangGraph BaseStore, LangMem, a Strands MemoryStore, or your own engine — the reducer imports none of them. See Long-Term Memory Hooks.
• Role alias normalisation: Understands user/assistant/agent in addition to human/ai — works with OpenAI-format dicts and agent framework outputs out of the box
• Framework-agnostic: Works with plain dicts, LangChain BaseMessage subclasses, or any duck-typed message object — zero dependencies

Install

pip install agentstate-reducer

Zero dependencies. Works with Python 3.10+.

Quick Start

from agentstate_reducer import MessageReducer

reducer = MessageReducer(min_messages=10, max_messages=20)

# Works with plain dicts
result = reducer.reduce(
    existing=[
        {"role": "system", "content": "You are helpful"},
        {"role": "human", "content": "Hello"},
        {"role": "ai", "content": "Hi there!"},
    ],
    new=[{"role": "human", "content": "New message"}],
)
print(result.surviving)  # Messages that remain
print(result.pruned)     # Messages that were removed

# Works with LangChain BaseMessage objects too
from langchain_core.messages import HumanMessage, AIMessage
result = reducer.reduce(
    existing=[HumanMessage(content="Hello")],
    new=[AIMessage(content="Hi!")],
)

LangGraph Integration

Use directly as a state annotation — the reducer runs automatically on every state update:

from agentstate_reducer import MessageReducer
from typing_extensions import Annotated, TypedDict

reducer = MessageReducer(min_messages=10, max_messages=20)

class MyState(TypedDict):
    messages: Annotated[list, reducer.as_langgraph_reducer()]

as_langgraph_reducer() returns a (existing, new) -> list function that LangGraph calls on every state merge.

Any LangGraph Checkpointer, Bounded (ReducingSaver)

Checkpoint storage grows on two axes: how many checkpoints you keep (prune, TTL) and how big each one is. The messages channel is stored whole on every step, so a thread's blobs grow quadratically with its length, and keep-latest never shrinks the latest one. ReducingSaver runs the reducer inside any saver's put, so every checkpoint holds at most max_messages messages and the pruned turns go to on_prune once each:

pip install "agentstate-reducer[langgraph]"

from langgraph.checkpoint.postgres import PostgresSaver
from agentstate_reducer import MessageReducer, ReducerConfig
from agentstate_reducer.langgraph import ReducingSaver

saver = ReducingSaver(PostgresSaver.from_conn_string(url), MessageReducer(config=ReducerConfig(max_messages=20)))
graph = builder.compile(checkpointer=saver)

Works with PostgresSaver, SQLite, MongoDB, Redis, in-memory. Wrapping one of the agentstate checkpointers (DynamoDB, Cosmos DB, Firestore) hands the reducer to the saver instead, so it still runs once; wrapping a saver that already has a different reducer raises. Optional capabilities (copy_thread, delete_for_runs, prune) are delegated only when the inner saver implements them. Not for DeltaChannel-backed message channels. Docs: Any Checkpointer, Bounded.

CosmosDB Checkpoint Integration

When using langgraph-checkpoint-cosmosdb, pass a MessageReducer to reduce messages at the persistence layer — useful when you don't control the state definition:

from agentstate_reducer import MessageReducer
from langgraph_checkpoint_cosmosdb import CosmosDBSaver

reducer = MessageReducer(min_messages=10, max_messages=20)
saver = CosmosDBSaver(
    database_name="mydb",
    container_name="checkpoints",
    reducer=reducer,           # applied before each checkpoint is stored
    messages_key="messages",   # which channel holds the message list (default)
)

Summarization

Provide a summarize_fn to generate a summary of pruned messages (e.g. via an LLM):

from agentstate_reducer import MessageReducer, ReducerConfig

def summarize(pruned_messages):
    # Call your LLM here
    return f"Summary of {len(pruned_messages)} pruned messages"

config = ReducerConfig(
    min_messages=10,
    max_messages=20,
    summarize_fn=summarize,
)
reducer = MessageReducer(config=config)
result = reducer.reduce(existing=messages)
print(result.summary)  # "Summary of 5 pruned messages"

result.summary is always returned when a summarize_fn is set — you can store or display it however you like.

Injecting the summary back into the conversation

By default the summary is only returned, not added to the surviving messages. Set inject_summary=True to also insert it in place of the pruned block, so the model sees it on the next turn:

config = ReducerConfig(
    min_messages=10,
    max_messages=20,
    summarize_fn=summarize,
    inject_summary=True,
)

The default injection is a human + ai pair — a human message carrying the summary followed by a short ai acknowledgement ("OK"):

[system] → [human: "Here is the summary of previous N messages: ..."] → [ai: "OK"] → [recent tail]

Why a pair? It keeps the conversation well-formed for providers that require strict role alternation (system → human → ai → …). Because both messages are ordinary human/ai messages at the front of the window, the next prune sweeps them into the new summary automatically — so there is always exactly one summary block, with no accumulation and no bookkeeping.

Custom shape. Override summary_message_factory — a Callable[[summary_text, pruned_count], list[message]] — to inject a single message or a provider-specific shape:

config = ReducerConfig(
    summarize_fn=summarize,
    inject_summary=True,
    # inject just one system message instead of the human/ai pair
    summary_message_factory=lambda text, n: [{"role": "system", "content": f"Summary ({n} msgs): {text}"}],
)

The summary is placed where the pruned messages were — after preserve_first, before the retained recent tail — never at index 0 (the system prompt is preserved). When used with a checkpoint/persistence saver that stores result.surviving, the injected summary is persisted automatically with no saver changes.

Long-Term Memory Hooks (on_prune)

Messages leaving the window are exactly the ones worth deciding about for long-term memory. on_prune hands them to any callable the moment they are pruned — the reducer never imports a store, a framework, or an LLM, so it stays zero-dependency and the same hook body works for every framework that runs the reducer (LangGraph checkpointers, CrewAI persistence, ...).

from agentstate_reducer import MessageReducer, ReducerConfig, Background

def remember(pruned, namespace):
    # namespace is whatever the framework integration forwarded (see below)
    for m in pruned:
        store.put(namespace, key=str(uuid4()), value={"role": m.type, "content": m.content})

reducer = MessageReducer(config=ReducerConfig(
    max_messages=20,
    on_prune=[remember],              # cheap hook: runs inline
    # on_prune=[Background(remember)]  # slow hook (LLM extraction): runs off the request path
))

Where the namespace comes from. reduce(..., namespace=X) forwards X untouched to every hook. Framework integrations look up ReducerConfig.namespace_key (default "memory_namespace") in their per-call config and pass that — e.g. in LangGraph:

graph.invoke(input, config={"configurable": {
    "thread_id": uuid4().hex,                       # short-term scope (checkpoint)
    "memory_namespace": ("memories", user.id),      # long-term scope (store)
}})

If the app never sets it, integrations fall back to a per-thread namespace. Old callers that pass no namespace get None.

Exactly once. Persistence layers often call reduce() several times per turn on overlapping lists (LangGraph writes a checkpoint per super-step). With dedupe_on_prune=True (default) each message id reaches the hooks once per reducer instance; ReducerResult.pruned is never filtered.

Background(fn, workers=2, max_pending=1000) runs a hook on a bounded worker pool: copies the list, drops (with a warning) when the backlog is full, swallows hook exceptions, drains at exit. On serverless runtimes call close() at the end of the handler.

Token-Budget Pruning

Instead of counting messages, you can prune to a token budget — useful when you want to stay within a model's context window or control cost. Set max_tokens and pruning switches from message-count mode to token mode.

from agentstate_reducer import MessageReducer, ReducerConfig

# Prune when the conversation exceeds 4000 tokens, down to ~2000
config = ReducerConfig(max_tokens=4000, target_tokens=2000)
reducer = MessageReducer(config=config)

result = reducer.reduce(existing=messages, new=new_messages)
# result.surviving stays within ~2000 tokens; whole messages only — never truncated

Whole messages only. The reducer never truncates message content — it drops whole messages, keeping the most recent ones that fit the budget, plus the preserved first message. This guarantees you never send a model a half-cut message.

max_tokens vs target_tokens. Pruning triggers when the total exceeds max_tokens, and reduces down to target_tokens (defaults to max_tokens if not set). Setting target_tokens lower than max_tokens creates hysteresis — prune at 4000, down to 2000 — so pruning runs less often.

How tokens are counted

The counter is resolved in three layers (highest priority first):

1. User-supplied token_counter — a Callable[[message], int] you pass on the config. Use this for exact, model-specific counting:

   import tiktoken
   enc = tiktoken.encoding_for_model("gpt-4o")
   config = ReducerConfig(
       max_tokens=4000,
       token_counter=lambda m: len(enc.encode(m.get("content", ""))),
   )
   

2. tiktoken — if installed (pip install "agentstate-reducer[tokens]"), the cl100k_base encoding is used automatically. Accurate for OpenAI-family models.
3. Character heuristic — len(content) / 4 plus a small per-message overhead. Dependency-free fallback, fine for approximate budgeting.

Install with tiktoken support:

pip install "agentstate-reducer[tokens]"

Token mode takes precedence over message-count mode: if max_tokens is set, max_messages/min_messages are ignored. preserve_first and cascade_tool_messages apply in both modes.

Pruning Behaviour

• Only ai/agent/assistant and human/user messages are candidates for pruning.
• system, tool, and function messages are never pruned directly.
• Index 0 is preserved by default (preserve_first=True) — typically the system prompt.
• When an ai message is pruned, all tool messages linked to it via tool_call_id are also pruned (cascade_tool_messages=True).
• Pruning is windowed: once len(existing + new) > max_messages, oldest eligible messages are removed until min_messages remain.

Role Aliases

The reducer normalises common role name variants before applying pruning rules, so you don't need to convert message formats:

This means OpenAI-format message lists (role: "user" / role: "assistant") work without any preprocessing.

Framework Compatibility

The adapter layer uses duck typing and class-name inspection — no langchain_core import required:

API Reference

MessageReducer(min_messages, max_messages, *, config)

ReducerConfig

reducer.reduce(existing, new, namespace=None) -> ReducerResult

reducer.as_langgraph_reducer() -> Callable

Returns a function with signature (existing, new) -> list for use with LangGraph's Annotated[list, fn] pattern.