|
|
"""Chain that runs an arbitrary python function.""" |
|
|
|
|
|
import functools |
|
|
import logging |
|
|
from typing import Any, Awaitable, Callable, Dict, List, Optional |
|
|
|
|
|
from langchain_core.callbacks import ( |
|
|
AsyncCallbackManagerForChainRun, |
|
|
CallbackManagerForChainRun, |
|
|
) |
|
|
from pydantic import Field |
|
|
|
|
|
from langchain.chains.base import Chain |
|
|
|
|
|
logger = logging.getLogger(__name__) |
|
|
|
|
|
|
|
|
class TransformChain(Chain): |
|
|
"""Chain that transforms the chain output. |
|
|
|
|
|
Example: |
|
|
.. code-block:: python |
|
|
|
|
|
from langchain.chains import TransformChain |
|
|
transform_chain = TransformChain(input_variables=["text"], |
|
|
output_variables["entities"], transform=func()) |
|
|
""" |
|
|
|
|
|
input_variables: List[str] |
|
|
"""The keys expected by the transform's input dictionary.""" |
|
|
output_variables: List[str] |
|
|
"""The keys returned by the transform's output dictionary.""" |
|
|
transform_cb: Callable[[Dict[str, str]], Dict[str, str]] = Field(alias="transform") |
|
|
"""The transform function.""" |
|
|
atransform_cb: Optional[Callable[[Dict[str, Any]], Awaitable[Dict[str, Any]]]] = ( |
|
|
Field(None, alias="atransform") |
|
|
) |
|
|
"""The async coroutine transform function.""" |
|
|
|
|
|
@staticmethod |
|
|
@functools.lru_cache |
|
|
def _log_once(msg: str) -> None: |
|
|
"""Log a message once. |
|
|
|
|
|
:meta private: |
|
|
""" |
|
|
logger.warning(msg) |
|
|
|
|
|
@property |
|
|
def input_keys(self) -> List[str]: |
|
|
"""Expect input keys. |
|
|
|
|
|
:meta private: |
|
|
""" |
|
|
return self.input_variables |
|
|
|
|
|
@property |
|
|
def output_keys(self) -> List[str]: |
|
|
"""Return output keys. |
|
|
|
|
|
:meta private: |
|
|
""" |
|
|
return self.output_variables |
|
|
|
|
|
def _call( |
|
|
self, |
|
|
inputs: Dict[str, str], |
|
|
run_manager: Optional[CallbackManagerForChainRun] = None, |
|
|
) -> Dict[str, str]: |
|
|
return self.transform_cb(inputs) |
|
|
|
|
|
async def _acall( |
|
|
self, |
|
|
inputs: Dict[str, Any], |
|
|
run_manager: Optional[AsyncCallbackManagerForChainRun] = None, |
|
|
) -> Dict[str, Any]: |
|
|
if self.atransform_cb is not None: |
|
|
return await self.atransform_cb(inputs) |
|
|
else: |
|
|
self._log_once( |
|
|
"TransformChain's atransform is not provided, falling" |
|
|
" back to synchronous transform" |
|
|
) |
|
|
return self.transform_cb(inputs) |
|
|
|
