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

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

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: trustgraph-base
-Version: 2.8.0
+Version: 2.8.2
 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.2}/pyproject.toml

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

{trustgraph_base-2.8.0 → trustgraph_base-2.8.2}/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.0 → trustgraph_base-2.8.2}/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.0 → trustgraph_base-2.8.2}/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.2/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.2}/trustgraph/base/async_processor.py

@@ -15,13 +15,14 @@ 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
 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.2}/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,22 +11,11 @@ logger = logging.getLogger(__name__)
 
 class AuditPublisher:
 
-    def __init__(self, component_name, processor_id=None, backend=None):
+    def __init__(self, component_name, async_backend=None):
         self.component_name = component_name
         self._handle = None
-
-        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
+        self._async_producer = None
+        self._async_backend = async_backend
 
     async def start(self, sender_pool=None):
         if sender_pool is not None:
@@ -36,8 +23,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,7 +52,7 @@ class AuditPublisher:
         try:
             if self._handle:
                 await self._handle.send(event)
-            elif self._legacy_producer:
-                await self._legacy_producer.send(event)
+            elif self._async_producer:
+                await self._async_producer.send(event)
         except Exception as e:
             logger.warning(f"Failed to emit audit event: {e}")

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

@@ -1,13 +1,13 @@
 
-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):
+class ConfigClient:
 
-    async def _request(self, timeout=CONFIG_TIMEOUT, **kwargs):
+    async def _request(self, timeout=None, **kwargs):
         resp = await self.request(
             ConfigRequest(**kwargs),
             timeout=timeout,
@@ -18,7 +18,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 +30,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 +39,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 +52,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 +61,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 +74,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 +84,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 +96,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.2/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.0 → trustgraph_base-2.8.2}/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.0 → trustgraph_base-2.8.2}/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.0 → trustgraph_base-2.8.2}/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(