mirror of
https://github.com/trustgraph-ai/trustgraph.git
synced 2026-07-08 21:02:12 +02:00
2a28ba0d8b
Adds a RabbitMQ backend as an alternative to Pulsar, selectable via PUBSUB_BACKEND=rabbitmq. Both backends implement the same PubSubBackend protocol — no application code changes needed to switch. RabbitMQ topology: - Single topic exchange per topicspace (e.g. 'tg') - Routing key derived from queue class and topic name - Shared consumers: named queue bound to exchange (competing, round-robin) - Exclusive consumers: anonymous auto-delete queue (broadcast, each gets every message). Used by Subscriber and config push consumer. - Thread-local producer connections (pika is not thread-safe) - Push-based consumption via basic_consume with process_data_events for heartbeat processing Consumer model changes: - Consumer class creates one backend consumer per concurrent task (required for pika thread safety, harmless for Pulsar) - Consumer class accepts consumer_type parameter - Subscriber passes consumer_type='exclusive' for broadcast semantics - Config push consumer uses consumer_type='exclusive' so every processor instance receives config updates - handle_one_from_queue receives consumer as parameter for correct per-connection ack/nack LibrarianClient: - New shared client class replacing duplicated librarian request-response code across 6+ services (chunking, decoders, RAG, etc.) - Uses stream-document instead of get-document-content for fetching document content in 1MB chunks (avoids broker message size limits) - Standalone object (self.librarian = LibrarianClient(...)) not a mixin - get-document-content marked deprecated in schema and OpenAPI spec Serialisation: - Extracted dataclass_to_dict/dict_to_dataclass to shared serialization.py (used by both Pulsar and RabbitMQ backends) Librarian queues: - Changed from flow class (persistent) back to request/response class now that stream-document eliminates large single messages - API upload chunk size reduced from 5MB to 3MB to stay under broker limits after base64 encoding Factory and CLI: - get_pubsub() handles 'rabbitmq' backend with RabbitMQ connection params - add_pubsub_args() includes RabbitMQ options (host, port, credentials) - add_pubsub_args(standalone=True) defaults to localhost for CLI tools - init_trustgraph skips Pulsar admin setup for non-Pulsar backends - tg-dump-queues and tg-monitor-prompts use backend abstraction - BaseClient and ConfigClient accept generic pubsub config
125 lines
3.2 KiB
Python
125 lines
3.2 KiB
Python
|
|
import uuid
|
|
import time
|
|
|
|
from .. exceptions import *
|
|
from ..base.pubsub import get_pubsub
|
|
|
|
# Default timeout for a request/response. In seconds.
|
|
DEFAULT_TIMEOUT=300
|
|
|
|
|
|
class BaseClient:
|
|
|
|
def __init__(
|
|
self,
|
|
subscriber=None,
|
|
input_queue=None,
|
|
output_queue=None,
|
|
input_schema=None,
|
|
output_schema=None,
|
|
**pubsub_config,
|
|
):
|
|
|
|
if input_queue == None: raise RuntimeError("Need input_queue")
|
|
if output_queue == None: raise RuntimeError("Need output_queue")
|
|
if input_schema == None: raise RuntimeError("Need input_schema")
|
|
if output_schema == None: raise RuntimeError("Need output_schema")
|
|
|
|
if subscriber == None:
|
|
subscriber = str(uuid.uuid4())
|
|
|
|
# Create backend using factory
|
|
self.backend = get_pubsub(**pubsub_config)
|
|
|
|
self.producer = self.backend.create_producer(
|
|
topic=input_queue,
|
|
schema=input_schema,
|
|
chunking_enabled=True,
|
|
)
|
|
|
|
self.consumer = self.backend.create_consumer(
|
|
topic=output_queue,
|
|
subscription=subscriber,
|
|
schema=output_schema,
|
|
consumer_type='shared',
|
|
)
|
|
|
|
self.input_schema = input_schema
|
|
self.output_schema = output_schema
|
|
|
|
def call(self, **args):
|
|
|
|
timeout = args.get("timeout", DEFAULT_TIMEOUT)
|
|
inspect = args.get("inspect", lambda x: True)
|
|
|
|
if "timeout" in args:
|
|
del args["timeout"]
|
|
|
|
if "inspect" in args:
|
|
del args["inspect"]
|
|
|
|
id = str(uuid.uuid4())
|
|
|
|
r = self.input_schema(**args)
|
|
|
|
end_time = time.time() + timeout
|
|
|
|
self.producer.send(r, properties={ "id": id })
|
|
|
|
while time.time() < end_time:
|
|
|
|
try:
|
|
msg = self.consumer.receive(timeout_millis=2500)
|
|
except TimeoutError:
|
|
continue
|
|
|
|
mid = msg.properties()["id"]
|
|
|
|
if mid == id:
|
|
|
|
value = msg.value()
|
|
|
|
if value.error:
|
|
|
|
self.consumer.acknowledge(msg)
|
|
|
|
if value.error.type == "llm-error":
|
|
raise LlmError(value.error.message)
|
|
|
|
elif value.error.type == "too-many-requests":
|
|
raise TooManyRequests(value.error.message)
|
|
|
|
elif value.error.type == "ParseError":
|
|
raise ParseError(value.error.message)
|
|
|
|
else:
|
|
|
|
raise RuntimeError(
|
|
f"{value.error.type}: {value.error.message}"
|
|
)
|
|
|
|
complete = inspect(value)
|
|
|
|
if not complete: continue
|
|
|
|
resp = msg.value()
|
|
self.consumer.acknowledge(msg)
|
|
return resp
|
|
|
|
# Ignore messages with wrong ID
|
|
self.consumer.acknowledge(msg)
|
|
|
|
raise TimeoutError("Timed out waiting for response")
|
|
|
|
def __del__(self):
|
|
|
|
if hasattr(self, "consumer"):
|
|
self.consumer.close()
|
|
|
|
if hasattr(self, "producer"):
|
|
self.producer.flush()
|
|
self.producer.close()
|
|
|
|
if hasattr(self, "backend"):
|
|
self.backend.close()
|