567-labs--instructor/tests/llm/test_genai/test_schema_conversion.py

"""Test schema conversion functions for Gemini."""

from enum import Enum
from typing import Optional
from pydantic import BaseModel

from instructor.v2.providers.gemini.utils import (
    map_to_gemini_function_schema,
    verify_no_unions,
)


class Priority(Enum):
    LOW = "low"
    MEDIUM = "medium"
    HIGH = "high"


class SimpleModel(BaseModel):
    name: str
    age: int
    is_active: bool


class OptionalModel(BaseModel):
    name: str
    age: Optional[int] = None
    description: Optional[str] = None


class EnumModel(BaseModel):
    name: str
    priority: Priority


class NestedModel(BaseModel):
    name: str
    items: list[str]
    details: SimpleModel


def test_simple_schema_conversion():
    """Test conversion strips extra pydantic fields like 'title'."""
    schema = SimpleModel.model_json_schema()
    result = map_to_gemini_function_schema(schema)

    # Input has 'title' fields that should be stripped out
    assert "title" in schema
    assert "title" in schema["properties"]["name"]

    # Output should be clean without title fields
    expected = {
        "type": "object",
        "properties": {
            "name": {"type": "string"},
            "age": {"type": "integer"},
            "is_active": {"type": "boolean"},
        },
        "required": ["name", "age", "is_active"],
    }

    assert result == expected


def test_optional_schema_conversion():
    """Test conversion transforms anyOf[T, null] to nullable fields."""
    schema = OptionalModel.model_json_schema()
    result = map_to_gemini_function_schema(schema)

    # Input should have anyOf with null type for optional fields
    assert schema["properties"]["age"]["anyOf"] == [
        {"type": "integer"},
        {"type": "null"},
    ]
    assert schema["properties"]["description"]["anyOf"] == [
        {"type": "string"},
        {"type": "null"},
    ]

    # Output should convert to nullable: true
    expected = {
        "type": "object",
        "properties": {
            "name": {"type": "string"},
            "age": {"type": "integer", "nullable": True},
            "description": {"type": "string", "nullable": True},
        },
        "required": ["name"],
    }

    assert result == expected


def test_enum_schema_conversion():
    """Test conversion resolves $refs and adds format: enum."""
    schema = EnumModel.model_json_schema()
    result = map_to_gemini_function_schema(schema)

    # Input should have $ref and $defs
    assert schema["properties"]["priority"]["$ref"] == "#/$defs/Priority"
    assert "$defs" in schema
    assert schema["$defs"]["Priority"]["enum"] == ["low", "medium", "high"]

    # Output should resolve the ref and add format: enum
    expected = {
        "type": "object",
        "properties": {
            "name": {"type": "string"},
            "priority": {
                "type": "string",
                "enum": ["low", "medium", "high"],
                "format": "enum",
            },
        },
        "required": ["name", "priority"],
    }

    assert result == expected


def test_nested_schema_conversion():
    """Test conversion of schema with nested objects."""
    schema = NestedModel.model_json_schema()
    result = map_to_gemini_function_schema(schema)

    expected = {
        "type": "object",
        "properties": {
            "name": {"type": "string"},
            "items": {"type": "array", "items": {"type": "string"}},
            "details": {
                "type": "object",
                "properties": {
                    "name": {"type": "string"},
                    "age": {"type": "integer"},
                    "is_active": {"type": "boolean"},
                },
                "required": ["name", "age", "is_active"],
            },
        },
        "required": ["name", "items", "details"],
    }

    assert result == expected


def test_verify_no_unions_valid():
    """Test verify_no_unions with valid schemas."""
    # Simple schema should pass
    simple_schema = SimpleModel.model_json_schema()
    assert verify_no_unions(simple_schema) is True

    # Optional schema should pass (Optional[T] is Union[T, None])
    optional_schema = OptionalModel.model_json_schema()
    assert verify_no_unions(optional_schema) is True


def test_verify_no_unions_invalid():
    """Test verify_no_unions with union schemas (now allowed)."""
    # Create a schema with a true union (not just Optional)
    invalid_schema = {
        "type": "object",
        "properties": {"value": {"anyOf": [{"type": "string"}, {"type": "integer"}]}},
    }
    assert verify_no_unions(invalid_schema) is True


def test_schema_without_refs():
    """Test schema conversion without $refs."""
    schema = {
        "type": "object",
        "properties": {"name": {"type": "string"}, "count": {"type": "integer"}},
        "required": ["name"],
    }

    result = map_to_gemini_function_schema(schema)

    expected = {
        "type": "object",
        "properties": {"name": {"type": "string"}, "count": {"type": "integer"}},
        "required": ["name"],
    }

    assert result == expected


def test_schema_with_description():
    """Test schema conversion preserves descriptions."""
    schema = {
        "type": "object",
        "description": "A test object",
        "properties": {"name": {"type": "string", "description": "The name field"}},
    }

    result = map_to_gemini_function_schema(schema)

    expected = {
        "type": "object",
        "description": "A test object",
        "properties": {"name": {"type": "string", "description": "The name field"}},
    }

    assert result == expected


def test_union_type_raises_error():
    """Test that union types are allowed in schema conversion."""
    # Create a model with a true union type (not Optional or Decimal)
    union_schema = {
        "type": "object",
        "properties": {"value": {"anyOf": [{"type": "string"}, {"type": "integer"}]}},
    }

    result = map_to_gemini_function_schema(union_schema)
    assert result["properties"]["value"]["anyOf"] == [
        {"type": "string"},
        {"type": "integer"},
    ]


def test_verify_no_unions_allows_optional():
    """Test that verify_no_unions allows Optional types."""
    # Schema with Optional field (Union with null)
    optional_schema = {
        "type": "object",
        "properties": {
            "name": {"type": "string"},
            "age": {"anyOf": [{"type": "integer"}, {"type": "null"}]},
        },
    }

    assert verify_no_unions(optional_schema) is True


def test_verify_no_unions_allows_decimal():
    """Test that verify_no_unions allows Decimal types (string | number)."""
    # Schema with Decimal field (Union of string and number)
    decimal_schema = {
        "type": "object",
        "properties": {
            "total": {"anyOf": [{"type": "number"}, {"type": "string"}]},
            "price": {
                "anyOf": [{"type": "string"}, {"type": "number"}]
            },  # Order shouldn't matter
        },
    }

    assert verify_no_unions(decimal_schema) is True


def test_verify_no_unions_rejects_other_unions():
    """Test that verify_no_unions allows non-Optional unions."""
    # Schema with unsupported union type (string | integer)
    union_schema = {
        "type": "object",
        "properties": {"value": {"anyOf": [{"type": "string"}, {"type": "integer"}]}},
    }

    assert verify_no_unions(union_schema) is True


def test_verify_no_unions_rejects_complex_unions():
    """Test that verify_no_unions allows complex union types."""
    # Schema with more than 2 types in union
    complex_union_schema = {
        "type": "object",
        "properties": {
            "value": {
                "anyOf": [{"type": "string"}, {"type": "integer"}, {"type": "boolean"}]
            }
        },
    }

    assert verify_no_unions(complex_union_schema) is True


def test_verify_no_unions_nested_schemas():
    """Test that verify_no_unions allows unions in nested schemas."""
    # Schema with nested object containing Decimal and Optional fields
    nested_schema = {
        "type": "object",
        "properties": {
            "receipt": {
                "type": "object",
                "properties": {
                    "total": {
                        "anyOf": [{"type": "number"}, {"type": "string"}]
                    },  # Decimal - should pass
                    "notes": {
                        "anyOf": [{"type": "string"}, {"type": "null"}]
                    },  # Optional - should pass
                },
            }
        },
    }

    assert verify_no_unions(nested_schema) is True

    # Schema with nested object containing unsupported union
    bad_nested_schema = {
        "type": "object",
        "properties": {
            "receipt": {
                "type": "object",
                "properties": {
                    "total": {
                        "anyOf": [{"type": "number"}, {"type": "string"}]
                    },  # Decimal - should pass
                    "status": {
                        "anyOf": [{"type": "string"}, {"type": "integer"}]
                    },  # Bad union - should fail
                },
            }
        },
    }

    assert verify_no_unions(bad_nested_schema) is True


def test_decimal_schema_conversion_succeeds():
    """Test that Decimal types (string | number) are successfully converted."""
    # Schema representing a Receipt with Decimal total field
    decimal_schema = {
        "type": "object",
        "title": "Receipt",
        "properties": {
            "total": {
                "anyOf": [{"type": "number"}, {"type": "string"}],
                "title": "Total",
            }
        },
        "required": ["total"],
    }

    # This should not raise an error now
    result = map_to_gemini_function_schema(decimal_schema)

    # The conversion should succeed and preserve the anyOf structure
    assert result["type"] == "object"
    assert result["properties"]["total"]["anyOf"] == [
        {"type": "number"},
        {"type": "string"},
    ]
    assert result["required"] == ["total"]
    # Title should be stripped out
    assert "title" not in result
    assert "title" not in result["properties"]["total"]