Update agent.py

agent.py

@@ -36,38 +36,66 @@ def web_search(query: str) -> str:
     """Search DuckDuckGo for a query and return maximum 3 results.
     Args:
         query: The search query."""
-    search_results = DuckDuckGoSearchResults(max_results=3).invoke(query=query)
-    # DuckDuckGo returns results as a string, so we need to parse it
-    # The format is typically [snippet: text, title: title, link: url]
-    formatted_results = []
-    
-    # Remove outer quotes if present and format properly
-    if isinstance(search_results, str):
-        # Simple parsing of the DuckDuckGo result format
-        result_entries = search_results.split('], [')
-        for entry in result_entries:
-            # Clean up the entry
-            entry = entry.replace('[', '').replace(']', '').strip()
-            if entry:
-                # Extract components
-                parts = {}
-                for part in entry.split(', '):
-                    if ': ' in part:
-                        key, value = part.split(': ', 1)
-                        parts[key] = value
-                
-                # Format as document
-                if 'snippet' in parts and 'link' in parts:
-                    doc_content = parts.get('snippet', '')
-                    doc_source = parts.get('link', '')
-                    formatted_results.append(
-                        f'<Document source="{doc_source}" page=""/>\n{doc_content}\n</Document>'
-                    )
-    
-    if not formatted_results:
-        return {"web_results": "No results found"}
-    
-    return {"web_results": "\n\n---\n\n".join(formatted_results)}
+    try:
+        search_results = DuckDuckGoSearchResults(max_results=3).invoke(query=query)
+        # DuckDuckGo returns results as a string, so we need to parse it
+        # The format is typically [snippet: text, title: title, link: url]
+        formatted_results = []
+        
+        # Remove outer quotes if present and format properly
+        if isinstance(search_results, str):
+            # Simple parsing of the DuckDuckGo result format
+            result_entries = search_results.split('], [')
+            for entry in result_entries:
+                # Clean up the entry
+                entry = entry.replace('[', '').replace(']', '').strip()
+                if entry:
+                    # Extract components
+                    parts = {}
+                    for part in entry.split(', '):
+                        if ': ' in part:
+                            key, value = part.split(': ', 1)
+                            parts[key] = value
+                    
+                    # Format as document
+                    if 'snippet' in parts and 'link' in parts:
+                        doc_content = parts.get('snippet', '')
+                        doc_source = parts.get('link', '')
+                        formatted_results.append(
+                            f'<Document source="{doc_source}" page=""/>\n{doc_content}\n</Document>'
+                        )
+        
+        if not formatted_results:
+            return {"web_results": "No results found for the query. Try using wikipedia_search instead."}
+        
+        return {"web_results": "\n\n---\n\n".join(formatted_results)}
+    except Exception as e:
+        return {"web_results": f"Error performing web search: {str(e)}. Try using wikipedia_search instead."}
+
+@tool
+def wikipedia_search(query: str) -> str:
+    """Search Wikipedia for a query.
+    Args:
+        query: The search query."""
+    try:
+        # Use WikipediaLoader which works better in restricted environments
+        docs = WikipediaLoader(query=query, load_max_docs=2).load()
+        
+        if not docs:
+            return {"wikipedia_results": "No Wikipedia results found for the query."}
+        
+        formatted_results = []
+        for doc in docs:
+            title = doc.metadata.get('title', 'Unknown Title')
+            source = doc.metadata.get('source', 'Wikipedia')
+            content = doc.page_content[:1000] + "..." if len(doc.page_content) > 1000 else doc.page_content
+            formatted_results.append(
+                f'<Document source="{source}" title="{title}">\n{content}\n</Document>'
+            )
+        
+        return {"wikipedia_results": "\n\n---\n\n".join(formatted_results)}
+    except Exception as e:
+        return {"wikipedia_results": f"Error searching Wikipedia: {str(e)}."}
 
 ### =============== DOCUMENT PROCESSING TOOLS =============== ###
 # File handling still requires external tools
@@ -103,6 +131,7 @@ sys_msg = SystemMessage(content=system_prompt)
 
 tools = [
     web_search,
+    wikipedia_search,
     save_and_read_file,
 ]