test_conversation.py — langchain Source File

Source Code

"""Test conversation chain and memory."""

import re
from typing import Any

import pytest
from langchain_core.callbacks import CallbackManagerForLLMRun
from langchain_core.language_models import LLM
from langchain_core.prompts.prompt import PromptTemplate
from typing_extensions import override

from langchain_classic.base_memory import BaseMemory
from langchain_classic.chains.conversation.base import ConversationChain
from langchain_classic.memory.buffer import ConversationBufferMemory
from langchain_classic.memory.buffer_window import ConversationBufferWindowMemory
from langchain_classic.memory.summary import ConversationSummaryMemory
from tests.unit_tests.llms.fake_llm import FakeLLM


class DummyLLM(LLM):
    last_prompt: str = ""

    def __init__(self, **kwargs: Any):
        super().__init__(**kwargs)

    @property
    def _llm_type(self) -> str:
        return "dummy"

    @override
    def _call(
        self,
        prompt: str,
        stop: list[str] | None = None,
        run_manager: CallbackManagerForLLMRun | None = None,
        **kwargs: Any,
    ) -> str:
        self.last_prompt = prompt
        return "dummy"


def test_memory_ai_prefix() -> None:
    """Test that ai_prefix in the memory component works."""
    memory = ConversationBufferMemory(memory_key="foo", ai_prefix="Assistant")
    memory.save_context({"input": "bar"}, {"output": "foo"})
    assert memory.load_memory_variables({}) == {"foo": "Human: bar\nAssistant: foo"}


def test_memory_human_prefix() -> None:
    """Test that human_prefix in the memory component works."""
    memory = ConversationBufferMemory(memory_key="foo", human_prefix="Friend")
    memory.save_context({"input": "bar"}, {"output": "foo"})
    assert memory.load_memory_variables({}) == {"foo": "Friend: bar\nAI: foo"}


async def test_memory_async() -> None:
    memory = ConversationBufferMemory(memory_key="foo", ai_prefix="Assistant")
    await memory.asave_context({"input": "bar"}, {"output": "foo"})
    assert await memory.aload_memory_variables({}) == {
        "foo": "Human: bar\nAssistant: foo",
// ... (112 more lines)