apunkt/trustgraph
1
0
Fork 0
mirror of https://github.com/trustgraph-ai/trustgraph.git synced 2026-07-24 04:31: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 8f0828c9a6
commit 2494b33f8e
2 changed files with 24 additions and 18 deletions
Download patch file Download diff file Expand all files Collapse all files

1
trustgraph-flow/pyproject.toml
View file

@ -36,7 +36,6 @@ dependencies = [
"pulsar-client", "pulsar-client",
"pymilvus", "pymilvus",
"pypdf", "pypdf",
"mistralai",
"pyyaml", "pyyaml",
"qdrant-client", "qdrant-client",
"rdflib", "rdflib",

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 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
import logging import logging
logger = logging.getLogger(__name__) logger = logging.getLogger(__name__)
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
@ -73,23 +69,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")
@ -98,7 +105,7 @@ class Processor(InputOutputProcessor):
# Used with Mistral doc upload # Used with Mistral doc upload
self.unique_id = str(uuid.uuid4()) self.unique_id = str(uuid.uuid4())
logger.info("PDF inited") logger.info("Mistral OCR processor initialized")
def ocr(self, blob): def ocr(self, blob):
@ -151,7 +158,7 @@ class Processor(InputOutputProcessor):
return markdown return markdown
async def on_message(self, msg, consumer): async def on_message(self, msg, consumer, flow):
logger.debug("PDF message received") logger.debug("PDF message received")
@ -166,14 +173,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)
logger.info("Done.") logger.info("Done.")
@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',
@ -183,5 +190,5 @@ class Processor(InputOutputProcessor):
def run(): def run():
Processor.launch(module, __doc__) Processor.launch(default_ident, __doc__)