summaryrefslogtreecommitdiff
path: root/src/ollama_web/tools/helper
diff options
context:
space:
mode:
Diffstat (limited to 'src/ollama_web/tools/helper')
-rw-r--r--src/ollama_web/tools/helper/__init__.py13
-rw-r--r--src/ollama_web/tools/helper/image.py65
-rw-r--r--src/ollama_web/tools/helper/pdf.py39
-rw-r--r--src/ollama_web/tools/helper/safe_http.py104
4 files changed, 221 insertions, 0 deletions
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/helper/image.py b/src/ollama_web/tools/helper/image.py
new file mode 100644
index 0000000..e89a96f
--- /dev/null
+++ b/src/ollama_web/tools/helper/image.py
@@ -0,0 +1,65 @@
+"""Image preprocessing helpers for ollama vision models."""
+
+from __future__ import annotations
+
+import io
+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",
+ "jpg": "JPEG",
+ "jpeg": "JPEG",
+ "gif": "GIF",
+ "bmp": "BMP",
+ "webp": "WEBP",
+ "tif": "TIFF",
+ "tiff": "TIFF",
+}
+
+
+def _format_from_name(name: str) -> str:
+ ext = Path(name).suffix.lstrip(".").lower()
+ return _IMAGE_FORMATS.get(ext, "JPEG")
+
+
+def resize_image(
+ data: bytes,
+ name: str = "image.jpg",
+ max_dimension: int = 1024,
+ quality: int = 85,
+) -> bytes:
+ """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:
+ if img.width * img.height > settings.max_image_pixels:
+ return data
+ processed = img.convert("RGB")
+
+ 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))
+ processed = processed.resize(new_size, Image.Resampling.LANCZOS)
+
+ fmt = _format_from_name(name)
+ output = io.BytesIO()
+ save_kwargs: dict[str, object] = {}
+ if fmt == "JPEG":
+ save_kwargs["quality"] = quality
+ save_kwargs["optimize"] = True
+ elif fmt == "WEBP":
+ save_kwargs["quality"] = quality
+ processed.save(output, format=fmt, **save_kwargs)
+ return output.getvalue()
+ except Exception: # noqa: BLE001
+ 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.")