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
#import litellm
#litellm._turn_on_debug()

# Создание модели
model_id = "Qwen/Qwen2.5-Coder-32B-Instruct"
hf_token = os.getenv('hf_token')
model = HfApiModel(model_id, token=hf_token)

#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,
#    # api_base="https://api.sambanova.ai/v1",
#    num_ctx=8192
#)

# Создание инструмента для посещения веб-страниц
class WebpageVisitorTool(Tool):
    name = "webpage_visitor"
    description = "Этот инструмент посещает веб-страницу и возвращает ее содержимое в формате Markdown."

    inputs = {
        "url": {
            "type": "string",
            "description": "URL веб-страницы, которую нужно посетить.",
        }
    }

    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)}"

# Создание веб-агента
#web_agent = ToolCallingAgent(
web_agent = CodeAgent(
    tools=[DuckDuckGoSearchTool(), WebpageVisitorTool()],
    model=model,
    max_steps=10,
    name="web_search_agent",
    description="Выполняет поиск в интернете"
)

# Создание менеджера-агента
manager_agent = CodeAgent(
    tools=[],
    model=model,
    managed_agents=[web_agent],
    additional_authorized_imports=["time", "numpy", "pandas"]
)

# Запуск системы
answer = manager_agent.run("Если обучение моделей языка продолжит масштабироваться с текущим темпом до 2030 года, какова будет потребляемая электрическая мощность в ГВт, необходимая для питания крупнейших тренировочных запусков к 2030 году? Что это будет соответствовать, по сравнению с некоторыми странами? Пожалуйста, предоставьте источник для любых использованных чисел.")
print(answer)