apunkt/trustgraph
1
0
Fork 0
mirror of https://github.com/trustgraph-ai/trustgraph.git synced 2026-07-18 01:31:02 +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

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