Add query/search agent with advanced mapping, validation, error reporting, and JSON support

query_search_agent.py

@@ -0,0 +1,124 @@
+"""Query/Search Agent Tool with Dynamic Column Mapping and JSON Support
+
+This agent provides query and search capabilities over registered schemas,
+performs dynamic mapping, validates target fields, and supports error handling
+and JSON round-trip.
+"""
+
+import json
+from typing import Any, Dict, List, Tuple, Optional
+from schema_registry import SchemaRegistry
+
+class QuerySearchAgent:
+    """Query/search agent leveraging the SchemaRegistry's mapping and validation."""
+    def __init__(self, registry: Optional[SchemaRegistry] = None):
+        self.registry = registry or SchemaRegistry()
+
+    def search(self, schema_name: str, query: Dict[str, Any], data_source: List[Dict[str, Any]], 
+               container_map: Optional[Dict[str, str]] = None) -> Tuple[List[Dict[str, Any]], List[str]]:
+        """Searches for items matching query constraints, supports dynamic container mapping."""
+        errors = []
+        results = []
+        
+        # Validate schema
+        schema = self.registry.get_schema(schema_name)
+        if not schema:
+            errors.append(f"Schema '{schema_name}' not found")
+            return [], errors
+        
+        # Validate query against schema
+        _, query_errors = self.registry.validate_data(schema_name, query)
+        if query_errors:
+            errors.extend(query_errors)
+            return [], errors
+        
+        # Search logic
+        for item in data_source:
+            match = True
+            for qkey, qval in query.items():
+                if item.get(qkey) != qval:
+                    match = False
+                    break
+            if match:
+                results.append(item)
+
+        # Apply dynamic mapping if requested
+        if container_map:
+            mapped_results = []
+            for result in results:
+                mapped_item = {}
+                for orig_field, mapped_field in container_map.items():
+                    mapped_item[mapped_field] = result.get(orig_field)
+                mapped_results.append(mapped_item)
+            results = mapped_results
+        
+        return results, errors
+
+    def advanced_query(self, source_schema: str, target_schema: str, query: Dict[str, Any], 
+                      data_source: List[Dict[str, Any]]) -> Tuple[List[Dict[str, Any]], List[str]]:
+        """Search and transform data from source to target schema with validation/errors/JSON support."""
+        errors = []
+        transformed_results = []
+        
+        for item in data_source:
+            # Validate source item
+            is_valid, validation_errors = self.registry.validate_data(source_schema, item)
+            if not is_valid:
+                errors.extend([f"Source error: {err}" for err in validation_errors])
+                continue
+            # Check query match
+            match = True
+            for qkey, qval in query.items():
+                if item.get(qkey) != qval:
+                    match = False
+                    break
+            if not match:
+                continue
+            # Transform to target schema
+            transformed, trans_errors = self.registry.transform_data(item, source_schema, target_schema)
+            if transformed:
+                transformed_results.append(transformed)
+            else:
+                errors.extend([f"Mapping error: {err}" for err in trans_errors])
+
+        return transformed_results, errors
+
+    def round_trip_json(self, obj: Any) -> Any:
+        """Verifies serialization and deserialization integrity (round-trip)."""
+        try:
+            dump = json.dumps(obj)
+            restore = json.loads(dump)
+            if dump == json.dumps(restore):
+                return restore
+            return None
+        except Exception as e:
+            return None
+
+
+if __name__ == "__main__":
+    # Example usage/test
+    registry = SchemaRegistry()
+    agent = QuerySearchAgent(registry)
+    users = [
+        {'id': 1, 'username': 'alice', 'email': 'alice@example.com', 'age': 30},
+        {'id': 2, 'username': 'bob', 'email': 'bob@example.com', 'age': 24},
+        {'id': 3, 'username': 'carol', 'email': 'carol@test.com', 'age': 29},
+    ]
+    query = {'username': 'alice'}
+    results, errors = agent.search('user', query, users)
+    print("Basic results:", results, "Errors:", errors)
+
+    # Advanced query with transformation and error reporting
+    registry.register_schema('short_user', {
+        'username': {'type': 'string', 'required': True},
+        'age': {'type': 'int', 'required': False}
+    }, 'Short user schema')
+    registry.register_mapping('user', 'short_user', {
+        'username': 'username',
+        'age': 'age',
+    })
+    adv_results, adv_errors = agent.advanced_query('user', 'short_user', query, users)
+    print("Advanced results:", adv_results, "Errors:", adv_errors)
+
+    # Round-trip JSON test
+    print("Round-trip JSON:", agent.round_trip_json(results))