mirror of
https://github.com/trustgraph-ai/trustgraph.git
synced 2026-07-24 12:41:02 +02:00
Structured data loader CLI
This commit is contained in:
Cyber MacGeddon
parent
c6cc8328db
commit
dc5b1002b4
1 changed files with 331 additions and 5 deletions
|
|
@ -501,14 +501,340 @@ def load_structured_data(
|
||||||
if not descriptor_file:
|
if not descriptor_file:
|
||||||
# Auto-generate descriptor if not provided
|
# Auto-generate descriptor if not provided
|
||||||
logger.info("No descriptor provided, auto-generating...")
|
logger.info("No descriptor provided, auto-generating...")
|
||||||
# TODO: Generate descriptor
|
logger.info(f"Schema name: {schema_name}")
|
||||||
print(f"Would auto-generate descriptor from {input_file}")
|
|
||||||
|
# Read sample data for descriptor generation
|
||||||
|
try:
|
||||||
|
with open(input_file, 'r', encoding='utf-8') as f:
|
||||||
|
sample_data = f.read(sample_chars)
|
||||||
|
logger.info(f"Read {len(sample_data)} characters for descriptor generation")
|
||||||
|
except Exception as e:
|
||||||
|
logger.error(f"Failed to read input file for descriptor generation: {e}")
|
||||||
|
raise
|
||||||
|
|
||||||
|
# Generate descriptor using TrustGraph prompt service
|
||||||
|
try:
|
||||||
|
from trustgraph.api import Api
|
||||||
|
from trustgraph.api.types import ConfigKey
|
||||||
|
|
||||||
|
api = Api(api_url)
|
||||||
|
config_api = api.config()
|
||||||
|
|
||||||
|
# Get available schemas
|
||||||
|
logger.info("Fetching available schemas for descriptor generation...")
|
||||||
|
schema_keys = config_api.list("schema")
|
||||||
|
logger.info(f"Found {len(schema_keys)} schemas: {schema_keys}")
|
||||||
|
|
||||||
|
if not schema_keys:
|
||||||
|
logger.warning("No schemas found in configuration")
|
||||||
|
print("No schemas available in TrustGraph configuration")
|
||||||
|
return
|
||||||
|
|
||||||
|
# Fetch each schema definition
|
||||||
|
schemas = []
|
||||||
|
config_keys = [ConfigKey(type="schema", key=key) for key in schema_keys]
|
||||||
|
schema_values = config_api.get(config_keys)
|
||||||
|
|
||||||
|
for value in schema_values:
|
||||||
|
try:
|
||||||
|
schema_def = json.loads(value.value) if isinstance(value.value, str) else value.value
|
||||||
|
schemas.append(schema_def)
|
||||||
|
logger.debug(f"Loaded schema: {value.key}")
|
||||||
|
except json.JSONDecodeError as e:
|
||||||
|
logger.warning(f"Failed to parse schema {value.key}: {e}")
|
||||||
|
continue
|
||||||
|
|
||||||
|
logger.info(f"Successfully loaded {len(schemas)} schema definitions")
|
||||||
|
|
||||||
|
# Generate descriptor using diagnose-structured-data prompt
|
||||||
|
flow_api = api.flow().id("default")
|
||||||
|
|
||||||
|
logger.info("Calling TrustGraph diagnose-structured-data prompt for descriptor generation...")
|
||||||
|
response = flow_api.prompt(
|
||||||
|
id="diagnose-structured-data",
|
||||||
|
variables={
|
||||||
|
"schemas": schemas,
|
||||||
|
"sample": sample_data
|
||||||
|
}
|
||||||
|
)
|
||||||
|
|
||||||
|
# Parse the generated descriptor
|
||||||
|
if isinstance(response, str):
|
||||||
|
try:
|
||||||
|
descriptor = json.loads(response)
|
||||||
|
except json.JSONDecodeError:
|
||||||
|
logger.error("Generated descriptor is not valid JSON")
|
||||||
|
raise ValueError("Failed to generate valid descriptor")
|
||||||
|
else:
|
||||||
|
descriptor = response
|
||||||
|
|
||||||
|
# Override schema_name if provided
|
||||||
|
if schema_name:
|
||||||
|
descriptor.setdefault('output', {})['schema_name'] = schema_name
|
||||||
|
|
||||||
|
logger.info("Successfully generated descriptor from data sample")
|
||||||
|
|
||||||
|
except ImportError as e:
|
||||||
|
logger.error(f"Failed to import TrustGraph API: {e}")
|
||||||
|
raise
|
||||||
|
except Exception as e:
|
||||||
|
logger.error(f"Failed to generate descriptor: {e}")
|
||||||
|
raise
|
||||||
|
else:
|
||||||
|
# Load existing descriptor
|
||||||
|
try:
|
||||||
|
with open(descriptor_file, 'r', encoding='utf-8') as f:
|
||||||
|
descriptor = json.load(f)
|
||||||
|
logger.info(f"Loaded descriptor configuration from {descriptor_file}")
|
||||||
|
except Exception as e:
|
||||||
|
logger.error(f"Failed to load descriptor file: {e}")
|
||||||
|
raise
|
||||||
|
|
||||||
logger.info(f"Processing {input_file} for import...")
|
logger.info(f"Processing {input_file} for import...")
|
||||||
# TODO: Implement full pipeline
|
|
||||||
print(f"Would process and import data from {input_file}")
|
# Parse data using the same logic as parse-only mode, but with full dataset
|
||||||
|
try:
|
||||||
|
format_info = descriptor.get('format', {})
|
||||||
|
format_type = format_info.get('type', 'csv').lower()
|
||||||
|
encoding = format_info.get('encoding', 'utf-8')
|
||||||
|
|
||||||
|
logger.info(f"Input format: {format_type}, encoding: {encoding}")
|
||||||
|
|
||||||
|
with open(input_file, 'r', encoding=encoding) as f:
|
||||||
|
raw_data = f.read()
|
||||||
|
|
||||||
|
logger.info(f"Read {len(raw_data)} characters from input file")
|
||||||
|
|
||||||
|
except Exception as e:
|
||||||
|
logger.error(f"Failed to read input file: {e}")
|
||||||
|
raise
|
||||||
|
|
||||||
|
# Parse data (reuse parse-only logic but process all records)
|
||||||
|
parsed_records = []
|
||||||
|
batch_size = descriptor.get('output', {}).get('options', {}).get('batch_size', 1000)
|
||||||
|
|
||||||
|
if format_type == 'csv':
|
||||||
|
import csv
|
||||||
|
from io import StringIO
|
||||||
|
|
||||||
|
options = format_info.get('options', {})
|
||||||
|
delimiter = options.get('delimiter', ',')
|
||||||
|
has_header = options.get('has_header', True) or options.get('header', True)
|
||||||
|
|
||||||
|
logger.info(f"CSV options - delimiter: '{delimiter}', has_header: {has_header}")
|
||||||
|
|
||||||
|
try:
|
||||||
|
reader = csv.DictReader(StringIO(raw_data), delimiter=delimiter)
|
||||||
|
if not has_header:
|
||||||
|
first_row = next(reader)
|
||||||
|
fieldnames = [f"field_{i+1}" for i in range(len(first_row))]
|
||||||
|
reader = csv.DictReader(StringIO(raw_data), fieldnames=fieldnames, delimiter=delimiter)
|
||||||
|
|
||||||
|
record_count = 0
|
||||||
|
for row in reader:
|
||||||
|
parsed_records.append(row)
|
||||||
|
record_count += 1
|
||||||
|
|
||||||
|
# Process in batches to avoid memory issues
|
||||||
|
if record_count % batch_size == 0:
|
||||||
|
logger.info(f"Parsed {record_count} records...")
|
||||||
|
|
||||||
|
except Exception as e:
|
||||||
|
logger.error(f"Failed to parse CSV data: {e}")
|
||||||
|
raise
|
||||||
|
|
||||||
|
elif format_type == 'json':
|
||||||
|
try:
|
||||||
|
data = json.loads(raw_data)
|
||||||
|
if isinstance(data, list):
|
||||||
|
parsed_records = data
|
||||||
|
elif isinstance(data, dict):
|
||||||
|
root_path = format_info.get('options', {}).get('root_path')
|
||||||
|
if root_path:
|
||||||
|
if root_path.startswith('$.'):
|
||||||
|
key = root_path[2:]
|
||||||
|
data = data.get(key, data)
|
||||||
|
|
||||||
|
if isinstance(data, list):
|
||||||
|
parsed_records = data
|
||||||
|
else:
|
||||||
|
parsed_records = [data]
|
||||||
|
|
||||||
|
except Exception as e:
|
||||||
|
logger.error(f"Failed to parse JSON data: {e}")
|
||||||
|
raise
|
||||||
|
|
||||||
|
elif format_type == 'xml':
|
||||||
|
import xml.etree.ElementTree as ET
|
||||||
|
|
||||||
|
options = format_info.get('options', {})
|
||||||
|
record_path = options.get('record_path', '//record')
|
||||||
|
field_attribute = options.get('field_attribute')
|
||||||
|
|
||||||
|
# Legacy support for old options format
|
||||||
|
if 'root_element' in options or 'record_element' in options:
|
||||||
|
root_element = options.get('root_element')
|
||||||
|
record_element = options.get('record_element', 'record')
|
||||||
|
if root_element:
|
||||||
|
record_path = f"//{root_element}/{record_element}"
|
||||||
|
else:
|
||||||
|
record_path = f"//{record_element}"
|
||||||
|
|
||||||
|
logger.info(f"XML options - record_path: '{record_path}', field_attribute: '{field_attribute}'")
|
||||||
|
|
||||||
|
try:
|
||||||
|
root = ET.fromstring(raw_data)
|
||||||
|
|
||||||
|
# Find record elements using XPath
|
||||||
|
xpath_expr = record_path
|
||||||
|
if xpath_expr.startswith('/ROOT/'):
|
||||||
|
xpath_expr = xpath_expr[6:]
|
||||||
|
elif xpath_expr.startswith('/'):
|
||||||
|
xpath_expr = '.' + xpath_expr
|
||||||
|
|
||||||
|
records = root.findall(xpath_expr)
|
||||||
|
logger.info(f"Found {len(records)} records using XPath: {record_path} (converted to: {xpath_expr})")
|
||||||
|
|
||||||
|
# Convert XML elements to dictionaries
|
||||||
|
for element in records:
|
||||||
|
record = {}
|
||||||
|
|
||||||
|
if field_attribute:
|
||||||
|
# Handle field elements with name attributes (UN data format)
|
||||||
|
for child in element:
|
||||||
|
if child.tag == 'field' and field_attribute in child.attrib:
|
||||||
|
field_name = child.attrib[field_attribute]
|
||||||
|
field_value = child.text.strip() if child.text else ""
|
||||||
|
record[field_name] = field_value
|
||||||
|
else:
|
||||||
|
# Handle standard XML structure
|
||||||
|
record.update(element.attrib)
|
||||||
|
|
||||||
|
for child in element:
|
||||||
|
if child.text:
|
||||||
|
record[child.tag] = child.text.strip()
|
||||||
|
else:
|
||||||
|
record[child.tag] = ""
|
||||||
|
|
||||||
|
if not record and element.text:
|
||||||
|
record['value'] = element.text.strip()
|
||||||
|
|
||||||
|
parsed_records.append(record)
|
||||||
|
|
||||||
|
except ET.ParseError as e:
|
||||||
|
logger.error(f"Failed to parse XML data: {e}")
|
||||||
|
raise
|
||||||
|
except Exception as e:
|
||||||
|
logger.error(f"Failed to process XML data: {e}")
|
||||||
|
raise
|
||||||
|
|
||||||
|
else:
|
||||||
|
raise ValueError(f"Unsupported format type: {format_type}")
|
||||||
|
|
||||||
|
logger.info(f"Successfully parsed {len(parsed_records)} records")
|
||||||
|
|
||||||
|
# Apply transformations and create TrustGraph objects
|
||||||
|
mappings = descriptor.get('mappings', [])
|
||||||
|
processed_records = []
|
||||||
|
schema_name = descriptor.get('output', {}).get('schema_name', 'default')
|
||||||
|
confidence = descriptor.get('output', {}).get('options', {}).get('confidence', 0.9)
|
||||||
|
|
||||||
|
logger.info(f"Applying {len(mappings)} field mappings...")
|
||||||
|
|
||||||
|
for record_num, record in enumerate(parsed_records, start=1):
|
||||||
|
processed_record = {}
|
||||||
|
|
||||||
|
for mapping in mappings:
|
||||||
|
source_field = mapping.get('source_field') or mapping.get('source')
|
||||||
|
target_field = mapping.get('target_field') or mapping.get('target')
|
||||||
|
|
||||||
|
if source_field in record:
|
||||||
|
value = record[source_field]
|
||||||
|
|
||||||
|
# Apply transforms
|
||||||
|
transforms = mapping.get('transforms', [])
|
||||||
|
for transform in transforms:
|
||||||
|
transform_type = transform.get('type')
|
||||||
|
|
||||||
|
if transform_type == 'trim' and isinstance(value, str):
|
||||||
|
value = value.strip()
|
||||||
|
elif transform_type == 'upper' and isinstance(value, str):
|
||||||
|
value = value.upper()
|
||||||
|
elif transform_type == 'lower' and isinstance(value, str):
|
||||||
|
value = value.lower()
|
||||||
|
elif transform_type == 'title_case' and isinstance(value, str):
|
||||||
|
value = value.title()
|
||||||
|
elif transform_type == 'to_int':
|
||||||
|
try:
|
||||||
|
value = int(value) if value != '' else None
|
||||||
|
except (ValueError, TypeError):
|
||||||
|
logger.warning(f"Failed to convert '{value}' to int in record {record_num}")
|
||||||
|
elif transform_type == 'to_float':
|
||||||
|
try:
|
||||||
|
value = float(value) if value != '' else None
|
||||||
|
except (ValueError, TypeError):
|
||||||
|
logger.warning(f"Failed to convert '{value}' to float in record {record_num}")
|
||||||
|
|
||||||
|
processed_record[target_field] = value
|
||||||
|
else:
|
||||||
|
logger.warning(f"Source field '{source_field}' not found in record {record_num}")
|
||||||
|
|
||||||
|
# Create TrustGraph ExtractedObject
|
||||||
|
output_record = {
|
||||||
|
"metadata": {
|
||||||
|
"id": f"import-{record_num}",
|
||||||
|
"metadata": [],
|
||||||
|
"user": "trustgraph",
|
||||||
|
"collection": "default"
|
||||||
|
},
|
||||||
|
"schema_name": schema_name,
|
||||||
|
"values": processed_record,
|
||||||
|
"confidence": confidence,
|
||||||
|
"source_span": ""
|
||||||
|
}
|
||||||
|
processed_records.append(output_record)
|
||||||
|
|
||||||
|
logger.info(f"Processed {len(processed_records)} records with transformations")
|
||||||
|
|
||||||
if dry_run:
|
if dry_run:
|
||||||
print("Dry run mode - no data will be imported")
|
print(f"Dry run mode - would import {len(processed_records)} records to TrustGraph")
|
||||||
|
print(f"Target schema: {schema_name}")
|
||||||
|
print(f"Sample record:")
|
||||||
|
if processed_records:
|
||||||
|
print(json.dumps(processed_records[0], indent=2))
|
||||||
|
return
|
||||||
|
|
||||||
|
# Import to TrustGraph using objects dispatcher
|
||||||
|
logger.info(f"Importing {len(processed_records)} records to TrustGraph...")
|
||||||
|
|
||||||
|
try:
|
||||||
|
import trustgraph.publish.writer as writer
|
||||||
|
|
||||||
|
# Initialize the objects writer
|
||||||
|
objects_writer = writer.Writer(api_url, "objects")
|
||||||
|
|
||||||
|
# Send records in batches
|
||||||
|
imported_count = 0
|
||||||
|
for i in range(0, len(processed_records), batch_size):
|
||||||
|
batch = processed_records[i:i + batch_size]
|
||||||
|
|
||||||
|
for record in batch:
|
||||||
|
objects_writer.write(record)
|
||||||
|
imported_count += 1
|
||||||
|
|
||||||
|
logger.info(f"Imported {imported_count}/{len(processed_records)} records...")
|
||||||
|
|
||||||
|
logger.info(f"Successfully imported {imported_count} records to TrustGraph")
|
||||||
|
print(f"Import completed: {imported_count} records imported to schema '{schema_name}'")
|
||||||
|
|
||||||
|
except ImportError as e:
|
||||||
|
logger.error(f"Failed to import TrustGraph writer: {e}")
|
||||||
|
print(f"Error: TrustGraph writer not available - {e}")
|
||||||
|
raise
|
||||||
|
except Exception as e:
|
||||||
|
logger.error(f"Failed to import data to TrustGraph: {e}")
|
||||||
|
print(f"Import failed: {e}")
|
||||||
|
raise
|
||||||
|
|
||||||
|
|
||||||
def main():
|
def main():
|
||||||
|
|
|
||||||
Loading…
Add table
Add a link
Reference in a new issue