apunkt/trustgraph
1
0
Fork 0
mirror of https://github.com/trustgraph-ai/trustgraph.git synced 2026-07-17 09:11:03 +02:00
Code Issues Projects Releases Packages Wiki Activity Actions

API finalised

Browse source
This commit is contained in:
Cyber MacGeddon 2025-04-16 16:11:18 +01:00
parent 91d28e9d56
commit b94b4b7389
10 changed files with 87 additions and 133 deletions
Download patch file Download diff file Expand all files Collapse all files

10
trustgraph-base/trustgraph/base/async_processor.py
View file

@ -48,9 +48,10 @@ class AsyncProcessor:
self.config_handlers = []
self.config_sub_task = self.subscribe(
self.config_push_queue, config_subscriber_id,
schema=ConfigPush, handler=self.on_config_change,
start_of_messages=True
flow = None,
queue = self.config_push_queue, subscriber = config_subscriber_id,
schema = ConfigPush, handler = self.on_config_change,
start_of_messages = True
)
self.running = True
@ -84,11 +85,12 @@ class AsyncProcessor:
await asyncio.sleep(2)
def subscribe(
self, queue, subscriber, schema, handler, metrics=None,
self, flow, queue, subscriber, schema, handler, metrics=None,
start_of_messages=False
):
return Consumer(
flow = flow,
taskgroup = self.taskgroup,
client = self.client,
queue = queue,

23
trustgraph-base/trustgraph/base/consumer.py
View file

@ -1,4 +1,5 @@
import _pulsar
import asyncio
import time
@ -7,13 +8,14 @@ from .. exceptions import TooManyRequests
class Consumer:
def __init__(
self, taskgroup, client, queue, subscriber, schema,
self, taskgroup, flow, client, queue, subscriber, schema,
handler, rate_limit_retry_time = 10,
rate_limit_timeout = 7200, metrics = None,
start_of_messages=False,
):
self.taskgroup = taskgroup
self.flow = flow
self.client = client
self.queue = queue
self.subscriber = subscriber
@ -28,6 +30,12 @@ class Consumer:
self.metrics = metrics
def __del__(self):
self.running = False
if hasattr(self, "consumer"):
self.consumer.close()
async def stop(self):
self.running = False
@ -58,7 +66,15 @@ class Consumer:
while self.running:
msg = await asyncio.to_thread(self.consumer.receive)
try:
msg = await asyncio.to_thread(
self.consumer.receive,
timeout_millis=2000
)
except _pulsar.Timeout:
continue
except Exception as e:
raise e
expiry = time.time() + self.rate_limit_timeout
@ -86,7 +102,7 @@ class Consumer:
if self.metrics:
with self.metrics.record_time():
await self.handler(msg, self)
await self.handler(msg, self, self.flow)
else:
await self.handler(msg, self.consumer)
@ -132,3 +148,4 @@ class Consumer:
if self.metrics:
self.metrics.state("stopped")

50
trustgraph-base/trustgraph/base/flow_processor.py
View file

@ -10,6 +10,9 @@ from .. base import AsyncProcessor, Consumer, Producer
from .. base import ProcessorMetrics, ConsumerMetrics, ProducerMetrics
class Flow:
pass
class FlowProcessor(AsyncProcessor):
def __init__(self, **params):
@ -31,16 +34,12 @@ class FlowProcessor(AsyncProcessor):
self.on_config(self.on_configuration)
self.consumers = {}
self.producers = {}
self.flows = {}
# These can be overriden by a derived class
self.consumer_spec = []
self.producer_spec = []
# "input", self.input_schema)
# "input_schema": Document,
# "output_schema": TextDocument,
self.config_spec = []
print("Service initialised.")
@ -50,9 +49,21 @@ class FlowProcessor(AsyncProcessor):
def register_producer(self, name, schema):
self.producer_spec.append((name, schema))
def register_config(self, name):
self.config_spec.append((name,))
async def start_flow(self, flow, defn):
producers = {}
flow_obj = Flow()
flow_obj.producer = {}
flow_obj.consumer = {}
flow_obj.config = {}
for spec in self.config_spec:
name = spec[0]
flow_obj.config[name] = defn[name]
setattr(flow_obj, name, defn[name])
for spec in self.producer_spec:
@ -68,9 +79,9 @@ class FlowProcessor(AsyncProcessor):
metrics = producer_metrics,
)
producers[name] = producer
flow_obj.producer[name] = producer
consumers = {}
setattr(flow_obj, name, producer)
for spec in self.consumer_spec:
@ -81,6 +92,7 @@ class FlowProcessor(AsyncProcessor):
)
consumer = self.subscribe(
flow = flow_obj,
queue = defn[name],
subscriber = self.subscriber,
schema = schema,
@ -92,25 +104,25 @@ class FlowProcessor(AsyncProcessor):
# metadata
consumer.id = self.id
consumer.name = name
consumer.flow = flow
consumer.q = producers
consumers[name] = consumer
consumer.flow = flow_obj
consumer.flow.name = flow
await consumer.start()
self.consumers[flow] = consumers
self.producers[flow] = producers
flow_obj.consumer[name] = consumer
setattr(flow_obj, name, consumer)
self.flows[flow] = flow_obj
print("Started flow: ", flow)
async def stop_flow(self, flow):
for c in self.consumers[flow]:
for c in self.flows[flow].consumer:
await c.stop()
del self.consumers[flow]
del self.producers[flow]
del self.flows[flow]
print("Stopped flow: ", flow, flush=True)
@ -125,7 +137,7 @@ class FlowProcessor(AsyncProcessor):
flow_config = json.loads(config["flows"][self.id])
wanted_keys = flow_config.keys()
current_keys = self.consumers.keys()
current_keys = self.flows.keys()
for key in wanted_keys:
if key not in current_keys:

4
trustgraph-base/trustgraph/base/producer.py
View file

@ -12,7 +12,8 @@ class Producer:
def __del__(self):
self.producer.close()
if hasattr(self, "producer"):
self.producer.close()
async def send(self, msg, properties={}):
@ -21,4 +22,3 @@ class Producer:
self.producer.send(msg, properties)

12
trustgraph-base/trustgraph/base/pubsub.py
View file

@ -51,7 +51,8 @@ class PulsarClient:
if self.client:
self.client.close()
def subscribe(self, queue, subscriber, schema, start_of_messages=False):
def subscribe(self, queue, subscriber, schema,
start_of_messages=False):
if start_of_messages:
pos = pulsar.InitialPosition.Earliest
@ -59,10 +60,11 @@ class PulsarClient:
pos = pulsar.InitialPosition.Latest
return self.client.subscribe(
queue, subscriber,
consumer_type=pulsar.ConsumerType.Shared,
schema=JsonSchema(schema),
initial_position=pos,
topic = queue,
subscription_name = subscriber,
consumer_type = pulsar.ConsumerType.Shared,
schema = JsonSchema(schema),
initial_position = pos,
)
def publish(self, queue, schema):

106
trustgraph-base/trustgraph/base/request_response.py
View file

@ -1,4 +1,5 @@
import json
from pulsar.schema import JsonSchema
from .. schema import Error
@ -8,109 +9,28 @@ from .. log_level import LogLevel
from .. base import AsyncProcessor, Consumer, Producer
from .. base import ProcessorMetrics, ConsumerMetrics, ProducerMetrics
from . flow_processor import FlowProcessor
class RequestResponseProcessor(AsyncProcessor):
class RequestResponseService(FlowProcessor):
def __init__(self, **params):
id = params.get("id")
self.subscriber = params.get("subscriber")
self.request_schema = params.get("request_schema")
self.response_schema = params.get("response_schema")
ProcessorMetrics(id=id).info(
{
"subscriber": self.subscriber,
"request_schema": self.request_schema.__name__,
"response_schema": self.response_schema.__name__,
}
)
super(RequestResponseService, self).__init__(**params)
super(Processor, self).__init__(
**params | {
"request_schema": self.request_schema.__name__,
"response_schema": self.response_schema.__name__,
}
)
self.on_config(self.on_configuration)
self.subs = {}
self.pubs = {}
# These can be overriden by a derived class
self.consumer_spec = [
("request", params.get("request_schema"), self.on_request)
]
self.producer_spec = [
("response", params.get("response_schema"))
]
print("Service initialised.")
async def start_handler(self, flow, defn):
request_metrics = ConsumerMetrics(self.id, flow)
response_metrics = ProducerMetrics(self.id, flow)
self.subs[flow] = self.subscribe(
queue = defn["input"]
subscriber = subscriber,
schema = self.request_schema,
handler = self.on_message,
metrics = request_metrics,
)
self.pubs[flow] = self.publish(
queue = defn["output"],
schema = self.response_schema,
metrics = response_metrics,
)
self.subs[flow].start()
print("Started flow for", flow)
async def on_configuration(self, config, version):
if "flows" not in config: return
if self.id in config["flows"]:
wanted_keys = config["flows"][self.id].keys()
current_keys = self.subs.keys()
for key in wanted_keys:
if key not in current_keys:
self.start_handler(key, config["flows"][key])
for key in current_keys:
if key not in wanted_keys:
self.stop_handler(key, config["flows"][key])
print("Handled config update")
async def start(self):
await self.request_sub.start()
async def on_message(self, msg, consumer):
try:
v = msg.value()
# Sender-produced ID
id = msg.properties()["id"]
print(f"Handling {id}...", flush=True)
consumer.acknowledge(msg)
except Exception as e:
resp = "error"
await self.response_pub.send(resp, properties={"id": id})
consumer.acknowledge(msg)
@staticmethod
def add_args(parser):
def add_args(parser, default_subscriber):
AsyncProcessor.add_args(parser)
FlowProcessor.add_args(parser)
def run():

4
trustgraph-flow/trustgraph/chunking/recursive/chunker.py
View file

@ -59,7 +59,7 @@ class Processor(FlowProcessor):
print("Chunker initialised", flush=True)
async def on_message(self, msg, consumer):
async def on_message(self, msg, consumer, flow):
v = msg.value()
print(f"Chunking {v.metadata.id}...", flush=True)
@ -81,7 +81,7 @@ class Processor(FlowProcessor):
id=consumer.id, flow=consumer.flow
).observe(len(chunk.page_content))
await consumer.q["output"].send(r)
await flow.producer["output"].send(r)
print("Done.", flush=True)

4
trustgraph-flow/trustgraph/chunking/token/chunker.py
View file

@ -58,7 +58,7 @@ class Processor(FlowProcessor):
print("Chunker initialised", flush=True)
async def on_message(self, msg, consumer):
async def on_message(self, msg, consumer, flow):
v = msg.value()
print(f"Chunking {v.metadata.id}...", flush=True)
@ -80,7 +80,7 @@ class Processor(FlowProcessor):
id=consumer.id, flow=consumer.flow
).observe(len(chunk.page_content))
await consumer.q["output"].send(r)
await flow.producer["output"].send(r)
print("Done.", flush=True)

3
trustgraph-flow/trustgraph/config/service/service.py
View file

@ -74,6 +74,7 @@ class Processor(AsyncProcessor):
)
self.subs = self.subscribe(
flow = None,
queue = request_queue,
subscriber = subscriber,
schema = request_schema,
@ -105,7 +106,7 @@ class Processor(AsyncProcessor):
print("Pushed version ", self.config.version)
async def on_message(self, msg, consumer):
async def on_message(self, msg, consumer, flow):
try:

4
trustgraph-flow/trustgraph/decoding/pdf/pdf_decoder.py
View file

@ -44,7 +44,7 @@ class Processor(FlowProcessor):
print("PDF inited", flush=True)
async def on_message(self, msg, consumer):
async def on_message(self, msg, consumer, flow):
print("PDF message received", flush=True)
@ -71,7 +71,7 @@ class Processor(FlowProcessor):
text=page.page_content.encode("utf-8"),
)
await consumer.q["output"].send(r)
await flow.producer["output"].send(r)
print("Done.", flush=True)