Update helper.py

helper.py

@@ -1,6 +1,8 @@
 from langchain_experimental.utilities import PythonREPL
 from langchain.tools import Tool
 from langchain_community.tools import TavilySearchResults
+from langchain_community.document_loaders import WikipediaLoader
+
 import re
 
 # For newer LangChain versions, sometimes it's directly from langchain.tools.python
@@ -234,100 +236,6 @@ from langchain_community.utilities import WikipediaAPIWrapper
 from langchain.tools import Tool # Ensure Tool is imported
 
 
-def get_wikipedia_tool() -> Tool:
-    """
-    Creates and returns a LangChain Tool for querying Wikipedia.
-    """
-    wikipedia_api_wrapper = WikipediaAPIWrapper(
-        top_k_results=3,
-        doc_content_chars_max=4000
-    )
-
-    wikipedia_query_tool = WikipediaQueryRun(api_wrapper=wikipedia_api_wrapper)
-
-    return Tool(
-        name="wikipedia_search_tool",
-        description="""
-        A specialized search tool for retrieving information from Wikipedia.
-        Use this tool when you need:
-        - Authoritative and well-established facts.
-        - Historical information, biographies, or scientific explanations.
-        - General knowledge about specific concepts, people, places, or events.
-        - Background information that is unlikely to be very recent.
-
-        **Input Format (CRITICAL):**
-        The input MUST be a concise and clear query string representing the topic you want to search on Wikipedia.
-        Think of a noun phrase or a short question that directly names the subject.
-        Example: "Albert Einstein"
-        Example: "Battle of Gettysburg"
-        Example: "photosynthesis process"
-
-        **DO NOT:**
-        - Ask natural language questions that are not search queries.
-        - Seek real-time information (e.g., current news, today's weather). For real-time data, use 'tavily_search'.
-        - Provide incomplete sentences or ambiguous terms.
-        - Expect this tool to perform calculations or access external websites beyond Wikipedia.
-
-        **Output:**
-        The tool returns a string containing snippets of relevant Wikipedia articles.
-        The output is limited in length to save tokens. If the answer is not found in the snippet,
-        you might need to refine your query or **consider using 'tavily_search' or other available tools for a broader search.**
-        """,
-        func=wikipedia_query_tool.run,
-    )
-
-wikipedia_search_tool = get_wikipedia_tool()
-
-import wikipedia
-
-def wikipedia_full_content(query: str) -> str:
-    """
-    Fetches the full content of the top Wikipedia article for a query.
-    If a section is specified in the query (e.g., "Mercedes Sosa Discography"),
-    it tries to extract that section.
-    """
-    try:
-        # Try to split query into page and section
-        if " section:" in query:
-            page_query, section = query.split(" section:", 1)
-        else:
-            page_query, section = query, None
-
-        results = wikipedia.search(page_query)
-        if not results:
-            return "No Wikipedia article found for your query."
-        page = wikipedia.page(results[0])
-        content = page.content
-
-        # If a section is specified, try to extract it
-        if section:
-            import re
-            # Simple regex to extract section
-            pattern = rf"==+\s*{re.escape(section.strip())}\s*==+(.*?)(==+|$)"
-            match = re.search(pattern, content, re.DOTALL | re.IGNORECASE)
-            if match:
-                section_content = match.group(1).strip()
-                return section_content[:2000]  # Limit for token safety
-            else:
-                return f"Section '{section}' not found. Returning start of article:\n\n{content[:2000]}"
-        else:
-            return content[:2000]  # Limit for token safety
-
-    except Exception as e:
-        return f"Wikipedia tool error: {e}"
-
-wikipedia_full_content_tool = Tool(
-    name="wikipedia_full_content_tool",
-    description="""
-    Fetches the full content (or a specific section) of a Wikipedia article for a given query.
-    Use this tool for questions about nominations, discographies, lists, or when you need more than a summary.
-    To get a section, use the format: "Page Title section: Section Name"
-    Example: "Mercedes Sosa section: Discography"
-    Example: "Dinosaur featured articles section: Featured article nominations"
-    """,
-    func=wikipedia_full_content,
-)
-
 import os
 from serpapi import GoogleSearch # Or use SerpApiClient for other engines
 from typing import Dict, Any
@@ -539,4 +447,27 @@ gemini_multimodal_tool = Tool(
     name="gemini_multimodal_tool",
     description=analyze_image_with_gemini.__doc__, # Use the docstring as description
     func=analyze_image_with_gemini,
+)
+
+
+def wiki_search(query: str) -> str:
+    """Search Wikipedia for a query and return maximum 2 results.
+    
+    Args:
+        query: The search query.
+    Returns:
+        A string with formatted Wikipedia search 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", "")}" page="{doc.metadata.get("page", "")}"/>\n{doc.page_content}\n</Document>'
+            for doc in search_docs
+        ])
+    return formatted_search_docs
+
+wikipedia_search_tool2 = Tool(
+    name="wikipedia_search_tool2",
+    description=wiki_search.__doc__,
+    func=wiki_search,
 )