|
|
"""Base classes for chain routing.""" |
|
|
|
|
|
from __future__ import annotations |
|
|
|
|
|
from abc import ABC |
|
|
from typing import Any, Dict, List, Mapping, NamedTuple, Optional |
|
|
|
|
|
from langchain_core.callbacks import ( |
|
|
AsyncCallbackManagerForChainRun, |
|
|
CallbackManagerForChainRun, |
|
|
Callbacks, |
|
|
) |
|
|
from pydantic import ConfigDict |
|
|
|
|
|
from langchain.chains.base import Chain |
|
|
|
|
|
|
|
|
class Route(NamedTuple): |
|
|
destination: Optional[str] |
|
|
next_inputs: Dict[str, Any] |
|
|
|
|
|
|
|
|
class RouterChain(Chain, ABC): |
|
|
"""Chain that outputs the name of a destination chain and the inputs to it.""" |
|
|
|
|
|
@property |
|
|
def output_keys(self) -> List[str]: |
|
|
return ["destination", "next_inputs"] |
|
|
|
|
|
def route(self, inputs: Dict[str, Any], callbacks: Callbacks = None) -> Route: |
|
|
""" |
|
|
Route inputs to a destination chain. |
|
|
|
|
|
Args: |
|
|
inputs: inputs to the chain |
|
|
callbacks: callbacks to use for the chain |
|
|
|
|
|
Returns: |
|
|
a Route object |
|
|
""" |
|
|
result = self(inputs, callbacks=callbacks) |
|
|
return Route(result["destination"], result["next_inputs"]) |
|
|
|
|
|
async def aroute( |
|
|
self, inputs: Dict[str, Any], callbacks: Callbacks = None |
|
|
) -> Route: |
|
|
result = await self.acall(inputs, callbacks=callbacks) |
|
|
return Route(result["destination"], result["next_inputs"]) |
|
|
|
|
|
|
|
|
class MultiRouteChain(Chain): |
|
|
"""Use a single chain to route an input to one of multiple candidate chains.""" |
|
|
|
|
|
router_chain: RouterChain |
|
|
"""Chain that routes inputs to destination chains.""" |
|
|
destination_chains: Mapping[str, Chain] |
|
|
"""Chains that return final answer to inputs.""" |
|
|
default_chain: Chain |
|
|
"""Default chain to use when none of the destination chains are suitable.""" |
|
|
silent_errors: bool = False |
|
|
"""If True, use default_chain when an invalid destination name is provided. |
|
|
Defaults to False.""" |
|
|
|
|
|
model_config = ConfigDict( |
|
|
arbitrary_types_allowed=True, |
|
|
extra="forbid", |
|
|
) |
|
|
|
|
|
@property |
|
|
def input_keys(self) -> List[str]: |
|
|
"""Will be whatever keys the router chain prompt expects. |
|
|
|
|
|
:meta private: |
|
|
""" |
|
|
return self.router_chain.input_keys |
|
|
|
|
|
@property |
|
|
def output_keys(self) -> List[str]: |
|
|
"""Will always return text key. |
|
|
|
|
|
:meta private: |
|
|
""" |
|
|
return [] |
|
|
|
|
|
def _call( |
|
|
self, |
|
|
inputs: Dict[str, Any], |
|
|
run_manager: Optional[CallbackManagerForChainRun] = None, |
|
|
) -> Dict[str, Any]: |
|
|
_run_manager = run_manager or CallbackManagerForChainRun.get_noop_manager() |
|
|
callbacks = _run_manager.get_child() |
|
|
route = self.router_chain.route(inputs, callbacks=callbacks) |
|
|
|
|
|
_run_manager.on_text( |
|
|
str(route.destination) + ": " + str(route.next_inputs), verbose=self.verbose |
|
|
) |
|
|
if not route.destination: |
|
|
return self.default_chain(route.next_inputs, callbacks=callbacks) |
|
|
elif route.destination in self.destination_chains: |
|
|
return self.destination_chains[route.destination]( |
|
|
route.next_inputs, callbacks=callbacks |
|
|
) |
|
|
elif self.silent_errors: |
|
|
return self.default_chain(route.next_inputs, callbacks=callbacks) |
|
|
else: |
|
|
raise ValueError( |
|
|
f"Received invalid destination chain name '{route.destination}'" |
|
|
) |
|
|
|
|
|
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() |
|
|
callbacks = _run_manager.get_child() |
|
|
route = await self.router_chain.aroute(inputs, callbacks=callbacks) |
|
|
|
|
|
await _run_manager.on_text( |
|
|
str(route.destination) + ": " + str(route.next_inputs), verbose=self.verbose |
|
|
) |
|
|
if not route.destination: |
|
|
return await self.default_chain.acall( |
|
|
route.next_inputs, callbacks=callbacks |
|
|
) |
|
|
elif route.destination in self.destination_chains: |
|
|
return await self.destination_chains[route.destination].acall( |
|
|
route.next_inputs, callbacks=callbacks |
|
|
) |
|
|
elif self.silent_errors: |
|
|
return await self.default_chain.acall( |
|
|
route.next_inputs, callbacks=callbacks |
|
|
) |
|
|
else: |
|
|
raise ValueError( |
|
|
f"Received invalid destination chain name '{route.destination}'" |
|
|
) |
|
|
|
