from smolagents import CodeAgent, HfApiModel, load_tool
import yaml
from smolagents import tool
from duckduckgo_search import DDGS

from tools.final_answer import FinalAnswerTool

from Gradio_UI import GradioUI


@tool
def DuckDuckGoSearchTool(query: str) -> str:
    """
    Инструмент для поиска информации в интернете с помощью DuckDuckGo.

    Args:
        query: Поисковый запрос.
    """
    with DDGS() as ddgs:
        results = [r for r in ddgs.text(query, max_results=5)]  # Ограничиваем до 5 результатов
        if not results:
            return "По вашему запросу ничего не найдено."

        formatted_results = "\n\n".join(
            f"**Заголовок:** {r['title']}\n**Ссылка:** {r['href']}\n**Краткое содержание:** {r['body']}"
            for r in results
        )
    return formatted_results


final_answer = FinalAnswerTool()

model = HfApiModel(
    max_tokens=2096,
    temperature=0.5,
    model_id='Qwen/Qwen2.5-Coder-32B-Instruct',  # Возможно, эта модель перегружена
    custom_role_conversions=None,
)

# Import tool from Hub
image_generation_tool = load_tool("agents-course/text-to-image", trust_remote_code=True)


with open("prompts.yaml", 'r') as stream:
    prompt_templates = yaml.safe_load(stream)

prompt_templates['system_prompt'] = system_prompt


agent = CodeAgent(
    model=model,
    tools=[DuckDuckGoSearchTool, final_answer, image_generation_tool],
    max_steps=6,
    verbosity_level=1,
    grammar=None,
    planning_interval=None,
    name=None,
    description=None,
    prompt_templates=prompt_templates
)


GradioUI(agent).launch()