summaryrefslogtreecommitdiff
path: root/src/ollama_web/tools
diff options
context:
space:
mode:
Diffstat (limited to 'src/ollama_web/tools')
-rw-r--r--src/ollama_web/tools/__init__.py21
-rw-r--r--src/ollama_web/tools/image.py72
2 files changed, 92 insertions, 1 deletions
diff --git a/src/ollama_web/tools/__init__.py b/src/ollama_web/tools/__init__.py
index c034c94..17ac505 100644
--- a/src/ollama_web/tools/__init__.py
+++ b/src/ollama_web/tools/__init__.py
@@ -1 +1,20 @@
-"""Tool package for ollama-web.""" \ No newline at end of file
+"""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 .registry import ToolRegistry, all_tool_callables, default_registry
+from .scrape import scrape_url
+from .search import web_search
+
+__all__ = [
+ "search_and_fetch",
+ "resize_image",
+ "extract_pdf_text",
+ "pdf_to_text",
+ "ToolRegistry",
+ "all_tool_callables",
+ "default_registry",
+ "scrape_url",
+ "web_search",
+]
diff --git a/src/ollama_web/tools/image.py b/src/ollama_web/tools/image.py
new file mode 100644
index 0000000..8bc14f5
--- /dev/null
+++ b/src/ollama_web/tools/image.py
@@ -0,0 +1,72 @@
+"""Image preprocessing helpers for ollama vision models."""
+
+from __future__ import annotations
+
+import io
+from pathlib import Path
+
+from PIL import Image
+
+
+_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.
+
+ 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.
+ """
+ if not data:
+ return data
+
+ try:
+ with Image.open(io.BytesIO(data)) as img:
+ # Preserve orientation if EXIF data exists.
+ img = img.convert("RGB")
+
+ width, height = img.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)
+
+ 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
+ img.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