tools.py

"""Tools available to the ReAct agent.

Three tools from the paper + a safe calculator:

    Search[entity]   — fetch the Wikipedia intro paragraph for `entity`
    Lookup[keyword]  — in the last Search result, find the next sentence
                       containing `keyword`  (like browser Ctrl+F)
    Calculate[expr]  — safely evaluate an arithmetic expression
    Finish[answer]   — end the episode and return `answer`

All network calls use the free MediaWiki API (no key required).
"""

from __future__ import annotations

import ast
import math
import operator
import re
import socket
import urllib.error
import urllib.parse
import urllib.request
import json

# ---------------------------------------------------------------------------
# Wikipedia search
# ---------------------------------------------------------------------------

_WIKI_API = "https://en.wikipedia.org/w/api.php"
_SENTENCE_RE = re.compile(r"(?<=[.!?])\s+")
_WIKI_TIMEOUT = int(__import__("os").environ.get("REACT_WIKI_TIMEOUT", "10"))

# Module-level state so Lookup[] can refer to the last Search[] result.
_last_search_sentences: list[str] = []
_lookup_index: int = 0


def _wiki_search(query: str, sentences: int = 5) -> str:
    """Return the first `sentences` sentences of the Wikipedia intro for `query`."""
    global _last_search_sentences, _lookup_index

    params = urllib.parse.urlencode({
        "action": "query",
        "prop": "extracts",
        "exintro": True,
        "explaintext": True,
        "redirects": True,
        "titles": query,
        "format": "json",
    })
    url = f"{_WIKI_API}?{params}"

    try:
        req = urllib.request.Request(url, headers={"User-Agent": "react-agent/0.1"})
        with urllib.request.urlopen(req, timeout=_WIKI_TIMEOUT) as resp:
            data = json.loads(resp.read().decode())
    except socket.timeout:
        return f"Wikipedia request timed out after {_WIKI_TIMEOUT}s for '{query}'."
    except (urllib.error.URLError, json.JSONDecodeError) as exc:
        return f"Could not reach Wikipedia: {exc}"

    pages = data.get("query", {}).get("pages", {})
    page = next(iter(pages.values()))

    if "missing" in page:
        # Try search-suggest fallback.
        return _wiki_search_suggest(query, sentences)

    text = page.get("extract", "").strip()
    if not text:
        return "No extract found for this article."

    # Split into sentences for Lookup[].
    all_sents = [s.strip() for s in _SENTENCE_RE.split(text) if s.strip()]
    _last_search_sentences = all_sents
    _lookup_index = 0

    snippet = " ".join(all_sents[:sentences])
    return snippet or text[:500]


def _wiki_search_suggest(query: str, sentences: int) -> str:
    """OpenSearch fallback: find the closest article title, then fetch it."""
    params = urllib.parse.urlencode({
        "action": "opensearch",
        "search": query,
        "limit": 1,
        "format": "json",
    })
    url = f"{_WIKI_API}?{params}"
    try:
        req = urllib.request.Request(url, headers={"User-Agent": "react-agent/0.1"})
        with urllib.request.urlopen(req, timeout=_WIKI_TIMEOUT) as resp:
            results = json.loads(resp.read().decode())
        if results[1]:
            return _wiki_search(results[1][0], sentences)
    except Exception:
        pass
    return f"Could not find a Wikipedia article for '{query}'."


def _wiki_lookup(keyword: str) -> str:
    """Return the next sentence in the last Search result containing `keyword`."""
    global _lookup_index

    if not _last_search_sentences:
        return "No previous Search result to look up in."

    keyword_lower = keyword.lower()
    for i in range(_lookup_index, len(_last_search_sentences)):
        if keyword_lower in _last_search_sentences[i].lower():
            _lookup_index = i + 1
            return _last_search_sentences[i]

    return f"No (more) results for '{keyword}' in the current article."


# ---------------------------------------------------------------------------
# Safe calculator
# ---------------------------------------------------------------------------

_SAFE_OPS = {
    ast.Add: operator.add,
    ast.Sub: operator.sub,
    ast.Mult: operator.mul,
    ast.Div: operator.truediv,
    ast.Pow: operator.pow,
    ast.USub: operator.neg,
    ast.UAdd: operator.pos,
    ast.Mod: operator.mod,
    ast.FloorDiv: operator.floordiv,
}

_SAFE_NAMES = {
    "abs": abs, "round": round, "sqrt": math.sqrt,
    "pi": math.pi, "e": math.e,
}


def _eval_node(node: ast.AST) -> float:
    if isinstance(node, ast.Constant):
        if isinstance(node.value, (int, float)):
            return float(node.value)
        raise ValueError(f"Unsupported constant: {node.value!r}")
    if isinstance(node, ast.Name):
        if node.id in _SAFE_NAMES:
            v = _SAFE_NAMES[node.id]
            if callable(v):
                raise ValueError(f"'{node.id}' needs arguments")
            return float(v)  # type: ignore[arg-type]
        raise ValueError(f"Unknown name: '{node.id}'")
    if isinstance(node, ast.BinOp):
        op = _SAFE_OPS.get(type(node.op))
        if op is None:
            raise ValueError(f"Unsupported operator: {type(node.op).__name__}")
        return op(_eval_node(node.left), _eval_node(node.right))
    if isinstance(node, ast.UnaryOp):
        op = _SAFE_OPS.get(type(node.op))
        if op is None:
            raise ValueError(f"Unsupported unary op: {type(node.op).__name__}")
        return op(_eval_node(node.operand))
    if isinstance(node, ast.Call):
        if isinstance(node.func, ast.Name) and node.func.id in _SAFE_NAMES:
            fn = _SAFE_NAMES[node.func.id]
            if callable(fn):
                args = [_eval_node(a) for a in node.args]
                return float(fn(*args))  # type: ignore[operator]
        raise ValueError(f"Disallowed function call: {ast.dump(node.func)}")
    raise ValueError(f"Unsupported AST node: {type(node).__name__}")


def _calculate(expr: str) -> str:
    expr = expr.strip().rstrip("=").strip()
    # Allow "x,xxx" comma-separated thousands.
    expr = expr.replace(",", "")
    try:
        tree = ast.parse(expr, mode="eval")
        result = _eval_node(tree.body)
        # Pretty-print: show int if exact, else float.
        if result == int(result) and abs(result) < 1e15:
            return str(int(result))
        return f"{result:.6g}"
    except Exception as exc:
        return f"Error: {exc}"


# ---------------------------------------------------------------------------
# Public dispatch
# ---------------------------------------------------------------------------

class ToolResult:
    def __init__(self, observation: str, done: bool = False, answer: str = ""):
        self.observation = observation
        self.done = done          # True when Finish[] is called
        self.answer = answer      # Set when done=True


_VALID_TOOLS = ("search", "lookup", "calculate", "finish")


def execute(action_name: str, action_input: str) -> ToolResult:
    """Dispatch an action and return a ToolResult.

    Tool names are matched case-insensitively. An empty action_input is
    handled gracefully by each tool rather than raising an exception.
    """
    name = action_name.strip().lower()
    inp = action_input.strip()

    if name == "search":
        if not inp:
            return ToolResult("Search requires a non-empty query.")
        obs = _wiki_search(inp)
        return ToolResult(obs)

    if name == "lookup":
        obs = _wiki_lookup(inp)
        return ToolResult(obs)

    if name == "calculate":
        if not inp:
            return ToolResult("Calculate requires a non-empty expression.")
        obs = _calculate(inp)
        return ToolResult(obs)

    if name == "finish":
        return ToolResult(observation="", done=True, answer=inp)

    valid = ", ".join(t.capitalize() for t in _VALID_TOOLS)
    return ToolResult(
        f"Unknown tool '{action_name}'. Valid tools: {valid}."
    )