test_pydantic_output_functions_parser_multiple_schemas() — langchain Function Reference

Source Code

libs/core/tests/unit_tests/output_parsers/test_openai_functions.py lines 167–196

def test_pydantic_output_functions_parser_multiple_schemas() -> None:
    """Test that the parser works if providing multiple pydantic schemas."""
    message = AIMessage(
        content="This is a test message",
        additional_kwargs={
            "function_call": {
                "name": "cookie",
                "arguments": json.dumps({"name": "value", "age": 10}),
            }
        },
    )
    chat_generation = ChatGeneration(message=message)

    class Cookie(BaseModel):
        """Test model."""

        name: str
        age: int

    class Dog(BaseModel):
        """Test model."""

        species: str

    # Full output
    parser = PydanticOutputFunctionsParser(
        pydantic_schema={"cookie": Cookie, "dog": Dog}
    )
    result = parser.parse_result([chat_generation])
    assert result == Cookie(name="value", age=10)