TestRunnableSequenceParallelTraceNesting Class — langchain Architecture

Source Code

libs/core/tests/unit_tests/runnables/test_tracing_interops.py lines 262–432

class TestRunnableSequenceParallelTraceNesting:
    @pytest.fixture(autouse=True)
    def _setup(self) -> None:
        self.tracer = _create_tracer_with_mocked_client()

    @staticmethod
    def _create_parent(
        other_thing: Callable[
            [int], Generator[int, None, None] | AsyncGenerator[int, None]
        ],
    ) -> RunnableLambda:
        @RunnableLambda
        def my_child_function(a: int) -> int:
            return a + 2

        parallel = RunnableParallel(
            chain_result=my_child_function.with_config(tags=["atag"]),
            other_thing=other_thing,
        )

        def before(x: int) -> int:
            return x

        def after(x: dict[str, Any]) -> int:
            return int(x["chain_result"])

        sequence = before | parallel | after
        if isasyncgenfunction(other_thing):

            @RunnableLambda
            async def parent(a: int) -> int:
                return await sequence.ainvoke(a)

        else:

            @RunnableLambda
            def parent(a: int) -> int:
                return sequence.invoke(a)

        return parent

    def _check_posts(self) -> None:
        posts = _get_posts(self.tracer.client)
        name_order = [
            "parent",
            "RunnableSequence",
            "before",
            "RunnableParallel<chain_result,other_thing>",
            ["my_child_function", "other_thing"],
            "after",
        ]
        expected_parents = {
            "parent": None,
            "RunnableSequence": "parent",
            "before": "RunnableSequence",
            "RunnableParallel<chain_result,other_thing>": "RunnableSequence",
            "my_child_function": "RunnableParallel<chain_result,other_thing>",
            "other_thing": "RunnableParallel<chain_result,other_thing>",
            "after": "RunnableSequence",
        }
        assert len(posts) == sum(
            1 if isinstance(n, str) else len(n) for n in name_order
        )
        prev_dotted_order = None
        dotted_order_map = {}
        id_map = {}
        parent_id_map = {}
        i = 0
        for name in name_order:
            if isinstance(name, list):
                for n in name:
                    matching_post = next(
                        p for p in posts[i : i + len(name)] if p["name"] == n
                    )
                    assert matching_post
                    dotted_order = matching_post["dotted_order"]
                    if prev_dotted_order is not None:
                        assert dotted_order > prev_dotted_order
                    dotted_order_map[n] = dotted_order
                    id_map[n] = matching_post["id"]
                    parent_id_map[n] = matching_post.get("parent_run_id")