agent.py

import os
import re
import requests
from smolagents import CodeAgent, ToolCallingAgent, WebSearchTool, tool
from smolagents.models.huggingface import HuggingFaceModel
from dotenv import load_dotenv
from markdownify import markdownify
from requests.exceptions import RequestException

def initialize_agent():
    # Load environment variables from .env file
    load_dotenv()

    # 1. Load the model
    # Make sure to set HF_TOKEN in your Hugging Face Space secrets
    model_name = "HuggingFaceH4/zephyr-7b-beta"
    try:
        model = HuggingFaceModel(model=model_name, token=os.environ.get("HF_TOKEN"))
    except Exception as e:
        print(f"Error loading model: {e}")
        model = None

    # 2. Define the tools
    @tool
    def visit_webpage(url: str) -> str:
        """Visits a webpage at the given URL and returns its content as a markdown string.

        Args:
            url: The URL of the webpage to visit.

        Returns:
            The content of the webpage converted to Markdown, or an error message if the request fails.
        """
        try:
            # Send a GET request to the URL
            response = requests.get(url)
            response.raise_for_status()  # Raise an exception for bad status codes

            # Convert the HTML content to Markdown
            markdown_content = markdownify(response.text).strip()

            # Remove multiple line breaks
            markdown_content = re.sub(r"\n{3,}", "\n\n", markdown_content)

            return markdown_content

        except RequestException as e:
            return f"Error fetching the webpage: {str(e)}"
        except Exception as e:
            return f"An unexpected error occurred: {str(e)}"

    # 3. Define the agents
    if model:
        web_agent = ToolCallingAgent(
            tools=[WebSearchTool(), visit_webpage],
            model=model,
            max_steps=10,
            name="web_search_agent",
            description="Runs web searches for you.",
        )

        manager_agent = CodeAgent(
            tools=[],
            model=model,
            managed_agents=[web_agent],
            additional_authorized_imports=["time", "numpy", "pandas"],
        )
        return manager_agent
    else:
        return None