rewards/sources.py

"""Compatibility wrapper for the structured research tool layer.

New code should use ``explainer_env.research.run_research_tool`` directly.
"""

from __future__ import annotations

try:
    from ..research import run_research_tool
except ImportError:  # pragma: no cover - supports direct test execution
    from research import run_research_tool


async def search(
    query: str,
    category_hint: str = "",
    difficulty: str = "easy",
    tool: str | None = None,
) -> str:
    """Run a research tool and render its result as text."""
    selected_tool = tool or _default_tool(query, category_hint, difficulty)
    result = await run_research_tool(selected_tool, query, category_hint)
    return result.render()


def _default_tool(query: str, category_hint: str, difficulty: str) -> str:
    text = f"{query} {category_hint}".lower()
    if any(term in text for term in ("api", "code", "plot", "marimo", "manim")):
        return "fetch_docs"
    if any(term in text for term in ("model", "dataset", "hugging face", "hf hub")):
        return "search_hf_hub"
    if difficulty == "hard":
        return "search_arxiv"
    return "search_wikipedia"