test_return_direct_integration_matrix() — langchain Function Reference

Source Code

libs/langchain_v1/tests/unit_tests/agents/test_return_direct_spec.py lines 76–118

def test_return_direct_integration_matrix(case: TestCase) -> None:
    poll_tool = _make_tool(return_direct=case.return_direct)

    model = ChatOpenAI(
        model="gpt-4o",
        temperature=0,
    )

    if case.response_format:
        agent = create_agent(
            model,
            tools=[poll_tool["tool"]],
            system_prompt=AGENT_PROMPT,
            response_format=ToolStrategy(case.response_format),
        )
    else:
        agent = create_agent(
            model,
            tools=[poll_tool["tool"]],
            system_prompt=AGENT_PROMPT,
        )

    result = agent.invoke(
        {
            "messages": [
                HumanMessage("Poll the job until it's done and tell me how many attempts it took.")
            ]
        }
    )

    # Count tool calls
    assert poll_tool["mock"].call_count == case.expected_tool_calls

    # Check last message content
    last_message = result["messages"][-1]
    assert last_message.content == case.expected_last_message

    # Check structured response
    if case.expected_structured_response is not None:
        structured_response_json = result["structured_response"]
        assert structured_response_json == case.expected_structured_response
    else:
        assert "structured_response" not in result