Update tools.py

tools.py

@@ -1,5 +1,6 @@
 from langchain_core.tools import tool
 import wikipediaapi
+import pandas as pd
 
 
 @tool
@@ -51,23 +52,39 @@ def divide(a: int, b: int) -> int:
 @tool
 def search_wikipedia(page_title: str, language: str) -> str:
     """
-    This tool allows searching wikipedia for information on a specific topic. It will return the resulting information in HTML format.
-    Information may be incomplete in some languages, so try to infer the best language for searching the information from the query.
-    For example, if you are looking for information on a spanish actor, search on the wikipedia in spanish.
+    This tool searches Wikipedia for a specific page and returns its text and any HTML tables it contains.
 
     Args:
-        page_title: The page title to search for
-        language: The language for searching the content. For example, "en" for english, "es" for spanish,  "fr" for french, etc.
+        page_title: Title of the Wikipedia page.
+        language: Language code (e.g., "en", "es", "fr").
 
-    Example:
-        {"message": "How many albums did Juanes produce between 2000 and 2005 ?"}
-        search_wikipedia("Juanes", "es")
+    Returns:
+        A string containing the page title, text, and any extracted tables in markdown format.
     """
-
     try:
-        wiki_wiki = wikipediaapi.Wikipedia('AIAgent (gabriel_abilleira@tutanota.com)', language, extract_format=wikipediaapi.ExtractFormat.HTML)
+        wiki_wiki = wikipediaapi.Wikipedia(
+            user_agent='AIAgent (gabriel_abilleira@tutanota.com)',
+            language=language,
+            extract_format=wikipediaapi.ExtractFormat.HTML
+        )
+
         page = wiki_wiki.page(page_title)
-    
-        return f"Title: {page.title}\n\n{page.text}"
+
+        if not page.exists():
+            return f"Error: Page '{page_title}' not found in language '{language}'."
+
+        # Use the URL to read tables
+        tables = pd.read_html(page.fullurl)
+        markdown_tables = []
+
+        for i, table in enumerate(tables):
+            if isinstance(table, pd.DataFrame):
+                markdown = table.to_markdown(index=False)
+                markdown_tables.append(f"\n---\n**Table {i + 1}:**\n{markdown}")
+
+        table_output = "\n".join(markdown_tables) if markdown_tables else "No tables found on this page."
+
+        return f"Title: {page.title}\n\nText: {page.summary}\n\n{table_output}"
+
     except Exception as e:
-        return f"Error: {str(e)}"
+        return f"Error retrieving Wikipedia content: {str(e)}"