From 7602fe6e7014b2e2447382cadfead1e1c20c0968 Mon Sep 17 00:00:00 2001 From: Cyber MacGeddon Date: Wed, 2 Oct 2024 10:50:45 +0100 Subject: [PATCH] Migrate to use metadata part of schema, also add metadata to triples & vecs --- trustgraph-base/trustgraph/schema/__init__.py | 1 + trustgraph-cli/scripts/tg-load-pdf | 28 +++++++++++++++++-- trustgraph-cli/scripts/tg-load-text | 28 +++++++++++++++++-- .../trustgraph/chunking/recursive/chunker.py | 14 ++++++---- .../trustgraph/chunking/token/chunker.py | 14 ++++++---- .../trustgraph/decoding/pdf/pdf_decoder.py | 14 ++++++---- .../embeddings/vectorize/vectorize.py | 8 +++--- .../extract/kg/definitions/extract.py | 10 +++---- .../extract/kg/relationships/extract.py | 23 +++++++++------ .../trustgraph/extract/kg/topics/extract.py | 10 +++---- .../trustgraph/extract/object/row/extract.py | 16 +++++------ 11 files changed, 113 insertions(+), 53 deletions(-) diff --git a/trustgraph-base/trustgraph/schema/__init__.py b/trustgraph-base/trustgraph/schema/__init__.py index 0cd5a370..f3cc5b60 100644 --- a/trustgraph-base/trustgraph/schema/__init__.py +++ b/trustgraph-base/trustgraph/schema/__init__.py @@ -7,6 +7,7 @@ from . object import * from . topic import * from . graph import * from . retrieval import * +from . metadata import * diff --git a/trustgraph-cli/scripts/tg-load-pdf b/trustgraph-cli/scripts/tg-load-pdf index 5d54da93..460a2f06 100755 --- a/trustgraph-cli/scripts/tg-load-pdf +++ b/trustgraph-cli/scripts/tg-load-pdf @@ -6,7 +6,7 @@ Loads a PDF document into TrustGraph processing. import pulsar from pulsar.schema import JsonSchema -from trustgraph.schema import Document, Source, document_ingest_queue +from trustgraph.schema import Document, document_ingest_queue, Metadata import base64 import hashlib import argparse @@ -15,12 +15,17 @@ import time from trustgraph.log_level import LogLevel +default_user = 'trustgraph' +default_collection = 'default' + class Loader: def __init__( self, pulsar_host, output_queue, + user, + collection, log_level, ): @@ -35,6 +40,9 @@ class Loader: chunking_enabled=True, ) + self.user = user + self.collection = collection + def load(self, files): for file in files: @@ -50,10 +58,12 @@ class Loader: id = hashlib.sha256(path.encode("utf-8")).hexdigest()[0:8] r = Document( - source=Source( + metadata=Metadata( source=path, title=path, id=id, + user=self.user, + collection=self.collection, ), data=base64.b64encode(data), ) @@ -90,6 +100,18 @@ def main(): help=f'Output queue (default: {default_output_queue})' ) + parser.add_argument( + '-u', '--user', + default=default_user, + help=f'User ID (default: {default_user})' + ) + + parser.add_argument( + '-c', '--collection', + default=default_collection, + help=f'Collection ID (default: {default_collection})' + ) + parser.add_argument( '-l', '--log-level', type=LogLevel, @@ -112,6 +134,8 @@ def main(): p = Loader( pulsar_host=args.pulsar_host, output_queue=args.output_queue, + user=args.user, + collection=args.collection, log_level=args.log_level, ) diff --git a/trustgraph-cli/scripts/tg-load-text b/trustgraph-cli/scripts/tg-load-text index 8137006c..e22af5b1 100755 --- a/trustgraph-cli/scripts/tg-load-text +++ b/trustgraph-cli/scripts/tg-load-text @@ -6,7 +6,7 @@ Loads a text document into TrustGraph processing. import pulsar from pulsar.schema import JsonSchema -from trustgraph.schema import TextDocument, Source, text_ingest_queue +from trustgraph.schema import TextDocument, text_ingest_queue, Metadata import base64 import hashlib import argparse @@ -15,12 +15,17 @@ import time from trustgraph.log_level import LogLevel +default_user = 'trustgraph' +default_collection = 'default' + class Loader: def __init__( self, pulsar_host, output_queue, + user, + collection, log_level, ): @@ -35,6 +40,9 @@ class Loader: chunking_enabled=True, ) + self.user = user + self.collection = collection + def load(self, files): for file in files: @@ -50,10 +58,12 @@ class Loader: id = hashlib.sha256(path.encode("utf-8")).hexdigest()[0:8] r = TextDocument( - source=Source( + metadata=Metadata( source=path, title=path, id=id, + user=self.user, + collection=self.collection, ), text=data, ) @@ -90,6 +100,18 @@ def main(): help=f'Output queue (default: {default_output_queue})' ) + parser.add_argument( + '-u', '--user', + default=default_user, + help=f'User ID (default: {default_user})' + ) + + parser.add_argument( + '-c', '--collection', + default=default_collection, + help=f'Collection ID (default: {default_collection})' + ) + parser.add_argument( '-l', '--log-level', type=LogLevel, @@ -112,6 +134,8 @@ def main(): p = Loader( pulsar_host=args.pulsar_host, output_queue=args.output_queue, + user=args.user, + collection=args.collection, log_level=args.log_level, ) diff --git a/trustgraph-flow/trustgraph/chunking/recursive/chunker.py b/trustgraph-flow/trustgraph/chunking/recursive/chunker.py index fe1a0cee..bec854bb 100755 --- a/trustgraph-flow/trustgraph/chunking/recursive/chunker.py +++ b/trustgraph-flow/trustgraph/chunking/recursive/chunker.py @@ -7,7 +7,7 @@ as text as separate output objects. from langchain_text_splitters import RecursiveCharacterTextSplitter from prometheus_client import Histogram -from ... schema import TextDocument, Chunk, Source +from ... schema import TextDocument, Chunk, Metadata from ... schema import text_ingest_queue, chunk_ingest_queue from ... log_level import LogLevel from ... base import ConsumerProducer @@ -55,7 +55,7 @@ class Processor(ConsumerProducer): def handle(self, msg): v = msg.value() - print(f"Chunking {v.source.id}...", flush=True) + print(f"Chunking {v.metadata.id}...", flush=True) texts = self.text_splitter.create_documents( [v.text.decode("utf-8")] @@ -63,13 +63,15 @@ class Processor(ConsumerProducer): for ix, chunk in enumerate(texts): - id = v.source.id + "-c" + str(ix) + id = v.metadata.id + "-c" + str(ix) r = Chunk( - source=Source( - source=v.source.source, + metadata=Metadata( + source=v.metadata.source, id=id, - title=v.source.title + title=v.metadata.title, + user=v.metadata.user, + collection=v.metadata.collection, ), chunk=chunk.page_content.encode("utf-8"), ) diff --git a/trustgraph-flow/trustgraph/chunking/token/chunker.py b/trustgraph-flow/trustgraph/chunking/token/chunker.py index c152b0fd..e7bc5667 100755 --- a/trustgraph-flow/trustgraph/chunking/token/chunker.py +++ b/trustgraph-flow/trustgraph/chunking/token/chunker.py @@ -7,7 +7,7 @@ as text as separate output objects. from langchain_text_splitters import TokenTextSplitter from prometheus_client import Histogram -from ... schema import TextDocument, Chunk, Source +from ... schema import TextDocument, Chunk, Metadata from ... schema import text_ingest_queue, chunk_ingest_queue from ... log_level import LogLevel from ... base import ConsumerProducer @@ -54,7 +54,7 @@ class Processor(ConsumerProducer): def handle(self, msg): v = msg.value() - print(f"Chunking {v.source.id}...", flush=True) + print(f"Chunking {v.metadata.id}...", flush=True) texts = self.text_splitter.create_documents( [v.text.decode("utf-8")] @@ -62,13 +62,15 @@ class Processor(ConsumerProducer): for ix, chunk in enumerate(texts): - id = v.source.id + "-c" + str(ix) + id = v.metadata.id + "-c" + str(ix) r = Chunk( - source=Source( - source=v.source.source, + metadata=Metadata( + source=v.metadata.source, id=id, - title=v.source.title + title=v.metadata.title, + user=v.metadata.user, + collection=v.metadata.collection, ), chunk=chunk.page_content.encode("utf-8"), ) diff --git a/trustgraph-flow/trustgraph/decoding/pdf/pdf_decoder.py b/trustgraph-flow/trustgraph/decoding/pdf/pdf_decoder.py index fffcaee0..7e9aeafa 100755 --- a/trustgraph-flow/trustgraph/decoding/pdf/pdf_decoder.py +++ b/trustgraph-flow/trustgraph/decoding/pdf/pdf_decoder.py @@ -8,7 +8,7 @@ import tempfile import base64 from langchain_community.document_loaders import PyPDFLoader -from ... schema import Document, TextDocument, Source +from ... schema import Document, TextDocument, Metadata from ... schema import document_ingest_queue, text_ingest_queue from ... log_level import LogLevel from ... base import ConsumerProducer @@ -45,7 +45,7 @@ class Processor(ConsumerProducer): v = msg.value() - print(f"Decoding {v.source.id}...", flush=True) + print(f"Decoding {v.metadata.id}...", flush=True) with tempfile.NamedTemporaryFile(delete_on_close=False) as fp: @@ -59,12 +59,14 @@ class Processor(ConsumerProducer): for ix, page in enumerate(pages): - id = v.source.id + "-p" + str(ix) + id = v.metadata.id + "-p" + str(ix) r = TextDocument( - source=Source( - source=v.source.source, - title=v.source.title, + metadata=Metadata( + source=v.metadata.source, + title=v.metadata.title, id=id, + user=v.metadata.user, + collection=v.metadata.collection, ), text=page.page_content.encode("utf-8"), ) diff --git a/trustgraph-flow/trustgraph/embeddings/vectorize/vectorize.py b/trustgraph-flow/trustgraph/embeddings/vectorize/vectorize.py index 3770fee2..4cf2af05 100755 --- a/trustgraph-flow/trustgraph/embeddings/vectorize/vectorize.py +++ b/trustgraph-flow/trustgraph/embeddings/vectorize/vectorize.py @@ -50,15 +50,15 @@ class Processor(ConsumerProducer): subscriber=module + "-emb", ) - def emit(self, source, chunk, vectors): + def emit(self, metadata, chunk, vectors): - r = ChunkEmbeddings(source=source, chunk=chunk, vectors=vectors) + r = ChunkEmbeddings(metadata=metadata, chunk=chunk, vectors=vectors) self.producer.send(r) def handle(self, msg): v = msg.value() - print(f"Indexing {v.source.id}...", flush=True) + print(f"Indexing {v.metadata.id}...", flush=True) chunk = v.chunk.decode("utf-8") @@ -67,7 +67,7 @@ class Processor(ConsumerProducer): vectors = self.embeddings.request(chunk) self.emit( - source=v.source, + metadata=v.metadata, chunk=chunk.encode("utf-8"), vectors=vectors ) diff --git a/trustgraph-flow/trustgraph/extract/kg/definitions/extract.py b/trustgraph-flow/trustgraph/extract/kg/definitions/extract.py index 06ba8e68..d528b74d 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, Source, Value +from .... schema import ChunkEmbeddings, Triple, Metadata, Value from .... schema import chunk_embeddings_ingest_queue, triples_store_queue from .... schema import prompt_request_queue from .... schema import prompt_response_queue @@ -69,15 +69,15 @@ class Processor(ConsumerProducer): return self.prompt.request_definitions(chunk) - def emit_edge(self, s, p, o): + def emit_edge(self, metadata, s, p, o): - t = Triple(s=s, p=p, o=o) + t = Triple(metadata=metadata, s=s, p=p, o=o) self.producer.send(t) def handle(self, msg): v = msg.value() - print(f"Indexing {v.source.id}...", flush=True) + print(f"Indexing {v.metadata.id}...", flush=True) chunk = v.chunk.decode("utf-8") @@ -101,7 +101,7 @@ class Processor(ConsumerProducer): s_value = Value(value=str(s_uri), is_uri=True) o_value = Value(value=str(o), is_uri=False) - self.emit_edge(s_value, DEFINITION_VALUE, o_value) + self.emit_edge(v.metadata, s_value, DEFINITION_VALUE, o_value) except Exception as e: print("Exception: ", e, flush=True) diff --git a/trustgraph-flow/trustgraph/extract/kg/relationships/extract.py b/trustgraph-flow/trustgraph/extract/kg/relationships/extract.py index 49ef9072..c1779143 100755 --- a/trustgraph-flow/trustgraph/extract/kg/relationships/extract.py +++ b/trustgraph-flow/trustgraph/extract/kg/relationships/extract.py @@ -9,7 +9,8 @@ import urllib.parse import os from pulsar.schema import JsonSchema -from .... schema import ChunkEmbeddings, Triple, GraphEmbeddings, Source, Value +from .... schema import ChunkEmbeddings, Triple, GraphEmbeddings +from .... schema import Metadata, Value from .... schema import chunk_embeddings_ingest_queue, triples_store_queue from .... schema import graph_embeddings_store_queue from .... schema import prompt_request_queue @@ -91,20 +92,20 @@ class Processor(ConsumerProducer): return self.prompt.request_relationships(chunk) - def emit_edge(self, s, p, o): + def emit_edge(self, metadata, s, p, o): - t = Triple(s=s, p=p, o=o) + t = Triple(metadata=metadata, s=s, p=p, o=o) self.producer.send(t) - def emit_vec(self, ent, vec): + def emit_vec(self, metadata, ent, vec): - r = GraphEmbeddings(entity=ent, vectors=vec) + r = GraphEmbeddings(metdata=metadata, entity=ent, vectors=vec) self.vec_prod.send(r) def handle(self, msg): v = msg.value() - print(f"Indexing {v.source.id}...", flush=True) + print(f"Indexing {v.metadata.id}...", flush=True) chunk = v.chunk.decode("utf-8") @@ -139,6 +140,7 @@ class Processor(ConsumerProducer): o_value = Value(value=str(o), is_uri=False) self.emit_edge( + v.metdata, s_value, p_value, o_value @@ -146,6 +148,7 @@ class Processor(ConsumerProducer): # Label for s self.emit_edge( + v.metdata, s_value, RDF_LABEL_VALUE, Value(value=str(s), is_uri=False) @@ -153,6 +156,7 @@ class Processor(ConsumerProducer): # Label for p self.emit_edge( + v.metdata, p_value, RDF_LABEL_VALUE, Value(value=str(p), is_uri=False) @@ -161,15 +165,16 @@ class Processor(ConsumerProducer): if rel.o_entity: # Label for o self.emit_edge( + v.metdata, o_value, RDF_LABEL_VALUE, Value(value=str(o), is_uri=False) ) - self.emit_vec(s_value, v.vectors) - self.emit_vec(p_value, v.vectors) + self.emit_vec(v.metadata, s_value, v.vectors) + self.emit_vec(v.metadata, p_value, v.vectors) if rel.o_entity: - self.emit_vec(o_value, v.vectors) + self.emit_vec(v.metadata, o_value, v.vectors) except Exception as e: print("Exception: ", e, flush=True) diff --git a/trustgraph-flow/trustgraph/extract/kg/topics/extract.py b/trustgraph-flow/trustgraph/extract/kg/topics/extract.py index e2ebe5b0..81e52669 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, Source, Value +from .... schema import ChunkEmbeddings, Triple, Metadata, Value from .... schema import chunk_embeddings_ingest_queue, triples_store_queue from .... schema import prompt_request_queue from .... schema import prompt_response_queue @@ -69,15 +69,15 @@ class Processor(ConsumerProducer): return self.prompt.request_topics(chunk) - def emit_edge(self, s, p, o): + def emit_edge(self, metadata, s, p, o): - t = Triple(s=s, p=p, o=o) + t = Triple(metadata=metadata, s=s, p=p, o=o) self.producer.send(t) def handle(self, msg): v = msg.value() - print(f"Indexing {v.source.id}...", flush=True) + print(f"Indexing {v.metadata.id}...", flush=True) chunk = v.chunk.decode("utf-8") @@ -101,7 +101,7 @@ class Processor(ConsumerProducer): s_value = Value(value=str(s_uri), is_uri=True) o_value = Value(value=str(o), is_uri=False) - self.emit_edge(s_value, DEFINITION_VALUE, o_value) + self.emit_edge(v. metadata, s_value, DEFINITION_VALUE, o_value) except Exception as e: print("Exception: ", e, flush=True) diff --git a/trustgraph-flow/trustgraph/extract/object/row/extract.py b/trustgraph-flow/trustgraph/extract/object/row/extract.py index aa53f2a6..185a59c3 100755 --- a/trustgraph-flow/trustgraph/extract/object/row/extract.py +++ b/trustgraph-flow/trustgraph/extract/object/row/extract.py @@ -8,7 +8,7 @@ import urllib.parse import os from pulsar.schema import JsonSchema -from .... schema import ChunkEmbeddings, Rows, ObjectEmbeddings, Source +from .... schema import ChunkEmbeddings, Rows, ObjectEmbeddings, Metadata from .... schema import RowSchema, Field from .... schema import chunk_embeddings_ingest_queue, rows_store_queue from .... schema import object_embeddings_store_queue @@ -124,24 +124,24 @@ class Processor(ConsumerProducer): def get_rows(self, chunk): return self.prompt.request_rows(self.schema, chunk) - def emit_rows(self, source, rows): + def emit_rows(self, metadata, rows): t = Rows( - source=source, row_schema=self.row_schema, rows=rows + metadata=metadata, row_schema=self.row_schema, rows=rows ) self.producer.send(t) - def emit_vec(self, source, name, vec, key_name, key): + def emit_vec(self, metadata, name, vec, key_name, key): r = ObjectEmbeddings( - source=source, vectors=vec, name=name, key_name=key_name, id=key + metadata=metadata, vectors=vec, name=name, key_name=key_name, id=key ) self.vec_prod.send(r) def handle(self, msg): v = msg.value() - print(f"Indexing {v.source.id}...", flush=True) + print(f"Indexing {v.metadata.id}...", flush=True) chunk = v.chunk.decode("utf-8") @@ -150,13 +150,13 @@ class Processor(ConsumerProducer): rows = self.get_rows(chunk) self.emit_rows( - source=v.source, + metadata=v.metadata, rows=rows ) for row in rows: self.emit_vec( - source=v.source, vec=v.vectors, + metadata=v.metadata, vec=v.vectors, name=self.schema.name, key_name=self.primary.name, key=row[self.primary.name] )