scrape_url / pdf_to_text / search_and_fetch はページング（継続トークン方式）に対応

mypy 6エラー修正

2 files changed, 232 insertions, 10 deletions

diff --git a/src/ollama_web/tools/helper/pagination.py b/src/ollama_web/tools/helper/pagination.py
new file mode 100644
index 0000000..7f18ac8
--- /dev/null
+++ b/src/ollama_web/tools/helper/pagination.py
@@ -0,0 +1,177 @@
+"""Pagination helpers for paged tool results.
+
+This module provides an opaque continuation-token mechanism so that LLM-driven
+tools can return large content in pages without requiring the model to remember
+numeric offsets. The token encodes the resumption state (URL + offset) as
+base64url(JSON) and is meant to be passed back verbatim by the model via the
+``continue_from`` parameter.
+
+Design rationale:
+- The token is opaque: the LLM should copy the ``next_token`` value from a
+  previous result and pass it to ``continue_from`` without interpreting it.
+- The state is encoded in the token itself, so the server stays stateless and
+  does not need to cache previous results.
+- A ``tool`` field guards against a token issued by one tool being replayed
+  against a different tool.
+"""
+
+from __future__ import annotations
+
+import base64
+import json
+from typing import Any
+
+# Token format version. Bump when the payload structure changes in a
+# backwards-incompatible way; ``parse_token`` rejects unknown versions.
+_TOKEN_VERSION = 1
+
+# Marker that separates the page body from the page-info footer in tool
+# results. Kept in sync with ``llm._trim_tool_result``.
+PAGE_INFO_MARKER = "\n---page-info---"
+
+
+class InvalidTokenError(ValueError):
+    """Raised when a continuation token is malformed or invalid."""
+
+
+def _b64url_encode(data: bytes) -> str:
+    return base64.urlsafe_b64encode(data).decode("ascii").rstrip("=")
+
+
+def _b64url_decode(token: str) -> bytes:
+    # Re-add any stripped padding before decoding.
+    padding = "=" * (-len(token) % 4)
+    return base64.urlsafe_b64decode(token + padding)
+
+
+def make_token(tool: str, url: str, offset: int, **extra: Any) -> str:
+    """Create an opaque continuation token.
+
+    Args:
+      tool: Name of the issuing tool (e.g. ``"scrape_url"``).
+      url: The URL the content was fetched from.
+      offset: The resumption offset (character or page index, tool-specific).
+      **extra: Additional tool-specific fields to embed in the token payload
+        (e.g. ``matched`` candidate list for ``search_and_fetch``).
+
+    Returns:
+      A base64url-encoded JSON string.
+    """
+    payload: dict[str, Any] = {
+        "v": _TOKEN_VERSION,
+        "tool": tool,
+        "url": url,
+        "offset": int(offset),
+    }
+    payload.update(extra)
+    raw = json.dumps(payload, ensure_ascii=False, separators=(",", ":")).encode("utf-8")
+    return _b64url_encode(raw)
+
+
+def parse_token(token: str | None, expected_tool: str) -> dict[str, Any]:
+    """Validate and decode a continuation token.
+
+    Args:
+      token: The opaque token string (typically from ``continue_from``).
+        ``None`` or an empty string raises :class:`InvalidTokenError`.
+      expected_tool: The tool name that should appear in the token's ``tool``
+        field. Mismatches raise :class:`InvalidTokenError`.
+
+    Returns:
+      The decoded payload dict (includes ``url``, ``offset`` and any extra
+      fields the issuing tool embedded).
+
+    Raises:
+      InvalidTokenError: If the token is empty, not valid base64, not valid
+        JSON, missing required fields, has an unsupported version, or names a
+        different tool than ``expected_tool``.
+    """
+    if not token or not token.strip():
+        raise InvalidTokenError("Empty continuation token.")
+
+    try:
+        raw = _b64url_decode(token.strip())
+    except (ValueError, base64.binascii.Error) as exc:  # type: ignore[attr-defined]
+        raise InvalidTokenError(f"Token is not valid base64url: {exc}") from exc
+
+    try:
+        payload = json.loads(raw.decode("utf-8"))
+    except (UnicodeDecodeError, json.JSONDecodeError) as exc:
+        raise InvalidTokenError(f"Token payload is not valid JSON: {exc}") from exc
+
+    if not isinstance(payload, dict):
+        raise InvalidTokenError("Token payload is not a JSON object.")
+
+    version = payload.get("v")
+    if version != _TOKEN_VERSION:
+        raise InvalidTokenError(f"Unsupported token version: {version!r}")
+
+    tool = payload.get("tool")
+    if tool != expected_tool:
+        raise InvalidTokenError(f"Token was issued for tool {tool!r}, expected {expected_tool!r}.")
+
+    if "url" not in payload or "offset" not in payload:
+        raise InvalidTokenError("Token payload missing 'url' or 'offset'.")
+
+    try:
+        payload["offset"] = int(payload["offset"])
+    except (TypeError, ValueError) as exc:
+        raise InvalidTokenError(f"Token 'offset' is not an integer: {exc}") from exc
+
+    return payload
+
+
+def format_paged_result(text: str, *, has_more: bool, next_token: str | None) -> str:
+    """Wrap page body with a page-info footer.
+
+    The footer is appended after :data:`PAGE_INFO_MARKER` so that
+    ``llm._trim_tool_result`` can protect it from truncation. When
+    ``has_more`` is False the ``next_token`` line is omitted.
+
+    Args:
+      text: The page body text.
+      has_more: Whether more content is available.
+      next_token: The continuation token to include when ``has_more`` is True.
+
+    Returns:
+      The formatted result string with the page-info footer.
+    """
+    if not has_more:
+        # No footer needed when there is nothing more to fetch. Keeping the
+        # body unchanged preserves backward compatibility for short results.
+        return text
+
+    if not next_token:
+        # Defensive: has_more without a token is a programming error; fall back
+        # to no footer to avoid confusing the model.
+        return text
+
+    footer = f"{PAGE_INFO_MARKER}\nhas_more: true\nnext_token: {next_token}"
+    return f"{text}{footer}"
+
+
+def parse_paged_result(result: str) -> tuple[str, bool, str | None]:
+    """Split a paged result into (body, has_more, next_token).
+
+    This is the inverse of :func:`format_paged_result` and is mainly useful in
+    tests and for chaining tool results.
+
+    Args:
+      result: The formatted result string from :func:`format_paged_result`.
+
+    Returns:
+      A tuple of (body_text, has_more, next_token_or_None).
+    """
+    if PAGE_INFO_MARKER not in result:
+        return result, False, None
+
+    body, _, footer = result.partition(PAGE_INFO_MARKER)
+    has_more = "has_more: true" in footer
+    next_token: str | None = None
+    if has_more:
+        for line in footer.splitlines():
+            line = line.strip()
+            if line.startswith("next_token:"):
+                next_token = line.split(":", 1)[1].strip()
+                break
+    return body, has_more, next_token
diff --git a/src/ollama_web/tools/helper/pdf.py b/src/ollama_web/tools/helper/pdf.py
index 51d3b3a..7b9fdac 100644
--- a/src/ollama_web/tools/helper/pdf.py
+++ b/src/ollama_web/tools/helper/pdf.py
@@ -9,31 +9,76 @@ from pypdf import PdfReader
 from ...config import settings
 
 
-def extract_pdf_text(data: bytes, max_chars: int | None = None) -> str:
-    """Extract text from PDF bytes."""
+def extract_pdf_text(
+    data: bytes,
+    max_chars: int | None = None,
+    page_offset: int = 0,
+) -> tuple[str, bool, int]:
+    """Extract text from PDF bytes, starting at ``page_offset``.
+
+    Args:
+      data: Raw PDF bytes.
+      max_chars: Maximum characters to return. Defaults to
+        ``settings.fetch_max_chars``.
+      page_offset: Zero-based index of the first page to extract. Pages before
+        this index are skipped. Used by the pagination layer to resume.
+
+    Returns:
+      A tuple ``(text, has_more, next_page_offset)`` where:
+        - ``text`` is the extracted text for this page (truncated to
+          ``max_chars``).
+        - ``has_more`` is True when there are remaining pages after the
+          extracted slice.
+        - ``next_page_offset`` is the page index to resume from on the next
+          call (only meaningful when ``has_more`` is True).
+    """
     if not data:
-        return "Empty PDF data."
+        return "Empty PDF data.", False, 0
 
     limit = max_chars if max_chars is not None and max_chars > 0 else settings.fetch_max_chars
 
     if len(data) > settings.max_fetch_mb * 1024 * 1024:
-        return "PDF exceeds the configured size limit."
+        return "PDF exceeds the configured size limit.", False, 0
 
     try:
         reader = PdfReader(stream=io.BytesIO(data))
     except Exception as exc:  # noqa: BLE001
-        return f"Failed to parse PDF: {exc}"
+        return f"Failed to parse PDF: {exc}", False, 0
+
+    total_pages = len(reader.pages)
+    if page_offset < 0:
+        page_offset = 0
+    if page_offset >= total_pages:
+        return "No text could be extracted from the PDF.", False, 0
+
+    # Cap the number of pages processed per call to max_pdf_pages.
+    end = min(page_offset + settings.max_pdf_pages, total_pages)
 
     parts: list[str] = []
-    for page in reader.pages[: settings.max_pdf_pages]:
+    collected = 0
+    last_extracted_page = page_offset
+    for page_index in range(page_offset, end):
         try:
-            text = page.extract_text()
-            if text:
-                parts.append(text)
+            text = reader.pages[page_index].extract_text()
         except Exception:  # noqa: BLE001
             continue
+        if text:
+            parts.append(text)
+            collected += len(text)
+            last_extracted_page = page_index
+        # Stop as soon as we reach the per-call character budget so the page
+        # boundary for the next token is accurate.
+        if collected >= limit:
+            break
 
     full = "\n\n".join(parts)
     if len(full) > limit:
         full = full[:limit].rsplit(" ", 1)[0] + "\n…[truncated]"
-    return full or "No text could be extracted from the PDF."
+
+    # has_more when there are pages beyond the last one we extracted.
+    has_more = (last_extracted_page + 1) < total_pages
+    next_page_offset = last_extracted_page + 1 if has_more else 0
+
+    if not full:
+        return "No text could be extracted from the PDF.", has_more, next_page_offset
+    return full, has_more, next_page_offset