mirror of
https://github.com/hwchase17/langchain.git
synced 2025-07-09 14:35:50 +00:00
Add callback function to RunnablePassthrough (#11564)
<!-- Thank you for contributing to LangChain! Replace this entire comment with: - **Description:** a description of the change, - **Issue:** the issue # it fixes (if applicable), - **Dependencies:** any dependencies required for this change, - **Tag maintainer:** for a quicker response, tag the relevant maintainer (see below), - **Twitter handle:** we announce bigger features on Twitter. If your PR gets announced, and you'd like a mention, we'll gladly shout you out! Please make sure your PR is passing linting and testing before submitting. Run `make format`, `make lint` and `make test` to check this locally. See contribution guidelines for more information on how to write/run tests, lint, etc: https://github.com/langchain-ai/langchain/blob/master/.github/CONTRIBUTING.md If you're adding a new integration, please include: 1. a test for the integration, preferably unit tests that do not rely on network access, 2. an example notebook showing its use. It lives in `docs/extras` directory. If no one reviews your PR within a few days, please @-mention one of @baskaryan, @eyurtsev, @hwchase17. -->
This commit is contained in:
parent
7f4734c0dd
commit
ca9de26f2b
@ -2,10 +2,12 @@
|
|||||||
from __future__ import annotations
|
from __future__ import annotations
|
||||||
|
|
||||||
import asyncio
|
import asyncio
|
||||||
|
import inspect
|
||||||
import threading
|
import threading
|
||||||
from typing import (
|
from typing import (
|
||||||
Any,
|
Any,
|
||||||
AsyncIterator,
|
AsyncIterator,
|
||||||
|
Awaitable,
|
||||||
Callable,
|
Callable,
|
||||||
Dict,
|
Dict,
|
||||||
Iterator,
|
Iterator,
|
||||||
@ -100,6 +102,26 @@ class RunnablePassthrough(RunnableSerializable[Input, Input]):
|
|||||||
|
|
||||||
input_type: Optional[Type[Input]] = None
|
input_type: Optional[Type[Input]] = None
|
||||||
|
|
||||||
|
func: Optional[Callable[[Input], None]] = None
|
||||||
|
|
||||||
|
afunc: Optional[Callable[[Input], Awaitable[None]]] = None
|
||||||
|
|
||||||
|
def __init__(
|
||||||
|
self,
|
||||||
|
func: Optional[
|
||||||
|
Union[Callable[[Input], None], Callable[[Input], Awaitable[None]]]
|
||||||
|
] = None,
|
||||||
|
afunc: Optional[Callable[[Input], Awaitable[None]]] = None,
|
||||||
|
*,
|
||||||
|
input_type: Optional[Type[Input]] = None,
|
||||||
|
**kwargs: Any,
|
||||||
|
) -> None:
|
||||||
|
if inspect.iscoroutinefunction(func):
|
||||||
|
afunc = func
|
||||||
|
func = None
|
||||||
|
|
||||||
|
super().__init__(func=func, afunc=afunc, input_type=input_type, **kwargs)
|
||||||
|
|
||||||
@classmethod
|
@classmethod
|
||||||
def is_lc_serializable(cls) -> bool:
|
def is_lc_serializable(cls) -> bool:
|
||||||
return True
|
return True
|
||||||
@ -140,6 +162,8 @@ class RunnablePassthrough(RunnableSerializable[Input, Input]):
|
|||||||
return RunnableAssign(RunnableParallel(kwargs))
|
return RunnableAssign(RunnableParallel(kwargs))
|
||||||
|
|
||||||
def invoke(self, input: Input, config: Optional[RunnableConfig] = None) -> Input:
|
def invoke(self, input: Input, config: Optional[RunnableConfig] = None) -> Input:
|
||||||
|
if self.func is not None:
|
||||||
|
self.func(input)
|
||||||
return self._call_with_config(identity, input, config)
|
return self._call_with_config(identity, input, config)
|
||||||
|
|
||||||
async def ainvoke(
|
async def ainvoke(
|
||||||
@ -148,6 +172,10 @@ class RunnablePassthrough(RunnableSerializable[Input, Input]):
|
|||||||
config: Optional[RunnableConfig] = None,
|
config: Optional[RunnableConfig] = None,
|
||||||
**kwargs: Optional[Any],
|
**kwargs: Optional[Any],
|
||||||
) -> Input:
|
) -> Input:
|
||||||
|
if self.afunc is not None:
|
||||||
|
await self.afunc(input, **kwargs)
|
||||||
|
elif self.func is not None:
|
||||||
|
self.func(input, **kwargs)
|
||||||
return await self._acall_with_config(aidentity, input, config)
|
return await self._acall_with_config(aidentity, input, config)
|
||||||
|
|
||||||
def transform(
|
def transform(
|
||||||
@ -156,7 +184,21 @@ class RunnablePassthrough(RunnableSerializable[Input, Input]):
|
|||||||
config: Optional[RunnableConfig] = None,
|
config: Optional[RunnableConfig] = None,
|
||||||
**kwargs: Any,
|
**kwargs: Any,
|
||||||
) -> Iterator[Input]:
|
) -> Iterator[Input]:
|
||||||
return self._transform_stream_with_config(input, identity, config)
|
if self.func is None:
|
||||||
|
for chunk in self._transform_stream_with_config(input, identity, config):
|
||||||
|
yield chunk
|
||||||
|
else:
|
||||||
|
final = None
|
||||||
|
|
||||||
|
for chunk in self._transform_stream_with_config(input, identity, config):
|
||||||
|
yield chunk
|
||||||
|
if final is None:
|
||||||
|
final = chunk
|
||||||
|
else:
|
||||||
|
final = final + chunk
|
||||||
|
|
||||||
|
if final is not None:
|
||||||
|
self.func(final, **kwargs)
|
||||||
|
|
||||||
async def atransform(
|
async def atransform(
|
||||||
self,
|
self,
|
||||||
@ -164,7 +206,47 @@ class RunnablePassthrough(RunnableSerializable[Input, Input]):
|
|||||||
config: Optional[RunnableConfig] = None,
|
config: Optional[RunnableConfig] = None,
|
||||||
**kwargs: Any,
|
**kwargs: Any,
|
||||||
) -> AsyncIterator[Input]:
|
) -> AsyncIterator[Input]:
|
||||||
async for chunk in self._atransform_stream_with_config(input, identity, config):
|
if self.afunc is None and self.func is None:
|
||||||
|
async for chunk in self._atransform_stream_with_config(
|
||||||
|
input, identity, config
|
||||||
|
):
|
||||||
|
yield chunk
|
||||||
|
else:
|
||||||
|
final = None
|
||||||
|
|
||||||
|
async for chunk in self._atransform_stream_with_config(
|
||||||
|
input, identity, config
|
||||||
|
):
|
||||||
|
yield chunk
|
||||||
|
if final is None:
|
||||||
|
final = chunk
|
||||||
|
else:
|
||||||
|
final = final + chunk
|
||||||
|
|
||||||
|
if final is not None:
|
||||||
|
if self.afunc is not None:
|
||||||
|
await self.afunc(final, **kwargs)
|
||||||
|
elif self.func is not None:
|
||||||
|
self.func(final, **kwargs)
|
||||||
|
|
||||||
|
def stream(
|
||||||
|
self,
|
||||||
|
input: Input,
|
||||||
|
config: Optional[RunnableConfig] = None,
|
||||||
|
**kwargs: Any,
|
||||||
|
) -> Iterator[Input]:
|
||||||
|
return self.transform(iter([input]), config, **kwargs)
|
||||||
|
|
||||||
|
async def astream(
|
||||||
|
self,
|
||||||
|
input: Input,
|
||||||
|
config: Optional[RunnableConfig] = None,
|
||||||
|
**kwargs: Any,
|
||||||
|
) -> AsyncIterator[Input]:
|
||||||
|
async def input_aiter() -> AsyncIterator[Input]:
|
||||||
|
yield input
|
||||||
|
|
||||||
|
async for chunk in self.atransform(input_aiter(), config, **kwargs):
|
||||||
yield chunk
|
yield chunk
|
||||||
|
|
||||||
|
|
||||||
|
@ -3281,7 +3281,11 @@
|
|||||||
"runnable",
|
"runnable",
|
||||||
"RunnablePassthrough"
|
"RunnablePassthrough"
|
||||||
],
|
],
|
||||||
"kwargs": {}
|
"kwargs": {
|
||||||
|
"func": null,
|
||||||
|
"afunc": null,
|
||||||
|
"input_type": null
|
||||||
|
}
|
||||||
},
|
},
|
||||||
"last": {
|
"last": {
|
||||||
"lc": 1,
|
"lc": 1,
|
||||||
|
@ -1016,6 +1016,34 @@ def test_configurable_fields_example() -> None:
|
|||||||
)
|
)
|
||||||
|
|
||||||
|
|
||||||
|
@pytest.mark.asyncio
|
||||||
|
async def test_passthrough_tap_async(mocker: MockerFixture) -> None:
|
||||||
|
fake = FakeRunnable()
|
||||||
|
mock = mocker.Mock()
|
||||||
|
|
||||||
|
seq: Runnable = fake | RunnablePassthrough(mock)
|
||||||
|
|
||||||
|
assert await seq.ainvoke("hello") == 5
|
||||||
|
assert mock.call_args_list == [mocker.call(5)]
|
||||||
|
mock.reset_mock()
|
||||||
|
|
||||||
|
assert [
|
||||||
|
part async for part in seq.astream("hello", dict(metadata={"key": "value"}))
|
||||||
|
] == [5]
|
||||||
|
assert mock.call_args_list == [mocker.call(5)]
|
||||||
|
mock.reset_mock()
|
||||||
|
|
||||||
|
assert seq.invoke("hello") == 5
|
||||||
|
assert mock.call_args_list == [mocker.call(5)]
|
||||||
|
mock.reset_mock()
|
||||||
|
|
||||||
|
assert [part for part in seq.stream("hello", dict(metadata={"key": "value"}))] == [
|
||||||
|
5
|
||||||
|
]
|
||||||
|
assert mock.call_args_list == [mocker.call(5)]
|
||||||
|
mock.reset_mock()
|
||||||
|
|
||||||
|
|
||||||
@pytest.mark.asyncio
|
@pytest.mark.asyncio
|
||||||
async def test_with_config(mocker: MockerFixture) -> None:
|
async def test_with_config(mocker: MockerFixture) -> None:
|
||||||
fake = FakeRunnable()
|
fake = FakeRunnable()
|
||||||
|
Loading…
Reference in New Issue
Block a user