108 lines
4.8 KiB
Python
108 lines
4.8 KiB
Python
# app/api/v1/endpoints/ocr.py
|
|
import logging
|
|
from typing import List
|
|
|
|
from fastapi import APIRouter, Depends, HTTPException, status, UploadFile, File
|
|
from google.api_core import exceptions as google_exceptions # Import Google API exceptions
|
|
|
|
from app.api.dependencies import get_current_user
|
|
from app.models import User as UserModel
|
|
from app.schemas.ocr import OcrExtractResponse
|
|
from app.core.gemini import extract_items_from_image_gemini, gemini_initialization_error # Import helper
|
|
|
|
logger = logging.getLogger(__name__)
|
|
router = APIRouter()
|
|
|
|
# Allowed image MIME types
|
|
ALLOWED_IMAGE_TYPES = ["image/jpeg", "image/png", "image/webp"]
|
|
MAX_FILE_SIZE_MB = 10 # Set a reasonable max file size
|
|
|
|
@router.post(
|
|
"/extract-items",
|
|
response_model=OcrExtractResponse,
|
|
summary="Extract List Items via OCR (Gemini)",
|
|
tags=["OCR"]
|
|
)
|
|
async def ocr_extract_items(
|
|
current_user: UserModel = Depends(get_current_user),
|
|
# Use File(...) for better metadata handling than UploadFile directly as type hint
|
|
image_file: UploadFile = File(..., description="Image file (JPEG, PNG, WEBP) of the shopping list or receipt."),
|
|
):
|
|
"""
|
|
Accepts an image upload, sends it to Gemini Flash with a prompt
|
|
to extract shopping list items, and returns the parsed items.
|
|
"""
|
|
# Check if Gemini client initialized correctly
|
|
if gemini_initialization_error:
|
|
logger.error("OCR endpoint called but Gemini client failed to initialize.")
|
|
raise HTTPException(
|
|
status_code=status.HTTP_503_SERVICE_UNAVAILABLE,
|
|
detail=f"OCR service unavailable: {gemini_initialization_error}"
|
|
)
|
|
|
|
logger.info(f"User {current_user.email} uploading image '{image_file.filename}' for OCR extraction.")
|
|
|
|
# --- File Validation ---
|
|
if image_file.content_type not in ALLOWED_IMAGE_TYPES:
|
|
logger.warning(f"Invalid file type uploaded by {current_user.email}: {image_file.content_type}")
|
|
raise HTTPException(
|
|
status_code=status.HTTP_400_BAD_REQUEST,
|
|
detail=f"Invalid file type. Allowed types: {', '.join(ALLOWED_IMAGE_TYPES)}",
|
|
)
|
|
|
|
# Simple size check (FastAPI/Starlette might handle larger limits via config)
|
|
# Read content first to get size accurately
|
|
contents = await image_file.read()
|
|
if len(contents) > MAX_FILE_SIZE_MB * 1024 * 1024:
|
|
logger.warning(f"File too large uploaded by {current_user.email}: {len(contents)} bytes")
|
|
raise HTTPException(
|
|
status_code=status.HTTP_413_REQUEST_ENTITY_TOO_LARGE,
|
|
detail=f"File size exceeds limit of {MAX_FILE_SIZE_MB} MB.",
|
|
)
|
|
# --- End File Validation ---
|
|
|
|
try:
|
|
# Call the Gemini helper function
|
|
extracted_items = await extract_items_from_image_gemini(image_bytes=contents)
|
|
|
|
logger.info(f"Successfully extracted {len(extracted_items)} items for user {current_user.email}.")
|
|
return OcrExtractResponse(extracted_items=extracted_items)
|
|
|
|
except ValueError as e:
|
|
# Handle errors from Gemini processing (blocked, empty response, etc.)
|
|
logger.warning(f"Gemini processing error for user {current_user.email}: {e}")
|
|
raise HTTPException(
|
|
status_code=status.HTTP_422_UNPROCESSABLE_ENTITY, # Or 400 Bad Request?
|
|
detail=f"Could not extract items from image: {e}",
|
|
)
|
|
except google_exceptions.ResourceExhausted as e:
|
|
# Specific handling for quota errors
|
|
logger.error(f"Gemini Quota Exceeded for user {current_user.email}: {e}", exc_info=True)
|
|
raise HTTPException(
|
|
status_code=status.HTTP_429_TOO_MANY_REQUESTS,
|
|
detail="OCR service quota exceeded. Please try again later.",
|
|
)
|
|
except google_exceptions.GoogleAPIError as e:
|
|
# Handle other Google API errors (e.g., invalid key, permissions)
|
|
logger.error(f"Gemini API Error for user {current_user.email}: {e}", exc_info=True)
|
|
raise HTTPException(
|
|
status_code=status.HTTP_503_SERVICE_UNAVAILABLE,
|
|
detail=f"OCR service error: {e}",
|
|
)
|
|
except RuntimeError as e:
|
|
# Catch initialization errors from get_gemini_client()
|
|
logger.error(f"Gemini client runtime error during OCR request: {e}")
|
|
raise HTTPException(
|
|
status_code=status.HTTP_503_SERVICE_UNAVAILABLE,
|
|
detail=f"OCR service configuration error: {e}"
|
|
)
|
|
except Exception as e:
|
|
# Catch any other unexpected errors
|
|
logger.exception(f"Unexpected error during OCR extraction for user {current_user.email}: {e}")
|
|
raise HTTPException(
|
|
status_code=status.HTTP_500_INTERNAL_SERVER_ERROR,
|
|
detail="An unexpected error occurred during item extraction.",
|
|
)
|
|
finally:
|
|
# Ensure file handle is closed (UploadFile uses SpooledTemporaryFile)
|
|
await image_file.close() |