File size: 6,460 Bytes
8ba64a4
 
9bb2fc2
8ba64a4
 
 
 
 
 
 
9645c29
 
 
 
 
e85027d
8ba64a4
 
 
 
 
 
 
9645c29
 
8ba64a4
9645c29
8ba64a4
9645c29
 
 
 
8ba64a4
9645c29
b02de59
 
 
8ba64a4
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
638f225
8ba64a4
 
2dc5702
8ba64a4
f124ae7
 
e3195a7
f124ae7
8ba64a4
 
 
e3195a7
8ba64a4
 
 
9645c29
 
 
 
 
 
 
 
8ba64a4
 
 
 
9645c29
8ba64a4
 
9645c29
8ba64a4
 
 
 
 
 
 
 
 
 
 
 
 
9645c29
8ba64a4
 
 
 
 
 
 
 
 
 
 
 
806f6b3
8ba64a4
 
 
9645c29
 
8ba64a4
 
9645c29
 
8ba64a4
e85027d
 
 
8ba64a4
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
e3195a7
8ba64a4
 
 
 
 
915961c
8ba64a4
 
 
 
 
 
 
 
 
 
806f6b3
8ba64a4
 
 
 
 
 
e85027d
 
 
8ba64a4
 
 
 
 
 
 
 
 
 
 
 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
import json
import os
from typing import Any, Dict, List, Type, Union

import openai
import weave
from openai import AsyncOpenAI
from pydantic import BaseModel

from app.utils.converter import product_data_to_str
from app.utils.image_processing import (
    get_data_format,
    get_image_base64_and_type,
    get_image_data,
)
from app.utils.logger import exception_to_str, setup_logger

from ..config import get_settings
from ..core import errors
from ..core.errors import BadRequestError, VendorError
from ..core.prompts import get_prompts
from .base import BaseAttributionService

ENV = os.getenv("ENV", "LOCAL")
if ENV == "LOCAL":  # local or demo
    weave_project_name = "cfai/attribution-exp"
elif ENV == "DEV":
    weave_project_name = "cfai/attribution-dev"
elif ENV == "UAT":
    weave_project_name = "cfai/attribution-uat"
elif ENV == "PROD":
    pass

if ENV != "PROD":
    #Disabled for now
    #weave.init(project_name=weave_project_name)
    print("something")
settings = get_settings()
prompts = get_prompts()
logger = setup_logger(__name__)


def get_response_format(json_schema: dict[str, any]) -> dict[str, any]:
    # OpenAI requires each $def have to have additionalProperties set to False
    json_schema["additionalProperties"] = False

    # check if the schema has a $defs key
    if "$defs" in json_schema:
        for keys in json_schema["$defs"].keys():
            json_schema["$defs"][keys]["additionalProperties"] = False
    response_format = {
        "type": "json_schema",
        "json_schema": {"strict": True, "name": "GarmentSchema", "schema": json_schema},
    }

    return response_format


class OpenAIService(BaseAttributionService):
    def __init__(self):
        self.client = AsyncOpenAI(api_key=settings.OPENAI_API_KEY)

    @weave.op
    async def extract_attributes(
        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,  # do not remove, this is for weave
        img_paths: List[str] = None,
        data: Dict[str, Any] = None,
    ) -> Dict[str, Any]:

        print("Prompt: ")
        print(prompts.EXTRACT_INFO_HUMAN_MESSAGE.replace("{product_data}", product_data_to_str(product_data)))
        
        text_content = [
            {
                "type": "text",
                "text": prompts.EXTRACT_INFO_HUMAN_MESSAGE.replace("{product_data}", product_data_to_str(product_data)),
            },
        ]
        if img_urls is not None:
            base64_data_list = []
            data_format_list = []

            for img_url in img_urls:
                base64_data, data_format = get_image_base64_and_type(img_url)
                base64_data_list.append(base64_data)
                data_format_list.append(data_format)

            image_content = [
                {
                    "type": "image_url",
                    "image_url": {
                        "url": f"data:image/{data_format};base64,{base64_data}",
                    },
                }
                for base64_data, data_format in zip(base64_data_list, data_format_list)
            ]
        elif img_paths is not None:
            image_content = [
                {
                    "type": "image_url",
                    "image_url": {
                        "url": f"data:image/{get_data_format(img_path)};base64,{get_image_data(img_path)}",
                    },
                }
                for img_path in img_paths
            ]

        try:
            logger.info("Extracting info via OpenAI...")
            response = await self.client.beta.chat.completions.parse(
                model=ai_model,
                messages=[
                    {
                        "role": "system",
                        "content": prompts.EXTRACT_INFO_SYSTEM_MESSAGE,
                    },
                    {
                        "role": "user",
                        "content": text_content + image_content,
                    },
                ],
                max_tokens=1000,
                response_format=attributes_model,
                logprobs=False,
                # top_logprobs=2,
                # temperature=0.0,
                top_p=1e-45,
            )
        except openai.BadRequestError as e:
            error_message = exception_to_str(e)
            raise BadRequestError(error_message)
        except Exception as e:
            raise VendorError(
                errors.VENDOR_THROW_ERROR.format(error_message=exception_to_str(e))
            )

        try:
            content = response.choices[0].message.content
            parsed_data = json.loads(content)
        except:
            raise VendorError(errors.VENDOR_ERROR_INVALID_JSON)

        return parsed_data

    @weave.op
    async def follow_schema(
        self, schema: Dict[str, Any], data: Dict[str, Any]
    ) -> Dict[str, Any]:
        logger.info("Following structure via OpenAI...")
        text_content = [
            {
                "type": "text",
                "text": prompts.FOLLOW_SCHEMA_HUMAN_MESSAGE.replace("{json_info}", json.dumps(data)),
            },
        ]

        try:
            response = await self.client.beta.chat.completions.parse(
                model=ai_model,
                messages=[
                    {
                        "role": "system",
                        "content": prompts.FOLLOW_SCHEMA_SYSTEM_MESSAGE,
                    },
                    {
                        "role": "user",
                        "content": text_content,
                    },
                ],
                max_tokens=1000,
                response_format=get_response_format(schema),
                logprobs=False,
                # top_logprobs=2,
                temperature=0.0,
            )
        except Exception as e:
            raise VendorError(
                errors.VENDOR_THROW_ERROR.format(error_message=exception_to_str(e))
            )

        if response.choices[0].message.refusal:
            logger.info("OpenAI refused to respond to the request")
            return {"status": "refused"}

        try:
            content = response.choices[0].message.content
            parsed_data = json.loads(content)
        except:
            raise ValueError(errors.VENDOR_ERROR_INVALID_JSON)

        return parsed_data