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"]
|