zwarm 1.0.0__py3-none-any.whl → 1.1.1__py3-none-any.whl

zwarm/core/compact.py

@@ -16,6 +16,13 @@ from typing import Any
 logger = logging.getLogger(__name__)
 
 
+def _get_attr(obj: Any, key: str, default: Any = None) -> Any:
+    """Get attribute from dict or object (handles both Pydantic models and dicts)."""
+    if isinstance(obj, dict):
+        return obj.get(key, default)
+    return getattr(obj, key, default)
+
+
 @dataclass
 class CompactionResult:
     """Result of a compaction operation."""
@@ -30,16 +37,17 @@ class CompactionResult:
         return self.removed_count > 0
 
 
-def estimate_tokens(messages: list[dict[str, Any]]) -> int:
+def estimate_tokens(messages: list[Any]) -> int:
     """
     Rough token estimate for messages.
 
     Uses ~4 chars per token as a simple heuristic.
     This is intentionally conservative.
+    Handles both dict messages and Pydantic model messages.
     """
     total_chars = 0
     for msg in messages:
-        content = msg.get("content", "")
+        content = _get_attr(msg, "content", "")
         if isinstance(content, str):
             total_chars += len(content)
         elif isinstance(content, list):
@@ -50,16 +58,22 @@ def estimate_tokens(messages: list[dict[str, Any]]) -> int:
                     total_chars += len(str(block.get("input", "")))
                 elif isinstance(block, str):
                     total_chars += len(block)
+                else:
+                    # Pydantic model block
+                    total_chars += len(str(_get_attr(block, "text", "")))
+                    total_chars += len(str(_get_attr(block, "input", "")))
 
         # Tool calls add tokens too
-        tool_calls = msg.get("tool_calls", [])
+        tool_calls = _get_attr(msg, "tool_calls", []) or []
         for tc in tool_calls:
-            total_chars += len(str(tc.get("function", {}).get("arguments", "")))
+            func = _get_attr(tc, "function", {}) or {}
+            args = _get_attr(func, "arguments", "") if isinstance(func, dict) else getattr(func, "arguments", "")
+            total_chars += len(str(args))
 
     return total_chars // 4
 
 
-def find_tool_groups(messages: list[dict[str, Any]]) -> list[tuple[int, int]]:
+def find_tool_groups(messages: list[Any]) -> list[tuple[int, int]]:
     """
     Find message index ranges that form tool call groups.
 
@@ -69,6 +83,7 @@ def find_tool_groups(messages: list[dict[str, Any]]) -> list[tuple[int, int]]:
 
     This handles both OpenAI format (role="tool") and Anthropic format
     (role="user" with tool_result content).
+    Also handles Pydantic model messages.
 
     Returns list of (start_idx, end_idx) tuples (inclusive).
     """
@@ -82,15 +97,16 @@ def find_tool_groups(messages: list[dict[str, Any]]) -> list[tuple[int, int]]:
         has_tool_calls = False
 
         # OpenAI format: tool_calls field
-        if msg.get("role") == "assistant" and msg.get("tool_calls"):
+        if _get_attr(msg, "role") == "assistant" and _get_attr(msg, "tool_calls"):
             has_tool_calls = True
 
         # Anthropic format: content blocks with type="tool_use"
-        if msg.get("role") == "assistant":
-            content = msg.get("content", [])
+        if _get_attr(msg, "role") == "assistant":
+            content = _get_attr(msg, "content", [])
             if isinstance(content, list):
                 for block in content:
-                    if isinstance(block, dict) and block.get("type") == "tool_use":
+                    block_type = _get_attr(block, "type", None)
+                    if block_type == "tool_use":
                         has_tool_calls = True
                         break
 
@@ -101,7 +117,7 @@ def find_tool_groups(messages: list[dict[str, Any]]) -> list[tuple[int, int]]:
             # Find all following tool responses
             while j < len(messages):
                 next_msg = messages[j]
-                role = next_msg.get("role", "")
+                role = _get_attr(next_msg, "role", "")
 
                 # OpenAI format: tool role
                 if role == "tool":
@@ -110,10 +126,10 @@ def find_tool_groups(messages: list[dict[str, Any]]) -> list[tuple[int, int]]:
 
                 # Anthropic format: user message with tool_result
                 if role == "user":
-                    content = next_msg.get("content", [])
+                    content = _get_attr(next_msg, "content", [])
                     if isinstance(content, list):
                         has_tool_result = any(
-                            isinstance(b, dict) and b.get("type") == "tool_result"
+                            _get_attr(b, "type", None) == "tool_result"
                             for b in content
                         )
                         if has_tool_result:
@@ -132,7 +148,7 @@ def find_tool_groups(messages: list[dict[str, Any]]) -> list[tuple[int, int]]:
 
 
 def compact_messages(
-    messages: list[dict[str, Any]],
+    messages: list[Any],
     keep_first_n: int = 2,
     keep_last_n: int = 10,
     max_tokens: int | None = None,
@@ -298,7 +314,7 @@ def compact_messages(
 
 
 def should_compact(
-    messages: list[dict[str, Any]],
+    messages: list[Any],
     max_tokens: int,
     threshold_pct: float = 0.85,
 ) -> bool:
@@ -306,6 +322,7 @@ def should_compact(
     Check if messages should be compacted.
 
     Returns True if estimated tokens exceed threshold percentage of max.
+    Handles both dict messages and Pydantic model messages.
     """
     current = estimate_tokens(messages)
     threshold = int(max_tokens * threshold_pct)

zwarm/core/test_compact.py

@@ -264,3 +264,49 @@ class TestEdgeCases:
         result = compact_messages(messages, keep_first_n=2, keep_last_n=2)
         assert not result.was_compacted
         assert result.messages == messages
+
+
+class TestPydanticModelMessages:
+    """Test handling of Pydantic model messages (not just dicts)."""
+
+    def test_estimate_tokens_with_objects(self):
+        """estimate_tokens should handle objects with attributes."""
+        class MockMessage:
+            def __init__(self, role, content):
+                self.role = role
+                self.content = content
+
+        messages = [
+            MockMessage("user", "Hello world"),
+            MockMessage("assistant", "Hi there!"),
+        ]
+        tokens = estimate_tokens(messages)
+        assert tokens > 0
+
+    def test_should_compact_with_objects(self):
+        """should_compact should handle objects with attributes."""
+        class MockMessage:
+            def __init__(self, role, content):
+                self.role = role
+                self.content = content
+
+        messages = [MockMessage("user", "x" * 4000)]
+        # Should not crash
+        result = should_compact(messages, max_tokens=500, threshold_pct=0.85)
+        assert result is True
+
+    def test_find_tool_groups_with_objects(self):
+        """find_tool_groups should handle objects with attributes."""
+        class MockMessage:
+            def __init__(self, role, content=None, tool_calls=None):
+                self.role = role
+                self.content = content
+                self.tool_calls = tool_calls
+
+        messages = [
+            MockMessage("user", "Task"),
+            MockMessage("assistant", "Done"),
+        ]
+        # Should not crash
+        groups = find_tool_groups(messages)
+        assert groups == []

{zwarm-1.0.0.dist-info → zwarm-1.1.1.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: zwarm
-Version: 1.0.0
+Version: 1.1.1
 Summary: Multi-Agent CLI Orchestration Research Platform
 Requires-Python: <3.14,>=3.13
 Requires-Dist: python-dotenv>=1.0.0

{zwarm-1.0.0.dist-info → zwarm-1.1.1.dist-info}/RECORD

@@ -9,12 +9,12 @@ zwarm/adapters/test_codex_mcp.py,sha256=vodQF5VUrM_F1GygUADtXYrA0kvAc7dWpT_Ff-Uo
 zwarm/cli/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 zwarm/cli/main.py,sha256=dU_8XLIYzcxikl_MgT6BmChNbqYq94dJ4Ra5STzm_KI,33109
 zwarm/core/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
-zwarm/core/compact.py,sha256=x3NYHvwDOLjNT7gcOG7mVoL4EDbWsG77rZAbjKZVOgQ,10109
+zwarm/core/compact.py,sha256=Y8C7Gs-5-WOU43WRvQ863Qzd5xtuEqR6Aw3r2p8_-i8,10907
 zwarm/core/config.py,sha256=H8XFsWeEehZnUsMf7nsM_YqjljFYaC43DLs8Xw4sOuQ,10360
 zwarm/core/environment.py,sha256=HVDpDZEpDSfyh9-wHZMzMKVUPKvioBkPVWeiME2JmFo,5435
 zwarm/core/models.py,sha256=PrC3okRBVJxISUa1Fax4KkagqLT6Xub-kTxC9drN0sY,10083
 zwarm/core/state.py,sha256=wMryIvXP-VDLh2b76A5taeL_9dm5k4jk4HnvHWgLqGE,7658
-zwarm/core/test_compact.py,sha256=CbVBkTaLV2VDEn51om0AbY8RgWM2czhxZQQt_U_NggU,10305
+zwarm/core/test_compact.py,sha256=WSdjCB5t4YMcknsrkmJIUsVOPY28s4y9GnDmu3Z4BFw,11878
 zwarm/core/test_config.py,sha256=26ozyiFOdjFF2c9Q-HDfFM6GOLfgw_5FZ55nTDMNYA8,4888
 zwarm/core/test_models.py,sha256=sWTIhMZvuLP5AooGR6y8OR2EyWydqVfhmGrE7NPBBnk,8450
 zwarm/prompts/__init__.py,sha256=FiaIOniLrIyfD3_osxT6I7FfyKjtctbf8jNs5QTPs_s,213
@@ -27,7 +27,7 @@ zwarm/watchers/builtin.py,sha256=52hyRREYYDsSuG-YKElXViSTyMmGySZaFreHc0pz-A4,124
 zwarm/watchers/manager.py,sha256=XZjBVeHjgCUlkTUeHqdvBvHoBC862U1ik0fG6nlRGog,5587
 zwarm/watchers/registry.py,sha256=A9iBIVIFNtO7KPX0kLpUaP8dAK7ozqWLA44ocJGnOw4,1219
 zwarm/watchers/test_watchers.py,sha256=zOsxumBqKfR5ZVGxrNlxz6KcWjkcdp0QhW9WB0_20zM,7855
-zwarm-1.0.0.dist-info/METADATA,sha256=aIyHVKHhTe8oo8thqCpMMErJo7S2ILWxvEYF3LaJYv8,14995
-zwarm-1.0.0.dist-info/WHEEL,sha256=WLgqFyCfm_KASv4WHyYy0P3pM_m7J5L9k2skdKLirC8,87
-zwarm-1.0.0.dist-info/entry_points.txt,sha256=u0OXq4q8d3yJ3EkUXwZfkS-Y8Lcy0F8cWrcQfoRxM6Q,46
-zwarm-1.0.0.dist-info/RECORD,,
+zwarm-1.1.1.dist-info/METADATA,sha256=9GoAYbxRjMSdbHIwrZPS_zdrrDSz7-Avnwy_8AqNEBM,14995
+zwarm-1.1.1.dist-info/WHEEL,sha256=WLgqFyCfm_KASv4WHyYy0P3pM_m7J5L9k2skdKLirC8,87
+zwarm-1.1.1.dist-info/entry_points.txt,sha256=u0OXq4q8d3yJ3EkUXwZfkS-Y8Lcy0F8cWrcQfoRxM6Q,46
+zwarm-1.1.1.dist-info/RECORD,,