trustgraph-base 2.3.14__tar.gz → 2.3.16__tar.gz

{trustgraph_base-2.3.14 → trustgraph_base-2.3.16}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: trustgraph-base
-Version: 2.3.14
+Version: 2.3.16
 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
@@ -13,6 +13,7 @@ Requires-Dist: prometheus-client
 Requires-Dist: requests
 Requires-Dist: python-logging-loki
 Requires-Dist: pika
+Requires-Dist: confluent-kafka
 Requires-Dist: pyyaml
 
 See https://trustgraph.ai/

{trustgraph_base-2.3.14 → trustgraph_base-2.3.16}/pyproject.toml

@@ -15,6 +15,7 @@ dependencies = [
     "requests",
     "python-logging-loki",
     "pika",
+    "confluent-kafka",
     "pyyaml",
 ]
 classifiers = [

{trustgraph_base-2.3.14 → trustgraph_base-2.3.16}/trustgraph/base/consumer.py

@@ -54,6 +54,17 @@ class Consumer:
         self.running = True
         self.consumer_task = None
 
+        # Kafka topics are created with 1 partition, so multiple
+        # consumers in the same group causes rebalance storms where
+        # no consumer can fetch. Cap to the backend's limit.
+        max_concurrency = getattr(backend, 'max_consumer_concurrency', None)
+        if isinstance(max_concurrency, int) and concurrency > max_concurrency:
+            logger.info(
+                f"Capping concurrency from {concurrency} to "
+                f"{max_concurrency} (backend limit)"
+            )
+            concurrency = max_concurrency
+
         self.concurrency = concurrency
 
         self.metrics = metrics

trustgraph_base-2.3.16/trustgraph/base/kafka_backend.py

@@ -0,0 +1,452 @@
+"""
+Kafka backend implementation for pub/sub abstraction.
+
+Each logical topic maps to a Kafka topic.  The topic name encodes
+the full identity:
+
+    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)
+
+The flow service manages topic lifecycle via AdminClient.
+
+Architecture:
+  Producer  -->  [Kafka topic]  -->  Consumer Group A  --> Consumer
+                                -->  Consumer Group A  --> Consumer
+                                -->  Consumer Group B  --> Consumer (response)
+"""
+
+import asyncio
+import json
+import logging
+import uuid
+from typing import Any
+
+from confluent_kafka import (
+    Producer as KafkaProducer,
+    Consumer as KafkaConsumer,
+    TopicPartition,
+    KafkaError,
+    KafkaException,
+)
+from confluent_kafka.admin import AdminClient, NewTopic
+
+from .backend import PubSubBackend, BackendProducer, BackendConsumer, Message
+from .serialization import dataclass_to_dict, dict_to_dataclass
+
+logger = logging.getLogger(__name__)
+
+# Retention defaults (milliseconds)
+LONG_RETENTION_MS = 7 * 24 * 60 * 60 * 1000   # 7 days
+SHORT_RETENTION_MS = 300 * 1000                 # 5 minutes
+
+
+class KafkaMessage:
+    """Wrapper for Kafka 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:
+        """Deserialize and return the message value as a dataclass."""
+        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:
+        """Return message properties from Kafka headers."""
+        headers = self._msg.headers() or []
+        return {
+            k: v.decode('utf-8') if isinstance(v, bytes) else v
+            for k, v in headers
+        }
+
+
+class KafkaBackendProducer:
+    """Publishes messages to a Kafka topic.
+
+    confluent-kafka Producer is thread-safe, so a single instance
+    can be shared across threads.
+    """
+
+    def __init__(self, bootstrap_servers, topic_name, durable):
+        self._topic_name = topic_name
+        self._durable = durable
+        self._producer = KafkaProducer({
+            'bootstrap.servers': bootstrap_servers,
+            'acks': 'all' if durable else '1',
+            'message.max.bytes': 10485760,
+        })
+
+    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
+
+        self._delivery_error = None
+
+        def _on_delivery(err, msg):
+            if err:
+                self._delivery_error = err
+
+        self._producer.produce(
+            topic=self._topic_name,
+            value=json_data,
+            headers=headers,
+            on_delivery=_on_delivery,
+        )
+        self._producer.flush()
+
+        if self._delivery_error:
+            raise KafkaException(self._delivery_error)
+
+    def flush(self) -> None:
+        self._producer.flush()
+
+    def close(self) -> None:
+        self._producer.flush()
+
+
+class KafkaBackendConsumer:
+    """Consumes from a Kafka topic using a consumer group.
+
+    Uses confluent-kafka Consumer.poll() for message delivery.
+    Not thread-safe — each instance must be used from a single thread,
+    which matches the ThreadPoolExecutor pattern in consumer.py.
+    """
+
+    def __init__(self, bootstrap_servers, topic_name, group_id,
+                 schema_cls, auto_offset_reset='latest'):
+        self._bootstrap_servers = bootstrap_servers
+        self._topic_name = topic_name
+        self._group_id = group_id
+        self._schema_cls = schema_cls
+        self._auto_offset_reset = auto_offset_reset
+        self._consumer = None
+
+    def _connect(self):
+        import time
+        t0 = time.monotonic()
+
+        def _on_assign(consumer, partitions):
+            elapsed = time.monotonic() - t0
+            logger.info(
+                f"Partition assignment for {self._topic_name}: "
+                f"{[p.partition for p in partitions]} "
+                f"after {elapsed:.1f}s"
+            )
+
+        def _on_revoke(consumer, partitions):
+            logger.info(
+                f"Partition revoke for {self._topic_name}: "
+                f"{[p.partition for p in partitions]}"
+            )
+
+        self._consumer = KafkaConsumer({
+            'bootstrap.servers': self._bootstrap_servers,
+            'group.id': self._group_id,
+            'auto.offset.reset': self._auto_offset_reset,
+            'enable.auto.commit': False,
+            'fetch.message.max.bytes': 10485760,
+            # Tighten group coordination timeouts for fast
+            # group join on single-member groups.
+            'session.timeout.ms': 6000,
+            'heartbeat.interval.ms': 1000,
+        })
+        self._consumer.subscribe(
+            [self._topic_name],
+            on_assign=_on_assign,
+            on_revoke=_on_revoke,
+        )
+        logger.info(
+            f"Kafka consumer subscribed: topic={self._topic_name}, "
+            f"group={self._group_id}"
+        )
+
+    def _is_alive(self):
+        return self._consumer is not None
+
+    def ensure_connected(self) -> None:
+        """Eagerly connect and subscribe.
+
+        For response/notify consumers this must be called before the
+        corresponding request is published, so that the consumer is
+        assigned a partition and will see the response message.
+        """
+        if not self._is_alive():
+            self._connect()
+
+            # Kick off group join. With auto.offset.reset=earliest
+            # on response/notify consumers, any messages published
+            # before assignment completes will be picked up once
+            # the consumer starts polling in receive().
+            self._consumer.poll(timeout=0.5)
+
+    def receive(self, timeout_millis: int = 2000) -> Message:
+        """Receive a message. Raises TimeoutError if none available."""
+        if not self._is_alive():
+            self._connect()
+
+        timeout_seconds = timeout_millis / 1000.0
+        msg = self._consumer.poll(timeout=timeout_seconds)
+
+        if msg is None:
+            raise TimeoutError("No message received within timeout")
+
+        if msg.error():
+            error = msg.error()
+            if error.code() == KafkaError._PARTITION_EOF:
+                raise TimeoutError("End of partition reached")
+            if error.code() == KafkaError.UNKNOWN_TOPIC_OR_PART:
+                raise TimeoutError("Topic not yet available")
+            raise KafkaException(error)
+
+        return KafkaMessage(msg, self._schema_cls)
+
+    def acknowledge(self, message: Message) -> None:
+        """Commit the message's offset (next offset to read)."""
+        if isinstance(message, KafkaMessage) and message._msg:
+            tp = TopicPartition(
+                message._msg.topic(),
+                message._msg.partition(),
+                message._msg.offset() + 1,
+            )
+            self._consumer.commit(offsets=[tp], asynchronous=False)
+
+    def negative_acknowledge(self, message: Message) -> None:
+        """Seek back to the message's offset for redelivery."""
+        if isinstance(message, KafkaMessage) and message._msg:
+            tp = TopicPartition(
+                message._msg.topic(),
+                message._msg.partition(),
+                message._msg.offset(),
+            )
+            self._consumer.seek(tp)
+
+    def unsubscribe(self) -> None:
+        if self._consumer:
+            try:
+                self._consumer.unsubscribe()
+            except Exception:
+                pass
+
+    def close(self) -> None:
+        if self._consumer:
+            try:
+                self._consumer.close()
+            except Exception:
+                pass
+            self._consumer = None
+
+
+class KafkaBackend:
+    """Kafka pub/sub backend using one topic per logical topic."""
+
+    def __init__(self, bootstrap_servers='localhost:9092',
+                 security_protocol='PLAINTEXT',
+                 sasl_mechanism=None,
+                 sasl_username=None,
+                 sasl_password=None):
+        self._bootstrap_servers = bootstrap_servers
+
+        # AdminClient config
+        self._admin_config = {
+            'bootstrap.servers': bootstrap_servers,
+        }
+
+        if security_protocol != 'PLAINTEXT':
+            self._admin_config['security.protocol'] = security_protocol
+            if sasl_mechanism:
+                self._admin_config['sasl.mechanism'] = sasl_mechanism
+            if sasl_username:
+                self._admin_config['sasl.username'] = sasl_username
+            if sasl_password:
+                self._admin_config['sasl.password'] = sasl_password
+
+        # Topics are created with 1 partition, so only 1 consumer
+        # per group can be active. Extra consumers cause rebalance
+        # storms that block message delivery.
+        self.max_consumer_concurrency = 1
+
+        logger.info(
+            f"Kafka backend: {bootstrap_servers} "
+            f"protocol={security_protocol}"
+        )
+
+    def _parse_topic(self, topic_id: str) -> tuple[str, str, bool]:
+        """
+        Parse topic identifier into Kafka topic name, class, and 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"
+            )
+
+        # Replace any remaining colons — flow topics can have
+        # extra segments (e.g. flow:tg:document-load:default)
+        # and Kafka rejects colons in topic names.
+        topic_name = f"{topicspace}.{cls}.{topic}".replace(':', '.')
+
+        return topic_name, cls, durable
+
+    def _retention_ms(self, cls):
+        """Return retention.ms for a topic class."""
+        if cls == 'flow':
+            return LONG_RETENTION_MS
+        return SHORT_RETENTION_MS
+
+    def create_producer(self, topic: str, schema: type,
+                        **options) -> BackendProducer:
+        topic_name, cls, durable = self._parse_topic(topic)
+        logger.debug(f"Creating producer: topic={topic_name}")
+        return KafkaBackendProducer(
+            self._bootstrap_servers, topic_name, durable,
+        )
+
+    def create_consumer(self, topic: str, subscription: str, schema: type,
+                        initial_position: str = 'latest',
+                        **options) -> BackendConsumer:
+        """Create a consumer subscribed to a Kafka topic.
+
+        Behaviour is determined by the topic's class prefix:
+          - flow: named consumer group, competing consumers
+          - request: named consumer group, competing consumers
+          - response: unique consumer group per instance
+          - notify: unique consumer group per instance
+        """
+        topic_name, cls, durable = self._parse_topic(topic)
+
+        if cls in ('response', 'notify'):
+            # Per-subscriber: unique group so every instance sees
+            # every message.  Filter by correlation ID happens at
+            # the Subscriber layer above.
+            # Use 'earliest' so that responses published before
+            # partition assignment completes are not missed.
+            # Each group is unique (UUID) with no committed offsets,
+            # so 'earliest' reads from the start of the topic.
+            # The correlation ID filter discards non-matching messages.
+            group_id = f"{subscription}-{uuid.uuid4()}"
+            auto_offset_reset = 'earliest'
+        else:
+            # Shared: named group, competing consumers
+            group_id = subscription
+            auto_offset_reset = (
+                'earliest' if initial_position == 'earliest'
+                else 'latest'
+            )
+
+        logger.debug(
+            f"Creating consumer: topic={topic_name}, "
+            f"group={group_id}, cls={cls}"
+        )
+
+        return KafkaBackendConsumer(
+            self._bootstrap_servers, topic_name, group_id,
+            schema, auto_offset_reset,
+        )
+
+    def _create_topic_sync(self, topic_name, retention_ms):
+        """Blocking topic creation via AdminClient."""
+        admin = AdminClient(self._admin_config)
+        new_topic = NewTopic(
+            topic_name,
+            num_partitions=1,
+            replication_factor=1,
+            config={
+                'retention.ms': str(retention_ms),
+            },
+        )
+        fs = admin.create_topics([new_topic])
+        for name, f in fs.items():
+            try:
+                f.result()
+                logger.info(f"Created topic: {name}")
+            except KafkaException as e:
+                # Topic already exists — idempotent
+                if e.args[0].code() == KafkaError.TOPIC_ALREADY_EXISTS:
+                    logger.debug(f"Topic already exists: {name}")
+                else:
+                    raise
+
+    async def create_topic(self, topic: str) -> None:
+        """Create a Kafka topic with appropriate retention."""
+        topic_name, cls, durable = self._parse_topic(topic)
+        retention_ms = self._retention_ms(cls)
+        await asyncio.to_thread(
+            self._create_topic_sync, topic_name, retention_ms,
+        )
+
+    def _delete_topic_sync(self, topic_name):
+        """Blocking topic deletion via AdminClient."""
+        admin = AdminClient(self._admin_config)
+        fs = admin.delete_topics([topic_name])
+        for name, f in fs.items():
+            try:
+                f.result()
+                logger.info(f"Deleted topic: {name}")
+            except KafkaException as e:
+                # Topic doesn't exist — idempotent
+                if e.args[0].code() == KafkaError.UNKNOWN_TOPIC_OR_PART:
+                    logger.debug(f"Topic not found: {name}")
+                else:
+                    raise
+            except Exception as e:
+                logger.debug(f"Topic delete for {name}: {e}")
+
+    async def delete_topic(self, topic: str) -> None:
+        """Delete a Kafka topic."""
+        topic_name, cls, durable = self._parse_topic(topic)
+        await asyncio.to_thread(self._delete_topic_sync, topic_name)
+
+    def _topic_exists_sync(self, topic_name):
+        """Blocking topic existence check via AdminClient."""
+        admin = AdminClient(self._admin_config)
+        metadata = admin.list_topics(timeout=10)
+        return topic_name in metadata.topics
+
+    async def topic_exists(self, topic: str) -> bool:
+        """Check whether a Kafka topic exists."""
+        topic_name, cls, durable = self._parse_topic(topic)
+        return await asyncio.to_thread(
+            self._topic_exists_sync, topic_name,
+        )
+
+    async def ensure_topic(self, topic: str) -> None:
+        """Ensure a topic exists, creating it if necessary."""
+        if not await self.topic_exists(topic):
+            await self.create_topic(topic)
+
+    def close(self) -> None:
+        pass

{trustgraph_base-2.3.14 → trustgraph_base-2.3.16}/trustgraph/base/pubsub.py

@@ -17,6 +17,12 @@ DEFAULT_RABBITMQ_USERNAME = os.getenv("RABBITMQ_USERNAME", 'guest')
 DEFAULT_RABBITMQ_PASSWORD = os.getenv("RABBITMQ_PASSWORD", 'guest')
 DEFAULT_RABBITMQ_VHOST = os.getenv("RABBITMQ_VHOST", '/')
 
+DEFAULT_KAFKA_BOOTSTRAP = os.getenv("KAFKA_BOOTSTRAP_SERVERS", 'kafka:9092')
+DEFAULT_KAFKA_PROTOCOL = os.getenv("KAFKA_SECURITY_PROTOCOL", 'PLAINTEXT')
+DEFAULT_KAFKA_SASL_MECHANISM = os.getenv("KAFKA_SASL_MECHANISM", None)
+DEFAULT_KAFKA_SASL_USERNAME = os.getenv("KAFKA_SASL_USERNAME", None)
+DEFAULT_KAFKA_SASL_PASSWORD = os.getenv("KAFKA_SASL_PASSWORD", None)
+
 
 def get_pubsub(**config: Any) -> Any:
     """
@@ -47,6 +53,25 @@ def get_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 .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}")
 
@@ -65,6 +90,7 @@ def add_pubsub_args(parser: ArgumentParser, standalone: bool = False) -> None:
     pulsar_host = STANDALONE_PULSAR_HOST if standalone else DEFAULT_PULSAR_HOST
     pulsar_listener = 'localhost' if standalone else None
     rabbitmq_host = 'localhost' if standalone else DEFAULT_RABBITMQ_HOST
+    kafka_bootstrap = 'localhost:9092' if standalone else DEFAULT_KAFKA_BOOTSTRAP
 
     parser.add_argument(
         '--pubsub-backend',
@@ -122,3 +148,34 @@ def add_pubsub_args(parser: ArgumentParser, standalone: bool = False) -> None:
         default=DEFAULT_RABBITMQ_VHOST,
         help=f'RabbitMQ vhost (default: {DEFAULT_RABBITMQ_VHOST})',
     )
+
+    # Kafka options
+    parser.add_argument(
+        '--kafka-bootstrap-servers',
+        default=kafka_bootstrap,
+        help=f'Kafka bootstrap servers (default: {kafka_bootstrap})',
+    )
+
+    parser.add_argument(
+        '--kafka-security-protocol',
+        default=DEFAULT_KAFKA_PROTOCOL,
+        help=f'Kafka security protocol (default: {DEFAULT_KAFKA_PROTOCOL})',
+    )
+
+    parser.add_argument(
+        '--kafka-sasl-mechanism',
+        default=DEFAULT_KAFKA_SASL_MECHANISM,
+        help='Kafka SASL mechanism',
+    )
+
+    parser.add_argument(
+        '--kafka-sasl-username',
+        default=DEFAULT_KAFKA_SASL_USERNAME,
+        help='Kafka SASL username',
+    )
+
+    parser.add_argument(
+        '--kafka-sasl-password',
+        default=DEFAULT_KAFKA_SASL_PASSWORD,
+        help='Kafka SASL password',
+    )

trustgraph_base-2.3.16/trustgraph/base_version.py

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

{trustgraph_base-2.3.14 → trustgraph_base-2.3.16}/trustgraph_base.egg-info/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: trustgraph-base
-Version: 2.3.14
+Version: 2.3.16
 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
@@ -13,6 +13,7 @@ Requires-Dist: prometheus-client
 Requires-Dist: requests
 Requires-Dist: python-logging-loki
 Requires-Dist: pika
+Requires-Dist: confluent-kafka
 Requires-Dist: pyyaml
 
 See https://trustgraph.ai/

{trustgraph_base-2.3.14 → trustgraph_base-2.3.16}/trustgraph_base.egg-info/SOURCES.txt

@@ -44,6 +44,7 @@ trustgraph/base/graph_embeddings_client.py
 trustgraph/base/graph_embeddings_query_service.py
 trustgraph/base/graph_embeddings_store_service.py
 trustgraph/base/graph_rag_client.py
+trustgraph/base/kafka_backend.py
 trustgraph/base/librarian_client.py
 trustgraph/base/llm_service.py
 trustgraph/base/logging.py

{trustgraph_base-2.3.14 → trustgraph_base-2.3.16}/trustgraph_base.egg-info/requires.txt

@@ -3,4 +3,5 @@ prometheus-client
 requests
 python-logging-loki
 pika
+confluent-kafka
 pyyaml

trustgraph_base-2.3.14/trustgraph/base_version.py

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