feat: updated the multiagent framework to include arxiv agent.

src/agent_hackathon/multiagent.py

@@ -1,7 +1,9 @@
-import asyncio
+# import asyncio
 from datetime import date
 
+import nest_asyncio
 from llama_index.core.agent.workflow import AgentWorkflow, ReActAgent
+from llama_index.core.tools import FunctionTool
 from llama_index.llms.huggingface_api import HuggingFaceInferenceAPI
 from llama_index.tools.duckduckgo import DuckDuckGoSearchToolSpec
 
@@ -11,6 +13,8 @@ from src.agent_hackathon.consts import PROJECT_ROOT_DIR
 from src.agent_hackathon.generate_arxiv_responses import ArxivResponseGenerator
 from src.agent_hackathon.logger import get_logger
 
+nest_asyncio.apply()
+
 # _ = load_dotenv(dotenv_path=find_dotenv(raise_error_if_not_found=False), override=True)
 
 logger = get_logger(log_name="multiagent", log_dir=PROJECT_ROOT_DIR / "logs")
@@ -28,32 +32,32 @@ class MultiAgentWorkflow:
             # provider="nebius",
             temperature=0.1,
             top_p=0.95,
-            max_tokens=8192,
             # api_key=os.getenv(key="NEBIUS_API_KEY"),
             # base_url="https://api.studio.nebius.com/v1/",
+            system_prompt="Don't just plan, but execute the plan until failure.",
         )
         self._generator = ArxivResponseGenerator(
             vector_store_path=PROJECT_ROOT_DIR / "db/arxiv_docs.db"
         )
-        # self._arxiv_rag_tool = FunctionTool.from_defaults(
-        #     fn=self._arxiv_rag,
-        #     name="arxiv_rag",
-        #     description="Retrieves arxiv research papers.",
-        #     return_direct=True,
-        # )
+        self._arxiv_rag_tool = FunctionTool.from_defaults(
+            fn=self._arxiv_rag,
+            name="arxiv_rag",
+            description="Retrieves arxiv research papers.",
+            return_direct=False,
+        )
         self._duckduckgo_search_tool = [
             tool
             for tool in DuckDuckGoSearchToolSpec().to_tool_list()
             if tool.metadata.name == "duckduckgo_full_search"
         ]
-        # self._arxiv_agent = ReActAgent(
-        #     name="arxiv_agent",
-        #     description="Retrieves information about arxiv research papers",
-        #     system_prompt="You are arxiv research paper agent, who retrieves information "
-        #     "about arxiv research papers.",
-        #     tools=[self._arxiv_rag_tool],
-        #     llm=self.llm,
-        # )
+        self._arxiv_agent = ReActAgent(
+            name="arxiv_agent",
+            description="Retrieves information about arxiv research papers",
+            system_prompt="You are arxiv research paper agent, who retrieves information "
+            "about arxiv research papers.",
+            tools=[self._arxiv_rag_tool],
+            llm=self.llm,
+        )
         self._websearch_agent = ReActAgent(
             name="web_search",
             description="Searches the web",
@@ -63,8 +67,8 @@ class MultiAgentWorkflow:
         )
 
         self._workflow = AgentWorkflow(
-            agents=[self._websearch_agent],
-            root_agent="web_search",
+            agents=[self._arxiv_agent, self._websearch_agent],
+            root_agent="arxiv_agent",
             timeout=180,
         )
         # AgentWorkflow.from_tools_or_functions(
@@ -116,29 +120,27 @@ class MultiAgentWorkflow:
         """
         logger.info("Running multi-agent workflow.")
         try:
-            research_papers = self._arxiv_rag(query=user_query)
             user_msg = (
-                f"search with the web search agent to find any relevant events related to: {user_query}.\n"
-                f" The web search results relevant to the current year: {date.today().year}. \n"
-            )
-            web_search_results = await self._workflow.run(user_msg=user_msg)
-            final_res = (
-                research_papers + "\n\n" + web_search_results.response.blocks[0].text
+                f"First, give me arxiv research papers about: {user_query}."
+                f"Then search with web search agent for any events related to : {user_query}.\n"
+                f"The web search results should be relevant to the current year: {date.today().year}."
+                "Return all the content from all the agents."
             )
+            results = await self._workflow.run(user_msg=user_msg)
             logger.info("Workflow run completed successfully.")
-            return final_res
+            return results
         except Exception as err:
             logger.error(f"Workflow run failed: {err}")
             raise
 
 
-if __name__ == "__main__":
-    USER_QUERY = "i want to learn more about nlp"
-    workflow = MultiAgentWorkflow()
-    logger.info("Starting workflow for user query.")
-    try:
-        result = asyncio.run(workflow.run(user_query=USER_QUERY))
-        logger.info("Workflow finished. Output below:")
-        print(result)
-    except Exception as err:
-        logger.error(f"Error during workflow execution: {err}")
+# if __name__ == "__main__":
+#     USER_QUERY = "i want to learn more about nlp"
+#     workflow = MultiAgentWorkflow()
+#     logger.info("Starting workflow for user query.")
+#     try:
+#         result = asyncio.run(workflow.run(user_query=USER_QUERY))
+#         logger.info("Workflow finished. Output below:")
+#         print(result)
+#     except Exception as err:
+#         logger.error(f"Error during workflow execution: {err}")