From 33f6b4dbb69e9366e312a8b44dbed9d22e7cba5b Mon Sep 17 00:00:00 2001 From: tanaka takashi Date: Sun, 23 Aug 2026 17:28:10 +0900 Subject: fix(tools): PDFページングを損失なしカーソル方式へ再設計 MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit PDF本文をページ単位で抽出した後に文字数上限で切断しつつ、継続位置だけ次ページへ進めていたため、ページ末尾が復元不能になる問題を修正する。ツール側の本文上限とLLM側の結果上限の不整合、25ページの処理上限、添付経路での継続情報欠落も共通契約で解消する。 - 共通PageResultと、メタデータ込みの結果上限内で本文量を決めるフォーマッターを追加 - v2トークンにPDFのページ番号とページ内文字位置を保持し、v1トークンの読み取り互換性を維持 - pdf_to_text、scrape_url、search_and_fetchを返却済み本文だけカーソルが進む方式へ統一 - max_pdf_pagesを文書全体ではなく1呼び出しの処理上限として扱い、長い単一ページと26ページ目以降の継続取得に対応 - attachment:をセッション限定で解決し、既存添付も保存済みプレビューではなく元PDFから再解析 - Web取得ファイルをセッションとURL単位で重複排除し、HTMLおよびテキスト経路の全文ページングも維持 - 正常なページング結果の後段切断を禁止し、結果上限・コンテキスト予算・ラウンド上限では未取得状態を明示 - 日英プロンプト、UIメッセージ、ツールスキーマ、変更履歴を実装へ同期 - 実PDF、30ページ超、長い単一ページ、旧トークン、セッション分離、コンテキスト境界の回帰テストを追加 検証: - pytest: 130 passed - ruff check src tests: passed - mypy src: passed - c2144.pdf: 19ページ、抽出本文46015文字（ページ区切り込み46051文字）を8回で重複・欠落なく完全再構成 --- src/ollama_web/tools/__init__.py | 3 +- src/ollama_web/tools/fetch.py | 290 ++++++++++++++++++++---------- src/ollama_web/tools/helper/__init__.py | 3 +- src/ollama_web/tools/helper/pagination.py | 192 +++++++++----------- src/ollama_web/tools/helper/pdf.py | 155 ++++++++++------ src/ollama_web/tools/pdf.py | 153 ++++++++++++---- src/ollama_web/tools/scrape.py | 70 +++++--- 7 files changed, 540 insertions(+), 326 deletions(-) (limited to 'src/ollama_web/tools') diff --git a/src/ollama_web/tools/__init__.py b/src/ollama_web/tools/__init__.py index f2a6b0b..05d4a04 100644 --- a/src/ollama_web/tools/__init__.py +++ b/src/ollama_web/tools/__init__.py @@ -2,7 +2,7 @@ from .fetch import search_and_fetch from .helper.image import resize_image -from .helper.pdf import extract_pdf_text +from .helper.pdf import extract_pdf_page, extract_pdf_text from .pdf import pdf_to_text from .registry import ToolRegistry, all_tool_callables, default_registry from .scrape import scrape_url @@ -12,6 +12,7 @@ __all__ = [ "search_and_fetch", "resize_image", "extract_pdf_text", + "extract_pdf_page", "pdf_to_text", "ToolRegistry", "all_tool_callables", diff --git a/src/ollama_web/tools/fetch.py b/src/ollama_web/tools/fetch.py index 65ec511..6ab17bb 100644 --- a/src/ollama_web/tools/fetch.py +++ b/src/ollama_web/tools/fetch.py @@ -2,6 +2,7 @@ from __future__ import annotations +import logging import mimetypes import re from pathlib import Path @@ -9,17 +10,24 @@ 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.pagination import ( + PageResult, + PageSizeError, + format_page_to_limit, + make_token, + parse_token, + text_page, +) +from .helper.pdf import extract_pdf_page, legacy_pdf_offset_to_cursor from .helper.safe_http import fetch_public_bytes -from .scrape import scrape_url +from .scrape import extract_readable_text 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]] = {} +_fetched_files: dict[str, dict[str, _FetchedFile]] = {} _FILETYPE_SUFFIXES: dict[str, str] = { @@ -37,6 +45,7 @@ _FILETYPE_SUFFIXES: dict[str, str] = { } _TOOL_NAME = "search_and_fetch" +logger = logging.getLogger("ollama_web.tools.fetch") def _slim_matched(matched: list[dict[str, str]]) -> list[dict[str, str]]: @@ -53,12 +62,6 @@ def _slim_matched(matched: list[dict[str, str]]) -> list[dict[str, str]]: return [{"title": m.get("title", ""), "href": m.get("href", "")} for m in matched] -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, @@ -69,7 +72,8 @@ def store_fetched_file( ) -> None: """Store a fetched file so the chat route can persist it to the session.""" if session_id: - _fetched_files.setdefault(session_id, []).append( + _fetched_files.setdefault(session_id, {}).setdefault( + url, { "name": name, "data": data, @@ -82,7 +86,7 @@ def store_fetched_file( 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, []) + return list(_fetched_files.pop(session_id, {}).values()) def _normalize_file_type(file_type: str) -> str: @@ -118,44 +122,10 @@ def _fetch_bytes(url: str) -> bytes: 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. - """ + """Decode a complete text-like file for character-offset pagination.""" 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 + raise ValueError("PDF files must use page-and-character cursor pagination.") try: text = data.decode("utf-8") @@ -166,10 +136,7 @@ def _full_text_for_file(url: str, data: bytes, file_type: str) -> str: 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) + text = extract_readable_text(data) return text @@ -223,19 +190,43 @@ def search_and_fetch( 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: + raw_matched = payload.get("matched", []) + if not isinstance(raw_matched, list) or not raw_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)) + resumed_matched = [ + {"title": str(item.get("title", "")), "href": str(item.get("href", ""))} + for item in raw_matched + if isinstance(item, dict) + ] + if payload.get("v") == 2 and payload.get("file_type") != ft: + return "Invalid continuation token: file type does not match." + try: + file_idx = int(payload.get("file_idx", 0)) + file_offset = int(payload.get("file_offset", payload.get("offset", 0))) + pdf_page_index = int(payload.get("pdf_page_index", 0)) + pdf_page_char_offset = int(payload.get("pdf_page_char_offset", 0)) + except (TypeError, ValueError) as exc: + return f"Invalid continuation token: cursor is not an integer: {exc}" + legacy_pdf_offset = ( + file_offset + if payload.get("v") == 1 and "pdf_page_index" not in payload + else None + ) + if min(file_idx, file_offset, pdf_page_index, pdf_page_char_offset) < 0: + return "Invalid continuation token: cursor values must not be negative." + if not resumed_matched or file_idx >= len(resumed_matched): + return "Invalid continuation token: file index is out of range." return _render_file_page( matched=resumed_matched, file_idx=file_idx, file_offset=file_offset, + pdf_page_index=pdf_page_index, + pdf_page_char_offset=pdf_page_char_offset, + legacy_pdf_offset=legacy_pdf_offset, ft=ft, per_file_chars=per_file_chars, session_id=session_id, @@ -265,6 +256,9 @@ def search_and_fetch( matched=matched, file_idx=0, file_offset=0, + pdf_page_index=0, + pdf_page_char_offset=0, + legacy_pdf_offset=None, ft=ft, per_file_chars=per_file_chars, session_id=session_id, @@ -277,6 +271,9 @@ def _render_file_page( matched: list[dict[str, str]], file_idx: int, file_offset: int, + pdf_page_index: int, + pdf_page_char_offset: int, + legacy_pdf_offset: int | None, ft: str, per_file_chars: int, session_id: str | None, @@ -288,8 +285,7 @@ def _render_file_page( re-run the web search. """ if file_idx >= len(matched): - # All candidates exhausted. - return format_paged_result("", has_more=False, next_token=None) + return "" item = matched[file_idx] url = item.get("href", "") @@ -305,59 +301,157 @@ def _render_file_page( 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=_slim_matched(matched), - file_idx=file_idx + 1, - file_offset=0, + has_more = file_idx + 1 < len(matched) + cursor = { + "file_idx": file_idx + 1, + "file_offset": 0, + "pdf_page_index": 0, + "pdf_page_char_offset": 0, + } + page = PageResult(body, cursor, cursor, has_more) + return format_page_to_limit( + lambda _budget: page, + lambda next_cursor: _make_fetch_token(matched, next_cursor, ft), + requested_chars=per_file_chars, + result_limit=settings.max_tool_result_chars, ) - 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 = ( + prefix = ( 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" + f"{p['separator']}\n" ) - if has_more_in_file: - # Continue within the same file. - next_token = make_token( - _TOOL_NAME, - url, - file_offset + per_file_chars, - matched=_slim_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) + if ft == "pdf": + if legacy_pdf_offset is not None: + try: + pdf_page_index, pdf_page_char_offset = legacy_pdf_offset_to_cursor( + data, legacy_pdf_offset + ) + except Exception as exc: # noqa: BLE001 + return f"Invalid continuation token: {exc}" + preview = extract_pdf_page(data, max_chars=min(per_file_chars, 1000)).text + + def page_factory(body_budget: int) -> PageResult: + source_page = extract_pdf_page( + data, + max_chars=body_budget, + page_index=pdf_page_index, + page_char_offset=pdf_page_char_offset, + ) + if source_page.has_more: + next_cursor = { + "file_idx": file_idx, + "file_offset": 0, + "pdf_page_index": source_page.next_cursor["page_index"], + "pdf_page_char_offset": source_page.next_cursor["page_char_offset"], + } + has_more = True + else: + next_cursor = { + "file_idx": file_idx + 1, + "file_offset": 0, + "pdf_page_index": 0, + "pdf_page_char_offset": 0, + } + has_more = file_idx + 1 < len(matched) + return PageResult( + source_page.text, + { + "file_idx": file_idx, + "file_offset": 0, + "pdf_page_index": pdf_page_index, + "pdf_page_char_offset": pdf_page_char_offset, + }, + next_cursor, + has_more, + source_page.warnings, + ) - # Move to the next candidate file. - if file_idx + 1 < len(matched): - next_token = make_token( - _TOOL_NAME, - url, - 0, - matched=_slim_matched(matched), - file_idx=file_idx + 1, - file_offset=0, + else: + full_text = _full_text_for_file(url, data, ft) + preview = full_text[: min(per_file_chars, 1000)] + + def page_factory(body_budget: int) -> PageResult: + source_page = text_page(full_text, file_offset, body_budget) + if source_page.has_more: + next_cursor = { + "file_idx": file_idx, + "file_offset": source_page.next_cursor["offset"], + "pdf_page_index": 0, + "pdf_page_char_offset": 0, + } + has_more = True + else: + next_cursor = { + "file_idx": file_idx + 1, + "file_offset": 0, + "pdf_page_index": 0, + "pdf_page_char_offset": 0, + } + has_more = file_idx + 1 < len(matched) + return PageResult( + source_page.text, + { + "file_idx": file_idx, + "file_offset": file_offset, + "pdf_page_index": 0, + "pdf_page_char_offset": 0, + }, + next_cursor, + has_more, + ) + + if session_id: + store_fetched_file(session_id, name, data, mime, preview, url) + + try: + result = format_page_to_limit( + page_factory, + lambda cursor: _make_fetch_token(matched, cursor, ft), + requested_chars=per_file_chars, + result_limit=settings.max_tool_result_chars, + prefix=prefix, + suffix="\n", + ) + logger.info( + ( + "fetch page type=%s file_idx=%d file_offset=%d pdf_page_index=%d " + "pdf_page_char_offset=%d result_chars=%d has_more=%s" + ), + ft, + file_idx, + file_offset, + pdf_page_index, + pdf_page_char_offset, + len(result), + "---page-info---" in result, ) - return format_paged_result(body, has_more=True, next_token=next_token) + return result + except (PageSizeError, ValueError) as exc: + return f"File pagination failed: {exc}" + - return format_paged_result(body, has_more=False, next_token=None) \ No newline at end of file +def _make_fetch_token( + matched: list[dict[str, str]], cursor: dict[str, int], file_type: str +) -> str: + file_idx = cursor["file_idx"] + url = matched[file_idx].get("href", "") if file_idx < len(matched) else "completed" + return make_token( + _TOOL_NAME, + url, + cursor["file_offset"], + matched=_slim_matched(matched), + file_idx=file_idx, + file_offset=cursor["file_offset"], + pdf_page_index=cursor["pdf_page_index"], + pdf_page_char_offset=cursor["pdf_page_char_offset"], + file_type=file_type, + ) diff --git a/src/ollama_web/tools/helper/__init__.py b/src/ollama_web/tools/helper/__init__.py index f201101..198a300 100644 --- a/src/ollama_web/tools/helper/__init__.py +++ b/src/ollama_web/tools/helper/__init__.py @@ -1,12 +1,13 @@ """Internal helper functions for tools.""" from .image import resize_image -from .pdf import extract_pdf_text +from .pdf import extract_pdf_page, extract_pdf_text from .safe_http import UnsafeURL, fetch_public_bytes, validate_public_http_url __all__ = [ "UnsafeURL", "extract_pdf_text", + "extract_pdf_page", "fetch_public_bytes", "resize_image", "validate_public_http_url", diff --git a/src/ollama_web/tools/helper/pagination.py b/src/ollama_web/tools/helper/pagination.py index 748338b..9d17823 100644 --- a/src/ollama_web/tools/helper/pagination.py +++ b/src/ollama_web/tools/helper/pagination.py @@ -1,32 +1,17 @@ -"""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. -""" +"""Lossless pagination helpers for tool results.""" from __future__ import annotations import base64 import json +from collections.abc import Callable +from dataclasses import dataclass, field 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 +_TOKEN_VERSION = 2 +_SUPPORTED_TOKEN_VERSIONS = {1, 2} +MAX_TOKEN_CHARS = 8192 -# 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---" @@ -34,29 +19,32 @@ class InvalidTokenError(ValueError): """Raised when a continuation token is malformed or invalid.""" +class PageSizeError(ValueError): + """Raised when page metadata alone cannot fit in the result budget.""" + + +@dataclass(frozen=True) +class PageResult: + """One lossless source slice and the cursor immediately after that slice.""" + + text: str + start_cursor: dict[str, int] + next_cursor: dict[str, int] + has_more: bool + warnings: tuple[str, ...] = field(default_factory=tuple) + + 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. - """ + """Create a version-2 opaque continuation token.""" payload: dict[str, Any] = { "v": _TOKEN_VERSION, "tool": tool, @@ -64,36 +52,23 @@ def make_token(tool: str, url: str, offset: int, **extra: Any) -> str: "offset": int(offset), } payload.update(extra) - # ensure_ascii=True keeps the token pure ASCII so truncation of the - # token by the model or tool-result trimming cannot split a multibyte - # UTF-8 sequence (which would raise UnicodeDecodeError on decode). raw = json.dumps(payload, ensure_ascii=True, separators=(",", ":")).encode("ascii") - return _b64url_encode(raw) + token = _b64url_encode(raw) + if len(token) > MAX_TOKEN_CHARS: + raise PageSizeError(f"Continuation token exceeds {MAX_TOKEN_CHARS} characters.") + return token 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``. - """ + """Validate and decode a version-1 or version-2 continuation token.""" if not token or not token.strip(): raise InvalidTokenError("Empty continuation token.") + token = token.strip() + if len(token) > MAX_TOKEN_CHARS: + raise InvalidTokenError(f"Continuation token exceeds {MAX_TOKEN_CHARS} characters.") try: - raw = _b64url_decode(token.strip()) + raw = _b64url_decode(token) except (ValueError, base64.binascii.Error) as exc: # type: ignore[attr-defined] raise InvalidTokenError(f"Token is not valid base64url: {exc}") from exc @@ -104,77 +79,86 @@ def parse_token(token: str | None, expected_tool: str) -> dict[str, Any]: 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 payload.get("v") not in _SUPPORTED_TOKEN_VERSIONS: + raise InvalidTokenError(f"Unsupported token version: {payload.get('v')!r}") + if payload.get("tool") != expected_tool: + raise InvalidTokenError( + f"Token was issued for tool {payload.get('tool')!r}, expected {expected_tool!r}." + ) if "url" not in payload or "offset" not in payload: raise InvalidTokenError("Token payload missing 'url' or 'offset'.") - + if not isinstance(payload["url"], str) or not payload["url"]: + raise InvalidTokenError("Token 'url' is not a non-empty string.") try: payload["offset"] = int(payload["offset"]) except (TypeError, ValueError) as exc: raise InvalidTokenError(f"Token 'offset' is not an integer: {exc}") from exc - + if payload["offset"] < 0: + raise InvalidTokenError("Token 'offset' must not be negative.") 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. - """ + """Append the protected continuation footer when unread content remains.""" 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}" + raise PageSizeError("A continuation token is required when has_more is true.") + return f"{text}{PAGE_INFO_MARKER}\nhas_more: true\nnext_token: {next_token}" + + +def format_page_to_limit( + page_factory: Callable[[int], PageResult], + token_factory: Callable[[dict[str, int]], str], + *, + requested_chars: int, + result_limit: int, + prefix: str = "", + suffix: str = "", +) -> str: + """Render the largest lossless page whose complete envelope fits the limit.""" + body_budget = max(1, min(requested_chars, result_limit)) + for _ in range(8): + page = page_factory(body_budget) + warning_text = "".join(f"[warning] {warning}\n" for warning in page.warnings) + body = f"{prefix}{warning_text}{page.text}{suffix}" + token = token_factory(page.next_cursor) if page.has_more else None + rendered = format_paged_result(body, has_more=page.has_more, next_token=token) + if len(rendered) <= result_limit: + return rendered + overflow = len(rendered) - result_limit + new_budget = body_budget - overflow - 16 + if new_budget >= body_budget: + new_budget = body_budget - 1 + if new_budget <= 0: + break + body_budget = new_budget + raise PageSizeError("Paged result metadata exceeds the configured result limit.") + + +def text_page(text: str, offset: int, max_chars: int) -> PageResult: + """Return a lossless character-offset page from ``text``.""" + if offset < 0 or offset > len(text): + raise ValueError("Text pagination offset is out of range.") + end = min(offset + max_chars, len(text)) + return PageResult( + text=text[offset:end], + start_cursor={"offset": offset}, + next_cursor={"offset": end}, + has_more=end < len(text), + ) 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). - """ + """Split a formatted result into body, has-more flag, and token.""" 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:"): + if line.strip().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 7b9fdac..243a837 100644 --- a/src/ollama_web/tools/helper/pdf.py +++ b/src/ollama_web/tools/helper/pdf.py @@ -1,4 +1,4 @@ -"""PDF text extraction helpers.""" +"""Lossless PDF text extraction helpers.""" from __future__ import annotations @@ -7,78 +7,117 @@ import io from pypdf import PdfReader from ...config import settings +from .pagination import PageResult -def extract_pdf_text( +def extract_pdf_page( 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). - """ + *, + max_chars: int, + page_index: int = 0, + page_char_offset: int = 0, + max_size_mb: int | None = None, +) -> PageResult: + """Extract a PDF slice using a page-and-character continuation cursor.""" + start = {"page_index": page_index, "page_char_offset": page_char_offset} if not 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.", False, 0 + return PageResult("Empty PDF data.", start, start, False) + size_limit_mb = settings.max_fetch_mb if max_size_mb is None else max_size_mb + if len(data) > size_limit_mb * 1024 * 1024: + return PageResult("PDF exceeds the configured size limit.", start, start, False) + if page_index < 0 or page_char_offset < 0: + return PageResult("Invalid PDF continuation cursor.", start, start, False) try: reader = PdfReader(stream=io.BytesIO(data)) except Exception as exc: # noqa: BLE001 - return f"Failed to parse PDF: {exc}", False, 0 + return PageResult(f"Failed to parse PDF: {exc}", start, start, False) 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) + if page_index > total_pages or (page_index == total_pages and page_char_offset): + return PageResult("Invalid PDF continuation cursor.", start, start, False) + if page_index == total_pages: + return PageResult("", start, start, False) parts: list[str] = [] - collected = 0 - last_extracted_page = page_offset - for page_index in range(page_offset, end): + warnings: list[str] = [] + remaining = max(1, max_chars) + current_page = page_index + current_char = page_char_offset + pages_processed = 0 + + while current_page < total_pages and pages_processed < settings.max_pdf_pages: try: - text = reader.pages[page_index].extract_text() - except Exception: # noqa: BLE001 + page_text = reader.pages[current_page].extract_text() or "" + except Exception as exc: # noqa: BLE001 + detail = str(exc).replace("\n", " ")[:120] + warnings.append(f"PDF page {current_page + 1} could not be extracted: {detail}") + if current_page + 1 < total_pages: + parts.append("\n\n") + current_page += 1 + current_char = 0 + pages_processed += 1 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: + + if current_char > len(page_text): + return PageResult("Invalid PDF continuation cursor.", start, start, False) + + available = page_text[current_char:] + take = min(remaining, len(available)) + if take: + parts.append(available[:take]) + current_char += take + remaining -= take + + if current_char < len(page_text): + break + + if current_page + 1 < total_pages: + parts.append("\n\n") + current_page += 1 + current_char = 0 + pages_processed += 1 + if remaining <= 0: break - full = "\n\n".join(parts) - if len(full) > limit: - full = full[:limit].rsplit(" ", 1)[0] + "\n…[truncated]" + next_cursor = {"page_index": current_page, "page_char_offset": current_char} + has_more = current_page < total_pages + text = "".join(parts) + if not text and not warnings and not has_more: + text = "No text could be extracted from the PDF." + return PageResult(text, start, next_cursor, has_more, tuple(warnings)) + + +def legacy_pdf_offset_to_cursor(data: bytes, flat_offset: int) -> tuple[int, int]: + """Translate a v1 ``\n\n``-joined character offset to the v2 PDF cursor.""" + if flat_offset < 0: + raise ValueError("Legacy PDF offset must not be negative.") + reader = PdfReader(stream=io.BytesIO(data)) + remaining = flat_offset + for page_index, page in enumerate(reader.pages): + try: + page_text = page.extract_text() or "" + except Exception: # noqa: BLE001 + page_text = "" + if remaining <= len(page_text): + return page_index, remaining + remaining -= len(page_text) + if page_index + 1 < len(reader.pages): + if remaining <= 2: + return page_index + 1, 0 + remaining -= 2 + if remaining == 0: + return len(reader.pages), 0 + raise ValueError("Legacy PDF offset is out of range.") - # 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 +def extract_pdf_text( + data: bytes, + max_chars: int | None = None, + page_offset: int = 0, +) -> tuple[str, bool, int]: + """Backward-compatible wrapper for page-only internal callers.""" + limit = max_chars if max_chars is not None and max_chars > 0 else settings.fetch_max_chars + page = extract_pdf_page(data, max_chars=limit, page_index=page_offset) + text = "".join(f"[warning] {warning}\n" for warning in page.warnings) + page.text + return text, page.has_more, page.next_cursor["page_index"] if page.has_more else 0 diff --git a/src/ollama_web/tools/pdf.py b/src/ollama_web/tools/pdf.py index 9fdb497..4faf8ac 100644 --- a/src/ollama_web/tools/pdf.py +++ b/src/ollama_web/tools/pdf.py @@ -2,67 +2,148 @@ from __future__ import annotations +import logging +from collections.abc import Callable +from typing import Any + from ..config import settings -from .helper.pagination import format_paged_result, make_token, parse_token -from .helper.pdf import extract_pdf_text +from .helper.pagination import ( + PageResult, + PageSizeError, + format_page_to_limit, + make_token, + parse_token, +) +from .helper.pdf import extract_pdf_page from .helper.safe_http import UnsafeURL, fetch_public_bytes _TOOL_NAME = "pdf_to_text" +_ATTACHMENT_PREFIX = "attachment:" +logger = logging.getLogger("ollama_web.tools.pdf") +AttachmentLoader = Callable[[str], bytes | None] -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. +def _cursor_value(payload: dict[str, Any], key: str, fallback: int = 0) -> int: + raw = payload.get(key, fallback) + if isinstance(raw, bool) or not isinstance(raw, (int, str)): + raise ValueError(f"Token {key!r} is not an integer.") + try: + value = int(raw) + except (TypeError, ValueError) as exc: + raise ValueError(f"Token {key!r} is not an integer.") from exc + if value < 0: + raise ValueError(f"Token {key!r} must not be negative.") + return value - Args: - url_or_path: File path or URL pointing to a PDF. - 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 for this page, followed by a ``---page-info---`` - footer when more pages are available, or an error message. - """ + +def _pdf_to_text( + url_or_path: str, + max_chars: int | None, + continue_from: str | None, + attachment_loader: AttachmentLoader | None, +) -> str: 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 + page_index = 0 + page_char_offset = 0 if continue_from: try: payload = parse_token(continue_from, _TOOL_NAME) + src = str(payload.get("url") or src) + # Version 1 used offset as the PDF page index. + page_index = _cursor_value(payload, "page_index", int(payload["offset"])) + page_char_offset = _cursor_value(payload, "page_char_offset") 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://"): + if src.startswith(_ATTACHMENT_PREFIX): + file_id = src.removeprefix(_ATTACHMENT_PREFIX) + data = attachment_loader(file_id) if attachment_loader is not None else None + if data is None: + return "Attachment PDF is unavailable in this session." + source_kind = "attachment" + elif src.startswith("http://") or src.startswith("https://"): try: - data, _, _ = fetch_public_bytes(src, timeout=settings.scrape_timeout) + data, final_url, _ = fetch_public_bytes(src, timeout=settings.scrape_timeout) + src = final_url or src except UnsafeURL as exc: return f"Blocked URL: {exc}" except Exception as exc: # noqa: BLE001 return f"Failed to fetch PDF: {exc}" + source_kind = "remote" else: return "Only http(s) PDF URLs are supported." - if data is None: - return "No PDF data loaded." + requested = max_chars if max_chars is not None and max_chars > 0 else settings.fetch_max_chars + + def page_factory(body_budget: int) -> PageResult: + return extract_pdf_page( + data, + max_chars=body_budget, + page_index=page_index, + page_char_offset=page_char_offset, + max_size_mb=(settings.max_upload_mb if source_kind == "attachment" else None), + ) + + def token_factory(cursor: dict[str, int]) -> str: + return make_token( + _TOOL_NAME, + src, + cursor["page_index"], + page_index=cursor["page_index"], + page_char_offset=cursor["page_char_offset"], + ) + + try: + result = format_page_to_limit( + page_factory, + token_factory, + requested_chars=requested, + result_limit=settings.max_tool_result_chars, + ) + except PageSizeError as exc: + return f"PDF pagination failed: {exc}" - text, has_more, next_page_offset = extract_pdf_text( - data, max_chars=max_chars, page_offset=page_offset + logger.info( + "pdf page source=%s page_index=%d page_char_offset=%d result_chars=%d has_more=%s", + source_kind, + page_index, + page_char_offset, + len(result), + "---page-info---" in result, ) - 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) + return result + + +def pdf_to_text( + url_or_path: str, max_chars: int | None = None, continue_from: str | None = None +) -> str: + """Extract a losslessly paged PDF from an HTTP(S) URL. + + Args: + url_or_path: HTTP(S) PDF URL. In chat, a session attachment URI may also + be supplied by the attachment context. Arbitrary filesystem paths are + never accepted. + max_chars: Requested maximum source characters for this page. The tool + may return fewer characters so metadata and its token remain intact. + continue_from: Opaque continuation token from the preceding result. + """ + return _pdf_to_text(url_or_path, max_chars, continue_from, None) + + +def make_session_pdf_tool(attachment_loader: AttachmentLoader) -> Callable[..., str]: + """Create a session-scoped ``pdf_to_text`` supporting attachment IDs.""" + + def session_pdf_to_text( + url_or_path: str, + max_chars: int | None = None, + continue_from: str | None = None, + ) -> str: + return _pdf_to_text(url_or_path, max_chars, continue_from, attachment_loader) + + session_pdf_to_text.__name__ = _TOOL_NAME + session_pdf_to_text.__doc__ = pdf_to_text.__doc__ + return session_pdf_to_text diff --git a/src/ollama_web/tools/scrape.py b/src/ollama_web/tools/scrape.py index 39dab83..2f85c38 100644 --- a/src/ollama_web/tools/scrape.py +++ b/src/ollama_web/tools/scrape.py @@ -2,13 +2,37 @@ from __future__ import annotations +import logging + from bs4 import BeautifulSoup from ..config import settings -from .helper.pagination import format_paged_result, make_token, parse_token +from .helper.pagination import ( + PageSizeError, + format_page_to_limit, + make_token, + parse_token, + text_page, +) from .helper.safe_http import UnsafeURL, fetch_public_bytes _TOOL_NAME = "scrape_url" +logger = logging.getLogger("ollama_web.tools.scrape") + + +def extract_readable_text(data: bytes, content_type: str = "") -> str: + """Decode and normalize the complete readable body of an HTML payload.""" + encoding = "utf-8" + if "charset=" in content_type: + encoding = content_type.rsplit("charset=", 1)[-1].split(";", 1)[0].strip() + html = data.decode(encoding or "utf-8", errors="replace") + soup = BeautifulSoup(html, "html.parser") + for tag in soup(["script", "style", "noscript", "template", "head", "nav", "footer", "aside"]): + tag.decompose() + container = soup.find("article") or soup.find("main") or soup + text = container.get_text(separator="\n") + lines = [line.strip() for line in text.splitlines()] + return "\n".join(line for line in lines if line) def scrape_url(url: str, max_chars: int | None = None, continue_from: str | None = None) -> str: @@ -47,43 +71,33 @@ def scrape_url(url: str, max_chars: int | None = None, continue_from: str | None # 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" - if "charset=" in content_type: - encoding = content_type.rsplit("charset=", 1)[-1].split(";", 1)[0].strip() - html = data.decode(encoding or "utf-8", errors="replace") url = final_url except UnsafeURL as exc: return f"Blocked URL: {exc}" except Exception as exc: # noqa: BLE001 return f"Failed to fetch {url}: {exc}" - soup = BeautifulSoup(html, "html.parser") - - # Remove non-content tags. - for tag in soup(["script", "style", "noscript", "template", "head", "nav", "footer", "aside"]): - tag.decompose() - - # Prefer article/main content when available. - container = soup.find("article") or soup.find("main") or soup - text = container.get_text(separator="\n") - # Collapse blank lines and trailing whitespace per line. - lines = [ln.strip() for ln in text.splitlines()] - cleaned = "\n".join(ln for ln in lines if ln) + cleaned = extract_readable_text(data, content_type) 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) + try: + result = format_page_to_limit( + lambda body_budget: text_page(cleaned, offset, body_budget), + lambda cursor: make_token(_TOOL_NAME, url, cursor["offset"]), + requested_chars=limit, + result_limit=settings.max_tool_result_chars, + ) + logger.info( + "scrape page offset=%d result_chars=%d has_more=%s", + offset, + len(result), + "---page-info---" in result, + ) + return result + except (PageSizeError, ValueError) as exc: + return f"Scrape pagination failed: {exc}" -- cgit v1.2.3