|
|
import warnings |
|
|
from typing import Any, Dict, List, Set |
|
|
|
|
|
from langchain_core.memory import BaseMemory |
|
|
from pydantic import field_validator |
|
|
|
|
|
from langchain.memory.chat_memory import BaseChatMemory |
|
|
|
|
|
|
|
|
class CombinedMemory(BaseMemory): |
|
|
"""Combining multiple memories' data together.""" |
|
|
|
|
|
memories: List[BaseMemory] |
|
|
"""For tracking all the memories that should be accessed.""" |
|
|
|
|
|
@field_validator("memories") |
|
|
@classmethod |
|
|
def check_repeated_memory_variable( |
|
|
cls, value: List[BaseMemory] |
|
|
) -> List[BaseMemory]: |
|
|
all_variables: Set[str] = set() |
|
|
for val in value: |
|
|
overlap = all_variables.intersection(val.memory_variables) |
|
|
if overlap: |
|
|
raise ValueError( |
|
|
f"The same variables {overlap} are found in multiple" |
|
|
"memory object, which is not allowed by CombinedMemory." |
|
|
) |
|
|
all_variables |= set(val.memory_variables) |
|
|
|
|
|
return value |
|
|
|
|
|
@field_validator("memories") |
|
|
@classmethod |
|
|
def check_input_key(cls, value: List[BaseMemory]) -> List[BaseMemory]: |
|
|
"""Check that if memories are of type BaseChatMemory that input keys exist.""" |
|
|
for val in value: |
|
|
if isinstance(val, BaseChatMemory): |
|
|
if val.input_key is None: |
|
|
warnings.warn( |
|
|
"When using CombinedMemory, " |
|
|
"input keys should be so the input is known. " |
|
|
f" Was not set on {val}" |
|
|
) |
|
|
return value |
|
|
|
|
|
@property |
|
|
def memory_variables(self) -> List[str]: |
|
|
"""All the memory variables that this instance provides.""" |
|
|
"""Collected from the all the linked memories.""" |
|
|
|
|
|
memory_variables = [] |
|
|
|
|
|
for memory in self.memories: |
|
|
memory_variables.extend(memory.memory_variables) |
|
|
|
|
|
return memory_variables |
|
|
|
|
|
def load_memory_variables(self, inputs: Dict[str, Any]) -> Dict[str, str]: |
|
|
"""Load all vars from sub-memories.""" |
|
|
memory_data: Dict[str, Any] = {} |
|
|
|
|
|
|
|
|
for memory in self.memories: |
|
|
data = memory.load_memory_variables(inputs) |
|
|
for key, value in data.items(): |
|
|
if key in memory_data: |
|
|
raise ValueError( |
|
|
f"The variable {key} is repeated in the CombinedMemory." |
|
|
) |
|
|
memory_data[key] = value |
|
|
|
|
|
return memory_data |
|
|
|
|
|
def save_context(self, inputs: Dict[str, Any], outputs: Dict[str, str]) -> None: |
|
|
"""Save context from this session for every memory.""" |
|
|
|
|
|
for memory in self.memories: |
|
|
memory.save_context(inputs, outputs) |
|
|
|
|
|
def clear(self) -> None: |
|
|
"""Clear context from this session for every memory.""" |
|
|
for memory in self.memories: |
|
|
memory.clear() |
|
|
|
