Upload agent.py

agent.py

@@ -59,6 +59,54 @@ def save_and_read_file(content: str, filename: Optional[str] = None) -> str:
     return f"File saved to {filepath}. You can read this file to process its contents."
 
 
+@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."""
+    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 arvix_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 {"arvix_results": formatted_search_docs}
+
+
 @tool
 def download_file_from_url(url: str, filename: Optional[str] = None) -> str:
     """
@@ -209,25 +257,27 @@ search_tool = DuckDuckGoSearchResults()
 #     temperature=0.2,
 #     max_tokens=4096,
 # )
-llm = ChatHuggingFace(
-    llm=HuggingFaceEndpoint(
-        repo_id="Qwen/Qwen3-4B",
-        # repo_id="meta-llama/Llama-3-70B-Instruct",
-        temperature=0,
-        huggingfacehub_api_token=os.environ["HUGGINGFACEHUB_API_TOKEN"],
-    ),
-    verbose=True,
-)
-# llm = ChatGoogleGenerativeAI(
-#     model="gemini-2.0-flash-exp", google_api_key=os.environ["GOOGLE_API_KEY"]
+# llm = ChatHuggingFace(
+#     llm=HuggingFaceEndpoint(
+#         repo_id="Qwen/Qwen3-4B",
+#         # repo_id="meta-llama/Llama-3-70B-Instruct",
+#         temperature=0,
+#         huggingfacehub_api_token=os.environ["HUGGINGFACEHUB_API_TOKEN"],
+#     ),
+#     verbose=True,
 # )
+llm = ChatGoogleGenerativeAI(
+    model="gemini-2.0-flash-exp", google_api_key=os.environ["GOOGLE_API_KEY"]
+)
 tools = [
     analyze_csv_file,
     analyze_excel_file,
     extract_text_from_image,
     download_file_from_url,
     save_and_read_file,
-    # search_tool,
+    web_search,
+    wiki_search,
+    arvix_search,
 ]
 # Bind the tools to the LLM
 model_with_tools = llm.bind_tools(tools)