from smolagents import CodeAgent, DuckDuckGoSearchTool , HfApiModel, load_tool, tool
from tools.final_answer import FinalAnswerTool
from tools.play_music import PlaySomeMusicMeastro
from typing import List
from Gradio_UI import GradioUI

import yaml
import wikipedia

@tool
def search_wikipedia_pages(search_string: str) ->  List[str]:
    """This tool helps to search for the names of potential names of wikipedia pages that relate to a search string.
    This doesn't guarantee that every returned string is a perfect, exact page title that wikipedia.page() will accept without any issue.
    This tool sends a query to the Wikipedia API's search endpoint. This endpoint performs a full-text search across Wikipedia's article titles and content.
    This will not retrieve the actual pages, but gives oversight of the different keywords on wikipedia to find more information.
    This tool will output the first 10 strings related to the search string as a list of strings.
    Args:
        search_string: the search string
    """
    try:
        page_titles = wikipedia.search(search_string, results=10)
        return page_titles
    except wikipedia.exceptions.PageError as err:
        s_error = f"PageError: {err}"
    except wikipedia.exceptions.DisambiguationError as err:
        s_error = f"DisambiguationError {err}"
    print (f"An error occured {s_error}")

@tool

def get_wikipedia_page_content(page_name: str, p_summary: bool = True) ->  str:
    """This tool helps to get the content of wikipedia pages by their name.
    If you use this tool based on the output of search_wikipedia_pages 
     than remember to try a next string that was the output of search_wikipedia_pages if you get a error: Page id
    The output is a string with the actual content of the wikipedia page.
    Args:
        page_name: the name of the page in wikipedia
        p_summary: do you want the full content or only the summary of the content
    """
    s_error = ""
    try:
        a_page = wikipedia.page(title=page_name)
        if p_summary: 
            return a_page.summary
        else:
            return a_page.content
    except wikipedia.exceptions.PageError as err:
        s_error = f"PageError: {err}"
    except wikipedia.exceptions.DisambiguationError as err:
        s_error = f"DisambiguationError {err}"
    print (f"An error occured {s_error}")
    
final_answer = FinalAnswerTool()

model = HfApiModel(
            max_tokens=2096,
            temperature=0.5,
            model_id='Qwen/Qwen2.5-Coder-32B-Instruct',# it is possible that this model may be overloaded
            #model_id='https://pflgm2locj2t89co.us-east-1.aws.endpoints.huggingface.cloud', 
            custom_role_conversions=None,
        )

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

agent = CodeAgent(
    model=model,
    tools=[final_answer, search_wikipedia_pages, get_wikipedia_page_content, PlaySomeMusicMeastro(), DuckDuckGoSearchTool()], 
    max_steps=12,
    verbosity_level=1,
    grammar=None,
    planning_interval=None,
    name=None,
    description=None,
    prompt_templates=prompt_templates,
    add_base_tools=True
)

GradioUI(agent).launch()