path: root/src/ollama_web/tools/fetch.py

"""Search-and-fetch tool that finds files on the web by type and returns text."""

from __future__ import annotations

import mimetypes
import re
from pathlib import Path
from typing import Any

from ..config import settings
from ..prompts import get_prompt, get_prompts
from .helper.pagination import format_paged_result, make_token, parse_token
from .helper.pdf import extract_pdf_text
from .helper.safe_http import fetch_public_bytes
from .scrape import scrape_url
from .search import web_search

_FetchedFile = dict[str, Any]

# Per-session temporary storage for files fetched by search_and_fetch.
# The chat route drains this after the tool call finishes.
_fetched_files: dict[str, list[_FetchedFile]] = {}


_FILETYPE_SUFFIXES: dict[str, str] = {
    "pdf": "pdf",
    "txt": "txt",
    "csv": "csv",
    "md": "md",
    "json": "json",
    "html": "html",
    "htm": "htm",
    "xml": "xml",
    "py": "py",
    "js": "js",
    "ts": "ts",
}

_TOOL_NAME = "search_and_fetch"


def set_session_context(session_id: str | None) -> None:
    """Set the active session id so fetched files can be persisted later."""
    # No-op marker kept for future hooks; storage is keyed by session_id directly.
    pass


def store_fetched_file(
    session_id: str,
    name: str,
    data: bytes,
    mime: str,
    text: str,
    url: str,
) -> None:
    """Store a fetched file so the chat route can persist it to the session."""
    if session_id:
        _fetched_files.setdefault(session_id, []).append(
            {
                "name": name,
                "data": data,
                "mime": mime,
                "text": text,
                "url": url,
            }
        )


def pop_fetched_files(session_id: str) -> list[_FetchedFile]:
    """Return and clear any files fetched for the given session."""
    return _fetched_files.pop(session_id, [])


def _normalize_file_type(file_type: str) -> str:
    ft = file_type.strip().lower().lstrip(".")
    return _FILETYPE_SUFFIXES.get(ft, ft)


def _filetype_query(query: str, file_type: str) -> str:
    return f"{query} filetype:{file_type}"


def _is_matching_url(url: str, file_type: str) -> bool:
    suffix = f".{file_type}"
    import httpx

    parsed = httpx.URL(url)
    path = parsed.path or ""
    if path.lower().endswith(suffix):
        return True
    # Some CDNs embed extension after query params.
    if re.search(rf"\.{re.escape(file_type)}(?:[?#]|$)", url, re.IGNORECASE):
        return True
    return False


def _guess_mime(name: str, file_type: str) -> str:
    mime, _ = mimetypes.guess_type(name)
    return mime or f"application/{file_type}"


def _fetch_bytes(url: str) -> bytes:
    data, _, _ = fetch_public_bytes(url, timeout=settings.scrape_timeout)
    return data


def _text_from_bytes(url: str, data: bytes, file_type: str, max_chars: int) -> str:
    import httpx

    name = Path(httpx.URL(url).path or "file").name or f"file.{file_type}"
    if not name or "." not in name:
        name = f"file.{file_type}"

    if file_type == "pdf":
        text, _, _ = extract_pdf_text(data, max_chars=max_chars, page_offset=0)
        return text

    # Decode text-like formats.
    try:
        text = data.decode("utf-8")
    except UnicodeDecodeError:
        try:
            text = data.decode("cp932")
        except UnicodeDecodeError:
            text = data.decode("utf-8", errors="replace")

    if file_type in {"html", "htm"}:
        text = scrape_url(url, max_chars=max_chars)

    if len(text) > max_chars:
        text = text[:max_chars].rsplit(" ", 1)[0] + "\n…[truncated]"
    return text


def _full_text_for_file(url: str, data: bytes, file_type: str) -> str:
    """Return the *full* extracted text for a file (no truncation).

    Used by the pagination path so we can slice arbitrary offsets. For PDFs we
    extract all pages; for text-like formats we decode the full payload.
    """
    if file_type == "pdf":
        # Extract with a very large char budget to get the full text.
        text, _, _ = extract_pdf_text(data, max_chars=10_000_000, page_offset=0)
        return text

    try:
        text = data.decode("utf-8")
    except UnicodeDecodeError:
        try:
            text = data.decode("cp932")
        except UnicodeDecodeError:
            text = data.decode("utf-8", errors="replace")

    if file_type in {"html", "htm"}:
        # Reuse scrape_url's cleaning but request a large budget to get the
        # full cleaned text. scrape_url itself now supports pagination, but
        # here we want the full body for manual slicing.
        text = scrape_url(url, max_chars=10_000_000)
    return text


def search_and_fetch(
    query: str,
    file_type: str,
    max_results: int | None = None,
    max_chars: int | None = None,
    session_id: str | None = None,
    continue_from: str | None = None,
) -> str:
    """Search the web for files of a given type and return their text content.

    Results are returned one file at a time. When a file is truncated or more
    candidate files remain, a ``next_token`` is appended that can be passed
    back via ``continue_from`` to resume.

    Args:
      query: Search keywords.
      file_type: Desired file extension (e.g. "pdf", "csv", "txt", "md", "json").
      max_results: Number of candidate URLs to fetch.
      max_chars: Max characters per file to return.
      session_id: If provided, fetched files are queued for session persistence.
      continue_from: Opaque continuation token returned in a previous call's
        ``next_token``. When provided, resumes from the encoded file index and
        in-file offset without re-running the search.

    Returns:
      A formatted string containing metadata and extracted text for the
      current file page, followed by a ``---page-info---`` footer when more
      content is available.
    """
    if not query or not query.strip():
        return get_prompt("fetch.empty_query")

    ft = _normalize_file_type(file_type)
    if not ft:
        return get_prompt("fetch.invalid_type")

    limit = min(
        max_results if max_results is not None and max_results > 0 else settings.fetch_max_results,
        settings.fetch_max_results,
    )
    per_file_chars = (
        max_chars if max_chars is not None and max_chars > 0 else settings.fetch_max_chars
    )

    # --- Pagination resume path ---
    if continue_from:
        try:
            payload = parse_token(continue_from, _TOOL_NAME)
        except ValueError as exc:
            return f"Invalid continuation token: {exc}"
        resumed_matched = payload.get("matched", [])
        if not isinstance(resumed_matched, list) or not resumed_matched:
            return (
                get_prompt("fetch.no_results", fallback="")
                .replace("{file_type}", ft)
                .replace("{query}", query)
            )
        file_idx = int(payload.get("file_idx", 0))
        file_offset = int(payload.get("offset", 0))
        return _render_file_page(
            matched=resumed_matched,
            file_idx=file_idx,
            file_offset=file_offset,
            ft=ft,
            per_file_chars=per_file_chars,
            session_id=session_id,
            query=query,
        )

    # --- Initial search path ---
    search_query = _filetype_query(query, ft)
    results = web_search(search_query, max_results=limit * 2)

    matched: list[dict[str, str]] = []
    for item in results:
        href = item.get("href", "")
        if _is_matching_url(href, ft):
            matched.append(item)
            if len(matched) >= limit:
                break

    if not matched:
        return (
            get_prompt("fetch.no_results", fallback="")
            .replace("{file_type}", ft)
            .replace("{query}", query)
        )

    return _render_file_page(
        matched=matched,
        file_idx=0,
        file_offset=0,
        ft=ft,
        per_file_chars=per_file_chars,
        session_id=session_id,
        query=query,
    )


def _render_file_page(
    *,
    matched: list[dict[str, str]],
    file_idx: int,
    file_offset: int,
    ft: str,
    per_file_chars: int,
    session_id: str | None,
    query: str,
) -> str:
    """Render a single file page and emit a continuation token if needed.

    The token encodes the candidate list so subsequent pages do not need to
    re-run the web search.
    """
    if file_idx >= len(matched):
        # All candidates exhausted.
        return format_paged_result("", has_more=False, next_token=None)

    item = matched[file_idx]
    url = item.get("href", "")
    title = item.get("title", "")
    p = get_prompts().get("fetch", {})

    try:
        data = _fetch_bytes(url)
    except Exception as exc:  # noqa: BLE001
        # On fetch failure, advance to the next candidate rather than aborting
        # the whole pagination chain.
        body = (
            f"=== {p['candidate']} {file_idx + 1} ===\n"
            f"URL: {url}\n{p['title']}: {title}\n{p['status']}: {p['fetch_failed']} ({exc})\n"
        )
        next_token = make_token(
            _TOOL_NAME,
            url,
            0,
            matched=matched,
            file_idx=file_idx + 1,
            file_offset=0,
        )
        return format_paged_result(body, has_more=True, next_token=next_token)

    full_text = _full_text_for_file(url, data, ft)
    page_text = full_text[file_offset : file_offset + per_file_chars]
    has_more_in_file = len(full_text) > file_offset + per_file_chars

    import httpx

    mime = _guess_mime(Path(httpx.URL(url).path or f"file.{ft}").name, ft)
    name = Path(httpx.URL(url).path or f"file.{ft}").name
    if session_id:
        store_fetched_file(session_id, name, data, mime, page_text, url)

    body = (
        f"=== {p['candidate']} {file_idx + 1} ===\n"
        f"URL: {url}\n"
        f"{p['title']}: {title}\n"
        f"{p['type']}: {mime}\n"
        f"{p['size']}: {len(data)} bytes\n"
        f"{p['chars']}: {len(page_text)}\n"
        f"{p['separator']}\n{page_text}\n"
    )

    if has_more_in_file:
        # Continue within the same file.
        next_token = make_token(
            _TOOL_NAME,
            url,
            file_offset + per_file_chars,
            matched=matched,
            file_idx=file_idx,
            file_offset=file_offset + per_file_chars,
        )
        return format_paged_result(body, has_more=True, next_token=next_token)

    # Move to the next candidate file.
    if file_idx + 1 < len(matched):
        next_token = make_token(
            _TOOL_NAME,
            url,
            0,
            matched=matched,
            file_idx=file_idx + 1,
            file_offset=0,
        )
        return format_paged_result(body, has_more=True, next_token=next_token)

    return format_paged_result(body, has_more=False, next_token=None)