from operator import itemgetter

from langchain.memory import ConversationBufferMemory
from langchain_core.prompts import ChatPromptTemplate, MessagesPlaceholder
from langchain_core.runnables import RunnableLambda, RunnablePassthrough
from langchain_openai import ChatOpenAI

model = ChatOpenAI()
prompt = ChatPromptTemplate.from_messages(
    [
        ("system", "你是一个乐于助人的机器人"),
        MessagesPlaceholder(variable_name="history"),
        ("human", "{input}"),
    ]
)

memory = ConversationBufferMemory(return_messages=True)

memory.load_memory_variables({})

chain = (
    RunnablePassthrough.assign(
        history=RunnableLambda(memory.load_memory_variables) | itemgetter("history")
    )
    | prompt
    | model
)

inputs = {"input": "你好我是tomie"}
response = chain.invoke(inputs)
response

#保存记忆
memory.save_context(inputs, {"output": response.content})
memory.load_memory_variables({})

inputs = {"input": "我叫什么名字?"}
response = chain.invoke(inputs)
response

from typing import Optional

from langchain_community.chat_message_histories import RedisChatMessageHistory
from langchain_community.chat_models import ChatOpenAI
from langchain_core.chat_history import BaseChatMessageHistory
from langchain_core.prompts import ChatPromptTemplate, MessagesPlaceholder
from langchain_core.runnables.history import RunnableWithMessageHistory

prompt = ChatPromptTemplate.from_messages(
    [
        ("system", "你是一个擅长{ability}的助手"),
        MessagesPlaceholder(variable_name="history"),
        ("human", "{question}"),
    ]
)

chain = prompt | ChatOpenAI(model="gpt-4-1106-preview",temperature=0)

chain_with_history = RunnableWithMessageHistory(
    chain,
    #使用redis存储聊天记录
    lambda session_id: RedisChatMessageHistory(session_id, url="redis://localhost:6379/0"),
    input_messages_key="question",
    history_messages_key="history",
)

#每次调用都会保存聊天记录，需要有对应的session_id
chain_with_history.invoke(
    {"ability": "历史", "question": "中国建都时间最长的城市是哪个?"},
    config={"configurable": {"session_id": "tomiezhang"}},
)

chain_with_history.invoke(
    {"ability": "历史", "question": "它有多少年建都历史？"},
    config={"configurable": {"session_id": "tomiezhang"}},
)