hari-huynh
Update memory for new chat session
fe81246
raw
history blame
1.85 kB
import chainlit as cl
from langchain.prompts import ChatPromptTemplate
from langchain.schema import StrOutputParser
from langchain.schema.runnable import Runnable
from langchain.schema.runnable.config import RunnableConfig
from react_agent_v2 import get_react_agent
from langchain.memory import ConversationBufferMemory
from langchain_community.chat_message_histories import ChatMessageHistory
from langchain_core.chat_history import BaseChatMessageHistory
from langchain_core.runnables.history import RunnableWithMessageHistory
@cl.on_chat_start
async def on_chat_start():
message_history = ChatMessageHistory()
memory = ConversationBufferMemory(
memory_key = "chat_history",
output_key = "output",
chat_memory = message_history,
return_message = True
)
agent_executor = get_react_agent(memory)
cl.user_session.set("runnable", agent_executor)
@cl.on_message
async def on_message(message: cl.Message):
# runnable = cl.user_session.get("runnable") # type: Runnable
#
# msg = cl.Message(content="")
#
# for chunk in await cl.make_async(runnable.stream)(
# {"input": message.content},
# config=RunnableConfig(callbacks=[cl.LangchainCallbackHandler()]),
# ):
# await msg.stream_token(chunk)
#
# await msg.send()
#
# # Get user input from the message
# user_input = message.content
#
# # Run the agent with user input and get the response
# response = await cl.make_async(agent_executor)(user_input)
#
# # Display the response to the user
# cl.message(response)
llm_chain = cl.user_session.get("runnable")
response = await llm_chain.ainvoke(
{"input": message.content}, callbacks = [cl.AsyncLangchainCallbackHandler()]
)
await cl.Message(response["output"].replace("`", "")).send()