trustgraph-base 2.8.1__tar.gz → 2.8.3__tar.gz

{trustgraph_base-2.8.1 → trustgraph_base-2.8.3}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: trustgraph-base
-Version: 2.8.1
+Version: 2.8.3
 Summary: TrustGraph provides a means to run a pipeline of flexible AI processing components in a flexible means to achieve a processing pipeline.
 Author-email: "trustgraph.ai" <security@trustgraph.ai>
 Project-URL: Homepage, https://github.com/trustgraph-ai/trustgraph

{trustgraph_base-2.8.1 → trustgraph_base-2.8.3}/trustgraph/base/__init__.py

@@ -1,10 +1,6 @@
 
-from . pubsub import get_pubsub, get_async_pubsub, add_pubsub_args
+from . pubsub import get_async_pubsub, add_pubsub_args
 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, SubscriberMetrics
 from . logging import add_logging_args, setup_logging
 from . workspace_processor import WorkspaceProcessor
@@ -12,10 +8,8 @@ from . flow_processor import FlowProcessor
 from . consumer_spec import ConsumerSpec
 from . parameter_spec import ParameterSpec
 from . producer_spec import ProducerSpec
-from . subscriber_spec import SubscriberSpec
 from . request_response_spec import RequestResponseSpec
 from . llm_service import LlmService, LlmResult, LlmChunk
-from . librarian_client import LibrarianClient
 from . librarian_spec import LibrarianSpec
 from . chunking_service import ChunkingService
 from . embeddings_service import EmbeddingsService

{trustgraph_base-2.8.1 → trustgraph_base-2.8.3}/trustgraph/base/agent_client.py

@@ -1,9 +1,9 @@
 
-from . request_response_spec import RequestResponse, RequestResponseSpec
+from . request_response_spec import RequestResponseSpec
 from .. schema import AgentRequest, AgentResponse
 from .. knowledge import Uri, Literal
 
-class AgentClient(RequestResponse):
+class AgentClient:
     async def invoke(self, question, plan=None, state=None,
                     history=[], think=None, observe=None, answer_callback=None,
                     timeout=300):

{trustgraph_base-2.8.1 → trustgraph_base-2.8.3}/trustgraph/base/async_processor.py

@@ -15,7 +15,7 @@ from .. schema import ConfigPush, ConfigRequest, ConfigResponse
 from .. schema import config_push_queue, config_request_queue
 from .. schema import config_response_queue
 from .. log_level import LogLevel
-from . pubsub import get_async_pubsub, get_pubsub, add_pubsub_args
+from . pubsub import get_async_pubsub, add_pubsub_args
 from . receiver_pool import ReceiverPool
 from . sender_pool import SenderPool
 from . request_response_client import RequestResponseClient
@@ -51,10 +51,12 @@ class AsyncProcessor:
         self.receiver_pool = ReceiverPool(
             backend=self.async_backend,
             concurrency=self.concurrency,
+            processor_id=self.id,
         )
 
         self.sender_pool = SenderPool(
             backend=self.async_backend,
+            processor_id=self.id,
         )
 
         self.config_push_queue = params.get(

{trustgraph_base-2.8.1 → trustgraph_base-2.8.3}/trustgraph/base/audit_publisher.py

@@ -4,8 +4,6 @@ import logging
 from datetime import datetime, timezone
 from uuid import uuid4
 
-from . producer import Producer
-from . metrics import ProducerMetrics
 from trustgraph.schema import AuditEvent, audit_events_queue
 
 logger = logging.getLogger(__name__)
@@ -13,26 +11,12 @@ logger = logging.getLogger(__name__)
 
 class AuditPublisher:
 
-    def __init__(self, component_name, processor_id=None, backend=None,
-                 async_backend=None):
+    def __init__(self, component_name, async_backend=None):
         self.component_name = component_name
         self._handle = None
         self._async_producer = None
         self._async_backend = async_backend
 
-        if backend is not None:
-            self._legacy_producer = Producer(
-                backend=backend,
-                topic=audit_events_queue,
-                schema=AuditEvent,
-                metrics=ProducerMetrics(
-                    processor=processor_id or component_name,
-                    producer="audit-events",
-                ),
-            )
-        else:
-            self._legacy_producer = None
-
     async def start(self, sender_pool=None):
         if sender_pool is not None:
             self._handle = await sender_pool.add_producer(
@@ -70,7 +54,5 @@ class AuditPublisher:
                 await self._handle.send(event)
             elif self._async_producer:
                 await self._async_producer.send(event)
-            elif self._legacy_producer:
-                await self._legacy_producer.send(event)
         except Exception as e:
             logger.warning(f"Failed to emit audit event: {e}")

{trustgraph_base-2.8.1 → trustgraph_base-2.8.3}/trustgraph/base/config_client.py

@@ -1,14 +1,11 @@
 
-from . request_response_spec import RequestResponse, RequestResponseSpec
+from . request_response_spec import RequestResponseSpec
 from .. schema import ConfigRequest, ConfigResponse, ConfigKey, ConfigValue
 
 CONFIG_TIMEOUT = 10
 
 
-class ConfigClient(RequestResponse):
-
-    def __init__(self, timeout=CONFIG_TIMEOUT, **kwargs):
-        super(ConfigClient, self).__init__(default_timeout=timeout, **kwargs)
+class ConfigClient:
 
     async def _request(self, timeout=None, **kwargs):
         resp = await self.request(

trustgraph_base-2.8.3/trustgraph/base/consumer_spec.py

@@ -0,0 +1,36 @@
+from __future__ import annotations
+
+from typing import Any
+
+from . spec import Spec
+
+class ConsumerSpec(Spec):
+    def __init__(self, name: str, schema: Any, handler: Any, concurrency: int = 1) -> None:
+        self.name = name
+        self.schema = schema
+        self.handler = handler
+        self.concurrency = concurrency
+
+    async def register(self, flow: Any, processor: Any, definition: dict[str, Any]) -> Any:
+
+        topic = definition["topics"][self.name]
+        subscription = (
+            processor.id + "--" + flow.workspace + "--" +
+            flow.name + "--" + self.name
+        )
+
+        handler = self.handler
+
+        async def pool_handler(message):
+            await handler(message, None, flow)
+
+        reg = await processor.receiver_pool.add_consumer(
+            topic=topic,
+            subscription=subscription,
+            schema=self.schema,
+            handler=pool_handler,
+        )
+
+        flow.consumer[self.name] = reg
+        return reg
+

{trustgraph_base-2.8.1 → trustgraph_base-2.8.3}/trustgraph/base/document_embeddings_client.py

@@ -1,14 +1,14 @@
 
 import logging
 
-from . request_response_spec import RequestResponse, RequestResponseSpec
+from . request_response_spec import RequestResponseSpec
 from .. schema import DocumentEmbeddingsRequest, DocumentEmbeddingsResponse
 from .. knowledge import Uri, Literal
 
 # Module logger
 logger = logging.getLogger(__name__)
 
-class DocumentEmbeddingsClient(RequestResponse):
+class DocumentEmbeddingsClient:
     async def query(self, vector, limit=20, collection="default", timeout=30):
 
         resp = await self.request(

{trustgraph_base-2.8.1 → trustgraph_base-2.8.3}/trustgraph/base/embeddings_client.py

@@ -1,8 +1,8 @@
 
-from . request_response_spec import RequestResponse, RequestResponseSpec
+from . request_response_spec import RequestResponseSpec
 from .. schema import EmbeddingsRequest, EmbeddingsResponse
 
-class EmbeddingsClient(RequestResponse):
+class EmbeddingsClient:
     async def embed(self, texts, timeout=300):
 
         resp = await self.request(

{trustgraph_base-2.8.1 → trustgraph_base-2.8.3}/trustgraph/base/graph_embeddings_client.py

@@ -4,7 +4,7 @@ from typing import Any
 
 import logging
 
-from . request_response_spec import RequestResponse, RequestResponseSpec
+from . request_response_spec import RequestResponseSpec
 from .. schema import GraphEmbeddingsRequest, GraphEmbeddingsResponse, IRI, LITERAL
 from .. knowledge import Uri, Literal
 
@@ -21,7 +21,7 @@ def to_value(x: Any) -> Any:
     # Fallback
     return Literal(x.value or x.iri)
 
-class GraphEmbeddingsClient(RequestResponse):
+class GraphEmbeddingsClient:
     async def query(self, vector, limit=20, collection="default", timeout=30):
 
         resp = await self.request(

{trustgraph_base-2.8.1 → trustgraph_base-2.8.3}/trustgraph/base/graph_rag_client.py

@@ -1,8 +1,8 @@
 
-from . request_response_spec import RequestResponse, RequestResponseSpec
+from . request_response_spec import RequestResponseSpec
 from .. schema import GraphRagQuery, GraphRagResponse
 
-class GraphRagClient(RequestResponse):
+class GraphRagClient:
     async def rag(self, query, collection="default",
                   chunk_callback=None, explain_callback=None,
                   parent_uri="",

{trustgraph_base-2.8.1 → trustgraph_base-2.8.3}/trustgraph/base/iam_client.py

@@ -1,7 +1,7 @@
 
 import json
 
-from . request_response_spec import RequestResponse, RequestResponseSpec
+from . request_response_spec import RequestResponseSpec
 from .. schema import (
     IamRequest, IamResponse,
     UserInput, WorkspaceInput, ApiKeyInput,
@@ -10,7 +10,7 @@ from .. schema import (
 IAM_TIMEOUT = 10
 
 
-class IamClient(RequestResponse):
+class IamClient:
     """Client for the IAM service request/response pub/sub protocol.
 
     Mirrors ``ConfigClient``: a thin wrapper around ``RequestResponse``

{trustgraph_base-2.8.1 → trustgraph_base-2.8.3}/trustgraph/base/keyword_index_client.py

@@ -1,13 +1,13 @@
 
 import logging
 
-from . request_response_spec import RequestResponse, RequestResponseSpec
+from . request_response_spec import RequestResponseSpec
 from .. schema import KeywordIndexRequest, KeywordIndexResponse
 
 # Module logger
 logger = logging.getLogger(__name__)
 
-class KeywordIndexClient(RequestResponse):
+class KeywordIndexClient:
     async def query(self, query, limit=20, collection="default", timeout=30):
 
         resp = await self.request(

{trustgraph_base-2.8.1 → trustgraph_base-2.8.3}/trustgraph/base/librarian_spec.py

@@ -4,7 +4,6 @@ import uuid
 from typing import Any
 
 from . spec import Spec
-from . librarian_client import LibrarianClient
 
 
 class LibrarianSpec(Spec):
@@ -13,24 +12,6 @@ class LibrarianSpec(Spec):
         self.request_name = request_name
         self.response_name = response_name
 
-    def add(self, flow: Any, processor: Any, definition: dict[str, Any]) -> None:
-
-        client = LibrarianClient(
-            id=flow.id,
-            backend=processor.pubsub,
-            taskgroup=processor.taskgroup,
-            librarian_request_queue=definition["topics"][self.request_name],
-            librarian_response_queue=definition["topics"][self.response_name],
-            librarian_subscriber=(
-                processor.id + "--" + flow.workspace + "--" +
-                flow.name + "--librarian--" + str(uuid.uuid4())
-            ),
-            workspace=flow.workspace,
-            flow_name=flow.name,
-        )
-
-        flow.librarian = client
-
     async def register(self, flow: Any, processor: Any, definition: dict[str, Any]) -> Any:
 
         from .async_librarian_client import AsyncLibrarianClient

{trustgraph_base-2.8.1 → trustgraph_base-2.8.3}/trustgraph/base/metrics.py

@@ -13,13 +13,11 @@ class ConsumerMetrics:
 
         self.processor = processor
         self.consumer = consumer
-        self.workspace = workspace
-        self.flow = flow
 
         if not hasattr(__class__, "state_metric"):
             __class__.state_metric = Enum(
                 'consumer_state', 'Consumer state',
-                ["processor", "workspace", "flow", "consumer"],
+                ["processor", "consumer"],
                 states=['stopped', 'running']
             )
 
@@ -32,31 +30,46 @@ class ConsumerMetrics:
         if not hasattr(__class__, "processing_metric"):
             __class__.processing_metric = Counter(
                 'processing_count', 'Processing count',
-                ["processor", "workspace", "flow", "consumer", "status"],
+                ["processor", "consumer", "status"],
             )
 
         if not hasattr(__class__, "rate_limit_metric"):
             __class__.rate_limit_metric = Counter(
                 'rate_limit_count', 'Rate limit event count',
-                ["processor", "workspace", "flow", "consumer"],
+                ["processor", "consumer"],
             )
 
+        __class__.request_metric.labels(
+            processor=self.processor, consumer=self.consumer,
+        )
+        __class__.processing_metric.labels(
+            processor=self.processor, consumer=self.consumer,
+            status="ok",
+        )
+        __class__.processing_metric.labels(
+            processor=self.processor, consumer=self.consumer,
+            status="error",
+        )
+        __class__.rate_limit_metric.labels(
+            processor=self.processor, consumer=self.consumer,
+        )
+
     def process(self, status: str) -> None:
         __class__.processing_metric.labels(
-            processor=self.processor, workspace=self.workspace,
-            flow=self.flow, consumer=self.consumer, status=status,
+            processor=self.processor,
+            consumer=self.consumer, status=status,
         ).inc()
 
     def rate_limit(self) -> None:
         __class__.rate_limit_metric.labels(
-            processor=self.processor, workspace=self.workspace,
-            flow=self.flow, consumer=self.consumer,
+            processor=self.processor,
+            consumer=self.consumer,
         ).inc()
 
     def state(self, state: str) -> None:
         __class__.state_metric.labels(
-            processor=self.processor, workspace=self.workspace,
-            flow=self.flow, consumer=self.consumer,
+            processor=self.processor,
+            consumer=self.consumer,
         ).state(state)
 
     def record_time(self) -> Any:
@@ -72,19 +85,21 @@ class ProducerMetrics:
 
         self.processor = processor
         self.producer = producer
-        self.workspace = workspace
-        self.flow = flow
 
         if not hasattr(__class__, "producer_metric"):
             __class__.producer_metric = Counter(
                 'producer_count', 'Output items produced',
-                ["processor", "workspace", "flow", "producer"],
+                ["processor", "producer"],
             )
 
+        __class__.producer_metric.labels(
+            processor=self.processor, producer=self.producer,
+        )
+
     def inc(self) -> None:
         __class__.producer_metric.labels(
-            processor=self.processor, workspace=self.workspace,
-            flow=self.flow, producer=self.producer,
+            processor=self.processor,
+            producer=self.producer,
         ).inc()
 
 class ProcessorMetrics:
@@ -111,42 +126,40 @@ class SubscriberMetrics:
 
         self.processor = processor
         self.subscriber = subscriber
-        self.workspace = workspace
-        self.flow = flow
 
         if not hasattr(__class__, "state_metric"):
             __class__.state_metric = Enum(
                 'subscriber_state', 'Subscriber state',
-                ["processor", "workspace", "flow", "subscriber"],
+                ["processor", "subscriber"],
                 states=['stopped', 'running']
             )
 
         if not hasattr(__class__, "received_metric"):
             __class__.received_metric = Counter(
                 'received_count', 'Received count',
-                ["processor", "workspace", "flow", "subscriber"],
+                ["processor", "subscriber"],
             )
 
         if not hasattr(__class__, "dropped_metric"):
             __class__.dropped_metric = Counter(
                 'dropped_count', 'Dropped messages count',
-                ["processor", "workspace", "flow", "subscriber"],
+                ["processor", "subscriber"],
             )
 
     def received(self) -> None:
         __class__.received_metric.labels(
-            processor=self.processor, workspace=self.workspace,
-            flow=self.flow, subscriber=self.subscriber,
+            processor=self.processor,
+            subscriber=self.subscriber,
         ).inc()
 
     def state(self, state: str) -> None:
         __class__.state_metric.labels(
-            processor=self.processor, workspace=self.workspace,
-            flow=self.flow, subscriber=self.subscriber,
+            processor=self.processor,
+            subscriber=self.subscriber,
         ).state(state)
 
     def dropped(self, state: str) -> None:
         __class__.dropped_metric.labels(
-            processor=self.processor, workspace=self.workspace,
-            flow=self.flow, subscriber=self.subscriber,
+            processor=self.processor,
+            subscriber=self.subscriber,
         ).inc()

trustgraph_base-2.8.3/trustgraph/base/producer_spec.py

@@ -0,0 +1,22 @@
+from __future__ import annotations
+
+from typing import Any
+
+from . spec import Spec
+
+class ProducerSpec(Spec):
+    def __init__(self, name: str, schema: Any) -> None:
+        self.name = name
+        self.schema = schema
+
+    async def register(self, flow: Any, processor: Any, definition: dict[str, Any]) -> Any:
+
+        topic = definition["topics"][self.name]
+
+        handle = await processor.sender_pool.add_producer(
+            topic=topic,
+            schema=self.schema,
+        )
+
+        flow.producer[self.name] = handle
+        return handle

{trustgraph_base-2.8.1 → trustgraph_base-2.8.3}/trustgraph/base/prompt_client.py

@@ -5,7 +5,7 @@ import inspect
 from dataclasses import dataclass
 from typing import Optional, Any
 
-from . request_response_spec import RequestResponse, RequestResponseSpec
+from . request_response_spec import RequestResponseSpec
 from .. schema import PromptRequest, PromptResponse
 
 @dataclass
@@ -18,7 +18,7 @@ class PromptResult:
     out_token: Optional[int] = None
     model: Optional[str] = None
 
-class PromptClient(RequestResponse):
+class PromptClient:
 
     async def prompt(self, id, variables, timeout=600, streaming=False, chunk_callback=None):
 

{trustgraph_base-2.8.1 → trustgraph_base-2.8.3}/trustgraph/base/pubsub.py

@@ -77,59 +77,6 @@ def get_async_pubsub(**config: Any) -> Any:
         )
 
 
-def get_pubsub(**config: Any) -> Any:
-    """
-    Factory function to create a pub/sub backend based on configuration.
-
-    Args:
-        config: Configuration dictionary from command-line args.
-                Key 'pubsub_backend' selects the backend (default: 'pulsar').
-
-    Returns:
-        Backend instance implementing the PubSubBackend protocol.
-    """
-    backend_type = config.get('pubsub_backend', 'pulsar')
-
-    if backend_type == 'pulsar':
-        from .pulsar_backend import PulsarBackend
-        return PulsarBackend(
-            host=config.get('pulsar_host', DEFAULT_PULSAR_HOST),
-            api_key=config.get('pulsar_api_key', DEFAULT_PULSAR_API_KEY),
-            listener=config.get('pulsar_listener'),
-            admin_url=config.get('pulsar_admin_url', DEFAULT_PULSAR_ADMIN_URL),
-        )
-    elif backend_type == 'rabbitmq':
-        from .rabbitmq_backend import RabbitMQBackend
-        return RabbitMQBackend(
-            host=config.get('rabbitmq_host', DEFAULT_RABBITMQ_HOST),
-            port=config.get('rabbitmq_port', DEFAULT_RABBITMQ_PORT),
-            username=config.get('rabbitmq_username', DEFAULT_RABBITMQ_USERNAME),
-            password=config.get('rabbitmq_password', DEFAULT_RABBITMQ_PASSWORD),
-            vhost=config.get('rabbitmq_vhost', DEFAULT_RABBITMQ_VHOST),
-        )
-    elif backend_type == 'kafka':
-        from .kafka_backend import KafkaBackend
-        return KafkaBackend(
-            bootstrap_servers=config.get(
-                'kafka_bootstrap_servers', DEFAULT_KAFKA_BOOTSTRAP,
-            ),
-            security_protocol=config.get(
-                'kafka_security_protocol', DEFAULT_KAFKA_PROTOCOL,
-            ),
-            sasl_mechanism=config.get(
-                'kafka_sasl_mechanism', DEFAULT_KAFKA_SASL_MECHANISM,
-            ),
-            sasl_username=config.get(
-                'kafka_sasl_username', DEFAULT_KAFKA_SASL_USERNAME,
-            ),
-            sasl_password=config.get(
-                'kafka_sasl_password', DEFAULT_KAFKA_SASL_PASSWORD,
-            ),
-        )
-    else:
-        raise ValueError(f"Unknown pub/sub backend: {backend_type}")
-
-
 STANDALONE_PULSAR_HOST = 'pulsar://localhost:6650'
 STANDALONE_PULSAR_ADMIN_URL = 'http://localhost:8080'
 

{trustgraph_base-2.8.1 → trustgraph_base-2.8.3}/trustgraph/base/receiver_pool.py

@@ -26,6 +26,7 @@ class ConsumerRegistration:
     backend_consumer: AsyncBackendConsumer
     receiver_task: asyncio.Task
     pool: 'ReceiverPool'
+    metrics: Any = None
 
     async def unregister(self):
         await self.pool.remove_consumer(self)
@@ -58,9 +59,13 @@ class ReceiverPool:
         await pool.stop()
     """
 
-    def __init__(self, backend: AsyncPubSubBackend, concurrency: int = 1):
+    def __init__(
+        self, backend: AsyncPubSubBackend, concurrency: int = 1,
+        processor_id: str | None = None,
+    ):
         self.backend = backend
         self.concurrency = concurrency
+        self.processor_id = processor_id
         self.work_queue = asyncio.Queue(maxsize=concurrency)
         self.registrations: list[ConsumerRegistration] = []
         self.worker_tasks: list[asyncio.Task] = []
@@ -93,6 +98,11 @@ class ReceiverPool:
                 await reg.receiver_task
             except BaseException:
                 pass
+
+        for reg in self.registrations:
+            if reg.metrics:
+                reg.metrics.state("stopped")
+
         self.registrations.clear()
 
         if not self.work_queue.empty():
@@ -136,8 +146,18 @@ class ReceiverPool:
             initial_position=initial_position,
         )
 
+        metrics = None
+        if self.processor_id:
+            from .metrics import ConsumerMetrics
+            metrics = ConsumerMetrics(
+                processor=self.processor_id, consumer=topic,
+            )
+            metrics.state("running")
+
+        instrumented = self._instrumented_handler(handler, metrics)
+
         receiver_task = asyncio.create_task(
-            self._receiver_loop(backend_consumer, handler),
+            self._receiver_loop(backend_consumer, instrumented),
             name=f"receiver-{topic}",
         )
 
@@ -147,12 +167,32 @@ class ReceiverPool:
             backend_consumer=backend_consumer,
             receiver_task=receiver_task,
             pool=self,
+            metrics=metrics,
         )
         self.registrations.append(reg)
 
         logger.info(f"Added consumer: {topic}")
         return reg
 
+    def _instrumented_handler(
+        self,
+        handler: Callable[..., Awaitable[None]],
+        metrics,
+    ) -> Callable[..., Awaitable[None]]:
+        if metrics is None:
+            return handler
+
+        async def wrapper(message):
+            with metrics.record_time():
+                try:
+                    await handler(message)
+                    metrics.process("ok")
+                except Exception:
+                    metrics.process("error")
+                    raise
+
+        return wrapper
+
     async def remove_consumer(self, reg: ConsumerRegistration):
         try:
             await reg.backend_consumer.close()
@@ -168,6 +208,9 @@ class ReceiverPool:
         if reg in self.registrations:
             self.registrations.remove(reg)
 
+        if reg.metrics:
+            reg.metrics.state("stopped")
+
         logger.info(f"Removed consumer: {reg.topic}")
 
     async def _receiver_loop(