Phoenix21 commited on
Commit
884caba
·
verified ·
1 Parent(s): 2065cb4

Update my_memory_logic.py

Browse files
Files changed (1) hide show
  1. 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., { "abc123": ChatMessageHistory(...) }
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 the RunnableWithMessageHistory (session-based chain)
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, # the Runnable from pipeline.py
33
- get_session_history, # fetches or creates ChatMessageHistory by session_id
34
- input_messages_key="input", # key in the dict for user's new query
35
- history_messages_key="chat_history", # key for existing chat logs
36
- output_messages_key="answer" # key for final output
37
  )
38
 
39
  ###############################################################################
40
- # 3) A convenience function to run a query with session-based memory
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"]