import re
import requests
from markdownify import markdownify
from requests.exceptions import RequestException
from smolagents import CodeAgent, DuckDuckGoSearchTool, FinalAnswerTool, HfApiModel, Tool, tool, VisitWebpageTool, ToolCallingAgent, LiteLLMModel
from huggingface_hub import login
import os
import openai

os.environ['SAMBANOVA_API_KEY'] = os.getenv('sambanova_token')
model = LiteLLMModel(
    model_id="sambanova/Qwen2.5-Coder-32B-Instruct",
    max_tokens=2096,
    temperature=0.1,
    num_ctx=8192
)

# Creating a tool for visiting web pages
class WebpageVisitorTool(Tool):
    name = "webpage_visitor"
    description = "This tool visits a web page and returns its content in Markdown format."

    inputs = {
        "url": {
            "type": "string",
            "description": "URL of the web page to visit.",
        }
    }

    output_type = "string"

    def forward(self, url: str) -> str:
        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)}"

# Creating a web agent
web_agent = CodeAgent(
    tools=[DuckDuckGoSearchTool(), WebpageVisitorTool()],
    model=model,
    max_steps=10,
    name="web_search_agent",
    description="Performs web search"
)

# Creating a manager agent
manager_agent = CodeAgent(
    tools=[],
    model=model,
    managed_agents=[web_agent],
    additional_authorized_imports=["time", "numpy", "pandas"]
)

# Running the system
answer = manager_agent.run("If language model training continues to scale at the current pace until 2030, what will be the electrical power consumption in GW required to power the largest training runs by 2030? What would this correspond to, compared to some countries? Please provide a source for any numbers used.")
print(answer)