agent/tools.py

"""The three tools the agent loop drives.

- search_docs: hybrid retrieval + hydrate → doc sections (repeatable)
- read_page:   whole-page hydrate (outline-first for oversized pages)
- ask_source:  refer to the real source (GitHub / DeepWiki) — never fabricates

Each returns a plain dict the planner LLM sees. ask_source is referral-only by
design: it points at where the implementation lives rather than claiming to
know it, so it works with no network access.
"""

from __future__ import annotations

from urllib.parse import quote

DEEPWIKI_URL = "https://deepwiki.com/pytorch/pytorch"
GH_CODE_SEARCH = "https://github.com/search?q=repo%3Apytorch%2Fpytorch+{q}&type=code"

# dropped from the GitHub code-search query: they add no signal and crowd out
# the discriminating terms (a code-search URL has a practical length limit)
_STOPWORDS = frozenset(
    "a an the is are was were be to of in on for how do does did i you it "
    "what when where why which that this with and or from can could should "
    "my me use using implement implemented implementation work works".split()
)
_MAX_SEARCH_TERMS = 12  # keep the URL sane while retaining the meaningful words


def _search_terms(question: str) -> str:
    """Meaningful words from the question, URL-encoded for GitHub code search.

    Drops stopwords (keeping code identifiers like torch.nn.Linear verbatim)
    and caps the count so the URL stays sane — without throwing away the words
    that discriminate the query (the old code kept only the first 6 words,
    which dropped the actual subject of longer questions).
    """
    words = [w for w in question.split() if w.lower().strip("?.,:;()") not in _STOPWORDS]
    kept = (words or question.split())[:_MAX_SEARCH_TERMS]
    return quote(" ".join(kept))


# content spaces the planner may restrict a search to (must match ingest's
# page_kind values); anything else from the model is ignored, not an error
SEARCH_KINDS = frozenset({"api", "tutorial", "guide"})


def search_docs(
    query: str, library: str | None = None, kind: str | None = None, k: int = 8
) -> dict:
    """Hybrid docs search over the English-only index.

    `kind` lets the planner choose the content space: 'api' searches only the
    reference pages (catalog questions — "what loss functions exist?"),
    'tutorial'/'guide' only the walkthroughs. Unknown values degrade to an
    unrestricted search rather than failing the tool call.
    """
    from index.hydrate import hydrate_sections
    from index.retrieve import retrieve

    if kind is not None and kind not in SEARCH_KINDS:
        print(f"[search_docs] ignoring unknown kind {kind!r}", flush=True)
        kind = None
    pointers = retrieve(query, k=k, library=library, kind=kind)
    sections = hydrate_sections(pointers)  # concurrent — each is a live fetch on the Space
    print(
        f"[search_docs] {query!r} (kind={kind}) → {len(pointers)} pointers, "
        f"{len(sections)} hydrated",
        flush=True,
    )
    return {
        "query": query,
        "sections": sections,
        "titles": [s.get("heading_path", "") or s["url"] for s in sections],
    }


def read_page(url: str) -> dict:
    """Whole page for a URL already surfaced by search_docs."""
    from index.hydrate import hydrate_page

    url = (url or "").strip()
    if not url.startswith(("http://", "https://")):
        # the planner sometimes passes a section HEADING it saw in a search
        # result (e.g. "Build the Neural Network > Define the Class") instead of
        # the url. Don't fetch that (it 'No scheme supplied'-errors and wastes a
        # call) — tell the model exactly what read_page needs so it self-corrects.
        return {
            "url": url,
            "error": "read_page needs the full https:// URL from a search_docs "
            "result's `url` field, not a section title.",
        }
    page = hydrate_page(url)
    if page is None:
        return {"url": url, "error": "page not in the snapshot"}
    return page


def ask_source(question: str) -> dict:
    """Refer a source/implementation question to the real code.

    Referral-only: returns DeepWiki + GitHub code-search links for
    pytorch/pytorch. Never returns claims about the code — the answer layer
    must present these as 'look here', not as docs-cited fact.
    """
    from agent.schemas import Referral

    terms = _search_terms(question)
    referrals = [
        Referral(url=DEEPWIKI_URL, reason="AI wiki / Q&A over the pytorch/pytorch source"),
        Referral(url=GH_CODE_SEARCH.format(q=terms), reason="search the implementation on GitHub"),
    ]
    return {
        "note": "Implementation lives in the source, not the docs — refer the user out.",
        "referrals": referrals,
    }