agent/tools.py

"""
Tool definitions (HF/OpenAI schema) and dispatcher for the VL agent.

Three tools:
  - search_web:    search the internet and return snippets
  - final_output:  signal the agent loop to stop and return an answer
  - abort:         signal the agent loop to stop and report a failure/reason
"""

import json
import logging
import os
import urllib.error
import urllib.parse
import urllib.request

from dotenv import load_dotenv

load_dotenv()

_BRAVE_WEB_SEARCH_URL = "https://api.search.brave.com/res/v1/web/search"

logger = logging.getLogger(__name__)
if not logger.handlers:
    _handler = logging.StreamHandler()
    _handler.setFormatter(logging.Formatter("%(levelname)s %(name)s: %(message)s"))
    logger.addHandler(_handler)
logger.setLevel(logging.INFO)
logger.propagate = False

# ---------------------------------------------------------------------------
# Tool schemas (passed to InferenceClient.chat_completion(tools=...))
# ---------------------------------------------------------------------------

AGENT_TOOLS = [
    {
        "type": "function",
        "function": {
            "name": "search_web",
            "description": (
                "Search the internet for up-to-date information. "
                "Use this when you need current facts, news, or data that you are "
                "not confident about from your training knowledge."
            ),
            "parameters": {
                "type": "object",
                "properties": {
                    "query": {
                        "type": "string",
                        "description": "The search query to look up on the web.",
                    }
                },
                "required": ["query"],
            },
        },
    },
    {
        "type": "function",
        "function": {
            "name": "final_output",
            "description": (
                "Deliver the final answer to the user. "
                "Call this once you have gathered enough information and are ready "
                "to give a complete, accurate response. The 'answer' field will be "
                "shown directly to the user."
            ),
            "parameters": {
                "type": "object",
                "properties": {
                    "answer": {
                        "type": "string",
                        "description": "The complete final answer to present to the user.",
                    }
                },
                "required": ["answer"],
            },
        },
    },
    {
        "type": "function",
        "function": {
            "name": "abort",
            "description": (
                "Abort the current task when it cannot be completed. "
                "Use this if the task is impossible, unsafe, or the user asked to stop."
            ),
            "parameters": {
                "type": "object",
                "properties": {
                    "reason": {
                        "type": "string",
                        "description": "Explanation of why the task is being aborted.",
                    }
                },
                "required": [],
            },
        },
    },
]


# ---------------------------------------------------------------------------
# Tool implementations
# ---------------------------------------------------------------------------

def _brave_api_key() -> str | None:
    return os.environ.get("BRAVE_SEARCH_API_KEY") or os.environ.get("BRAVE_API_KEY")


def _search_web(query: str, max_results: int = 5) -> str:
    """Run a Brave Web Search and return formatted snippets."""
    count = max(1, min(max_results, 20))
    logger.info("search_web input: query=%r max_results=%d", query, count)

    def _finish(out: str) -> str:
        logger.info("search_web output:\n%s", out)
        return out

    token = _brave_api_key()
    if not token:
        return _finish(
            "Search unavailable: set BRAVE_SEARCH_API_KEY or BRAVE_API_KEY "
            "in a .env file or the process environment."
        )

    params = urllib.parse.urlencode({"q": query, "count": str(count)})
    url = f"{_BRAVE_WEB_SEARCH_URL}?{params}"
    req = urllib.request.Request(
        url,
        headers={
            "X-Subscription-Token": token,
            "Accept": "application/json",
        },
        method="GET",
    )
    try:
        with urllib.request.urlopen(req, timeout=30) as resp:
            payload = json.loads(resp.read().decode())
    except urllib.error.HTTPError as exc:
        detail = exc.read().decode(errors="replace")[:800]
        return _finish(f"Brave Search API error ({exc.code}): {detail}")
    except urllib.error.URLError as exc:
        return _finish(f"Search request failed: {exc.reason or exc}")
    except json.JSONDecodeError as exc:
        return _finish(f"Search returned invalid JSON: {exc}")

    results = (payload.get("web") or {}).get("results") or []
    if not results:
        return _finish("No results found for the given query.")

    blocks = []
    for item in results[:count]:
        title = item.get("title") or ""
        body = item.get("description") or ""
        href = item.get("url") or ""
        blocks.append(f"**{title}**\n{body}\nSource: {href}")

    return _finish("\n\n---\n\n".join(blocks))


def _final_output(answer: str) -> str:
    """No-op implementation; the orchestrator reads 'answer' directly."""
    return "Answer delivered."


def _abort(reason: str = "") -> str:
    """No-op implementation; the orchestrator reads 'reason' directly."""
    return f"Aborted: {reason}" if reason else "Task aborted."


# ---------------------------------------------------------------------------
# Dispatcher
# ---------------------------------------------------------------------------

def dispatch_tool(tool_name: str, arguments: dict | str) -> str:
    """
    Call the named tool with the given arguments and return the result string.

    `arguments` may arrive as a JSON string (from the model) or already as a dict.
    """
    if isinstance(arguments, str):
        try:
            arguments = json.loads(arguments)
        except json.JSONDecodeError:
            arguments = {}

    if tool_name == "search_web":
        query = arguments.get("query", "")
        if not query:
            msg = "Error: 'query' parameter is required for search_web."
            logger.info("search_web input: query=%r (missing)", query)
            logger.info("search_web output:\n%s", msg)
            return msg
        return _search_web(query)

    if tool_name == "final_output":
        answer = arguments.get("answer", "")
        return _final_output(answer)

    if tool_name == "abort":
        reason = arguments.get("reason", "")
        return _abort(reason)

    return f"Error: unknown tool '{tool_name}'."