apis/baseimg2html.py

"""
baseimg2html: 画像からHTMLコンポーネントを生成。
HF版: VLM (Qwen2.5-VL) を使用。Vertex AI は使用しない。
"""

import os
from typing import List

from pydantic import BaseModel

from src.utils.tracer import customtracer


class newHTMLs(BaseModel):
    HTMLs: List[str]


@customtracer
def baseimg2html(p: str, base64img: str, gcp_key: str = "default") -> dict:
    """
    input1 (text): OCRテキストやページの説明
    input2 (text): base64エンコードされた画像
    input3 (text): default
    output1 (json): HTMLコンポーネントのリスト（{"HTMLs": [...]}）

    NOTE: HF版は VLM ベース。Vertex AI / GCP は使用しない。
    """
    from src.clients.llm_client import LLMClient

    client = LLMClient()

    system_prompt = (
        "あなたはHTMLとCSSの達人です。"
        "画像とテキスト情報を基に、適切なHTMLコンポーネントを生成してください。"
        "各コンポーネントは完全に機能するHTMLとして生成してください。"
    )

    prompt = (
        system_prompt + "\n\n[テキスト情報]\n" + str(p)
        + "\n\n画像のコンポーネントをHTMLリストとして出力してください。"
    )

    # Strip base64 data-URI prefix if present
    if base64img and "," in base64img:
        base64img = base64img.split(",", 1)[1]

    result = client.call(
        prompt=prompt,
        schema=newHTMLs,
        model="Qwen/Qwen2.5-VL-72B-Instruct",
        images=[base64img] if base64img else None,
        temperature=0.1,
    )
    return result.model_dump()