diff options
Diffstat (limited to 'src/ollama_web/tools/helper')
| -rw-r--r-- | src/ollama_web/tools/helper/__init__.py | 3 | ||||
| -rw-r--r-- | src/ollama_web/tools/helper/pagination.py | 192 | ||||
| -rw-r--r-- | src/ollama_web/tools/helper/pdf.py | 155 |
3 files changed, 187 insertions, 163 deletions
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 |