Spaces:
Running
Running
File size: 3,043 Bytes
8817c45 83d839e |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 |
from langchain_core.memory import BaseMemory
from langchain_core.messages import BaseMessage, AIMessage, HumanMessage
from langchain_core.pydantic_v1 import Field
from typing import List
from datetime import datetime
class EnhancedInMemoryHistory(BaseMemory):
"""
Custom memory class for storing chat history with timestamps.
Attributes:
messages (List[BaseMessage]): A list of messages exchanged between the user and the bot.
timestamps (List[datetime]): A list of timestamps when the messages were exchanged.
"""
messages: List[BaseMessage] = Field(default_factory=list)
timestamps: List[datetime] = Field(default_factory=list)
@property
def memory_variables(self):
"""Returns a list of memory variables (history) used in the conversation chain."""
return ["history"]
def add_messages(self, messages: List[BaseMessage]):
"""
Adds new messages to the memory and timestamps them.
Parameters:
messages (List[BaseMessage]): A list of messages to add to the memory.
"""
current_time = datetime.now()
self.messages.extend(messages)
self.timestamps.extend([current_time] * len(messages))
def get_recent_messages(self, limit: int = 5):
"""
Retrieves the most recent messages.
Parameters:
limit (int): Number of recent messages to retrieve (default is 5).
Returns:
List[BaseMessage]: A list of the most recent messages.
"""
return self.messages[-limit:]
def clear(self):
"""Clears all messages and timestamps from the memory."""
self.messages = []
self.timestamps = []
def load_memory_variables(self, inputs: dict):
"""
Loads memory variables for the conversation chain.
Parameters:
inputs (dict): Input data for the conversation chain.
Returns:
dict: A dictionary with the conversation history.
"""
return {"history": "\n".join([msg.content for msg in self.messages])}
def save_context(self, inputs: dict, outputs: str):
"""
Saves the context of the conversation by storing the user input and bot output.
Parameters:
inputs (dict): The user input.
outputs (str): The bot's response.
"""
self.add_messages([HumanMessage(content=inputs['input']), AIMessage(content=str(outputs))])
def clear_memory(self):
"""Clears the memory (alias for the clear function)."""
self.clear()
# Store sessions in a dictionary
store = {}
def get_by_session_id(session_id: str) -> EnhancedInMemoryHistory:
"""
Retrieves or creates a new memory session by session ID.
Parameters:
session_id (str): The session ID for the chat session.
Returns:
EnhancedInMemoryHistory: The memory object for the session.
"""
if session_id not in store:
store[session_id] = EnhancedInMemoryHistory()
return store[session_id] |