summaryrefslogtreecommitdiff
path: root/src/ollama_web/tools/pdf.py
blob: 657edf950bd367100119c8e37aea55c07c99f05c (plain)
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
"""PDF tool exposed to Ollama."""

from __future__ import annotations

import logging
from collections.abc import Callable
from typing import Any

from ..config import settings
from .helper.pagination import (
    PageResult,
    PageSizeError,
    format_page_to_limit,
    make_token,
    parse_paged_result,
    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]

_RESUME_RULE = (
    "\nresume_tool: pdf_to_text"
    "\nresume_page_index: {page_index}"
    "\nresume_page_char_offset: {page_char_offset}"
    "\nnext_token_is_latest: true"
    "\nsupersedes_previous_tokens: true"
    "\nresume_with_continue_from_only: true"
    "\npreserve_next_token_on_tool_limit: true"
)


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


def _pdf_to_text(
    url_or_path: str | None,
    max_chars: int | None,
    continue_from: str | None,
    attachment_loader: AttachmentLoader | None,
) -> str:
    src = (url_or_path or "").strip()
    page_index = 0
    page_char_offset = 0
    if continue_from:
        try:
            payload = parse_token(continue_from, _TOOL_NAME)
            # A continuation token is self-contained. Its source is
            # authoritative so a later conversation cannot accidentally pair
            # the cursor with a stale attachment ID or URL.
            src = str(payload["url"])
            # 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}\n"
                "restart_from_beginning: false\n"
                "required_action: reuse_latest_next_token"
            )
    elif not src:
        return "Invalid PDF source."

    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, 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."

    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_limit = settings.max_tool_result_chars
        render_limit = result_limit
        result = ""
        for _ in range(4):
            result = format_page_to_limit(
                page_factory,
                token_factory,
                requested_chars=requested,
                result_limit=render_limit,
            )
            _, has_more, next_token = parse_paged_result(result)
            if not has_more or next_token is None:
                break
            next_payload = parse_token(next_token, _TOOL_NAME)
            resume_metadata = _RESUME_RULE.format(
                page_index=_cursor_value(next_payload, "page_index", next_payload["offset"]),
                page_char_offset=_cursor_value(next_payload, "page_char_offset"),
            )
            overflow = len(result) + len(resume_metadata) - result_limit
            if overflow <= 0:
                result += resume_metadata
                break
            render_limit -= overflow + 8
            if render_limit <= 0:
                raise PageSizeError("PDF continuation metadata exceeds the result limit.")
        else:
            raise PageSizeError("PDF continuation metadata could not fit the result limit.")
    except PageSizeError as exc:
        return f"PDF pagination failed: {exc}"

    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,
    )
    return result


def pdf_to_text(
    url_or_path: str | None = None,
    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 for the first call. In chat, a session
        attachment URI may also be supplied by the attachment context.
        Omit this argument when continue_from is supplied. 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, self-contained continuation token from the most
        recent result for this PDF. It restores both the source and cursor.
    """
    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 | None = None,
        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