apunkt/trustgraph
1
0
Fork 0
mirror of https://github.com/trustgraph-ai/trustgraph.git synced 2026-04-27 17:36:23 +02:00
Code Issues Projects Releases Packages Wiki Activity Actions

Fixed Mistral OCR to use current API (#448)

Browse source 
* Fixed Mistral OCR to use current API

* Added PDF decoder tests
This commit is contained in:
cybermaggedon 2025-08-04 10:08:16 +01:00 • committed by GitHub
parent 8f0828c9a6
commit 7e0d831026
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
5 changed files with 549 additions and 18 deletions
Download patch file Download diff file Expand all files Collapse all files

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

@ -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__)