test_higher_order_lambda_runnable_async() — langchain Function Reference

Architecture documentation for the test_higher_order_lambda_runnable_async() function in test_runnable.py from the langchain codebase.

Function python LangChainCore Runnables calls 1

Entity Profile

LangChainCore→ Runnables→ test_higher_order_lambda_runnable_async() — langchain Function Reference

Dependency Diagram

graph TD
  d3230587_b866_4309_a37d_9fdfe73dd2ed["test_higher_order_lambda_runnable_async()"]
  26df6ad8_0189_51d0_c3c1_6c3248893ff5["test_runnable.py"]
  d3230587_b866_4309_a37d_9fdfe73dd2ed -->|defined in| 26df6ad8_0189_51d0_c3c1_6c3248893ff5
  8652094c_ec57_c551_fc44_9566d00cf872["abatch()"]
  d3230587_b866_4309_a37d_9fdfe73dd2ed -->|calls| 8652094c_ec57_c551_fc44_9566d00cf872
  style d3230587_b866_4309_a37d_9fdfe73dd2ed fill:#6366f1,stroke:#818cf8,color:#fff

Relationship Graph

Source Code

libs/core/tests/unit_tests/runnables/test_runnable.py lines 2993–3056

async def test_higher_order_lambda_runnable_async(mocker: MockerFixture) -> None:
    math_chain = ChatPromptTemplate.from_template(
        "You are a math genius. Answer the question: {question}"
    ) | FakeListLLM(responses=["4"])
    english_chain = ChatPromptTemplate.from_template(
        "You are an english major. Answer the question: {question}"
    ) | FakeListLLM(responses=["2"])
    input_map = RunnableParallel(
        key=lambda x: x["key"],
        input={"question": lambda x: x["question"]},
    )

    def router(value: dict[str, Any]) -> Runnable:
        if value["key"] == "math":
            return itemgetter("input") | math_chain
        if value["key"] == "english":
            return itemgetter("input") | english_chain
        msg = f"Unknown key: {value['key']}"
        raise ValueError(msg)

    chain: Runnable = input_map | router

    result = await chain.ainvoke({"key": "math", "question": "2 + 2"})
    assert result == "4"

    result2 = await chain.abatch(
        [
            {"key": "math", "question": "2 + 2"},
            {"key": "english", "question": "2 + 2"},
        ]
    )
    assert result2 == ["4", "2"]

    # Test ainvoke
    async def arouter(params: dict[str, Any]) -> Runnable:
        if params["key"] == "math":
            return itemgetter("input") | math_chain
        if params["key"] == "english":
            return itemgetter("input") | english_chain
        msg = f"Unknown key: {params['key']}"
        raise ValueError(msg)

    achain: Runnable = input_map | arouter
    math_spy = mocker.spy(math_chain.__class__, "ainvoke")
    tracer = FakeTracer()
    assert (
        await achain.ainvoke(
            {"key": "math", "question": "2 + 2"}, {"callbacks": [tracer]}
        )
        == "4"
    )
    assert math_spy.call_args.args[1] == {
        "key": "math",
        "input": {"question": "2 + 2"},
    }
    assert len([r for r in tracer.runs if r.parent_run_id is None]) == 1
    parent_run = next(r for r in tracer.runs if r.parent_run_id is None)
    assert len(parent_run.child_runs) == 2
    router_run = parent_run.child_runs[1]
    assert router_run.name == "arouter"
    assert len(router_run.child_runs) == 1
    math_run = router_run.child_runs[0]
    assert math_run.name == "RunnableSequence"
    assert len(math_run.child_runs) == 3