File size: 2,029 Bytes
9eaf8c4 b5e8391 6705f79 d25ef9b 0dad17d d25ef9b 0dad17d b5e8391 0dad17d 884caba 0dad17d 884caba 0dad17d 884caba 9eaf8c4 0dad17d 884caba 0dad17d e6a4ae7 0dad17d e6a4ae7 0dad17d e6a4ae7 |
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 |
import os
# Import the PipelineRunnable from pipeline.py
from pipeline import pipeline_runnable
from langchain_core.chat_history import BaseChatMessageHistory
from langchain_community.chat_message_histories import ChatMessageHistory
from langchain_core.runnables.history import RunnableWithMessageHistory
###############################################################################
# 1) In-memory store: session_id -> ChatMessageHistory
###############################################################################
store = {} # e.g. { "abc123": ChatMessageHistory() }
def get_session_history(session_id: str) -> BaseChatMessageHistory:
if session_id not in store:
store[session_id] = ChatMessageHistory()
return store[session_id]
###############################################################################
# 2) Create RunnableWithMessageHistory
###############################################################################
conversational_rag_chain = RunnableWithMessageHistory(
pipeline_runnable,
get_session_history,
input_messages_key="input",
history_messages_key="chat_history",
output_messages_key="answer"
)
###############################################################################
# 3) Convenience function for session-based memory
###############################################################################
def run_with_session_memory(user_query: str, session_id: str) -> str:
if not user_query:
raise ValueError("User query is missing. Ensure the 'input' key is provided.")
# Prepare input dictionary with the correct structure
input_data = {"input": user_query}
# Ensure the session ID is passed correctly in the configuration
response = conversational_rag_chain.invoke(
input_data, # Pass the dictionary containing the 'input' key
config={"configurable": {"session_id": session_id}} # Pass session ID in the config
)
return response.get("answer", "No answer returned from the chain.")
|