test_sequential.py — langchain Source File

Source Code

"""Test pipeline functionality."""

import re

import pytest
from langchain_core.callbacks.manager import (
    AsyncCallbackManagerForChainRun,
    CallbackManagerForChainRun,
)
from typing_extensions import override

from langchain_classic.chains.base import Chain
from langchain_classic.chains.sequential import SequentialChain, SimpleSequentialChain
from langchain_classic.memory import ConversationBufferMemory
from langchain_classic.memory.simple import SimpleMemory
from tests.unit_tests.callbacks.fake_callback_handler import FakeCallbackHandler


class FakeChain(Chain):
    """Fake Chain for testing purposes."""

    input_variables: list[str]
    output_variables: list[str]

    @property
    def input_keys(self) -> list[str]:
        """Input keys this chain returns."""
        return self.input_variables

    @property
    def output_keys(self) -> list[str]:
        """Input keys this chain returns."""
        return self.output_variables

    @override
    def _call(
        self,
        inputs: dict[str, str],
        run_manager: CallbackManagerForChainRun | None = None,
    ) -> dict[str, str]:
        outputs = {}
        for var in self.output_variables:
            variables = [inputs[k] for k in self.input_variables]
            outputs[var] = f"{' '.join(variables)}foo"
        return outputs

    @override
    async def _acall(
        self,
        inputs: dict[str, str],
        run_manager: AsyncCallbackManagerForChainRun | None = None,
    ) -> dict[str, str]:
        outputs = {}
        for var in self.output_variables:
            variables = [inputs[k] for k in self.input_variables]
            outputs[var] = f"{' '.join(variables)}foo"
        return outputs


def test_sequential_usage_single_inputs() -> None:
// ... (205 more lines)