"""Tests for tool registry and built-in tools."""
from __future__ import annotations
import asyncio
import json
from io import BytesIO
from unittest.mock import MagicMock, patch
import pytest
from ollama_web.tools.helper.pagination import (
InvalidTokenError,
format_paged_result,
make_token,
parse_paged_result,
parse_token,
)
from ollama_web.tools.registry import ToolRegistry, default_registry
# ---------------------------------------------------------------------------
# Registry tests (existing)
# ---------------------------------------------------------------------------
def test_registry_register_and_execute():
reg = ToolRegistry()
def add(a: int, b: int) -> int:
"""Add two numbers.
Args:
a: First number.
b: Second number.
Returns:
The sum.
"""
return a + b
reg.register(add)
assert reg.get("add") is add
result = reg.execute("add", {"a": 2, "b": 3})
assert result == "5"
def test_registry_execute_unknown():
reg = ToolRegistry()
out = reg.execute("nonexistent", {})
assert "error" in json.loads(out)
def test_registry_execute_invalid_json_args():
reg = ToolRegistry()
def echo(x: str) -> str:
return x
reg.register(echo)
out = reg.execute("echo", "{not json")
assert "error" in json.loads(out)
def test_registry_execute_handles_exceptions():
reg = ToolRegistry()
def boom() -> None:
raise ValueError("boom")
reg.register(boom)
out = reg.execute("boom", {})
assert "error" in json.loads(out)
def test_default_registry_has_tools():
reg = default_registry()
assert "web_search" in reg.names
assert "scrape_url" in reg.names
assert "pdf_to_text" in reg.names
assert "search_and_fetch" in reg.names
# ---------------------------------------------------------------------------
# web_search tests (existing)
# ---------------------------------------------------------------------------
def test_web_search_empty_query():
from ollama_web.tools.search import web_search
assert web_search("") == []
def test_web_search_mocked():
from ollama_web.tools import search as search_mod
fake = [{"title": "T", "href": "http://x", "body": "B"}]
class FakeDDGS:
def __enter__(self):
return self
def __exit__(self, *a):
return None
def text(self, q, max_results=5):
return fake
monkey = MagicMock()
monkey.DDGS = FakeDDGS
with patch.object(search_mod, "DDGS", FakeDDGS):
from ollama_web.tools.search import web_search
result = web_search("test", max_results=1)
assert len(result) == 1
assert result[0]["title"] == "T"
def test_web_search_returns_empty_list_for_no_provider_results() -> None:
from ollama_web.tools import search as search_mod
class EmptyDDGS:
def __enter__(self):
return self
def __exit__(self, *args):
return None
def text(self, query, max_results=5):
return []
with patch.object(search_mod, "DDGS", EmptyDDGS):
assert search_mod.web_search("nothing found") == []
def test_web_search_logs_and_propagates_provider_errors(
caplog: pytest.LogCaptureFixture,
) -> None:
from ollama_web.tools import search as search_mod
class FailingDDGS:
def __enter__(self):
return self
def __exit__(self, *args):
return None
def text(self, query, max_results=5):
raise RuntimeError("403 Forbidden")
with (
patch.object(search_mod, "DDGS", FailingDDGS),
caplog.at_level("ERROR", logger="ollama_web.tools.search"),
):
with pytest.raises(RuntimeError, match="403 Forbidden"):
search_mod.web_search("private search terms")
assert "DuckDuckGo web search failed" in caplog.text
assert "private search terms" not in caplog.text
registry = ToolRegistry()
registry.register(search_mod.web_search)
with patch.object(search_mod, "DDGS", FailingDDGS):
result = json.loads(registry.execute("web_search", {"query": "test"}))
assert result == {"error": "RuntimeError: 403 Forbidden"}
# ---------------------------------------------------------------------------
# scrape_url tests (existing + pagination)
# ---------------------------------------------------------------------------
def test_scrape_url_empty():
from ollama_web.tools.scrape import scrape_url
assert scrape_url("") == "Invalid URL."
def _fake_fetch_public_bytes(html: bytes, content_type: str = "text/html; charset=utf-8"):
"""Return a fake fetch_public_bytes that yields the given HTML."""
return MagicMock(return_value=(html, "https://example.com/page", content_type))
def test_scrape_url_no_continue_from_returns_first_page():
"""continue_from=None returns the first page with has_more + next_token."""
from ollama_web.tools import scrape as scrape_mod
# Build HTML longer than max_chars so pagination kicks in.
body = "" + ("A" * 100) + ""
html = f"
{body}"
with patch.object(
scrape_mod, "fetch_public_bytes", _fake_fetch_public_bytes(html.encode("utf-8"))
):
result = scrape_mod.scrape_url("https://example.com/page", max_chars=40)
body_text, has_more, next_token = parse_paged_result(result)
assert len(body_text) == 40
assert has_more is True
assert next_token is not None
def test_scrape_url_continue_from_returns_next_page():
"""Passing next_token returns the continuation page."""
from ollama_web.tools import scrape as scrape_mod
body = "" + ("A" * 100) + ""
html = f"{body}"
with patch.object(
scrape_mod, "fetch_public_bytes", _fake_fetch_public_bytes(html.encode("utf-8"))
):
first = scrape_mod.scrape_url("https://example.com/page", max_chars=40)
_, _, next_token = parse_paged_result(first)
with patch.object(
scrape_mod, "fetch_public_bytes", _fake_fetch_public_bytes(html.encode("utf-8"))
):
second = scrape_mod.scrape_url(
"https://example.com/page", max_chars=40, continue_from=next_token
)
body_text, has_more, _ = parse_paged_result(second)
# Second page should start at offset 40.
assert body_text == "A" * 40
assert has_more is True
def test_scrape_url_pagination_covers_full_text():
"""Iterating all pages reconstructs the full cleaned text."""
from ollama_web.tools import scrape as scrape_mod
body = "" + ("A" * 100) + ""
html = f"{body}"
collected = ""
token = None
for _ in range(10):
with patch.object(
scrape_mod, "fetch_public_bytes", _fake_fetch_public_bytes(html.encode("utf-8"))
):
result = scrape_mod.scrape_url(
"https://example.com/page", max_chars=40, continue_from=token
)
body_text, has_more, token = parse_paged_result(result)
collected += body_text
if not has_more:
break
assert collected == "A" * 100
def test_scrape_url_continue_from_invalid_token():
"""An invalid continuation token returns an error message."""
from ollama_web.tools.scrape import scrape_url
result = scrape_url("https://example.com/page", continue_from="not-a-valid-token")
assert "Invalid continuation token" in result
def test_scrape_url_short_text_no_next_token():
"""Text shorter than max_chars has no next_token."""
from ollama_web.tools import scrape as scrape_mod
html = b"short"
with patch.object(scrape_mod, "fetch_public_bytes", _fake_fetch_public_bytes(html)):
result = scrape_mod.scrape_url("https://example.com/page", max_chars=4000)
body_text, has_more, next_token = parse_paged_result(result)
assert body_text == "short"
assert has_more is False
assert next_token is None
# ---------------------------------------------------------------------------
# pdf_to_text tests (pagination)
# ---------------------------------------------------------------------------
def _make_multipage_pdf(page_texts: list[str]) -> bytes:
"""Build a text-bearing PDF fixture without an additional dependency."""
from pypdf import PdfWriter
from pypdf.generic import DecodedStreamObject, DictionaryObject, NameObject
writer = PdfWriter()
font = DictionaryObject(
{
NameObject("/Type"): NameObject("/Font"),
NameObject("/Subtype"): NameObject("/Type1"),
NameObject("/BaseFont"): NameObject("/Helvetica"),
}
)
font_ref = writer._add_object(font)
for text in page_texts:
page = writer.add_blank_page(width=600, height=800)
page[NameObject("/Resources")] = DictionaryObject(
{NameObject("/Font"): DictionaryObject({NameObject("/F1"): font_ref})}
)
escaped = text.replace("\\", "\\\\").replace("(", "\\(").replace(")", "\\)")
stream = DecodedStreamObject()
stream.set_data(f"BT /F1 10 Tf 20 760 Td ({escaped}) Tj ET".encode("latin-1"))
page[NameObject("/Contents")] = writer._add_object(stream)
buf = BytesIO()
writer.write(buf)
return buf.getvalue()
def test_pdf_to_text_only_https_supported():
"""Local paths are rejected (regression)."""
from ollama_web.tools.pdf import pdf_to_text
result = pdf_to_text("/local/path.pdf")
assert "Only http(s) PDF URLs are supported." in result
def test_pdf_to_text_continue_from_invalid_token():
"""Invalid continuation token returns an error."""
from ollama_web.tools.pdf import pdf_to_text
result = pdf_to_text("https://example.com/doc.pdf", continue_from="bad-token")
assert "Invalid continuation token" in result
assert "restart_from_beginning: false" in result
assert "required_action: reuse_latest_next_token" in result
def test_pdf_to_text_rejects_negative_page_character_cursor():
from ollama_web.tools.pdf import pdf_to_text
token = make_token(
"pdf_to_text",
"https://example.com/doc.pdf",
0,
page_index=0,
page_char_offset=-1,
)
result = pdf_to_text("https://example.com/doc.pdf", continue_from=token)
assert "Invalid continuation token" in result
def test_pdf_to_text_pagination_roundtrip(monkeypatch):
"""Real PDF pagination reconstructs long pages and documents over 25 pages."""
from pypdf import PdfReader
from ollama_web.tools import pdf as pdf_mod
pages = [f"PAGE{i:02d}-" + ("X" * (7000 if i == 0 else 80)) for i in range(30)]
data = _make_multipage_pdf(pages)
expected = "\n\n".join(
(page.extract_text() or "") for page in PdfReader(BytesIO(data)).pages
)
# Patch fetch_public_bytes to return dummy bytes.
monkeypatch.setattr(
pdf_mod,
"fetch_public_bytes",
lambda url, timeout=None: (data, url, "application/pdf"),
)
collected = ""
token = None
for _ in range(50):
result = pdf_mod.pdf_to_text(
"https://example.com/doc.pdf" if token is None else "",
max_chars=8000,
continue_from=token,
)
assert len(result) <= 6000
body_text, has_more, token = parse_paged_result(result)
assert "next_token_is_latest:" not in body_text
if has_more:
assert "resume_tool: pdf_to_text" in result
assert "next_token_is_latest: true" in result
assert "supersedes_previous_tokens: true" in result
assert "resume_with_continue_from_only: true" in result
assert "preserve_next_token_on_tool_limit: true" in result
else:
assert "resume_tool:" not in result
collected += body_text
if not has_more:
break
assert collected == expected
def test_pdf_to_text_continuation_token_source_is_authoritative(monkeypatch):
"""A continuation restores its URL and does not require url_or_path."""
from ollama_web.tools import pdf as pdf_mod
data = _make_multipage_pdf(["SOURCE-" + "S" * 9000])
fetched: list[str] = []
def fake_fetch(url: str, timeout=None):
fetched.append(url)
return data, url, "application/pdf"
monkeypatch.setattr(pdf_mod, "fetch_public_bytes", fake_fetch)
first = pdf_mod.pdf_to_text("https://example.com/source.pdf")
_, has_more, token = parse_paged_result(first)
assert has_more is True
assert token is not None
resumed = pdf_mod.pdf_to_text(continue_from=token)
body, _, _ = parse_paged_result(resumed)
assert body
assert fetched == [
"https://example.com/source.pdf",
"https://example.com/source.pdf",
]
# Even a conflicting argument cannot redirect a valid continuation.
pdf_mod.pdf_to_text("https://attacker.example/wrong.pdf", continue_from=token)
assert fetched[-1] == "https://example.com/source.pdf"
def test_pdf_to_text_ollama_schema_allows_token_only_continuation():
from ollama._utils import convert_function_to_tool
from ollama_web.tools.pdf import pdf_to_text
schema = convert_function_to_tool(pdf_to_text).function.parameters
assert "url_or_path" not in (schema.required or [])
assert "continue_from" not in (schema.required or [])
def test_session_pdf_tool_reads_original_bytes_and_is_session_scoped(tmp_path):
from pypdf import PdfReader
from ollama_web.sessions import SessionStore
from ollama_web.tools.pdf import make_session_pdf_tool
data = _make_multipage_pdf(["FIRST-" + "A" * 7000, "SECOND-" + "B" * 200])
expected = "\n\n".join(
(page.extract_text() or "") for page in PdfReader(BytesIO(data)).pages
)
store = SessionStore(tmp_path)
session = store.create()
other_session = store.create()
chat_file = store.add_file(
session["id"],
"document.pdf",
data,
"application/pdf",
"legacy truncated preview",
)
tool = make_session_pdf_tool(lambda file_id: store.get_file_data(session["id"], file_id))
token = None
collected = ""
for _ in range(10):
result = tool(
f"attachment:{chat_file.id}" if token is None else "",
continue_from=token,
)
body, has_more, token = parse_paged_result(result)
collected += body
if not has_more:
break
assert collected == expected
other_tool = make_session_pdf_tool(
lambda file_id: store.get_file_data(other_session["id"], file_id)
)
assert "unavailable" in other_tool(f"attachment:{chat_file.id}").lower()
def test_search_and_fetch_pdf_pagination_is_lossless_and_persists_once(monkeypatch):
from pypdf import PdfReader
from ollama_web.tools import fetch as fetch_mod
data = _make_multipage_pdf([f"P{i:02d}-" + "Z" * 300 for i in range(30)])
expected = "\n\n".join(
(page.extract_text() or "") for page in PdfReader(BytesIO(data)).pages
)
monkeypatch.setattr(
fetch_mod,
"web_search",
lambda query, max_results=None: [
{"title": "PDF", "href": "https://example.com/document.pdf", "body": ""}
],
)
monkeypatch.setattr(
fetch_mod,
"fetch_public_bytes",
lambda url, timeout=None: (data, url, "application/pdf"),
)
token = None
collected = ""
for _ in range(20):
result = fetch_mod.search_and_fetch(
"query", "pdf", max_results=1, session_id="session", continue_from=token
)
assert len(result) <= 6000
body, has_more, token = parse_paged_result(result)
collected += body.split("---\n", 1)[1].removesuffix("\n")
if not has_more:
break
assert collected == expected
assert len(fetch_mod.pop_fetched_files("session")) == 1
def test_search_and_fetch_accepts_legacy_v1_pdf_offset(monkeypatch):
import base64
from pypdf import PdfReader
from ollama_web.tools import fetch as fetch_mod
url = "https://example.com/legacy.pdf"
data = _make_multipage_pdf(["LEGACY-" + "Q" * 500])
expected = PdfReader(BytesIO(data)).pages[0].extract_text() or ""
matched = [{"title": "Legacy", "href": url}]
payload = {
"v": 1,
"tool": "search_and_fetch",
"url": url,
"offset": 100,
"file_offset": 100,
"file_idx": 0,
"matched": matched,
}
token = base64.urlsafe_b64encode(
json.dumps(payload, separators=(",", ":")).encode("ascii")
).decode("ascii").rstrip("=")
monkeypatch.setattr(
fetch_mod, "fetch_public_bytes", lambda fetched_url, timeout=None: (data, url, "")
)
result = fetch_mod.search_and_fetch("query", "pdf", continue_from=token)
body, has_more, _ = parse_paged_result(result)
assert body.split("---\n", 1)[1].removesuffix("\n") == expected[100:]
assert has_more is False
def test_attachment_context_keeps_pdf_continuation_and_registry_can_resume(tmp_path):
from ollama_web.routes.chat import _attachment_context, _build_registry
from ollama_web.sessions import SessionStore
data = _make_multipage_pdf(["ATTACHED-" + "R" * 8000])
store = SessionStore(tmp_path)
session = store.create()
chat_file = store.add_file(
session["id"], "attached.pdf", data, "application/pdf", "old preview"
)
context = _attachment_context(
store, session["id"], [chat_file.id], total_limit=5000
)
assert len(context) <= 5000
assert f'uri="attachment:{chat_file.id}"' in context
assert "next_token:" in context
registry = asyncio.run(_build_registry(session["id"], store))
first_result = context.split("{source}