trustgraph 0.2.4__tar.gz → 0.3.1__tar.gz

{trustgraph-0.2.4 → trustgraph-0.3.1}/PKG-INFO

@@ -1,9 +1,9 @@
 Metadata-Version: 2.1
 Name: trustgraph
-Version: 0.2.4
+Version: 0.3.1
 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.2.4.tar.gz
+Download-URL: https://github.com/trustgraph-ai/trustgraph/archive/refs/tags/v0.3.1.tar.gz
 Author: trustgraph.ai
 Author-email: security@trustgraph.ai
 Classifier: Programming Language :: Python :: 3

{trustgraph-0.2.4 → trustgraph-0.3.1}/setup.py

@@ -4,7 +4,7 @@ import os
 with open("README.md", "r") as fh:
     long_description = fh.read()
 
-version = "0.2.4"
+version = "0.3.1"
 
 setuptools.setup(
     name="trustgraph",

trustgraph-0.3.1/trustgraph/base/processor.py

@@ -0,0 +1,268 @@
+
+import os
+import argparse
+import pulsar
+import time
+from pulsar.schema import JsonSchema
+
+from .. log_level import LogLevel
+
+class BaseProcessor:
+
+    default_pulsar_host = os.getenv("PULSAR_HOST", 'pulsar://pulsar:6650')
+
+    def __init__(
+            self,
+            pulsar_host=default_pulsar_host,
+            log_level=LogLevel.INFO,
+    ):
+
+        self.client = None
+
+        if pulsar_host == None:
+            pulsar_host = default_pulsar_host
+
+        self.pulsar_host = pulsar_host
+
+        self.client = pulsar.Client(
+            pulsar_host,
+            logger=pulsar.ConsoleLogger(log_level.to_pulsar())
+        )
+
+    def __del__(self):
+
+        if self.client:
+            self.client.close()
+
+    @staticmethod
+    def add_args(parser):
+
+        parser.add_argument(
+            '-p', '--pulsar-host',
+            default=__class__.default_pulsar_host,
+            help=f'Pulsar host (default: {__class__.default_pulsar_host})',
+        )
+
+        parser.add_argument(
+            '-l', '--log-level',
+            type=LogLevel,
+            default=LogLevel.INFO,
+            choices=list(LogLevel),
+            help=f'Output queue (default: info)'
+        )
+
+    def run(self):
+        raise RuntimeError("Something should have implemented the run method")
+
+    @classmethod
+    def start(cls, prog, doc):
+
+        while True:
+
+            parser = argparse.ArgumentParser(
+                prog=prog,
+                description=doc
+            )
+
+            cls.add_args(parser)
+
+            args = parser.parse_args()
+            args = vars(args)
+
+            try:
+
+                p = cls(**args)
+                p.run()
+
+            except Exception as e:
+
+                print("Exception:", e, flush=True)
+                print("Will retry...", flush=True)
+
+                time.sleep(10)
+
+class Consumer(BaseProcessor):
+
+    def __init__(
+            self,
+            pulsar_host=None,
+            log_level=LogLevel.INFO,
+            input_queue="input",
+            subscriber="subscriber",
+            input_schema=None,
+    ):
+
+        super(Consumer, self).__init__(
+            pulsar_host=pulsar_host,
+            log_level=log_level,
+        )
+
+        if input_schema == None:
+            raise RuntimeError("input_schema must be specified")
+
+        self.consumer = self.client.subscribe(
+            input_queue, subscriber,
+            schema=JsonSchema(input_schema),
+        )
+
+    def run(self):
+
+        while True:
+
+            msg = self.consumer.receive()
+
+            try:
+
+                self.handle(msg)
+
+                # Acknowledge successful processing of the message
+                self.consumer.acknowledge(msg)
+
+            except Exception as e:
+
+                print("Exception:", e, flush=True)
+
+                # Message failed to be processed
+                self.consumer.negative_acknowledge(msg)
+
+    @staticmethod
+    def add_args(parser, default_input_queue, default_subscriber):
+
+        BaseProcessor.add_args(parser)
+
+        parser.add_argument(
+            '-i', '--input-queue',
+            default=default_input_queue,
+            help=f'Input queue (default: {default_input_queue})'
+        )
+
+        parser.add_argument(
+            '-s', '--subscriber',
+            default=default_subscriber,
+            help=f'Queue subscriber name (default: {default_subscriber})'
+        )
+
+class ConsumerProducer(BaseProcessor):
+
+    def __init__(
+            self,
+            pulsar_host=None,
+            log_level=LogLevel.INFO,
+            input_queue="input",
+            output_queue="output",
+            subscriber="subscriber",
+            input_schema=None,
+            output_schema=None,
+    ):
+
+        super(ConsumerProducer, self).__init__(
+            pulsar_host=pulsar_host,
+            log_level=log_level,
+        )
+
+        if input_schema == None:
+            raise RuntimeError("input_schema must be specified")
+
+        if output_schema == None:
+            raise RuntimeError("output_schema must be specified")
+
+        self.consumer = self.client.subscribe(
+            input_queue, subscriber,
+            schema=JsonSchema(input_schema),
+        )
+
+        self.producer = self.client.create_producer(
+            topic=output_queue,
+            schema=JsonSchema(output_schema),
+        )
+
+    def run(self):
+
+        while True:
+
+            msg = self.consumer.receive()
+
+            try:
+
+                resp = self.handle(msg)
+
+                # Acknowledge successful processing of the message
+                self.consumer.acknowledge(msg)
+
+            except Exception as e:
+
+                print("Exception:", e, flush=True)
+
+                # Message failed to be processed
+                self.consumer.negative_acknowledge(msg)
+
+    def send(self, msg, properties={}):
+
+        self.producer.send(msg, properties)
+
+    @staticmethod
+    def add_args(
+            parser, default_input_queue, default_subscriber,
+            default_output_queue,
+    ):
+
+        BaseProcessor.add_args(parser)
+
+        parser.add_argument(
+            '-i', '--input-queue',
+            default=default_input_queue,
+            help=f'Input queue (default: {default_input_queue})'
+        )
+
+        parser.add_argument(
+            '-s', '--subscriber',
+            default=default_subscriber,
+            help=f'Queue subscriber name (default: {default_subscriber})'
+        )
+
+        parser.add_argument(
+            '-o', '--output-queue',
+            default=default_output_queue,
+            help=f'Output queue (default: {default_output_queue})'
+        )
+
+class Producer(BaseProcessor):
+
+    def __init__(
+            self,
+            pulsar_host=None,
+            log_level=LogLevel.INFO,
+            output_queue="output",
+            output_schema=None,
+    ):
+
+        super(Producer, self).__init__(
+            pulsar_host=pulsar_host,
+            log_level=log_level,
+        )
+
+        if output_schema == None:
+            raise RuntimeError("output_schema must be specified")
+
+        self.producer = self.client.create_producer(
+            topic=output_queue,
+            schema=JsonSchema(output_schema),
+        )
+
+    def send(self, msg, properties={}):
+
+        self.producer.send(msg, properties)
+
+    @staticmethod
+    def add_args(
+            parser, default_input_queue, default_subscriber,
+            default_output_queue,
+    ):
+
+        BaseProcessor.add_args(parser)
+
+        parser.add_argument(
+            '-o', '--output-queue',
+            default=default_output_queue,
+            help=f'Output queue (default: {default_output_queue})'
+        )

trustgraph-0.3.1/trustgraph/chunker/recursive/chunker.py

@@ -0,0 +1,99 @@
+
+"""
+Simple decoder, accepts text documents on input, outputs chunks from the
+as text as separate output objects.
+"""
+
+from langchain_text_splitters import RecursiveCharacterTextSplitter
+
+
+from ... schema import TextDocument, Chunk, Source
+from ... log_level import LogLevel
+from ... base import ConsumerProducer
+
+default_input_queue = 'text-doc-load'
+default_output_queue = 'chunk-load'
+default_subscriber = 'chunker-recursive'
+
+class Processor(ConsumerProducer):
+
+    def __init__(
+            self,
+            pulsar_host=None,
+            input_queue=default_input_queue,
+            output_queue=default_output_queue,
+            subscriber=default_subscriber,
+            log_level=LogLevel.INFO,
+            chunk_size=2000,
+            chunk_overlap=100,
+    ):
+
+        super(Processor, self).__init__(
+            pulsar_host=pulsar_host,
+            log_level=log_level,
+            input_queue=input_queue,
+            output_queue=output_queue,
+            subscriber=subscriber,
+            input_schema=TextDocument,
+            output_schema=Chunk,
+        )
+
+        self.text_splitter = RecursiveCharacterTextSplitter(
+            chunk_size=chunk_size,
+            chunk_overlap=chunk_overlap,
+            length_function=len,
+            is_separator_regex=False,
+        )
+
+    def handle(self, msg):
+
+        v = msg.value()
+        print(f"Chunking {v.source.id}...", flush=True)
+
+        texts = self.text_splitter.create_documents(
+            [v.text.decode("utf-8")]
+        )
+
+        for ix, chunk in enumerate(texts):
+
+            id = v.source.id + "-c" + str(ix)
+
+            r = Chunk(
+                source=Source(
+                    source=v.source.source,
+                    id=id,
+                    title=v.source.title
+                ),
+                chunk=chunk.page_content.encode("utf-8"),
+            )
+
+            self.send(r)
+
+            print("Done.", flush=True)
+
+    @staticmethod
+    def add_args(parser):
+
+        ConsumerProducer.add_args(
+            parser, default_input_queue, default_subscriber,
+            default_output_queue,
+        )
+
+        parser.add_argument(
+            '-z', '--chunk-size',
+            type=int,
+            default=2000,
+            help=f'Chunk size (default: 2000)'
+        )
+
+        parser.add_argument(
+            '-v', '--chunk-overlap',
+            type=int,
+            default=100,
+            help=f'Chunk overlap (default: 100)'
+        )
+
+def run():
+
+    Processor.start('chunker', __doc__)
+

trustgraph-0.3.1/trustgraph/decoder/pdf/pdf_decoder.py

@@ -0,0 +1,87 @@
+
+"""
+Simple decoder, accepts PDF documents on input, outputs pages from the
+PDF document as text as separate output objects.
+"""
+
+import tempfile
+import base64
+from langchain_community.document_loaders import PyPDFLoader
+
+from ... schema import Document, TextDocument, Source
+from ... log_level import LogLevel
+from ... base import ConsumerProducer
+
+default_input_queue = 'document-load'
+default_output_queue = 'text-doc-load'
+default_subscriber = 'pdf-decoder'
+
+class Processor(ConsumerProducer):
+
+    def __init__(
+            self,
+            pulsar_host=None,
+            input_queue=default_input_queue,
+            output_queue=default_output_queue,
+            subscriber=default_subscriber,
+            log_level=LogLevel.INFO,
+    ):
+
+        super(Processor, self).__init__(
+            pulsar_host=pulsar_host,
+            log_level=log_level,
+            input_queue=input_queue,
+            output_queue=output_queue,
+            subscriber=subscriber,
+            input_schema=Document,
+            output_schema=TextDocument,
+        )
+
+        print("PDF inited")
+
+    def handle(self, msg):
+
+        print("PDF message received")
+
+        v = msg.value()
+
+        print(f"Decoding {v.source.id}...", flush=True)
+
+        with tempfile.NamedTemporaryFile(delete_on_close=False) as fp:
+
+            fp.write(base64.b64decode(v.data))
+            fp.close()
+
+            with open(fp.name, mode='rb') as f:
+
+                loader = PyPDFLoader(fp.name)
+                pages = loader.load()
+
+                for ix, page in enumerate(pages):
+
+                    id = v.source.id + "-p" + str(ix)
+                    r = TextDocument(
+                        source=Source(
+                            source=v.source.source,
+                            title=v.source.title,
+                            id=id,
+                        ),
+                        text=page.page_content.encode("utf-8"),
+                    )
+
+                    self.send(r)
+
+        print("Done.", flush=True)
+
+    @staticmethod
+    def add_args(parser):
+
+        ConsumerProducer.add_args(
+            parser, default_input_queue, default_subscriber,
+            default_output_queue,
+        )
+
+def run():
+
+    Processor.start("pdf-decoder", __doc__)
+

trustgraph-0.3.1/trustgraph/embeddings/hf/hf.py

@@ -0,0 +1,77 @@
+
+"""
+Embeddings service, applies an embeddings model selected from HuggingFace.
+Input is text, output is embeddings vector.
+"""
+
+from langchain_huggingface import HuggingFaceEmbeddings
+
+from ... schema import EmbeddingsRequest, EmbeddingsResponse
+from ... log_level import LogLevel
+from ... base import ConsumerProducer
+
+default_input_queue = 'embeddings'
+default_output_queue = 'embeddings-response'
+default_subscriber = 'embeddings-hf'
+default_model="all-MiniLM-L6-v2"
+
+class Processor(ConsumerProducer):
+
+    def __init__(
+            self,
+            pulsar_host=None,
+            input_queue=default_input_queue,
+            output_queue=default_output_queue,
+            subscriber=default_subscriber,
+            log_level=LogLevel.INFO,
+            model=default_model,
+    ):
+
+        super(Processor, self).__init__(
+            pulsar_host=pulsar_host,
+            log_level=log_level,
+            input_queue=input_queue,
+            output_queue=output_queue,
+            subscriber=subscriber,
+            input_schema=EmbeddingsRequest,
+            output_schema=EmbeddingsResponse,
+        )
+
+        self.embeddings = HuggingFaceEmbeddings(model_name=model)
+
+    def handle(self, msg):
+
+        v = msg.value()
+
+        # Sender-produced ID
+        id = msg.properties()["id"]
+
+        print(f"Handling input {id}...", flush=True)
+
+        text = v.text
+        embeds = self.embeddings.embed_documents([text])
+
+        print("Send response...", flush=True)
+        r = EmbeddingsResponse(vectors=embeds)
+        self.producer.send(r, properties={"id": id})
+
+        print("Done.", flush=True)
+
+    @staticmethod
+    def add_args(parser):
+
+        ConsumerProducer.add_args(
+            parser, default_input_queue, default_subscriber,
+            default_output_queue,
+        )
+
+        parser.add_argument(
+            '-m', '--model',
+            default="all-MiniLM-L6-v2",
+            help=f'LLM model (default: all-MiniLM-L6-v2)'
+        )
+
+def run():
+
+    Processor.start("embeddings-hf", __doc__)
+

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

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

trustgraph-0.3.1/trustgraph/embeddings/ollama/processor.py

@@ -0,0 +1,86 @@
+
+"""
+Embeddings service, applies an embeddings model selected from HuggingFace.
+Input is text, output is embeddings vector.
+"""
+from langchain_community.embeddings import OllamaEmbeddings
+
+from ... schema import EmbeddingsRequest, EmbeddingsResponse
+from ... log_level import LogLevel
+from ... base import ConsumerProducer
+
+default_input_queue = 'embeddings'
+default_output_queue = 'embeddings-response'
+default_subscriber = 'embeddings-ollama'
+default_model="mxbai-embed-large"
+default_ollama = 'http://localhost:11434'
+
+class Processor(ConsumerProducer):
+
+    def __init__(
+            self,
+            pulsar_host=None,
+            input_queue=default_input_queue,
+            output_queue=default_output_queue,
+            subscriber=default_subscriber,
+            log_level=LogLevel.INFO,
+            model=default_model,
+            ollama=default_ollama,
+    ):
+
+        super(Processor, self).__init__(
+            pulsar_host=pulsar_host,
+            log_level=log_level,
+            input_queue=input_queue,
+            output_queue=output_queue,
+            subscriber=subscriber,
+            input_schema=EmbeddingsRequest,
+            output_schema=EmbeddingsResponse,
+        )
+
+        self.embeddings = OllamaEmbeddings(base_url=ollama, model=model)
+
+    def handle(self, msg):
+
+        v = msg.value()
+
+        # Sender-produced ID
+
+        id = msg.properties()["id"]
+
+        print(f"Handling input {id}...", flush=True)
+
+        text = v.text
+        embeds = self.embeddings.embed_query([text])
+
+        print("Send response...", flush=True)
+        r = EmbeddingsResponse(vectors=[embeds])
+
+        self.producer.send(r, properties={"id": id})
+
+        print("Done.", flush=True)
+
+    @staticmethod
+    def add_args(parser):
+
+        ConsumerProducer.add_args(
+            parser, default_input_queue, default_subscriber,
+            default_output_queue,
+        )
+
+        parser.add_argument(
+            '-m', '--model',
+            default=default_model,
+            help=f'Embeddings model (default: {default_model})'
+        )
+
+        parser.add_argument(
+            '-r', '--ollama',
+            default=default_ollama,
+            help=f'ollama (default: {default_ollama})'
+        )
+
+def run():
+
+    Processor.start('embeddings-ollama', __doc__)
+