feat: mistral ocr instrumentation

Author: CakeCrusher

python/instrumentation/openinference-instrumentation-mistralai/README.md

@@ -76,6 +76,74 @@ Now simply run the python file and observe the traces in Phoenix.
 python your_file.py
 ```
 
+## OCR and Input Image Tracing
+
+The MistralAI instrumentation automatically traces input images and documents passed to the OCR API, following OpenInference semantic conventions. This includes:
+
+### Supported Input Types
+
+- **HTTP Image URLs**: `https://example.com/image.jpg`
+- **Base64 Images**: `data:image/jpeg;base64,{base64_data}`  
+- **PDF URLs**: `https://example.com/document.pdf`
+- **Base64 PDFs**: `data:application/pdf;base64,{base64_data}`
+
+### Trace Attributes
+
+For **image inputs**, the instrumentation creates:
+- `input.message_content.type`: `"image"`
+- `input.message_content.image.image.url`: The image URL or base64 data URL
+- `input.message_content.image.metadata`: JSON metadata including source, encoding type, and MIME type
+
+For **document inputs**, the instrumentation creates:
+- `input.message_content.type`: `"document"`  
+- `input.document.url`: The document URL or base64 data URL
+- `input.document.metadata`: JSON metadata including source, encoding type, and MIME type
+
+### Example Usage
+
+```python
+import base64
+import os
+from mistralai import Mistral
+from openinference.instrumentation.mistralai import MistralAIInstrumentor
+
+# Set up instrumentation
+MistralAIInstrumentor().instrument()
+
+client = Mistral(api_key=os.environ["MISTRAL_API_KEY"])
+
+# OCR with HTTP image URL
+response = client.ocr.process(
+    model="mistral-ocr-latest",
+    document={
+        "type": "image_url",
+        "image_url": "https://example.com/receipt.png"
+    },
+    include_image_base64=True
+)
+
+# OCR with base64 image  
+with open("image.jpg", "rb") as f:
+    base64_image = base64.b64encode(f.read()).decode('utf-8')
+
+response = client.ocr.process(
+    model="mistral-ocr-latest", 
+    document={
+        "type": "image_url",
+        "image_url": f"data:image/jpeg;base64,{base64_image}"
+    },
+    include_image_base64=True
+)
+```
+
+### Privacy and Configuration
+
+Input image tracing works seamlessly with [TraceConfig](https://github.com/Arize-ai/openinference/tree/main/python/openinference-instrumentation#tracing-configuration) for:
+
+- **Image size limits**: Control maximum base64 image length with `base64_image_max_length`
+- **Privacy controls**: Hide input images with `hide_inputs` or `hide_input_images`
+- **MIME type detection**: Automatic detection and proper formatting of image data URLs
+
 ## More Info
 
 * [More info on OpenInference and Phoenix](https://docs.arize.com/phoenix)

python/instrumentation/openinference-instrumentation-mistralai/examples/ocr_test.py

@@ -0,0 +1,72 @@
+#!/usr/bin/env python3
+
+import os
+
+from dotenv import load_dotenv
+from mistralai import Mistral
+from phoenix.otel import register
+
+from openinference.instrumentation.mistralai import MistralAIInstrumentor
+
+load_dotenv()
+tracer = register(
+    project_name="mistral-ocr",
+    endpoint=os.getenv("PHOENIX_COLLECTOR_ENDPOINT"),
+)
+
+# Initialize instrumentation
+MistralAIInstrumentor().instrument(tracer_provider=tracer)
+
+# Initialize client
+client = Mistral(api_key=os.environ.get("MISTRAL_API_KEY"))
+
+
+def test_ocr_with_working_image():
+    """Test OCR with a working image URL that should display in Phoenix"""
+
+    # Using a reliable image URL - this is a simple diagram/chart
+    image_url = "https://upload.wikimedia.org/wikipedia/commons/d/d1/Ai_lizard.png"
+    try:
+        print("🔍 Testing OCR with working image URL...")
+        print(f"Image URL: {image_url}")
+
+        ocr_response = client.ocr.process(
+            model="mistral-ocr-latest",
+            document={
+                "type": "image_url",
+                "image_url": image_url,
+            },
+            include_image_base64=True,
+        )
+
+        print("✅ OCR completed successfully!")
+
+        if hasattr(ocr_response, "pages") and ocr_response.pages:
+            print(f"📄 Pages processed: {len(ocr_response.pages)}")
+
+            for i, page in enumerate(ocr_response.pages):
+                print(f"\n--- Page {i + 1} ---")
+                if hasattr(page, "markdown") and page.markdown:
+                    print("📝 Markdown content:")
+                    print(
+                        page.markdown[:200] + "..." if len(page.markdown) > 200 else page.markdown
+                    )
+
+                if hasattr(page, "images") and page.images:
+                    print(f"🖼️  Extracted images: {len(page.images)}")
+                    for j, img in enumerate(page.images):
+                        if hasattr(img, "id"):
+                            print(f"  - Image {j + 1}: {img.id}")
+
+        print("\n🔗 View traces in your Phoenix project")
+
+    except Exception as e:
+        print(f"❌ Error: {e}")
+        print("Make sure:")
+        print("1. MISTRAL_API_KEY environment variable is set")
+        print("2. You have sufficient credits")
+        print("3. The image URL is accessible")
+
+
+if __name__ == "__main__":
+    test_ocr_with_working_image()

python/instrumentation/openinference-instrumentation-mistralai/src/openinference/instrumentation/mistralai/__init__.py

@@ -14,6 +14,10 @@
     _AsyncStreamChatWrapper,
     _SyncChatWrapper,
 )
+from openinference.instrumentation.mistralai._ocr_wrapper import (
+    _AsyncOCRWrapper,
+    _SyncOCRWrapper,
+)
 from openinference.instrumentation.mistralai.package import _instruments
 from openinference.instrumentation.mistralai.version import __version__
 
@@ -38,6 +42,8 @@ class MistralAIInstrumentor(BaseInstrumentor):  # type: ignore
         "_original_sync_stream_agent_method",
         "_original_async_agent_method",
         "_original_async_stream_agent_method",
+        "_original_sync_ocr_method",
+        "_original_async_ocr_method",
     )
 
     def instrumentation_dependencies(self) -> Collection[str]:
@@ -59,6 +65,12 @@ def _instrument(self, **kwargs: Any) -> None:
             import mistralai
             from mistralai.agents import Agents
             from mistralai.chat import Chat
+
+            Ocr = None
+            try:
+                from mistralai.ocr import Ocr
+            except ImportError:
+                print("Outdated version of mistralai: currently version does not support Ocr")
         except ImportError as err:
             raise Exception(
                 "Could not import mistralai. Please install with `pip install mistralai`."
@@ -72,6 +84,10 @@ def _instrument(self, **kwargs: Any) -> None:
         self._original_sync_stream_agent_method = Agents.stream
         self._original_async_agent_method = Agents.complete_async
         self._original_async_stream_agent_method = Agents.stream_async
+        if Ocr is not None:
+            self._original_sync_ocr_method = Ocr.process
+            self._original_async_ocr_method = Ocr.process_async
+
         wrap_function_wrapper(
             module="mistralai.chat",
             name="Chat.complete",
@@ -120,6 +136,20 @@ def _instrument(self, **kwargs: Any) -> None:
             wrapper=_AsyncStreamChatWrapper("MistralAsyncClient.agents", self._tracer, mistralai),
         )
 
+        # Instrument OCR methods
+        if Ocr is not None:
+            wrap_function_wrapper(
+                module="mistralai.ocr",
+                name="Ocr.process",
+                wrapper=_SyncOCRWrapper("MistralClient.ocr", self._tracer, mistralai),
+            )
+
+            wrap_function_wrapper(
+                module="mistralai.ocr",
+                name="Ocr.process_async",
+                wrapper=_AsyncOCRWrapper("MistralAsyncClient.ocr", self._tracer, mistralai),
+            )
+
     def _uninstrument(self, **kwargs: Any) -> None:
         from mistralai.agents import Agents
         from mistralai.chat import Chat
@@ -132,3 +162,11 @@ def _uninstrument(self, **kwargs: Any) -> None:
         Agents.stream = self._original_sync_stream_agent_method  # type: ignore
         Agents.complete_async = self._original_async_agent_method  # type: ignore
         Agents.stream_async = self._original_async_stream_agent_method  # type: ignore
+        try:
+            from mistralai.ocr import Ocr
+
+            Ocr.process = self._original_sync_ocr_method  # type: ignore
+            Ocr.process_async = self._original_async_ocr_method  # type: ignore
+        except ImportError:
+            # OCR module not available, nothing to uninstrument
+            pass