Updated agent.py to accommodate errors when trying to use LlamaIndex WikipediaReader

agent.py

@@ -8,6 +8,8 @@ from llama_index.readers.wikipedia import WikipediaReader
 from llama_index.readers.web import SimpleWebPageReader
 from llama_index.core.schema import Document
 
+import wikipedia
+
 # Load environment variables
 load_dotenv()
 
@@ -63,20 +65,66 @@ def wikipedia_search(query: str, num_results: int = 2) -> str:
         A formatted string with the search results
     """
     try:
+        # First try with LlamaIndex WikipediaReader
         reader = WikipediaReader()
         docs = reader.load_data(query=query, max_docs=num_results)
-        if not docs:
-            return f"No Wikipedia results found for '{query}'."
-        
-        results = []
-        for i, doc in enumerate(docs, 1):
-            title = doc.metadata.get("title", "Unknown Title")
-            content = doc.text[:1000] + "..." if len(doc.text) > 1000 else doc.text
-            results.append(f"Result {i}: {title}\n{content}\n")
-        
-        return "\n".join(results)
+        if docs:
+            results = []
+            for i, doc in enumerate(docs, 1):
+                title = doc.metadata.get("title", "Unknown Title")
+                content = doc.text[:1000] + "..." if len(doc.text) > 1000 else doc.text
+                results.append(f"Result {i}: {title}\n{content}\n")
+            
+            return "\n".join(results)
+        else:
+            # If no results from LlamaIndex, try with direct Wikipedia package
+            print(f"No results from LlamaIndex WikipediaReader for '{query}', trying direct Wikipedia package...")
+            return _fallback_wikipedia_search(query, num_results)
+            
     except Exception as e:
-        return f"Error searching Wikipedia: {str(e)}"
+        print(f"Error with LlamaIndex WikipediaReader: {str(e)}")
+        # Fall back to direct Wikipedia package
+        print(f"Falling back to direct Wikipedia package...")
+        try:
+            return _fallback_wikipedia_search(query, num_results)
+        except Exception as fallback_error:
+            print(f"Fallback also failed: {fallback_error}")
+            return f"Error searching Wikipedia: Unable to retrieve information about '{query}'. Please try a different search term or approach."
+
+def _fallback_wikipedia_search(query: str, num_results: int = 2) -> str:
+    """
+    Fallback implementation using the direct Wikipedia package.
+    """
+    # First search for pages
+    search_results = wikipedia.search(query, results=num_results)
+    
+    if not search_results:
+        return f"No Wikipedia results found for '{query}'."
+    
+    results = []
+    for i, page_title in enumerate(search_results, 1):
+        try:
+            # Get the page content
+            page = wikipedia.page(page_title)
+            title = page.title
+            # Get a summary instead of full content
+            content = page.summary[:1000] + "..." if len(page.summary) > 1000 else page.summary
+            results.append(f"Result {i}: {title}\n{content}\n")
+        except wikipedia.exceptions.DisambiguationError as e:
+            # Handle disambiguation pages
+            options = e.options[:5]  # Limit to 5 options
+            results.append(f"Result {i}: Multiple options found for '{page_title}':\n" + 
+                          "\n".join([f"- {opt}" for opt in options]))
+        except wikipedia.exceptions.PageError:
+            # Skip pages that don't exist
+            continue
+        except Exception as e:
+            results.append(f"Result {i}: Error retrieving information for '{page_title}': {str(e)}")
+    
+    if not results:
+        return f"Could not retrieve valid information for '{query}'."
+    
+    return "\n".join(results)
 
 def web_search(url: str) -> str:
     """