diff --git a/src/backend/langflow/interface/custom_lists.py b/src/backend/langflow/interface/custom_lists.py index 59f163338..32d59619b 100644 --- a/src/backend/langflow/interface/custom_lists.py +++ b/src/backend/langflow/interface/custom_lists.py @@ -6,15 +6,8 @@ from langchain.agents import agent_toolkits from langchain.chat_models import ChatOpenAI ## Memory -# from langchain.memory.buffer_window import ConversationBufferWindowMemory -# from langchain.memory.chat_memory import ChatMessageHistory -# from langchain.memory.combined import CombinedMemory -# from langchain.memory.entity import ConversationEntityMemory -# from langchain.memory.kg import ConversationKGMemory -# from langchain.memory.readonly import ReadOnlySharedMemory -# from langchain.memory.simple import SimpleMemory -# from langchain.memory.summary import ConversationSummaryMemory -# from langchain.memory.summary_buffer import ConversationSummaryBufferMemory +from langchain import memory + ## Document Loaders from langchain.document_loaders import ( AirbyteJSONLoader, @@ -104,23 +97,6 @@ llm_type_to_cls_dict = llms.type_to_cls_dict llm_type_to_cls_dict["openai-chat"] = ChatOpenAI # type: ignore -## Memory - -memory_type_to_cls_dict: dict[str, Any] = { - # "CombinedMemory": CombinedMemory, - # "ConversationBufferWindowMemory": ConversationBufferWindowMemory, - # "ConversationBufferMemory": ConversationBufferMemory, - # "SimpleMemory": SimpleMemory, - # "ConversationSummaryBufferMemory": ConversationSummaryBufferMemory, - # "ConversationKGMemory": ConversationKGMemory, - # "ConversationEntityMemory": ConversationEntityMemory, - # "ConversationSummaryMemory": ConversationSummaryMemory, - # "ChatMessageHistory": ChatMessageHistory, - # "ConversationStringBufferMemory": ConversationStringBufferMemory, - # "ReadOnlySharedMemory": ReadOnlySharedMemory, -} - - ## Chain # from langchain.chains.loading import type_to_loader_dict # from langchain.chains.conversation.base import ConversationChain @@ -142,6 +118,14 @@ toolkit_type_to_cls_dict: dict[str, Any] = { if not toolkit_name.islower() } +## Memory + + +memory_type_to_cls_dict: dict[str, Any] = { + memory_name: import_class(f"langchain.memory.{memory_name}") + for memory_name in memory.__all__ +} + wrapper_type_to_cls_dict: dict[str, Any] = { wrapper.__name__: wrapper for wrapper in [requests.RequestsWrapper] diff --git a/src/backend/langflow/interface/importing/utils.py b/src/backend/langflow/interface/importing/utils.py index f054ddc26..b8717f804 100644 --- a/src/backend/langflow/interface/importing/utils.py +++ b/src/backend/langflow/interface/importing/utils.py @@ -36,10 +36,16 @@ def import_by_type(_type: str, name: str) -> Any: "chains": import_chain, "toolkits": import_toolkit, "wrappers": import_wrapper, + "memory": import_memory, } return func_dict[_type](name) +def import_memory(memory: str) -> Any: + """Import memory from memory name""" + return import_module(f"from langchain.memory import {memory}") + + def import_class(class_path: str) -> Any: """Import class from class path""" module_path, class_name = class_path.rsplit(".", 1)