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

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

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: trustgraph-base
-Version: 2.8.0
+Version: 2.8.1
 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
@@ -15,6 +15,7 @@ Requires-Dist: python-logging-loki
 Requires-Dist: pika
 Requires-Dist: aio-pika
 Requires-Dist: confluent-kafka
+Requires-Dist: aiokafka
 Requires-Dist: pyyaml
 Requires-Dist: websockets
 

{trustgraph_base-2.8.0 → trustgraph_base-2.8.1}/pyproject.toml

@@ -17,6 +17,7 @@ dependencies = [
     "pika",
     "aio-pika",
     "confluent-kafka",
+    "aiokafka",
     "pyyaml",
     "websockets",
 ]

{trustgraph_base-2.8.0 → trustgraph_base-2.8.1}/trustgraph/base/async_config_client.py

@@ -12,7 +12,7 @@ class AsyncConfigClient:
 
     @classmethod
     async def create(cls, backend, request_topic, response_topic,
-                     subscription=None):
+                     subscription=None, timeout=CONFIG_TIMEOUT):
         client = await RequestResponseClient.create(
             backend=backend,
             request_topic=request_topic,
@@ -20,6 +20,7 @@ class AsyncConfigClient:
             request_schema=ConfigRequest,
             response_schema=ConfigResponse,
             subscription=subscription,
+            default_timeout=timeout,
         )
         return cls(client)
 

trustgraph_base-2.8.1/trustgraph/base/async_kafka_backend.py

@@ -0,0 +1,304 @@
+"""
+Async Kafka backend implementation.
+
+Uses aiokafka for fully async Kafka operations. Each logical topic
+maps to a Kafka topic, matching the sync Kafka backend's topology:
+
+    class:topicspace:topic  →  topicspace.class.topic
+
+Producers publish to the topic directly.
+Consumers use consumer groups for competing-consumer semantics:
+
+  - flow / request: named consumer group (competing consumers)
+  - response / notify: unique consumer group per instance, filtering
+    messages by correlation ID (all subscribers see all messages)
+"""
+
+import asyncio
+import json
+import logging
+import uuid
+from typing import Any
+
+from aiokafka import AIOKafkaProducer, AIOKafkaConsumer
+from aiokafka.admin import AIOKafkaAdminClient, NewTopic
+from aiokafka.errors import TopicAlreadyExistsError, UnknownTopicOrPartitionError
+
+from .async_backend import AsyncPubSubBackend, AsyncBackendProducer
+from .async_backend import AsyncBackendConsumer, Message
+from .serialization import dataclass_to_dict, dict_to_dataclass
+
+logger = logging.getLogger(__name__)
+
+LONG_RETENTION_MS = 7 * 24 * 60 * 60 * 1000   # 7 days
+SHORT_RETENTION_MS = 300 * 1000                 # 5 minutes
+
+
+class AsyncKafkaMessage:
+    """Wrapper for aiokafka messages to match Message protocol."""
+
+    def __init__(self, msg, schema_cls):
+        self._msg = msg
+        self._schema_cls = schema_cls
+        self._value = None
+
+    def value(self) -> Any:
+        if self._value is None:
+            data_dict = json.loads(self._msg.value.decode('utf-8'))
+            self._value = dict_to_dataclass(data_dict, self._schema_cls)
+        return self._value
+
+    def properties(self) -> dict:
+        headers = self._msg.headers or []
+        return {
+            k: v.decode('utf-8') if isinstance(v, bytes) else v
+            for k, v in headers
+        }
+
+
+class AsyncKafkaProducer:
+    """Async Kafka producer using aiokafka."""
+
+    def __init__(self, producer, topic_name):
+        self._producer = producer
+        self._topic_name = topic_name
+
+    async def send(self, message: Any, properties: dict = {}) -> None:
+        data_dict = dataclass_to_dict(message)
+        json_data = json.dumps(data_dict).encode('utf-8')
+
+        headers = [
+            (k, str(v).encode('utf-8'))
+            for k, v in properties.items()
+        ] if properties else None
+
+        await self._producer.send_and_wait(
+            self._topic_name,
+            value=json_data,
+            headers=headers,
+        )
+
+    async def close(self) -> None:
+        try:
+            await self._producer.stop()
+        except BaseException:
+            pass
+
+
+class AsyncKafkaConsumer:
+    """Async Kafka consumer using aiokafka."""
+
+    def __init__(self, consumer, schema_cls):
+        self._consumer = consumer
+        self._schema_cls = schema_cls
+
+    async def receive(self) -> Message:
+        msg = await self._consumer.__anext__()
+        return AsyncKafkaMessage(msg, self._schema_cls)
+
+    async def acknowledge(self, message: Message) -> None:
+        if isinstance(message, AsyncKafkaMessage):
+            await self._consumer.commit()
+
+    async def negative_acknowledge(self, message: Message) -> None:
+        pass
+
+    async def close(self) -> None:
+        try:
+            await self._consumer.stop()
+        except BaseException:
+            pass
+
+
+class AsyncKafkaBackend:
+    """Async Kafka pub/sub backend using aiokafka.
+
+    Fully async — zero OS threads for receiving, sending, and
+    consumer/producer lifecycle.
+    """
+
+    def __init__(
+        self, bootstrap_servers='localhost:9092',
+        security_protocol='PLAINTEXT',
+        sasl_mechanism=None,
+        sasl_username=None,
+        sasl_password=None,
+    ):
+        self._bootstrap_servers = bootstrap_servers
+        self._security_protocol = security_protocol
+        self._sasl_mechanism = sasl_mechanism
+        self._sasl_username = sasl_username
+        self._sasl_password = sasl_password
+
+        self.max_consumer_concurrency = 1
+
+        logger.info(
+            f"Async Kafka backend: {bootstrap_servers} "
+            f"protocol={security_protocol}"
+        )
+
+    def _client_kwargs(self):
+        kwargs = {
+            'bootstrap_servers': self._bootstrap_servers,
+            'security_protocol': self._security_protocol,
+        }
+        if self._sasl_mechanism:
+            kwargs['sasl_mechanism'] = self._sasl_mechanism
+        if self._sasl_username:
+            kwargs['sasl_plain_username'] = self._sasl_username
+        if self._sasl_password:
+            kwargs['sasl_plain_password'] = self._sasl_password
+        return kwargs
+
+    def _parse_topic(self, topic_id: str) -> tuple[str, str, bool]:
+        """Parse topic identifier into Kafka topic name, class, durability.
+
+        Format: class:topicspace:topic
+        Returns: (topic_name, class, durable)
+        """
+        if ':' not in topic_id:
+            return f'tg.flow.{topic_id}', 'flow', True
+
+        parts = topic_id.split(':', 2)
+        if len(parts) != 3:
+            raise ValueError(
+                f"Invalid topic format: {topic_id}, "
+                f"expected class:topicspace:topic"
+            )
+
+        cls, topicspace, topic = parts
+
+        if cls == 'flow':
+            durable = True
+        elif cls in ('request', 'response', 'notify'):
+            durable = False
+        else:
+            raise ValueError(
+                f"Invalid topic class: {cls}, "
+                f"expected flow, request, response, or notify"
+            )
+
+        topic_name = f"{topicspace}.{cls}.{topic}".replace(':', '.')
+
+        return topic_name, cls, durable
+
+    def _retention_ms(self, cls):
+        if cls == 'flow':
+            return LONG_RETENTION_MS
+        return SHORT_RETENTION_MS
+
+    async def create_producer(
+        self, topic: str, schema: type, **options,
+    ) -> AsyncBackendProducer:
+        topic_name, cls, durable = self._parse_topic(topic)
+
+        producer = AIOKafkaProducer(
+            **self._client_kwargs(),
+            acks='all' if durable else 1,
+            max_request_size=10485760,
+        )
+        await producer.start()
+
+        logger.debug(f"Created async producer: topic={topic_name}")
+
+        return AsyncKafkaProducer(producer, topic_name)
+
+    async def create_consumer(
+        self, topic: str, subscription: str, schema: type,
+        initial_position: str = 'latest', **options,
+    ) -> AsyncBackendConsumer:
+        topic_name, cls, durable = self._parse_topic(topic)
+
+        if cls in ('response', 'notify'):
+            group_id = f"{subscription}-{uuid.uuid4()}"
+            auto_offset_reset = 'earliest'
+        else:
+            group_id = subscription
+            auto_offset_reset = (
+                'earliest' if initial_position == 'earliest'
+                else 'latest'
+            )
+
+        consumer = AIOKafkaConsumer(
+            topic_name,
+            **self._client_kwargs(),
+            group_id=group_id,
+            auto_offset_reset=auto_offset_reset,
+            enable_auto_commit=False,
+            max_partition_fetch_bytes=10485760,
+            session_timeout_ms=6000,
+            heartbeat_interval_ms=1000,
+        )
+        await consumer.start()
+
+        logger.debug(
+            f"Created async consumer: topic={topic_name}, "
+            f"group={group_id}, cls={cls}"
+        )
+
+        return AsyncKafkaConsumer(consumer, schema)
+
+    async def create_topic(self, topic: str) -> None:
+        topic_name, cls, durable = self._parse_topic(topic)
+        retention_ms = self._retention_ms(cls)
+
+        admin = AIOKafkaAdminClient(**self._client_kwargs())
+        try:
+            await admin.start()
+            new_topic = NewTopic(
+                topic_name,
+                num_partitions=1,
+                replication_factor=1,
+                topic_configs={
+                    'retention.ms': str(retention_ms),
+                },
+            )
+            await admin.create_topics([new_topic])
+            logger.info(f"Created topic: {topic_name}")
+        except TopicAlreadyExistsError:
+            logger.debug(f"Topic already exists: {topic_name}")
+        finally:
+            await admin.close()
+
+    async def delete_topic(self, topic: str) -> None:
+        topic_name, cls, durable = self._parse_topic(topic)
+
+        admin = AIOKafkaAdminClient(**self._client_kwargs())
+        try:
+            await admin.start()
+            await admin.delete_topics([topic_name])
+            logger.info(f"Deleted topic: {topic_name}")
+        except UnknownTopicOrPartitionError:
+            logger.debug(f"Topic not found: {topic_name}")
+        except BaseException as e:
+            logger.debug(f"Topic delete for {topic_name}: {e}")
+        finally:
+            try:
+                await admin.close()
+            except BaseException:
+                pass
+
+    async def topic_exists(self, topic: str) -> bool:
+        topic_name, cls, durable = self._parse_topic(topic)
+
+        admin = AIOKafkaAdminClient(**self._client_kwargs())
+        try:
+            await admin.start()
+            metadata = await admin.describe_topics([topic_name])
+            return len(metadata) > 0
+        except UnknownTopicOrPartitionError:
+            return False
+        except BaseException:
+            return False
+        finally:
+            try:
+                await admin.close()
+            except BaseException:
+                pass
+
+    async def ensure_topic(self, topic: str) -> None:
+        if not await self.topic_exists(topic):
+            await self.create_topic(topic)
+
+    async def close(self) -> None:
+        logger.info("Async Kafka backend closed")

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

@@ -22,6 +22,7 @@ from . request_response_client import RequestResponseClient
 from . logging import add_logging_args, setup_logging
 
 default_config_queue = config_push_queue
+default_config_timeout = 60
 
 logger = logging.getLogger(__name__)
 
@@ -60,6 +61,10 @@ class AsyncProcessor:
             "config_push_queue", default_config_queue
         )
 
+        self.config_timeout = int(params.get(
+            "config_timeout", default_config_timeout
+        ))
+
         self.config_handlers = []
         self.workspace_handlers = []
         self.config_version = 0
@@ -83,6 +88,7 @@ class AsyncProcessor:
             response_topic=config_response_queue,
             request_schema=ConfigRequest,
             response_schema=ConfigResponse,
+            default_timeout=self.config_timeout,
         )
 
     async def _fetch_type_workspace(self, client, workspace, config_type):
@@ -92,7 +98,6 @@ class AsyncProcessor:
                 workspace=workspace,
                 type=config_type,
             ),
-            timeout=60,
         )
         if resp.error:
             raise RuntimeError(f"Config error: {resp.error.message}")
@@ -104,7 +109,6 @@ class AsyncProcessor:
                 operation="getkeys-all-ws",
                 type=config_type,
             ),
-            timeout=60,
         )
         if resp.error:
             raise RuntimeError(f"Config error: {resp.error.message}")
@@ -433,6 +437,14 @@ class AsyncProcessor:
             help='Number of concurrent workers (default: 1)',
         )
 
+        parser.add_argument(
+            '--config-timeout',
+            type=int,
+            default=default_config_timeout,
+            help=f'Config request timeout in seconds '
+                 f'(default: {default_config_timeout})',
+        )
+
         parser.add_argument(
             '--metrics',
             action=argparse.BooleanOptionalAction,

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

@@ -13,9 +13,12 @@ logger = logging.getLogger(__name__)
 
 class AuditPublisher:
 
-    def __init__(self, component_name, processor_id=None, backend=None):
+    def __init__(self, component_name, processor_id=None, backend=None,
+                 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(
@@ -36,8 +39,19 @@ class AuditPublisher:
                 topic=audit_events_queue,
                 schema=AuditEvent,
             )
+        elif self._async_backend is not None:
+            self._async_producer = await self._async_backend.create_producer(
+                topic=audit_events_queue,
+                schema=AuditEvent,
+            )
 
     async def stop(self):
+        if self._async_producer:
+            try:
+                await self._async_producer.close()
+            except BaseException:
+                pass
+            self._async_producer = None
         if self._handle:
             await self._handle.unregister()
 
@@ -54,6 +68,8 @@ class AuditPublisher:
         try:
             if self._handle:
                 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:

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

@@ -7,7 +7,10 @@ CONFIG_TIMEOUT = 10
 
 class ConfigClient(RequestResponse):
 
-    async def _request(self, timeout=CONFIG_TIMEOUT, **kwargs):
+    def __init__(self, timeout=CONFIG_TIMEOUT, **kwargs):
+        super(ConfigClient, self).__init__(default_timeout=timeout, **kwargs)
+
+    async def _request(self, timeout=None, **kwargs):
         resp = await self.request(
             ConfigRequest(**kwargs),
             timeout=timeout,
@@ -18,7 +21,7 @@ class ConfigClient(RequestResponse):
             )
         return resp
 
-    async def get(self, workspace, type, key, timeout=CONFIG_TIMEOUT):
+    async def get(self, workspace, type, key, timeout=None):
         """Get a single config value. Returns the value string or None."""
         resp = await self._request(
             operation="get",
@@ -30,7 +33,7 @@ class ConfigClient(RequestResponse):
             return resp.values[0].value
         return None
 
-    async def put(self, workspace, type, key, value, timeout=CONFIG_TIMEOUT):
+    async def put(self, workspace, type, key, value, timeout=None):
         """Put a single config value."""
         await self._request(
             operation="put",
@@ -39,7 +42,7 @@ class ConfigClient(RequestResponse):
             timeout=timeout,
         )
 
-    async def put_many(self, workspace, values, timeout=CONFIG_TIMEOUT):
+    async def put_many(self, workspace, values, timeout=None):
         """Put multiple config values in a single request within a
         single workspace. values is a list of (type, key, value) tuples."""
         await self._request(
@@ -52,7 +55,7 @@ class ConfigClient(RequestResponse):
             timeout=timeout,
         )
 
-    async def delete(self, workspace, type, key, timeout=CONFIG_TIMEOUT):
+    async def delete(self, workspace, type, key, timeout=None):
         """Delete a single config key."""
         await self._request(
             operation="delete",
@@ -61,7 +64,7 @@ class ConfigClient(RequestResponse):
             timeout=timeout,
         )
 
-    async def delete_many(self, workspace, keys, timeout=CONFIG_TIMEOUT):
+    async def delete_many(self, workspace, keys, timeout=None):
         """Delete multiple config keys in a single request within a
         single workspace. keys is a list of (type, key) tuples."""
         await self._request(
@@ -74,7 +77,7 @@ class ConfigClient(RequestResponse):
             timeout=timeout,
         )
 
-    async def keys(self, workspace, type, timeout=CONFIG_TIMEOUT):
+    async def keys(self, workspace, type, timeout=None):
         """List all keys for a config type within a workspace."""
         resp = await self._request(
             operation="list",
@@ -84,7 +87,7 @@ class ConfigClient(RequestResponse):
         )
         return resp.directory
 
-    async def get_all(self, workspace, timeout=CONFIG_TIMEOUT):
+    async def get_all(self, workspace, timeout=None):
         """Return every config entry in ``workspace`` as a nested dict
         ``{type: {key: value}}``.  Values are returned as the raw
         strings stored by config-svc (typically JSON); callers parse
@@ -96,7 +99,7 @@ class ConfigClient(RequestResponse):
         )
         return resp.config
 
-    async def workspaces_for_type(self, type, timeout=CONFIG_TIMEOUT):
+    async def workspaces_for_type(self, type, timeout=None):
         """Return the set of distinct workspaces with any config of
         the given type."""
         resp = await self._request(

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

@@ -29,8 +29,8 @@ def get_async_pubsub(**config: Any) -> Any:
     """
     Factory function to create an async pub/sub backend.
 
-    Currently only Pulsar is supported. RabbitMQ and Kafka async
-    backends are not yet implemented.
+    Creates the appropriate async backend based on the 'pubsub_backend'
+    config key. Supported: pulsar, rabbitmq, kafka.
     """
     backend_type = config.get('pubsub_backend', 'pulsar')
 
@@ -51,10 +51,29 @@ def get_async_pubsub(**config: Any) -> Any:
             password=config.get('rabbitmq_password', DEFAULT_RABBITMQ_PASSWORD),
             vhost=config.get('rabbitmq_vhost', DEFAULT_RABBITMQ_VHOST),
         )
+    elif backend_type == 'kafka':
+        from .async_kafka_backend import AsyncKafkaBackend
+        return AsyncKafkaBackend(
+            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"Async backend not yet supported for: {backend_type}. "
-            f"Supported: 'pulsar', 'rabbitmq'."
+            f"Supported: 'pulsar', 'rabbitmq', 'kafka'."
         )
 
 

{trustgraph_base-2.8.0 → trustgraph_base-2.8.1}/trustgraph/base/request_response_client.py

@@ -33,12 +33,13 @@ class RequestResponseClient:
         await client.close()
     """
 
-    def __init__(self):
+    def __init__(self, default_timeout=60):
         self.producer = None
         self.consumer = None
         self.receiver_task = None
         self.pending: dict[str, asyncio.Future] = {}
         self.running = False
+        self.default_timeout = default_timeout
 
     @classmethod
     async def create(
@@ -49,8 +50,9 @@ class RequestResponseClient:
         request_schema: type,
         response_schema: type,
         subscription: str | None = None,
+        default_timeout: float = 60,
     ) -> 'RequestResponseClient':
-        client = cls()
+        client = cls(default_timeout=default_timeout)
 
         client.producer = await backend.create_producer(
             topic=request_topic,
@@ -80,9 +82,12 @@ class RequestResponseClient:
         return client
 
     async def request(
-        self, message: Any, timeout: float = 60,
+        self, message: Any, timeout: float | None = None,
         properties: dict | None = None,
     ) -> Any:
+        if timeout is None:
+            timeout = self.default_timeout
+
         request_id = str(uuid.uuid4())
         future = asyncio.get_event_loop().create_future()
         self.pending[request_id] = future
@@ -100,9 +105,12 @@ class RequestResponseClient:
             raise
 
     async def request_stream(
-        self, message: Any, timeout: float = 300,
+        self, message: Any, timeout: float | None = None,
         properties: dict | None = None,
     ):
+        if timeout is None:
+            timeout = self.default_timeout
+
         request_id = str(uuid.uuid4())
         queue = asyncio.Queue()
         self.pending[request_id] = queue

{trustgraph_base-2.8.0 → trustgraph_base-2.8.1}/trustgraph/base/request_response_spec.py

@@ -21,8 +21,11 @@ class RequestResponse(Subscriber):
             request_metrics,
             response_topic, response_schema,
             response_metrics,
+            default_timeout=300,
     ):
 
+        self.default_timeout = default_timeout
+
         super(RequestResponse, self).__init__(
             backend = backend,
             subscription = subscription,
@@ -47,7 +50,10 @@ class RequestResponse(Subscriber):
         await self.producer.stop()
         await super(RequestResponse, self).stop()
 
-    async def request(self, req, timeout=300, recipient=None):
+    async def request(self, req, timeout=None, recipient=None):
+
+        if timeout is None:
+            timeout = self.default_timeout
 
         id = str(uuid.uuid4())
 

trustgraph_base-2.8.1/trustgraph/base_version.py

@@ -0,0 +1 @@
+__version__ = "2.8.1"

{trustgraph_base-2.8.0 → trustgraph_base-2.8.1}/trustgraph_base.egg-info/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: trustgraph-base
-Version: 2.8.0
+Version: 2.8.1
 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
@@ -15,6 +15,7 @@ Requires-Dist: python-logging-loki
 Requires-Dist: pika
 Requires-Dist: aio-pika
 Requires-Dist: confluent-kafka
+Requires-Dist: aiokafka
 Requires-Dist: pyyaml
 Requires-Dist: websockets
 

{trustgraph_base-2.8.0 → trustgraph_base-2.8.1}/trustgraph_base.egg-info/SOURCES.txt

@@ -26,6 +26,7 @@ trustgraph/base/agent_client.py
 trustgraph/base/agent_service.py
 trustgraph/base/async_backend.py
 trustgraph/base/async_config_client.py
+trustgraph/base/async_kafka_backend.py
 trustgraph/base/async_librarian_client.py
 trustgraph/base/async_processor.py
 trustgraph/base/async_pulsar_backend.py

{trustgraph_base-2.8.0 → trustgraph_base-2.8.1}/trustgraph_base.egg-info/requires.txt

@@ -5,5 +5,6 @@ python-logging-loki
 pika
 aio-pika
 confluent-kafka
+aiokafka
 pyyaml
 websockets

trustgraph_base-2.8.0/trustgraph/base_version.py

@@ -1 +0,0 @@
-__version__ = "2.8.0"