from smolagents import CodeAgent, DuckDuckGoSearchTool, HfApiModel, load_tool, tool
import yaml
from tools.final_answer import FinalAnswerTool
from Gradio_UI import GradioUI

# Import the DuckDuckGo search tool
search_tool = DuckDuckGoSearchTool(max_results=5)

# Custom summarization tool
@tool
def extract_summary_from_text(text: str) -> str:
    """Extracts a short summary from the given text.
    Args:
        text: The input text to summarize.
    """
    return text[:300] + "..." if len(text) > 300 else text

# Hugging Face API model
model = HfApiModel(
    max_tokens=2096,
    temperature=0.5,
    model_id='Qwen/Qwen2.5-Coder-32B-Instruct',
    custom_role_conversions=None,
)

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

# ✅ Use CodeAgent instead of WebAgent
agent = CodeAgent(
    model=model,
    tools=[search_tool, extract_summary_from_text, FinalAnswerTool()],  # Add tools
    max_steps=6,
    verbosity_level=1,
    grammar=None,
    planning_interval=None,
    name="WebAgent",
    description="A web-searching AI agent",
    prompt_templates=prompt_templates
)

# Launch Gradio UI
GradioUI(agent).launch()