from langchain_google_genai import ChatGoogleGenerativeAI
from langchain_community.tools.tavily_search import TavilySearchResults
from langgraph.prebuilt import create_react_agent
from langgraph.checkpoint.memory import MemorySaver
from langchain_core.messages import HumanMessage, AIMessage
import re

model = ChatGoogleGenerativeAI(
    model="gemini-2.0-flash",
    temperature=0,
)

search_tool = TavilySearchResults(max_results=5)

memory = MemorySaver()
agent = create_react_agent(model, [search_tool], checkpointer=memory)


def run_agent(question: str) -> str:
    prompt = f"""
You are a helpful assistant that answers requested questions using tools.
I will give you a question at the end. Report your thoughts, and give the final answer with the following template:
FINAL ANSWER: [YOUR FINAL ANSWER].
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.

Remember that your answer should start with "FINAL ANSWER: " and be followed by the answer.

The question is:

{question}
"""
    for step in agent.stream(
        {"messages": [HumanMessage(content=prompt)]},
        {"configurable": {"thread_id": "tid1"}},
        stream_mode="values",
    ):
        step["messages"][-1].pretty_print()

        # check if the step is a final answer
        if isinstance(step["messages"][-1], AIMessage):
            # check if the answer is final, and extract it
            final_answer = re.search(
                r"FINAL ANSWER: (.*)",
                step["messages"][-1].content,
            )
            if final_answer:
                return final_answer.group(1).strip()

    return "No final answer found."