Standard core complete

trustgraph/chunker/recursive/chunker.py

@@ -4,28 +4,19 @@ Simple decoder, accepts text documents on input, outputs chunks from the
 as text as separate output objects.
 """
 
-import pulsar
-from pulsar.schema import JsonSchema
-import tempfile
-import base64
-import os
-import argparse
-from langchain_text_splitters import RecursiveCharacterTextSplitter
-import time
-
 from ... schema import TextDocument, Chunk, Source
 from ... log_level import LogLevel
+from ... base import ConsumerProducer
 
-default_pulsar_host = os.getenv("PULSAR_HOST", 'pulsar://pulsar:6650')
 default_input_queue = 'text-doc-load'
 default_output_queue = 'chunk-load'
 default_subscriber = 'chunker-recursive'
 
-class Processor:
+class Processor(ConsumerProducer):
 
     def __init__(
             self,
-            pulsar_host=default_pulsar_host,
+            pulsar_host=None,
             input_queue=default_input_queue,
             output_queue=default_output_queue,
             subscriber=default_subscriber,
@@ -34,21 +25,14 @@ class Processor:
             chunk_overlap=100,
     ):
 
-        self.client = None
+        super(Processor, self).__init__(
-        
+            pulsar_host=pulsar_host,
-        self.client = pulsar.Client(
+            log_level=log_level,
-            pulsar_host,
+            input_queue=input_queue,
-            logger=pulsar.ConsoleLogger(log_level.to_pulsar())
+            output_queue=output_queue,
-        )
+            subscriber=subscriber,
-
+            input_schema=TextDocument,
-        self.consumer = self.client.subscribe(
+            output_schema=Chunk,
-            input_queue, subscriber,
-            schema=JsonSchema(TextDocument),
-        )
-
-        self.producer = self.client.create_producer(
-            topic=output_queue,
-            schema=JsonSchema(Chunk),
         )
 
         self.text_splitter = RecursiveCharacterTextSplitter(
@@ -58,18 +42,7 @@ class Processor:
             is_separator_regex=False,
         )
 
-        print("Chunker inited")
+    def handle(self, msg):
-
-    def run(self):
-
-        print("Chunker running")
-
-        while True:
-
-            msg = self.consumer.receive()
-            print("Chunker message received")
-
-            try:
 
         v = msg.value()
         print(f"Chunking {v.source.id}...", flush=True)
@@ -91,61 +64,16 @@ class Processor:
                 chunk=chunk.page_content.encode("utf-8"),
             )
 
-                    self.producer.send(r)
+            self.send(r)
-
-                # Acknowledge successful processing of the message
-                self.consumer.acknowledge(msg)
 
             print("Done.", flush=True)
 
-            except Exception as e:
+    @staticmethod
-                print(e, flush=True)
+    def add_args(parser):
 
-                # Message failed to be processed
+        ConsumerProducer.add_args(
-                self.consumer.negative_acknowledge(msg)
+            parser, default_input_queue, default_subscriber,
-
+            default_output_queue,
-    def __del__(self):
-
-        if self.client:
-            self.client.close()
-
-def run():
-
-    parser = argparse.ArgumentParser(
-        prog='pdf-decoder',
-        description=__doc__,
-    )
-
-    parser.add_argument(
-        '-p', '--pulsar-host',
-        default=default_pulsar_host,
-        help=f'Pulsar host (default: {default_pulsar_host})',
-    )
-
-    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})'
-    )
-
-    parser.add_argument(
-        '-l', '--log-level',
-        type=LogLevel,
-        default=LogLevel.INFO,
-        choices=list(LogLevel),
-        help=f'Output queue (default: info)'
         )
 
         parser.add_argument(
@@ -162,30 +90,7 @@ def run():
             help=f'Chunk overlap (default: 100)'
         )
 
-    args = parser.parse_args()
+def run():
-
-
-    while True:
-
-        try:
-
-            p = Processor(
-                pulsar_host=args.pulsar_host,
-                input_queue=args.input_queue,
-                output_queue=args.output_queue,
-                subscriber=args.subscriber,
-                log_level=args.log_level,
-                chunk_size=args.chunk_size,
-                chunk_overlap=args.chunk_overlap,
-            )
-
-            p.run()
-
-        except Exception as e:
-
-            print("Exception:", e, flush=True)
-            print("Will retry...", flush=True)
-
-        time.sleep(10)
 
+    Processor.start('chunker', __doc__)
 

trustgraph/decoder/pdf/pdf_decoder.py

@@ -4,71 +4,45 @@ Simple decoder, accepts PDF documents on input, outputs pages from the
 PDF document as text as separate output objects.
 """
 
-import pulsar
-from pulsar.schema import JsonSchema
 from langchain_community.document_loaders import PyPDFLoader
-import tempfile
-import base64
-import os
-import argparse
-import time
 
 from ... schema import Document, TextDocument, Source
 from ... log_level import LogLevel
+from ... base import ConsumerProducer
 
-default_pulsar_host = os.getenv("PULSAR_HOST", 'pulsar://pulsar:6650')
 default_input_queue = 'document-load'
 default_output_queue = 'text-doc-load'
 default_subscriber = 'pdf-decoder'
 
-class Processor:
+class Processor(ConsumerProducer):
 
     def __init__(
             self,
-            pulsar_host=default_pulsar_host,
+            pulsar_host=None,
             input_queue=default_input_queue,
             output_queue=default_output_queue,
             subscriber=default_subscriber,
             log_level=LogLevel.INFO,
     ):
 
-        self.client = None
+        super(Processor, self).__init__(
-
+            pulsar_host=pulsar_host,
-        self.client = pulsar.Client(
+            log_level=log_level,
-            pulsar_host,
+            input_queue=input_queue,
-            logger=pulsar.ConsoleLogger(log_level.to_pulsar())
+            output_queue=output_queue,
-        )
+            subscriber=subscriber,
-
+            input_schema=Document,
-        self.consumer = self.client.subscribe(
+            output_schema=TextDocument,
-            input_queue, subscriber,
-            schema=JsonSchema(Document),
-        )
-
-        self.producer = self.client.create_producer(
-            topic=output_queue,
-            schema=JsonSchema(TextDocument),
         )
 
         print("PDF inited")
 
-        print("Pulsar", pulsar_host)
+    def handle(self, msg):
-        print("Input", input_queue)
-        print("Output", output_queue)
-        print("Subscriber", subscriber)
-
-    def run(self):
-
-        print("PDF running")
-
-        while True:
-
-            msg = self.consumer.receive()
 
         print("PDF message received")
 
-            try:
-
         v = msg.value()
+
         print(f"Decoding {v.source.id}...", flush=True)
 
         with tempfile.NamedTemporaryFile(delete_on_close=False) as fp:
@@ -93,82 +67,19 @@ class Processor:
                         text=page.page_content.encode("utf-8"),
                     )
 
-                            self.producer.send(r)
+                    self.send(r)
-
-                # Acknowledge successful processing of the message
-                self.consumer.acknowledge(msg)
 
         print("Done.", flush=True)
 
-            except Exception as e:
+    @staticmethod
-                print(e, flush=True)
+    def add_args(parser):
 
-                # Message failed to be processed
+        ConsumerProducer.add_args(
-                self.consumer.negative_acknowledge(msg)
+            parser, default_input_queue, default_subscriber,
-
+            default_output_queue,
-    def __del__(self):
+        )
-
-        if self.client:
-            self.client.close()
 
 def run():
 
-    parser = argparse.ArgumentParser(
+    Processor.start("pdf-decoder", __doc__)
-        prog='pdf-decoder',
-        description=__doc__,
-    )
-
-    parser.add_argument(
-        '-p', '--pulsar-host',
-        default=default_pulsar_host,
-        help=f'Pulsar host (default: {default_pulsar_host})',
-    )
-
-    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})'
-    )
-
-    parser.add_argument(
-        '-l', '--log-level',
-        type=LogLevel,
-        default=LogLevel.INFO,
-        choices=list(LogLevel),
-        help=f'Output queue (default: info)'
-    )
-
-    args = parser.parse_args()
-
-    while True:
-
-        try:
-            p = Processor(
-                pulsar_host=args.pulsar_host,
-                input_queue=args.input_queue,
-                output_queue=args.output_queue,
-                subscriber=args.subscriber,
-                log_level=args.log_level,
-            )
-
-            p.run()
-
-        except Exception as e:
-
-            print("Exception:", e, flush=True)
-            print("Will retry...", flush=True)
-
-        time.sleep(10)
 

trustgraph/embeddings/vectorize/vectorize.py

@@ -4,49 +4,34 @@ Vectorizer, calls the embeddings service to get embeddings for a chunk.
 Input is text chunk, output is chunk and vectors.
 """
 
-import pulsar
-from pulsar.schema import JsonSchema
-import tempfile
-import base64
-import os
-import argparse
-import time
-
 from ... schema import Chunk, VectorsChunk
 from ... embeddings_client import EmbeddingsClient
 from ... log_level import LogLevel
+from ... base import ConsumerProducer
 
-default_pulsar_host = os.getenv("PULSAR_HOST", 'pulsar://pulsar:6650')
 default_input_queue = 'chunk-load'
 default_output_queue = 'vectors-chunk-load'
 default_subscriber = 'embeddings-vectorizer'
 
-class Processor:
+class Processor(ConsumerProducer):
 
     def __init__(
             self,
-            pulsar_host=default_pulsar_host,
+            pulsar_host=None,
             input_queue=default_input_queue,
             output_queue=default_output_queue,
             subscriber=default_subscriber,
             log_level=LogLevel.INFO,
     ):
 
-        self.client = None
+        super(Processor, self).__init__(
-        
+            pulsar_host=pulsar_host,
-        self.client = pulsar.Client(
+            log_level=log_level,
-            pulsar_host,
+            input_queue=input_queue,
-            logger=pulsar.ConsoleLogger(log_level.to_pulsar())
+            output_queue=output_queue,
-        )
+            subscriber=subscriber,
-
+            input_schema=Chunk,
-        self.consumer = self.client.subscribe(
+            output_schema=VectorsChunk,
-            input_queue, subscriber,
-            schema=JsonSchema(Chunk),
-        )
-
-        self.producer = self.client.create_producer(
-            topic=output_queue,
-            schema=JsonSchema(VectorsChunk),
         )
 
         self.embeddings = EmbeddingsClient(pulsar_host=pulsar_host)
@@ -56,13 +41,7 @@ class Processor:
         r = VectorsChunk(source=source, chunk=chunk, vectors=vectors)
         self.producer.send(r)
 
-    def run(self):
+    def handle(self, msg):
-
-        while True:
-
-            msg = self.consumer.receive()
-
-            try:
 
         v = msg.value()
         print(f"Indexing {v.source.id}...", flush=True)
@@ -84,80 +63,15 @@ class Processor:
 
         print("Done.", flush=True)
 
-                # Acknowledge successful processing of the message
+    @staticmethod
-                self.consumer.acknowledge(msg)
+    def add_args(parser):
 
-            except Exception as e:
+        ConsumerProducer.add_args(
-
+            parser, default_input_queue, default_subscriber,
-                print("Exception:", e, flush=True)
+            default_output_queue,
-
+        )
-                # Message failed to be processed
-                self.consumer.negative_acknowledge(msg)
-
-    def __del__(self):
-
-        if self.client:
-            self.client.close()
 
 def run():
 
-    parser = argparse.ArgumentParser(
+    Processor.start("embeddings-vectorize", __doc__)
-        prog='embeddings-vectorizer',
-        description=__doc__,
-    )
-
-    parser.add_argument(
-        '-p', '--pulsar-host',
-        default=default_pulsar_host,
-        help=f'Pulsar host (default: {default_pulsar_host})',
-    )
-
-    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})'
-    )
-
-    parser.add_argument(
-        '-l', '--log-level',
-        type=LogLevel,
-        default=LogLevel.INFO,
-        choices=list(LogLevel),
-        help=f'Output queue (default: info)'
-    )
-
-    args = parser.parse_args()
-
-    while True:
-
-        try:
-
-            p = Processor(
-                pulsar_host=args.pulsar_host,
-                input_queue=args.input_queue,
-                output_queue=args.output_queue,
-                subscriber=args.subscriber,
-                log_level=args.log_level,
-            )
-
-            p.run()
-
-        except Exception as e:
-
-            print("Exception:", e, flush=True)
-            print("Will retry...", flush=True)
-
-        time.sleep(10)
 

trustgraph/rag/graph/rag.py

@@ -4,30 +4,22 @@ Simple RAG service, performs query using graph RAG an LLM.
 Input is query, output is response.
 """
 
-import pulsar
-from pulsar.schema import JsonSchema
-import tempfile
-import base64
-import os
-import argparse
-import time
-
 from ... schema import GraphRagQuery, GraphRagResponse
 from ... log_level import LogLevel
 from ... graph_rag import GraphRag
+from ... base import ConsumerProducer
 
-default_pulsar_host = os.getenv("PULSAR_HOST", 'pulsar://pulsar:6650')
 default_input_queue = 'graph-rag-query'
 default_output_queue = 'graph-rag-response'
 default_subscriber = 'graph-rag'
 default_graph_hosts = [ 'localhost' ]
 default_vector_store = 'http://localhost:19530'
 
-class Processor:
+class Processor(ConsumerProducer):
 
     def __init__(
             self,
-            pulsar_host=default_pulsar_host,
+            pulsar_host=None,
             input_queue=default_input_queue,
             output_queue=default_output_queue,
             subscriber=default_subscriber,
@@ -39,21 +31,14 @@ class Processor:
             max_sg_size=3000,
     ):
 
-        self.client = None
+        super(Processor, self).__init__(
-
+            pulsar_host=pulsar_host,
-        self.client = pulsar.Client(
+            log_level=log_level,
-            pulsar_host,
+            input_queue=input_queue,
-            logger=pulsar.ConsoleLogger(log_level.to_pulsar())
+            output_queue=output_queue,
-        )
+            subscriber=subscriber,
-
+            input_schema=TextCompletionRequest,
-        self.consumer = self.client.subscribe(
+            output_schema=TextCompletionResponse,
-            input_queue, subscriber,
-            schema=JsonSchema(GraphRagQuery),
-        )
-
-        self.producer = self.client.create_producer(
-            topic=output_queue,
-            schema=JsonSchema(GraphRagResponse),
         )
 
         self.rag = GraphRag(
@@ -66,13 +51,7 @@ class Processor:
             max_sg_size=max_sg_size,
         )
 
-    def run(self):
+    def handle(self, msg):
-
-        while True:
-
-            msg = self.consumer.receive()
-
-            try:
 
         v = msg.value()
 
@@ -90,58 +69,12 @@ class Processor:
 
         print("Done.", flush=True)
 
-                # Acknowledge successful processing of the message
+    @staticmethod
-                self.consumer.acknowledge(msg)
+    def add_args(parser):
 
-            except Exception as e:
+        ConsumerProducer.add_args(
-
+            parser, default_input_queue, default_subscriber,
-                print("Exception:", e, flush=True)
+            default_output_queue,
-
-                # Message failed to be processed
-                self.consumer.negative_acknowledge(msg)
-
-    def __del__(self):
-
-        if self.client:
-            self.client.close()
-
-def run():
-
-    parser = argparse.ArgumentParser(
-        prog='graph-rag',
-        description=__doc__,
-    )
-
-    parser.add_argument(
-        '-p', '--pulsar-host',
-        default=default_pulsar_host,
-        help=f'Pulsar host (default: {default_pulsar_host})',
-    )
-
-    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})'
-    )
-
-    parser.add_argument(
-        '-l', '--log-level',
-        type=LogLevel,
-        default=LogLevel.INFO,
-        choices=list(LogLevel),
-        help=f'Output queue (default: info)'
         )
 
         parser.add_argument(
@@ -177,31 +110,8 @@ def run():
             help=f'Max subgraph size (default: 3000)'
         )
 
-    args = parser.parse_args()
+def run():
 
-    while True:
+    Processor.start('graph-rag', __doc__)
 
-        try:
-
-            p = Processor(
-                pulsar_host=args.pulsar_host,
-                input_queue=args.input_queue,
-                output_queue=args.output_queue,
-                subscriber=args.subscriber,
-                log_level=args.log_level,
-                graph_hosts=args.graph_hosts.split(","),
-                vector_store=args.vector_store,
-                entity_limit=args.entity_limit,
-                triple_limit=args.triple_limit,
-                max_sg_size=args.max_subgraph_size,
-            )
-
-            p.run()
-
-        except Exception as e:
-
-            print("Exception:", e, flush=True)
-            print("Will retry...", flush=True)
-
-        time.sleep(10)