diff options
Diffstat (limited to 'src/ollama_web/tools')
| -rw-r--r-- | src/ollama_web/tools/__init__.py | 5 | ||||
| -rw-r--r-- | src/ollama_web/tools/fetch.py | 42 | ||||
| -rw-r--r-- | src/ollama_web/tools/helper/__init__.py | 13 | ||||
| -rw-r--r-- | src/ollama_web/tools/helper/image.py (renamed from src/ollama_web/tools/image.py) | 31 | ||||
| -rw-r--r-- | src/ollama_web/tools/helper/pdf.py | 39 | ||||
| -rw-r--r-- | src/ollama_web/tools/helper/safe_http.py | 104 | ||||
| -rw-r--r-- | src/ollama_web/tools/pdf.py | 66 | ||||
| -rw-r--r-- | src/ollama_web/tools/scrape.py | 22 |
8 files changed, 213 insertions, 109 deletions
diff --git a/src/ollama_web/tools/__init__.py b/src/ollama_web/tools/__init__.py index 17ac505..f2a6b0b 100644 --- a/src/ollama_web/tools/__init__.py +++ b/src/ollama_web/tools/__init__.py @@ -1,8 +1,9 @@ """Tool package for ollama-web.""" from .fetch import search_and_fetch -from .image import resize_image -from .pdf import extract_pdf_text, pdf_to_text +from .helper.image import resize_image +from .helper.pdf import extract_pdf_text +from .pdf import pdf_to_text from .registry import ToolRegistry, all_tool_callables, default_registry from .scrape import scrape_url from .search import web_search diff --git a/src/ollama_web/tools/fetch.py b/src/ollama_web/tools/fetch.py index 62670b2..e5af37a 100644 --- a/src/ollama_web/tools/fetch.py +++ b/src/ollama_web/tools/fetch.py @@ -2,19 +2,16 @@ from __future__ import annotations -import base64 import mimetypes import re from pathlib import Path from typing import Any -import httpx - from ..config import settings -from .pdf import extract_pdf_text +from .helper.pdf import extract_pdf_text +from .helper.safe_http import fetch_public_bytes from .scrape import scrape_url - _FetchedFile = dict[str, Any] # Per-session temporary storage for files fetched by search_and_fetch. @@ -43,7 +40,14 @@ def set_session_context(session_id: str | None) -> None: pass -def store_fetched_file(session_id: str, name: str, data: bytes, mime: str, text: str, url: str) -> None: +def store_fetched_file( + session_id: str, + name: str, + data: bytes, + mime: str, + text: str, + url: str, +) -> None: """Store a fetched file so the chat route can persist it to the session.""" if session_id: _fetched_files.setdefault(session_id, []).append( @@ -73,6 +77,8 @@ def _filetype_query(query: str, file_type: str) -> str: def _is_matching_url(url: str, file_type: str) -> bool: suffix = f".{file_type}" + import httpx + parsed = httpx.URL(url) path = parsed.path or "" if path.lower().endswith(suffix): @@ -89,17 +95,13 @@ def _guess_mime(name: str, file_type: str) -> str: def _fetch_bytes(url: str) -> bytes: - with httpx.Client( - timeout=settings.scrape_timeout, - follow_redirects=True, - headers={"User-Agent": "ollama-web/0.1 (+https://github.com/local)"}, - ) as client: - resp = client.get(url) - resp.raise_for_status() - return resp.content + data, _, _ = fetch_public_bytes(url, timeout=settings.scrape_timeout) + return data def _text_from_bytes(url: str, data: bytes, file_type: str, max_chars: int) -> str: + import httpx + name = Path(httpx.URL(url).path or "file").name or f"file.{file_type}" if not name or "." not in name: name = f"file.{file_type}" @@ -151,10 +153,11 @@ def search_and_fetch( if not ft: return "Invalid file type." - limit = ( + limit = min( max_results if max_results is not None and max_results > 0 - else settings.fetch_max_results + else settings.fetch_max_results, + settings.fetch_max_results, ) per_file_chars = ( max_chars if max_chars is not None and max_chars > 0 else settings.fetch_max_chars @@ -183,10 +186,15 @@ def search_and_fetch( try: data = _fetch_bytes(url) except Exception as exc: # noqa: BLE001 - outputs.append(f"=== 候補 {idx} ===\nURL: {url}\nタイトル: {title}\n状態: 取得失敗 ({exc})\n") + outputs.append( + f"=== 候補 {idx} ===\n" + f"URL: {url}\nタイトル: {title}\n状態: 取得失敗 ({exc})\n" + ) continue text = _text_from_bytes(url, data, ft, per_file_chars) + import httpx + mime = _guess_mime(Path(httpx.URL(url).path or f"file.{ft}").name, ft) name = Path(httpx.URL(url).path or f"file.{ft}").name if session_id: diff --git a/src/ollama_web/tools/helper/__init__.py b/src/ollama_web/tools/helper/__init__.py new file mode 100644 index 0000000..f201101 --- /dev/null +++ b/src/ollama_web/tools/helper/__init__.py @@ -0,0 +1,13 @@ +"""Internal helper functions for tools.""" + +from .image import resize_image +from .pdf import extract_pdf_text +from .safe_http import UnsafeURL, fetch_public_bytes, validate_public_http_url + +__all__ = [ + "UnsafeURL", + "extract_pdf_text", + "fetch_public_bytes", + "resize_image", + "validate_public_http_url", +] diff --git a/src/ollama_web/tools/image.py b/src/ollama_web/tools/helper/image.py index 8bc14f5..e89a96f 100644 --- a/src/ollama_web/tools/image.py +++ b/src/ollama_web/tools/helper/image.py @@ -7,6 +7,10 @@ from pathlib import Path from PIL import Image +from ...config import settings + +Image.MAX_IMAGE_PIXELS = settings.max_image_pixels + _IMAGE_FORMATS: dict[str, str] = { "png": "PNG", @@ -31,31 +35,21 @@ def resize_image( max_dimension: int = 1024, quality: int = 85, ) -> bytes: - """Resize and re-encode an image to reduce payload size for ollama. - - Args: - data: Raw image bytes. - name: Original file name, used to guess the output format. - max_dimension: Maximum width or height in pixels. - quality: JPEG/WebP quality when re-encoding to a lossy format. - - Returns: - Resized image bytes. If the input is already small, it may be returned - unchanged when re-encoding is not needed. - """ + """Resize and re-encode an image to reduce payload size for ollama.""" if not data: return data try: with Image.open(io.BytesIO(data)) as img: - # Preserve orientation if EXIF data exists. - img = img.convert("RGB") + if img.width * img.height > settings.max_image_pixels: + return data + processed = img.convert("RGB") - width, height = img.size + width, height = processed.size if width > max_dimension or height > max_dimension: ratio = min(max_dimension / width, max_dimension / height) new_size = (int(width * ratio), int(height * ratio)) - img = img.resize(new_size, Image.Resampling.LANCZOS) + processed = processed.resize(new_size, Image.Resampling.LANCZOS) fmt = _format_from_name(name) output = io.BytesIO() @@ -65,8 +59,7 @@ def resize_image( save_kwargs["optimize"] = True elif fmt == "WEBP": save_kwargs["quality"] = quality - img.save(output, format=fmt, **save_kwargs) + processed.save(output, format=fmt, **save_kwargs) return output.getvalue() except Exception: # noqa: BLE001 - # Fall back to the original bytes if anything goes wrong. - return data
\ No newline at end of file + return data diff --git a/src/ollama_web/tools/helper/pdf.py b/src/ollama_web/tools/helper/pdf.py new file mode 100644 index 0000000..51d3b3a --- /dev/null +++ b/src/ollama_web/tools/helper/pdf.py @@ -0,0 +1,39 @@ +"""PDF text extraction helpers.""" + +from __future__ import annotations + +import io + +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.""" + 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 + + if len(data) > settings.max_fetch_mb * 1024 * 1024: + return "PDF exceeds the configured size limit." + + 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[: settings.max_pdf_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." diff --git a/src/ollama_web/tools/helper/safe_http.py b/src/ollama_web/tools/helper/safe_http.py new file mode 100644 index 0000000..17157db --- /dev/null +++ b/src/ollama_web/tools/helper/safe_http.py @@ -0,0 +1,104 @@ +"""Safe outbound HTTP fetching for LLM tools.""" + +from __future__ import annotations + +import ipaddress +import socket + +import httpx + +from ...config import settings + +_METADATA_IPS = { + ipaddress.ip_address("169.254.169.254"), +} + + +class UnsafeURL(ValueError): + """Raised when an outbound URL is not allowed.""" + + +def _is_public_ip(ip: ipaddress.IPv4Address | ipaddress.IPv6Address) -> bool: + if ip in _METADATA_IPS: + return False + return not ( + ip.is_private + or ip.is_loopback + or ip.is_link_local + or ip.is_multicast + or ip.is_reserved + or ip.is_unspecified + ) + + +def validate_public_http_url(url: str) -> httpx.URL: + """Return a parsed URL after rejecting private network targets.""" + parsed = httpx.URL(url) + if parsed.scheme not in {"http", "https"}: + raise UnsafeURL("Only http(s) URLs are allowed.") + host = parsed.host + if not host: + raise UnsafeURL("URL host is required.") + if host.lower() == "localhost" or host.lower().endswith(".localhost"): + raise UnsafeURL("Localhost URLs are not allowed.") + + try: + ip = ipaddress.ip_address(host) + if not _is_public_ip(ip): + raise UnsafeURL("Private or local IP addresses are not allowed.") + return parsed + except ValueError: + pass + + try: + infos = socket.getaddrinfo( + host, + parsed.port or (443 if parsed.scheme == "https" else 80), + type=socket.SOCK_STREAM, + ) + except socket.gaierror as exc: + raise UnsafeURL(f"Could not resolve host: {host}") from exc + + addresses = {item[4][0] for item in infos} + if not addresses: + raise UnsafeURL(f"Could not resolve host: {host}") + for address in addresses: + ip = ipaddress.ip_address(address) + if not _is_public_ip(ip): + raise UnsafeURL("Resolved private or local IP addresses are not allowed.") + return parsed + + +def fetch_public_bytes( + url: str, + *, + max_bytes: int | None = None, + timeout: float | None = None, +) -> tuple[bytes, str, str]: + """Fetch bytes from a public http(s) URL with redirect and size checks.""" + limit = max_bytes or settings.max_fetch_mb * 1024 * 1024 + current = validate_public_http_url(url) + headers = {"User-Agent": "ollama-web/0.1 (+https://github.com/local)"} + with httpx.Client(timeout=timeout or settings.scrape_timeout, headers=headers) as client: + for _ in range(6): + with client.stream("GET", current, follow_redirects=False) as resp: + if resp.is_redirect: + location = resp.headers.get("location") + if not location: + raise UnsafeURL("Redirect without Location header.") + current = validate_public_http_url(str(current.join(location))) + continue + resp.raise_for_status() + declared = resp.headers.get("content-length") + if declared and int(declared) > limit: + raise UnsafeURL("Remote response exceeds the configured size limit.") + chunks: list[bytes] = [] + total = 0 + for chunk in resp.iter_bytes(): + total += len(chunk) + if total > limit: + raise UnsafeURL("Remote response exceeds the configured size limit.") + chunks.append(chunk) + content_type = resp.headers.get("content-type", "") + return b"".join(chunks), str(resp.url), content_type + raise UnsafeURL("Too many redirects.") diff --git a/src/ollama_web/tools/pdf.py b/src/ollama_web/tools/pdf.py index 815c053..373506c 100644 --- a/src/ollama_web/tools/pdf.py +++ b/src/ollama_web/tools/pdf.py @@ -1,48 +1,10 @@ -"""PDF text extraction helpers.""" +"""PDF tool exposed to Ollama.""" 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." +from .helper.pdf import extract_pdf_text +from .helper.safe_http import UnsafeURL, fetch_public_bytes def pdf_to_text(url_or_path: str, max_chars: int | None = None) -> str: @@ -61,28 +23,14 @@ def pdf_to_text(url_or_path: str, max_chars: int | None = None) -> str: 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 + data, _, _ = fetch_public_bytes(src, timeout=settings.scrape_timeout) + except UnsafeURL as exc: + return f"Blocked URL: {exc}" 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}" + return "Only http(s) PDF URLs are supported." if data is None: return "No PDF data loaded." diff --git a/src/ollama_web/tools/scrape.py b/src/ollama_web/tools/scrape.py index d20956d..d5fa2f4 100644 --- a/src/ollama_web/tools/scrape.py +++ b/src/ollama_web/tools/scrape.py @@ -2,10 +2,10 @@ from __future__ import annotations -import httpx from bs4 import BeautifulSoup from ..config import settings +from .helper.safe_http import UnsafeURL, fetch_public_bytes def scrape_url(url: str, max_chars: int | None = None) -> str: @@ -27,17 +27,15 @@ def scrape_url(url: str, max_chars: int | None = None) -> str: return "Invalid URL." limit = max_chars if max_chars is not None and max_chars > 0 else settings.scrape_max_chars - headers = {"User-Agent": "ollama-web/0.1 (+https://github.com/local)"} - try: - with httpx.Client( - timeout=settings.scrape_timeout, - follow_redirects=True, - headers=headers, - ) as client: - resp = client.get(url) - resp.raise_for_status() - html = resp.text + data, final_url, content_type = fetch_public_bytes(url, timeout=settings.scrape_timeout) + encoding = "utf-8" + if "charset=" in content_type: + encoding = content_type.rsplit("charset=", 1)[-1].split(";", 1)[0].strip() + html = data.decode(encoding or "utf-8", errors="replace") + url = final_url + except UnsafeURL as exc: + return f"Blocked URL: {exc}" except Exception as exc: # noqa: BLE001 return f"Failed to fetch {url}: {exc}" @@ -56,4 +54,4 @@ def scrape_url(url: str, max_chars: int | None = None) -> str: if len(cleaned) > limit: cleaned = cleaned[:limit].rsplit(" ", 1)[0] + "\n…[truncated]" - return cleaned or "No readable text content found."
\ No newline at end of file + return cleaned or "No readable text content found." |