Update app.py

app.py

@@ -10,7 +10,10 @@ from collections import Counter
 import re
 from io import BytesIO
 from youtube_transcript_api import YouTubeTranscriptApi
-
+from langchain_community.tools.tavily_search import TavilySearchResults
+from langchain_community.document_loaders import WikipediaLoader
+from langchain_community.utilities import WikipediaAPIWrapper
+from langchain_community.document_loaders import ArxivLoader
 
 # (Keep Constants as is)
 # --- Constants ---
@@ -21,17 +24,68 @@ load_dotenv()
 
 
 
-import wikipedia
 
 @tool
-class WikipediaSearchTool(Tool):
-    def use(self, query: str) -> str:
-        try:
-            return wikipedia.summary(query, sentences=2)
-        except wikipedia.exceptions.DisambiguationError as e:
-            return f"Disambiguation Error: {e.options}"
-        except wikipedia.exceptions.PageError:
-            return "No page found."
+def wiki_search(query: str) -> str:
+    """Search Wikipedia for a query and return up to 2 results."""
+    search_docs = WikipediaLoader(query=query, load_max_docs=2).load()
+
+    formatted_search_docs = "\n\n---\n\n".join(
+        [
+            f'<Document source="{doc.metadata.get("source", "Wikipedia")}" page="{doc.metadata.get("page", "")}"/>\n{doc.page_content}\n</Document>'
+            for doc in search_docs
+        ]
+    )
+    return formatted_search_docs
+    
+    
+@tool
+def wikidata_query(query: str) -> str:
+    """
+    Run a SPARQL query on Wikidata and return results.
+    """
+    endpoint_url = "https://query.wikidata.org/sparql"
+    headers = {
+        "Accept": "application/sparql-results+json"
+    }
+    response = requests.get(endpoint_url, headers=headers, params={"query": query})
+    data = response.json()
+    return json.dumps(data, indent=2)
+
+
+@tool
+def web_search(query: str) -> str:
+    """Search Tavily for a query and return up to 3 results."""
+    tavily_key = os.getenv("TAVILY_API_KEY")
+    
+    if not tavily_key:
+        return "Error: Tavily API key not set."
+
+    search_tool = TavilySearchResults(tavily_api_key=tavily_key, max_results=3)
+    search_docs = search_tool.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 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 formatted_search_docs
 
 
 
@@ -173,7 +227,7 @@ class BasicAgent:
         )
 
         search_tool = DuckDuckGoSearchTool()
-        wiki_search_tool = search_wikipedia()
+        wiki_search_tool = wiki_search()
         str_reverse_tool = StringReverseTool()
         keywords_extract_tool = KeywordsExtractorTool()
         speech_to_text_tool = SpeechToTextTool()