trustgraph-flow 1.2.6__tar.gz → 1.2.8__tar.gz

{trustgraph_flow-1.2.6 → trustgraph_flow-1.2.8}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: trustgraph-flow
-Version: 1.2.6
+Version: 1.2.8
 Summary: TrustGraph provides a means to run a pipeline of flexible AI processing components in a flexible means to achieve a processing pipeline.
 Author-email: "trustgraph.ai" <security@trustgraph.ai>
 Project-URL: Homepage, https://github.com/trustgraph-ai/trustgraph
@@ -34,7 +34,6 @@ Requires-Dist: prometheus-client
 Requires-Dist: pulsar-client
 Requires-Dist: pymilvus
 Requires-Dist: pypdf
-Requires-Dist: mistralai
 Requires-Dist: pyyaml
 Requires-Dist: qdrant-client
 Requires-Dist: rdflib

{trustgraph_flow-1.2.6 → trustgraph_flow-1.2.8}/pyproject.toml

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

{trustgraph_flow-1.2.6 → trustgraph_flow-1.2.8}/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 = "pdf-decoder"
 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__)
 

trustgraph_flow-1.2.8/trustgraph/flow_version.py

@@ -0,0 +1 @@
+__version__ = "1.2.8"

{trustgraph_flow-1.2.6 → trustgraph_flow-1.2.8}/trustgraph_flow.egg-info/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: trustgraph-flow
-Version: 1.2.6
+Version: 1.2.8
 Summary: TrustGraph provides a means to run a pipeline of flexible AI processing components in a flexible means to achieve a processing pipeline.
 Author-email: "trustgraph.ai" <security@trustgraph.ai>
 Project-URL: Homepage, https://github.com/trustgraph-ai/trustgraph
@@ -34,7 +34,6 @@ Requires-Dist: prometheus-client
 Requires-Dist: pulsar-client
 Requires-Dist: pymilvus
 Requires-Dist: pypdf
-Requires-Dist: mistralai
 Requires-Dist: pyyaml
 Requires-Dist: qdrant-client
 Requires-Dist: rdflib

{trustgraph_flow-1.2.6 → trustgraph_flow-1.2.8}/trustgraph_flow.egg-info/requires.txt

@@ -24,7 +24,6 @@ prometheus-client
 pulsar-client
 pymilvus
 pypdf
-mistralai
 pyyaml
 qdrant-client
 rdflib

trustgraph_flow-1.2.6/trustgraph/flow_version.py

@@ -1 +0,0 @@
-__version__ = "1.2.6"