apis/base64img2component.py

from openai import os
from src.clients.llm_client import LLMClient
import json 
import pandas as pd
from pydantic import BaseModel
from enum import Enum
import base64
from io import BytesIO
from PIL import Image
from functools import cache
from datetime import datetime
import pytz
from src.utils.tracer import customtracer

def _ask_raw_hf(messages, model, response_format=None):
    """Compatibility wrapper: routes OpenAI-style messages through HF LLMClient."""
    from src.clients.llm_client import LLMClient
    import json, re

    client = LLMClient()

    # Extract system prompt and user content from messages list
    system_prompt = None
    user_text = ""
    images = []
    for msg in messages:
        role = msg.get("role", "")
        c = msg.get("content", "")
        if role == "system":
            if isinstance(c, str):
                system_prompt = c
        elif role == "user":
            if isinstance(c, str):
                user_text = c
            elif isinstance(c, list):
                for part in c:
                    if isinstance(part, dict):
                        if part.get("type") == "text":
                            user_text += part.get("text", "")
                        elif part.get("type") == "image_url":
                            url = part.get("image_url", {}).get("url", "")
                            if url.startswith("data:"):
                                images.append(url.split(",", 1)[1] if "," in url else url)
                            else:
                                images.append(url)

    if response_format is not None and hasattr(response_format, "model_json_schema"):
        result = client.call(
            prompt=user_text,
            schema=response_format,
            model=model,
            system_prompt=system_prompt,
            images=images if images else None,
            temperature=0,
        )
        import json
        return json.dumps(result.model_dump(), ensure_ascii=False)
    else:
        return client.call_raw(
            prompt=user_text,
            model=model,
            system_prompt=system_prompt,
            images=images if images else None,
        )

class UIoption(str, Enum):
    element1 = "バナー/動画"
    element2 = "CTA"
    element3 = "テキスト"
    element4 = "フォーム"

class Component(BaseModel):
    component_large: str
    component_middle: str
    component_small: list[str]
    UIelement: UIoption

class Components(BaseModel):
    components: list[Component]

def ask_raw(messages):
    client = LLMClient()
    # HF: beta.parse not available; use _ask_raw_hf instead
    response = client.chat.completions.create(
        model='meta-llama/Llama-3.3-70B-Instruct',
        messages=messages,
        top_p=1,
        frequency_penalty=0,
        presence_penalty=0,
        response_format=Components,
        temperature=0
    )
    return response.choices[0].message.content

@customtracer
def base64img2component(p, image64, openai_key=os.environ.get('OPENAI_KEY')):
    """
    input1 (text): OCR text extracted from LP screenshot (long string)
    input2 (text): スクショ
    input3 (text): default
    output1 (json): components list
    """
    print(datetime.now(pytz.timezone('Asia/Tokyo')).strftime("%Y-%m-%d %H:%M:%S"), f"base64img2component:", image64[0:30]) 
    
    if openai_key == "default":
        os.environ['OPENAI_API_KEY'] = os.environ.get('OPENAI_KEY')
    else:
        os.environ['OPENAI_API_KEY'] = openai_key
    
    messages=[
        {
        "role": "system",
        "content": """
■ コンポーネント要素のアウトプットのサンプル
[
{"component_large":"商品/サービスの特徴","component_middle":"アンカー", "component_small":[], "UIelement":"テキスト"},
{"component_large":"FAQ/よくある質問","component_middle":"よくある質問", "component_small":["自宅外出ずに何か書類が届くことはありますか？","家族割などの割引はありますか？"], "UIelement":"表組み"}
]
"""
        },
        {
        "role": "user",
        "content":  [{"type": "text", "text":p}]
        },
    ]

    messages[1]["content"].insert(0, {"type": "image_url", "image_url": {"url":"data:image/png;base64,"+image64}})
    # Propagate OpenAI auth errors explicitly so caller can display message.
    try:
        return ask_raw(messages)
    except openai.AuthenticationError as e:
        # Raise RuntimeError with clear message for API key / auth issues.
        # Caller (BE_Origin side) can catch and display this message to user.
        raise RuntimeError(f"[base64img2component] OpenAI AuthenticationError: {e}") from e