trustgraph 0.5.2__tar.gz → 0.5.5__tar.gz

{trustgraph-0.5.2 → trustgraph-0.5.5}/PKG-INFO

@@ -1,9 +1,9 @@
 Metadata-Version: 2.1
 Name: trustgraph
-Version: 0.5.2
+Version: 0.5.5
 Summary: TrustGraph provides a means to run a pipeline of flexible AI processing components in a flexible means to achieve a processing pipeline.
 Home-page: https://github.com/trustgraph-ai/trustgraph
-Download-URL: https://github.com/trustgraph-ai/trustgraph/archive/refs/tags/v0.5.2.tar.gz
+Download-URL: https://github.com/trustgraph-ai/trustgraph/archive/refs/tags/v0.5.5.tar.gz
 Author: trustgraph.ai
 Author-email: security@trustgraph.ai
 Classifier: Programming Language :: Python :: 3

trustgraph-0.5.5/scripts/concat-parquet

@@ -0,0 +1,45 @@
+#!/usr/bin/env python3
+
+"""
+Concatenates multiple parquet files into a single parquet output
+"""
+
+import pyarrow as pa
+import pyarrow.parquet as pq
+import pandas as pd
+import sys
+import argparse
+
+parser = argparse.ArgumentParser(
+    prog="combine-parquet",
+    description=__doc__
+)
+
+parser.add_argument(
+    '-i', '--input',
+    nargs='*',
+    help=f'Input files'
+)
+
+parser.add_argument(
+    '-o', '--output',
+    help=f'Output files'
+)
+
+args = parser.parse_args()
+
+df = None
+
+for file in args.input:
+
+    part = pq.read_table(file).to_pandas()
+
+    if df is None:
+        df = part
+    else:
+        df = pd.concat([df, part], ignore_index=True)
+
+if df is not None:
+
+    table = pa.Table.from_pandas(df)
+    pq.write_table(table, args.output)

trustgraph-0.5.5/scripts/dump-parquet

@@ -0,0 +1,24 @@
+#!/usr/bin/env python3
+
+import pyarrow as pa
+import pyarrow.csv as pc
+import pyarrow.parquet as pq
+import pandas as pd
+import sys
+
+df = None
+
+for file in sys.argv[1:]:
+
+    part = pq.read_table(file).to_pandas()
+
+    if df is None:
+        df = part
+    else:
+        df = pd.concat([df, part], ignore_index=True)
+
+if df is not None:
+
+    table = pa.Table.from_pandas(df)
+    pc.write_csv(table, sys.stdout.buffer)
+

trustgraph-0.5.5/scripts/ge-dump-parquet

@@ -0,0 +1,6 @@
+#!/usr/bin/env python3
+
+from trustgraph.dump.graph_embeddings.parquet import run
+
+run()
+

trustgraph-0.5.5/scripts/load-graph-embeddings

@@ -0,0 +1,145 @@
+#!/usr/bin/env python3
+
+"""
+Loads Graph embeddings into TrustGraph processing.
+"""
+
+import pulsar
+from pulsar.schema import JsonSchema
+from trustgraph.schema import GraphEmbeddings, Value
+from trustgraph.schema import graph_embeddings_store_queue
+import argparse
+import os
+import time
+import pyarrow as pa
+import pyarrow.parquet as pq
+
+from trustgraph.log_level import LogLevel
+
+class Loader:
+
+    def __init__(
+            self,
+            pulsar_host,
+            output_queue,
+            log_level,
+            file,
+    ):
+
+        self.client = pulsar.Client(
+            pulsar_host,
+            logger=pulsar.ConsoleLogger(log_level.to_pulsar())
+        )
+
+        self.producer = self.client.create_producer(
+            topic=output_queue,
+            schema=JsonSchema(GraphEmbeddings),
+            chunking_enabled=True,
+        )
+
+        self.file = file
+
+    def run(self):
+
+        try:
+
+            path = self.file
+
+            print("Reading file...")
+            table = pq.read_table(path)
+            print("Loaded.")
+
+            names = set(table.column_names)
+
+            if "embeddings" not in names:
+                print("No 'embeddings' column")
+
+            if "entity" not in names:
+                print("No 'entity' column")
+
+            embc = table.column("embeddings")
+            entc = table.column("entity")
+
+            for emb, ent in zip(embc, entc):
+
+                b = emb.as_py()
+                n = ent.as_py()
+
+                r = GraphEmbeddings(
+                    vectors=b,
+                    entity=Value(
+                        value=n,
+                        is_uri=n.startswith("https:")
+                    )
+                )
+
+                self.producer.send(r)
+
+        except Exception as e:
+            print(e, flush=True)
+            
+    def __del__(self):
+        self.client.close()
+
+def main():
+
+    parser = argparse.ArgumentParser(
+        prog='loader',
+        description=__doc__,
+    )
+
+    default_pulsar_host = os.getenv("PULSAR_HOST", 'pulsar://localhost:6650')
+    default_output_queue = graph_embeddings_store_queue
+
+    parser.add_argument(
+        '-p', '--pulsar-host',
+        default=default_pulsar_host,
+        help=f'Pulsar host (default: {default_pulsar_host})',
+    )
+
+    parser.add_argument(
+        '-o', '--output-queue',
+        default=default_output_queue,
+        help=f'Output queue (default: {default_output_queue})'
+    )
+
+    parser.add_argument(
+        '-l', '--log-level',
+        type=LogLevel,
+        default=LogLevel.ERROR,
+        choices=list(LogLevel),
+        help=f'Output queue (default: info)'
+    )
+
+    parser.add_argument(
+        '-f', '--file',
+        required=True,
+        help=f'File to load'
+    )
+
+    args = parser.parse_args()
+
+    while True:
+
+        try:
+            p = Loader(
+                pulsar_host=args.pulsar_host,
+                output_queue=args.output_queue,
+                log_level=args.log_level,
+                file=args.file,
+            )
+
+            p.run()
+
+            print("File loaded.")
+            break
+
+        except Exception as e:
+
+            print("Exception:", e, flush=True)
+            print("Will retry...", flush=True)
+
+        time.sleep(10)
+
+main()
+

trustgraph-0.5.5/scripts/load-triples

@@ -0,0 +1,144 @@
+#!/usr/bin/env python3
+
+"""
+Loads Graph embeddings into TrustGraph processing.
+"""
+
+import pulsar
+from pulsar.schema import JsonSchema
+from trustgraph.schema import Triple, Value
+from trustgraph.schema import triples_store_queue
+import argparse
+import os
+import time
+import pyarrow as pa
+import pyarrow.parquet as pq
+
+from trustgraph.log_level import LogLevel
+
+class Loader:
+
+    def __init__(
+            self,
+            pulsar_host,
+            output_queue,
+            log_level,
+            file,
+    ):
+
+        self.client = pulsar.Client(
+            pulsar_host,
+            logger=pulsar.ConsoleLogger(log_level.to_pulsar())
+        )
+
+        self.producer = self.client.create_producer(
+            topic=output_queue,
+            schema=JsonSchema(Triple),
+            chunking_enabled=True,
+        )
+
+        self.file = file
+
+    def run(self):
+
+        try:
+
+            path = self.file
+
+            print("Reading file...")
+            table = pq.read_table(path)
+            print("Loaded.")
+
+            names = set(table.column_names)
+
+            if "s" not in names:
+                print("No 's' column")
+
+            if "p" not in names:
+                print("No 'p' column")
+
+            if "o" not in names:
+                print("No 'o' column")
+
+            sc = table.column("s")
+            pc = table.column("p")
+            oc = table.column("o")
+
+            for s, p, o in zip(sc, pc, oc):
+
+                r = Triple(
+                    s=Value(value=s.as_py(), is_uri=True),
+                    p=Value(value=p.as_py(), is_uri=True),
+                    o=Value(value=o.as_py(), is_uri=o.as_py().startswith("https:"))
+                )
+
+                self.producer.send(r)
+
+        except Exception as e:
+            print(e, flush=True)
+            
+    def __del__(self):
+        self.client.close()
+
+def main():
+
+    parser = argparse.ArgumentParser(
+        prog='loader',
+        description=__doc__,
+    )
+
+    default_pulsar_host = os.getenv("PULSAR_HOST", 'pulsar://localhost:6650')
+    default_output_queue = triples_store_queue
+
+    parser.add_argument(
+        '-p', '--pulsar-host',
+        default=default_pulsar_host,
+        help=f'Pulsar host (default: {default_pulsar_host})',
+    )
+
+    parser.add_argument(
+        '-o', '--output-queue',
+        default=default_output_queue,
+        help=f'Output queue (default: {default_output_queue})'
+    )
+
+    parser.add_argument(
+        '-l', '--log-level',
+        type=LogLevel,
+        default=LogLevel.ERROR,
+        choices=list(LogLevel),
+        help=f'Output queue (default: info)'
+    )
+
+    parser.add_argument(
+        '-f', '--file',
+        required=True,
+        help=f'File to load'
+    )
+
+    args = parser.parse_args()
+
+    while True:
+
+        try:
+            p = Loader(
+                pulsar_host=args.pulsar_host,
+                output_queue=args.output_queue,
+                log_level=args.log_level,
+                file=args.file,
+            )
+
+            p.run()
+
+            print("File loaded.")
+            break
+
+        except Exception as e:
+
+            print("Exception:", e, flush=True)
+            print("Will retry...", flush=True)
+
+        time.sleep(10)
+
+main()
+

{trustgraph-0.5.2 → trustgraph-0.5.5}/scripts/loader

@@ -6,7 +6,7 @@ Loads a PDF documented into TrustGraph processing.
 
 import pulsar
 from pulsar.schema import JsonSchema
-from trustgraph.schema import Document, Source
+from trustgraph.schema import Document, Source, document_ingest_queue
 import base64
 import hashlib
 import argparse
@@ -72,7 +72,7 @@ def main():
     )
 
     default_pulsar_host = os.getenv("PULSAR_HOST", 'pulsar://localhost:6650')
-    default_output_queue = 'document-load'
+    default_output_queue = document_ingest_queue
 
     parser.add_argument(
         '-p', '--pulsar-host',

{trustgraph-0.5.2 → trustgraph-0.5.5}/setup.py

@@ -4,7 +4,7 @@ import os
 with open("README.md", "r") as fh:
     long_description = fh.read()
 
-version = "0.5.2"
+version = "0.5.5"
 
 setuptools.setup(
     name="trustgraph",
@@ -48,9 +48,12 @@ setuptools.setup(
     ],
     scripts=[
         "scripts/chunker-recursive",
+        "scripts/concat-parquet",
+        "scripts/dump-parquet",
         "scripts/embeddings-hf",
         "scripts/embeddings-ollama",
         "scripts/embeddings-vectorize",
+        "scripts/ge-dump-parquet",
         "scripts/ge-write-milvus",
         "scripts/graph-rag",
         "scripts/graph-show",
@@ -58,6 +61,8 @@ setuptools.setup(
         "scripts/init-pulsar-manager",
         "scripts/kg-extract-definitions",
         "scripts/kg-extract-relationships",
+        "scripts/load-graph-embeddings",
+        "scripts/load-triples",
         "scripts/loader",
         "scripts/pdf-decoder",
         "scripts/query",
@@ -66,8 +71,7 @@ setuptools.setup(
         "scripts/text-completion-claude",
         "scripts/text-completion-ollama",
         "scripts/text-completion-vertexai",
-        "scripts/triples-write-cassandra",
-        "scripts/dump-parquet",
         "scripts/triples-dump-parquet",
+        "scripts/triples-write-cassandra",
     ]
 )

{trustgraph-0.5.2 → trustgraph-0.5.5}/trustgraph/base/base_processor.py

@@ -60,10 +60,10 @@ class BaseProcessor:
         )
 
         parser.add_argument(
-            '-M', '--metrics-enabled',
-            type=bool,
+            '--metrics',
+            action=argparse.BooleanOptionalAction,
             default=True,
-            help=f'Pulsar host (default: true)',
+            help=f'Metrics enabled (default: true)',
         )
 
         parser.add_argument(
@@ -79,20 +79,22 @@ class BaseProcessor:
     @classmethod
     def start(cls, prog, doc):
 
-        while True:
+        parser = argparse.ArgumentParser(
+            prog=prog,
+            description=doc
+        )
 
-            parser = argparse.ArgumentParser(
-                prog=prog,
-                description=doc
-            )
+        cls.add_args(parser)
 
-            cls.add_args(parser)
+        args = parser.parse_args()
+        args = vars(args)
 
-            args = parser.parse_args()
-            args = vars(args)
+        print(args)
 
-            if args["metrics_enabled"]:
-                start_http_server(args["metrics_port"])
+        if args["metrics"]:
+            start_http_server(args["metrics_port"])
+
+        while True:
 
             try:
 

trustgraph-0.5.5/trustgraph/dump/graph_embeddings/parquet/processor.py

@@ -0,0 +1,87 @@
+
+"""
+Write graph embeddings to parquet files in a directory.
+"""
+
+import pulsar
+import base64
+import os
+import argparse
+import time
+
+from .... trustgraph import TrustGraph
+from .... schema import GraphEmbeddings
+from .... schema import graph_embeddings_store_queue
+from .... log_level import LogLevel
+from .... base import Consumer
+
+from . writer import ParquetWriter
+
+module = ".".join(__name__.split(".")[1:-1])
+
+default_input_queue = graph_embeddings_store_queue
+default_subscriber = module
+default_graph_host='localhost'
+default_directory = "."
+default_file_template = "graph-embeds-{id}.parquet"
+default_rotation_time = 60
+        
+class Processor(Consumer):
+
+    def __init__(self, **params):
+        
+        input_queue = params.get("input_queue", default_input_queue)
+        subscriber = params.get("subscriber", default_subscriber)
+        directory = params.get("directory", default_directory)
+        file_template = params.get("file_template", default_file_template)
+        rotation_time = params.get("rotation_time", default_rotation_time)
+
+        super(Processor, self).__init__(
+            **params | {
+                "input_queue": input_queue,
+                "subscriber": subscriber,
+                "input_schema": GraphEmbeddings,
+            }
+        )
+
+        self.writer = ParquetWriter(directory, file_template, rotation_time)
+
+    def __del__(self):
+        if hasattr(self, "writer"):
+            del self.writer
+
+    def handle(self, msg):
+
+        v = msg.value()
+        self.writer.write(v.vectors, v.entity.value)
+
+    @staticmethod
+    def add_args(parser):
+
+        Consumer.add_args(
+            parser, default_input_queue, default_subscriber,
+        )
+
+        parser.add_argument(
+            '-d', '--directory',
+            default=default_directory,
+            help=f'Directory to write to (default: {default_directory})'
+        )
+
+        parser.add_argument(
+            '-f', '--file-template',
+            default=default_file_template,
+            help=f'Directory to write to (default: {default_file_template})'
+        )
+
+        parser.add_argument(
+            '-t', '--rotation-time',
+            type=int,
+            default=default_rotation_time,
+            help=f'Rotation time / seconds (default: {default_rotation_time})'
+        )
+
+def run():
+
+    Processor.start(module, __doc__)
+

trustgraph-0.5.5/trustgraph/dump/graph_embeddings/parquet/writer.py

@@ -0,0 +1,94 @@
+
+import threading
+import queue
+import time
+import uuid
+import pyarrow as pa
+import pyarrow.parquet as pq
+
+class ParquetWriter:
+
+    def __init__(self, directory, file_template, rotation_time):
+        self.directory = directory
+        self.file_template = file_template
+        self.rotation_time = rotation_time
+
+        self.q = queue.Queue()
+
+        self.running = True
+
+        self.thread = threading.Thread(target=(self.writer_thread))
+        self.thread.start()
+
+    def writer_thread(self):
+
+        items = []
+
+        timeout = None
+
+        while self.running:
+
+            try:
+
+                item = self.q.get(timeout=1)
+
+                if timeout == None:
+                    timeout = time.time() + self.rotation_time
+
+                items.append(item)
+
+            except queue.Empty:
+                pass
+
+            if timeout:
+                if time.time() > timeout:
+
+                    self.write_file(items)
+                    timeout = None
+                    items = []
+
+    def write_file(self, items):
+
+        try:
+
+            schema = pa.schema([
+                pa.field('embeddings', pa.list_(pa.list_(pa.float64()))),
+                pa.field('entity', pa.string()),
+            ])
+
+            fname = self.file_template.format(id=str(uuid.uuid4()))
+            path = f"{self.directory}/{fname}"
+
+            writer = pq.ParquetWriter(path, schema)
+
+            batch = pa.record_batch(
+                [
+                    [i[0] for i in items],
+                    [i[1] for i in items],
+                ],
+                names=['embeddings', 'entity']
+            )
+
+            writer.write_batch(batch)
+
+            writer.close()
+
+            print(f"Wrote {path}.")
+
+        except Exception as e:
+
+            print("Parquet write:", e)
+
+    def write(self, embeds, ent):
+        self.q.put((embeds, ent))
+
+    def __del__(self):
+
+        self.running = False
+
+        if hasattr(self, "q"):
+            self.thread.join()
+
+            
+
+        

trustgraph-0.5.5/trustgraph/embeddings/ollama/__init__.py

@@ -0,0 +1,3 @@
+
+from . processor import *
+

{trustgraph-0.5.2 → trustgraph-0.5.5}/trustgraph/embeddings/vectorize/vectorize.py

@@ -36,7 +36,7 @@ class Processor(ConsumerProducer):
 
         self.embeddings = EmbeddingsClient(
             pulsar_host=self.pulsar_host,
-            subscriber=module + "emb",
+            subscriber=module + "-emb",
         )
 
     def emit(self, source, chunk, vectors):

{trustgraph-0.5.2 → trustgraph-0.5.5}/trustgraph/kg/extract_definitions/extract.py

@@ -87,6 +87,9 @@ class Processor(ConsumerProducer):
 
                 o = defn["definition"]
 
+                if s == "": continue
+                if o == "": continue
+
                 s_value = Value(value=str(s_uri), is_uri=True)
                 o_value = Value(value=str(o), is_uri=False)