Phoenix21 commited on
Commit
6705f79
·
verified ·
1 Parent(s): a79a41b

Update my_memory_logic.py

Browse files
Files changed (1) hide show
  1. my_memory_logic.py +24 -21
my_memory_logic.py CHANGED
@@ -2,52 +2,55 @@
2
 
3
  import os
4
 
5
- # Import your actual RAG chain (or pipeline) from pipeline.py
6
- # We'll assume `rag_chain` is exposed by pipeline.py
7
- from pipeline import run_with_chain_context
8
- # We'll import the session-based classes from langchain_core
9
- # If they're in different modules, adjust accordingly.
 
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) We'll keep an in-memory store mapping session_id -> ChatMessageHistory
16
  ###############################################################################
17
  store = {} # e.g., { "abc123": ChatMessageHistory(...) }
18
 
19
  def get_session_history(session_id: str) -> BaseChatMessageHistory:
20
  """
21
- Retrieve or create a ChatMessageHistory object for the given session_id.
22
- This ensures each session_id has its own conversation history.
23
  """
24
  if session_id not in store:
25
  store[session_id] = ChatMessageHistory()
26
  return store[session_id]
27
 
28
  ###############################################################################
29
- # 2) Create the RunnableWithMessageHistory (session-based chain)
30
  ###############################################################################
31
- # This wraps your `rag_chain` so it automatically reads/writes
32
- # conversation history from get_session_history for each session.
 
 
 
 
33
  conversational_rag_chain = RunnableWithMessageHistory(
34
- run_with_chain_context, # the main chain from pipeline.py
35
- get_session_history, # fetches or creates ChatMessageHistory by session_id
36
- input_messages_key="input", # key in the dict for user's new query
37
- history_messages_key="chat_history", # key for existing chat logs
38
- output_messages_key="answer" # key for final output
39
  )
40
 
41
  ###############################################################################
42
- # 3) A convenience function to run a query with session-based memory
43
  ###############################################################################
44
  def run_with_session_memory(user_query: str, session_id: str) -> str:
45
  """
46
- A helper that calls our `conversational_rag_chain`
47
- with a given session_id. Returns the final 'answer'.
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={