diff options
Diffstat (limited to 'src/ollama_web/tools/helper')
| -rw-r--r-- | src/ollama_web/tools/helper/pagination.py | 177 | ||||
| -rw-r--r-- | src/ollama_web/tools/helper/pdf.py | 65 |
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 |