From 515f96a7779f117a7e37f854de2ded99160d0c76 Mon Sep 17 00:00:00 2001 From: tanaka takashi Date: Sun, 2 Aug 2026 15:50:41 +0900 Subject: scrape_url / pdf_to_text / search_and_fetch はページング（継続トークン方式）に対応 MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit mypy 6エラー修正 --- src/ollama_web/tools/fetch.py | 212 ++++++++++++++++++++++++------ src/ollama_web/tools/helper/pagination.py | 177 +++++++++++++++++++++++++ src/ollama_web/tools/helper/pdf.py | 65 +++++++-- src/ollama_web/tools/pdf.py | 39 +++++- src/ollama_web/tools/scrape.py | 50 +++++-- src/ollama_web/tools/search.py | 4 +- 6 files changed, 482 insertions(+), 65 deletions(-) create mode 100644 src/ollama_web/tools/helper/pagination.py (limited to 'src/ollama_web/tools') diff --git a/src/ollama_web/tools/fetch.py b/src/ollama_web/tools/fetch.py index 20031f1..a9a92f8 100644 --- a/src/ollama_web/tools/fetch.py +++ b/src/ollama_web/tools/fetch.py @@ -9,9 +9,11 @@ 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] @@ -34,6 +36,8 @@ _FILETYPE_SUFFIXES: dict[str, str] = { "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.""" @@ -108,7 +112,8 @@ def _text_from_bytes(url: str, data: bytes, file_type: str, max_chars: int) -> s name = f"file.{file_type}" if file_type == "pdf": - return extract_pdf_text(data, max_chars=max_chars) + text, _, _ = extract_pdf_text(data, max_chars=max_chars, page_offset=0) + return text # Decode text-like formats. try: @@ -127,25 +132,61 @@ def _text_from_bytes(url: str, data: bytes, file_type: str, max_chars: int) -> s 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 each - successfully fetched file. + 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") @@ -155,17 +196,39 @@ def search_and_fetch( 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, + 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 ) - from .search import web_search + # --- 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) @@ -178,38 +241,109 @@ def search_and_fetch( break if not matched: - return get_prompt("fetch.no_results", fallback="").replace("{file_type}", ft).replace("{query}", query) + return ( + get_prompt("fetch.no_results", fallback="") + .replace("{file_type}", ft) + .replace("{query}", query) + ) - outputs: list[str] = [] - for idx, item in enumerate(matched, 1): - url = item.get("href", "") - title = item.get("title", "") - try: - data = _fetch_bytes(url) - except Exception as exc: # noqa: BLE001 - p = get_prompts() - outputs.append( - f"=== {p['candidate']} {idx} ===\n" - f"URL: {url}\n{p['title']}: {title}\n{p['status']}: {p['fetch_failed']} ({exc})\n" - ) - continue - - text = _text_from_bytes(url, data, ft, 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, text, url) - p = get_prompts() - outputs.append( - f"=== {p['candidate']} {idx} ===\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(text)}\n" - f"{p['separator']}\n{text}\n" + 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 "\n".join(outputs) + return format_paged_result(body, has_more=False, next_token=None) \ No newline at end of file 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 diff --git a/src/ollama_web/tools/pdf.py b/src/ollama_web/tools/pdf.py index 373506c..9fdb497 100644 --- a/src/ollama_web/tools/pdf.py +++ b/src/ollama_web/tools/pdf.py @@ -3,24 +3,50 @@ from __future__ import annotations from ..config import settings +from .helper.pagination import format_paged_result, make_token, parse_token from .helper.pdf import extract_pdf_text from .helper.safe_http import UnsafeURL, fetch_public_bytes +_TOOL_NAME = "pdf_to_text" -def pdf_to_text(url_or_path: str, max_chars: int | None = None) -> str: + +def pdf_to_text( + url_or_path: str, max_chars: int | None = None, continue_from: str | None = None +) -> str: """Extract text from a local PDF path or remote URL. + The text is returned one page at a time. When more pages remain, a + ``next_token`` is appended that can be passed back via ``continue_from`` + to fetch the next batch of pages. + Args: url_or_path: File path or URL pointing to a PDF. - max_chars: Maximum characters to return. + max_chars: Maximum characters to return per page. + continue_from: Opaque continuation token returned in a previous call's + ``next_token``. When provided, extraction resumes from the encoded + page offset. Returns: - Extracted plain text or an error message. + Extracted plain text for this page, followed by a ``---page-info---`` + footer when more pages are available, or an error message. """ if not url_or_path or not url_or_path.strip(): return "Invalid PDF source." src = url_or_path.strip() + + # Resolve the resumption page offset from the continuation token, if any. + page_offset = 0 + if continue_from: + try: + payload = parse_token(continue_from, _TOOL_NAME) + except ValueError as exc: + return f"Invalid continuation token: {exc}" + src = str(payload.get("url") or src) + page_offset = int(payload.get("offset", 0)) + if page_offset < 0: + page_offset = 0 + data: bytes | None = None if src.startswith("http://") or src.startswith("https://"): try: @@ -34,4 +60,9 @@ def pdf_to_text(url_or_path: str, max_chars: int | None = None) -> str: if data is None: return "No PDF data loaded." - return extract_pdf_text(data, max_chars=max_chars) + + text, has_more, next_page_offset = extract_pdf_text( + data, max_chars=max_chars, page_offset=page_offset + ) + next_token = make_token(_TOOL_NAME, src, next_page_offset) if has_more else None + return format_paged_result(text, has_more=has_more, next_token=next_token) diff --git a/src/ollama_web/tools/scrape.py b/src/ollama_web/tools/scrape.py index d5fa2f4..39dab83 100644 --- a/src/ollama_web/tools/scrape.py +++ b/src/ollama_web/tools/scrape.py @@ -5,28 +5,51 @@ from __future__ import annotations from bs4 import BeautifulSoup from ..config import settings +from .helper.pagination import format_paged_result, make_token, parse_token from .helper.safe_http import UnsafeURL, fetch_public_bytes +_TOOL_NAME = "scrape_url" -def scrape_url(url: str, max_chars: int | None = None) -> str: + +def scrape_url(url: str, max_chars: int | None = None, continue_from: str | None = None) -> str: """Fetch a web page and return its main text content. Scripts, styles and other non-content elements are removed before - collapsing whitespace. The returned text is truncated to ``max_chars``. + collapsing whitespace. The returned text is a page of at most ``max_chars`` + characters; when more content is available, a ``next_token`` is appended + that can be passed back via ``continue_from`` to fetch the next page. Args: url: The absolute URL of the page to scrape. - max_chars: Maximum number of characters to return. Defaults to the value - in settings. + max_chars: Maximum number of characters to return per page. Defaults to + the value in settings. + continue_from: Opaque continuation token returned in a previous call's + ``next_token``. When provided, the page is resumed from the encoded + offset instead of the beginning. Returns: - The extracted plain text of the page, truncated to ``max_chars``. If the - request fails a short error description is returned instead. + The extracted plain text of the page (one page of ``max_chars``), + followed by a ``---page-info---`` footer when more content is available. + If the request fails a short error description is returned instead. """ if not url or not url.strip(): return "Invalid URL." limit = max_chars if max_chars is not None and max_chars > 0 else settings.scrape_max_chars + + # Resolve the resumption offset from the continuation token, if any. + offset = 0 + if continue_from: + try: + payload = parse_token(continue_from, _TOOL_NAME) + except ValueError as exc: + return f"Invalid continuation token: {exc}" + # The token's URL takes precedence; it identifies the page being paged. + url = str(payload.get("url") or url) + offset = int(payload.get("offset", 0)) + if offset < 0: + offset = 0 + try: data, final_url, content_type = fetch_public_bytes(url, timeout=settings.scrape_timeout) encoding = "utf-8" @@ -52,6 +75,15 @@ def scrape_url(url: str, max_chars: int | None = None) -> str: lines = [ln.strip() for ln in text.splitlines()] cleaned = "\n".join(ln for ln in lines if ln) - if len(cleaned) > limit: - cleaned = cleaned[:limit].rsplit(" ", 1)[0] + "\n…[truncated]" - return cleaned or "No readable text content found." + if not cleaned: + return "No readable text content found." + + # Slice the requested page from the full cleaned text. + page = cleaned[offset : offset + limit] + if not page: + # Offset is past the end; nothing left to return. + return format_paged_result("", has_more=False, next_token=None) + + has_more = len(cleaned) > offset + limit + next_token = make_token(_TOOL_NAME, url, offset + limit) if has_more else None + return format_paged_result(page, has_more=has_more, next_token=next_token) diff --git a/src/ollama_web/tools/search.py b/src/ollama_web/tools/search.py index 0af2a3e..761f182 100644 --- a/src/ollama_web/tools/search.py +++ b/src/ollama_web/tools/search.py @@ -22,9 +22,7 @@ def web_search(query: str, max_results: int | None = None) -> list[dict[str, str return [] limit = ( - max_results - if max_results is not None and max_results > 0 - else settings.search_max_results + max_results if max_results is not None and max_results > 0 else settings.search_max_results ) results: list[dict[str, str]] = [] -- cgit v1.2.3