import base64
from openai import OpenAI
import os

client = OpenAI(base_url="https://integrate.api.nvidia.com/v1", api_key=os.environ["NIM_API_KEY"])

def get_elements_from_screenshot(image_path):
    with open(image_path, "rb") as f:
        image_b64 = base64.b64encode(f.read()).decode()

    completion = client.chat.completions.create(
        model="nvidia/nemotron-3-nano-omni-30b-a3b-reasoning",
        messages=[{
            "role": "user",
            "content": [
                {"type": "text", "text": "Return ONLY valid JSON identifying all interactive elements (name, box, actionable)."},
                {"type": "image_url", "image_url": {"url": f"data:image/jpeg;base64,{image_b64}"}}
            ]
        }],
        temperature=0.1
    )
    return completion.choices[0].message.content