thanhnt-cf's picture
step 1
0dd08cb
raw
history blame
2.03 kB
from abc import ABC, abstractmethod
from typing import Any, Dict, List, Type, Union
from pydantic import BaseModel
from app.utils.converter import to_snake_case
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_model: Type[BaseModel],
ai_model: str,
img_urls: List[str],
product_taxonomy: str,
product_data: Dict[str, Union[str, List[str]]],
pil_images: List[Any] = None,
img_paths: List[str] = None,
) -> Dict[str, Any]:
# validate_json_schema(schema)
# create mappings for keys of attributes, to make the key following naming convention of python variables
schema = attributes_model.model_json_schema()
data = await self.extract_attributes(
attributes_model,
ai_model,
img_urls,
product_taxonomy if product_taxonomy != "" else "main",
product_data,
# pil_images=pil_images, # temporarily removed to 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