Spaces:
Sleeping
Sleeping
File size: 6,266 Bytes
db70dbc f1280b9 db70dbc f1280b9 db70dbc f1280b9 db70dbc f1280b9 db70dbc 97858fd |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 |
"""Tool definitions and utility functions for the agent"""
from typing import List, Dict, Any
import os
from dotenv import load_dotenv
from llama_index.core.tools import BaseTool, FunctionTool
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()
# --- Text Processing Tools ---
def text_reverser(text: str) -> str:
"""
Reverse the given text. Useful for answering questions that are written backwards.
Args:
text: The text to reverse
Returns:
The reversed text
"""
return text[::-1]
# --- Math Tools ---
def simple_calculator(operation: str, a: float, b: float) -> float:
"""
Perform a simple calculation.
Args:
operation: One of 'add', 'subtract', 'multiply', 'divide'
a: First number
b: Second number
Returns:
The result of the calculation
"""
if operation == "add":
return a + b
elif operation == "subtract":
return a - b
elif operation == "multiply":
return a * b
elif operation == "divide":
if b == 0:
raise ValueError("Cannot divide by zero")
return a / b
else:
raise ValueError(f"Unknown operation: {operation}")
# --- Information Retrieval Tools ---
def wikipedia_search(query: str, num_results: int = 2) -> str:
"""
Search Wikipedia for information.
Args:
query: The search query
num_results: Number of results to return (default: 2)
Returns:
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 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:
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:
"""
Fetch and extract content from a specific web page.
Args:
url: The URL of the web page to search
Returns:
The extracted content from the web page
"""
try:
reader = SimpleWebPageReader()
docs = reader.load_data(urls=[url])
if not docs:
return f"No content found for URL: {url}"
# Just return the content of the first document
return docs[0].text
except Exception as e:
return f"Error retrieving web page: {str(e)}"
# --- Tool Selection and Routing ---
def get_tools() -> List[BaseTool]:
"""Create and return a list of tools for the agent."""
text_reverser_tool = FunctionTool.from_defaults(
fn=text_reverser,
name="text_reverser",
description="Reverses the given text. Useful for processing reversed questions or text.",
)
calculator_tool = FunctionTool.from_defaults(
fn=simple_calculator,
name="calculator",
description="Performs simple calculations: add, subtract, multiply, divide.",
)
wikipedia_tool = FunctionTool.from_defaults(
fn=wikipedia_search,
name="wikipedia_search",
description="Searches Wikipedia for information on a topic.",
)
web_tool = FunctionTool.from_defaults(
fn=web_search,
name="web_search",
description="Fetches and extracts content from a specific web page.",
)
return [
text_reverser_tool,
calculator_tool,
wikipedia_tool,
web_tool
]
if __name__ == "__main__":
print("This module defines tools for the agent. Run app.py or standalone_debug.py to test the agent.") |