OCR-Space Function.py

"""
title: OCR.Space Text Extraction
author: cloph-dsp
version: 1.0
description: Extract text from images in chat messages using OCR.Space API. 
             Get your key here -> https://ocr.space/OCRAPI
"""

from pydantic import BaseModel, Field
from typing import Callable, Awaitable, Any, Optional
import asyncio
import base64
import json
import requests
import re


class Filter:
    class Valves(BaseModel):
        priority: int = Field(
            default=0, description="Priority level for filtering operations."
        )
        API_KEY: str = Field(
            default="", description="OCR.Space API Key"
        )
        MAX_RETRIES: int = Field(default=3, description="Maximum retry attempts for OCR processing")
        LANGUAGE: str = Field(
            default="eng", 
            description="OCR language code (e.g., eng, ger, spa, fra, ita, por, etc.)"
        )

    def __init__(self):
        self.valves = self.Valves()

    def _extract_base64_content(self, image_url: str) -> str:
        """Extract base64 content from data URL."""
        # Check if it's a data URL (base64)
        match = re.match(r'data:image/([a-zA-Z]+);base64,(.+)', image_url)
        if match:
            return match.group(2)
        return image_url

    async def _process_image_ocr(self, image_url: str, event_emitter) -> str:
        """Process image through OCR.Space API with retry mechanism."""
        retries = 0
        while retries < self.valves.MAX_RETRIES:
            try:
                # Extract base64 content if it's a data URL
                image_data = self._extract_base64_content(image_url)
                
                # Determine if it's a URL or base64 data
                if image_data == image_url:  # It's a URL
                    ocr_result = requests.post(
                        'https://api.ocr.space/parse/image',
                        data={
                            'url': image_url,
                            'apikey': self.valves.API_KEY,
                            'language': self.valves.LANGUAGE,
                            'isOverlayRequired': False
                        },
                    )
                else:  # It's base64 data
                    ocr_result = requests.post(
                        'https://api.ocr.space/parse/image',
                        data={
                            'base64Image': f'data:image/jpeg;base64,{image_data}',
                            'apikey': self.valves.API_KEY,
                            'language': self.valves.LANGUAGE,
                            'isOverlayRequired': False
                        },
                    )
                
                result = json.loads(ocr_result.content.decode())
                
                if result.get("OCRExitCode") == 1:  # Success
                    extracted_text = ""
                    for text_result in result.get("ParsedResults", []):
                        extracted_text += text_result.get("ParsedText", "")
                    
                    if not extracted_text:
                        raise Exception("No text extracted from the image")
                    
                    return extracted_text
                else:
                    error_message = result.get("ErrorMessage", "Unknown OCR error")
                    raise Exception(f"OCR error: {error_message}")
                
            except Exception as e:
                retries += 1
                if retries < self.valves.MAX_RETRIES:
                    await event_emitter(
                        {
                            "type": "status",
                            "data": {
                                "description": f"⚠️ OCR processing failed, retrying ({retries}/{self.valves.MAX_RETRIES})...",
                                "done": False,
                            },
                        }
                    )
                    await asyncio.sleep(2**retries)  # Exponential backoff
                else:
                    await event_emitter(
                        {
                            "type": "status",
                            "data": {
                                "description": f"❌ OCR failed after {self.valves.MAX_RETRIES} attempts: {str(e)}",
                                "done": True,
                            },
                        }
                    )
                    raise
        
        raise Exception(f"Failed to process image after {self.valves.MAX_RETRIES} attempts")

    def _find_image_in_messages(self, messages):
        """Find images in the user messages."""
        for m_index, message in enumerate(messages):
            if message["role"] == "user" and isinstance(message.get("content"), list):
                for c_index, content in enumerate(message["content"]):
                    if content["type"] == "image_url":
                        return m_index, c_index, content["image_url"]["url"]
        return None

    async def inlet(
        self,
        body: dict,
        __event_emitter__: Callable[[Any], Awaitable[None]],
        __user__: Optional[dict] = None,
        __model__: Optional[dict] = None,
    ) -> dict:
        messages = body.get("messages", [])
        
        # Check for API key
        if not self.valves.API_KEY or self.valves.API_KEY == "helloworld":
            # Add a system message informing about missing API key
            await __event_emitter__(
                {
                    "type": "status",
                    "data": {
                        "description": "❌ OCR feature disabled: No valid API key configured",
                        "done": True,
                    },
                }
            )
            # Add a message to the conversation
            if messages and messages[-1]["role"] == "user":
                content = messages[-1].get("content", "")
                if isinstance(content, list):
                    # Add a text note about the missing API key to the existing content
                    for item in content:
                        if item["type"] == "image_url":
                            messages[-1]["content"] = [
                                {"type": "text", "text": "I notice you've uploaded an image, but OCR processing is not available (missing API key)."}
                            ]
            return body
        
        # Look for images in messages
        image_info = self._find_image_in_messages(messages)
        if not image_info:
            return body
        
        message_index, content_index, image_url = image_info
        
        try:
            # Show processing status
            await __event_emitter__(
                {
                    "type": "status",
                    "data": {
                        "description": "🔍 Processing image with OCR.Space...",
                        "done": False,
                    },
                }
            )
            
            # Extract text from image
            extracted_text = await self._process_image_ocr(image_url, __event_emitter__)
            
            # Update status to complete
            await __event_emitter__(
                {
                    "type": "status",
                    "data": {
                        "description": "✅ Text successfully extracted from image",
                        "done": True,
                    },
                }
            )
            
            # Format the extracted text with context
            ocr_context = (
                f"[System note: The following text was extracted from an uploaded image via OCR. "
                f"Use it to inform your response.]\n\n{extracted_text}"
            )
            
            # Update the message content
            user_message = ""
            if isinstance(messages[message_index]["content"], list):
                for item in messages[message_index]["content"]:
                    if item["type"] == "text":
                        user_message = item["text"]
                        break
            
            # Replace the image with the extracted text
            messages[message_index]["content"] = [
                {
                    "type": "text",
                    "text": f"{ocr_context}\n\nUser query: {user_message}" if user_message else ocr_context
                }
            ]
            
            body["messages"] = messages
            
        except Exception as e:
            await __event_emitter__(
                {
                    "type": "status",
                    "data": {
                        "description": f"❌ OCR processing failed: {str(e)}",
                        "done": True,
                    },
                }
            )
            
            # Replace the image with an error message
            if isinstance(messages[message_index]["content"], list):
                user_message = ""
                for item in messages[message_index]["content"]:
                    if item["type"] == "text":
                        user_message = item["text"]
                        break
                
                messages[message_index]["content"] = [
                    {
                        "type": "text",
                        "text": f"[Image processing failed: {str(e)}]\n\n{user_message}" if user_message else f"[Image processing failed: {str(e)}]"
                    }
                ]
                
                body["messages"] = messages
        
        return body

    async def stream(self, event: dict) -> dict:
        # No modifications needed for streaming events
        return event

    async def outlet(
        self,
        body: dict,
        __event_emitter__: Callable[[Any], Awaitable[None]],
        __user__: Optional[dict] = None,
        __model__: Optional[dict] = None,
    ) -> dict:
        # No modifications needed for the response
        return body