Update my_memory_logic.py
Browse files- my_memory_logic.py +24 -21
my_memory_logic.py
CHANGED
@@ -2,52 +2,55 @@
|
|
2 |
|
3 |
import os
|
4 |
|
5 |
-
# Import
|
6 |
-
#
|
7 |
-
|
8 |
-
|
9 |
-
|
|
|
10 |
from langchain_core.chat_history import BaseChatMessageHistory
|
11 |
from langchain_community.chat_message_histories import ChatMessageHistory
|
12 |
from langchain_core.runnables.history import RunnableWithMessageHistory
|
13 |
|
14 |
###############################################################################
|
15 |
-
# 1)
|
16 |
###############################################################################
|
17 |
store = {} # e.g., { "abc123": ChatMessageHistory(...) }
|
18 |
|
19 |
def get_session_history(session_id: str) -> BaseChatMessageHistory:
|
20 |
"""
|
21 |
-
Retrieve or create a ChatMessageHistory
|
22 |
-
This ensures each session_id has its own conversation
|
23 |
"""
|
24 |
if session_id not in store:
|
25 |
store[session_id] = ChatMessageHistory()
|
26 |
return store[session_id]
|
27 |
|
28 |
###############################################################################
|
29 |
-
# 2)
|
30 |
###############################################################################
|
31 |
-
#
|
32 |
-
#
|
|
|
|
|
|
|
|
|
33 |
conversational_rag_chain = RunnableWithMessageHistory(
|
34 |
-
run_with_chain_context,
|
35 |
-
get_session_history,
|
36 |
-
input_messages_key="input",
|
37 |
-
history_messages_key="chat_history",
|
38 |
-
output_messages_key="answer"
|
39 |
)
|
40 |
|
41 |
###############################################################################
|
42 |
-
# 3) A convenience function
|
43 |
###############################################################################
|
44 |
def run_with_session_memory(user_query: str, session_id: str) -> str:
|
45 |
"""
|
46 |
-
|
47 |
-
|
48 |
"""
|
49 |
-
# We invoke the chain with the user query;
|
50 |
-
# the chain automatically updates the session’s chat history.
|
51 |
response = conversational_rag_chain.invoke(
|
52 |
{"input": user_query},
|
53 |
config={
|
|
|
2 |
|
3 |
import os
|
4 |
|
5 |
+
# Import the "run_with_chain_context" function from pipeline.py
|
6 |
+
# This function must accept a dict with { "input": ..., "chat_history": ... }
|
7 |
+
# and return a dict with { "answer": ... }.
|
8 |
+
from pipeline import run_with_chain_context
|
9 |
+
|
10 |
+
# For session-based chat history
|
11 |
from langchain_core.chat_history import BaseChatMessageHistory
|
12 |
from langchain_community.chat_message_histories import ChatMessageHistory
|
13 |
from langchain_core.runnables.history import RunnableWithMessageHistory
|
14 |
|
15 |
###############################################################################
|
16 |
+
# 1) In-Memory Store: session_id -> ChatMessageHistory
|
17 |
###############################################################################
|
18 |
store = {} # e.g., { "abc123": ChatMessageHistory(...) }
|
19 |
|
20 |
def get_session_history(session_id: str) -> BaseChatMessageHistory:
|
21 |
"""
|
22 |
+
Retrieve (or create) a ChatMessageHistory for the given session_id.
|
23 |
+
This ensures each session_id has its own conversation transcripts.
|
24 |
"""
|
25 |
if session_id not in store:
|
26 |
store[session_id] = ChatMessageHistory()
|
27 |
return store[session_id]
|
28 |
|
29 |
###############################################################################
|
30 |
+
# 2) Build a RunnableWithMessageHistory that wraps "run_with_chain_context"
|
31 |
###############################################################################
|
32 |
+
# "run_with_chain_context" must be a function returning a dict,
|
33 |
+
# e.g. { "answer": "... final string ..." }
|
34 |
+
# input_messages_key -> "input"
|
35 |
+
# history_messages_key -> "chat_history"
|
36 |
+
# output_messages_key -> "answer"
|
37 |
+
|
38 |
conversational_rag_chain = RunnableWithMessageHistory(
|
39 |
+
run_with_chain_context, # from pipeline.py
|
40 |
+
get_session_history,
|
41 |
+
input_messages_key="input",
|
42 |
+
history_messages_key="chat_history",
|
43 |
+
output_messages_key="answer"
|
44 |
)
|
45 |
|
46 |
###############################################################################
|
47 |
+
# 3) A convenience function that calls our chain with session-based memory
|
48 |
###############################################################################
|
49 |
def run_with_session_memory(user_query: str, session_id: str) -> str:
|
50 |
"""
|
51 |
+
Calls the 'conversational_rag_chain' with a given session_id and user_query.
|
52 |
+
This returns the final 'answer' from run_with_chain_context.
|
53 |
"""
|
|
|
|
|
54 |
response = conversational_rag_chain.invoke(
|
55 |
{"input": user_query},
|
56 |
config={
|