summaryrefslogtreecommitdiff
path: root/tests
diff options
context:
space:
mode:
authortanaka takashi <isazaka@gmail.com>2026-08-23 17:28:10 +0900
committertanaka takashi <isazaka@gmail.com>2026-08-23 17:28:39 +0900
commit33f6b4dbb69e9366e312a8b44dbed9d22e7cba5b (patch)
tree29a7480317af94fefe9f69a53095af2f233207ce /tests
parent68969bfbeea68ded7ff863d229760562072f910b (diff)
fix(tools): PDFページングを損失なしカーソル方式へ再設計
PDF本文をページ単位で抽出した後に文字数上限で切断しつつ、継続位置だけ次ページへ進めていたため、ページ末尾が復元不能になる問題を修正する。ツール側の本文上限とLLM側の結果上限の不整合、25ページの処理上限、添付経路での継続情報欠落も共通契約で解消する。 - 共通PageResultと、メタデータ込みの結果上限内で本文量を決めるフォーマッターを追加 - v2トークンにPDFのページ番号とページ内文字位置を保持し、v1トークンの読み取り互換性を維持 - pdf_to_text、scrape_url、search_and_fetchを返却済み本文だけカーソルが進む方式へ統一 - max_pdf_pagesを文書全体ではなく1呼び出しの処理上限として扱い、長い単一ページと26ページ目以降の継続取得に対応 - attachment:<file_id>をセッション限定で解決し、既存添付も保存済みプレビューではなく元PDFから再解析 - Web取得ファイルをセッションとURL単位で重複排除し、HTMLおよびテキスト経路の全文ページングも維持 - 正常なページング結果の後段切断を禁止し、結果上限・コンテキスト予算・ラウンド上限では未取得状態を明示 - 日英プロンプト、UIメッセージ、ツールスキーマ、変更履歴を実装へ同期 - 実PDF、30ページ超、長い単一ページ、旧トークン、セッション分離、コンテキスト境界の回帰テストを追加 検証: - pytest: 130 passed - ruff check src tests: passed - mypy src: passed - c2144.pdf: 19ページ、抽出本文46015文字(ページ区切り込み46051文字)を8回で重複・欠落なく完全再構成
Diffstat (limited to 'tests')
-rw-r--r--tests/test_llm.py109
-rw-r--r--tests/test_tools.py343
2 files changed, 411 insertions, 41 deletions
diff --git a/tests/test_llm.py b/tests/test_llm.py
index 0e1b146..edaaef2 100644
--- a/tests/test_llm.py
+++ b/tests/test_llm.py
@@ -236,3 +236,112 @@ def test_stream_chat_with_tools_disables_mcp_tools_for_final_answer_after_limit(
assert any(event.get("type") == "status" for event in events)
assert any(event == {"type": "delta", "content": "final answer"} for event in events)
assert events[-1] == {"type": "done"}
+
+
+def test_paged_source_limit_emits_incomplete_warning(
+ monkeypatch: pytest.MonkeyPatch,
+) -> None:
+ from ollama_web import llm
+
+ registry = ToolRegistry()
+
+ def scrape_url(url: str) -> str:
+ """Return a paged result."""
+ return "body\n---page-info---\nhas_more: true\nnext_token: token"
+
+ registry.register(scrape_url)
+
+ class FakeClient:
+ def __init__(self) -> None:
+ self.calls: list[dict[str, Any]] = []
+
+ def chat(self, **kwargs: Any) -> list[SimpleNamespace]:
+ self.calls.append(kwargs)
+ if len(self.calls) == 1:
+ return [
+ SimpleNamespace(
+ message={
+ "content": "",
+ "tool_calls": [
+ {
+ "function": {
+ "name": "scrape_url",
+ "arguments": {"url": "https://example.com"},
+ }
+ }
+ ],
+ }
+ )
+ ]
+ return [SimpleNamespace(message={"content": "partial answer"})]
+
+ client = FakeClient()
+ monkeypatch.setattr(llm, "get_client", lambda _host=None: client)
+ monkeypatch.setattr(llm, "MAX_TOOL_ROUNDS", 1)
+ events = list(
+ llm.stream_chat_with_tools(
+ [Message(role="user", content="read all")],
+ model="test-model",
+ registry=registry,
+ capabilities={"tools"},
+ language="en",
+ )
+ )
+ assert any(event.get("type") == "warning" for event in events)
+ final_instruction = str(client.calls[-1]["messages"][-1].content)
+ assert "unread content" in final_instruction
+
+
+def test_paged_tool_is_not_executed_when_context_budget_is_exhausted(
+ monkeypatch: pytest.MonkeyPatch,
+) -> None:
+ from ollama_web import llm
+
+ executed: list[str] = []
+ registry = ToolRegistry()
+
+ def pdf_to_text(url_or_path: str) -> str:
+ """Read a PDF."""
+ executed.append(url_or_path)
+ return "should not execute"
+
+ registry.register(pdf_to_text)
+
+ class FakeClient:
+ def __init__(self) -> None:
+ self.calls = 0
+
+ def chat(self, **kwargs: Any) -> list[SimpleNamespace]:
+ self.calls += 1
+ if self.calls == 1:
+ return [
+ SimpleNamespace(
+ message={
+ "content": "",
+ "tool_calls": [
+ {
+ "function": {
+ "name": "pdf_to_text",
+ "arguments": {"url_or_path": "https://example.com/a.pdf"},
+ }
+ }
+ ],
+ }
+ )
+ ]
+ return [SimpleNamespace(message={"content": "partial"})]
+
+ monkeypatch.setattr(llm, "get_client", lambda _host=None: FakeClient())
+ monkeypatch.setattr(llm, "MAX_TOOL_CONTEXT_CHARS", 1000)
+ monkeypatch.setattr(llm, "MAX_TOOL_ROUNDS", 1)
+ events = list(
+ llm.stream_chat_with_tools(
+ [Message(role="user", content="read all")],
+ model="test-model",
+ registry=registry,
+ capabilities={"tools"},
+ language="en",
+ )
+ )
+ assert executed == []
+ assert any("budget exhausted" in event.get("result", "") for event in events)
diff --git a/tests/test_tools.py b/tests/test_tools.py
index 387cbe2..95f6514 100644
--- a/tests/test_tools.py
+++ b/tests/test_tools.py
@@ -2,7 +2,9 @@
from __future__ import annotations
+import asyncio
import json
+from io import BytesIO
from unittest.mock import MagicMock, patch
import pytest
@@ -268,20 +270,30 @@ def test_scrape_url_short_text_no_next_token():
# ---------------------------------------------------------------------------
-def _make_multipage_pdf(num_pages: int, text_per_page: str = "Page content. ") -> bytes:
- """Build an in-memory multi-page PDF using pypdf's writer."""
+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()
- for _ in range(num_pages):
- writer.add_blank_page(width=200, height=200)
- # pypdf cannot easily inject text into blank pages; we rely on
- # extract_text returning empty for blank pages. To get real text we use
- # a reportlab-free approach: attach a content stream. For test purposes
- # we instead patch extract_pdf_text to return deterministic text.
- import io
-
- buf = io.BytesIO()
+ 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()
@@ -302,47 +314,221 @@ def test_pdf_to_text_continue_from_invalid_token():
assert "Invalid continuation 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):
- """Pagination across pages reconstructs all page text."""
+ """Real PDF pagination reconstructs long pages and documents over 25 pages."""
+ from pypdf import PdfReader
+
from ollama_web.tools import pdf as pdf_mod
- from ollama_web.tools.helper import pdf as pdf_helper
-
- # Patch extract_pdf_text to return deterministic pages.
- pages = [f"PAGE{i} " + ("X" * 50) for i in range(6)]
-
- def fake_extract(data, max_chars=None, page_offset=0):
- limit = max_chars if max_chars and max_chars > 0 else 10_000
- collected = ""
- last = page_offset
- for idx in range(page_offset, len(pages)):
- collected += pages[idx]
- last = idx
- if len(collected) >= limit:
- break
- has_more = (last + 1) < len(pages)
- return collected, has_more, last + 1 if has_more else 0
-
- monkeypatch.setattr(pdf_helper, "extract_pdf_text", fake_extract)
- monkeypatch.setattr(pdf_mod, "extract_pdf_text", fake_extract)
+
+ 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: (b"%PDF-1.4", "", "application/pdf"),
+ lambda url, timeout=None: (data, url, "application/pdf"),
)
collected = ""
token = None
- for _ in range(10):
+ for _ in range(50):
result = pdf_mod.pdf_to_text(
- "https://example.com/doc.pdf", max_chars=100, continue_from=token
+ "https://example.com/doc.pdf", max_chars=8000, continue_from=token
)
+ assert len(result) <= 6000
body_text, has_more, token = parse_paged_result(result)
collected += body_text
if not has_more:
break
- assert collected == "".join(pages)
+ assert collected == expected
+
+
+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}", 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("<attached-file", 1)[1]
+ token = first_result.split("next_token:", 1)[1].splitlines()[0].strip()
+ resumed = registry.execute(
+ "pdf_to_text",
+ {"url_or_path": f"attachment:{chat_file.id}", "continue_from": token},
+ )
+ assert "unavailable" not in resumed.lower()
+
+
+def test_uploaded_pdf_can_be_resumed_from_original_file(tmp_path, monkeypatch):
+ from starlette.testclient import TestClient
+
+ from ollama_web.app import create_app
+ from ollama_web.config import settings
+ from ollama_web.routes.chat import _build_registry
+
+ data = _make_multipage_pdf(["UPLOAD-" + "U" * 8000])
+ monkeypatch.setattr(settings, "data_dir", str(tmp_path))
+ monkeypatch.setattr(settings, "pin", "123456")
+ app = create_app()
+ client = TestClient(app)
+ login = client.post("/api/auth/login", json={"pin": "123456"})
+ csrf = login.headers["x-csrf-token"]
+ session_id = client.post(
+ "/api/sessions", headers={"X-CSRF-Token": csrf}
+ ).json()["id"]
+ uploaded = client.post(
+ f"/api/sessions/{session_id}/files",
+ headers={"X-CSRF-Token": csrf},
+ files={"files": ("upload.pdf", data, "application/pdf")},
+ ).json()["files"][0]
+ assert len(uploaded["text"]) <= settings.max_tool_result_chars
+
+ registry = asyncio.run(_build_registry(session_id, app.state.session_store))
+ first = registry.execute("pdf_to_text", {"url_or_path": f"attachment:{uploaded['id']}"})
+ _, has_more, token = parse_paged_result(first)
+ assert has_more is True
+ resumed = registry.execute(
+ "pdf_to_text",
+ {"url_or_path": f"attachment:{uploaded['id']}", "continue_from": token},
+ )
+ assert "unavailable" not in resumed.lower()
# ---------------------------------------------------------------------------
@@ -367,6 +553,27 @@ def test_search_and_fetch_continue_from_invalid_token():
assert "Invalid continuation token" in result
+def test_search_and_fetch_v2_token_is_bound_to_file_type(monkeypatch):
+ from ollama_web.tools import fetch as fetch_mod
+
+ monkeypatch.setattr(
+ fetch_mod,
+ "web_search",
+ lambda query, max_results=None: [
+ {"title": "Text", "href": "https://example.com/file.txt", "body": ""}
+ ],
+ )
+ monkeypatch.setattr(
+ fetch_mod,
+ "fetch_public_bytes",
+ lambda url, timeout=None: (b"A" * 200, url, "text/plain"),
+ )
+ first = fetch_mod.search_and_fetch("query", "txt", max_chars=50)
+ _, _, token = parse_paged_result(first)
+ result = fetch_mod.search_and_fetch("query", "pdf", continue_from=token)
+ assert "file type does not match" in result
+
+
def test_search_and_fetch_pagination_across_files(monkeypatch):
"""Pagination walks through multiple candidate files in order."""
from ollama_web.tools import fetch as fetch_mod
@@ -403,6 +610,37 @@ def test_search_and_fetch_pagination_across_files(monkeypatch):
assert seen_titles == ["File1", "File2"]
+def test_search_and_fetch_long_html_uses_full_cleaned_source(monkeypatch):
+ from ollama_web.tools import fetch as fetch_mod
+
+ source = "".join(f"paragraph-{index:04d} " for index in range(1200))
+ html = f"<html><body><main>{source}</main></body></html>".encode()
+ monkeypatch.setattr(
+ fetch_mod,
+ "web_search",
+ lambda query, max_results=None: [
+ {"title": "Long HTML", "href": "https://example.com/page.html", "body": ""}
+ ],
+ )
+ monkeypatch.setattr(
+ fetch_mod,
+ "fetch_public_bytes",
+ lambda url, timeout=None: (html, url, "text/html; charset=utf-8"),
+ )
+
+ token = None
+ collected = ""
+ for _ in range(20):
+ result = fetch_mod.search_and_fetch(
+ "query", "html", max_results=1, continue_from=token
+ )
+ body, has_more, token = parse_paged_result(result)
+ collected += body.split("---\n", 1)[1].removesuffix("\n")
+ if not has_more:
+ break
+ assert collected == source.strip()
+
+
def test_make_token_ensure_ascii_true():
"""Tokens are pure ASCII even with non-ASCII content (ensure_ascii=True)."""
token = make_token("scrape_url", "https://example.com/日本語", 0, title="論文–要約")
@@ -531,7 +769,32 @@ def test_make_token_and_parse_token_roundtrip():
assert payload["url"] == "https://example.com/page"
assert payload["offset"] == 4000
assert payload["tool"] == "scrape_url"
+ assert payload["v"] == 2
+
+
+def test_parse_token_accepts_legacy_v1_payload():
+ import base64
+
+ raw = json.dumps(
+ {"v": 1, "tool": "pdf_to_text", "url": "https://example.com/a.pdf", "offset": 3}
+ ).encode("ascii")
+ token = base64.urlsafe_b64encode(raw).decode("ascii").rstrip("=")
+ payload = parse_token(token, "pdf_to_text")
assert payload["v"] == 1
+ assert payload["offset"] == 3
+
+
+def test_parse_token_rejects_negative_and_oversized_tokens():
+ import base64
+
+ raw = json.dumps(
+ {"v": 1, "tool": "scrape_url", "url": "https://example.com", "offset": -1}
+ ).encode("ascii")
+ negative = base64.urlsafe_b64encode(raw).decode("ascii").rstrip("=")
+ with pytest.raises(InvalidTokenError):
+ parse_token(negative, "scrape_url")
+ with pytest.raises(InvalidTokenError):
+ parse_token("A" * 9000, "scrape_url")
def test_parse_token_invalid_base64():
@@ -609,18 +872,16 @@ def test_format_paged_result_marker_at_end():
# ---------------------------------------------------------------------------
-def test_trim_tool_result_preserves_page_info():
+def test_trim_tool_result_rejects_oversized_paged_result_atomically():
from ollama_web.llm import _trim_tool_result
body = "A" * 5000
footer = "\n---page-info---\nhas_more: true\nnext_token: tok123"
result = body + footer
trimmed = _trim_tool_result(result, 200)
- assert "---page-info---" in trimmed
- assert "next_token: tok123" in trimmed
- assert "has_more: true" in trimmed
- # Body should be truncated.
- assert len(trimmed) < len(result)
+ assert "Repeat the same tool call" in trimmed
+ assert "tok123" not in trimmed
+ assert "A" * 20 not in trimmed
def test_trim_tool_result_no_page_info_unchanged():