# -------------------------------------------------------------------------------- NOTEBOOK-CELL: CODE
# -*- coding: utf-8 -*-
import dataiku
import pandas as pd

folder = dataiku.Folder("e23xCOAh")

LLM_ID = dataiku.get_custom_variables()["LLM_id"]
project = dataiku.api_client().get_default_project()
llm = project.get_llm(LLM_ID)

def get_image(folder, image_path):
    """
    Get an image from a managed folder.
    """
    with folder.get_download_stream(image_path) as stream:
        return stream.read()

# -------------------------------------------------------------------------------- NOTEBOOK-CELL: CODE
results = {"filepath": [], "caption": []}

for f in folder.list_paths_in_partition():
    completion = llm.new_completion()
    
    mp_message = completion.new_multipart_message()
    mp_message.with_text("Provide a precise caption of the image.")
    mp_message.with_inline_image(get_image(folder, f))
    mp_message.add()
    
    completion.settings["maxOutputTokens"] = 300
    completion.settings["temperature"] = 0.1
    resp = completion.execute()
    
    results["filepath"].append(f)
    results["caption"].append(resp.text)

df = pd.DataFrame.from_dict(results)

# -------------------------------------------------------------------------------- NOTEBOOK-CELL: CODE
dataiku.Dataset("image_descriptions").write_with_schema(df)