mirror of
https://github.com/hwchase17/langchain.git
synced 2025-09-02 03:26:17 +00:00
Harrison/improve memory (#432)
add AI prefix add new type of memory Co-authored-by: Jason <chisanch@usc.edu>
This commit is contained in:
31
tests/integration_tests/chains/test_memory.py
Normal file
31
tests/integration_tests/chains/test_memory.py
Normal file
@@ -0,0 +1,31 @@
|
||||
"""Test memory functionality."""
|
||||
from langchain.chains.conversation.memory import ConversationSummaryBufferMemory
|
||||
from tests.unit_tests.llms.fake_llm import FakeLLM
|
||||
|
||||
|
||||
def test_summary_buffer_memory_no_buffer_yet() -> None:
|
||||
"""Test ConversationSummaryBufferMemory when no inputs put in buffer yet."""
|
||||
memory = ConversationSummaryBufferMemory(llm=FakeLLM(), memory_key="baz")
|
||||
output = memory.load_memory_variables({})
|
||||
assert output == {"baz": ""}
|
||||
|
||||
|
||||
def test_summary_buffer_memory_buffer_only() -> None:
|
||||
"""Test ConversationSummaryBufferMemory when only buffer."""
|
||||
memory = ConversationSummaryBufferMemory(llm=FakeLLM(), memory_key="baz")
|
||||
memory.save_context({"input": "bar"}, {"output": "foo"})
|
||||
assert memory.buffer == ["Human: bar\nAI: foo"]
|
||||
output = memory.load_memory_variables({})
|
||||
assert output == {"baz": "Human: bar\nAI: foo"}
|
||||
|
||||
|
||||
def test_summary_buffer_memory_summary() -> None:
|
||||
"""Test ConversationSummaryBufferMemory when only buffer."""
|
||||
memory = ConversationSummaryBufferMemory(
|
||||
llm=FakeLLM(), memory_key="baz", max_token_limit=13
|
||||
)
|
||||
memory.save_context({"input": "bar"}, {"output": "foo"})
|
||||
memory.save_context({"input": "bar1"}, {"output": "foo1"})
|
||||
assert memory.buffer == ["Human: bar1\nAI: foo1"]
|
||||
output = memory.load_memory_variables({})
|
||||
assert output == {"baz": "foo\nHuman: bar1\nAI: foo1"}
|
Reference in New Issue
Block a user