| """Chain pipeline where the outputs of one step feed directly into next.""" |
|
|
| from typing import Any, Dict, List, Optional |
|
|
| from langchain_core.callbacks import ( |
| AsyncCallbackManagerForChainRun, |
| CallbackManagerForChainRun, |
| ) |
| from langchain_core.utils.input import get_color_mapping |
| from pydantic import ConfigDict, model_validator |
| from typing_extensions import Self |
|
|
| from langchain.chains.base import Chain |
|
|
|
|
| class SequentialChain(Chain): |
| """Chain where the outputs of one chain feed directly into next.""" |
|
|
| chains: List[Chain] |
| input_variables: List[str] |
| output_variables: List[str] |
| return_all: bool = False |
|
|
| model_config = ConfigDict( |
| arbitrary_types_allowed=True, |
| extra="forbid", |
| ) |
|
|
| @property |
| def input_keys(self) -> List[str]: |
| """Return expected input keys to the chain. |
| |
| :meta private: |
| """ |
| return self.input_variables |
|
|
| @property |
| def output_keys(self) -> List[str]: |
| """Return output key. |
| |
| :meta private: |
| """ |
| return self.output_variables |
|
|
| @model_validator(mode="before") |
| @classmethod |
| def validate_chains(cls, values: Dict) -> Any: |
| """Validate that the correct inputs exist for all chains.""" |
| chains = values["chains"] |
| input_variables = values["input_variables"] |
| memory_keys = list() |
| if "memory" in values and values["memory"] is not None: |
| """Validate that prompt input variables are consistent.""" |
| memory_keys = values["memory"].memory_variables |
| if set(input_variables).intersection(set(memory_keys)): |
| overlapping_keys = set(input_variables) & set(memory_keys) |
| raise ValueError( |
| f"The input key(s) {''.join(overlapping_keys)} are found " |
| f"in the Memory keys ({memory_keys}) - please use input and " |
| f"memory keys that don't overlap." |
| ) |
|
|
| known_variables = set(input_variables + memory_keys) |
|
|
| for chain in chains: |
| missing_vars = set(chain.input_keys).difference(known_variables) |
| if chain.memory: |
| missing_vars = missing_vars.difference(chain.memory.memory_variables) |
|
|
| if missing_vars: |
| raise ValueError( |
| f"Missing required input keys: {missing_vars}, " |
| f"only had {known_variables}" |
| ) |
| overlapping_keys = known_variables.intersection(chain.output_keys) |
| if overlapping_keys: |
| raise ValueError( |
| f"Chain returned keys that already exist: {overlapping_keys}" |
| ) |
|
|
| known_variables |= set(chain.output_keys) |
|
|
| if "output_variables" not in values: |
| if values.get("return_all", False): |
| output_keys = known_variables.difference(input_variables) |
| else: |
| output_keys = chains[-1].output_keys |
| values["output_variables"] = output_keys |
| else: |
| missing_vars = set(values["output_variables"]).difference(known_variables) |
| if missing_vars: |
| raise ValueError( |
| f"Expected output variables that were not found: {missing_vars}." |
| ) |
|
|
| return values |
|
|
| def _call( |
| self, |
| inputs: Dict[str, str], |
| run_manager: Optional[CallbackManagerForChainRun] = None, |
| ) -> Dict[str, str]: |
| known_values = inputs.copy() |
| _run_manager = run_manager or CallbackManagerForChainRun.get_noop_manager() |
| for i, chain in enumerate(self.chains): |
| callbacks = _run_manager.get_child() |
| outputs = chain(known_values, return_only_outputs=True, callbacks=callbacks) |
| known_values.update(outputs) |
| return {k: known_values[k] for k in self.output_variables} |
|
|
| async def _acall( |
| self, |
| inputs: Dict[str, Any], |
| run_manager: Optional[AsyncCallbackManagerForChainRun] = None, |
| ) -> Dict[str, Any]: |
| known_values = inputs.copy() |
| _run_manager = run_manager or AsyncCallbackManagerForChainRun.get_noop_manager() |
| callbacks = _run_manager.get_child() |
| for i, chain in enumerate(self.chains): |
| outputs = await chain.acall( |
| known_values, return_only_outputs=True, callbacks=callbacks |
| ) |
| known_values.update(outputs) |
| return {k: known_values[k] for k in self.output_variables} |
|
|
|
|
| class SimpleSequentialChain(Chain): |
| """Simple chain where the outputs of one step feed directly into next.""" |
|
|
| chains: List[Chain] |
| strip_outputs: bool = False |
| input_key: str = "input" |
| output_key: str = "output" |
|
|
| model_config = ConfigDict( |
| arbitrary_types_allowed=True, |
| extra="forbid", |
| ) |
|
|
| @property |
| def input_keys(self) -> List[str]: |
| """Expect input key. |
| |
| :meta private: |
| """ |
| return [self.input_key] |
|
|
| @property |
| def output_keys(self) -> List[str]: |
| """Return output key. |
| |
| :meta private: |
| """ |
| return [self.output_key] |
|
|
| @model_validator(mode="after") |
| def validate_chains(self) -> Self: |
| """Validate that chains are all single input/output.""" |
| for chain in self.chains: |
| if len(chain.input_keys) != 1: |
| raise ValueError( |
| "Chains used in SimplePipeline should all have one input, got " |
| f"{chain} with {len(chain.input_keys)} inputs." |
| ) |
| if len(chain.output_keys) != 1: |
| raise ValueError( |
| "Chains used in SimplePipeline should all have one output, got " |
| f"{chain} with {len(chain.output_keys)} outputs." |
| ) |
| return self |
|
|
| def _call( |
| self, |
| inputs: Dict[str, str], |
| run_manager: Optional[CallbackManagerForChainRun] = None, |
| ) -> Dict[str, str]: |
| _run_manager = run_manager or CallbackManagerForChainRun.get_noop_manager() |
| _input = inputs[self.input_key] |
| color_mapping = get_color_mapping([str(i) for i in range(len(self.chains))]) |
| for i, chain in enumerate(self.chains): |
| _input = chain.run(_input, callbacks=_run_manager.get_child(f"step_{i+1}")) |
| if self.strip_outputs: |
| _input = _input.strip() |
| _run_manager.on_text( |
| _input, color=color_mapping[str(i)], end="\n", verbose=self.verbose |
| ) |
| return {self.output_key: _input} |
|
|
| async def _acall( |
| self, |
| inputs: Dict[str, Any], |
| run_manager: Optional[AsyncCallbackManagerForChainRun] = None, |
| ) -> Dict[str, Any]: |
| _run_manager = run_manager or AsyncCallbackManagerForChainRun.get_noop_manager() |
| _input = inputs[self.input_key] |
| color_mapping = get_color_mapping([str(i) for i in range(len(self.chains))]) |
| for i, chain in enumerate(self.chains): |
| _input = await chain.arun( |
| _input, callbacks=_run_manager.get_child(f"step_{i+1}") |
| ) |
| if self.strip_outputs: |
| _input = _input.strip() |
| await _run_manager.on_text( |
| _input, color=color_mapping[str(i)], end="\n", verbose=self.verbose |
| ) |
| return {self.output_key: _input} |
|
|
