From 33a4647155bae4b57b0359603d84a625ef383730 Mon Sep 17 00:00:00 2001 From: Cyber MacGeddon Date: Wed, 16 Oct 2024 16:56:36 +0100 Subject: [PATCH] Change graph writers to ingest Triples instead of Triple. This is just an interface change, triples are still passed one-at-a-time from extractors, but this permits sending in batches in the future. --- trustgraph-base/trustgraph/schema/graph.py | 8 +++----- trustgraph-base/trustgraph/schema/types.py | 5 +++++ .../extract/kg/definitions/extract.py | 9 ++++++--- .../extract/kg/relationships/extract.py | 11 +++++++---- .../trustgraph/extract/kg/topics/extract.py | 9 ++++++--- .../storage/triples/cassandra/write.py | 17 +++++++++-------- .../trustgraph/storage/triples/neo4j/write.py | 16 +++++++++------- 7 files changed, 45 insertions(+), 30 deletions(-) diff --git a/trustgraph-base/trustgraph/schema/graph.py b/trustgraph-base/trustgraph/schema/graph.py index 107478a4..2d108a30 100644 --- a/trustgraph-base/trustgraph/schema/graph.py +++ b/trustgraph-base/trustgraph/schema/graph.py @@ -1,7 +1,7 @@ from pulsar.schema import Record, Bytes, String, Boolean, Integer, Array, Double -from . types import Error, Value +from . types import Error, Value, Triple from . topic import topic from . metadata import Metadata @@ -41,11 +41,9 @@ graph_embeddings_response_queue = topic( # Graph triples -class Triple(Record): +class Triples(Record): metadata = Metadata() - s = Value() - p = Value() - o = Value() + triples = Array(Triple()) triples_store_queue = topic('triples-store') diff --git a/trustgraph-base/trustgraph/schema/types.py b/trustgraph-base/trustgraph/schema/types.py index 4cad70ac..b75a0884 100644 --- a/trustgraph-base/trustgraph/schema/types.py +++ b/trustgraph-base/trustgraph/schema/types.py @@ -10,6 +10,11 @@ class Value(Record): is_uri = Boolean() type = String() +class Triple(Record): + s = Value() + p = Value() + o = Value() + class Field(Record): name = String() # int, string, long, bool, float, double diff --git a/trustgraph-flow/trustgraph/extract/kg/definitions/extract.py b/trustgraph-flow/trustgraph/extract/kg/definitions/extract.py index d528b74d..9e3e7344 100755 --- a/trustgraph-flow/trustgraph/extract/kg/definitions/extract.py +++ b/trustgraph-flow/trustgraph/extract/kg/definitions/extract.py @@ -7,7 +7,7 @@ get entity definitions which are output as graph edges. import urllib.parse import json -from .... schema import ChunkEmbeddings, Triple, Metadata, Value +from .... schema import ChunkEmbeddings, Triple, Triples, Metadata, Value from .... schema import chunk_embeddings_ingest_queue, triples_store_queue from .... schema import prompt_request_queue from .... schema import prompt_response_queue @@ -44,7 +44,7 @@ class Processor(ConsumerProducer): "output_queue": output_queue, "subscriber": subscriber, "input_schema": ChunkEmbeddings, - "output_schema": Triple, + "output_schema": Triples, "prompt_request_queue": pr_request_queue, "prompt_response_queue": pr_response_queue, } @@ -71,7 +71,10 @@ class Processor(ConsumerProducer): def emit_edge(self, metadata, s, p, o): - t = Triple(metadata=metadata, s=s, p=p, o=o) + t = Triples( + metadata=metadata, + triples=[Triple(s=s, p=p, o=o)], + ) self.producer.send(t) def handle(self, msg): diff --git a/trustgraph-flow/trustgraph/extract/kg/relationships/extract.py b/trustgraph-flow/trustgraph/extract/kg/relationships/extract.py index cb80e47f..0cade2f9 100755 --- a/trustgraph-flow/trustgraph/extract/kg/relationships/extract.py +++ b/trustgraph-flow/trustgraph/extract/kg/relationships/extract.py @@ -9,7 +9,7 @@ import urllib.parse import os from pulsar.schema import JsonSchema -from .... schema import ChunkEmbeddings, Triple, GraphEmbeddings +from .... schema import ChunkEmbeddings, Triple, Triples, GraphEmbeddings from .... schema import Metadata, Value from .... schema import chunk_embeddings_ingest_queue, triples_store_queue from .... schema import graph_embeddings_store_queue @@ -50,7 +50,7 @@ class Processor(ConsumerProducer): "output_queue": output_queue, "subscriber": subscriber, "input_schema": ChunkEmbeddings, - "output_schema": Triple, + "output_schema": Triples, "prompt_request_queue": pr_request_queue, "prompt_response_queue": pr_response_queue, } @@ -69,7 +69,7 @@ class Processor(ConsumerProducer): "prompt_response_queue": pr_response_queue, "subscriber": subscriber, "input_schema": ChunkEmbeddings.__name__, - "output_schema": Triple.__name__, + "output_schema": Triples.__name__, "vector_schema": GraphEmbeddings.__name__, }) @@ -94,7 +94,10 @@ class Processor(ConsumerProducer): def emit_edge(self, metadata, s, p, o): - t = Triple(metadata=metadata, s=s, p=p, o=o) + t = Triples( + metadata=metadata, + triples=[Triple(s=s, p=p, o=o)], + ) self.producer.send(t) def emit_vec(self, metadata, ent, vec): diff --git a/trustgraph-flow/trustgraph/extract/kg/topics/extract.py b/trustgraph-flow/trustgraph/extract/kg/topics/extract.py index 81e52669..8dfc3e6e 100755 --- a/trustgraph-flow/trustgraph/extract/kg/topics/extract.py +++ b/trustgraph-flow/trustgraph/extract/kg/topics/extract.py @@ -7,7 +7,7 @@ get entity definitions which are output as graph edges. import urllib.parse import json -from .... schema import ChunkEmbeddings, Triple, Metadata, Value +from .... schema import ChunkEmbeddings, Triple, Triples, Metadata, Value from .... schema import chunk_embeddings_ingest_queue, triples_store_queue from .... schema import prompt_request_queue from .... schema import prompt_response_queue @@ -44,7 +44,7 @@ class Processor(ConsumerProducer): "output_queue": output_queue, "subscriber": subscriber, "input_schema": ChunkEmbeddings, - "output_schema": Triple, + "output_schema": Triples, "prompt_request_queue": pr_request_queue, "prompt_response_queue": pr_response_queue, } @@ -71,7 +71,10 @@ class Processor(ConsumerProducer): def emit_edge(self, metadata, s, p, o): - t = Triple(metadata=metadata, s=s, p=p, o=o) + t = Triples( + metadata=metadata, + triples=[Triple(s=s, p=p, o=o)], + ) self.producer.send(t) def handle(self, msg): diff --git a/trustgraph-flow/trustgraph/storage/triples/cassandra/write.py b/trustgraph-flow/trustgraph/storage/triples/cassandra/write.py index e8074c80..6825d452 100755 --- a/trustgraph-flow/trustgraph/storage/triples/cassandra/write.py +++ b/trustgraph-flow/trustgraph/storage/triples/cassandra/write.py @@ -10,7 +10,7 @@ import argparse import time from .... direct.cassandra import TrustGraph -from .... schema import Triple +from .... schema import Triple, Triples from .... schema import triples_store_queue from .... log_level import LogLevel from .... base import Consumer @@ -33,7 +33,7 @@ class Processor(Consumer): **params | { "input_queue": input_queue, "subscriber": subscriber, - "input_schema": Triple, + "input_schema": Triples, "graph_host": graph_host, } ) @@ -62,12 +62,13 @@ class Processor(Consumer): raise e self.table = table - - self.tg.insert( - v.s.value, - v.p.value, - v.o.value - ) + + for t in v.triples: + self.tg.insert( + t.s.value, + t.p.value, + t.o.value + ) @staticmethod def add_args(parser): diff --git a/trustgraph-flow/trustgraph/storage/triples/neo4j/write.py b/trustgraph-flow/trustgraph/storage/triples/neo4j/write.py index d20ecdd2..c346e5c4 100755 --- a/trustgraph-flow/trustgraph/storage/triples/neo4j/write.py +++ b/trustgraph-flow/trustgraph/storage/triples/neo4j/write.py @@ -116,14 +116,16 @@ class Processor(Consumer): v = msg.value() - self.create_node(v.s.value) + for t in v.triples: - if v.o.is_uri: - self.create_node(v.o.value) - self.relate_node(v.s.value, v.p.value, v.o.value) - else: - self.create_literal(v.o.value) - self.relate_literal(v.s.value, v.p.value, v.o.value) + self.create_node(t.s.value) + + if v.o.is_uri: + self.create_node(t.o.value) + self.relate_node(t.s.value, t.p.value, t.o.value) + else: + self.create_literal(t.o.value) + self.relate_literal(t.s.value, t.p.value, t.o.value) @staticmethod def add_args(parser):