Update my_memory_logic.py
Browse files- my_memory_logic.py +10 -26
my_memory_logic.py
CHANGED
|
@@ -12,46 +12,30 @@ from langchain_core.runnables.history import RunnableWithMessageHistory
|
|
| 12 |
###############################################################################
|
| 13 |
# 1) In-memory store: session_id -> ChatMessageHistory
|
| 14 |
###############################################################################
|
| 15 |
-
store = {} # e.g.
|
| 16 |
|
| 17 |
def get_session_history(session_id: str) -> BaseChatMessageHistory:
|
| 18 |
-
"""
|
| 19 |
-
Retrieve or create a ChatMessageHistory object for the given session_id.
|
| 20 |
-
This ensures each session_id has its own conversation history.
|
| 21 |
-
"""
|
| 22 |
if session_id not in store:
|
| 23 |
store[session_id] = ChatMessageHistory()
|
| 24 |
return store[session_id]
|
| 25 |
|
| 26 |
###############################################################################
|
| 27 |
-
# 2) Create
|
| 28 |
###############################################################################
|
| 29 |
-
# This wraps your `pipeline_runnable` so it automatically reads/writes
|
| 30 |
-
# conversation history from get_session_history for each session.
|
| 31 |
conversational_rag_chain = RunnableWithMessageHistory(
|
| 32 |
-
pipeline_runnable,
|
| 33 |
-
get_session_history,
|
| 34 |
-
input_messages_key="input",
|
| 35 |
-
history_messages_key="chat_history",
|
| 36 |
-
output_messages_key="answer"
|
| 37 |
)
|
| 38 |
|
| 39 |
###############################################################################
|
| 40 |
-
# 3)
|
| 41 |
###############################################################################
|
| 42 |
def run_with_session_memory(user_query: str, session_id: str) -> str:
|
| 43 |
-
"""
|
| 44 |
-
A helper that calls our `conversational_rag_chain`
|
| 45 |
-
with a given session_id. Returns the final 'answer'.
|
| 46 |
-
"""
|
| 47 |
-
# We invoke the chain with the user query;
|
| 48 |
-
# the chain automatically updates the session’s chat history.
|
| 49 |
response = conversational_rag_chain.invoke(
|
| 50 |
{"input": user_query},
|
| 51 |
-
config={
|
| 52 |
-
"configurable": {
|
| 53 |
-
"session_id": session_id
|
| 54 |
-
}
|
| 55 |
-
}
|
| 56 |
)
|
| 57 |
-
return response["answer"]
|
|
|
|
| 12 |
###############################################################################
|
| 13 |
# 1) In-memory store: session_id -> ChatMessageHistory
|
| 14 |
###############################################################################
|
| 15 |
+
store = {} # e.g. { "abc123": ChatMessageHistory() }
|
| 16 |
|
| 17 |
def get_session_history(session_id: str) -> BaseChatMessageHistory:
|
|
|
|
|
|
|
|
|
|
|
|
|
| 18 |
if session_id not in store:
|
| 19 |
store[session_id] = ChatMessageHistory()
|
| 20 |
return store[session_id]
|
| 21 |
|
| 22 |
###############################################################################
|
| 23 |
+
# 2) Create RunnableWithMessageHistory
|
| 24 |
###############################################################################
|
|
|
|
|
|
|
| 25 |
conversational_rag_chain = RunnableWithMessageHistory(
|
| 26 |
+
pipeline_runnable,
|
| 27 |
+
get_session_history,
|
| 28 |
+
input_messages_key="input",
|
| 29 |
+
history_messages_key="chat_history",
|
| 30 |
+
output_messages_key="answer"
|
| 31 |
)
|
| 32 |
|
| 33 |
###############################################################################
|
| 34 |
+
# 3) Convenience function for session-based memory
|
| 35 |
###############################################################################
|
| 36 |
def run_with_session_memory(user_query: str, session_id: str) -> str:
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 37 |
response = conversational_rag_chain.invoke(
|
| 38 |
{"input": user_query},
|
| 39 |
+
config={"configurable": {"session_id": session_id}}
|
|
|
|
|
|
|
|
|
|
|
|
|
| 40 |
)
|
| 41 |
+
return response["answer"]
|