summaryrefslogtreecommitdiff
path: root/src/ollama_web/tools/pdf.py
blob: 4faf8ac166728f8cd381ba66dadcf7ebaa6cc92a (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
"""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_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 _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,
    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()
    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}"

    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 = 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}"

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