diff --git a/src/lfx/src/lfx/memory/__init__.py b/src/lfx/src/lfx/memory/__init__.py index cf798e67d2..2ea2151a3d 100644 --- a/src/lfx/src/lfx/memory/__init__.py +++ b/src/lfx/src/lfx/memory/__init__.py @@ -1,60 +1,77 @@ -"""Memory management for lfx with dynamic loading. +"""Memory management for lfx with dynamic dispatch. -This module automatically chooses between the full langflow implementation -(when available) and the lfx implementation (when standalone). +Routes memory operations to either the full langflow implementation (when +langflow is installed AND a real database service is registered) or the lfx +stub implementation (standalone / noop DB). + +Dispatch is evaluated at call time, not import time, because the database +service is typically registered *after* this module is first imported (e.g., +from Component class definitions loaded before graph setup). An import-time +decision can't distinguish "langflow is importable" from "a real DB is wired", +and picking the langflow backend with a NoopDatabaseService yields silent +no-op inserts followed by spurious "Message with id X not found" errors on +update. """ -from lfx.utils.langflow_utils import has_langflow_memory +from __future__ import annotations + +from typing import Any + +from lfx.utils.langflow_utils import has_langflow_db_backend + + +def _impl(): + if has_langflow_db_backend(): + from langflow import memory as impl + else: + from lfx.memory import stubs as impl + return impl + + +def aadd_messages(*args: Any, **kwargs: Any): + return _impl().aadd_messages(*args, **kwargs) + + +def aadd_messagetables(*args: Any, **kwargs: Any): + return _impl().aadd_messagetables(*args, **kwargs) + + +def add_messages(*args: Any, **kwargs: Any): + return _impl().add_messages(*args, **kwargs) + + +def adelete_messages(*args: Any, **kwargs: Any): + return _impl().adelete_messages(*args, **kwargs) + + +def aget_messages(*args: Any, **kwargs: Any): + return _impl().aget_messages(*args, **kwargs) + + +def astore_message(*args: Any, **kwargs: Any): + return _impl().astore_message(*args, **kwargs) + + +def aupdate_messages(*args: Any, **kwargs: Any): + return _impl().aupdate_messages(*args, **kwargs) + + +def delete_message(*args: Any, **kwargs: Any): + return _impl().delete_message(*args, **kwargs) + + +def delete_messages(*args: Any, **kwargs: Any): + return _impl().delete_messages(*args, **kwargs) + + +def get_messages(*args: Any, **kwargs: Any): + return _impl().get_messages(*args, **kwargs) + + +def store_message(*args: Any, **kwargs: Any): + return _impl().store_message(*args, **kwargs) -# Import the appropriate implementation -if has_langflow_memory(): - try: - # Import full langflow implementation - from langflow.memory import ( - aadd_messages, - aadd_messagetables, - add_messages, - adelete_messages, - aget_messages, - astore_message, - aupdate_messages, - delete_message, - delete_messages, - get_messages, - store_message, - ) - except ImportError: - # Fallback to lfx implementation if langflow import fails - from lfx.memory.stubs import ( - aadd_messages, - aadd_messagetables, - add_messages, - adelete_messages, - aget_messages, - astore_message, - aupdate_messages, - delete_message, - delete_messages, - get_messages, - store_message, - ) -else: - # Use lfx implementation - from lfx.memory.stubs import ( - aadd_messages, - aadd_messagetables, - add_messages, - adelete_messages, - aget_messages, - astore_message, - aupdate_messages, - delete_message, - delete_messages, - get_messages, - store_message, - ) -# Export the available functions __all__ = [ "aadd_messages", "aadd_messagetables", diff --git a/src/lfx/src/lfx/utils/langflow_utils.py b/src/lfx/src/lfx/utils/langflow_utils.py index 88c8824518..4db2d82676 100644 --- a/src/lfx/src/lfx/utils/langflow_utils.py +++ b/src/lfx/src/lfx/utils/langflow_utils.py @@ -50,3 +50,22 @@ def has_langflow_memory(): _LangflowModule.set_available(is_langflow_available) return is_langflow_available + + +def has_langflow_db_backend() -> bool: + """Return True iff langflow-backed memory calls have a real DB to hit. + + Requires both langflow to be importable AND the registered database + service to be a non-noop implementation. Evaluated on every call because + the database service is typically registered *after* this module is first + imported (e.g., from Component class definitions loaded before graph setup). + """ + if not has_langflow_memory(): + return False + from lfx.services.database.service import NoopDatabaseService + from lfx.services.deps import get_db_service + + try: + return not isinstance(get_db_service(), NoopDatabaseService) + except Exception: # noqa: BLE001 + return False diff --git a/src/lfx/tests/unit/memory/test_memory_dispatch.py b/src/lfx/tests/unit/memory/test_memory_dispatch.py new file mode 100644 index 0000000000..7a61017e08 --- /dev/null +++ b/src/lfx/tests/unit/memory/test_memory_dispatch.py @@ -0,0 +1,118 @@ +"""Regression tests for lfx.memory runtime dispatch. + +Original bug: when langflow was installed alongside lfx but `lfx run` had +only a NoopDatabaseService registered, `lfx.memory` bound at import time to +`langflow.memory` (because the `langflow` package was importable). The +langflow-backed `aupdate_messages` then called `session.get(...)` on a +NoopSession, which always returns `None`, raising spurious +"Message with id X not found" errors mid-stream. +""" + +from __future__ import annotations + +import uuid + +import pytest +from lfx.services.database.service import NoopDatabaseService +from lfx.utils.langflow_utils import has_langflow_db_backend + + +class _FakeRealDbService: + """Stand-in for any non-noop DatabaseService implementation.""" + + +class TestHasLangflowDbBackend: + def test_returns_false_when_langflow_not_importable(self, monkeypatch): + monkeypatch.setattr("lfx.utils.langflow_utils.has_langflow_memory", lambda: False) + assert has_langflow_db_backend() is False + + def test_returns_false_with_noop_db_service(self, monkeypatch): + monkeypatch.setattr("lfx.utils.langflow_utils.has_langflow_memory", lambda: True) + monkeypatch.setattr("lfx.services.deps.get_db_service", lambda: NoopDatabaseService()) + assert has_langflow_db_backend() is False + + def test_returns_true_with_real_db_service(self, monkeypatch): + monkeypatch.setattr("lfx.utils.langflow_utils.has_langflow_memory", lambda: True) + monkeypatch.setattr("lfx.services.deps.get_db_service", lambda: _FakeRealDbService()) + assert has_langflow_db_backend() is True + + def test_returns_false_when_get_db_service_raises(self, monkeypatch): + monkeypatch.setattr("lfx.utils.langflow_utils.has_langflow_memory", lambda: True) + + def boom(): + msg = "service manager exploded" + raise RuntimeError(msg) + + monkeypatch.setattr("lfx.services.deps.get_db_service", boom) + assert has_langflow_db_backend() is False + + +class TestMemoryDispatch: + def test_dispatches_to_stubs_when_no_real_db(self, monkeypatch): + import lfx.memory as memory_mod + from lfx.memory import stubs + + monkeypatch.setattr("lfx.memory.has_langflow_db_backend", lambda: False) + assert memory_mod._impl() is stubs + + def test_dispatches_to_langflow_when_real_db(self, monkeypatch): + pytest.importorskip("langflow.memory") + import langflow.memory as langflow_memory + import lfx.memory as memory_mod + + monkeypatch.setattr("lfx.memory.has_langflow_db_backend", lambda: True) + assert memory_mod._impl() is langflow_memory + + def test_dispatch_is_evaluated_per_call(self, monkeypatch): + """Dispatch must read the backend state each call, not cache at import. + + The database service is often registered *after* lfx.memory is imported + (components load first, services register during graph setup), so + memoizing the dispatcher would bind to whatever state existed at + component-module load time. + """ + import lfx.memory as memory_mod + from lfx.memory import stubs + + state = {"real": False} + monkeypatch.setattr("lfx.memory.has_langflow_db_backend", lambda: state["real"]) + + assert memory_mod._impl() is stubs + state["real"] = True + pytest.importorskip("langflow.memory") + import langflow.memory as langflow_memory + + assert memory_mod._impl() is langflow_memory + + +class TestAupdateMessagesRegression: + """Direct regression for the original 'Message with id X not found' crash.""" + + @pytest.mark.asyncio + async def test_aupdate_messages_does_not_raise_against_noop_session(self, monkeypatch): + """Regression: route to stubs (no-op) instead of raising via langflow.memory. + + With langflow importable but only a NoopDatabaseService registered, + aupdate_messages must route to stubs and succeed silently rather than + trigger langflow.memory's strict existence check against NoopSession. + """ + try: + from langflow.schema.message import Message + except ImportError: + from lfx.schema.message import Message + + # Force the noop-DB branch even if a real DB happens to be registered in + # this test environment. + monkeypatch.setattr("lfx.services.deps.get_db_service", lambda: NoopDatabaseService()) + + from lfx.memory import aupdate_messages + + msg = Message( + id=str(uuid.uuid4()), + text="hello", + sender="AI", + sender_name="Test", + session_id="test-session", + ) + result = await aupdate_messages(msg) + assert isinstance(result, list)