apunkt/trustgraph
1
0
Fork 0
mirror of https://github.com/trustgraph-ai/trustgraph.git synced 2026-07-24 20:51:02 +02:00
Code Issues Projects Releases Packages Wiki Activity Actions 2

Fixed Mistral OCR to use current API

Browse source
This commit is contained in:
Cyber MacGeddon 2025-08-04 09:54:44 +01:00
parent 81c7c1181b
commit 0074d822b9
1 changed files with 24 additions and 16 deletions
Download patch file Download diff file Expand all files Collapse all files

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

@ -15,13 +15,9 @@ from mistralai import DocumentURLChunk, ImageURLChunk, TextChunk
from mistralai.models import OCRResponse from mistralai.models import OCRResponse
from ... schema import Document, TextDocument, Metadata from ... schema import Document, TextDocument, Metadata
from ... schema import document_ingest_queue, text_ingest_queue from ... base import FlowProcessor, ConsumerSpec, ProducerSpec
from ... log_level import LogLevel
from ... base import InputOutputProcessor
module = "ocr" default_ident = "mistral-ocr"
default_subscriber = module
default_api_key = os.getenv("MISTRAL_TOKEN") default_api_key = os.getenv("MISTRAL_TOKEN")
pages_per_chunk = 5 pages_per_chunk = 5
@ -69,23 +65,34 @@ def get_combined_markdown(ocr_response: OCRResponse) -> str:
return "\n\n".join(markdowns) return "\n\n".join(markdowns)
class Processor(InputOutputProcessor): class Processor(FlowProcessor):
def __init__(self, **params): def __init__(self, **params):
id = params.get("id") id = params.get("id", default_ident)
subscriber = params.get("subscriber", default_subscriber)
api_key = params.get("api_key", default_api_key) api_key = params.get("api_key", default_api_key)
super(Processor, self).__init__( super(Processor, self).__init__(
**params | { **params | {
"id": id, "id": id,
"subscriber": subscriber,
"input_schema": Document,
"output_schema": TextDocument,
} }
) )
self.register_specification(
ConsumerSpec(
name = "input",
schema = Document,
handler = self.on_message,
)
)
self.register_specification(
ProducerSpec(
name = "output",
schema = TextDocument,
)
)
if api_key is None: if api_key is None:
raise RuntimeError("Mistral API key not specified") raise RuntimeError("Mistral API key not specified")
@ -95,6 +102,7 @@ class Processor(InputOutputProcessor):
self.unique_id = str(uuid.uuid4()) self.unique_id = str(uuid.uuid4())
print("PDF inited") print("PDF inited")
logger.info("Mistral OCR processor initialized")
def ocr(self, blob): def ocr(self, blob):
@ -147,7 +155,7 @@ class Processor(InputOutputProcessor):
return markdown return markdown
async def on_message(self, msg, consumer): async def on_message(self, msg, consumer, flow):
print("PDF message received") print("PDF message received")
@ -162,14 +170,14 @@ class Processor(InputOutputProcessor):
text=markdown.encode("utf-8"), text=markdown.encode("utf-8"),
) )
await consumer.q.output.send(r) await flow("output").send(r)
print("Done.", flush=True) print("Done.", flush=True)
@staticmethod @staticmethod
def add_args(parser): def add_args(parser):
InputOutputProcessor.add_args(parser, default_subscriber) FlowProcessor.add_args(parser)
parser.add_argument( parser.add_argument(
'-k', '--api-key', '-k', '--api-key',
@ -179,5 +187,5 @@ class Processor(InputOutputProcessor):
def run(): def run():
Processor.launch(module, __doc__) Processor.launch(default_ident, __doc__)