Project: build a research agent

Write three plain Python functions for an AI agent, using the httpx library: (1) search_wikipedia(query: str) calling the MediaWiki API (action=query, list=search) and returning titles and snippets as a dict; (2) read_page(title: str) returning the plain-text introduction of a page; (3) save_note(note: str, source_title: str) storing notes in a module-level list. Each function needs type hints and a clear docstring. Tools must return an error dict instead of raising exceptions. Cap the page text length.

import re
import sys

import httpx
from google import genai
from google.genai import types

client = genai.Client()
MODEL = "gemini-flash-latest"
API = "https://en.wikipedia.org/w/api.php"
# Wikimedia asks every client to send a User-Agent with contact details. Use your own email.
HEADERS = {"User-Agent": "ResearchAgentDemo/0.1 (learning project; contact: your-email@example.com)"}

notes = []  # the agent's scratchpad: (note, source_title)


def search_wikipedia(query: str) -> dict:
    """Search Wikipedia. Returns up to 5 page titles with a short snippet each."""
    params = {"action": "query", "list": "search", "srsearch": query,
              "srlimit": 5, "format": "json"}
    try:
        r = httpx.get(API, params=params, headers=HEADERS, timeout=15)
        r.raise_for_status()
        hits = r.json()["query"]["search"]
    except Exception as e:  # tools should return errors, not crash the loop
        return {"error": f"search failed: {e}"}
    results = [{"title": h["title"], "snippet": re.sub(r"<[^>]+>", "", h["snippet"])} for h in hits]
    return {"results": results}


def read_page(title: str) -> dict:
    """Read the introduction of a Wikipedia page by its exact title."""
    params = {"action": "query", "prop": "extracts", "exintro": 1, "explaintext": 1,
              "titles": title, "format": "json", "redirects": 1}
    try:
        r = httpx.get(API, params=params, headers=HEADERS, timeout=15)
        r.raise_for_status()
        pages = r.json()["query"]["pages"]
        page = next(iter(pages.values()))
    except Exception as e:
        return {"error": f"could not read page: {e}"}
    text = page.get("extract")
    if not text:
        return {"error": f"no page found with title '{title}'. Try search_wikipedia."}
    return {"title": page["title"], "text": text[:4000]}  # cap size to control cost


def save_note(note: str, source_title: str) -> dict:
    """Save one fact you want to use in the final report, with its Wikipedia page title."""
    notes.append((note, source_title))
    return {"saved": True, "notes_so_far": len(notes)}


TOOLS = {f.__name__: f for f in (search_wikipedia, read_page, save_note)}

SYSTEM = """You are a careful research assistant.
Plan first: list 2-4 sub-questions. Then use the tools to answer them.
Rules:
- Only state facts you found with the tools. Save each useful fact with save_note.
- Prefer reading a page over trusting a search snippet.
- When you have enough, stop calling tools and write the final report.
Final report format: a short title, 3-6 bullet points of findings, then a 'Sources' list
of the Wikipedia page titles you used. Say clearly what you could not find."""


def research(question, max_steps=10):
    notes.clear()
    config = types.GenerateContentConfig(
        system_instruction=SYSTEM,
        tools=list(TOOLS.values()),
        automatic_function_calling=types.AutomaticFunctionCallingConfig(disable=True),
    )
    contents = [types.Content(role="user", parts=[types.Part(text=question)])]
    for step in range(1, max_steps + 1):
        resp = client.models.generate_content(model=MODEL, contents=contents, config=config)
        calls = resp.function_calls
        if not calls:
            return resp.text
        contents.append(resp.candidates[0].content)
        results = []
        for call in calls:
            args = dict(call.args)
            print(f"[{step}] {call.name} {args}")
            output = TOOLS[call.name](**args) if call.name in TOOLS else {"error": "unknown tool"}
            results.append(types.Part.from_function_response(name=call.name, response=output))
        contents.append(types.Content(role="user", parts=results))
    # Out of steps: ask for a report from what we have, with no tools.
    contents.append(types.Content(role="user", parts=[types.Part(
        text="Step limit reached. Write the final report now using only your saved notes.")]))
    final = client.models.generate_content(
        model=MODEL, contents=contents,
        config=types.GenerateContentConfig(system_instruction=SYSTEM))
    return final.text


if __name__ == "__main__":
    q = " ".join(sys.argv[1:]) or "How does UPI work and who runs it?"
    print("\n" + research(q))

Here is the system prompt for a research agent that uses Wikipedia tools: [paste SYSTEM]. Act as a strict reviewer. List the ways this agent could (1) loop without finishing, (2) state facts it did not read, (3) waste tokens. For each, suggest one change to the prompt or to the loop code.

python research_agent.py "How does UPI work and who runs it?"

    # inside the loop, replace the line that runs the tool with:
            key = (call.name, str(sorted(args.items())))
            if key in seen:
                output = {"error": "You already did this. Try a different query or finish."}
            else:
                seen.add(key)
                output = TOOLS[call.name](**args) if call.name in TOOLS else {"error": "unknown tool"}

Below is a research report written by an AI agent, followed by the Wikipedia introduction text it read. For each bullet in the report, mark it SUPPORTED, NOT SUPPORTED or PARTLY SUPPORTED by the text, and quote the supporting sentence.

Report:
[paste report]

Source text:
[paste page text]