Spaces:
Sleeping
Sleeping
File size: 4,886 Bytes
9989568 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 | """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))
|