"""
This app.py consitutes an application where users can pose queries to
an agent about plants.
"""
from smolagents import CodeAgent,DuckDuckGoSearchTool, HfApiModel,load_tool,tool
import datetime
import requests
import pytz
import yaml
from tools.final_answer import FinalAnswerTool
from huggingface_hub import InferenceClient

from Gradio_UI import GradioUI



@tool
def describe_plant_image(user_query: str, image_url: str) -> str:
    """
    Describe a plant image and answer the user's query.
    Only to be used if and only if a user provides an image_url or if
    prior chat messages have retrieved image_url(s).

    Args:
        user_query: Any relevant text pertaining to the image_url that has been based off the chat
        image_url: The image url that has been provided by the user 
    """
    system_prompt = (
        "You are an LLM assistant that analyzes plant images and:\n"
        "1. Identifies the plant if possible\n"
        "2. Describes key visible characteristics\n"
        "3. Answers the user's question clearly and concisely"
    )
    vl_model = InferenceClient(model="CohereLabs/aya-vision-32b:cohere")

    response = vl_model.chat.completions.create(
        messages=[
            {
                "role": "system",
                "content": system_prompt,
            },
            {
                "role": "user",
                "content": [
                    {"type": "text", "text": user_query},
                    {
                        "type": "image_url",
                        "image_url": {"url": image_url},
                    },
                ],
            },
        ]
    )

    return response.choices[0].message.content

query_internet = DuckDuckGoSearchTool()
final_answer = FinalAnswerTool()

# If the agent does not answer, the model is overloaded, please use another model or the following Hugging Face Endpoint that also contains qwen2.5 coder:
# model_id='https://pflgm2locj2t89co.us-east-1.aws.endpoints.huggingface.cloud' 

model = HfApiModel(
max_tokens=2096,
temperature=0.5,
model_id='Qwen/Qwen2.5-Coder-32B-Instruct',# it is possible that this model may be overloaded
custom_role_conversions=None,
)


# Import tool from Hub
image_generation_tool = load_tool("agents-course/text-to-image", trust_remote_code=True)

with open("prompts.yaml", 'r') as stream:
    prompt_templates = yaml.safe_load(stream)
    
agent = CodeAgent(
    model=model,
    tools=[query_internet, describe_plant_image, final_answer], ## add your tools here (don't remove final answer)
    max_steps=6,
    verbosity_level=1,
    grammar=None,
    planning_interval=None,
    name=None,
    description=None,
    prompt_templates=prompt_templates
)


GradioUI(agent).launch()