app/services/base.py

from abc import ABC, abstractmethod
from typing import Any, Dict, List, Type

from pydantic import BaseModel

from app.schemas.schema_tools import (
    convert_attribute_to_model,
    validate_json_data,
    validate_json_schema,
)


class BaseAttributionService(ABC):
    @abstractmethod
    async def extract_attributes(
        self,
        attributes_model: Type[BaseModel],
        ai_model: str,
        img_urls: List[str],
        product_taxonomy: str,
        pil_images: List[Any] = None,
    ) -> Dict[str, Any]:
        pass

    @abstractmethod
    async def follow_schema(
        self, schema: Dict[str, Any], data: Dict[str, Any]
    ) -> Dict[str, Any]:
        pass

    async def extract_attributes_with_validation(
        self,
        attributes: Dict[str, Any],
        ai_model: str,
        img_urls: List[str],
        product_taxonomy: str,
        product_data: Dict[str, str],
        pil_images: List[Any] = None,
        img_paths: List[str] = None,
    ) -> Dict[str, Any]:
        # validate_json_schema(schema)
        attributes_model = convert_attribute_to_model(attributes)
        schema = attributes_model.model_json_schema()
        data = await self.extract_attributes(
            attributes_model,
            ai_model,
            img_urls,
            product_taxonomy,
            product_data,
            # pil_images=pil_images, # temporarily removed for save cost
            img_paths=img_paths,
        )
        validate_json_data(data, schema)
        return data

    async def follow_schema_with_validation(
        self, schema: Dict[str, Any], data: Dict[str, Any]
    ) -> Dict[str, Any]:
        validate_json_schema(schema)
        data = await self.follow_schema(schema, data)
        validate_json_data(data, schema)
        return data