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 tweak for input_output

Browse source
This commit is contained in:
Cyber MacGeddon 2025-04-16 12:55:47 +01:00
parent 510cfb1103
commit 8f89336db3
7 changed files with 105 additions and 70 deletions
Download patch file Download diff file Expand all files Collapse all files

92
trustgraph-base/trustgraph/base/input_output.py
View file

@ -16,85 +16,101 @@ class InputOutputProcessor(AsyncProcessor):
self.id = params.get("id")
self.subscriber = params.get("subscriber")
self.input_schema = params.get("input_schema")
self.output_schema = params.get("output_schema")
ProcessorMetrics(id=self.id).info(
{
"subscriber": self.subscriber,
"input_schema": self.input_schema.__name__,
"output_schema": self.output_schema.__name__,
}
)
super(InputOutputProcessor, self).__init__(
**params | {
"id": self.id,
"input_schema": self.input_schema.__name__,
"output_schema": self.output_schema.__name__,
}
)
self.on_config(self.on_configuration)
self.consumers = {}
self.producers = {}
# These can be overriden by a derived class
self.consumer_spec = ("input", self.input_schema)
self.producer_spec = [
("output", self.output_schema)
]
self.consumer_spec = []
self.producer_spec = []
# "input", self.input_schema)
# "input_schema": Document,
# "output_schema": TextDocument,
print("Service initialised.")
def register_consumer(self, name, schema, handler):
self.consumer_spec.append((name, schema, handler))
def register_producer(self, name, schema):
self.producer_spec.append((name, schema))
async def start_flow(self, flow, defn):
consumer_tag = self.consumer_spec[0]
consumer_schema = self.consumer_spec[1]
consumer_metrics = ConsumerMetrics(self.id, f"{flow}-{consumer_tag}")
consumer = self.subscribe(
queue = defn[consumer_tag],
subscriber = self.subscriber,
schema = consumer_schema,
handler = self.on_message,
metrics = consumer_metrics,
)
class Queues:
pass
consumer.q = Queues()
consumer.id = self.id
consumer.flow = flow
producers = {}
for spec in self.producer_spec:
producer_tag = spec[0]
producer_schema = spec[1]
name, schema = spec
producer_metrics = ProducerMetrics(
self.id, f"{flow}-{producer_tag}"
self.id, f"{flow}-{name}"
)
producer = self.publish(
queue = defn[producer_tag],
schema = spec[1],
queue = defn[name],
schema = schema,
metrics = producer_metrics,
)
setattr(consumer.q, producer_tag, producer)
producers[name] = producer
self.consumers[flow] = consumer
consumers = {}
await consumer.start()
for spec in self.consumer_spec:
name, schema, handler = spec
consumer_metrics = ConsumerMetrics(
self.id, f"{flow}-{name}"
)
consumer = self.subscribe(
queue = defn[name],
subscriber = self.subscriber,
schema = schema,
handler = handler,
metrics = consumer_metrics,
)
# Consumer handle gets access to producers and other
# metadata
consumer.id = self.id
consumer.name = name
consumer.flow = flow
consumer.q = producers
consumers[name] = consumer
await consumer.start()
self.consumers[flow] = consumers
self.producers[flow] = producers
print("Started flow: ", flow)
async def stop_flow(self, flow):
await self.consumers[flow].stop()
for c in self.consumers[flow]:
await c.stop()
del self.consumers[flow]
del self.producers[flow]
print("Stopped flow: ", flow, flush=True)

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

@ -10,6 +10,10 @@ class Producer:
self.metrics = metrics
def __del__(self):
self.producer.close()
async def send(self, msg, properties={}):
if self.metrics:
@ -17,6 +21,4 @@ class Producer:
self.producer.send(msg, properties)
# FIXME
# __class__.output_metric.inc()

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

@ -28,8 +28,6 @@ class Processor(InputOutputProcessor):
super(Processor, self).__init__(
**params | {
"subscriber": subscriber,
"input_schema": TextDocument,
"output_schema": Chunk,
}
)
@ -48,6 +46,17 @@ class Processor(InputOutputProcessor):
is_separator_regex=False,
)
self.register_consumer(
name = "input",
schema = TextDocument,
handler = self.on_message,
)
self.register_producer(
name = "output",
schema = Chunk,
)
print("Chunker initialised", flush=True)
async def on_message(self, msg, consumer):
@ -72,7 +81,7 @@ class Processor(InputOutputProcessor):
id=consumer.id, flow=consumer.flow
).observe(len(chunk.page_content))
await consumer.q.output.send(r)
await consumer.q["output"].send(r)
print("Done.", flush=True)

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

@ -28,8 +28,6 @@ class Processor(InputOutputProcessor):
super(Processor, self).__init__(
**params | {
"subscriber": subscriber,
"input_schema": TextDocument,
"output_schema": Chunk,
}
)
@ -47,6 +45,17 @@ class Processor(InputOutputProcessor):
chunk_overlap=chunk_overlap,
)
self.register_consumer(
name = "input",
schema = TextDocument,
handler = self.on_message,
)
self.register_producer(
name = "output",
schema = Chunk,
)
print("Chunker initialised", flush=True)
async def on_message(self, msg, consumer):
@ -71,7 +80,7 @@ class Processor(InputOutputProcessor):
id=consumer.id, flow=consumer.flow
).observe(len(chunk.page_content))
await consumer.q.output.send(r)
await consumer.q["output"].send(r)
print("Done.", flush=True)

21
trustgraph-flow/trustgraph/decoding/mistral_ocr/processor.py
View file

@ -17,12 +17,10 @@ from mistralai.models import OCRResponse
from ... schema import Document, TextDocument, Metadata
from ... schema import document_ingest_queue, text_ingest_queue
from ... log_level import LogLevel
from ... base import ConsumerProducer
from ... base import InputOutputProcessor
module = "ocr"
default_input_queue = document_ingest_queue
default_output_queue = text_ingest_queue
default_subscriber = module
default_api_key = os.getenv("MISTRAL_TOKEN")
@ -71,19 +69,17 @@ def get_combined_markdown(ocr_response: OCRResponse) -> str:
return "\n\n".join(markdowns)
class Processor(ConsumerProducer):
class Processor(InputOutputProcessor):
def __init__(self, **params):
input_queue = params.get("input_queue", default_input_queue)
output_queue = params.get("output_queue", default_output_queue)
id = params.get("id")
subscriber = params.get("subscriber", default_subscriber)
api_key = params.get("api_key", default_api_key)
super(Processor, self).__init__(
**params | {
"input_queue": input_queue,
"output_queue": output_queue,
"id": id,
"subscriber": subscriber,
"input_schema": Document,
"output_schema": TextDocument,
@ -151,7 +147,7 @@ class Processor(ConsumerProducer):
return markdown
async def handle(self, msg):
async def on_message(self, msg, consumer):
print("PDF message received")
@ -166,17 +162,14 @@ class Processor(ConsumerProducer):
text=markdown.encode("utf-8"),
)
await self.send(r)
await consumer.q.output.send(r)
print("Done.", flush=True)
@staticmethod
def add_args(parser):
ConsumerProducer.add_args(
parser, default_input_queue, default_subscriber,
default_output_queue,
)
InputOutputProcessor.add_args(parser, default_subscriber)
parser.add_argument(
'-k', '--api-key',

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

@ -22,17 +22,26 @@ class Processor(InputOutputProcessor):
def __init__(self, **params):
id = params.get("id")
subscriber = params.get("subscriber")
subscriber = params.get("subscriber", default_subscriber)
super(Processor, self).__init__(
**params | {
"id": id,
"subscriber": subscriber,
"input_schema": Document,
"output_schema": TextDocument,
}
)
self.register_consumer(
name = "input",
schema = Document,
handler = self.on_message,
)
self.register_producer(
name = "output",
schema = TextDocument,
)
print("PDF inited", flush=True)
async def on_message(self, msg, consumer):
@ -62,7 +71,7 @@ class Processor(InputOutputProcessor):
text=page.page_content.encode("utf-8"),
)
await consumer.q.output.send(r)
await consumer.q["output"].send(r)
print("Done.", flush=True)

9
trustgraph-flow/trustgraph/embeddings/document_embeddings/embeddings.py
View file

@ -11,20 +11,17 @@ from ... schema import document_embeddings_store_queue
from ... schema import embeddings_request_queue, embeddings_response_queue
from ... clients.embeddings_client import EmbeddingsClient
from ... log_level import LogLevel
from ... base import ConsumerProducer
from ... base import InputOutputProcessor
module = "document-embeddings"
default_input_queue = chunk_ingest_queue
default_output_queue = document_embeddings_store_queue
default_subscriber = module
class Processor(ConsumerProducer):
class Processor(InputOutputProcessor):
def __init__(self, **params):
input_queue = params.get("input_queue", default_input_queue)
output_queue = params.get("output_queue", default_output_queue)
id = params.get("id")
subscriber = params.get("subscriber", default_subscriber)
emb_request_queue = params.get(
"embeddings_request_queue", embeddings_request_queue