langgraph/libs/prebuilt/tests/test_responses.py

"""Unit tests for langgraph.prebuilt.responses module."""

from typing import Union

import pytest
from pydantic import BaseModel

from langgraph.prebuilt.responses import (
    OutputToolBinding,
    ToolOutput,
    _SchemaSpec,
)


class _TestModel(BaseModel):
    """A test model for structured output."""

    name: str
    age: int
    email: str = "default@example.com"


class CustomModel(BaseModel):
    """Custom model with a custom docstring."""

    value: float
    description: str


class EmptyDocModel(BaseModel):
    # No custom docstring, should have no description in tool
    data: str


class TestUsingToolStrategy:
    """Test UsingToolStrategy dataclass."""

    def test_basic_creation(self):
        """Test basic UsingToolStrategy creation."""
        strategy = ToolOutput(schema=_TestModel)
        assert strategy.schema == _TestModel
        assert strategy.tool_message_content == "ok!"
        assert len(strategy.schema_specs) == 1

    def test_multiple_schemas(self):
        """Test UsingToolStrategy with multiple schemas."""
        strategy = ToolOutput(schema=Union[_TestModel, CustomModel])
        assert len(strategy.schema_specs) == 2
        assert strategy.schema_specs[0].schema == _TestModel
        assert strategy.schema_specs[1].schema == CustomModel

    def test_schema_with_tool_message_content(self):
        """Test UsingToolStrategy with tool message content."""
        strategy = ToolOutput(schema=_TestModel, tool_message_content="custom message")
        assert strategy.schema == _TestModel
        assert strategy.tool_message_content == "custom message"
        assert len(strategy.schema_specs) == 1


class TestOutputToolBinding:
    """Test OutputToolBinding dataclass and its methods."""

    def test_from_schema_spec_basic(self):
        """Test basic OutputToolBinding creation from SchemaSpec."""
        schema_spec = _SchemaSpec(schema=_TestModel)
        tool_binding = OutputToolBinding.from_schema_spec(schema_spec)

        assert tool_binding.schema == _TestModel
        assert tool_binding.schema_kind == "pydantic"
        assert tool_binding.tool is not None
        assert tool_binding.tool.name == "_TestModel"

    def test_from_schema_spec_with_custom_name(self):
        """Test OutputToolBinding creation with custom name."""
        schema_spec = _SchemaSpec(schema=_TestModel, name="custom_tool_name")
        tool_binding = OutputToolBinding.from_schema_spec(schema_spec)
        assert tool_binding.tool.name == "custom_tool_name"

    def test_from_schema_spec_with_custom_description(self):
        """Test OutputToolBinding creation with custom description."""
        schema_spec = _SchemaSpec(
            schema=_TestModel, description="Custom tool description"
        )
        tool_binding = OutputToolBinding.from_schema_spec(schema_spec)

        assert tool_binding.tool.description == "Custom tool description"

    def test_from_schema_spec_with_model_docstring(self):
        """Test OutputToolBinding creation using model docstring as description."""
        schema_spec = _SchemaSpec(schema=CustomModel)
        tool_binding = OutputToolBinding.from_schema_spec(schema_spec)

        assert tool_binding.tool.description == "Custom model with a custom docstring."

    @pytest.mark.skip(
        reason="Need to fix bug in langchain-core for inheritance of doc-strings."
    )
    def test_from_schema_spec_empty_docstring(self):
        """Test OutputToolBinding creation with model that has default docstring."""

        # Create a model with the same docstring as BaseModel
        class DefaultDocModel(BaseModel):
            # This should have the same docstring as BaseModel
            pass

        schema_spec = _SchemaSpec(schema=DefaultDocModel)
        tool_binding = OutputToolBinding.from_schema_spec(schema_spec)

        # Should use empty description when model has default BaseModel docstring
        assert tool_binding.tool.description == ""

    def test_parse_payload_pydantic_success(self):
        """Test successful parsing for Pydantic model."""
        schema_spec = _SchemaSpec(schema=_TestModel)
        tool_binding = OutputToolBinding.from_schema_spec(schema_spec)

        tool_args = {"name": "John", "age": 30}
        result = tool_binding.parse(tool_args)

        assert isinstance(result, _TestModel)
        assert result.name == "John"
        assert result.age == 30
        assert result.email == "default@example.com"  # default value

    def test_parse_payload_pydantic_validation_error(self):
        """Test parsing failure for invalid Pydantic data."""
        schema_spec = _SchemaSpec(schema=_TestModel)
        tool_binding = OutputToolBinding.from_schema_spec(schema_spec)

        # Missing required field 'name'
        tool_args = {"age": 30}

        with pytest.raises(ValueError, match="Failed to parse tool args to _TestModel"):
            tool_binding.parse(tool_args)

    def test_parse_payload_invalid_kind(self):
        """Test parsing with invalid kind."""
        from unittest.mock import Mock

        mock_tool = Mock()

        tool_binding = OutputToolBinding(
            schema=_TestModel,
            schema_kind="invalid_kind",  # type: ignore
            tool=mock_tool,
        )

        with pytest.raises(ValueError, match="Unsupported schema kind: invalid_kind"):
            tool_binding.parse({"name": "test", "age": 25})

    def test_parse_payload_invalid_pydantic_schema(self):
        """Test parsing with invalid schema for pydantic kind."""
        from unittest.mock import Mock

        mock_tool = Mock()

        # Create tool binding with dict schema but pydantic kind
        tool_binding = OutputToolBinding(
            schema={"type": "object"}, schema_kind="pydantic", tool=mock_tool
        )

        with pytest.raises(
            ValueError, match="Expected Pydantic model class for 'pydantic' kind"
        ):
            tool_binding.parse({"name": "test", "age": 25})


class TestEdgeCases:
    """Test edge cases and error conditions."""

    def test_empty_schemas_list(self) -> None:
        """Test UsingToolStrategy with empty schemas list."""
        strategy = ToolOutput(EmptyDocModel)
        assert len(strategy.schema_specs) == 1

    @pytest.mark.skip(
        reason="Need to fix bug in langchain-core for inheritance of doc-strings."
    )
    def test_base_model_doc_constant(self) -> None:
        """Test that BASE_MODEL_DOC constant is set correctly."""
        binding = OutputToolBinding.from_schema_spec(_SchemaSpec(EmptyDocModel))
        assert binding.tool.name == "EmptyDocModel"
        assert (
            binding.tool.description[:5] == ""
        )  # Should be empty for default docstring