langchain[patch]: add aprep_output method to langchain/chains/base.py (#20748)

## Description

Add `aprep_output` method to `langchain/chains/base.py`. Some downstream
`ChatMessageHistory` objects that use async connections require an async
way to append to the context.

It turned out that `ainvoke()` was calling `prep_output` which is
synchronous.

---------

Co-authored-by: Bagatur <22008038+baskaryan@users.noreply.github.com>
This commit is contained in:
Alex Lee 2024-04-24 15:16:25 -07:00 committed by GitHub
parent 43c041cda5
commit 243ba71b28
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194

View File

@ -204,7 +204,7 @@ class Chain(RunnableSerializable[Dict[str, Any], Dict[str, Any]], ABC):
if new_arg_supported
else await self._acall(inputs)
)
final_outputs: Dict[str, Any] = self.prep_outputs(
final_outputs: Dict[str, Any] = await self.aprep_outputs(
inputs, outputs, return_only_outputs
)
except BaseException as e:
@ -458,6 +458,32 @@ class Chain(RunnableSerializable[Dict[str, Any], Dict[str, Any]], ABC):
else:
return {**inputs, **outputs}
async def aprep_outputs(
self,
inputs: Dict[str, str],
outputs: Dict[str, str],
return_only_outputs: bool = False,
) -> Dict[str, str]:
"""Validate and prepare chain outputs, and save info about this run to memory.
Args:
inputs: Dictionary of chain inputs, including any inputs added by chain
memory.
outputs: Dictionary of initial chain outputs.
return_only_outputs: Whether to only return the chain outputs. If False,
inputs are also added to the final outputs.
Returns:
A dict of the final chain outputs.
"""
self._validate_outputs(outputs)
if self.memory is not None:
await self.memory.asave_context(inputs, outputs)
if return_only_outputs:
return outputs
else:
return {**inputs, **outputs}
def prep_inputs(self, inputs: Union[Dict[str, Any], Any]) -> Dict[str, str]:
"""Prepare chain inputs, including adding inputs from memory.