from smolagents import CodeAgent, DuckDuckGoSearchTool, HfApiModel, load_tool
import yaml
from tools.final_answer import FinalAnswerTool
from tools.weather import get_weather
from Gradio_UI import GradioUI
import os

# Initialize tools
final_answer = FinalAnswerTool()
duck_duck_go_search = DuckDuckGoSearchTool()

# Get the Hugging Face token from environment variable
hf_token = os.getenv('HUGGINGFACE_TOKEN')
if not hf_token:
    raise ValueError("Please set the HUGGINGFACE_TOKEN environment variable")

model = HfApiModel(
    max_tokens=2096,
    temperature=0.5,
    model_id='Qwen/Qwen2.5-Coder-32B-Instruct',
    custom_role_conversions=None,
    token=hf_token,
)

# Define prompt templates directly in the code
prompt_templates = {
    "system_prompt": """You are a helpful AI assistant that can use tools to solve problems.
You have access to the following tools:
{{tools}}

When you want to use a tool, use this format:
```python
get_weather(location="New York")
```
or
```python
duck_duck_go_search(query="latest news about artificial intelligence")
```

When you have a final answer, use:
```python
final_answer("Your detailed response here")
```""",
    "user": "{{query}}",
    "tool_response": """The result of your tool call was:
{{observation}}

What would you like to do next?""",
    "final_answer": {
        "pre_messages": "Please provide your final answer.",
        "post_messages": "Let me know if you have any other questions!"
    }
}

# We're creating our CodeAgent with multiple tools
agent = CodeAgent(
    model=model,
    tools=[
        get_weather,  # Pass the function directly
        final_answer,
        duck_duck_go_search,
    ],
    max_steps=6,
    verbosity_level=1,
    grammar=None,
    planning_interval=None,
    name=None,
    description=None,
    prompt_templates=prompt_templates
)

GradioUI(agent).launch()