File size: 2,705 Bytes
5b2b2db 7b899a2 9eaf8c4 7b899a2 9b11728 7b899a2 0dad17d 7b899a2 0dad17d 7b899a2 0dad17d 7b899a2 0dad17d 7b899a2 0dad17d 7b899a2 0dad17d 9b11728 7b899a2 9eaf8c4 0dad17d 7b899a2 0dad17d 7b899a2 0dad17d |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 |
# my_memory_logic.py
import os
# Import your actual RAG chain (or pipeline) from pipeline.py
# We'll assume `rag_chain` is exposed by pipeline.py
from pipeline import run_with_chain_context
# We'll import the session-based classes from langchain_core
# If they're in different modules, adjust accordingly.
from langchain_core.chat_history import BaseChatMessageHistory
from langchain_community.chat_message_histories import ChatMessageHistory
from langchain_core.runnables.history import RunnableWithMessageHistory
###############################################################################
# 1) We'll keep an in-memory store mapping session_id -> ChatMessageHistory
###############################################################################
store = {} # e.g., { "abc123": ChatMessageHistory(...) }
def get_session_history(session_id: str) -> BaseChatMessageHistory:
"""
Retrieve or create a ChatMessageHistory object for the given session_id.
This ensures each session_id has its own conversation history.
"""
if session_id not in store:
store[session_id] = ChatMessageHistory()
return store[session_id]
###############################################################################
# 2) Create the RunnableWithMessageHistory (session-based chain)
###############################################################################
# This wraps your `rag_chain` so it automatically reads/writes
# conversation history from get_session_history for each session.
conversational_rag_chain = RunnableWithMessageHistory(
run_with_chain_context, # the main chain from pipeline.py
get_session_history, # fetches or creates ChatMessageHistory by session_id
input_messages_key="input", # key in the dict for user's new query
history_messages_key="chat_history", # key for existing chat logs
output_messages_key="answer" # key for final output
)
###############################################################################
# 3) A convenience function to run a query with session-based memory
###############################################################################
def run_with_session_memory(user_query: str, session_id: str) -> str:
"""
A helper that calls our `conversational_rag_chain`
with a given session_id. Returns the final 'answer'.
"""
# We invoke the chain with the user query;
# the chain automatically updates the session’s chat history.
response = conversational_rag_chain.invoke(
{"input": user_query},
config={
"configurable": {
"session_id": session_id
}
}
)
return response["answer"]
|