summaryrefslogtreecommitdiff
path: root/src/ollama_web/tools/pdf.py
blob: 815c053cc7e5dc58c8d812398d05edd270b596d3 (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
"""PDF text extraction helpers."""

from __future__ import annotations

import io
from pathlib import Path

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.

    Args:
      data: Raw PDF file bytes.
      max_chars: Maximum characters to return (default from settings).

    Returns:
      Extracted plain text. Errors are returned as short strings.
    """
    if not data:
        return "Empty PDF data."

    limit = max_chars if max_chars is not None and max_chars > 0 else settings.fetch_max_chars

    try:
        reader = PdfReader(stream=io.BytesIO(data))
    except Exception as exc:  # noqa: BLE001
        return f"Failed to parse PDF: {exc}"

    parts: list[str] = []
    for page in reader.pages:
        try:
            text = page.extract_text()
            if text:
                parts.append(text)
        except Exception:  # noqa: BLE001
            continue

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


def pdf_to_text(url_or_path: str, max_chars: int | None = None) -> str:
    """Extract text from a local PDF path or remote URL.

    Args:
      url_or_path: File path or URL pointing to a PDF.
      max_chars: Maximum characters to return.

    Returns:
      Extracted plain text or an error message.
    """
    if not url_or_path or not url_or_path.strip():
        return "Invalid PDF source."

    src = url_or_path.strip()
    data: bytes | None = None
    if src.startswith("http://") or src.startswith("https://"):
        import httpx
        from ..config import settings as cfg  # local import to avoid top-level cycles

        try:
            with httpx.Client(
                timeout=cfg.scrape_timeout,
                follow_redirects=True,
                headers={"User-Agent": "ollama-web/0.1 (+https://github.com/local)"},
            ) as client:
                resp = client.get(src)
                resp.raise_for_status()
                data = resp.content
        except Exception as exc:  # noqa: BLE001
            return f"Failed to fetch PDF: {exc}"
    else:
        path = Path(src)
        if not path.exists():
            return f"PDF not found: {src}"
        try:
            data = path.read_bytes()
        except Exception as exc:  # noqa: BLE001
            return f"Failed to read PDF: {exc}"

    if data is None:
        return "No PDF data loaded."
    return extract_pdf_text(data, max_chars=max_chars)