feat: :sparkles: create search supervisor

agents.py

@@ -1,4 +1,14 @@
-from tools import general_tools, file_agent_tools, data_agent_tools, math_agent_tools, analyze_video_tools, youtube_transcript_tools
+from tools import (
+    general_tools, 
+    file_agent_tools, 
+    data_agent_tools, 
+    math_agent_tools, 
+    analyze_video_tools, 
+    youtube_transcript_tools, 
+    google_search,
+    wiki_search,
+    arxiv_search
+)
 from langgraph.prebuilt import create_react_agent
 from langgraph.checkpoint.memory import MemorySaver
 from langchain_openai import ChatOpenAI
@@ -59,6 +69,27 @@ transcript_agent = create_react_agent(
     prompt="You analyze audio/speech content in videos. Use tools to get transcripts."
 )
 
+wiki_agent = create_react_agent(
+    model=llm,
+    tools=[wiki_search],
+    name="wiki_analyst", 
+    prompt="You search information from wikipedia."
+)
+
+google_agent = create_react_agent(
+    model=llm,
+    tools=[google_search],
+    name="google_search_analyst", 
+    prompt="You search information from google search."
+)
+
+arxiv_agent = create_react_agent(
+    model=llm,
+    tools=[arxiv_search],
+    name="arxiv_analyst", 
+    prompt="You search information from arxiv."
+)
+
 excel_prompt = """You are a supervisor. You coordinate file_reader, calculator, and data_processor to solve problems step by step. 
 Do not do calculations or file reading yourself, use the tools.
 Report your thoughts, and finish your answer with the following template: FINAL ANSWER: [YOUR FINAL ANSWER]. 
@@ -76,6 +107,15 @@ If you are asked for a number, don't use comma to write your number neither use
 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.
 """
+
+search_analyzer_prompt = """You coordinate different search agents to answer questions. 
+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.
+"""
+
 # Supervisor
 excel_supervisor = create_supervisor(
     [file_agent, math_agent, data_agent],
@@ -88,4 +128,11 @@ video_supervisor = create_supervisor(
     [video_agent, transcript_agent],
     model=llm,
     prompt=video_analyzer_prompt
+).compile()
+
+# search supervisor
+search_supervisor = create_supervisor(
+    [wiki_agent, google_agent, arxiv_agent],
+    model=llm,
+    prompt=search_analyzer_prompt
 ).compile()

qa_graph.py

@@ -32,6 +32,21 @@ def get_file_type(file_path: str) -> str:
     else:
         return "unknown"
 
+def answer_qery(question: str, thread_id: str = "default") -> str:
+    """Ask the agent a question."""
+    config = {"configurable": {"thread_id": thread_id}, "recursion_limit": 50}
+    
+    try:
+        result = video_supervisor.invoke({
+            "messages": [
+                {"role": "user", "content": question}
+            ]
+        })
+        print(result)
+        return result["messages"][-1].content
+    except Exception as e:
+        return f"Error: {str(e)}"
+    
 def ask_question(question: str, thread_id: str = "default") -> str:
     """Ask the agent a question."""
     config = {"configurable": {"thread_id": thread_id}, "recursion_limit": 100}
@@ -87,65 +102,6 @@ def ask_question_youtube(question: Question) -> str:
     print(result)
     return result["messages"][-1].content
 
-test = [
-    {
-    "task_id": "8e867cd7-cff9-4e6c-867a-ff5ddc2550be",
-    "question": "How many studio albums were published by Mercedes Sosa between 2000 and 2009 (included)? You can use the latest 2022 version of english wikipedia.",
-    "Level": "1",
-    "file_name": ""
-  },
-#     {
-#     "task_id": "cca530fc-4052-43b2-b130-b30968d8aa44",
-#     "question": "Review the chess position provided in the image. It is black's turn. Provide the correct next move for black which guarantees a win. Please provide your response in algebraic notation.",
-#     "Level": "1",
-#     "file_name": "cca530fc-4052-43b2-b130-b30968d8aa44.png"
-#   },
-#             {
-#     "task_id": "1f975693-876d-457b-a649-393859e79bf3",
-#     "question": "Hi, I was out sick from my classes on Friday, so I'm trying to figure out what I need to study for my Calculus mid-term next week. My friend from class sent me an audio recording of Professor Willowbrook giving out the recommended reading for the test, but my headphones are broken :(\n\nCould you please listen to the recording for me and tell me the page numbers I'm supposed to go over? I've attached a file called Homework.mp3 that has the recording. Please provide just the page numbers as a comma-delimited list. And please provide the list in ascending order.",
-#     "Level": "1",
-#     "file_name": "1f975693-876d-457b-a649-393859e79bf3.mp3"
-#   },
-#               {
-#     "task_id": "7bd855d8-463d-4ed5-93ca-5fe35145f733",
-#     "question": "The attached Excel file contains the sales of menu items for a local fast-food chain. What were the total sales that the chain made from food (not including drinks)? Express your answer in USD with two decimal places.",
-#     "Level": "1",
-#     "file_name": "7bd855d8-463d-4ed5-93ca-5fe35145f733.xlsx"
-#   },
-#                 {
-#     "task_id": "f918266a-b3e0-4914-865d-4faa564f1aef",
-#     "question": "What is the final numeric output from the attached Python code?",
-#     "Level": "1",
-#     "file_name": "f918266a-b3e0-4914-865d-4faa564f1aef.py"
-#   },
-#                   {
-#     "task_id": "cabe07ed-9eca-40ea-8ead-410ef5e83f91",
-#     "question": "What is the surname of the equine veterinarian mentioned in 1.E Exercises from the chemistry materials licensed by Marisa Alviar-Agnew & Henry Agnew under the CK-12 license in LibreText's Introductory Chemistry materials as compiled 08/21/2023?",
-#     "Level": "1",
-#     "file_name": ""
-#   },
-#                     {
-#     "task_id": "9d191bce-651d-4746-be2d-7ef8ecadb9c2",
-#     "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?\"",
-#     "Level": "1",
-#     "file_name": ""
-#   },
-#                       {
-#     "task_id": "a1e91b78-d3d8-4675-bb8d-62741b4b68a6",
-#     "question": "In the video https://www.youtube.com/watch?v=L1vXCYZAYYM, what is the highest number of bird species to be on camera simultaneously?",
-#     "Level": "1",
-#     "file_name": ""
-#   },        
-            ]
-
-questions = [Question(**item) for item in test]
-for q in questions:
-    print(q.question)
-    print(q.file_name)
-    print(q.local_file_path)
-
-
-
 # State
 class State(TypedDict):
     question: Question
@@ -159,7 +115,7 @@ def ask_question_node(state: State) -> dict:
     thread_id = f"test_{question_obj.task_id}"
     
     # Call your existing function
-    answer = ask_question(question_obj.question, thread_id)
+    answer = answer_qery(question_obj.question, thread_id)
     
     # Return dict to update state
     return {"answer": answer}
@@ -229,17 +185,31 @@ def build_graph():
     return react_graph
 
 if __name__ == "__main__":
-    for i, question in enumerate(questions):
-        print(f"\n{i}. {question.question}")
-        
-        react_graph = build_graph()
-        # Invoke the graph and capture the result
-        result = react_graph.invoke({
-            "question": question,
-            "decision": "",
-            "answer": ""
-        })
-        print(result)
+    test = [
+        {
+        "task_id": "8e867cd7-cff9-4e6c-867a-ff5ddc2550be",
+        "question": "How many studio albums were published by Mercedes Sosa between 2000 and 2009 (included)? You can use the latest 2022 version of english wikipedia.",
+        "Level": "1",
+        "file_name": ""
+        },     
+                ]
+
+    questions = [Question(**item) for item in test]
+    for q in questions:
+        print(q.question)
+        print(q.file_name)
+        print(q.local_file_path)
+        for i, question in enumerate(questions):
+            print(f"\n{i}. {question.question}")
         
-        print(f"Answer: {result['answer']}")
-        print("-" * 50)
+            react_graph = build_graph()
+            # Invoke the graph and capture the result
+            result = react_graph.invoke({
+                "question": question,
+                "decision": "",
+                "answer": ""
+            })
+            print(result)
+            
+            print(f"Answer: {result['answer']}")
+            print("-" * 50)

tools.py

@@ -340,10 +340,22 @@ def google_search():
 )
     google_search = GoogleSearchRun(api_wrapper=api_wrapper)
     return google_search
+
+@tool
+def wiki_search():
+    """Google search tool"""
+    api_wrapper = WikipediaAPIWrapper()
+    search = WikipediaQueryRun(api_wrapper=api_wrapper)
+    return search
+
+@tool
+def arxiv_search():
+    """Google search tool"""
+    api_wrapper = ArxivAPIWrapper()
+    search = ArxivQueryRun(api_wrapper=api_wrapper)
+    return search
 def general_tools():    
     tools = [
-        ArxivQueryRun(api_wrapper=ArxivAPIWrapper()),
-        google_search,
         analyze_image,
         read_python_file,
         transcribe_audio,
@@ -369,3 +381,11 @@ def math_agent_tools():
 def data_agent_tools():
     tools = [extract_values, filter_rows]
     return tools
+
+def search_agen_tools():
+    tools = [
+        google_search,
+        ArxivQueryRun(api_wrapper=ArxivAPIWrapper()),
+        WikipediaQueryRun(api_wrapper=WikipediaAPIWrapper())
+    ]
+    return tools