memwal 0.1.11.dev6__tar.gz → 0.1.11.dev7__tar.gz

{memwal-0.1.11.dev6 → memwal-0.1.11.dev7}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.5
 Name: memwal
-Version: 0.1.11.dev6
+Version: 0.1.11.dev7
 Summary: Python SDK for Walrus Memory — Privacy-first AI memory with Ed25519 signing
 Project-URL: Homepage, https://memory.walrus.xyz
 Project-URL: Documentation, https://memory.walrus.xyz

{memwal-0.1.11.dev6 → memwal-0.1.11.dev7}/memwal/__init__.py

@@ -128,4 +128,4 @@ __all__ = [
     "RecallManualResult",
 ]
 
-__version__ = "0.1.11.dev6"
+__version__ = "0.1.11.dev7"

{memwal-0.1.11.dev6 → memwal-0.1.11.dev7}/memwal/middleware.py

@@ -188,7 +188,7 @@ class _PendingSaves:
 
             try:
                 asyncio.run(_drain(list(self._tasks)))
-            except RuntimeError:
+            except (RuntimeError, ValueError):
                 logger.warning(
                     "Walrus Memory flush_sync() could not await %d pending "
                     "task(s) bound to a different event loop (mixing sync "
@@ -366,12 +366,32 @@ def with_memwal_langchain(
                     last_human_idx = i
                     break
 
-            guard_msg = SystemMessage(content=_UNTRUSTED_MEMORY_SYSTEM_INSTRUCTION)
             memory_msg = HumanMessage(content=memory_context)
             if last_human_idx > 0:
-                result[last_human_idx:last_human_idx] = [guard_msg, memory_msg]
+                result.insert(last_human_idx, memory_msg)
             else:
-                result[0:0] = [guard_msg, memory_msg]
+                result.insert(0, memory_msg)
+
+            if result and isinstance(result[0], SystemMessage):
+                existing_content = result[0].content
+                if isinstance(existing_content, list):
+                    new_content = list(existing_content) + [
+                        {"type": "text", "text": f"\n\n{_UNTRUSTED_MEMORY_SYSTEM_INSTRUCTION}"}
+                    ]
+                else:
+                    text_str = str(existing_content or "")
+                    new_content = f"{text_str}\n\n{_UNTRUSTED_MEMORY_SYSTEM_INSTRUCTION}".strip()
+                if hasattr(result[0], "model_copy"):
+                    result[0] = result[0].model_copy(update={"content": new_content})
+                else:
+                    result[0] = SystemMessage(
+                        content=new_content,
+                        additional_kwargs=getattr(result[0], "additional_kwargs", {}),
+                        id=getattr(result[0], "id", None),
+                        name=getattr(result[0], "name", None),
+                    )
+            else:
+                result.insert(0, SystemMessage(content=_UNTRUSTED_MEMORY_SYSTEM_INSTRUCTION))
 
             return result
         except Exception as e:
@@ -625,21 +645,36 @@ def _inject_openai_memory(
     messages: List[Dict[str, Any]],
     memory_context: str,
 ) -> List[Dict[str, Any]]:
-    """Insert a fixed system guard plus memory content in a user message."""
+    """Insert a fixed system guard at index 0 and memory content in a user message."""
+    result = [dict(m) if isinstance(m, dict) else m for m in messages]
     last_user_idx = -1
-    for i in range(len(messages) - 1, -1, -1):
-        if isinstance(messages[i], dict) and messages[i].get("role") == "user":
+    for i in range(len(result) - 1, -1, -1):
+        if isinstance(result[i], dict) and result[i].get("role") == "user":
             last_user_idx = i
             break
 
-    memory_messages: List[Dict[str, Any]] = [
-        {"role": "system", "content": _UNTRUSTED_MEMORY_SYSTEM_INSTRUCTION},
-        {"role": "user", "content": memory_context},
-    ]
-
+    memory_message = {"role": "user", "content": memory_context}
     if last_user_idx > 0:
-        messages[last_user_idx:last_user_idx] = memory_messages
+        result.insert(last_user_idx, memory_message)
+    else:
+        result.insert(0, memory_message)
+
+    if result and isinstance(result[0], dict) and result[0].get("role") == "system":
+        existing_content = result[0].get("content")
+        if isinstance(existing_content, list):
+            result[0] = {
+                **result[0],
+                "content": existing_content + [
+                    {"type": "text", "text": f"\n\n{_UNTRUSTED_MEMORY_SYSTEM_INSTRUCTION}"}
+                ],
+            }
+        else:
+            text_str = str(existing_content or "")
+            result[0] = {
+                **result[0],
+                "content": f"{text_str}\n\n{_UNTRUSTED_MEMORY_SYSTEM_INSTRUCTION}".strip(),
+            }
     else:
-        messages[0:0] = memory_messages
+        result.insert(0, {"role": "system", "content": _UNTRUSTED_MEMORY_SYSTEM_INSTRUCTION})
 
-    return messages
+    return result

{memwal-0.1.11.dev6 → memwal-0.1.11.dev7}/pyproject.toml

@@ -4,7 +4,7 @@ build-backend = "hatchling.build"
 
 [project]
 name = "memwal"
-version = "0.1.11.dev6"
+version = "0.1.11.dev7"
 description = "Python SDK for Walrus Memory — Privacy-first AI memory with Ed25519 signing"
 readme = "README.md"
 license = "MIT"

{memwal-0.1.11.dev6 → memwal-0.1.11.dev7}/tests/test_middleware.py

@@ -240,27 +240,53 @@ class TestFormatMemories:
 class TestInjectOpenAIMemory:
     """Tests for _inject_openai_memory."""
 
-    def test_inserts_before_last_user_message(self) -> None:
+    def test_folds_guard_into_system_message_and_inserts_memory_before_user_message(self) -> None:
         messages = [
             {"role": "system", "content": "You are helpful."},
             {"role": "user", "content": "question"},
         ]
         result = _inject_openai_memory(messages, "memory context")
-        # Should be: existing-system | fixed-guard | memory-user | user
+        # Should be: existing-system-with-guard | memory-user | user
+        assert len(result) == 3
         assert result[0]["role"] == "system"
-        assert result[1]["role"] == "system"
-        assert "untrusted data" in result[1]["content"]
-        assert result[2] == {"role": "user", "content": "memory context"}
-        assert result[3]["role"] == "user"
+        assert "You are helpful." in result[0]["content"]
+        assert "untrusted data" in result[0]["content"]
+        assert result[1] == {"role": "user", "content": "memory context"}
+        assert result[2]["role"] == "user"
+        assert result[2]["content"] == "question"
+
+    def test_multi_turn_never_injects_system_message_after_index_0(self) -> None:
+        messages = [
+            {"role": "system", "content": "You are helpful."},
+            {"role": "user", "content": "first"},
+            {"role": "assistant", "content": "reply"},
+            {"role": "user", "content": "second"},
+        ]
+        result = _inject_openai_memory(messages, "recalled fact")
+        # System message only at index 0
+        assert result[0]["role"] == "system"
+        assert "You are helpful." in result[0]["content"]
+        assert "untrusted data" in result[0]["content"]
+
+        # No system message after index 0
+        assert not any(m.get("role") == "system" for m in result[1:])
+
+        # Memory is injected as user message right before "second"
+        assert result[1] == {"role": "user", "content": "first"}
+        assert result[2] == {"role": "assistant", "content": "reply"}
+        assert result[3] == {"role": "user", "content": "recalled fact"}
+        assert result[4] == {"role": "user", "content": "second"}
 
     def test_inserts_at_start_when_only_user_message(self) -> None:
         messages = [{"role": "user", "content": "question"}]
         result = _inject_openai_memory(messages, "context")
-        # last_user_idx == 0, so prepend guard + memory-data messages
+        # When no system message exists, prepends guard system message at index 0
         assert result[0]["role"] == "system"
+        assert "untrusted data" in result[0]["content"]
         assert result[1]["role"] == "user"
         assert result[1]["content"] == "context"
         assert result[2]["role"] == "user"
+        assert result[2]["content"] == "question"
 
     def test_does_not_mutate_original(self) -> None:
         original = [{"role": "user", "content": "q"}]
@@ -335,6 +361,102 @@ class TestWithMemWalLangChain:
         assert "User loves coffee" in human_msgs[0].content
         assert "BEGIN_UNTRUSTED_WALRUS_MEMORY_" in human_msgs[0].content
 
+    @respx.mock
+    async def test_multi_turn_never_injects_system_message_after_index_0(self) -> None:
+        """In a multi-turn conversation, SystemMessage stays strictly at index 0."""
+        _mock_seal_session_prereqs()
+        from langchain_core.messages import AIMessage, HumanMessage, SystemMessage
+        from langchain_core.outputs import ChatGeneration, ChatResult
+
+        llm = self._make_llm()
+        captured: list = []
+
+        async def _capture_agenerate(messages_batch, *a, **kw):
+            captured.extend(messages_batch)
+            return ChatResult(
+                generations=[ChatGeneration(message=AIMessage(content="ok"))]
+            )
+
+        llm._agenerate = _capture_agenerate
+
+        respx.post(_RECALL_URL).mock(
+            return_value=_mock_recall([
+                {"blob_id": "b1", "text": "User loves coffee", "distance": 0.05}
+            ])
+        )
+        respx.post(_ANALYZE_URL).mock(return_value=_mock_analyze())
+
+        smart_llm = with_memwal_langchain(
+            llm, key=_KEY_HEX, account_id=_ACCOUNT_ID, server_url=_SERVER, auto_save=False
+        )
+        messages = [
+            SystemMessage(content="You are helpful."),
+            HumanMessage(content="I like coffee"),
+            AIMessage(content="Noted!"),
+            HumanMessage(content="What do I like?"),
+        ]
+        await smart_llm._agenerate([messages])
+
+        assert len(captured) == 1
+        injected = captured[0]
+        # SystemMessage strictly at index 0, with guard folded in
+        assert isinstance(injected[0], SystemMessage)
+        assert "You are helpful." in injected[0].content
+        assert "untrusted data" in injected[0].content
+
+        # No SystemMessage after index 0
+        assert not any(isinstance(m, SystemMessage) for m in injected[1:])
+
+        # The recalled memory is in a HumanMessage before the final query
+        assert isinstance(injected[3], HumanMessage)
+        assert "User loves coffee" in injected[3].content
+        assert isinstance(injected[4], HumanMessage)
+        assert injected[4].content == "What do I like?"
+
+    @respx.mock
+    async def test_multimodal_system_message_appends_block_and_preserves_metadata(self) -> None:
+        """A list-content SystemMessage appends a text block and retains its metadata."""
+        _mock_seal_session_prereqs()
+        from langchain_core.messages import AIMessage, HumanMessage, SystemMessage
+        from langchain_core.outputs import ChatGeneration, ChatResult
+
+        llm = self._make_llm()
+        captured: list = []
+
+        async def _capture_agenerate(messages_batch, *a, **kw):
+            captured.extend(messages_batch)
+            return ChatResult(
+                generations=[ChatGeneration(message=AIMessage(content="ok"))]
+            )
+
+        llm._agenerate = _capture_agenerate
+
+        respx.post(_RECALL_URL).mock(
+            return_value=_mock_recall([
+                {"blob_id": "b1", "text": "User loves coffee", "distance": 0.05}
+            ])
+        )
+        respx.post(_ANALYZE_URL).mock(return_value=_mock_analyze())
+
+        smart_llm = with_memwal_langchain(
+            llm, key=_KEY_HEX, account_id=_ACCOUNT_ID, server_url=_SERVER, auto_save=False
+        )
+        blocks = [
+            {"type": "text", "text": "Base instructions", "cache_control": {"type": "ephemeral"}}
+        ]
+        sys_msg = SystemMessage(content=blocks, id="sys-123", additional_kwargs={"source": "test"})
+        await smart_llm._agenerate([[sys_msg, HumanMessage("What do I like?")]])
+
+        assert len(captured) == 1
+        injected = captured[0]
+        assert isinstance(injected[0], SystemMessage)
+        assert injected[0].id == "sys-123"
+        assert injected[0].additional_kwargs == {"source": "test"}
+        assert isinstance(injected[0].content, list)
+        assert len(injected[0].content) == 2
+        assert injected[0].content[0] == blocks[0]
+        assert "untrusted data" in injected[0].content[1]["text"]
+
     @respx.mock
     async def test_no_memories_found_no_injection(self) -> None:
         """When no memories match, the message list is passed unchanged."""