Refactor Wikipedia tool to simplify search functionality and remove unused parameters

tools/wikipedia.py

@@ -3,141 +3,12 @@ from langchain_community.document_loaders import WikipediaLoader
 from typing import Optional
 
 @tool
-def fetch_wikipedia_content(query: str, top_k: int = 2, max_chars_per_doc: int = 16000, specific_section: Optional[str] = None) -> dict[str, str]:
-    """Wikipedia search with optional section-based content management.
-    
-    Args:
-        query: The search query.
-        top_k (Optional): The number of top results to return (default: 2).
-        max_chars_per_doc (Optional): Maximum characters per document (default: 16000).
-        specific_section (Optional): If provided, return only this specific section (use section name from sections list).
-    """
-    # Load documents with full content to analyze sections
-    search_docs = WikipediaLoader(
-        query=query, 
-        load_max_docs=top_k, 
-        doc_content_chars_max=128000  # Load full content to analyze sections
-    ).load()
-    
-    formatted_docs = []
-    for doc in search_docs:
-        content = doc.page_content
-        source = doc.metadata.get("source", "")
-        page = doc.metadata.get("page", "")
-        
-        if specific_section:
-            # Return only the requested section
-            section_content = _extract_specific_section(content, specific_section)
-            if section_content:
-                formatted_doc = f'<Document source="{source}" page="{page}" section="{specific_section}"/>\n{section_content}\n</Document>'
-            else:
-                formatted_doc = f'<Document source="{source}" page="{page}"/>\nSection "{specific_section}" not found. Available sections: {_get_section_names(content)}\n</Document>'
-        else:
-            # Process with section-based distribution
-            processed_content, sections_info = _process_content_by_sections(content, max_chars_per_doc)
-            formatted_doc = f'<Document source="{source}" page="{page}"/>\n{sections_info}\n\n{processed_content}\n</Document>'
-        
-        formatted_docs.append(formatted_doc)
-    
-    return {"wiki_results": "\n\n---\n\n".join(formatted_docs)}
-
-def _get_sections(content: str) -> list[tuple[str, str]]:
-    """Parse Wikipedia content into top-level sections only.
-    
-    Returns:
-        List of (section_name, section_content) tuples
+def wiki_search(query: str) -> str:
     """
-    sections = []
-    lines = content.split('\n')
-    current_section = "Introduction"
-    current_content = []
-    
-    for line in lines:
-        # Check if this is a TOP-LEVEL section header (== but not ===)
-        if line.strip().startswith('==') and not line.strip().startswith('==='):
-            # Save previous section
-            if current_content:
-                sections.append((current_section, '\n'.join(current_content)))
-            
-            # Start new section
-            current_section = line.strip().replace('=', '').strip()
-            current_content = []
-        else:
-            current_content.append(line)
-    
-    # Add the last section
-    if current_content:
-        sections.append((current_section, '\n'.join(current_content)))
-    
-    return sections
-
-def _get_section_names(content: str) -> str:
-    """Get a comma-separated list of section names."""
-    sections = _get_sections(content)
-    return ", ".join([section[0] for section in sections])
+    Search Wikipedia for a given query and return the summary.
 
-def _extract_specific_section(content: str, section_name: str) -> Optional[str]:
-    """Extract a specific section from Wikipedia content."""
-    sections = _get_sections(content)
-    
-    # Look for exact match first
-    for name, content_part in sections:
-        if name.lower() == section_name.lower():
-            return content_part
-    
-    # Look for partial match
-    for name, content_part in sections:
-        if section_name.lower() in name.lower():
-            return content_part
-    
-    return None
-
-def _process_content_by_sections(content: str, max_chars: int) -> tuple[str, str]:
-    """Process content by distributing characters equally across sections.
-    
-    Returns:
-        Tuple of (processed_content, sections_info)
+    Args:
+        query (str): The search query to find relevant Wikipedia articles.
     """
-    sections = _get_sections(content)
-    
-    # If content is under limit, return as is
-    total_chars = len(content)
-    if total_chars <= max_chars:
-        sections_info = f"Available sections: {_get_section_names(content)}"
-        return content, sections_info
-    
-    # Calculate chars per section
-    num_sections = len(sections)
-    if num_sections == 0:
-        return content[:max_chars] + "...", "No sections found, content truncated"
-    
-    chars_per_section = max_chars // num_sections
-    reserve_chars = max_chars % num_sections  # Extra chars to distribute
-    
-    processed_sections = []
-    
-    for i, (section_name, section_content) in enumerate(sections):
-        # Give some sections a few extra characters from the reserve
-        section_limit = chars_per_section + (1 if i < reserve_chars else 0)
-        
-        if len(section_content) > section_limit:
-            truncated_content = section_content[:section_limit-3] + "..."
-        else:
-            truncated_content = section_content
-        
-        processed_sections.append(f"== {section_name} ==\n{truncated_content}")
-    
-    sections_info = f"Content truncated due to {max_chars} char limit. Available sections: {_get_section_names(content)}. Use specific_section parameter to load any section in full."
-    
-    return "\n\n".join(processed_sections), sections_info
-
-if __name__ == "__main__":
-    # Example usage
-    query = "Albert Einstein"
-    result = fetch_wikipedia_content.invoke(query, top_k=1, max_chars_per_doc=16000)
-    print(result['wiki_results'])
-
-    # Very small page example
-    query = "Nico Ditch"
-    result = fetch_wikipedia_content.invoke(query, top_k=1, max_chars_per_doc=16000)
-    print(result['wiki_results'])
+    loader = WikipediaLoader(query=query, load_max_docs=2).load()
+    return "\n---\n".join([f'<Document source="{doc.metadata["source"]}"/>\n{doc.page_content}\n</Document>' for doc in loader])