apunkt/trustgraph
1
0
Fork 0
mirror of https://github.com/trustgraph-ai/trustgraph.git synced 2026-07-24 20:51:02 +02:00
Code Issues Projects Releases Packages Wiki Activity Actions 2

Object batching

Browse source
This commit is contained in:
Cyber MacGeddon 2025-09-05 15:48:44 +01:00
parent 162859f1d6
commit 2191de0017
1 changed files with 32 additions and 5 deletions
Download patch file Download diff file Expand all files Collapse all files

37
trustgraph-cli/trustgraph/cli/load_structured_data.py
View file

@ -804,7 +804,18 @@ def load_structured_data(
print(f"Target schema: {schema_name}") print(f"Target schema: {schema_name}")
print(f"Sample record:") print(f"Sample record:")
if processed_records: if processed_records:
print(json.dumps(processed_records[0], indent=2)) # Show what the batched format will look like
sample_batch = processed_records[:min(3, len(processed_records))]
batch_values = [record["values"] for record in sample_batch]
first_record = processed_records[0]
batched_sample = {
"metadata": first_record["metadata"],
"schema_name": first_record["schema_name"],
"values": batch_values,
"confidence": first_record["confidence"],
"source_span": first_record["source_span"]
}
print(json.dumps(batched_sample, indent=2))
return return
# Import to TrustGraph using objects import endpoint via WebSocket # Import to TrustGraph using objects import endpoint via WebSocket
@ -828,10 +839,26 @@ def load_structured_data(
async with connect(objects_url) as ws: async with connect(objects_url) as ws:
imported_count = 0 imported_count = 0
for record in processed_records: # Process records in batches
# Send individual ExtractedObject records for i in range(0, len(processed_records), batch_size):
await ws.send(json.dumps(record)) batch_records = processed_records[i:i + batch_size]
imported_count += 1
# Extract values from each record in the batch
batch_values = [record["values"] for record in batch_records]
# Create batched ExtractedObject message using first record as template
first_record = batch_records[0]
batched_record = {
"metadata": first_record["metadata"],
"schema_name": first_record["schema_name"],
"values": batch_values, # Array of value dictionaries
"confidence": first_record["confidence"],
"source_span": first_record["source_span"]
}
# Send batched ExtractedObject
await ws.send(json.dumps(batched_record))
imported_count += len(batch_records)
if imported_count % 100 == 0: if imported_count % 100 == 0:
logger.info(f"Imported {imported_count}/{len(processed_records)} records...") logger.info(f"Imported {imported_count}/{len(processed_records)} records...")