""" 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)