From 2494b33f8e741eaf9ffd27a310af7a865235f2a6 Mon Sep 17 00:00:00 2001
From: Cyber MacGeddon <cybermaggedon@gmail.com>
Date: Mon, 4 Aug 2025 09:54:44 +0100
Subject: [PATCH] Fixed Mistral OCR to use current API

---
 trustgraph-flow/pyproject.toml                |  1 -
 .../decoding/mistral_ocr/processor.py         | 41 +++++++++++--------
 2 files changed, 24 insertions(+), 18 deletions(-)

diff --git a/trustgraph-flow/pyproject.toml b/trustgraph-flow/pyproject.toml
index c7eef10b..911c91a0 100644
--- a/trustgraph-flow/pyproject.toml
+++ b/trustgraph-flow/pyproject.toml
@@ -36,7 +36,6 @@ dependencies = [
     "pulsar-client",
     "pymilvus",
     "pypdf",
-    "mistralai",
     "pyyaml",
     "qdrant-client",
     "rdflib",
diff --git a/trustgraph-flow/trustgraph/decoding/mistral_ocr/processor.py b/trustgraph-flow/trustgraph/decoding/mistral_ocr/processor.py
index 4bacd278..9532fa0f 100755
--- a/trustgraph-flow/trustgraph/decoding/mistral_ocr/processor.py
+++ b/trustgraph-flow/trustgraph/decoding/mistral_ocr/processor.py
@@ -15,17 +15,13 @@ from mistralai import DocumentURLChunk, ImageURLChunk, TextChunk
 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 InputOutputProcessor
+from ... base import FlowProcessor, ConsumerSpec, ProducerSpec
 
 import logging
 
 logger = logging.getLogger(__name__)
 
-module = "ocr"
-
-default_subscriber = module
+default_ident = "mistral-ocr"
 default_api_key = os.getenv("MISTRAL_TOKEN")
 
 pages_per_chunk = 5
@@ -73,23 +69,34 @@ def get_combined_markdown(ocr_response: OCRResponse) -> str:
 
     return "\n\n".join(markdowns)
 
-class Processor(InputOutputProcessor):
+class Processor(FlowProcessor):
 
     def __init__(self, **params):
 
-        id = params.get("id")
-        subscriber = params.get("subscriber", default_subscriber)
+        id = params.get("id", default_ident)
         api_key = params.get("api_key", default_api_key)
 
         super(Processor, self).__init__(
             **params | {
                 "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:
             raise RuntimeError("Mistral API key not specified")
 
@@ -98,7 +105,7 @@ class Processor(InputOutputProcessor):
         # Used with Mistral doc upload
         self.unique_id = str(uuid.uuid4())
 
-        logger.info("PDF inited")
+        logger.info("Mistral OCR processor initialized")
 
     def ocr(self, blob):
 
@@ -151,7 +158,7 @@ class Processor(InputOutputProcessor):
 
             return markdown
 
-    async def on_message(self, msg, consumer):
+    async def on_message(self, msg, consumer, flow):
 
         logger.debug("PDF message received")
 
@@ -166,14 +173,14 @@ class Processor(InputOutputProcessor):
             text=markdown.encode("utf-8"),
         )
 
-        await consumer.q.output.send(r)
+        await flow("output").send(r)
 
         logger.info("Done.")
 
     @staticmethod
     def add_args(parser):
 
-        InputOutputProcessor.add_args(parser, default_subscriber)
+        FlowProcessor.add_args(parser)
 
         parser.add_argument(
             '-k', '--api-key',
@@ -183,5 +190,5 @@ class Processor(InputOutputProcessor):
 
 def run():
 
-    Processor.launch(module, __doc__)
+    Processor.launch(default_ident, __doc__)