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

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

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

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

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

{trustgraph_base-2.7.14 → trustgraph_base-2.8.0}/trustgraph/base/__init__.py

@@ -1,5 +1,5 @@
 
-from . pubsub import get_pubsub, add_pubsub_args
+from . pubsub import get_pubsub, get_async_pubsub, add_pubsub_args
 from . async_processor import AsyncProcessor
 from . consumer import Consumer
 from . producer import Producer
@@ -51,4 +51,11 @@ from . row_embeddings_query_client import RowEmbeddingsQueryClientSpec
 from . collection_config_handler import CollectionConfigHandler
 from . audit_publisher import AuditPublisher
 from . schema_compatibility import is_strict_mode_compatible
+from . async_backend import (
+    AsyncPubSubBackend, AsyncBackendConsumer, AsyncBackendProducer, Message,
+)
+from . receiver_pool import ReceiverPool, ConsumerRegistration
+from . sender_pool import SenderPool, ProducerHandle
+from . request_response_client import RequestResponseClient
+from . async_librarian_client import AsyncLibrarianClient
 

trustgraph_base-2.8.0/trustgraph/base/async_backend.py

@@ -0,0 +1,119 @@
+"""
+Async backend abstraction interfaces for pub/sub systems.
+
+These protocols replace the synchronous protocols in backend.py.
+All methods that perform broker I/O are async. The backend handles
+its own threading internally — the application layer never creates
+threads for broker operations.
+"""
+
+from typing import Protocol, Any, runtime_checkable
+
+
+@runtime_checkable
+class Message(Protocol):
+    """Protocol for a received message."""
+
+    def value(self) -> Any:
+        """Get the deserialized message content."""
+        ...
+
+    def properties(self) -> dict:
+        """Get message properties/metadata."""
+        ...
+
+
+@runtime_checkable
+class AsyncBackendConsumer(Protocol):
+    """Protocol for an async backend consumer.
+
+    receive() is awaitable with no timeout — shutdown is signalled
+    by cancelling the awaiting task. acknowledge() and
+    negative_acknowledge() are async so that backends with threading
+    constraints (e.g. RabbitMQ) can route ack calls internally.
+    """
+
+    async def receive(self) -> Message:
+        """Await the next message.
+
+        Does not time out. Shutdown is signalled by cancelling the
+        awaiting task.
+        """
+        ...
+
+    async def acknowledge(self, message: Message) -> None:
+        """Acknowledge successful processing of a message."""
+        ...
+
+    async def negative_acknowledge(self, message: Message) -> None:
+        """Negative acknowledge — triggers redelivery."""
+        ...
+
+    async def close(self) -> None:
+        """Close the consumer and release broker resources."""
+        ...
+
+
+@runtime_checkable
+class AsyncBackendProducer(Protocol):
+    """Protocol for an async backend producer.
+
+    send() is async so the backend can perform broker I/O without
+    blocking the event loop.
+    """
+
+    async def send(self, message: Any, properties: dict = {}) -> None:
+        """Send a message with optional properties."""
+        ...
+
+    async def close(self) -> None:
+        """Flush pending messages and close the producer."""
+        ...
+
+
+@runtime_checkable
+class AsyncPubSubBackend(Protocol):
+    """Protocol for an async pub/sub backend.
+
+    Factory methods are async since they perform broker roundtrips.
+    """
+
+    async def create_consumer(
+        self, topic: str, subscription: str, schema: type,
+        initial_position: str = 'latest', **options,
+    ) -> AsyncBackendConsumer:
+        """Create a consumer subscribed to a topic.
+
+        Returns only when the subscription is active and ready to
+        receive messages.
+        """
+        ...
+
+    async def create_producer(
+        self, topic: str, schema: type, **options,
+    ) -> AsyncBackendProducer:
+        """Create a producer bound to a topic.
+
+        Returns only when the producer is ready to send.
+        """
+        ...
+
+    async def create_topic(self, topic: str) -> None:
+        """Create broker-side resources for a logical topic."""
+        ...
+
+    async def delete_topic(self, topic: str) -> None:
+        """Delete a topic and discard any in-flight messages."""
+        ...
+
+    async def topic_exists(self, topic: str) -> bool:
+        """Check whether a topic exists."""
+        ...
+
+    async def ensure_topic(self, topic: str) -> None:
+        """Ensure a topic exists, creating it if necessary."""
+        ...
+
+    async def close(self) -> None:
+        """Close the backend and release all resources."""
+        ...

trustgraph_base-2.8.0/trustgraph/base/async_config_client.py

@@ -0,0 +1,121 @@
+
+from .. schema import ConfigRequest, ConfigResponse, ConfigKey, ConfigValue
+from . request_response_client import RequestResponseClient
+
+CONFIG_TIMEOUT = 10
+
+
+class AsyncConfigClient:
+
+    def __init__(self, client):
+        self._client = client
+
+    @classmethod
+    async def create(cls, backend, request_topic, response_topic,
+                     subscription=None):
+        client = await RequestResponseClient.create(
+            backend=backend,
+            request_topic=request_topic,
+            response_topic=response_topic,
+            request_schema=ConfigRequest,
+            response_schema=ConfigResponse,
+            subscription=subscription,
+        )
+        return cls(client)
+
+    async def start(self):
+        pass
+
+    async def stop(self):
+        await self._client.close()
+
+    async def close(self):
+        await self._client.close()
+
+    async def request(self, message, timeout=CONFIG_TIMEOUT):
+        return await self._client.request(message, timeout=timeout)
+
+    async def _request(self, timeout=CONFIG_TIMEOUT, **kwargs):
+        resp = await self._client.request(
+            ConfigRequest(**kwargs),
+            timeout=timeout,
+        )
+        if resp.error:
+            raise RuntimeError(
+                f"{resp.error.type}: {resp.error.message}"
+            )
+        return resp
+
+    async def get(self, workspace, type, key, timeout=CONFIG_TIMEOUT):
+        resp = await self._request(
+            operation="get",
+            workspace=workspace,
+            keys=[ConfigKey(type=type, key=key)],
+            timeout=timeout,
+        )
+        if resp.values and len(resp.values) > 0:
+            return resp.values[0].value
+        return None
+
+    async def put(self, workspace, type, key, value, timeout=CONFIG_TIMEOUT):
+        await self._request(
+            operation="put",
+            workspace=workspace,
+            values=[ConfigValue(type=type, key=key, value=value)],
+            timeout=timeout,
+        )
+
+    async def put_many(self, workspace, values, timeout=CONFIG_TIMEOUT):
+        await self._request(
+            operation="put",
+            workspace=workspace,
+            values=[
+                ConfigValue(type=t, key=k, value=v)
+                for t, k, v in values
+            ],
+            timeout=timeout,
+        )
+
+    async def delete(self, workspace, type, key, timeout=CONFIG_TIMEOUT):
+        await self._request(
+            operation="delete",
+            workspace=workspace,
+            keys=[ConfigKey(type=type, key=key)],
+            timeout=timeout,
+        )
+
+    async def delete_many(self, workspace, keys, timeout=CONFIG_TIMEOUT):
+        await self._request(
+            operation="delete",
+            workspace=workspace,
+            keys=[
+                ConfigKey(type=t, key=k)
+                for t, k in keys
+            ],
+            timeout=timeout,
+        )
+
+    async def keys(self, workspace, type, timeout=CONFIG_TIMEOUT):
+        resp = await self._request(
+            operation="list",
+            workspace=workspace,
+            type=type,
+            timeout=timeout,
+        )
+        return resp.directory
+
+    async def get_all(self, workspace, timeout=CONFIG_TIMEOUT):
+        resp = await self._request(
+            operation="config",
+            workspace=workspace,
+            timeout=timeout,
+        )
+        return resp.config
+
+    async def workspaces_for_type(self, type, timeout=CONFIG_TIMEOUT):
+        resp = await self._request(
+            operation="getvalues-all-ws",
+            type=type,
+            timeout=timeout,
+        )
+        return {v.workspace for v in resp.values if v.workspace}

trustgraph_base-2.8.0/trustgraph/base/async_librarian_client.py

@@ -0,0 +1,260 @@
+"""
+Async librarian client using the async pub/sub backend directly.
+
+Replaces LibrarianClient's Consumer+Producer (which create OS threads)
+with async backend producers/consumers and a dedicated receiver
+coroutine. Supports both single-response and streaming operations.
+"""
+
+import asyncio
+import base64
+import logging
+import uuid
+
+from ..schema import LibrarianRequest, LibrarianResponse, DocumentMetadata
+
+logger = logging.getLogger(__name__)
+
+
+class AsyncLibrarianClient:
+
+    def __init__(self):
+        self._producer = None
+        self._consumer = None
+        self._receiver_task = None
+        self._pending = {}
+        self._streams = {}
+        self.running = False
+
+    @classmethod
+    async def create(
+        cls, backend, request_topic, response_topic,
+        subscription=None,
+    ):
+        client = cls()
+
+        client._producer = await backend.create_producer(
+            topic=request_topic,
+            schema=LibrarianRequest,
+        )
+
+        if subscription is None:
+            subscription = f"librarian-{uuid.uuid4().hex[:12]}"
+
+        client._consumer = await backend.create_consumer(
+            topic=response_topic,
+            subscription=subscription,
+            schema=LibrarianResponse,
+            initial_position='latest',
+        )
+
+        client.running = True
+        client._receiver_task = asyncio.create_task(
+            client._response_loop(),
+            name=f"librarian-receiver-{response_topic}",
+        )
+
+        logger.info(
+            f"AsyncLibrarianClient created: "
+            f"req={request_topic}, resp={response_topic}"
+        )
+        return client
+
+    async def start(self):
+        pass
+
+    async def _response_loop(self):
+        try:
+            while self.running:
+                msg = await self._consumer.receive()
+                response = msg.value()
+                request_id = msg.properties().get("id")
+
+                if request_id:
+                    if request_id in self._pending:
+                        future = self._pending.pop(request_id)
+                        if not future.done():
+                            future.set_result(response)
+                    elif request_id in self._streams:
+                        await self._streams[request_id].put(response)
+
+                await self._consumer.acknowledge(msg)
+
+        except asyncio.CancelledError:
+            raise
+        except Exception as e:
+            logger.error(
+                f"Librarian response loop error: {e}", exc_info=True,
+            )
+
+    async def request(self, request, timeout=120):
+        request_id = str(uuid.uuid4())
+
+        future = asyncio.get_event_loop().create_future()
+        self._pending[request_id] = future
+
+        try:
+            await self._producer.send(
+                request, properties={"id": request_id},
+            )
+            response = await asyncio.wait_for(future, timeout=timeout)
+
+            if response.error:
+                raise RuntimeError(
+                    f"Librarian error: {response.error.type}: "
+                    f"{response.error.message}"
+                )
+
+            return response
+
+        except asyncio.TimeoutError:
+            self._pending.pop(request_id, None)
+            raise RuntimeError("Timeout waiting for librarian response")
+
+    async def stream(self, request, timeout=120):
+        request_id = str(uuid.uuid4())
+
+        q = asyncio.Queue()
+        self._streams[request_id] = q
+
+        try:
+            await self._producer.send(
+                request, properties={"id": request_id},
+            )
+
+            chunks = []
+            while True:
+                response = await asyncio.wait_for(
+                    q.get(), timeout=timeout,
+                )
+
+                if response.error:
+                    raise RuntimeError(
+                        f"Librarian error: {response.error.type}: "
+                        f"{response.error.message}"
+                    )
+
+                chunks.append(response)
+
+                if response.is_final:
+                    break
+
+            return chunks
+
+        except asyncio.TimeoutError:
+            self._streams.pop(request_id, None)
+            raise RuntimeError("Timeout waiting for librarian stream")
+        finally:
+            self._streams.pop(request_id, None)
+
+    async def stop(self):
+        self.running = False
+
+        if self._receiver_task:
+            self._receiver_task.cancel()
+            try:
+                await self._receiver_task
+            except asyncio.CancelledError:
+                pass
+
+        for future in self._pending.values():
+            if not future.done():
+                future.cancel()
+        self._pending.clear()
+
+        for q in self._streams.values():
+            await q.put(None)
+        self._streams.clear()
+
+        if self._producer:
+            try:
+                await self._producer.close()
+            except Exception as e:
+                logger.warning(f"Error closing librarian producer: {e}")
+
+        if self._consumer:
+            try:
+                await self._consumer.close()
+            except Exception as e:
+                logger.warning(f"Error closing librarian consumer: {e}")
+
+        logger.info("AsyncLibrarianClient closed")
+
+    async def fetch_document_content(self, document_id, timeout=120):
+        req = LibrarianRequest(
+            operation="stream-document",
+            document_id=document_id,
+        )
+        chunks = await self.stream(req, timeout=timeout)
+
+        raw = b""
+        for chunk in chunks:
+            if chunk.content:
+                if isinstance(chunk.content, bytes):
+                    raw += base64.b64decode(chunk.content)
+                else:
+                    raw += base64.b64decode(
+                        chunk.content.encode("utf-8")
+                    )
+
+        return base64.b64encode(raw)
+
+    async def fetch_document_text(self, document_id, timeout=120):
+        content = await self.fetch_document_content(
+            document_id, timeout=timeout,
+        )
+        return base64.b64decode(content).decode("utf-8")
+
+    async def fetch_document_metadata(self, document_id, timeout=120):
+        req = LibrarianRequest(
+            operation="get-document-metadata",
+            document_id=document_id,
+        )
+        response = await self.request(req, timeout=timeout)
+        return response.document_metadata
+
+    async def save_child_document(self, doc_id, parent_id, content,
+                                  document_type="chunk", title=None,
+                                  kind="text/plain", timeout=120):
+        if isinstance(content, str):
+            content = content.encode("utf-8")
+
+        doc_metadata = DocumentMetadata(
+            id=doc_id,
+            kind=kind,
+            title=title or doc_id,
+            parent_id=parent_id,
+            document_type=document_type,
+        )
+
+        req = LibrarianRequest(
+            operation="add-child-document",
+            document_metadata=doc_metadata,
+            content=base64.b64encode(content).decode("utf-8"),
+        )
+
+        await self.request(req, timeout=timeout)
+        return doc_id
+
+    async def save_document(self, doc_id, content, title=None,
+                            document_type="answer", kind="text/plain",
+                            timeout=120):
+        if isinstance(content, str):
+            content = content.encode("utf-8")
+
+        doc_metadata = DocumentMetadata(
+            id=doc_id,
+            kind=kind,
+            title=title or doc_id,
+            document_type=document_type,
+        )
+
+        req = LibrarianRequest(
+            operation="add-document",
+            document_id=doc_id,
+            document_metadata=doc_metadata,
+            content=base64.b64encode(content).decode("utf-8"),
+        )
+
+        await self.request(req, timeout=timeout)
+        return doc_id