langchain/langchain/chains/base.py

"""Base interface that all chains should implement."""
from abc import ABC, abstractmethod
from typing import Any, Dict, List, Optional
from langchain.logger import PrintLogger
import uuid

from pydantic import BaseModel, Extra, root_validator

from langchain.logger import Logger, CONTEXT_KEY

class Memory(BaseModel, ABC):
    """Base interface for memory in chains."""

    class Config:
        """Configuration for this pydantic object."""

        extra = Extra.forbid
        arbitrary_types_allowed = True

    @property
    @abstractmethod
    def memory_variables(self) -> List[str]:
        """Input keys this memory class will load dynamically."""

    @abstractmethod
    def load_memory_variables(self, inputs: Dict[str, Any]) -> Dict[str, str]:
        """Return key-value pairs given the text input to the chain."""

    @abstractmethod
    def save_context(self, inputs: Dict[str, Any], outputs: Dict[str, str]) -> None:
        """Save the context of this model run to memory."""


class Chain(BaseModel, ABC):
    """Base interface that all chains should implement."""

    memory: Optional[Memory] = None

    verbose: bool = False
    """Whether to print out response text."""
    logger: Optional[Logger] = None

    @property
    @abstractmethod
    def input_keys(self) -> List[str]:
        """Input keys this chain expects."""

    @property
    @abstractmethod
    def output_keys(self) -> List[str]:
        """Output keys this chain expects."""

    @root_validator()
    def add_logger(cls, values: Dict) -> Dict:
        """Add a printing logger if verbose=True and none provided."""
        if values["verbose"] and values["logger"] is None:
            values["logger"] = PrintLogger()
        return values

    class Config:
        """Configuration for this pydantic object."""

        extra = Extra.forbid
        arbitrary_types_allowed = True

    def _validate_inputs(self, inputs: Dict[str, str]) -> None:
        """Check that all inputs are present."""
        missing_keys = set(self.input_keys).difference(inputs)
        if missing_keys:
            raise ValueError(f"Missing some input keys: {missing_keys}")

    def _validate_outputs(self, outputs: Dict[str, str]) -> None:
        if set(outputs) != set(self.output_keys):
            raise ValueError(
                f"Did not get output keys that were expected. "
                f"Got: {set(outputs)}. Expected: {set(self.output_keys)}."
            )

    @abstractmethod
    def _call(self, inputs: Dict[str, str]) -> Dict[str, str]:
        """Run the logic of this chain and return the output."""

    def __call__(
        self, inputs: Dict[str, Any], return_only_outputs: bool = False
    ) -> Dict[str, str]:
        """Run the logic of this chain and add to output if desired.

        Args:
            inputs: Dictionary of inputs.
            return_only_outputs: boolean for whether to return only outputs in the
                response. If True, only new keys generated by this chain will be
                returned. If False, both input keys and new keys generated by this
                chain will be returned. Defaults to False.

        """
        if CONTEXT_KEY not in inputs:
            inputs[CONTEXT_KEY] = {}
        if "id" not in inputs[CONTEXT_KEY]:
            inputs[CONTEXT_KEY]["id"] = str(uuid.uuid4())

        if self.memory is not None:
            external_context = self.memory.load_memory_variables(inputs)
            inputs = dict(inputs, **external_context)
        self._validate_inputs(inputs)
        if self.logger:
            self.logger.log_start_of_chain(inputs)
        outputs = self._call(inputs)
        self._validate_outputs(outputs)
        outputs[CONTEXT_KEY] = inputs[CONTEXT_KEY]
        if self.logger:
            self.logger.log_end_of_chain(outputs)
        if self.memory is not None:
            self.memory.save_context(inputs, outputs)
        if return_only_outputs:
            return outputs
        else:
            return {**inputs, **outputs}

    def apply(self, input_list: List[Dict[str, Any]]) -> List[Dict[str, str]]:
        """Call the chain on all inputs in the list."""
        return [self(inputs) for inputs in input_list]

    def run(self, text: str) -> str:
        """Run text in, text out (if applicable)."""
        if len(self.input_keys) != 1:
            raise ValueError(
                f"`run` not supported when there is not exactly "
                f"one input key, got {self.input_keys}."
            )
        if len(self.output_keys) != 1:
            raise ValueError(
                f"`run` not supported when there is not exactly "
                f"one output key, got {self.output_keys}."
            )
        return self({self.input_keys[0]: text})[self.output_keys[0]]