from smolagents import CodeAgent, HfApiModel, tool
import requests
import yaml

# Define a research tool using EXA API
@tool
def search_cape_town_weather() -> str:
    """Searches for the latest weather in Cape Town using EXA search API."""
    api_key = "fe5ee653-7f0d-493c-99ae-8e63d46133d9"  # Replace with your free EXA API key
    query = "current weather in Cape Town site:bbc.com OR site:cnn.com OR site:news24.com"
    url = f"https://api.exa.ai/v1/search?q={query}&api_key={api_key}"

    try:
        response = requests.get(url)
        data = response.json()

        if "results" in data and len(data["results"]) > 0:
            top_result = data["results"][0]  # Get the first search result
            title = top_result["title"]
            snippet = top_result["snippet"]
            link = top_result["url"]
            return f"Latest Cape Town weather report:\n{title}\n{snippet}\nMore info: {link}"
        else:
            return "No weather updates found. Try again later."
    except Exception as e:
        return f"Error fetching data: {str(e)}"

# Load model from Hugging Face
model = HfApiModel(
    model_id='mistralai/Mistral-7B-Instruct',  # Change to your preferred model
    max_tokens=200,
    temperature=0.5
)

# Load prompt templates
with open("prompts.yaml", "r") as stream:
    prompt_templates = yaml.safe_load(stream)

# Define the AI Agent
agent = CodeAgent(
    model=model,
    tools=[search_cape_town_weather],  # Add the EXA search tool
    max_steps=3,
    verbosity_level=1,
    prompt_templates=prompt_templates
)

# Test the agent (Optional)
if __name__ == "__main__":
    print(agent.run("Find the latest weather update for Cape Town."))