_atransform_stream_with_config() — langchain Function Reference

Source Code

libs/core/langchain_core/runnables/base.py lines 2359–2464

    async def _atransform_stream_with_config(
        self,
        inputs: AsyncIterator[Input],
        transformer: Callable[[AsyncIterator[Input]], AsyncIterator[Output]]
        | Callable[
            [AsyncIterator[Input], AsyncCallbackManagerForChainRun],
            AsyncIterator[Output],
        ]
        | Callable[
            [AsyncIterator[Input], AsyncCallbackManagerForChainRun, RunnableConfig],
            AsyncIterator[Output],
        ],
        config: RunnableConfig | None,
        run_type: str | None = None,
        **kwargs: Any | None,
    ) -> AsyncIterator[Output]:
        """Transform a stream with config.

        Helper method to transform an Async `Iterator` of `Input` values into an
        Async `Iterator` of `Output` values, with callbacks.

        Use this to implement `astream` or `atransform` in `Runnable` subclasses.

        """
        # Extract defers_inputs from kwargs if present
        defers_inputs = kwargs.pop("defers_inputs", False)

        # tee the input so we can iterate over it twice
        input_for_tracing, input_for_transform = atee(inputs, 2)
        # Start the input iterator to ensure the input Runnable starts before this one
        final_input: Input | None = await anext(input_for_tracing, None)
        final_input_supported = True
        final_output: Output | None = None
        final_output_supported = True

        config = ensure_config(config)
        callback_manager = get_async_callback_manager_for_config(config)
        run_manager = await callback_manager.on_chain_start(
            None,
            {"input": ""},
            run_type=run_type,
            name=config.get("run_name") or self.get_name(),
            run_id=config.pop("run_id", None),
            defers_inputs=defers_inputs,
        )
        try:
            child_config = patch_config(config, callbacks=run_manager.get_child())
            if accepts_config(transformer):
                kwargs["config"] = child_config
            if accepts_run_manager(transformer):
                kwargs["run_manager"] = run_manager
            with set_config_context(child_config) as context:
                iterator_ = context.run(transformer, input_for_transform, **kwargs)  # type: ignore[arg-type]

                if stream_handler := next(
                    (
                        cast("_StreamingCallbackHandler", h)
                        for h in run_manager.handlers
                        # instance check OK here, it's a mixin
                        if isinstance(h, _StreamingCallbackHandler)
                    ),
                    None,
                ):
                    # populates streamed_output in astream_log() output if needed
                    iterator = stream_handler.tap_output_aiter(
                        run_manager.run_id, iterator_
                    )
                else:
                    iterator = iterator_
                try:
                    while True:
                        chunk = await coro_with_context(anext(iterator), context)
                        yield chunk
                        if final_output_supported:
                            if final_output is None:
                                final_output = chunk
                            else:
                                try:
                                    final_output = final_output + chunk
                                except TypeError:
                                    final_output = chunk