trustgraph-base 0.22.4__tar.gz → 0.22.6__tar.gz

{trustgraph-base-0.22.4 → trustgraph-base-0.22.6}/PKG-INFO

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

trustgraph-base-0.22.6/trustgraph/base/__init__.py

@@ -0,0 +1,31 @@
+
+from . pubsub import PulsarClient
+from . async_processor import AsyncProcessor
+from . consumer import Consumer
+from . producer import Producer
+from . publisher import Publisher
+from . subscriber import Subscriber
+from . metrics import ProcessorMetrics, ConsumerMetrics, ProducerMetrics
+from . flow_processor import FlowProcessor
+from . consumer_spec import ConsumerSpec
+from . setting_spec import SettingSpec
+from . producer_spec import ProducerSpec
+from . subscriber_spec import SubscriberSpec
+from . request_response_spec import RequestResponseSpec
+from . llm_service import LlmService, LlmResult
+from . embeddings_service import EmbeddingsService
+from . embeddings_client import EmbeddingsClientSpec
+from . text_completion_client import TextCompletionClientSpec
+from . prompt_client import PromptClientSpec
+from . triples_store_service import TriplesStoreService
+from . graph_embeddings_store_service import GraphEmbeddingsStoreService
+from . document_embeddings_store_service import DocumentEmbeddingsStoreService
+from . triples_query_service import TriplesQueryService
+from . graph_embeddings_query_service import GraphEmbeddingsQueryService
+from . document_embeddings_query_service import DocumentEmbeddingsQueryService
+from . graph_embeddings_client import GraphEmbeddingsClientSpec
+from . triples_client import TriplesClientSpec
+from . document_embeddings_client import DocumentEmbeddingsClientSpec
+from . agent_service import AgentService
+from . graph_rag_client import GraphRagClientSpec
+

trustgraph-base-0.22.6/trustgraph/base/agent_client.py

@@ -0,0 +1,39 @@
+
+from . request_response_spec import RequestResponse, RequestResponseSpec
+from .. schema import AgentRequest, AgentResponse
+from .. knowledge import Uri, Literal
+
+class AgentClient(RequestResponse):
+    async def request(self, recipient, question, plan=None, state=None,
+                    history=[], timeout=300):
+
+        resp = await self.request(
+            AgentRequest(
+                question = question,
+                plan = plan,
+                state = state,
+                history = history,
+            ),
+            recipient=recipient,
+            timeout=timeout,
+        )
+
+        print(resp, flush=True)
+
+        if resp.error:
+            raise RuntimeError(resp.error.message)
+
+        return resp
+
+class GraphEmbeddingsClientSpec(RequestResponseSpec):
+    def __init__(
+            self, request_name, response_name,
+    ):
+        super(GraphEmbeddingsClientSpec, self).__init__(
+            request_name = request_name,
+            request_schema = GraphEmbeddingsRequest,
+            response_name = response_name,
+            response_schema = GraphEmbeddingsResponse,
+            impl = GraphEmbeddingsClient,
+        )
+

trustgraph-base-0.22.6/trustgraph/base/agent_service.py

@@ -0,0 +1,100 @@
+
+"""
+Agent manager service completion base class
+"""
+
+import time
+from prometheus_client import Histogram
+
+from .. schema import AgentRequest, AgentResponse, Error
+from .. exceptions import TooManyRequests
+from .. base import FlowProcessor, ConsumerSpec, ProducerSpec
+
+default_ident = "agent-manager"
+
+class AgentService(FlowProcessor):
+
+    def __init__(self, **params):
+
+        id = params.get("id")
+
+        super(AgentService, self).__init__(**params | { "id": id })
+
+        self.register_specification(
+            ConsumerSpec(
+                name = "request",
+                schema = AgentRequest,
+                handler = self.on_request
+            )
+        )
+
+        self.register_specification(
+            ProducerSpec(
+                name = "next",
+                schema = AgentRequest
+            )
+        )
+
+        self.register_specification(
+            ProducerSpec(
+                name = "response",
+                schema = AgentResponse
+            )
+        )
+
+    async def on_request(self, msg, consumer, flow):
+
+        try:
+
+            request = msg.value()
+
+            # Sender-produced ID
+            id = msg.properties()["id"]
+
+            async def respond(resp):
+
+                await flow("response").send(
+                    resp,
+                    properties={"id": id}
+                )
+
+            async def next(resp):
+
+                await flow("next").send(
+                    resp,
+                    properties={"id": id}
+                )
+
+            await self.agent_request(
+                request = request, respond = respond, next = next,
+                flow = flow
+            )
+
+        except TooManyRequests as e:
+            raise e
+
+        except Exception as e:
+
+            # Apart from rate limits, treat all exceptions as unrecoverable
+            print(f"on_request Exception: {e}")
+
+            print("Send error response...", flush=True)
+
+            await flow.producer["response"].send(
+                AgentResponse(
+                    error=Error(
+                        type = "agent-error",
+                        message = str(e),
+                    ),
+                    thought = None,
+                    observation = None,
+                    answer = None,
+                ),
+                properties={"id": id}
+            )
+
+    @staticmethod
+    def add_args(parser):
+
+        FlowProcessor.add_args(parser)
+

trustgraph-base-0.22.6/trustgraph/base/async_processor.py

@@ -0,0 +1,254 @@
+
+# Base class for processors.  Implements:
+# - Pulsar client, subscribe and consume basic
+# - the async startup logic
+# - Initialising metrics
+
+import asyncio
+import argparse
+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
+from .. exceptions import TooManyRequests
+from . pubsub import PulsarClient
+from . producer import Producer
+from . consumer import Consumer
+from . metrics import ProcessorMetrics
+
+default_config_queue = config_push_queue
+
+# Async processor
+class AsyncProcessor:
+
+    def __init__(self, **params):
+
+        # Store the identity
+        self.id = params.get("id")
+
+        # Register a pulsar client
+        self.pulsar_client = PulsarClient(**params)
+
+        # Initialise metrics, records the parameters
+        ProcessorMetrics(id=self.id).info({
+            k: str(params[k])
+            for k in params
+            if k != "id"
+        })
+
+        # The processor runs all activity in a taskgroup, it's mandatory
+        # that this is provded
+        self.taskgroup = params.get("taskgroup")
+        if self.taskgroup is None:
+            raise RuntimeError("Essential taskgroup missing")
+
+        # Get the configuration topic
+        self.config_push_queue = params.get(
+            "config_push_queue", default_config_queue
+        )
+
+        # This records registered configuration handlers
+        self.config_handlers = []
+
+        # Create a random ID for this subscription to the configuration
+        # service
+        config_subscriber_id = str(uuid.uuid4())
+
+        # Subscribe to config queue
+        self.config_sub_task = Consumer(
+
+            taskgroup = self.taskgroup,
+            client = self.client,
+            subscriber = config_subscriber_id,
+            flow = None,
+
+            topic = self.config_push_queue,
+            schema = ConfigPush,
+
+            handler = self.on_config_change,
+
+            # This causes new subscriptions to view the entire history of
+            # configuration
+            start_of_messages = True
+        )
+
+        self.running = True
+
+    # This is called to start dynamic behaviour.  An over-ride point for
+    # extra functionality
+    async def start(self):
+        await self.config_sub_task.start()
+
+    # This is called to stop all threads.  An over-ride point for extra
+    # functionality
+    def stop(self):
+        self.client.close()
+        self.running = False
+
+    # Returns the pulsar host
+    @property
+    def pulsar_host(self): return self.client.pulsar_host
+
+    # Returns the pulsar client
+    @property
+    def client(self): return self.pulsar_client.client
+
+    # Register a new event handler for configuration change
+    def register_config_handler(self, handler):
+        self.config_handlers.append(handler)
+
+    # Called when a new configuration message push occurs
+    async def on_config_change(self, message, consumer):
+
+        # Get configuration data and version number
+        config = message.value().config
+        version = message.value().version
+
+        # Acknowledge the message
+        consumer.acknowledge(message)
+
+        # Invoke message handlers
+        print("Config change event", config, version, flush=True)
+        for ch in self.config_handlers:
+            await ch(config, version)
+
+    # This is the 'main' body of the handler.  It is a point to override
+    # if needed.  By default does nothing.  Processors are implemented
+    # by adding consumer/producer functionality so maybe nothing is needed
+    # in the run() body
+    async def run(self):
+        while self.running:
+            await asyncio.sleep(2)
+
+    # Startup fabric.  This runs in 'async' mode, creates a taskgroup and
+    # runs the producer.
+    @classmethod
+    async def launch_async(cls, args):
+
+        try:
+
+            # Create a taskgroup.  This seems complicated, when an exception
+            # occurs, unhandled it looks like it cancels all threads in the
+            # taskgroup.  Needs the exception to be caught in the right
+            # place.
+            async with asyncio.TaskGroup() as tg:
+
+
+                    # Create a processor instance, and include the taskgroup
+                    # as a paramter.  A processor identity ident is used as
+                    # - subscriber name
+                    # - an identifier for flow configuration
+                    p = cls(**args | { "taskgroup": tg })
+
+                    # Start the processor
+                    await p.start()
+
+                    # Run the processor
+                    task = tg.create_task(p.run())
+
+                    # The taskgroup causes everything to wait until
+                    # all threads have stopped
+
+        # This is here to output a debug message, shouldn't be needed.
+        except Exception as e:
+            print("Exception, closing taskgroup", flush=True)
+            raise e
+
+    # Startup fabric.  launch calls launch_async in async mode.
+    @classmethod
+    def launch(cls, ident, doc):
+
+        # Start assembling CLI arguments
+        parser = argparse.ArgumentParser(
+            prog=ident,
+            description=doc
+        )
+        
+        parser.add_argument(
+            '--id',
+            default=ident,
+            help=f'Configuration identity (default: {ident})',
+        )
+
+        # Invoke the class-specific add_args, which manages adding all the
+        # command-line arguments
+        cls.add_args(parser)
+
+        # Parse arguments
+        args = parser.parse_args()
+        args = vars(args)
+
+        # Debug
+        print(args, flush=True)
+
+        # Start the Prometheus metrics service if needed
+        if args["metrics"]:
+            start_http_server(args["metrics_port"])
+
+        # Loop forever, exception handler
+        while True:
+
+            print("Starting...", flush=True)
+
+            try:
+
+                # Launch the processor in an asyncio handler
+                asyncio.run(cls.launch_async(
+                    args
+                ))
+
+            except KeyboardInterrupt:
+                print("Keyboard interrupt.", flush=True)
+                return
+
+            except _pulsar.Interrupted:
+                print("Pulsar Interrupted.", flush=True)
+                return
+
+            # Exceptions from a taskgroup come in as an exception group
+            except ExceptionGroup as e:
+
+                print("Exception group:", flush=True)
+
+                for se in e.exceptions:
+                    print("  Type:", type(se), flush=True)
+                    print(f"  Exception: {se}", flush=True)
+
+            except Exception as e:
+                print("Type:", type(e), flush=True)
+                print("Exception:", e, flush=True)
+
+            # Retry occurs here
+            print("Will retry...", flush=True)
+            time.sleep(4)
+            print("Retrying...", flush=True)
+
+    # The command-line arguments are built using a stack of add_args
+    # invocations
+    @staticmethod
+    def add_args(parser):
+
+        PulsarClient.add_args(parser)
+
+        parser.add_argument(
+            '--config-push-queue',
+            default=default_config_queue,
+            help=f'Config push queue {default_config_queue}',
+        )
+
+        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)',
+        )

trustgraph-base-0.22.6/trustgraph/base/consumer.py

@@ -0,0 +1,197 @@
+
+from pulsar.schema import JsonSchema
+import pulsar
+import _pulsar
+import asyncio
+import time
+
+from .. exceptions import TooManyRequests
+
+class Consumer:
+
+    def __init__(
+            self, taskgroup, flow, client, topic, subscriber, schema,
+            handler, 
+            metrics = None,
+            start_of_messages=False,
+            rate_limit_retry_time = 10, rate_limit_timeout = 7200,
+            reconnect_time = 5,
+    ):
+
+        self.taskgroup = taskgroup
+        self.flow = flow
+        self.client = client
+        self.topic = topic
+        self.subscriber = subscriber
+        self.schema = schema
+        self.handler = handler
+
+        self.rate_limit_retry_time = rate_limit_retry_time
+        self.rate_limit_timeout = rate_limit_timeout
+
+        self.reconnect_time = 5
+
+        self.start_of_messages = start_of_messages
+
+        self.running = True
+        self.task = None
+
+        self.metrics = metrics
+
+        self.consumer = None
+
+    def __del__(self):
+        self.running = False
+
+        if hasattr(self, "consumer"):
+            if self.consumer:
+                self.consumer.close()
+
+    async def stop(self):
+
+        self.running = False
+        await self.task
+
+    async def start(self):
+
+        self.running = True
+
+        # Puts it in the stopped state, the run thread should set running
+        if self.metrics:
+            self.metrics.state("stopped")
+
+        self.task = self.taskgroup.create_task(self.run())
+
+    async def run(self):
+
+        while self.running:
+
+            if self.metrics:
+                self.metrics.state("stopped")
+
+            try:
+
+                print(self.topic, "subscribing...", flush=True)
+
+                if self.start_of_messages:
+                    pos = pulsar.InitialPosition.Earliest
+                else:
+                    pos = pulsar.InitialPosition.Latest
+
+                self.consumer = await asyncio.to_thread(
+                    self.client.subscribe,
+                    topic = self.topic,
+                    subscription_name = self.subscriber,
+                    schema = JsonSchema(self.schema),
+                    initial_position = pos,
+                    consumer_type = pulsar.ConsumerType.Shared,
+                )
+
+            except Exception as e:
+
+                print("consumer subs Exception:", e, flush=True)
+                await asyncio.sleep(self.reconnect_time)
+                continue
+
+            print(self.topic, "subscribed", flush=True)
+
+            if self.metrics:
+                self.metrics.state("running")
+
+            try:
+
+                await self.consume()
+
+                if self.metrics:
+                    self.metrics.state("stopped")
+
+            except Exception as e:
+
+                print("consumer loop exception:", e, flush=True)
+                self.consumer.close()
+                self.consumer = None
+                await asyncio.sleep(self.reconnect_time)
+                continue
+
+    async def consume(self):
+
+        while self.running:
+
+            try:
+                msg = await asyncio.to_thread(
+                    self.consumer.receive,
+                    timeout_millis=2000
+                )
+            except _pulsar.Timeout:
+                continue
+            except Exception as e:
+                raise e
+
+            expiry = time.time() + self.rate_limit_timeout
+
+            # This loop is for retry on rate-limit / resource limits
+            while self.running:
+
+                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)
+
+                    if self.metrics:
+                        self.metrics.process("error")
+
+                    # Break out of retry loop, processes next message
+                    break
+
+                try:
+
+                    print("Handle...", flush=True)
+
+                    if self.metrics:
+
+                        with self.metrics.record_time():
+                            await self.handler(msg, self, self.flow)
+
+                    else:
+                        await self.handler(msg, self.consumer)
+
+                    print("Handled.", flush=True)
+
+                    # Acknowledge successful processing of the message
+                    self.consumer.acknowledge(msg)
+
+                    if self.metrics:
+                        self.metrics.process("success")
+
+                    # Break out of retry loop
+                    break
+
+                except TooManyRequests:
+
+                    print("TooManyRequests: will retry...", flush=True)
+
+                    if self.metrics:
+                        self.metrics.rate_limit()
+
+                    # Sleep
+                    await asyncio.sleep(self.rate_limit_retry_time)
+
+                    # Contine from retry loop, just causes a reprocessing
+                    continue
+
+                except Exception as e:
+
+                    print("consume exception:", e, flush=True)
+
+                    # Message failed to be processed, this causes it to
+                    # be retried
+                    self.consumer.negative_acknowledge(msg)
+
+                    if self.metrics:
+                        self.metrics.process("error")
+
+                    # Break out of retry loop, processes next message
+                    break

trustgraph-base-0.22.6/trustgraph/base/consumer_spec.py

@@ -0,0 +1,36 @@
+
+from . metrics import ConsumerMetrics
+from . consumer import Consumer
+from . spec import Spec
+
+class ConsumerSpec(Spec):
+    def __init__(self, name, schema, handler):
+        self.name = name
+        self.schema = schema
+        self.handler = handler
+
+    def add(self, flow, processor, definition):
+
+        consumer_metrics = ConsumerMetrics(
+            flow.id, f"{flow.name}-{self.name}"
+        )
+
+        consumer = Consumer(
+            taskgroup = processor.taskgroup,
+            flow = flow,
+            client = processor.client,
+            topic = definition[self.name],
+            subscriber = processor.id + "--" + self.name,
+            schema = self.schema,
+            handler = self.handler,
+            metrics = consumer_metrics,
+        )
+
+        # Consumer handle gets access to producers and other
+        # metadata
+        consumer.id = flow.id
+        consumer.name = self.name
+        consumer.flow = flow
+
+        flow.consumer[self.name] = consumer
+