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

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

View file

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

View file

@ -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,
)

View file

@ -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,
)

View file

@ -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"),
)

View file

@ -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"),
)

View file

@ -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"),
)

View file

@ -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
)

View file

@ -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)

View file

@ -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)

View file

@ -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)

View file

@ -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]
)