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. Provide an answer like so: YOUR FINAL ANSWER. DO NOT SAY ANYTHING EXCEPT YOUR FINAL ANSWER. 
Only use the provided tools if your own knowledge would be insufficient. If you are unable to answer the question using the provided tools, give your best guess.
Your final answer should be a number OR as few words as possible OR a comma separated list of numbers and/or strings.
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.1-8b-instant",
            api_key=os.environ["GROQ_API_KEY"],
            base_url="https://api.groq.com/openai/v1/", 
            temperature=0.7, 
            max_tokens=512
        )
        print("CustomAgent initialized.")

    def __call__(self, question: str) -> str:
        """
        this is where the magic happens
        """
        print(f"Agent received question (first 50 chars): {question[:50]}...")
        @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(query: str) -> str:
        """
        Performs a web search using an external web search API. Note that better results may be obtained from varying the wording of the query.
        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}")