trustgraph-base 0.22.9__tar.gz → 0.22.11__tar.gz

{trustgraph-base-0.22.9 → trustgraph-base-0.22.11}/PKG-INFO

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

{trustgraph-base-0.22.9 → trustgraph-base-0.22.11}/trustgraph/api/api.py

@@ -562,233 +562,3 @@ class Api:
         except:
             raise ProtocolException(f"Response not formatted correctly")
 
-    def flow_list_classes(self):
-
-        # The input consists of system and prompt strings
-        input = {
-            "operation": "list-classes",
-        }
-
-        url = f"{self.url}flow"
-
-        # Invoke the API, input is passed as JSON
-        resp = requests.post(url, json=input)
-
-        # Should be a 200 status code
-        if resp.status_code != 200:
-            raise ProtocolException(f"Status code {resp.status_code}")
-
-        try:
-            # Parse the response as JSON
-            object = resp.json()
-        except:
-            raise ProtocolException(f"Expected JSON response")
-
-        self.check_error(object)
-
-        try:
-            return object["class-names"]
-        except:
-            raise ProtocolException(f"Response not formatted correctly")
-
-    def flow_get_class(self, class_name):
-
-        # The input consists of system and prompt strings
-        input = {
-            "operation": "get-class",
-            "class-name": class_name,
-        }
-
-        url = f"{self.url}flow"
-
-        # Invoke the API, input is passed as JSON
-        resp = requests.post(url, json=input)
-
-        # Should be a 200 status code
-        if resp.status_code != 200:
-            raise ProtocolException(f"Status code {resp.status_code}")
-
-        try:
-            # Parse the response as JSON
-            object = resp.json()
-        except:
-            raise ProtocolException(f"Expected JSON response")
-
-        self.check_error(object)
-
-        try:
-            return json.loads(object["class-definition"])
-        except Exception as e:
-            print(e)
-            raise ProtocolException(f"Response not formatted correctly")
-
-    def flow_put_class(self, class_name, definition):
-
-        # The input consists of system and prompt strings
-        input = {
-            "operation": "put-class",
-            "class-name": class_name,
-            "class-definition": json.dumps(definition),
-        }
-
-        url = f"{self.url}flow"
-
-        # Invoke the API, input is passed as JSON
-        resp = requests.post(url, json=input)
-
-        # Should be a 200 status code
-        if resp.status_code != 200:
-            raise ProtocolException(f"Status code {resp.status_code}")
-
-        try:
-            # Parse the response as JSON
-            object = resp.json()
-        except:
-            raise ProtocolException(f"Expected JSON response")
-
-        self.check_error(object)
-
-        return
-
-    def flow_delete_class(self, class_name):
-
-        # The input consists of system and prompt strings
-        input = {
-            "operation": "delete-class",
-            "class-name": class_name,
-        }
-
-        url = f"{self.url}flow"
-
-        # Invoke the API, input is passed as JSON
-        resp = requests.post(url, json=input)
-
-        # Should be a 200 status code
-        if resp.status_code != 200:
-            raise ProtocolException(f"Status code {resp.status_code}")
-
-        try:
-            # Parse the response as JSON
-            object = resp.json()
-        except:
-            raise ProtocolException(f"Expected JSON response")
-
-        self.check_error(object)
-
-        return
-
-    def flow_list(self):
-
-        # The input consists of system and prompt strings
-        input = {
-            "operation": "list-flows",
-        }
-
-        url = f"{self.url}flow"
-
-        # Invoke the API, input is passed as JSON
-        resp = requests.post(url, json=input)
-
-        # Should be a 200 status code
-        if resp.status_code != 200:
-            raise ProtocolException(f"Status code {resp.status_code}")
-
-        try:
-            # Parse the response as JSON
-            object = resp.json()
-        except:
-            raise ProtocolException(f"Expected JSON response")
-
-        self.check_error(object)
-
-        try:
-            return object["flow-ids"]
-        except:
-            raise ProtocolException(f"Response not formatted correctly")
-
-    def flow_get(self, id):
-
-        # The input consists of system and prompt strings
-        input = {
-            "operation": "get-flow",
-            "flow-id": id,
-        }
-
-        url = f"{self.url}flow"
-
-        # Invoke the API, input is passed as JSON
-        resp = requests.post(url, json=input)
-
-        # Should be a 200 status code
-        if resp.status_code != 200:
-            raise ProtocolException(f"Status code {resp.status_code}")
-
-        try:
-            # Parse the response as JSON
-            object = resp.json()
-        except:
-            raise ProtocolException(f"Expected JSON response")
-
-        self.check_error(object)
-
-        try:
-            return json.loads(object["flow"])
-        except:
-            raise ProtocolException(f"Response not formatted correctly")
-
-    def flow_start(self, class_name, id, description):
-
-        # The input consists of system and prompt strings
-        input = {
-            "operation": "start-flow",
-            "flow-id": id,
-            "class-name": class_name,
-            "description": description,
-        }
-
-        url = f"{self.url}flow"
-
-        # Invoke the API, input is passed as JSON
-        resp = requests.post(url, json=input)
-
-        # Should be a 200 status code
-        if resp.status_code != 200:
-            raise ProtocolException(f"Status code {resp.status_code}")
-
-        try:
-            # Parse the response as JSON
-            object = resp.json()
-        except:
-            raise ProtocolException(f"Expected JSON response")
-
-        self.check_error(object)
-
-        return
-
-    def flow_stop(self, id):
-
-        # The input consists of system and prompt strings
-        input = {
-            "operation": "stop-flow",
-            "flow-id": id,
-        }
-
-        url = f"{self.url}flow"
-
-        # Invoke the API, input is passed as JSON
-        resp = requests.post(url, json=input)
-
-        # Should be a 200 status code
-        if resp.status_code != 200:
-            raise ProtocolException(f"Status code {resp.status_code}")
-
-        try:
-            # Parse the response as JSON
-            object = resp.json()
-        except:
-            raise ProtocolException(f"Expected JSON response")
-
-        self.check_error(object)
-
-        return
-

trustgraph-base-0.22.11/trustgraph/base/__init__.py

@@ -0,0 +1,8 @@
+
+from . base_processor import BaseProcessor
+from . consumer import Consumer
+from . producer import Producer
+from . consumer_producer import ConsumerProducer
+from . publisher import Publisher
+from . subscriber import Subscriber
+

trustgraph-base-0.22.11/trustgraph/base/base_processor.py

@@ -0,0 +1,210 @@
+
+import asyncio
+import os
+import argparse
+import pulsar
+from pulsar.schema import JsonSchema
+import _pulsar
+import time
+import uuid
+from prometheus_client import start_http_server, Info
+
+from .. schema import ConfigPush, config_push_queue
+from .. log_level import LogLevel
+
+default_config_queue = config_push_queue
+config_subscriber_id = str(uuid.uuid4())
+
+class BaseProcessor:
+
+    default_pulsar_host = os.getenv("PULSAR_HOST", 'pulsar://pulsar:6650')
+    default_pulsar_api_key = os.getenv("PULSAR_API_KEY", None)
+
+    def __init__(self, **params):
+
+        self.client = None
+
+        if not hasattr(__class__, "params_metric"):
+            __class__.params_metric = Info(
+                'params', 'Parameters configuration'
+            )
+
+        # FIXME: Maybe outputs information it should not
+        __class__.params_metric.info({
+            k: str(params[k])
+            for k in params
+        })
+
+        pulsar_host = params.get("pulsar_host", self.default_pulsar_host)
+        pulsar_listener = params.get("pulsar_listener", None)
+        pulsar_api_key = params.get("pulsar_api_key", None)
+        log_level = params.get("log_level", LogLevel.INFO)
+
+        self.config_push_queue = params.get(
+            "config_push_queue",
+            default_config_queue
+        )
+
+        self.pulsar_host = pulsar_host
+        self.pulsar_api_key = pulsar_api_key
+
+        if pulsar_api_key:
+            auth = pulsar.AuthenticationToken(pulsar_api_key)
+            self.client = pulsar.Client(
+            pulsar_host,
+            authentication=auth,
+            logger=pulsar.ConsoleLogger(log_level.to_pulsar())
+            )
+        else:
+            self.client = pulsar.Client(
+            pulsar_host,
+            listener_name=pulsar_listener,
+            logger=pulsar.ConsoleLogger(log_level.to_pulsar())
+            )
+
+        self.pulsar_listener = pulsar_listener
+
+        self.config_subscriber = self.client.subscribe(
+            self.config_push_queue, config_subscriber_id,
+            consumer_type=pulsar.ConsumerType.Shared,
+            initial_position=pulsar.InitialPosition.Earliest,
+            schema=JsonSchema(ConfigPush),         
+        )
+
+    def __del__(self):
+
+        if hasattr(self, "client"):
+            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(
+            '--pulsar-api-key',
+            default=__class__.default_pulsar_api_key,
+            help=f'Pulsar API key',
+        )
+
+        parser.add_argument(
+            '--config-push-queue',
+            default=default_config_queue,
+            help=f'Config push queue {default_config_queue}',
+        )
+
+        parser.add_argument(
+            '--pulsar-listener',
+            help=f'Pulsar listener (default: none)',
+        )
+
+        parser.add_argument(
+            '-l', '--log-level',
+            type=LogLevel,
+            default=LogLevel.INFO,
+            choices=list(LogLevel),
+            help=f'Output queue (default: info)'
+        )
+
+        parser.add_argument(
+            '--metrics',
+            action=argparse.BooleanOptionalAction,
+            default=True,
+            help=f'Metrics enabled (default: true)',
+        )
+
+        parser.add_argument(
+            '-P', '--metrics-port',
+            type=int,
+            default=8000,
+            help=f'Pulsar host (default: 8000)',
+        )
+
+    async def start(self):
+        pass
+
+    async def run_config_queue(self):
+
+        if self.module == "config.service":
+            print("I am config-svc, not looking at config queue", flush=True)
+            return
+
+        print("Config thread running", flush=True)
+        
+        while True:
+
+            try:
+                msg = await asyncio.to_thread(
+                    self.config_subscriber.receive, timeout_millis=2000
+                )
+            except pulsar.Timeout:
+                continue
+
+            v = msg.value()
+            print("Got config version", v.version, flush=True)
+
+            await self.on_config(v.version, v.config)
+
+    async def on_config(self, version, config):
+        pass
+
+    async def run(self):
+        raise RuntimeError("Something should have implemented the run method")
+
+    @classmethod
+    async def launch_async(cls, args, prog):
+        p = cls(**args)
+        p.module = prog
+        await p.start()
+
+        task1 = asyncio.create_task(p.run_config_queue())
+        task2 = asyncio.create_task(p.run())
+
+        await asyncio.gather(task1, task2)
+
+    @classmethod
+    def launch(cls, prog, doc):
+
+        parser = argparse.ArgumentParser(
+            prog=prog,
+            description=doc
+        )
+
+        cls.add_args(parser)
+
+        args = parser.parse_args()
+        args = vars(args)
+
+        print(args)
+
+        if args["metrics"]:
+            start_http_server(args["metrics_port"])
+
+        while True:
+
+            try:
+
+                asyncio.run(cls.launch_async(args, prog))
+
+            except KeyboardInterrupt:
+                print("Keyboard interrupt.")
+                return
+
+            except _pulsar.Interrupted:
+                print("Pulsar Interrupted.")
+                return
+
+            except Exception as e:
+
+                print(type(e))
+
+                print("Exception:", e, flush=True)
+                print("Will retry...", flush=True)
+
+                time.sleep(4)
+

trustgraph-base-0.22.11/trustgraph/base/consumer.py

@@ -0,0 +1,173 @@
+
+import asyncio
+from pulsar.schema import JsonSchema
+import pulsar
+from prometheus_client import Histogram, Info, Counter, Enum
+import time
+
+from . base_processor import BaseProcessor
+from .. exceptions import TooManyRequests
+
+default_rate_limit_retry = 10
+default_rate_limit_timeout = 7200
+
+class Consumer(BaseProcessor):
+
+    def __init__(self, **params):
+
+        if not hasattr(__class__, "state_metric"):
+            __class__.state_metric = Enum(
+                'processor_state', 'Processor state',
+                states=['starting', 'running', 'stopped']
+            )
+            __class__.state_metric.state('starting')
+
+        __class__.state_metric.state('starting')
+
+        super(Consumer, self).__init__(**params)
+
+        self.input_queue = params.get("input_queue")
+        self.subscriber = params.get("subscriber")
+        self.input_schema = params.get("input_schema")
+
+        self.rate_limit_retry = params.get(
+            "rate_limit_retry", default_rate_limit_retry
+        )
+        self.rate_limit_timeout = params.get(
+            "rate_limit_timeout", default_rate_limit_timeout
+        )
+
+        if self.input_schema == None:
+            raise RuntimeError("input_schema must be specified")
+
+        if not hasattr(__class__, "request_metric"):
+            __class__.request_metric = Histogram(
+                'request_latency', 'Request latency (seconds)'
+            )
+
+        if not hasattr(__class__, "pubsub_metric"):
+            __class__.pubsub_metric = Info(
+                'pubsub', 'Pub/sub configuration'
+            )
+
+        if not hasattr(__class__, "processing_metric"):
+            __class__.processing_metric = Counter(
+                'processing_count', 'Processing count', ["status"]
+            )
+
+        if not hasattr(__class__, "rate_limit_metric"):
+            __class__.rate_limit_metric = Counter(
+                'rate_limit_count', 'Rate limit event count',
+            )
+
+        __class__.pubsub_metric.info({
+            "input_queue": self.input_queue,
+            "subscriber": self.subscriber,
+            "input_schema": self.input_schema.__name__,
+            "rate_limit_retry": str(self.rate_limit_retry),
+            "rate_limit_timeout": str(self.rate_limit_timeout),
+        })
+
+        self.consumer = self.client.subscribe(
+            self.input_queue, self.subscriber,
+            consumer_type=pulsar.ConsumerType.Shared,
+            schema=JsonSchema(self.input_schema),
+        )
+
+        print("Initialised consumer.", flush=True)
+
+    async def run(self):
+
+        __class__.state_metric.state('running')
+
+        while True:
+
+            msg = await asyncio.to_thread(self.consumer.receive)
+
+            expiry = time.time() + self.rate_limit_timeout
+
+            # This loop is for retry on rate-limit / resource limits
+            while True:
+
+                if time.time() > expiry:
+
+                    print("Gave up waiting for rate-limit retry", flush=True)
+
+                    # Message failed to be processed, this causes it to
+                    # be retried
+                    self.consumer.negative_acknowledge(msg)
+
+                    __class__.processing_metric.labels(status="error").inc()
+
+                    # Break out of retry loop, processes next message
+                    break
+
+                try:
+
+                    with __class__.request_metric.time():
+                        await self.handle(msg)
+
+                    # Acknowledge successful processing of the message
+                    self.consumer.acknowledge(msg)
+
+                    __class__.processing_metric.labels(status="success").inc()
+
+                    # Break out of retry loop
+                    break
+
+                except TooManyRequests:
+
+                    print("TooManyRequests: will retry...", flush=True)
+
+                    __class__.rate_limit_metric.inc()
+
+                    # Sleep
+                    time.sleep(self.rate_limit_retry)
+
+                    # Contine from retry loop, just causes a reprocessing
+                    continue
+                
+                except Exception as e:
+
+                    print("Exception:", e, flush=True)
+
+                    # Message failed to be processed, this causes it to
+                    # be retried
+                    self.consumer.negative_acknowledge(msg)
+
+                    __class__.processing_metric.labels(status="error").inc()
+
+                    # Break out of retry loop, processes next message
+                    break
+
+    @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})'
+        )
+
+        parser.add_argument(
+            '--rate-limit-retry',
+            type=int,
+            default=default_rate_limit_retry,
+            help=f'Rate limit retry (default: {default_rate_limit_retry})'
+        )
+
+        parser.add_argument(
+            '--rate-limit-timeout',
+            type=int,
+            default=default_rate_limit_timeout,
+            help=f'Rate limit timeout (default: {default_rate_limit_timeout})'
+        )
+

trustgraph-base-0.22.11/trustgraph/base/consumer_producer.py

@@ -0,0 +1,62 @@
+
+from pulsar.schema import JsonSchema
+import pulsar
+from prometheus_client import Histogram, Info, Counter, Enum
+import time
+
+from . consumer import Consumer
+from .. exceptions import TooManyRequests
+
+class ConsumerProducer(Consumer):
+
+    def __init__(self, **params):
+
+        super(ConsumerProducer, self).__init__(**params)
+
+        self.output_queue = params.get("output_queue")
+        self.output_schema = params.get("output_schema")
+
+        if not hasattr(__class__, "output_metric"):
+            __class__.output_metric = Counter(
+                'output_count', 'Output items created'
+            )
+
+        __class__.pubsub_metric.info({
+            "input_queue": self.input_queue,
+            "output_queue": self.output_queue,
+            "subscriber": self.subscriber,
+            "input_schema": self.input_schema.__name__,
+            "output_schema": self.output_schema.__name__,
+            "rate_limit_retry": str(self.rate_limit_retry),
+            "rate_limit_timeout": str(self.rate_limit_timeout),
+        })
+
+        if self.output_schema == None:
+            raise RuntimeError("output_schema must be specified")
+
+        self.producer = self.client.create_producer(
+            topic=self.output_queue,
+            schema=JsonSchema(self.output_schema),
+            chunking_enabled=True,
+        )
+
+        print("Initialised consumer/producer.")
+
+    async def send(self, msg, properties={}):
+        self.producer.send(msg, properties)
+        __class__.output_metric.inc()
+
+    @staticmethod
+    def add_args(
+            parser, default_input_queue, default_subscriber,
+            default_output_queue,
+    ):
+
+        Consumer.add_args(parser, default_input_queue, default_subscriber)
+
+        parser.add_argument(
+            '-o', '--output-queue',
+            default=default_output_queue,
+            help=f'Output queue (default: {default_output_queue})'
+        )
+