from magentic import FunctionCall, OpenaiChatModel, prompt_chain
import requests

from BasicAgent import BasicAgent

import os


class CustomAgent(BasicAgent):
    def __init__(self):
        super().__init__()
        self._system_prompt = """\
You are a general AI assistant. I will ask you a question. Report your thoughts and provide an answer using, if appropriate, the tools at your disposal.
Your final answer should be a number OR as few words as possible OR a comma separated list of numbers and/or strings. Do not provide anything except your final answer.
If you are asked for a number, don't use comma to write your number neither use units such as $ or percent sign unless specified otherwise. 
If you are asked for a string, don't use articles, neither abbreviations (e.g. for cities), and write the digits in plain text unless specified otherwise.
If you are asked for a comma separated list, apply the above rules depending of whether the element to be put in the list is a number or a string.
Question: {question}
""" # System prompt, adapted from the GAIA team's example at https://huggingface.co/spaces/gaia-benchmark/leaderboard
        self._model = OpenaiChatModel(
            model="llama-3.3-70b-versatile",
            api_key=os.environ["GROQ_API_KEY"],
            base_url="https://api.groq.com/openai/v1/", 
            temperature=0.2, 
            max_tokens=512
        )
        print("CustomAgent initialized.")

    def __call__(self, question: str) -> str:
        """
        this is where the magic happens
        """
        super().__call__(question)
        @prompt_chain(
            template=self._system_prompt,
            model=self._model,
            functions=[self.websearch],
        )
        def _magic(question: str) -> FunctionCall | str:
            """
            An agent, using websearch to gather information and
            answer a user's query.
            """
            ...
        result = _magic(question)
        print("Agent returning actual answer:", result)
        return result
    
    @staticmethod
    def websearch(self, query: str) -> str:
        """
        Performs a web search using an external web search API.
        Args:
            query (str): The search query string.
        Returns:
            str: The response from the web search API as a string.
        """
        try:
            r = requests.post(
                f"{os.environ['WEBSEARCH_API_URL']}/tools/websearch",
                json={"query": query}
            )
            r.raise_for_status()
            response = r.text
        except Exception as e:
            print("Error:", e)
            response = "An error has occurred while using websearch, please inform the user of this"
        print(f"Agent searched for {query}. Response: {response}")