src/plant_identifier/identifier.py

import os
import shutil
import tempfile
from typing import List, Union, BinaryIO

from smolagents import CodeAgent, LiteLLMModel

from .plantnet_tool import get_species


def identify_plant(
    images: List[Union[str, BinaryIO]],
    openai_api_key: str,
    plantnet_api_key: str,
    model_id: str = "openai/gpt-4o-mini",
) -> str:
    """
    Identify a plant species from one or more images.

    Parameters
    ----------
    images:
        List of file paths OR file-like objects (opened in binary mode).
    openai_api_key:
        OpenAI API key used by LiteLLMModel.
    plantnet_api_key:
        PlantNet API key used by get_species().
    model_id:
        LiteLLM model identifier.

    Returns
    -------
    str
        A French summary with common name and scientific name.
    """

    # Work in a temporary folder so it runs cleanly on Hugging Face Spaces.
    tmp_dir = tempfile.mkdtemp(prefix="plant_id_")
    saved_files: list[str] = []

    try:
        # Normalise input images to local files
        for idx, img in enumerate(images):
            if isinstance(img, str):
                saved_files.append(img)
            else:
                file_path = os.path.join(tmp_dir, f"uploaded_{idx}.jpg")
                with open(file_path, "wb") as f:
                    shutil.copyfileobj(img, f)
                saved_files.append(file_path)

        model = LiteLLMModel(model_id=model_id, api_key=openai_api_key)
        agent = CodeAgent(tools=[get_species], model=model, add_base_tools=False)

        prompt = f"""
L'API key PlantNet se trouve dans la variable {plantnet_api_key}.
Lance une identification de plante à partir de la liste d'images {saved_files} en utilisant le tool get_species.

Donne une liste à puces des trois meilleurs résultats d'identification de l'espèce végétale, en donnant pour chacun le nom commum, le nom scientifique et la probabilité d'identification sous forme de pourcentage.

Si le score du résultat retourné est inférieur à 0.5, alors retourne l'espèce avec un message d'erreur indiquant que l'identification n'a pas pu être effectuée de manière précise
et qu'il faut réessayer avec plus d'images.
"""
        result = agent.run(prompt)
        return result

    finally:
        # Clean up temp folder
        try:
            shutil.rmtree(tmp_dir)
        except Exception:
            pass