adding pg support for external message persistance

This commit is contained in:
zhenjianpeng 2023-06-28 15:20:06 +08:00
commit d157fffee3
4 changed files with 46 additions and 0 deletions

View file

@ -69,6 +69,7 @@ memories:
- ConversationBufferMemory - ConversationBufferMemory
- ConversationSummaryMemory - ConversationSummaryMemory
- ConversationKGMemory - ConversationKGMemory
- PostgresChatMessageHistory
prompts: prompts:
- PromptTemplate - PromptTemplate
- FewShotPromptTemplate - FewShotPromptTemplate

View file

@ -21,6 +21,9 @@ CUSTOM_NODES = {
"utilities": { "utilities": {
"SQLDatabase": frontend_node.agents.SQLDatabaseNode(), "SQLDatabase": frontend_node.agents.SQLDatabaseNode(),
}, },
"memories": {
"PostgresChatMessageHistory": frontend_node.memories.PostgresChatMessageHistoryFrontendNode(),
},
"chains": { "chains": {
"SeriesCharacterChain": frontend_node.chains.SeriesCharacterChainNode(), "SeriesCharacterChain": frontend_node.chains.SeriesCharacterChainNode(),
"TimeTravelGuideChain": frontend_node.chains.TimeTravelGuideChainNode(), "TimeTravelGuideChain": frontend_node.chains.TimeTravelGuideChainNode(),

View file

@ -225,6 +225,7 @@ def load_agent_executor(agent_class: type[agent_module.Agent], params, **kwargs)
"""Load agent executor from agent class, tools and chain""" """Load agent executor from agent class, tools and chain"""
allowed_tools: Sequence[BaseTool] = params.get("allowed_tools", []) allowed_tools: Sequence[BaseTool] = params.get("allowed_tools", [])
llm_chain = params["llm_chain"] llm_chain = params["llm_chain"]
memory = params["memory"]
# if allowed_tools is not a list or set, make it a list # if allowed_tools is not a list or set, make it a list
if not isinstance(allowed_tools, (list, set)) and isinstance( if not isinstance(allowed_tools, (list, set)) and isinstance(
allowed_tools, BaseTool allowed_tools, BaseTool
@ -237,6 +238,7 @@ def load_agent_executor(agent_class: type[agent_module.Agent], params, **kwargs)
return AgentExecutor.from_agent_and_tools( return AgentExecutor.from_agent_and_tools(
agent=agent, agent=agent,
tools=allowed_tools, tools=allowed_tools,
memory=memory,
**kwargs, **kwargs,
) )

View file

@ -2,6 +2,7 @@ from typing import Optional
from langflow.template.field.base import TemplateField from langflow.template.field.base import TemplateField
from langflow.template.frontend_node.base import FrontendNode from langflow.template.frontend_node.base import FrontendNode
from langflow.template.template.base import Template
class MemoryFrontendNode(FrontendNode): class MemoryFrontendNode(FrontendNode):
@ -64,3 +65,42 @@ class MemoryFrontendNode(FrontendNode):
field.value = "" field.value = ""
if field.name == "memory_key": if field.name == "memory_key":
field.value = "chat_history" field.value = "chat_history"
class PostgresChatMessageHistoryFrontendNode(MemoryFrontendNode):
name: str = "PostgresChatMessageHistory"
template: Template = Template(
type_name="PostgresChatMessageHistory",
fields=[
TemplateField(
field_type="str",
required=True,
placeholder="",
is_list=False,
show=True,
multiline=False,
name="session_id",
),
TemplateField(
field_type="str",
required=True,
show=True,
name="connection_string",
),
TemplateField(
field_type="str",
required=True,
placeholder="",
is_list=False,
show=True,
multiline=False,
value="message_store",
name="table_name",
),
],
)
description: str = "Memory store with Postgres"
base_classes: list[str] = [
"PostgresChatMessageHistory",
"BaseChatMessageHistory"
]