app/request_handler/extract_handler.py

from io import BytesIO

import requests
from fastapi import HTTPException
from PIL import Image

from app.config import get_settings
from app.core.errors import BadRequestError, VendorError
from app.schemas.requests import ExtractionRequest
from app.schemas.responses import APIResponse
from app.services.factory import AIServiceFactory
from app.utils.logger import setup_logger

logger = setup_logger(__name__)
settings = get_settings()


async def handle_extract(request: ExtractionRequest):
    request.max_attempts = max(request.max_attempts, 1)
    request.max_attempts = min(request.max_attempts, 5)

    for attempt in range(1, request.max_attempts + 1):
        try:
            logger.info(f"Attempt: {attempt}")
            if request.ai_model in settings.OPENAI_MODELS:
                ai_vendor = "openai"
            elif request.ai_model in settings.ANTHROPIC_MODELS:
                ai_vendor = "anthropic"
            else:
                raise ValueError(
                    f"Invalid AI model: {request.ai_model}, only support {settings.SUPPORTED_MODELS}"
                )
            service = AIServiceFactory.get_service(ai_vendor)

            pil_images = []
            for url in request.img_urls:
                try:
                    response = requests.get(url)
                    response.raise_for_status()
                    image = Image.open(BytesIO(response.content))
                    pil_images.append(image)
                except Exception as e:
                    print(e)
                    logger.error(f"Failed to download or process image from {url}: {e}")
                    raise HTTPException(
                        status_code=400,
                        detail=f"Failed to process image from {url}",
                        headers={"attempt": attempt},
                    )

            json_attributes = await service.extract_attributes_with_validation(
                request.attributes,
                request.ai_model,
                request.img_urls,
                request.product_taxonomy,
                request.product_data,
                pil_images=pil_images,
            )
            break
        except BadRequestError as e:
            logger.error("Bad request error: ", e)
            raise HTTPException(
                status_code=400, detail=str(e), headers={"attempt": attempt}
            )
        except ValueError as e:
            logger.error("Value error: ", e)
            raise HTTPException(
                status_code=400, detail=str(e), headers={"attempt": attempt}
            )
        except VendorError as e:
            logger.error("Vendor error: ", e)
            if attempt == request.max_attempts:
                raise HTTPException(
                    status_code=500, detail=str(e), headers={"attempt": attempt}
                )
            else:
                if request.ai_model in settings.ANTHROPIC_MODELS:
                    request.ai_model = settings.OPENAI_MODELS[
                        0
                    ]  # switch to OpenAI, and try again if max_attempts not reached
                    logger.info(
                        f"Switching from anthropic to {request.ai_model} for attempt {attempt + 1}"
                    )
                elif request.ai_model in settings.OPENAI_MODELS:
                    request.ai_model = settings.ANTHROPIC_MODELS[
                        0
                    ]  # switch to anthropic, and try again if max_attempts not reached
                    logger.info(
                        f"Switching from OpenAI to {request.ai_model} for attempt {attempt + 1}"
                    )

        except HTTPException as e:
            logger.error("HTTP exception: ", e)
            raise e
        except Exception as e:
            logger.error("Exception: ", e)
            if (
                "overload" in str(e).lower()
                and request.ai_model in settings.ANTHROPIC_MODELS
            ):
                request.ai_model = settings.OPENAI_MODELS[
                    0
                ]  # switch to OpenAI, and try again if max_attempts not reached
            if attempt == request.max_attempts:
                raise HTTPException(
                    status_code=500,
                    detail="Internal server error",
                    headers={"attempt": attempt},
                )

    return json_attributes, attempt