test_openai_functions_router() — langchain Function Reference

Source Code

libs/langchain/tests/unit_tests/runnables/test_openai_functions.py lines 44–97

def test_openai_functions_router(
    snapshot: SnapshotAssertion,
    mocker: MockerFixture,
) -> None:
    revise = mocker.Mock(
        side_effect=lambda kw: f"Revised draft: no more {kw['notes']}!",
    )
    accept = mocker.Mock(side_effect=lambda kw: f"Accepted draft: {kw['draft']}!")

    router = OpenAIFunctionsRouter(
        {
            "revise": revise,
            "accept": accept,
        },
        functions=[
            {
                "name": "revise",
                "description": "Sends the draft for revision.",
                "parameters": {
                    "type": "object",
                    "properties": {
                        "notes": {
                            "type": "string",
                            "description": "The editor's notes to guide the revision.",
                        },
                    },
                },
            },
            {
                "name": "accept",
                "description": "Accepts the draft.",
                "parameters": {
                    "type": "object",
                    "properties": {
                        "draft": {
                            "type": "string",
                            "description": "The draft to accept.",
                        },
                    },
                },
            },
        ],
    )

    model = FakeChatOpenAI()

    chain = model.bind(functions=router.functions) | router

    assert router.functions == snapshot

    assert chain.invoke("Something about turtles?") == "Accepted draft: turtles!"

    revise.assert_not_called()
    accept.assert_called_once_with({"draft": "turtles"})