winebox 0.1.2__py3-none-any.whl → 0.1.4__py3-none-any.whl

winebox/services/image_storage.py

@@ -4,42 +4,171 @@ import uuid
 from pathlib import Path
 
 import aiofiles
-from fastapi import UploadFile
+from fastapi import HTTPException, UploadFile, status
 
 from winebox.config import settings
 
+# Allowed MIME types for image uploads
+ALLOWED_MIME_TYPES = {
+    "image/jpeg",
+    "image/png",
+    "image/gif",
+    "image/webp",
+}
+
+# File extension to MIME type mapping
+EXTENSION_MIME_MAP = {
+    ".jpg": "image/jpeg",
+    ".jpeg": "image/jpeg",
+    ".png": "image/png",
+    ".gif": "image/gif",
+    ".webp": "image/webp",
+}
+
+# Magic byte signatures for image formats
+# Each entry is (magic_bytes, offset, detected_extension)
+IMAGE_MAGIC_SIGNATURES = [
+    # JPEG: starts with FF D8 FF
+    (b"\xff\xd8\xff", 0, ".jpg"),
+    # PNG: starts with 89 50 4E 47 0D 0A 1A 0A
+    (b"\x89PNG\r\n\x1a\n", 0, ".png"),
+    # GIF87a and GIF89a
+    (b"GIF87a", 0, ".gif"),
+    (b"GIF89a", 0, ".gif"),
+    # WebP: starts with RIFF....WEBP
+    (b"RIFF", 0, ".webp"),  # Additional check for WEBP at offset 8
+]
+
+
+def detect_image_type(content: bytes) -> str | None:
+    """Detect image type from file content using magic bytes.
+
+    Args:
+        content: The file content bytes.
+
+    Returns:
+        The detected extension (e.g., ".jpg") or None if not a valid image.
+    """
+    if len(content) < 12:
+        return None
+
+    for magic, offset, ext in IMAGE_MAGIC_SIGNATURES:
+        if content[offset:offset + len(magic)] == magic:
+            # Special case for WebP: verify WEBP signature at offset 8
+            if ext == ".webp":
+                if content[8:12] != b"WEBP":
+                    continue
+            return ext
+
+    return None
+
+
+class FileSizeExceededError(Exception):
+    """Raised when uploaded file exceeds size limit."""
+
+    pass
+
+
+class InvalidFileTypeError(Exception):
+    """Raised when uploaded file has invalid type."""
+
+    pass
+
+
+class InvalidMagicBytesError(Exception):
+    """Raised when file content doesn't match a valid image format."""
+
+    pass
+
 
 class ImageStorageService:
     """Service for storing and managing wine label images."""
 
-    def __init__(self, storage_path: Path | None = None) -> None:
+    def __init__(
+        self,
+        storage_path: Path | None = None,
+        max_size_bytes: int | None = None,
+    ) -> None:
         """Initialize the image storage service.
 
         Args:
             storage_path: Path to store images. Defaults to config setting.
+            max_size_bytes: Maximum file size in bytes. Defaults to config setting.
         """
         self.storage_path = storage_path or settings.image_storage_path
         self.storage_path.mkdir(parents=True, exist_ok=True)
+        self.max_size_bytes = max_size_bytes or settings.max_upload_size_bytes
+
+    def _validate_extension(self, filename: str | None) -> str:
+        """Validate and return the file extension.
+
+        Args:
+            filename: The original filename.
+
+        Returns:
+            Valid extension (e.g., ".jpg").
+
+        Raises:
+            InvalidFileTypeError: If extension is not allowed.
+        """
+        if not filename:
+            return ".jpg"
+
+        ext = Path(filename).suffix.lower()
+        if ext not in EXTENSION_MIME_MAP:
+            raise InvalidFileTypeError(
+                f"Invalid file type. Allowed types: {', '.join(EXTENSION_MIME_MAP.keys())}"
+            )
+        return ext
 
     async def save_image(self, upload_file: UploadFile) -> str:
-        """Save an uploaded image file.
+        """Save an uploaded image file with size, type, and content validation.
 
         Args:
             upload_file: The uploaded file from FastAPI.
 
         Returns:
             The filename of the saved image.
-        """
-        # Generate unique filename
-        ext = Path(upload_file.filename or "image.jpg").suffix.lower()
-        if ext not in [".jpg", ".jpeg", ".png", ".gif", ".webp"]:
-            ext = ".jpg"
 
+        Raises:
+            HTTPException: If file exceeds size limit, has invalid type, or
+                          content doesn't match a valid image format.
+        """
+        # Validate extension first
+        try:
+            declared_ext = self._validate_extension(upload_file.filename)
+        except InvalidFileTypeError as e:
+            raise HTTPException(
+                status_code=status.HTTP_400_BAD_REQUEST,
+                detail=str(e),
+            )
+
+        # Read content with size limit check
+        content = await upload_file.read()
+        if len(content) > self.max_size_bytes:
+            max_mb = self.max_size_bytes / (1024 * 1024)
+            raise HTTPException(
+                status_code=status.HTTP_413_CONTENT_TOO_LARGE,
+                detail=f"File size exceeds maximum allowed size of {max_mb:.1f} MB",
+            )
+
+        # Validate magic bytes - ensure file content matches a valid image format
+        detected_ext = detect_image_type(content)
+        if detected_ext is None:
+            raise HTTPException(
+                status_code=status.HTTP_400_BAD_REQUEST,
+                detail="Invalid file content. File does not appear to be a valid image.",
+            )
+
+        # Use the detected extension (more reliable than declared extension)
+        # This prevents attacks where malicious files are renamed to .jpg/.png
+        ext = detected_ext
+
+        # Generate unique filename with detected extension
         filename = f"{uuid.uuid4()}{ext}"
         file_path = self.storage_path / filename
 
         # Save file
-        content = await upload_file.read()
         async with aiofiles.open(file_path, "wb") as f:
             await f.write(content)
 

winebox/services/ocr.py

@@ -1,5 +1,6 @@
 """OCR service for extracting text from wine label images."""
 
+import io
 import logging
 from pathlib import Path
 
@@ -66,6 +67,42 @@ class OCRService:
             logger.error(f"OCR extraction failed: {e}")
             return ""
 
+    async def extract_text_from_bytes(self, image_data: bytes) -> str:
+        """Extract text from image bytes without saving to disk.
+
+        Args:
+            image_data: Raw image data as bytes.
+
+        Returns:
+            Extracted text from the image.
+        """
+        try:
+            import pytesseract
+
+            # Open image from bytes
+            image = Image.open(io.BytesIO(image_data))
+
+            # Preprocess image for better OCR results
+            # Convert to grayscale
+            if image.mode != "L":
+                image = image.convert("L")
+
+            # Extract text
+            text = pytesseract.image_to_string(
+                image,
+                lang="eng",
+                config="--psm 6",  # Assume uniform block of text
+            )
+
+            return text.strip()
+
+        except ImportError:
+            logger.error("pytesseract is not installed")
+            return ""
+        except Exception as e:
+            logger.error(f"OCR extraction failed: {e}")
+            return ""
+
     async def extract_text_with_confidence(
         self, image_path: str | Path
     ) -> tuple[str, float]:

winebox/services/vision.py

@@ -0,0 +1,278 @@
+"""Claude Vision service for wine label analysis."""
+
+import base64
+import json
+import logging
+import os
+from typing import Any
+
+from winebox.config import settings
+
+logger = logging.getLogger(__name__)
+
+WINE_ANALYSIS_PROMPT = """Analyze this wine label image and extract the following information.
+Return ONLY a valid JSON object with these fields (use null for any field you cannot determine):
+
+{
+    "name": "The wine name/title",
+    "winery": "The winery or producer name",
+    "vintage": 2020,
+    "grape_variety": "The grape variety (e.g., Cabernet Sauvignon, Chardonnay)",
+    "region": "The wine region (e.g., Napa Valley, Bordeaux)",
+    "country": "The country of origin",
+    "alcohol_percentage": 13.5,
+    "raw_text": "All readable text from the label, preserving line breaks"
+}
+
+Important:
+- vintage should be a number (year) or null
+- alcohol_percentage should be a number or null
+- Extract ALL visible text for raw_text, including small print
+- If you see multiple wines or labels, focus on the main/primary one
+- Be thorough - wine labels often have text in multiple locations"""
+
+
+class ClaudeVisionService:
+    """Service for analyzing wine labels using Claude's vision capabilities."""
+
+    def __init__(self) -> None:
+        """Initialize the Claude Vision service."""
+        self._default_client = None
+
+    def _get_system_api_key(self) -> str | None:
+        """Get the system-wide API key from settings or environment."""
+        return settings.anthropic_api_key or os.getenv("ANTHROPIC_API_KEY")
+
+    def _get_client(self, user_api_key: str | None = None):
+        """Get an Anthropic client, using user key if provided, else system key."""
+        try:
+            import anthropic
+
+            # Use user's API key if provided, otherwise use system key
+            api_key = user_api_key or self._get_system_api_key()
+            if not api_key:
+                raise ValueError("No Anthropic API key configured")
+
+            # If using system key, cache the client
+            if not user_api_key:
+                if self._default_client is None:
+                    self._default_client = anthropic.Anthropic(api_key=api_key)
+                return self._default_client
+
+            # Create a new client for user-specific key
+            return anthropic.Anthropic(api_key=api_key)
+        except ImportError:
+            logger.error("anthropic package is not installed")
+            raise
+
+    @property
+    def client(self):
+        """Lazy-load the default Anthropic client (for backward compatibility)."""
+        return self._get_client()
+
+    def is_available(self, user_api_key: str | None = None) -> bool:
+        """Check if Claude Vision is available.
+
+        Args:
+            user_api_key: Optional user-specific API key to check.
+        """
+        try:
+            api_key = user_api_key or self._get_system_api_key()
+            return bool(api_key) and settings.use_claude_vision
+        except Exception:
+            return False
+
+    async def analyze_label(
+        self,
+        image_data: bytes,
+        media_type: str = "image/jpeg",
+        user_api_key: str | None = None,
+    ) -> dict[str, Any]:
+        """Analyze a wine label image using Claude Vision.
+
+        Args:
+            image_data: Raw image data as bytes.
+            media_type: MIME type of the image (image/jpeg, image/png, etc.)
+            user_api_key: Optional user-specific API key.
+
+        Returns:
+            Dictionary with parsed wine information.
+        """
+        try:
+            # Encode image to base64
+            image_base64 = base64.standard_b64encode(image_data).decode("utf-8")
+
+            # Get client (uses user key if provided, else system key)
+            client = self._get_client(user_api_key)
+
+            # Call Claude API with vision
+            message = client.messages.create(
+                model="claude-sonnet-4-20250514",
+                max_tokens=1024,
+                messages=[
+                    {
+                        "role": "user",
+                        "content": [
+                            {
+                                "type": "image",
+                                "source": {
+                                    "type": "base64",
+                                    "media_type": media_type,
+                                    "data": image_base64,
+                                },
+                            },
+                            {
+                                "type": "text",
+                                "text": WINE_ANALYSIS_PROMPT,
+                            },
+                        ],
+                    }
+                ],
+            )
+
+            # Extract the response text
+            response_text = message.content[0].text
+
+            # Parse JSON from response
+            # Handle case where Claude might wrap JSON in markdown code blocks
+            if "```json" in response_text:
+                response_text = response_text.split("```json")[1].split("```")[0]
+            elif "```" in response_text:
+                response_text = response_text.split("```")[1].split("```")[0]
+
+            result = json.loads(response_text.strip())
+
+            # Ensure all expected fields exist
+            return {
+                "name": result.get("name"),
+                "winery": result.get("winery"),
+                "vintage": result.get("vintage"),
+                "grape_variety": result.get("grape_variety"),
+                "region": result.get("region"),
+                "country": result.get("country"),
+                "alcohol_percentage": result.get("alcohol_percentage"),
+                "raw_text": result.get("raw_text", ""),
+            }
+
+        except json.JSONDecodeError as e:
+            logger.error(f"Failed to parse Claude response as JSON: {e}")
+            logger.debug(f"Response was: {response_text}")
+            return self._empty_result()
+        except Exception as e:
+            logger.error(f"Claude Vision analysis failed: {e}")
+            return self._empty_result()
+
+    async def analyze_labels(
+        self,
+        front_image_data: bytes,
+        back_image_data: bytes | None = None,
+        front_media_type: str = "image/jpeg",
+        back_media_type: str = "image/jpeg",
+        user_api_key: str | None = None,
+    ) -> dict[str, Any]:
+        """Analyze front and back wine label images.
+
+        Args:
+            front_image_data: Front label image data.
+            back_image_data: Optional back label image data.
+            front_media_type: MIME type of front image.
+            back_media_type: MIME type of back image.
+            user_api_key: Optional user-specific API key.
+
+        Returns:
+            Combined analysis results.
+        """
+        try:
+            # Get client (uses user key if provided, else system key)
+            client = self._get_client(user_api_key)
+
+            # Build message content with images
+            content = [
+                {
+                    "type": "image",
+                    "source": {
+                        "type": "base64",
+                        "media_type": front_media_type,
+                        "data": base64.standard_b64encode(front_image_data).decode("utf-8"),
+                    },
+                },
+                {
+                    "type": "text",
+                    "text": "Front label:" if back_image_data else WINE_ANALYSIS_PROMPT,
+                },
+            ]
+
+            if back_image_data:
+                content.extend([
+                    {
+                        "type": "image",
+                        "source": {
+                            "type": "base64",
+                            "media_type": back_media_type,
+                            "data": base64.standard_b64encode(back_image_data).decode("utf-8"),
+                        },
+                    },
+                    {
+                        "type": "text",
+                        "text": "Back label:",
+                    },
+                    {
+                        "type": "text",
+                        "text": WINE_ANALYSIS_PROMPT.replace(
+                            "this wine label image",
+                            "these wine label images (front and back)"
+                        ),
+                    },
+                ])
+
+            # Call Claude API
+            message = client.messages.create(
+                model="claude-sonnet-4-20250514",
+                max_tokens=1024,
+                messages=[{"role": "user", "content": content}],
+            )
+
+            response_text = message.content[0].text
+
+            # Parse JSON
+            if "```json" in response_text:
+                response_text = response_text.split("```json")[1].split("```")[0]
+            elif "```" in response_text:
+                response_text = response_text.split("```")[1].split("```")[0]
+
+            result = json.loads(response_text.strip())
+
+            return {
+                "name": result.get("name"),
+                "winery": result.get("winery"),
+                "vintage": result.get("vintage"),
+                "grape_variety": result.get("grape_variety"),
+                "region": result.get("region"),
+                "country": result.get("country"),
+                "alcohol_percentage": result.get("alcohol_percentage"),
+                "raw_text": result.get("raw_text", ""),
+                "front_label_text": result.get("raw_text", ""),
+                "back_label_text": None,  # Combined in raw_text
+            }
+
+        except json.JSONDecodeError as e:
+            logger.error(f"Failed to parse Claude response as JSON: {e}")
+            return self._empty_result()
+        except Exception as e:
+            logger.error(f"Claude Vision analysis failed: {e}")
+            return self._empty_result()
+
+    def _empty_result(self) -> dict[str, Any]:
+        """Return an empty result dictionary."""
+        return {
+            "name": None,
+            "winery": None,
+            "vintage": None,
+            "grape_variety": None,
+            "region": None,
+            "country": None,
+            "alcohol_percentage": None,
+            "raw_text": "",
+            "front_label_text": "",
+            "back_label_text": None,
+        }