agentstate-reducer 0.1.0__py3-none-any.whl

agentstate_reducer/__init__.py

@@ -0,0 +1,14 @@
+"""
+agentstate-reducer: Framework-agnostic message reducer for AI agent state management.
+
+Works with LangGraph, CrewAI, and plain dicts.
+"""
+
+from .models import ReducerConfig, ReducerResult
+from .reducer import MessageReducer
+
+__all__ = [
+    "MessageReducer",
+    "ReducerConfig",
+    "ReducerResult",
+]

agentstate_reducer/adapters.py

@@ -0,0 +1,128 @@
+"""
+Adapters for extracting fields from messages in a framework-agnostic way.
+
+Works with:
+- Plain dicts:  {"role": "ai", "content": "...", "id": "...", "tool_calls": [...]}
+- LangChain BaseMessage subclasses: AIMessage, HumanMessage, ToolMessage, etc.
+- Any object with the expected attributes (duck typing)
+
+This module does NOT import langchain_core. It uses duck typing and class-name
+inspection so the package remains zero-dependency.
+"""
+
+from typing import Any, Dict, List, Optional
+
+
+# ── Role mapping from LangChain class names ──
+
+_CLASS_NAME_TO_ROLE = {
+    "human": "human",
+    "humanmessage": "human",
+    "ai": "ai",
+    "aimessage": "ai",
+    "tool": "tool",
+    "toolmessage": "tool",
+    "system": "system",
+    "systemmessage": "system",
+    "function": "function",
+    "functionmessage": "function",
+    "chatmessage": "chat",
+}
+
+
+def get_role(msg: Any) -> str:
+    """
+    Extract the role/type from a message.
+
+    Handles:
+    - dict with "role" key:  {"role": "ai", ...}
+    - dict with "type" key:  {"type": "ai", ...}  (LangChain serialized format)
+    - LangChain BaseMessage: uses class name → role mapping
+    - Any object with a .type attribute
+
+    Returns:
+        Lowercase role string: "human", "ai", "tool", "system", etc.
+    """
+    if isinstance(msg, dict):
+        return (msg.get("role") or msg.get("type") or "").lower()
+
+    # Try class name mapping (covers all LangChain message types)
+    class_name = type(msg).__name__.lower()
+    role = _CLASS_NAME_TO_ROLE.get(class_name)
+    if role:
+        return role
+
+    # Fallback: .type attribute (LangChain BaseMessage has this)
+    return getattr(msg, "type", "unknown").lower()
+
+
+def get_id(msg: Any) -> Optional[str]:
+    """
+    Extract the unique message ID.
+
+    Handles:
+    - dict: msg["id"]
+    - Object: msg.id
+    """
+    if isinstance(msg, dict):
+        return msg.get("id")
+    return getattr(msg, "id", None)
+
+
+def get_content(msg: Any) -> str:
+    """
+    Extract message content.
+
+    Handles:
+    - dict: msg["content"]
+    - Object: msg.content
+    """
+    if isinstance(msg, dict):
+        return msg.get("content", "")
+    return getattr(msg, "content", "")
+
+
+def get_tool_calls(msg: Any) -> List[Dict]:
+    """
+    Extract tool_calls from an AI message.
+
+    Handles:
+    - dict: msg["tool_calls"] → list of dicts
+    - LangChain AIMessage: msg.tool_calls → list of ToolCall objects or dicts
+
+    Returns:
+        List of dicts with at least {"id": ...} for each tool call.
+        Returns empty list if no tool_calls found.
+    """
+    if isinstance(msg, dict):
+        return msg.get("tool_calls", [])
+
+    raw = getattr(msg, "tool_calls", [])
+    if not raw:
+        return []
+
+    result = []
+    for tc in raw:
+        if isinstance(tc, dict):
+            result.append(tc)
+        else:
+            # LangChain ToolCall object → convert to dict
+            result.append({
+                "id": getattr(tc, "id", None),
+                "name": getattr(tc, "name", None),
+                "args": getattr(tc, "args", {}),
+            })
+    return result
+
+
+def get_tool_call_id(msg: Any) -> Optional[str]:
+    """
+    Extract tool_call_id from a ToolMessage.
+
+    Handles:
+    - dict: msg["tool_call_id"]
+    - LangChain ToolMessage: msg.tool_call_id
+    """
+    if isinstance(msg, dict):
+        return msg.get("tool_call_id")
+    return getattr(msg, "tool_call_id", None)

agentstate_reducer/models.py

@@ -0,0 +1,53 @@
+"""
+Shared data models for the agentstate-reducer package.
+
+These are framework-agnostic — no imports from langchain, crewai, etc.
+"""
+
+from dataclasses import dataclass, field
+from typing import Any, Callable, List, Optional
+
+
+@dataclass
+class ReducerConfig:
+    """
+    Configuration for the MessageReducer.
+
+    Attributes:
+        min_messages:          Number of messages to retain after pruning.
+        max_messages:          Pruning triggers when message count exceeds this.
+                               Set to None to disable pruning.
+        preserve_first:        If True, index 0 is never pruned (system message).
+        cascade_tool_messages: If True, when an AIMessage is pruned, also prune
+                               any ToolMessages linked to it via tool_call_id.
+        summarize_fn:          Optional callable(pruned_messages) -> str.
+                               Called with the list of pruned messages so you can
+                               generate a summary (e.g., via LLM) of what was removed.
+    """
+
+    min_messages: int = 10
+    max_messages: Optional[int] = 20
+    preserve_first: bool = True
+    cascade_tool_messages: bool = True
+    summarize_fn: Optional[Callable[[List[Any]], str]] = None
+
+
+@dataclass
+class ReducerResult:
+    """
+    Output of a reduce operation.
+
+    Both downstream consumers need to know what survived AND what was pruned:
+    - langgraph-checkpoint-cosmosdb: stores `surviving`, may log `pruned`
+    - crewai-persistence-cosmosdb:   stores `surviving`, may call summarize on `pruned`
+
+    Attributes:
+        surviving: Messages that remain after pruning.
+        pruned:    Messages that were removed.
+        summary:   If a summarize_fn was configured and pruning occurred,
+                   this contains the summary string.
+    """
+
+    surviving: List[Any] = field(default_factory=list)
+    pruned: List[Any] = field(default_factory=list)
+    summary: Optional[str] = None

agentstate_reducer/reducer.py

@@ -0,0 +1,208 @@
+"""
+Core message reducer — framework-agnostic.
+
+Works with:
+- Plain dicts:      {"role": "ai", "content": "...", "id": "..."}
+- LangChain types:  AIMessage, HumanMessage, ToolMessage, SystemMessage
+- Any object with role/id/tool_calls attributes (duck typing)
+
+Ported from langgraph-reducer's Reducer class with these improvements:
+- Works with both LangChain types and plain dicts
+- Returns ReducerResult(surviving, pruned, summary) so callers decide
+  what to do with pruned messages
+- Configurable via ReducerConfig dataclass
+- Provides as_langgraph_reducer() bridge for Annotated[list, fn] usage
+
+Preserves original behaviors:
+1. Skip index 0 (system message) during pruning
+2. When pruning an AIMessage, also prune linked ToolMessages (cascade)
+3. Windowed pruning: prune when len > max_messages, keep min_messages
+"""
+
+import logging
+from typing import Any, List, Optional, Set
+
+from .adapters import get_role, get_tool_call_id, get_tool_calls
+from .models import ReducerConfig, ReducerResult
+
+logger = logging.getLogger(__name__)
+
+
+class MessageReducer:
+    """
+    Framework-agnostic message reducer.
+
+    Concatenates existing + new messages, then prunes oldest messages
+    when the total exceeds ``max_messages``, retaining ``min_messages``.
+
+    Pruning rules:
+    - Index 0 is preserved (configurable via ``preserve_first``).
+    - Only ``ai`` and ``human`` messages are pruned.
+    - When an ``ai`` message is pruned, any ``tool`` messages linked to
+      it (by ``tool_call_id``) are also pruned (cascade).
+
+    Args:
+        min_messages: Number of messages to keep after pruning.
+        max_messages: Threshold to trigger pruning. ``None`` disables pruning.
+        config:       Optional ``ReducerConfig`` (overrides min/max params).
+
+    Example with plain dicts::
+
+        reducer = MessageReducer(min_messages=5, max_messages=10)
+        result = reducer.reduce(
+            existing=[{"role": "system", "content": "You are helpful"}],
+            new=[{"role": "human", "content": "Hello"}],
+        )
+        # result.surviving = [...]
+        # result.pruned = [...]
+
+    Example with LangChain messages::
+
+        from langchain_core.messages import HumanMessage, AIMessage
+        reducer = MessageReducer(min_messages=5, max_messages=10)
+        result = reducer.reduce(
+            existing=[SystemMessage(content="...")],
+            new=[HumanMessage(content="Hello")],
+        )
+    """
+
+    def __init__(
+        self,
+        min_messages: int = 0,
+        max_messages: Optional[int] = None,
+        *,
+        config: Optional[ReducerConfig] = None,
+    ):
+        if config is not None:
+            self.config = config
+        else:
+            self.config = ReducerConfig(
+                min_messages=min_messages,
+                max_messages=max_messages,
+            )
+
+    @property
+    def min_messages(self) -> int:
+        return self.config.min_messages
+
+    @property
+    def max_messages(self) -> Optional[int]:
+        return self.config.max_messages
+
+    def reduce(
+        self,
+        existing: Optional[List[Any]] = None,
+        new: Optional[List[Any]] = None,
+    ) -> ReducerResult:
+        """
+        Concatenate existing + new messages, then prune if over threshold.
+
+        Args:
+            existing: Current message list. Defaults to empty list.
+            new:      New messages to append. Defaults to empty list.
+
+        Returns:
+            ``ReducerResult`` with surviving messages, pruned messages,
+            and optional summary.
+        """
+        if existing is None:
+            existing = []
+        if new is None:
+            new = []
+
+        messages = list(existing) + list(new)
+        max_msgs = self.config.max_messages
+
+        # No pruning needed
+        if max_msgs is None or len(messages) <= max_msgs:
+            return ReducerResult(surviving=messages, pruned=[])
+
+        # ── Identify indices to prune ──
+        to_delete: Set[int] = set()
+        excess_count = len(messages) - self.config.min_messages
+        start_idx = 1 if self.config.preserve_first else 0
+
+        # Iterate over the pruning window: from start_idx to excess_count
+        for i, msg in enumerate(messages[start_idx:excess_count], start=start_idx):
+            role = get_role(msg)
+
+            if role in ("ai", "human"):
+                to_delete.add(i)
+
+            # ToolMessage cascade: when pruning an AIMessage, also prune
+            # any ToolMessages that reference its tool_call IDs
+            if role == "ai" and self.config.cascade_tool_messages:
+                for tc in get_tool_calls(messages[i]):
+                    tc_id = (
+                        tc.get("id")
+                        if isinstance(tc, dict)
+                        else getattr(tc, "id", None)
+                    )
+                    if tc_id is None:
+                        continue
+                    for j in range(i + 1, len(messages)):
+                        if (
+                            get_role(messages[j]) == "tool"
+                            and get_tool_call_id(messages[j]) == tc_id
+                        ):
+                            to_delete.add(j)
+
+        # ── Split into surviving and pruned ──
+        pruned = [messages[i] for i in sorted(to_delete)]
+        surviving = [msg for i, msg in enumerate(messages) if i not in to_delete]
+
+        logger.debug(
+            "Reduced messages from %d to %d (pruned %d)",
+            len(messages),
+            len(surviving),
+            len(pruned),
+        )
+
+        # ── Optional summarization ──
+        summary = None
+        if pruned and self.config.summarize_fn is not None:
+            try:
+                summary = self.config.summarize_fn(pruned)
+            except Exception as exc:
+                logger.warning("Summarization failed: %s", exc)
+
+        return ReducerResult(
+            surviving=surviving,
+            pruned=pruned,
+            summary=summary,
+        )
+
+    def as_langgraph_reducer(self):
+        """
+        Return a function compatible with LangGraph's ``Annotated[list, fn]`` pattern.
+
+        The returned function has the signature ``(existing, new) -> merged_list``
+        which is what LangGraph expects for a custom reducer annotated on a
+        state field.
+
+        Usage::
+
+            reducer = MessageReducer(min_messages=10, max_messages=20)
+
+            class MyState(TypedDict):
+                messages: Annotated[list, reducer.as_langgraph_reducer()]
+
+        This is the bridge that lets ``PrunableStateFactory`` keep working
+        unchanged when migrating from ``langgraph-reducer``.
+        """
+
+        def _reduce(existing=None, new=None):
+            if existing is None:
+                existing = []
+            result = self.reduce(existing, new)
+            return result.surviving
+
+        return _reduce
+
+    def __repr__(self) -> str:
+        return (
+            f"MessageReducer(min_messages={self.config.min_messages}, "
+            f"max_messages={self.config.max_messages}, "
+            f"preserve_first={self.config.preserve_first}, "
+            f"cascade_tool_messages={self.config.cascade_tool_messages})"
+        )

agentstate_reducer-0.1.0.dist-info/METADATA

@@ -0,0 +1,131 @@
+Metadata-Version: 2.4
+Name: agentstate-reducer
+Version: 0.1.0
+Summary: Framework-agnostic message reducer for AI agent state management. Works with LangGraph, CrewAI, and plain dicts.
+Project-URL: Homepage, https://github.com/skamalj/agentstate-reducer
+Project-URL: Repository, https://github.com/skamalj/agentstate-reducer
+Author-email: Kamal <skamalj@github.com>
+License-Expression: MIT
+Classifier: License :: OSI Approved :: MIT License
+Classifier: Programming Language :: Python :: 3
+Classifier: Topic :: Scientific/Engineering :: Artificial Intelligence
+Requires-Python: >=3.10
+Provides-Extra: dev
+Requires-Dist: langchain-core>=0.2; extra == 'dev'
+Requires-Dist: pytest>=7.0; extra == 'dev'
+Description-Content-Type: text/markdown
+
+# 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`
+- **System message preservation**: Index 0 (system prompt) is never pruned
+- **ToolMessage cascade**: When an AI message is pruned, linked ToolMessages are pruned too
+- **Optional summarization**: Get a callback with pruned messages to generate summaries
+
+## Install
+
+```bash
+pip install agentstate-reducer
+```
+
+**Zero dependencies.** Works with Python 3.10+.
+
+## Quick Start
+
+```python
+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 with `PrunableStateFactory` or directly as a state annotation:
+
+```python
+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()]
+```
+
+## CrewAI Integration
+
+Use with a CosmosDB persistence backend:
+
+```python
+from agentstate_reducer import MessageReducer
+
+reducer = MessageReducer(min_messages=10, max_messages=20)
+# Pass to your persistence class which calls reducer.reduce()
+# on save_state() to prune before storing
+```
+
+## Summarization
+
+```python
+from agentstate_reducer import MessageReducer, ReducerConfig
+
+def summarize(pruned_messages):
+    # Call your LLM to summarize what was removed
+    return f"Summary of {len(pruned_messages)} 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 messages"
+```
+
+## API
+
+### `MessageReducer(min_messages, max_messages, *, config)`
+
+| Param | Default | Description |
+|---|---|---|
+| `min_messages` | `0` | Messages to retain after pruning |
+| `max_messages` | `None` | Threshold to trigger pruning (`None` = never prune) |
+| `config` | `None` | `ReducerConfig` object (overrides other params) |
+
+### `reducer.reduce(existing, new) -> ReducerResult`
+
+| Field | Type | Description |
+|---|---|---|
+| `surviving` | `list` | Messages that remain |
+| `pruned` | `list` | Messages that were removed |
+| `summary` | `str \| None` | Summary from `summarize_fn` if configured |
+
+### `reducer.as_langgraph_reducer() -> Callable`
+
+Returns a function with signature `(existing, new) -> list` for use with LangGraph's `Annotated[list, fn]` pattern.

agentstate_reducer-0.1.0.dist-info/RECORD

@@ -0,0 +1,7 @@
+agentstate_reducer/__init__.py,sha256=ZG4H-nc_p9ZLYaQ8Hw8kHMROsTZJDIkNNyGztR9z1kI,307
+agentstate_reducer/adapters.py,sha256=oFkfCKV5PVgJamMbWDcfRb9h2p6QFNYWnj2ydVecgys,3406
+agentstate_reducer/models.py,sha256=CmIJbmvzSh25p2OoZxFF_TP2fgHAIOdJQ-vCTcoUW-M,1951
+agentstate_reducer/reducer.py,sha256=Ro9h24WYvpyAma1Ui45pFvKmvLPE72LBYmOPcQkltAs,7166
+agentstate_reducer-0.1.0.dist-info/METADATA,sha256=gR1DO4C9Dfo7jwAgBknbDpyhf4vvmu8sTjAX330llps,4056
+agentstate_reducer-0.1.0.dist-info/WHEEL,sha256=mffPy8wBnZQn2VnJUU5jE99KsxaSfiyMHV9Yt0aLVxs,87
+agentstate_reducer-0.1.0.dist-info/RECORD,,

agentstate_reducer-0.1.0.dist-info/WHEEL

@@ -0,0 +1,4 @@
+Wheel-Version: 1.0
+Generator: hatchling 1.30.1
+Root-Is-Purelib: true
+Tag: py3-none-any