amd-gaia 0.15.0__py3-none-any.whl → 0.15.2__py3-none-any.whl

gaia/llm/vlm_client.py

@@ -1,382 +1,382 @@
-#!/usr/bin/env python3
-# Copyright(C) 2024-2025 Advanced Micro Devices, Inc. All rights reserved.
-# SPDX-License-Identifier: MIT
-
-"""
-Vision-Language Model (VLM) client for extracting text from images.
-
-Handles model loading/unloading and image-to-text extraction via Lemonade server.
-"""
-
-import base64
-import logging
-import os
-from typing import Optional
-
-from dotenv import load_dotenv
-
-# Load environment variables from .env file
-load_dotenv()
-
-# Default Lemonade server URL (can be overridden via LEMONADE_BASE_URL env var)
-DEFAULT_LEMONADE_URL = "http://localhost:8000/api/v1"
-
-logger = logging.getLogger(__name__)
-
-# Magic bytes for common image formats
-IMAGE_SIGNATURES = {
-    b"\x89PNG\r\n\x1a\n": "image/png",
-    b"\xff\xd8\xff": "image/jpeg",
-    b"GIF87a": "image/gif",
-    b"GIF89a": "image/gif",
-    b"RIFF": "image/webp",  # WebP starts with RIFF...WEBP
-    b"BM": "image/bmp",
-}
-
-
-def detect_image_mime_type(image_bytes: bytes) -> str:
-    """
-    Detect MIME type from image bytes using magic number signatures.
-
-    Args:
-        image_bytes: Raw image bytes
-
-    Returns:
-        MIME type string (e.g., "image/jpeg", "image/png")
-        Defaults to "image/png" if format not detected.
-    """
-    for signature, mime_type in IMAGE_SIGNATURES.items():
-        if image_bytes.startswith(signature):
-            # Special case: WebP needs additional check for WEBP marker
-            if signature == b"RIFF" and len(image_bytes) >= 12:
-                if image_bytes[8:12] != b"WEBP":
-                    continue
-            return mime_type
-
-    # Default to PNG if format not detected
-    logger.debug("Could not detect image format, defaulting to image/png")
-    return "image/png"
-
-
-class VLMClient:
-    """
-    VLM client for extracting text from images using Lemonade server.
-
-    Handles:
-    - Model loading (default: Qwen3-VL-4B-Instruct-GGUF)
-    - Image-to-markdown conversion
-    - State tracking for VLM processing
-    """
-
-    def __init__(
-        self,
-        vlm_model: str = "Qwen3-VL-4B-Instruct-GGUF",
-        base_url: Optional[str] = None,
-        auto_load: bool = True,
-    ):
-        """
-        Initialize VLM client.
-
-        Args:
-            vlm_model: Vision model to use for image extraction
-            base_url: Lemonade server API URL (defaults to LEMONADE_BASE_URL env var)
-            auto_load: Automatically load VLM model on first use
-        """
-        # Use provided base_url, fall back to env var, then default
-        if base_url is None:
-            base_url = os.getenv("LEMONADE_BASE_URL", DEFAULT_LEMONADE_URL)
-        from urllib.parse import urlparse
-
-        from gaia.llm.lemonade_client import LemonadeClient
-
-        self.vlm_model = vlm_model
-        self.base_url = base_url
-
-        # Parse base_url to extract host and port for LemonadeClient
-        parsed = urlparse(base_url)
-        host = parsed.hostname or "localhost"
-        port = parsed.port or 8000
-
-        # Get base server URL (without /api/v1) for user-facing messages
-        self.server_url = f"http://{host}:{port}"
-
-        self.client = LemonadeClient(model=vlm_model, host=host, port=port)
-        self.auto_load = auto_load
-        self.vlm_loaded = False
-
-        logger.debug(f"VLM Client initialized: {self.vlm_model} at {self.server_url}")
-
-    def check_availability(self) -> bool:
-        """
-        Check if VLM model is available on Lemonade server.
-
-        Returns:
-            True if model is available, False otherwise
-        """
-        try:
-            models_response = self.client.list_models()
-            available_models = [
-                m.get("id", "") for m in models_response.get("data", [])
-            ]
-
-            if self.vlm_model in available_models:
-                logger.debug(f"VLM model available: {self.vlm_model}")
-                return True
-            else:
-                logger.warning(f"❌ VLM model not found: {self.vlm_model}")
-                logger.warning("")
-                logger.warning("📥 To download this model:")
-                logger.warning(f"   1. Open Lemonade Model Manager ({self.server_url})")
-                logger.warning(f"   2. Search for: {self.vlm_model}")
-                logger.warning("   3. Click 'Download' to install the model")
-                logger.warning("")
-                logger.warning(
-                    f"   Available models: {', '.join(available_models[:3])}..."
-                )
-                return False
-
-        except Exception as e:
-            logger.error(f"Failed to check VLM availability: {e}")
-            logger.error(
-                f"   Make sure Lemonade server is running at {self.server_url}"
-            )
-            return False
-
-    def _ensure_vlm_loaded(self) -> bool:
-        """
-        Ensure VLM model is loaded, load it if necessary.
-
-        The model will be automatically downloaded if not available (handled by
-        lemonade_client.chat_completions with auto_download=True).
-
-        Returns:
-            True if VLM is loaded, False if loading failed
-        """
-        if self.vlm_loaded:
-            return True
-
-        if not self.auto_load:
-            logger.warning("VLM not loaded and auto_load=False")
-            return False
-
-        try:
-            logger.debug(f"Loading VLM model: {self.vlm_model}")
-            # Load model (auto-download handled by lemonade_client, may take hours)
-            self.client.load_model(self.vlm_model, timeout=60, auto_download=True)
-            self.vlm_loaded = True
-            logger.debug(f"VLM model loaded: {self.vlm_model}")
-            return True
-
-        except Exception as e:
-            logger.error(f"Failed to load VLM model: {e}")
-            logger.error(
-                f"   Make sure Lemonade server is running at {self.server_url}"
-            )
-            return False
-
-    def extract_from_image(
-        self,
-        image_bytes: bytes,
-        image_num: int = 1,
-        page_num: int = 1,
-        prompt: Optional[str] = None,
-    ) -> str:
-        """
-        Extract text from an image using VLM.
-
-        Args:
-            image_bytes: Image as PNG/JPEG bytes
-            image_num: Image number on page (for logging)
-            page_num: Page number (for logging)
-            prompt: Custom extraction prompt (optional)
-
-        Returns:
-            Extracted text in markdown format
-        """
-        # Ensure VLM is loaded
-        if not self._ensure_vlm_loaded():
-            error_msg = "VLM model not available"
-            logger.error(error_msg)
-            return f"[VLM extraction failed: {error_msg}]"
-
-        # Encode image as base64 and detect MIME type
-        # Note: Image size optimization happens in pdf_utils.py during extraction
-        image_b64 = base64.b64encode(image_bytes).decode("utf-8")
-        mime_type = detect_image_mime_type(image_bytes)
-
-        # Default prompt for text extraction
-        if not prompt:
-            prompt = """You are an OCR system. Extract ALL visible text from this image exactly as it appears.
-
-Instructions:
-1. Extract EVERY word you see - don't skip or paraphrase
-2. Preserve exact formatting (headings, bold, bullets, tables)
-3. If it's a table, format as markdown table
-4. If it's a chart, describe what you see: [CHART: ...]
-5. Do NOT add placeholders like "[Insert ...]"  - only extract actual text
-6. Do NOT generate or invent content - only extract what you see
-
-Output format: Clean markdown with the ACTUAL text from the image."""
-
-        # Format message with image (OpenAI vision format)
-        messages = [
-            {
-                "role": "user",
-                "content": [
-                    {"type": "text", "text": prompt},
-                    {
-                        "type": "image_url",
-                        "image_url": {"url": f"data:{mime_type};base64,{image_b64}"},
-                    },
-                ],
-            }
-        ]
-
-        try:
-            import time
-
-            start_time = time.time()
-
-            logger.debug(
-                f"VLM extracting from image {image_num} on page {page_num} ({mime_type})..."
-            )
-            logger.debug(
-                f"   Image: {mime_type}, {len(image_b64)} chars base64 ({len(image_bytes)} bytes raw)"
-            )
-
-            # Call VLM using chat completions endpoint
-            response = self.client.chat_completions(
-                model=self.vlm_model,
-                messages=messages,
-                temperature=0.1,  # Low temp for accurate extraction
-                max_completion_tokens=2048,  # Allow detailed extraction
-                timeout=300,  # VLM needs more time for complex forms (5 min)
-            )
-
-            elapsed = time.time() - start_time
-
-            # Extract text from response
-            if (
-                isinstance(response, dict)
-                and "choices" in response
-                and len(response["choices"]) > 0
-            ):
-                extracted_text = response["choices"][0]["message"]["content"]
-                size_kb = len(image_bytes) / 1024
-                logger.debug(
-                    f"Extracted {len(extracted_text)} chars from image {image_num} "
-                    f"in {elapsed:.2f}s ({size_kb:.0f}KB image)"
-                )
-                return extracted_text
-            else:
-                # Check for specific error types and provide helpful messages
-                error_msg = self._parse_vlm_error(response)
-                logger.error(error_msg)
-                return f"[VLM extraction failed: {error_msg}]"
-
-        except Exception as e:
-            logger.error(
-                f"VLM extraction failed for page {page_num}, image {image_num}: {e}"
-            )
-            import traceback
-
-            logger.debug(traceback.format_exc())
-            return f"[VLM extraction failed: {str(e)}]"
-
-    def _parse_vlm_error(self, response: dict) -> str:
-        """Parse VLM error response and return a helpful error message."""
-        if not isinstance(response, dict):
-            return f"Unexpected response type: {type(response)}"
-
-        # Check for nested error structure from Lemonade
-        error = response.get("error", {})
-        if isinstance(error, dict):
-            details = error.get("details", {})
-            inner_response = (
-                details.get("response", {}) if isinstance(details, dict) else {}
-            )
-            inner_error = (
-                inner_response.get("error", {})
-                if isinstance(inner_response, dict)
-                else {}
-            )
-
-            # Context size error
-            if inner_error.get("type") == "exceed_context_size_error":
-                n_ctx = inner_error.get("n_ctx", "unknown")
-                n_prompt = inner_error.get("n_prompt_tokens", "unknown")
-                return (
-                    f"Context size too small! Image requires {n_prompt} tokens "
-                    f"but model context is only {n_ctx}. "
-                    f"To fix: Right-click Lemonade tray icon → Settings → "
-                    f"set Context Size to 32768, then restart the model."
-                )
-
-            # Other backend errors
-            if error.get("type") == "backend_error":
-                msg = inner_error.get(
-                    "message", error.get("message", "Unknown backend error")
-                )
-                return f"Backend error: {msg}"
-
-        return f"Unexpected response format: {response}"
-
-    def extract_from_page_images(self, images: list, page_num: int) -> list:
-        """
-        Extract text from multiple images on a page.
-
-        Args:
-            images: List of image dicts with 'image_bytes', 'width', 'height', etc.
-            page_num: Page number
-
-        Returns:
-            List of dicts:
-            [
-                {
-                    "image_num": 1,
-                    "text": "extracted markdown",
-                    "dimensions": "800x600",
-                    "size_kb": 45.2
-                },
-                ...
-            ]
-        """
-        results = []
-
-        for img_idx, img_data in enumerate(images, 1):
-            extracted_text = self.extract_from_image(
-                image_bytes=img_data["image_bytes"],
-                image_num=img_idx,
-                page_num=page_num,
-            )
-
-            results.append(
-                {
-                    "image_num": img_idx,
-                    "text": extracted_text,
-                    "dimensions": f"{img_data['width']}x{img_data['height']}",
-                    "size_kb": img_data["size_kb"],
-                }
-            )
-
-        return results
-
-    def cleanup(self):
-        """
-        Cleanup VLM resources.
-
-        Call this after batch processing to mark VLM as unloaded.
-        Note: Model remains loaded on server; this just updates local state.
-        """
-        if self.vlm_loaded:
-            logger.info("🧹 VLM processing complete")
-            self.vlm_loaded = False
-
-    def __enter__(self):
-        """Context manager entry - ensure VLM loaded."""
-        self._ensure_vlm_loaded()
-        return self
-
-    def __exit__(self, exc_type, exc_val, exc_tb):
-        """Context manager exit - cleanup VLM state."""
-        self.cleanup()
+#!/usr/bin/env python3
+# Copyright(C) 2025-2026 Advanced Micro Devices, Inc. All rights reserved.
+# SPDX-License-Identifier: MIT
+
+"""
+Vision-Language Model (VLM) client for extracting text from images.
+
+Handles model loading/unloading and image-to-text extraction via Lemonade server.
+"""
+
+import base64
+import logging
+import os
+from typing import Optional
+
+from dotenv import load_dotenv
+
+# Load environment variables from .env file
+load_dotenv()
+
+# Default Lemonade server URL (can be overridden via LEMONADE_BASE_URL env var)
+DEFAULT_LEMONADE_URL = "http://localhost:8000/api/v1"
+
+logger = logging.getLogger(__name__)
+
+# Magic bytes for common image formats
+IMAGE_SIGNATURES = {
+    b"\x89PNG\r\n\x1a\n": "image/png",
+    b"\xff\xd8\xff": "image/jpeg",
+    b"GIF87a": "image/gif",
+    b"GIF89a": "image/gif",
+    b"RIFF": "image/webp",  # WebP starts with RIFF...WEBP
+    b"BM": "image/bmp",
+}
+
+
+def detect_image_mime_type(image_bytes: bytes) -> str:
+    """
+    Detect MIME type from image bytes using magic number signatures.
+
+    Args:
+        image_bytes: Raw image bytes
+
+    Returns:
+        MIME type string (e.g., "image/jpeg", "image/png")
+        Defaults to "image/png" if format not detected.
+    """
+    for signature, mime_type in IMAGE_SIGNATURES.items():
+        if image_bytes.startswith(signature):
+            # Special case: WebP needs additional check for WEBP marker
+            if signature == b"RIFF" and len(image_bytes) >= 12:
+                if image_bytes[8:12] != b"WEBP":
+                    continue
+            return mime_type
+
+    # Default to PNG if format not detected
+    logger.debug("Could not detect image format, defaulting to image/png")
+    return "image/png"
+
+
+class VLMClient:
+    """
+    VLM client for extracting text from images using Lemonade server.
+
+    Handles:
+    - Model loading (default: Qwen3-VL-4B-Instruct-GGUF)
+    - Image-to-markdown conversion
+    - State tracking for VLM processing
+    """
+
+    def __init__(
+        self,
+        vlm_model: str = "Qwen3-VL-4B-Instruct-GGUF",
+        base_url: Optional[str] = None,
+        auto_load: bool = True,
+    ):
+        """
+        Initialize VLM client.
+
+        Args:
+            vlm_model: Vision model to use for image extraction
+            base_url: Lemonade server API URL (defaults to LEMONADE_BASE_URL env var)
+            auto_load: Automatically load VLM model on first use
+        """
+        # Use provided base_url, fall back to env var, then default
+        if base_url is None:
+            base_url = os.getenv("LEMONADE_BASE_URL", DEFAULT_LEMONADE_URL)
+        from urllib.parse import urlparse
+
+        from gaia.llm.lemonade_client import LemonadeClient
+
+        self.vlm_model = vlm_model
+        self.base_url = base_url
+
+        # Parse base_url to extract host and port for LemonadeClient
+        parsed = urlparse(base_url)
+        host = parsed.hostname or "localhost"
+        port = parsed.port or 8000
+
+        # Get base server URL (without /api/v1) for user-facing messages
+        self.server_url = f"http://{host}:{port}"
+
+        self.client = LemonadeClient(model=vlm_model, host=host, port=port)
+        self.auto_load = auto_load
+        self.vlm_loaded = False
+
+        logger.debug(f"VLM Client initialized: {self.vlm_model} at {self.server_url}")
+
+    def check_availability(self) -> bool:
+        """
+        Check if VLM model is available on Lemonade server.
+
+        Returns:
+            True if model is available, False otherwise
+        """
+        try:
+            models_response = self.client.list_models()
+            available_models = [
+                m.get("id", "") for m in models_response.get("data", [])
+            ]
+
+            if self.vlm_model in available_models:
+                logger.debug(f"VLM model available: {self.vlm_model}")
+                return True
+            else:
+                logger.warning(f"❌ VLM model not found: {self.vlm_model}")
+                logger.warning("")
+                logger.warning("📥 To download this model:")
+                logger.warning(f"   1. Open Lemonade Model Manager ({self.server_url})")
+                logger.warning(f"   2. Search for: {self.vlm_model}")
+                logger.warning("   3. Click 'Download' to install the model")
+                logger.warning("")
+                logger.warning(
+                    f"   Available models: {', '.join(available_models[:3])}..."
+                )
+                return False
+
+        except Exception as e:
+            logger.error(f"Failed to check VLM availability: {e}")
+            logger.error(
+                f"   Make sure Lemonade server is running at {self.server_url}"
+            )
+            return False
+
+    def _ensure_vlm_loaded(self) -> bool:
+        """
+        Ensure VLM model is loaded, load it if necessary.
+
+        The model will be automatically downloaded if not available (handled by
+        lemonade_client.chat_completions with auto_download=True).
+
+        Returns:
+            True if VLM is loaded, False if loading failed
+        """
+        if self.vlm_loaded:
+            return True
+
+        if not self.auto_load:
+            logger.warning("VLM not loaded and auto_load=False")
+            return False
+
+        try:
+            logger.debug(f"Loading VLM model: {self.vlm_model}")
+            # Load model (auto-download handled by lemonade_client, may take hours)
+            self.client.load_model(self.vlm_model, timeout=60, auto_download=True)
+            self.vlm_loaded = True
+            logger.debug(f"VLM model loaded: {self.vlm_model}")
+            return True
+
+        except Exception as e:
+            logger.error(f"Failed to load VLM model: {e}")
+            logger.error(
+                f"   Make sure Lemonade server is running at {self.server_url}"
+            )
+            return False
+
+    def extract_from_image(
+        self,
+        image_bytes: bytes,
+        image_num: int = 1,
+        page_num: int = 1,
+        prompt: Optional[str] = None,
+    ) -> str:
+        """
+        Extract text from an image using VLM.
+
+        Args:
+            image_bytes: Image as PNG/JPEG bytes
+            image_num: Image number on page (for logging)
+            page_num: Page number (for logging)
+            prompt: Custom extraction prompt (optional)
+
+        Returns:
+            Extracted text in markdown format
+        """
+        # Ensure VLM is loaded
+        if not self._ensure_vlm_loaded():
+            error_msg = "VLM model not available"
+            logger.error(error_msg)
+            return f"[VLM extraction failed: {error_msg}]"
+
+        # Encode image as base64 and detect MIME type
+        # Note: Image size optimization happens in pdf_utils.py during extraction
+        image_b64 = base64.b64encode(image_bytes).decode("utf-8")
+        mime_type = detect_image_mime_type(image_bytes)
+
+        # Default prompt for text extraction
+        if not prompt:
+            prompt = """You are an OCR system. Extract ALL visible text from this image exactly as it appears.
+
+Instructions:
+1. Extract EVERY word you see - don't skip or paraphrase
+2. Preserve exact formatting (headings, bold, bullets, tables)
+3. If it's a table, format as markdown table
+4. If it's a chart, describe what you see: [CHART: ...]
+5. Do NOT add placeholders like "[Insert ...]"  - only extract actual text
+6. Do NOT generate or invent content - only extract what you see
+
+Output format: Clean markdown with the ACTUAL text from the image."""
+
+        # Format message with image (OpenAI vision format)
+        messages = [
+            {
+                "role": "user",
+                "content": [
+                    {"type": "text", "text": prompt},
+                    {
+                        "type": "image_url",
+                        "image_url": {"url": f"data:{mime_type};base64,{image_b64}"},
+                    },
+                ],
+            }
+        ]
+
+        try:
+            import time
+
+            start_time = time.time()
+
+            logger.debug(
+                f"VLM extracting from image {image_num} on page {page_num} ({mime_type})..."
+            )
+            logger.debug(
+                f"   Image: {mime_type}, {len(image_b64)} chars base64 ({len(image_bytes)} bytes raw)"
+            )
+
+            # Call VLM using chat completions endpoint
+            response = self.client.chat_completions(
+                model=self.vlm_model,
+                messages=messages,
+                temperature=0.1,  # Low temp for accurate extraction
+                max_completion_tokens=2048,  # Allow detailed extraction
+                timeout=300,  # VLM needs more time for complex forms (5 min)
+            )
+
+            elapsed = time.time() - start_time
+
+            # Extract text from response
+            if (
+                isinstance(response, dict)
+                and "choices" in response
+                and len(response["choices"]) > 0
+            ):
+                extracted_text = response["choices"][0]["message"]["content"]
+                size_kb = len(image_bytes) / 1024
+                logger.debug(
+                    f"Extracted {len(extracted_text)} chars from image {image_num} "
+                    f"in {elapsed:.2f}s ({size_kb:.0f}KB image)"
+                )
+                return extracted_text
+            else:
+                # Check for specific error types and provide helpful messages
+                error_msg = self._parse_vlm_error(response)
+                logger.error(error_msg)
+                return f"[VLM extraction failed: {error_msg}]"
+
+        except Exception as e:
+            logger.error(
+                f"VLM extraction failed for page {page_num}, image {image_num}: {e}"
+            )
+            import traceback
+
+            logger.debug(traceback.format_exc())
+            return f"[VLM extraction failed: {str(e)}]"
+
+    def _parse_vlm_error(self, response: dict) -> str:
+        """Parse VLM error response and return a helpful error message."""
+        if not isinstance(response, dict):
+            return f"Unexpected response type: {type(response)}"
+
+        # Check for nested error structure from Lemonade
+        error = response.get("error", {})
+        if isinstance(error, dict):
+            details = error.get("details", {})
+            inner_response = (
+                details.get("response", {}) if isinstance(details, dict) else {}
+            )
+            inner_error = (
+                inner_response.get("error", {})
+                if isinstance(inner_response, dict)
+                else {}
+            )
+
+            # Context size error
+            if inner_error.get("type") == "exceed_context_size_error":
+                n_ctx = inner_error.get("n_ctx", "unknown")
+                n_prompt = inner_error.get("n_prompt_tokens", "unknown")
+                return (
+                    f"Context size too small! Image requires {n_prompt} tokens "
+                    f"but model context is only {n_ctx}. "
+                    f"To fix: Right-click Lemonade tray icon → Settings → "
+                    f"set Context Size to 32768, then restart the model."
+                )
+
+            # Other backend errors
+            if error.get("type") == "backend_error":
+                msg = inner_error.get(
+                    "message", error.get("message", "Unknown backend error")
+                )
+                return f"Backend error: {msg}"
+
+        return f"Unexpected response format: {response}"
+
+    def extract_from_page_images(self, images: list, page_num: int) -> list:
+        """
+        Extract text from multiple images on a page.
+
+        Args:
+            images: List of image dicts with 'image_bytes', 'width', 'height', etc.
+            page_num: Page number
+
+        Returns:
+            List of dicts:
+            [
+                {
+                    "image_num": 1,
+                    "text": "extracted markdown",
+                    "dimensions": "800x600",
+                    "size_kb": 45.2
+                },
+                ...
+            ]
+        """
+        results = []
+
+        for img_idx, img_data in enumerate(images, 1):
+            extracted_text = self.extract_from_image(
+                image_bytes=img_data["image_bytes"],
+                image_num=img_idx,
+                page_num=page_num,
+            )
+
+            results.append(
+                {
+                    "image_num": img_idx,
+                    "text": extracted_text,
+                    "dimensions": f"{img_data['width']}x{img_data['height']}",
+                    "size_kb": img_data["size_kb"],
+                }
+            )
+
+        return results
+
+    def cleanup(self):
+        """
+        Cleanup VLM resources.
+
+        Call this after batch processing to mark VLM as unloaded.
+        Note: Model remains loaded on server; this just updates local state.
+        """
+        if self.vlm_loaded:
+            logger.info("🧹 VLM processing complete")
+            self.vlm_loaded = False
+
+    def __enter__(self):
+        """Context manager entry - ensure VLM loaded."""
+        self._ensure_vlm_loaded()
+        return self
+
+    def __exit__(self, exc_type, exc_val, exc_tb):
+        """Context manager exit - cleanup VLM state."""
+        self.cleanup()