Create gaia_agent.py

gaia_agent.py

@@ -0,0 +1,200 @@
+from langchain_ollama import ChatOllama
+from langchain_core.messages import SystemMessage, HumanMessage
+from langgraph.graph import START, StateGraph, MessagesState
+from langgraph.prebuilt import ToolNode, tools_condition
+from langchain_community.tools import DuckDuckGoSearchRun
+from langchain_community.tools.tavily_search import TavilySearchResults
+from langchain_community.document_loaders import WikipediaLoader
+from langchain_community.document_loaders import ArxivLoader
+from langchain_core.tools import tool
+from langchain_huggingface import HuggingFaceEndpoint, ChatHuggingFace
+from math import sqrt
+
+### =============== MATHEMATICAL TOOLS =============== ###
+
+
+@tool
+def multiply(a: float, b: float) -> float:
+    """
+    Multiplies two numbers.
+    Args:
+        a (float): the first number
+        b (float): the second number
+    """
+    return a * b
+
+
+@tool
+def add(a: float, b: float) -> float:
+    """
+    Adds two numbers.
+    Args:
+        a (float): the first number
+        b (float): the second number
+    """
+    return a + b
+
+
+@tool
+def subtract(a: float, b: float) -> int:
+    """
+    Subtracts two numbers.
+    Args:
+        a (float): the first number
+        b (float): the second number
+    """
+    return a - b
+
+
+@tool
+def divide(a: float, b: float) -> float:
+    """
+    Divides two numbers.
+    Args:
+        a (float): the first float number
+        b (float): the second float number
+    """
+    if b == 0:
+        raise ValueError("Cannot divided by zero.")
+    return a / b
+
+
+@tool
+def modulus(a: int, b: int) -> int:
+    """
+    Get the modulus of two numbers.
+    Args:
+        a (int): the first number
+        b (int): the second number
+    """
+    return a % b
+
+
+@tool
+def power(a: float, b: float) -> float:
+    """
+    Get the power of two numbers.
+    Args:
+        a (float): the first number
+        b (float): the second number
+    """
+    return a**b
+
+
+@tool
+def square_root(a: float) -> float | complex:
+    """
+    Get the square root of a number.
+    Args:
+        a (float): the number to get the square root of
+    """
+    if a >= 0:
+        return a**0.5
+    return sqrt(a)
+
+
+### =============== BROWSER TOOLS =============== ###
+
+search_tool = DuckDuckGoSearchRun()
+
+
+@tool
+def wiki_search(query: str) -> str:
+    """Search Wikipedia for a query and return maximum 2 results.
+    Args:
+        query: The search query."""
+    search_docs = WikipediaLoader(query=query, load_max_docs=2).load()
+    formatted_search_docs = "\n\n---\n\n".join(
+        [
+            f'<Document source="{doc.metadata["source"]}" page="{doc.metadata.get("page", "")}"/>\n{doc.page_content}\n</Document>'
+            for doc in search_docs
+        ]
+    )
+    return {"wiki_results": formatted_search_docs}
+
+
+@tool
+def web_search(query: str) -> str:
+    """Search Tavily for a query and return maximum 3 results.
+    Args:
+        query: The search query."""
+    print(f"I'm running web_search with {query = }")
+    search_docs = TavilySearchResults(max_results=3).invoke(query=query)
+    formatted_search_docs = "\n\n---\n\n".join(
+        [
+            f'<Document source="{doc.metadata["source"]}" page="{doc.metadata.get("page", "")}"/>\n{doc.page_content}\n</Document>'
+            for doc in search_docs
+        ]
+    )
+    return {"web_results": formatted_search_docs}
+
+
+@tool
+def arxiv_search(query: str) -> str:
+    """Search Arxiv for a query and return maximum 3 result.
+    Args:
+        query: The search query."""
+    search_docs = ArxivLoader(query=query, load_max_docs=3).load()
+    formatted_search_docs = "\n\n---\n\n".join(
+        [
+            f'<Document source="{doc.metadata["source"]}" page="{doc.metadata.get("page", "")}"/>\n{doc.page_content[:1000]}\n</Document>'
+            for doc in search_docs
+        ]
+    )
+    return {"arxiv_results": formatted_search_docs}
+
+
+tools = [multiply, add, subtract, divide, modulus, power, square_root, web_search, arxiv_search, wiki_search]
+
+
+GAIA_SYSTEM_PROMPT = """
+You are a general AI assistant. I will ask you a question. Report your thoughts, and finish your 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.
+"""
+sys_msg = SystemMessage(content=GAIA_SYSTEM_PROMPT)
+
+
+# Build graph function
+def build_graph(provider: str = "ollama"):
+    """Build the graph"""
+    # Load environment variables from .env file
+    if provider == "ollama":
+        chat = ChatOllama(model="llama3.1")
+    elif provider == "huggingface":
+        llm = HuggingFaceEndpoint(
+            repo_id="Qwen/Qwen2.5-Coder-32B-Instruct"
+        )
+
+        chat = ChatHuggingFace(llm=llm, verbose=True)
+    else:
+        raise ValueError("Invalid provider. Choose 'ollama' or 'huggingface'.")
+    # Bind tools to LLM
+    chat_with_tools = chat.bind_tools(tools)
+
+    # Node
+    def assistant(state: MessagesState):
+        """Assistant node"""
+        print([sys_msg] + state["messages"])
+        return {"messages": [chat_with_tools.invoke([sys_msg] + state["messages"])]}
+
+    builder = StateGraph(MessagesState)
+    builder.add_node("assistant", assistant)
+    builder.add_node("tools", ToolNode(tools))
+    builder.add_edge(START, "assistant")
+    builder.add_conditional_edges(
+        "assistant",
+        tools_condition,
+    )
+    builder.add_edge("tools", "assistant")
+
+    return builder.compile()
+
+
+# test
+if __name__ == "__main__":
+    question = "Examine the video at https://www.youtube.com/watch?v=1htKBjuUWec.\n\nWhat does Teal'c say in response to the question \"Isn't that hot?\""
+    # fixed_answer = "extremely"
+    graph = build_graph(provider="huggingface")
+    messages = [HumanMessage(content=question)]
+    messages = graph.invoke({"messages": messages})
+    for m in messages["messages"]:
+        m.pretty_print()