forked from Archives/langchain
You cannot select more than 25 topics
Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.
38 lines
1.1 KiB
Python
38 lines
1.1 KiB
Python
import pytest
|
|
|
|
from langchain.chains.conversation.memory import (
|
|
ConversationBufferMemory,
|
|
ConversationBufferWindowMemory,
|
|
ConversationSummaryMemory,
|
|
)
|
|
from langchain.memory import ReadOnlySharedMemory, SimpleMemory
|
|
from langchain.schema import BaseMemory
|
|
from tests.unit_tests.llms.fake_llm import FakeLLM
|
|
|
|
|
|
def test_simple_memory() -> None:
|
|
"""Test SimpleMemory."""
|
|
memory = SimpleMemory(memories={"baz": "foo"})
|
|
|
|
output = memory.load_memory_variables({})
|
|
|
|
assert output == {"baz": "foo"}
|
|
assert ["baz"] == memory.memory_variables
|
|
|
|
|
|
@pytest.mark.parametrize(
|
|
"memory",
|
|
[
|
|
ConversationBufferMemory(memory_key="baz"),
|
|
ConversationSummaryMemory(llm=FakeLLM(), memory_key="baz"),
|
|
ConversationBufferWindowMemory(memory_key="baz"),
|
|
],
|
|
)
|
|
def test_readonly_memory(memory: BaseMemory) -> None:
|
|
read_only_memory = ReadOnlySharedMemory(memory=memory)
|
|
memory.save_context({"input": "bar"}, {"output": "foo"})
|
|
|
|
assert read_only_memory.load_memory_variables({}) == memory.load_memory_variables(
|
|
{}
|
|
)
|