examples/validation_usage.py

"""
Example usage of Pydantic validation in agent tools.

This file demonstrates how to integrate output validation into your tools.
"""

from tools.base import BaseAgentTool
from core.validation import (
    validate_tool_output,
    ensure_tool_output_schema,
    ToolOutput
)


# ============================================================================
# Method 1: Manual validation using _validate_output()
# ============================================================================

class ExampleTool1(BaseAgentTool):
    """
    Example tool using manual validation.

    This approach gives you full control over when validation happens.
    """

    name = "example_tool_1"
    description = "Example tool with manual validation"
    inputs = {
        "query": {"type": "string", "description": "Query string"}
    }
    output_type = "string"

    def forward(self, query: str) -> str:
        """Execute tool with manual validation"""
        try:
            # Your tool logic here
            result = self.process_query(query)

            # Format success (existing method)
            response = self._format_success(
                result,
                metadata={"query": query}
            )

            # Optional: Validate output before returning
            # This ensures output conforms to ToolOutput schema
            return self._validate_output(response)

        except Exception as e:
            # Error formatting already creates valid ToolOutput format
            return self._format_error(
                e,
                recovery_hint="Check query format",
                fallback_action="Try simpler query"
            )

    def process_query(self, query: str):
        """Mock processing logic"""
        return {"data": f"Processed: {query}"}


# ============================================================================
# Method 2: Using the decorator for automatic validation
# ============================================================================

class ExampleTool2(BaseAgentTool):
    """
    Example tool using decorator for automatic validation.

    The decorator wraps the entire forward() method and ensures
    all outputs (success or error) conform to ToolOutput schema.
    """

    name = "example_tool_2"
    description = "Example tool with decorator validation"
    inputs = {
        "data": {"type": "string", "description": "Input data"}
    }
    output_type = "string"

    @ensure_tool_output_schema
    def forward(self, data: str) -> str:
        """Execute tool with automatic validation via decorator"""
        # Your tool logic here
        # No need to handle validation manually - decorator handles it

        result = self.process_data(data)

        # Return using existing format methods or raw dict
        return self._format_success(
            result,
            metadata={"processed": True}
        )

    def process_data(self, data: str):
        """Mock processing logic"""
        if data == "error":
            raise ValueError("Invalid data")
        return {"processed": data.upper()}


# ============================================================================
# Method 3: Direct validation without base class methods
# ============================================================================

class ExampleTool3(BaseAgentTool):
    """
    Example tool using direct Pydantic validation.

    This approach bypasses _format_success and constructs
    ToolOutput objects directly.
    """

    name = "example_tool_3"
    description = "Example tool with direct Pydantic validation"
    inputs = {
        "value": {"type": "integer", "description": "Input value"}
    }
    output_type = "string"

    def forward(self, value: int) -> str:
        """Execute tool with direct Pydantic validation"""
        try:
            # Your tool logic here
            result = value * 2

            # Create ToolOutput directly
            output = ToolOutput(
                success=True,
                result=result,
                metadata={
                    "operation": "multiply",
                    "factor": 2
                }
            )

            # Return as JSON
            return output.model_dump_json(indent=2)

        except Exception as e:
            # Create error ToolOutput directly
            error_output = ToolOutput(
                success=False,
                error=str(e),
                error_type=type(e).__name__,
                recovery_hint="Check input value",
                metadata={"input_value": value}
            )

            return error_output.model_dump_json(indent=2)


# ============================================================================
# Method 4: Handling existing tool outputs (backward compatibility)
# ============================================================================

def wrap_existing_tool_output(tool_output: str) -> str:
    """
    Wrap existing tool output in validation.

    Useful for integrating validation into existing tools
    without modifying their code.

    Args:
        tool_output: Raw output from existing tool

    Returns:
        Validated and potentially repaired output
    """
    validated = validate_tool_output(tool_output)
    return validated.model_dump_json(indent=2)


# ============================================================================
# Example usage
# ============================================================================

if __name__ == "__main__":
    import json

    print("=" * 60)
    print("Validation Usage Examples")
    print("=" * 60)

    # Example 1: Manual validation
    print("\n1. Manual validation:")
    tool1 = ExampleTool1()
    result1 = tool1.forward("test query")
    print(result1)

    # Example 2: Decorator validation
    print("\n2. Decorator validation:")
    tool2 = ExampleTool2()
    result2 = tool2.forward("test data")
    print(result2)

    # Test error case
    print("\n2b. Decorator validation (error case):")
    result2_error = tool2.forward("error")
    print(result2_error)

    # Example 3: Direct Pydantic
    print("\n3. Direct Pydantic validation:")
    tool3 = ExampleTool3()
    result3 = tool3.forward(21)
    print(result3)

    # Example 4: Wrapping existing output
    print("\n4. Wrapping existing output:")
    existing_output = '{"success": true, "result": "legacy data"}'
    wrapped = wrap_existing_tool_output(existing_output)
    print(wrapped)

    # Example 5: Repairing malformed output
    print("\n5. Repairing malformed output:")
    malformed = '{"success": true, "result": incomplete'
    repaired = wrap_existing_tool_output(malformed)
    print(repaired)

    print("\n" + "=" * 60)
    print("All examples completed successfully")
    print("=" * 60)