Spaces:
Sleeping
Sleeping
Add query/search agent with advanced mapping, validation, error reporting, and JSON support
9989568 verified | """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)) | |