update llm, search tools

agent.py

@@ -14,23 +14,19 @@ from langchain_core.messages.ai import subtract_usage
 
 from langchain.tools import Tool
 from langchain_core.tools import tool
-from langchain_community.tools import WikipediaQueryRun
-from langchain_community.utilities import WikipediaAPIWrapper
-from langchain_community.tools import DuckDuckGoSearchResults
-from langchain_community.utilities import DuckDuckGoSearchAPIWrapper
-from langchain_community.utilities import ArxivAPIWrapper
+from langchain_community.tools.tavily_search import TavilySearchResults
+from langchain_community.document_loaders import WikipediaLoader
+from langchain_community.document_loaders import ArxivLoader
 from langchain_community.retrievers import BM25Retriever
 
 from langgraph.prebuilt import ToolNode, tools_condition
 
+# load system prompt
+from prompts import system_prompt
+
 
 # load environment variables
 load_dotenv()
-HUGGINGFACEHUB_API_TOKEN = os.getenv("HUGGINGFACEHUB_API_TOKEN")
-print(f"DEBUG: HUGGINGFACEHUB_API_TOKEN = {HUGGINGFACEHUB_API_TOKEN}")
-GOOGLE_API_KEY = os.getenv("GOOGLE_API_KEY")
-print(f"DEBUG: GOOGLE_API_KEY = {GOOGLE_API_KEY}")
-GROQ_API_KEY = os.getenv("GROQ_API_KEY")
 
 
 # maths tool
@@ -92,53 +88,56 @@ def modulus(a:int, b:int) -> int:
 
 # wikipedia search tool
 @tool
-def search_wiki(query: str) -> Dict[str, str]:
-    """search wikipedia with a query
-    args:
-        query: a search query
-    """
-    docs = WikipediaQueryRun(api_wrapper=WikipediaAPIWrapper())
-    docs.run(query)
-    formatted_result = f'<Document source="{docs.metadata["source"]}" page="{docs.metadata.get("page", "")}"/>\n{docs.page_content}\n</Document>'
-    return formatted_result
+def search_wiki(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}
 
 
 # internet search tool
 @tool
-def search_web(query: str) -> Dict[str, str]:
-    """search internet with a query
-    args:
-        query: a search query
-    """
-    wrapper = DuckDuckGoSearchAPIWrapper(region="en-us", max_results=2)
-    docs = DuckDuckGoSearchResults(api_wrapper=wrapper)
-    docs.invoke(query)
-    formatted_result = f'<Document source="{docs.metadata["source"]}" page="{docs.metadata.get("page", "")}"/>\n{docs.page_content}\n</Document>'
-    return formatted_result
+def search_web(query: str) -> str:
+    """Search Tavily for a query and return maximum 3 results.
+
+    Args:
+        query: The search 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}
 
 
 # ArXiv search tool
 @tool
-def search_arxiv(query: str) -> Dict[str, str]:
-    """search ArXiv for the paper with the given identifier
-    args:
-        query: a search identifier
-    """
-    arxiv = ArxivAPIWrapper()
-    docs = arxiv.run(query)
-    formatted_result = f'<Document source="{docs.metadata["source"]}" page="{docs.metadata.get("page", "")}"/>\n{docs.page_content}\n</Document>'
-    return formatted_result
+def search_arxiv(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 {"arvix_results": formatted_search_docs}
 
 
 # build retriever
 # bm25_retriever = BM25Retriever.from_documents(docs)
 
 
-# load system prompt from file
-with open("system_prompt.txt", "r", encoding="utf-8") as f:
-    system_prompt = f.read()
-
-
 # init system message
 sys_msg = SystemMessage(content=system_prompt)
 
@@ -159,7 +158,7 @@ tools = [
 def build_graph():
     # llm
     llm = ChatGroq(
-        model="qwen-qwq-32b",
+        model="meta-llama/llama-4-scout-17b-16e-instruct",
         temperature=0,
     )
     print(f"DEBUG: llm object = {llm}")

system_prompt.txt → prompts.py

@@ -1,5 +1,6 @@
-You are a helpful assistant tasked with answering questions using a set of tools.
+system_prompt = """You are a helpful assistant tasked with answering questions using a set of tools.
 Now, 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.
 Your answer should only start with "FINAL ANSWER: ", then follows with the answer.
+"""

pyproject.toml

@@ -18,5 +18,6 @@ dependencies = [
     "pandas>=2.2.3",
     "rank-bm25>=0.2.2",
     "requests>=2.32.3",
+    "tavily-python>=0.7.2",
     "wikipedia>=1.4.0",
 ]

requirements.txt

@@ -11,4 +11,5 @@ langgraph
 pandas
 rank-bm25
 requests
+tavily-python
 wikipedia