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

Migrate to use metadata part of schema, also add metadata to triples & vecs

Browse source
This commit is contained in:
Cyber MacGeddon 2024-10-02 10:50:45 +01:00
parent c86ad2d5b8
commit 7602fe6e70
11 changed files with 113 additions and 53 deletions
Download patch file Download diff file Expand all files Collapse all files

1
trustgraph-base/trustgraph/schema/__init__.py
View file

@ -7,6 +7,7 @@ from . object import *
from . topic import * from . topic import *
from . graph import * from . graph import *
from . retrieval import * from . retrieval import *
from . metadata import *

28
trustgraph-cli/scripts/tg-load-pdf
View file

@ -6,7 +6,7 @@ Loads a PDF document into TrustGraph processing.
import pulsar import pulsar
from pulsar.schema import JsonSchema 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 base64
import hashlib import hashlib
import argparse import argparse
@ -15,12 +15,17 @@ import time
from trustgraph.log_level import LogLevel from trustgraph.log_level import LogLevel
default_user = 'trustgraph'
default_collection = 'default'
class Loader: class Loader:
def __init__( def __init__(
self, self,
pulsar_host, pulsar_host,
output_queue, output_queue,
user,
collection,
log_level, log_level,
): ):
@ -35,6 +40,9 @@ class Loader:
chunking_enabled=True, chunking_enabled=True,
) )
self.user = user
self.collection = collection
def load(self, files): def load(self, files):
for file in files: for file in files:
@ -50,10 +58,12 @@ class Loader:
id = hashlib.sha256(path.encode("utf-8")).hexdigest()[0:8] id = hashlib.sha256(path.encode("utf-8")).hexdigest()[0:8]
r = Document( r = Document(
source=Source( metadata=Metadata(
source=path, source=path,
title=path, title=path,
id=id, id=id,
user=self.user,
collection=self.collection,
), ),
data=base64.b64encode(data), data=base64.b64encode(data),
) )
@ -90,6 +100,18 @@ def main():
help=f'Output queue (default: {default_output_queue})' 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( parser.add_argument(
'-l', '--log-level', '-l', '--log-level',
type=LogLevel, type=LogLevel,
@ -112,6 +134,8 @@ def main():
p = Loader( p = Loader(
pulsar_host=args.pulsar_host, pulsar_host=args.pulsar_host,
output_queue=args.output_queue, output_queue=args.output_queue,
user=args.user,
collection=args.collection,
log_level=args.log_level, log_level=args.log_level,
) )

28
trustgraph-cli/scripts/tg-load-text
View file

@ -6,7 +6,7 @@ Loads a text document into TrustGraph processing.
import pulsar import pulsar
from pulsar.schema import JsonSchema 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 base64
import hashlib import hashlib
import argparse import argparse
@ -15,12 +15,17 @@ import time
from trustgraph.log_level import LogLevel from trustgraph.log_level import LogLevel
default_user = 'trustgraph'
default_collection = 'default'
class Loader: class Loader:
def __init__( def __init__(
self, self,
pulsar_host, pulsar_host,
output_queue, output_queue,
user,
collection,
log_level, log_level,
): ):
@ -35,6 +40,9 @@ class Loader:
chunking_enabled=True, chunking_enabled=True,
) )
self.user = user
self.collection = collection
def load(self, files): def load(self, files):
for file in files: for file in files:
@ -50,10 +58,12 @@ class Loader:
id = hashlib.sha256(path.encode("utf-8")).hexdigest()[0:8] id = hashlib.sha256(path.encode("utf-8")).hexdigest()[0:8]
r = TextDocument( r = TextDocument(
source=Source( metadata=Metadata(
source=path, source=path,
title=path, title=path,
id=id, id=id,
user=self.user,
collection=self.collection,
), ),
text=data, text=data,
) )
@ -90,6 +100,18 @@ def main():
help=f'Output queue (default: {default_output_queue})' 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( parser.add_argument(
'-l', '--log-level', '-l', '--log-level',
type=LogLevel, type=LogLevel,
@ -112,6 +134,8 @@ def main():
p = Loader( p = Loader(
pulsar_host=args.pulsar_host, pulsar_host=args.pulsar_host,
output_queue=args.output_queue, output_queue=args.output_queue,
user=args.user,
collection=args.collection,
log_level=args.log_level, log_level=args.log_level,
) )

14
trustgraph-flow/trustgraph/chunking/recursive/chunker.py
View file

@ -7,7 +7,7 @@ as text as separate output objects.
from langchain_text_splitters import RecursiveCharacterTextSplitter from langchain_text_splitters import RecursiveCharacterTextSplitter
from prometheus_client import Histogram 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 ... schema import text_ingest_queue, chunk_ingest_queue
from ... log_level import LogLevel from ... log_level import LogLevel
from ... base import ConsumerProducer from ... base import ConsumerProducer
@ -55,7 +55,7 @@ class Processor(ConsumerProducer):
def handle(self, msg): def handle(self, msg):
v = msg.value() 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( texts = self.text_splitter.create_documents(
[v.text.decode("utf-8")] [v.text.decode("utf-8")]
@ -63,13 +63,15 @@ class Processor(ConsumerProducer):
for ix, chunk in enumerate(texts): for ix, chunk in enumerate(texts):
id = v.source.id + "-c" + str(ix) id = v.metadata.id + "-c" + str(ix)
r = Chunk( r = Chunk(
source=Source( metadata=Metadata(
source=v.source.source, source=v.metadata.source,
id=id, 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"), chunk=chunk.page_content.encode("utf-8"),
) )

14
trustgraph-flow/trustgraph/chunking/token/chunker.py
View file

@ -7,7 +7,7 @@ as text as separate output objects.
from langchain_text_splitters import TokenTextSplitter from langchain_text_splitters import TokenTextSplitter
from prometheus_client import Histogram 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 ... schema import text_ingest_queue, chunk_ingest_queue
from ... log_level import LogLevel from ... log_level import LogLevel
from ... base import ConsumerProducer from ... base import ConsumerProducer
@ -54,7 +54,7 @@ class Processor(ConsumerProducer):
def handle(self, msg): def handle(self, msg):
v = msg.value() 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( texts = self.text_splitter.create_documents(
[v.text.decode("utf-8")] [v.text.decode("utf-8")]
@ -62,13 +62,15 @@ class Processor(ConsumerProducer):
for ix, chunk in enumerate(texts): for ix, chunk in enumerate(texts):
id = v.source.id + "-c" + str(ix) id = v.metadata.id + "-c" + str(ix)
r = Chunk( r = Chunk(
source=Source( metadata=Metadata(
source=v.source.source, source=v.metadata.source,
id=id, 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"), chunk=chunk.page_content.encode("utf-8"),
) )

14
trustgraph-flow/trustgraph/decoding/pdf/pdf_decoder.py
View file

@ -8,7 +8,7 @@ import tempfile
import base64 import base64
from langchain_community.document_loaders import PyPDFLoader 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 ... schema import document_ingest_queue, text_ingest_queue
from ... log_level import LogLevel from ... log_level import LogLevel
from ... base import ConsumerProducer from ... base import ConsumerProducer
@ -45,7 +45,7 @@ class Processor(ConsumerProducer):
v = msg.value() 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: with tempfile.NamedTemporaryFile(delete_on_close=False) as fp:
@ -59,12 +59,14 @@ class Processor(ConsumerProducer):
for ix, page in enumerate(pages): for ix, page in enumerate(pages):
id = v.source.id + "-p" + str(ix) id = v.metadata.id + "-p" + str(ix)
r = TextDocument( r = TextDocument(
source=Source( metadata=Metadata(
source=v.source.source, source=v.metadata.source,
title=v.source.title, title=v.metadata.title,
id=id, id=id,
user=v.metadata.user,
collection=v.metadata.collection,
), ),
text=page.page_content.encode("utf-8"), text=page.page_content.encode("utf-8"),
) )

8
trustgraph-flow/trustgraph/embeddings/vectorize/vectorize.py
View file

@ -50,15 +50,15 @@ class Processor(ConsumerProducer):
subscriber=module + "-emb", 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) self.producer.send(r)
def handle(self, msg): def handle(self, msg):
v = msg.value() 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") chunk = v.chunk.decode("utf-8")
@ -67,7 +67,7 @@ class Processor(ConsumerProducer):
vectors = self.embeddings.request(chunk) vectors = self.embeddings.request(chunk)
self.emit( self.emit(
source=v.source, metadata=v.metadata,
chunk=chunk.encode("utf-8"), chunk=chunk.encode("utf-8"),
vectors=vectors vectors=vectors
) )

10
trustgraph-flow/trustgraph/extract/kg/definitions/extract.py
View file

@ -7,7 +7,7 @@ get entity definitions which are output as graph edges.
import urllib.parse import urllib.parse
import json 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 chunk_embeddings_ingest_queue, triples_store_queue
from .... schema import prompt_request_queue from .... schema import prompt_request_queue
from .... schema import prompt_response_queue from .... schema import prompt_response_queue
@ -69,15 +69,15 @@ class Processor(ConsumerProducer):
return self.prompt.request_definitions(chunk) 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) self.producer.send(t)
def handle(self, msg): def handle(self, msg):
v = msg.value() 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") chunk = v.chunk.decode("utf-8")
@ -101,7 +101,7 @@ class Processor(ConsumerProducer):
s_value = Value(value=str(s_uri), is_uri=True) s_value = Value(value=str(s_uri), is_uri=True)
o_value = Value(value=str(o), is_uri=False) 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: except Exception as e:
print("Exception: ", e, flush=True) print("Exception: ", e, flush=True)

23
trustgraph-flow/trustgraph/extract/kg/relationships/extract.py
View file

@ -9,7 +9,8 @@ import urllib.parse
import os import os
from pulsar.schema import JsonSchema 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 chunk_embeddings_ingest_queue, triples_store_queue
from .... schema import graph_embeddings_store_queue from .... schema import graph_embeddings_store_queue
from .... schema import prompt_request_queue from .... schema import prompt_request_queue
@ -91,20 +92,20 @@ class Processor(ConsumerProducer):
return self.prompt.request_relationships(chunk) 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) 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) self.vec_prod.send(r)
def handle(self, msg): def handle(self, msg):
v = msg.value() 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") chunk = v.chunk.decode("utf-8")
@ -139,6 +140,7 @@ class Processor(ConsumerProducer):
o_value = Value(value=str(o), is_uri=False) o_value = Value(value=str(o), is_uri=False)
self.emit_edge( self.emit_edge(
v.metdata,
s_value, s_value,
p_value, p_value,
o_value o_value
@ -146,6 +148,7 @@ class Processor(ConsumerProducer):
# Label for s # Label for s
self.emit_edge( self.emit_edge(
v.metdata,
s_value, s_value,
RDF_LABEL_VALUE, RDF_LABEL_VALUE,
Value(value=str(s), is_uri=False) Value(value=str(s), is_uri=False)
@ -153,6 +156,7 @@ class Processor(ConsumerProducer):
# Label for p # Label for p
self.emit_edge( self.emit_edge(
v.metdata,
p_value, p_value,
RDF_LABEL_VALUE, RDF_LABEL_VALUE,
Value(value=str(p), is_uri=False) Value(value=str(p), is_uri=False)
@ -161,15 +165,16 @@ class Processor(ConsumerProducer):
if rel.o_entity: if rel.o_entity:
# Label for o # Label for o
self.emit_edge( self.emit_edge(
v.metdata,
o_value, o_value,
RDF_LABEL_VALUE, RDF_LABEL_VALUE,
Value(value=str(o), is_uri=False) Value(value=str(o), is_uri=False)
) )
self.emit_vec(s_value, v.vectors) self.emit_vec(v.metadata, s_value, v.vectors)
self.emit_vec(p_value, v.vectors) self.emit_vec(v.metadata, p_value, v.vectors)
if rel.o_entity: if rel.o_entity:
self.emit_vec(o_value, v.vectors) self.emit_vec(v.metadata, o_value, v.vectors)
except Exception as e: except Exception as e:
print("Exception: ", e, flush=True) print("Exception: ", e, flush=True)

10
trustgraph-flow/trustgraph/extract/kg/topics/extract.py
View file

@ -7,7 +7,7 @@ get entity definitions which are output as graph edges.
import urllib.parse import urllib.parse
import json 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 chunk_embeddings_ingest_queue, triples_store_queue
from .... schema import prompt_request_queue from .... schema import prompt_request_queue
from .... schema import prompt_response_queue from .... schema import prompt_response_queue
@ -69,15 +69,15 @@ class Processor(ConsumerProducer):
return self.prompt.request_topics(chunk) 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) self.producer.send(t)
def handle(self, msg): def handle(self, msg):
v = msg.value() 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") chunk = v.chunk.decode("utf-8")
@ -101,7 +101,7 @@ class Processor(ConsumerProducer):
s_value = Value(value=str(s_uri), is_uri=True) s_value = Value(value=str(s_uri), is_uri=True)
o_value = Value(value=str(o), is_uri=False) 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: except Exception as e:
print("Exception: ", e, flush=True) print("Exception: ", e, flush=True)

16
trustgraph-flow/trustgraph/extract/object/row/extract.py
View file

@ -8,7 +8,7 @@ import urllib.parse
import os import os
from pulsar.schema import JsonSchema 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 RowSchema, Field
from .... schema import chunk_embeddings_ingest_queue, rows_store_queue from .... schema import chunk_embeddings_ingest_queue, rows_store_queue
from .... schema import object_embeddings_store_queue from .... schema import object_embeddings_store_queue
@ -124,24 +124,24 @@ class Processor(ConsumerProducer):
def get_rows(self, chunk): def get_rows(self, chunk):
return self.prompt.request_rows(self.schema, chunk) return self.prompt.request_rows(self.schema, chunk)
def emit_rows(self, source, rows): def emit_rows(self, metadata, rows):
t = 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) 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( 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) self.vec_prod.send(r)
def handle(self, msg): def handle(self, msg):
v = msg.value() 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") chunk = v.chunk.decode("utf-8")
@ -150,13 +150,13 @@ class Processor(ConsumerProducer):
rows = self.get_rows(chunk) rows = self.get_rows(chunk)
self.emit_rows( self.emit_rows(
source=v.source, metadata=v.metadata,
rows=rows rows=rows
) )
for row in rows: for row in rows:
self.emit_vec( self.emit_vec(
source=v.source, vec=v.vectors, metadata=v.metadata, vec=v.vectors,
name=self.schema.name, key_name=self.primary.name, name=self.schema.name, key_name=self.primary.name,
key=row[self.primary.name] key=row[self.primary.name]
) )