import asyncio
import aiohttp
from ai_query import stream_text
from ai_query.providers import google, tool, Field
@tool(description="Search Wikipedia for articles matching a query")
async def search_wikipedia(
query: str = Field(description="The search query")
) -> str:
"""Search Wikipedia and return article titles."""
url = "https://en.wikipedia.org/w/api.php"
params = {
"action": "query",
"list": "search",
"srsearch": query,
"format": "json",
"srlimit": 5
}
async with aiohttp.ClientSession() as session:
async with session.get(url, params=params) as response:
data = await response.json()
results = data.get("query", {}).get("search", [])
if not results:
return f"No results found for '{query}'"
titles = [r["title"] for r in results]
return f"Found articles: {', '.join(titles)}"
@tool(description="Get the content of a Wikipedia article")
async def get_article(
title: str = Field(description="The exact article title")
) -> str:
"""Fetch the introduction of a Wikipedia article."""
url = "https://en.wikipedia.org/w/api.php"
params = {
"action": "query",
"titles": title,
"prop": "extracts",
"exintro": True,
"explaintext": True,
"format": "json"
}
async with aiohttp.ClientSession() as session:
async with session.get(url, params=params) as response:
data = await response.json()
pages = data.get("query", {}).get("pages", {})
for page in pages.values():
extract = page.get("extract", "")
if extract:
# Truncate to avoid token limits
return extract[:2000]
return f"Could not find article: {title}"
async def main():
print("Wikipedia Research Agent")
print("=" * 40)
result = stream_text(
model=google("gemini-2.0-flash"),
system="""You are a research assistant. When asked about a topic:
1. Search Wikipedia for relevant articles
2. Read the most relevant article(s)
3. Provide a comprehensive summary""",
prompt="Tell me about the history of Python programming language",
tools={
"search_wikipedia": search_wikipedia,
"get_article": get_article
}
)
async for chunk in result.text_stream:
print(chunk, end="", flush=True)
usage = await result.usage
print(f"\n\n[Tokens used: {usage.total_tokens}]")
if __name__ == "__main__":
asyncio.run(main())
