|
|
"""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} |
|
|
|
