From 8181429f77fcc518fadcd99e2fddcbf110581108 Mon Sep 17 00:00:00 2001 From: HoangDucBach Date: Sat, 15 Aug 2026 09:28:44 +0700 Subject: [PATCH] fix(python-sdk): recall memories in sync LangChain calls inside a running loop MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit with_memwal_langchain() advertises recall-before-call, but the sync _generate() wrapper appended the untouched message list whenever an event loop was already running — the normal case in notebooks and async application hosts. The LLM call still succeeded, so callers got a plain answer with no memory context and no warning: Walrus Memory looked connected while recall never ran. Route injection through the module's existing _run_blocking() helper, which moves the coroutine onto a worker thread, so the sync LangChain wrapper now behaves like the sync OpenAI one instead of silently degrading. Fixes #607 --- .../python-sdk-memwal/memwal/middleware.py | 22 +++--- .../tests/test_middleware.py | 78 +++++++++++++++++++ 2 files changed, 87 insertions(+), 13 deletions(-) diff --git a/packages/python-sdk-memwal/memwal/middleware.py b/packages/python-sdk-memwal/memwal/middleware.py index 8121a4a2b..903cfdebf 100644 --- a/packages/python-sdk-memwal/memwal/middleware.py +++ b/packages/python-sdk-memwal/memwal/middleware.py @@ -403,24 +403,20 @@ async def patched_agenerate( return result + def _run_memwal(coro_factory: Callable[[], Any]) -> Any: + # Keep httpx clients bound to the short-lived loop that uses them. + return _run_blocking(lambda: _with_fresh_http_client(memwal, coro_factory())) + def patched_generate( messages: List[List[BaseMessage]], *args: Any, **kwargs: Any ) -> ChatResult: - # For sync generate, we inject memories synchronously via asyncio.run - import asyncio - + # Inject via _run_blocking so recall still runs when a loop is already + # running (notebooks, async hosts) — the same helper the sync OpenAI + # wrapper uses. Passing the messages through untouched there made + # Walrus Memory look connected while recall silently never ran. enriched = [] for msg_list in messages: - try: - loop = asyncio.get_running_loop() - except RuntimeError: - loop = None - - if loop is not None and loop.is_running(): - # Already in async context -- cannot use asyncio.run - enriched.append(msg_list) - else: - enriched.append(asyncio.run(_inject_memories(msg_list))) + enriched.append(_run_memwal(lambda: _inject_memories(msg_list))) result = original_generate(enriched, *args, **kwargs) diff --git a/packages/python-sdk-memwal/tests/test_middleware.py b/packages/python-sdk-memwal/tests/test_middleware.py index ef11c80a8..47bad26cc 100644 --- a/packages/python-sdk-memwal/tests/test_middleware.py +++ b/packages/python-sdk-memwal/tests/test_middleware.py @@ -447,6 +447,84 @@ async def test_no_user_message_no_recall(self) -> None: await smart_llm._agenerate([[SystemMessage("only system")]]) assert not recall_route.called + def _capturing_generate(self, captured: list): + """Replacement for llm._generate that records the batch it received.""" + from langchain_core.messages import AIMessage + from langchain_core.outputs import ChatGeneration, ChatResult + + def _generate(messages_batch, *a, **kw): + captured.extend(messages_batch) + return ChatResult( + generations=[ChatGeneration(message=AIMessage(content="ok"))] + ) + + return _generate + + @respx.mock + def test_sync_generate_injects_memories(self) -> None: + """The plain sync path (no running loop) recalls and injects.""" + _mock_seal_session_prereqs() + from langchain_core.messages import HumanMessage + + llm = self._make_llm() + captured: list = [] + llm._generate = self._capturing_generate(captured) + + recall_route = respx.post(_RECALL_URL).mock( + return_value=_mock_recall([ + {"blob_id": "b1", "text": "User loves coffee", "distance": 0.05} + ]) + ) + + smart_llm = with_memwal_langchain( + llm, key=_KEY_HEX, account_id=_ACCOUNT_ID, server_url=_SERVER, auto_save=False + ) + smart_llm._generate([[HumanMessage("What do I drink?")]]) + + assert recall_route.called + assert any("User loves coffee" in m.content for m in captured[0]) + + @respx.mock + async def test_sync_generate_injects_memories_inside_running_loop(self) -> None: + """GH #607: sync _generate must still recall when a loop is already + running (notebooks, async hosts). + + It used to append the untouched msg_list in that branch, so callers got + a normal LLM answer with no memory context and no warning -- Walrus + Memory looked connected while recall never ran. + """ + _mock_seal_session_prereqs() + from langchain_core.messages import HumanMessage, SystemMessage + + llm = self._make_llm() + captured: list = [] + llm._generate = self._capturing_generate(captured) + + recall_route = respx.post(_RECALL_URL).mock( + return_value=_mock_recall([ + {"blob_id": "b1", "text": "User loves coffee", "distance": 0.05} + ]) + ) + + smart_llm = with_memwal_langchain( + llm, key=_KEY_HEX, account_id=_ACCOUNT_ID, server_url=_SERVER, auto_save=False + ) + + # An async test body is itself a running loop -- the exact condition + # that used to disable injection. + assert asyncio.get_running_loop().is_running() + smart_llm._generate([[HumanMessage("What do I drink?")]]) + + assert recall_route.called, "recall was skipped inside a running loop" + assert len(captured) == 1 + injected = captured[0] + assert len(injected) == 3, "guard + memory message were not injected" + assert any( + isinstance(m, HumanMessage) and "User loves coffee" in m.content + for m in injected + ) + assert any(isinstance(m, SystemMessage) for m in injected) + def test_wraps_a_real_pydantic_backed_chat_model(self) -> None: """with_memwal_langchain must work on a real BaseChatModel, not just a MagicMock. LangChain chat models are Pydantic models that reject