diff options
| author | tanaka takashi <isazaka@gmail.com> | 2026-08-23 17:28:10 +0900 |
|---|---|---|
| committer | tanaka takashi <isazaka@gmail.com> | 2026-08-23 17:28:39 +0900 |
| commit | 33f6b4dbb69e9366e312a8b44dbed9d22e7cba5b (patch) | |
| tree | 29a7480317af94fefe9f69a53095af2f233207ce | |
| parent | 68969bfbeea68ded7ff863d229760562072f910b (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回で重複・欠落なく完全再構成
| -rw-r--r-- | CHANGELOG.md | 9 | ||||
| -rw-r--r-- | docs/tools-schema.md | 66 | ||||
| -rw-r--r-- | src/ollama_web/i18n/messages/en.json | 3 | ||||
| -rw-r--r-- | src/ollama_web/i18n/messages/ja.json | 3 | ||||
| -rw-r--r-- | src/ollama_web/llm.py | 116 | ||||
| -rw-r--r-- | src/ollama_web/prompts/en.json | 4 | ||||
| -rw-r--r-- | src/ollama_web/prompts/ja.json | 4 | ||||
| -rw-r--r-- | src/ollama_web/routes/chat.py | 86 | ||||
| -rw-r--r-- | src/ollama_web/routes/sessions.py | 17 | ||||
| -rw-r--r-- | src/ollama_web/tools/__init__.py | 3 | ||||
| -rw-r--r-- | src/ollama_web/tools/fetch.py | 290 | ||||
| -rw-r--r-- | src/ollama_web/tools/helper/__init__.py | 3 | ||||
| -rw-r--r-- | src/ollama_web/tools/helper/pagination.py | 192 | ||||
| -rw-r--r-- | src/ollama_web/tools/helper/pdf.py | 155 | ||||
| -rw-r--r-- | src/ollama_web/tools/pdf.py | 153 | ||||
| -rw-r--r-- | src/ollama_web/tools/scrape.py | 70 | ||||
| -rw-r--r-- | tests/test_llm.py | 109 | ||||
| -rw-r--r-- | tests/test_tools.py | 343 |
18 files changed, 1166 insertions, 460 deletions
diff --git a/CHANGELOG.md b/CHANGELOG.md index c3bad6d..6b61e61 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -7,6 +7,15 @@ ## [Unreleased] +### Changed + +- PDFページングをページ番号とページ内文字位置の損失なしカーソルへ変更し、ローカル添付PDFもセッション限定URIで継続取得できるようにしました。 + +### Fixed + +- ツール内部の8,000文字ページがLLM層で6,000文字へ再切断され、継続カーソルとの間に本文欠落が生じる問題を修正しました。 +- 25ページを超えるPDFの後半が`search_and_fetch`で取得されない問題と、継続ごとに同じ取得ファイルが重複保存される問題を修正しました。 + ## [0.5.1] - 2026-08-16 ### Changed diff --git a/docs/tools-schema.md b/docs/tools-schema.md index b9cf4b9..874640a 100644 --- a/docs/tools-schema.md +++ b/docs/tools-schema.md @@ -23,7 +23,7 @@ ollama-web は LLM（Ollama）にツール呼び出し（function calling）機 ツールの登録・実行は [`ToolRegistry`](src/ollama_web/tools/registry.py:17) が担い、チャットループ（[`llm.py`](src/ollama_web/llm.py:1)）経由で LLM から呼び出されます。 -`scrape_url` / `pdf_to_text` / `search_and_fetch` は**ページング（継続トークン方式）に対応**しています。`web_search` は API 呼び出し回数とキャッシュ要件の懸念からページング対象外です。 +`scrape_url` / `pdf_to_text` / `search_and_fetch` は**損失なしページング（継続トークン方式）に対応**しています。本文・メタデータ・継続トークンを含む結果全体をLLM側の上限内に収めてからカーソルを確定します。`web_search` はページング対象外です。 ## ツール一覧 @@ -114,13 +114,13 @@ next_token: {token} ## 3. pdf_to_text - **実装**: [`pdf_to_text()`](src/ollama_web/tools/pdf.py:10) -- **依存**: [`extract_pdf_text()`](src/ollama_web/tools/helper/pdf.py:12)（[`pypdf`](https://pypdf.readthedocs.io/) 使用） +- **依存**: [`extract_pdf_page()`](src/ollama_web/tools/helper/pdf.py:12)（[`pypdf`](https://pypdf.readthedocs.io/) 使用） ### スキーマ | パラメータ | 型 | 必須 | デフォルト | 説明 | |------------|----|------|------------|------| -| `url_or_path` | `str` | ✅ | — | PDF の URL またはローカルパス。空文字列は `"Invalid PDF source."` を返す | +| `url_or_path` | `str` | ✅ | — | HTTP(S) PDF URL。チャット内では現在のセッションの `attachment:<file_id>` も利用可能。任意のローカルパスは拒否 | | `max_chars` | `int \| None` | ❌ | [`settings.fetch_max_chars`](src/ollama_web/config.py:52) (=8000) | 1ページあたりの返却文字数上限 | | `continue_from` | `str \| None` | ❌ | `None` | 前回戻り値の `next_token`。指定時は続きページを取得 | @@ -130,26 +130,26 @@ next_token: {token} ### 制約 -- **`http://` / `https://` のみサポート**。ローカルパスは `"Only http(s) PDF URLs are supported."` を返す。 -- [`extract_pdf_text()`](src/ollama_web/tools/helper/pdf.py:12) 側の制約: +- 通常は **`http://` / `https://` のみサポート**。チャット内のセッションスコープ付きラッパーだけが `attachment:<file_id>` を解決する。 +- [`extract_pdf_page()`](src/ollama_web/tools/helper/pdf.py:12) 側の制約: - [`settings.max_fetch_mb`](src/ollama_web/config.py:53) (=10MB) 超過時は `"PDF exceeds the configured size limit."` - - 1回あたり [`settings.max_pdf_pages`](src/ollama_web/config.py:54) (=25) ページまで抽出 - - `page_offset` から `max_pdf_pages` 件を対象に抽出（ページング時） - - `max_chars` 超過時は単語境界で切り詰め `…[truncated]` を付与 + - 1回あたり [`settings.max_pdf_pages`](src/ollama_web/config.py:54) (=25) ページまで処理するが、26ページ目以降は継続可能 + - `(page_index, page_char_offset)` から再開し、ページ途中の未返却文字を次回返す + - `max_chars` は要求上限で、結果全体を6000文字以内に収めるため本文が短くなる場合がある ### 処理フロー -1. `continue_from` があればトークンから `(url, page_offset)` を復元 +1. `continue_from` があればトークンから `(url, page_index, page_char_offset)` を復元 2. PDF を全件取得・解析 -3. `page_offset` ページ目から順にテキスト抽出し `max_chars` に達したら停止 -4. 残りページがあれば `has_more: true` + `next_token`（`page_offset = 最後に抽出したページ + 1`）を付与 +3. カーソル位置から、結果全体の上限内で実際に返せる文字数だけ抽出 +4. 未返却文字または残りページがあれば、その直後を指す `next_token` を付与 --- ## 4. search_and_fetch - **実装**: [`search_and_fetch()`](src/ollama_web/tools/fetch.py:130) -- **依存**: [`web_search()`](src/ollama_web/tools/search.py:10)、[`fetch_public_bytes()`](src/ollama_web/tools/helper/safe_http.py:72)、[`extract_pdf_text()`](src/ollama_web/tools/helper/pdf.py:12)、[`scrape_url()`](src/ollama_web/tools/scrape.py:11) +- **依存**: [`web_search()`](src/ollama_web/tools/search.py:10)、[`fetch_public_bytes()`](src/ollama_web/tools/helper/safe_http.py:72)、[`extract_pdf_page()`](src/ollama_web/tools/helper/pdf.py:12)、[`scrape_url()`](src/ollama_web/tools/scrape.py:11) ### スキーマ @@ -209,11 +209,11 @@ URL: {url} ### helper/pdf.py -- **実装**: [`extract_pdf_text()`](src/ollama_web/tools/helper/pdf.py:12) +- **実装**: [`extract_pdf_page()`](src/ollama_web/tools/helper/pdf.py:12) - `pypdf.PdfReader` でページごとに `extract_text()`。例外ページはスキップ。 -- `page_offset` パラメータで指定ページから抽出（ページング対応）。 -- 1回あたり `max_pdf_pages` 件まで、`max_chars` で切り詰め。 -- 戻り値は `(text, has_more, next_page_offset)` のタプル。 +- ページ番号とページ内文字位置の複合カーソルで再開する。 +- `max_pdf_pages` は1回の作業量だけを制限し、文書全体は制限しない。 +- 戻り値は本文、開始・次カーソル、`has_more`、ページ警告を持つ `PageResult`。 ### helper/image.py @@ -259,27 +259,24 @@ flowchart LR ## ツール結果の切り詰め（truncation）ポリシー -ツール結果は LLM コンテキストに戻される前に**多段階**で切り詰められる。ページング対応ツールの `---page-info---` フッターは保護され、本文側が優先的に削られる。 +非ページング結果はLLMコンテキストへ戻す前に切り詰められる。ページング対応ツールは本文・メタデータ・トークンを含む結果全体を事前に上限内へ収め、LLM層では切り詰めない。境界違反時は本文と進行済みトークンを返さず、同じ呼び出しを再試行できるエラーにする。 ```mermaid flowchart TD - T[ツール実行] --> S1[ステージ1: ツール内部 max_chars/max_pdf_pages で切り詰め] - S1 --> S2[ステージ2: llm.py _trim_tool_result で max_tool_result_chars=6000 に即時切り詰め] - S2 --> S3[ステージ3: 累積が MAX_TOOL_CONTEXT_CHARS=24000 超過時、各結果をさらに削る] - S3 --> S4[ステージ4: MAX_TOOL_ROUNDS=5 で連続ツール呼び出しを打ち切り] + T[ツール実行] --> S1[ページングツールが結果全体を6000文字以内に整形] + S1 --> S2[非ページング結果だけを llm.py で切り詰め] + S2 --> S3[累積コンテキスト予算到達前に次のページ呼び出しを停止] + S3 --> S4[MAX_TOOL_ROUNDS=5 で一般ツール呼び出しを打ち切り] ``` | ステージ | 場所 | 関数/定数 | 制限値 | |----------|------|-----------|--------| -| 1 | 各ツール内部 | `max_chars` / `max_pdf_pages` | ツールごと（デフォルト 4000〜8000文字 / 25ページ） | -| 2 | [`llm.py`](src/ollama_web/llm.py:1) | [`_trim_tool_result()`](src/ollama_web/llm.py:1) / [`max_tool_result_chars`](src/ollama_web/config.py:58) | **6000文字**（各ツール結果ごとに即時適用） | -| 3 | [`llm.py`](src/ollama_web/llm.py:1) | [`MAX_TOOL_CONTEXT_CHARS`](src/ollama_web/llm.py:1) | **24000文字**（累積超過時に各結果を再削り） | +| 1 | 各ページングツール | `format_page_to_limit()` | **6000文字**（フッター込み） | +| 2 | [`llm.py`](src/ollama_web/llm.py:1) | [`_trim_tool_result()`](src/ollama_web/llm.py:1) | 非ページング結果のみ切り詰め | +| 3 | [`llm.py`](src/ollama_web/llm.py:1) | [`MAX_TOOL_CONTEXT_CHARS`](src/ollama_web/llm.py:1) | **24000文字**（次のページを実行前に停止） | | 4 | [`llm.py`](src/ollama_web/llm.py:1) | [`MAX_TOOL_ROUNDS`](src/ollama_web/llm.py:1) | **5ラウンド**（連続ツール呼び出し上限） | -`_trim_tool_result()` は: -- `---page-info---` フッターを含む結果は、フッターを保護し本文のみを切り詰める（段落境界 `\n\n` 優先、次に単語境界） -- フッター無し時は従来ロジック（段落境界優先、次に単語境界） -- 末尾に `…[tool result truncated]` を付与 +`_trim_tool_result()` はページング結果が上限を超えた場合、本文を部分返却せず再試行可能な境界エラーへ置換する。フッター無しの結果は従来どおり段落・単語境界で切り詰める。 --- @@ -296,7 +293,7 @@ flowchart TD ```json { - "v": 1, + "v": 2, "tool": "scrape_url", "url": "https://example.com/page", "offset": 4000 @@ -307,13 +304,16 @@ flowchart TD ```json { - "v": 1, + "v": 2, "tool": "search_and_fetch", "url": "https://example.com/file.pdf", "offset": 8000, "matched": [{"title": "...", "href": "..."}], "file_idx": 1, - "file_offset": 8000 + "file_offset": 8000, + "pdf_page_index": 3, + "pdf_page_char_offset": 1200, + "file_type": "pdf" } ``` @@ -328,7 +328,7 @@ next_token: {token} ``` - `has_more: false` のときは `next_token` 行を省略（フッター自体も付与しない） -- `---page-info---` は [`_trim_tool_result()`](src/ollama_web/llm.py:1) の切り詰め対象外 +- 正常なページング結果全体が上限内に収まるため、`---page-info---` を含む結果は切り詰め対象外 ### LLM誤用対策 @@ -368,5 +368,3 @@ sequenceDiagram Tool->>Tool: slice offset:offset+max_chars Tool-->>LLM: text + has_more:false ``` - -実装の詳細設計は [`docs/pagination-plan.md`](pagination-plan.md) を参照。
\ No newline at end of file diff --git a/src/ollama_web/i18n/messages/en.json b/src/ollama_web/i18n/messages/en.json index 42bce33..74e0c00 100644 --- a/src/ollama_web/i18n/messages/en.json +++ b/src/ollama_web/i18n/messages/en.json @@ -76,6 +76,7 @@ "connecting": "Connecting…", "disconnected": "Disconnected", "tool_limit_reached": "Tool call limit reached. Summarizing findings…", + "pagination_incomplete": "The source still has unread content. The answer uses only the retrieved portion.", "thinking": "Thinking process", "executing_tool": "Executing {label}…", "running": "Running…", @@ -139,4 +140,4 @@ "invalid_number": "Enter a number." } } -}
\ No newline at end of file +} diff --git a/src/ollama_web/i18n/messages/ja.json b/src/ollama_web/i18n/messages/ja.json index a476ca2..f775ba6 100644 --- a/src/ollama_web/i18n/messages/ja.json +++ b/src/ollama_web/i18n/messages/ja.json @@ -76,6 +76,7 @@ "connecting": "接続中...", "disconnected": "切断しました", "tool_limit_reached": "ツール呼び出し上限に達しました。これまでの調査結果を整理中…", + "pagination_incomplete": "資料に未取得部分があります。取得済み範囲のみで回答します。", "thinking": "思考プロセス", "executing_tool": "{label} を実行中…", "running": "実行中…", @@ -139,4 +140,4 @@ "invalid_number": "数値を入力してください。" } } -}
\ No newline at end of file +} diff --git a/src/ollama_web/llm.py b/src/ollama_web/llm.py index 187b5e7..20485eb 100644 --- a/src/ollama_web/llm.py +++ b/src/ollama_web/llm.py @@ -14,6 +14,7 @@ from .config import settings from .i18n import t from .mcp import redact_secrets from .prompts import get_prompt +from .tools.helper.pagination import parse_paged_result, parse_token from .tools.registry import ToolRegistry, default_registry # Maximum successive tool-call rounds before we force a final answer. @@ -176,29 +177,22 @@ def _token_counts(response: Any) -> dict[str, int] | None: def _trim_tool_result(result: str, limit: int) -> str: """Truncate ``result`` to ``limit`` chars while keeping it readable. - When the result contains a ``---page-info---`` footer (emitted by paged - tools), the footer is preserved and only the body is truncated so the - model can still see the ``next_token`` needed to fetch the continuation. + Paged tools must fit their own complete envelope before advancing a cursor. + An oversized paged result is therefore rejected atomically instead of + being truncated into an unrecoverable content gap. """ if len(result) <= limit: return result - # Protect the page-info footer from truncation. ``split`` keeps the - # marker on the footer side so the ``---page-info---`` line is preserved. + # A paged tool must size its own body before advancing its cursor. Trimming + # it here would create an unrecoverable gap between visible text and the + # next token, so fail the call atomically and let the model repeat it. if _PAGE_INFO_MARKER in result: - body, footer = result.split(_PAGE_INFO_MARKER, 1) - footer = _PAGE_INFO_MARKER + footer - body_limit = limit - len(footer) - if body_limit <= 0: - # The footer alone exceeds the limit; keep it verbatim and drop the - # body entirely so the model can still continue paging. - return "…[tool result truncated]" + footer - if len(body) <= body_limit: - return result - cutoff = body[:body_limit].rsplit("\n\n", 1)[0] - if len(cutoff) < body_limit * 0.5: - cutoff = body[:body_limit].rsplit(" ", 1)[0] - return cutoff + "\n\n…[tool result truncated]" + footer + return ( + "Paged tool result exceeded the context-safe result limit. " + "Repeat the same tool call with a smaller max_chars value; no part " + "of this page was consumed." + ) # No footer: legacy truncation behaviour. cutoff = result[:limit].rsplit("\n\n", 1)[0] @@ -208,6 +202,23 @@ def _trim_tool_result(result: str, limit: int) -> str: return cutoff + "\n\n…[tool result truncated]" +def _paged_source_key(name: str, arguments: Any) -> tuple[str, str]: + """Identify one paged source without exposing its content.""" + if not isinstance(arguments, dict): + return name, "unknown" + if name == "search_and_fetch": + return name, f"{arguments.get('query', '')}|{arguments.get('file_type', '')}" + token = arguments.get("continue_from") + if isinstance(token, str) and token: + try: + payload = parse_token(token, name) + return name, str(payload.get("url", "unknown")) + except ValueError: + pass + source = arguments.get("url_or_path", arguments.get("url", "unknown")) + return name, str(source) + + def stream_chat_with_tools( messages: list[Message], model: str, @@ -242,6 +253,7 @@ def stream_chat_with_tools( messages = [_embed_images_in_content(m) if _is_image_message(m) else m for m in messages] cumulative_tool_chars = 0 + outstanding_paged_sources: set[tuple[str, str]] = set() MAX_EMPTY_RETRIES = 2 empty_retries = 0 @@ -339,6 +351,11 @@ def stream_chat_with_tools( messages.append(msg) if not tool_calls: + if outstanding_paged_sources: + yield { + "type": "warning", + "message": t("status.pagination_incomplete", lang=language), + } if is_empty: # Ran out of tool rounds while still empty; surface a warning. yield { @@ -358,14 +375,41 @@ def stream_chat_with_tools( for call in tool_calls: name = call.get("name", "") args = call.get("arguments", {}) + paged_source = _paged_source_key(name, args) yield {"type": "tool_start", "name": name, "arguments": _safe_tool_arguments(args)} t1 = time.time() - try: - result = reg.execute(name, args if isinstance(args, dict) else json.dumps(args)) - except Exception as exc: # noqa: BLE001 - logger.warning("tool %s execution failed in %.2fs: %s", name, time.time() - t1, exc) - result = f"Tool execution failed: {exc}" + remaining_context = MAX_TOOL_CONTEXT_CHARS - ( + cumulative_tool_chars + total_result_len + ) + is_builtin_paged = name in {"scrape_url", "pdf_to_text", "search_and_fetch"} + if is_builtin_paged and remaining_context < settings.max_tool_result_chars: + result = ( + "Tool context budget exhausted before this paged call was executed. " + "Unread content remains; do not claim that the source was read in full." + ) + outstanding_paged_sources.add(paged_source) + logger.info("paged tool skipped name=%s reason=context_budget", name) + else: + try: + result = reg.execute( + name, args if isinstance(args, dict) else json.dumps(args) + ) + except Exception as exc: # noqa: BLE001 + logger.warning( + "tool %s execution failed in %.2fs: %s", name, time.time() - t1, exc + ) + result = f"Tool execution failed: {exc}" + raw_is_paged = _PAGE_INFO_MARKER in result + if raw_is_paged: + _, _, next_token = parse_paged_result(result) + if next_token and isinstance(args, dict): + paged_source = _paged_source_key( + name, {**args, "continue_from": next_token} + ) + outstanding_paged_sources.add(paged_source) + elif is_builtin_paged: + outstanding_paged_sources.discard(paged_source) # Always truncate per-tool result immediately so it does not bloat # the context even before we reach the cumulative limit. result = _safe_tool_result(_trim_tool_result(result, settings.max_tool_result_chars)) @@ -379,21 +423,6 @@ def stream_chat_with_tools( tool_results.append((name, result)) total_result_len += len(result) - # If the accumulated tool results are still too large, truncate each - # one more aggressively to keep ollama's context processing fast. - if total_result_len > MAX_TOOL_CONTEXT_CHARS: - budget = MAX_TOOL_CONTEXT_CHARS - len(tool_results) * 200 # reserve footer per tool - max_per_tool = budget // max(1, len(tool_results)) - logger.warning( - "tool results too large (%d chars), truncating to %d chars per tool", - total_result_len, - max_per_tool, - ) - trimmed_results: list[tuple[str, str]] = [] - for name, result in tool_results: - trimmed_results.append((name, _trim_tool_result(result, max_per_tool))) - tool_results = trimmed_results - cumulative_tool_chars += sum(len(r) for _, r in tool_results) for name, result in tool_results: messages.append(Message(role="tool", tool_name=name, content=result)) @@ -410,10 +439,19 @@ def stream_chat_with_tools( "type": "status", "message": t("status.tool_limit_reached", lang=language), } + limit_instruction = get_prompt("tool_limit_instruction", lang=language) + if outstanding_paged_sources: + limit_instruction += "\n\n" + get_prompt( + "pagination_incomplete_instruction", lang=language + ) + yield { + "type": "warning", + "message": t("status.pagination_incomplete", lang=language), + } messages.append( Message( role="user", - content=get_prompt("tool_limit_instruction", lang=language), + content=limit_instruction, ) ) t0 = time.time() diff --git a/src/ollama_web/prompts/en.json b/src/ollama_web/prompts/en.json index 42e2e68..1d35f52 100644 --- a/src/ollama_web/prompts/en.json +++ b/src/ollama_web/prompts/en.json @@ -2,7 +2,9 @@ "tool_system": "You can use web search, scraping, and file fetch tools. Do not run more searches than necessary in a single response; review the results before conducting the next investigation. Tool results are summarized, so grasp the important points and use them effectively. If a tool result is truncated, you can fetch the continuation using the next_token at the end of the result. Pass the next_token value verbatim to the continue_from parameter. Omit continue_from when no continuation is needed.", "current_time": "Current time: {time}", "tool_limit_instruction": "The tool call limit has been reached. Based on the search results and conversation so far, please summarize the final answer the user requested.", + "pagination_incomplete_instruction": "A paged source still has unread content. Answer only from the retrieved portion and do not claim that the full source was reviewed.", "attachment_prefix": "The following is the content of files attached by the user:", + "attachment_paging": "A PDF body is an initial preview. When has_more is true, pass the attached-file uri and next_token to pdf_to_text for the needed continuation. Do not claim to have read the full file while unread content remains.", "fetch": { "candidate": "Candidate", "title": "Title", @@ -16,4 +18,4 @@ "empty_query": "Empty query.", "invalid_type": "Invalid file type." } -}
\ No newline at end of file +} diff --git a/src/ollama_web/prompts/ja.json b/src/ollama_web/prompts/ja.json index 88af710..6bcba63 100644 --- a/src/ollama_web/prompts/ja.json +++ b/src/ollama_web/prompts/ja.json @@ -2,7 +2,9 @@ "tool_system": "あなたはWeb検索・スクレイピング・ファイル取得ツールを利用できます。一度の回答で必要以上に多くの検索を同時に行わず、結果を確認してから次の調査を行ってください。各ツール結果は要約されて提供されるため、重要なポイントを把握して活用してください。ツール結果が途切れている場合、結果末尾の next_token を使って続きを取得できます。next_token の値は解釈せず、そのまま continue_from パラメータに渡してください。続きが必要ない場合は continue_from を省略してください。", "current_time": "現在時刻: {time}", "tool_limit_instruction": "ツール呼び出しの上限に達しました。これまでの検索結果と会話を元に、ユーザーが求めた回答を最終的にまとめてください。", + "pagination_incomplete_instruction": "ページングされた資料に未取得部分があります。取得済み範囲だけを根拠に回答し、全文を確認したとは述べないでください。", "attachment_prefix": "以下はユーザーが添付したファイルの内容です：", + "attachment_paging": "PDFは先頭プレビューです。has_more: true の場合、attached-file の uri と next_token を pdf_to_text に渡し、必要な範囲を続けて取得してください。未取得部分が残る場合は全文を読んだと主張しないでください。", "fetch": { "candidate": "候補", "title": "タイトル", @@ -16,4 +18,4 @@ "empty_query": "Empty query.", "invalid_type": "Invalid file type." } -}
\ No newline at end of file +} diff --git a/src/ollama_web/routes/chat.py b/src/ollama_web/routes/chat.py index e5b90d3..d94b567 100644 --- a/src/ollama_web/routes/chat.py +++ b/src/ollama_web/routes/chat.py @@ -16,7 +16,6 @@ from ollama._types import Image as OllamaImage from starlette.requests import Request from starlette.responses import StreamingResponse -from ..config import settings from ..i18n import t from ..mcp import collect_mcp_tools, make_mcp_executor from ..prompts import get_prompt @@ -24,6 +23,7 @@ from ..sessions import SessionStore, is_valid_id from ..settings_store import load_app_settings from ..tools.fetch import pop_fetched_files from ..tools.helper.image import resize_image +from ..tools.pdf import make_session_pdf_tool from ..tools.registry import ToolRegistry, default_registry logger = logging.getLogger("ollama_web.chat") @@ -65,7 +65,9 @@ def _parse_messages(raw: list[dict[str, Any]]) -> list[Message]: return out -async def _build_registry(session_id: str | None) -> ToolRegistry: +async def _build_registry( + session_id: str | None, store: SessionStore | None = None +) -> ToolRegistry: """Return a registry with built-in tools and any configured MCP tools.""" reg = default_registry() @@ -81,6 +83,20 @@ async def _build_registry(session_id: str | None) -> ToolRegistry: reg._tools["search_and_fetch"] = wrapped # noqa: SLF001 + if session_id is not None and store is not None: + + def load_pdf_attachment(file_id: str) -> bytes | None: + chat_file = store.get_file(session_id, file_id) + if chat_file is None or chat_file.is_image: + return None + if chat_file.mime.lower() != "application/pdf" and not chat_file.name.lower().endswith( + ".pdf" + ): + return None + return store.get_file_data(session_id, file_id) + + reg._tools["pdf_to_text"] = make_session_pdf_tool(load_pdf_attachment) # noqa: SLF001 + try: mcp_tools = await collect_mcp_tools() except Exception as exc: # noqa: BLE001 @@ -105,20 +121,57 @@ async def _build_registry(session_id: str | None) -> ToolRegistry: return reg -def _attachment_context(store: SessionStore, session_id: str, file_ids: list[str]) -> str: - parts: list[str] = [] +def _attachment_context( + store: SessionStore, + session_id: str, + file_ids: list[str], + *, + total_limit: int, +) -> str: + if total_limit <= 0: + return "" + chat_files = [] for file_id in file_ids: chat_file = store.get_file(session_id, file_id) - if chat_file is None or chat_file.is_image: - continue - text = _limit_text(chat_file.text or "", settings.max_attachment_text_chars) - parts.append( + if chat_file is not None and not chat_file.is_image: + chat_files.append(chat_file) + if not chat_files: + return "" + + parts: list[str] = [] + prefix = get_prompt("attachment_prefix") + paging_hint = get_prompt( + "attachment_paging", + fallback=( + "PDF previews may have unread content. Use pdf_to_text with the attachment URI " + "and next_token until has_more is absent." + ), + ) + used = len(prefix) + len(paging_hint) + 4 + per_file = max(256, min(4000, (total_limit - used) // max(1, len(chat_files)))) + pdf_tool = make_session_pdf_tool( + lambda file_id: store.get_file_data(session_id, file_id) + if store.get_file(session_id, file_id) is not None + else None + ) + + for chat_file in chat_files: + source_uri = f"attachment:{chat_file.id}" + if chat_file.mime.lower() == "application/pdf" or chat_file.name.lower().endswith(".pdf"): + text = pdf_tool(source_uri, max_chars=per_file) + else: + text = _limit_text(chat_file.text or "", per_file) + block = ( f'<attached-file name="{html.escape(chat_file.name)}" id="{chat_file.id}" ' - f'source="{chat_file.source}">\n{text}\n</attached-file>' + f'source="{chat_file.source}" uri="{source_uri}">\n{text}\n</attached-file>' ) + if used + len(block) + 2 > total_limit: + break + parts.append(block) + used += len(block) + 2 if not parts: return "" - return "\n\n".join([get_prompt("attachment_prefix"), *parts]) + return "\n\n".join([prefix, paging_hint, *parts]) def _collect_images(store: SessionStore, session_id: str, file_ids: list[str]) -> list[bytes]: @@ -217,7 +270,16 @@ async def _chat_event_stream( attachment_ctx = "" if session_id and file_ids: - attachment_ctx = _attachment_context(store, session_id, file_ids) + attachment_budget = max( + 0, + request.app.state.settings.max_message_chars - len(user_content) - 256, + ) + attachment_ctx = _attachment_context( + store, + session_id, + file_ids, + total_limit=attachment_budget, + ) if attachment_ctx and user_content: # Inject attachment context into the latest user message. @@ -247,7 +309,7 @@ async def _chat_event_stream( from ollama_web import llm # local import to avoid circular at module load - registry = await _build_registry(session_id) + registry = await _build_registry(session_id, store) total_image_bytes = sum( len(img.value) for img in (getattr(messages[-1], "images", None) or []) diff --git a/src/ollama_web/routes/sessions.py b/src/ollama_web/routes/sessions.py index e1e669c..73f79d6 100644 --- a/src/ollama_web/routes/sessions.py +++ b/src/ollama_web/routes/sessions.py @@ -11,7 +11,7 @@ from starlette.requests import Request from starlette.responses import JSONResponse from ..sessions import _IMAGE_MIMES, _IMAGE_SUFFIXES, SessionStore, is_valid_id -from ..tools.helper.pdf import extract_pdf_text +from ..tools.helper.pdf import extract_pdf_page logger = logging.getLogger("ollama_web.sessions") @@ -127,8 +127,17 @@ async def upload_file(request: Request) -> JSONResponse: # Images are passed to ollama via Message.images (base64), not as text. text = "" elif content_type == "application/pdf" or name.lower().endswith(".pdf"): - # extract_pdf_text now returns (text, has_more, next_page_offset). - text = extract_pdf_text(data)[0] + # Store only a bounded preview. The chat-scoped PDF tool reads the + # original bytes on demand and exposes the lossless continuation. + preview_limit = min( + request.app.state.settings.max_attachment_text_chars, + request.app.state.settings.max_tool_result_chars, + ) + text = extract_pdf_page( + data, + max_chars=preview_limit, + max_size_mb=request.app.state.settings.max_upload_mb, + ).text else: # Try decode as text for text-like uploads. try: @@ -172,4 +181,4 @@ async def clear_messages(request: Request) -> JSONResponse: session = store.clear_messages(session_id) if session is None: return JSONResponse({"error": "Session not found"}, status_code=404) - return JSONResponse({"cleared": True, "session": session})
\ No newline at end of file + return JSONResponse({"cleared": True, "session": session}) diff --git a/src/ollama_web/tools/__init__.py b/src/ollama_web/tools/__init__.py index f2a6b0b..05d4a04 100644 --- a/src/ollama_web/tools/__init__.py +++ b/src/ollama_web/tools/__init__.py @@ -2,7 +2,7 @@ from .fetch import search_and_fetch from .helper.image import resize_image -from .helper.pdf import extract_pdf_text +from .helper.pdf import extract_pdf_page, extract_pdf_text from .pdf import pdf_to_text from .registry import ToolRegistry, all_tool_callables, default_registry from .scrape import scrape_url @@ -12,6 +12,7 @@ __all__ = [ "search_and_fetch", "resize_image", "extract_pdf_text", + "extract_pdf_page", "pdf_to_text", "ToolRegistry", "all_tool_callables", diff --git a/src/ollama_web/tools/fetch.py b/src/ollama_web/tools/fetch.py index 65ec511..6ab17bb 100644 --- a/src/ollama_web/tools/fetch.py +++ b/src/ollama_web/tools/fetch.py @@ -2,6 +2,7 @@ from __future__ import annotations +import logging import mimetypes import re from pathlib import Path @@ -9,17 +10,24 @@ from typing import Any from ..config import settings from ..prompts import get_prompt, get_prompts -from .helper.pagination import format_paged_result, make_token, parse_token -from .helper.pdf import extract_pdf_text +from .helper.pagination import ( + PageResult, + PageSizeError, + format_page_to_limit, + make_token, + parse_token, + text_page, +) +from .helper.pdf import extract_pdf_page, legacy_pdf_offset_to_cursor from .helper.safe_http import fetch_public_bytes -from .scrape import scrape_url +from .scrape import extract_readable_text from .search import web_search _FetchedFile = dict[str, Any] # Per-session temporary storage for files fetched by search_and_fetch. # The chat route drains this after the tool call finishes. -_fetched_files: dict[str, list[_FetchedFile]] = {} +_fetched_files: dict[str, dict[str, _FetchedFile]] = {} _FILETYPE_SUFFIXES: dict[str, str] = { @@ -37,6 +45,7 @@ _FILETYPE_SUFFIXES: dict[str, str] = { } _TOOL_NAME = "search_and_fetch" +logger = logging.getLogger("ollama_web.tools.fetch") def _slim_matched(matched: list[dict[str, str]]) -> list[dict[str, str]]: @@ -53,12 +62,6 @@ def _slim_matched(matched: list[dict[str, str]]) -> list[dict[str, str]]: return [{"title": m.get("title", ""), "href": m.get("href", "")} for m in matched] -def set_session_context(session_id: str | None) -> None: - """Set the active session id so fetched files can be persisted later.""" - # No-op marker kept for future hooks; storage is keyed by session_id directly. - pass - - def store_fetched_file( session_id: str, name: str, @@ -69,7 +72,8 @@ def store_fetched_file( ) -> None: """Store a fetched file so the chat route can persist it to the session.""" if session_id: - _fetched_files.setdefault(session_id, []).append( + _fetched_files.setdefault(session_id, {}).setdefault( + url, { "name": name, "data": data, @@ -82,7 +86,7 @@ def store_fetched_file( def pop_fetched_files(session_id: str) -> list[_FetchedFile]: """Return and clear any files fetched for the given session.""" - return _fetched_files.pop(session_id, []) + return list(_fetched_files.pop(session_id, {}).values()) def _normalize_file_type(file_type: str) -> str: @@ -118,44 +122,10 @@ def _fetch_bytes(url: str) -> bytes: 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}" - - if file_type == "pdf": - text, _, _ = extract_pdf_text(data, max_chars=max_chars, page_offset=0) - return text - - # Decode text-like formats. - try: - text = data.decode("utf-8") - except UnicodeDecodeError: - try: - text = data.decode("cp932") - except UnicodeDecodeError: - text = data.decode("utf-8", errors="replace") - - if file_type in {"html", "htm"}: - text = scrape_url(url, max_chars=max_chars) - - if len(text) > max_chars: - text = text[:max_chars].rsplit(" ", 1)[0] + "\n…[truncated]" - return text - - def _full_text_for_file(url: str, data: bytes, file_type: str) -> str: - """Return the *full* extracted text for a file (no truncation). - - Used by the pagination path so we can slice arbitrary offsets. For PDFs we - extract all pages; for text-like formats we decode the full payload. - """ + """Decode a complete text-like file for character-offset pagination.""" if file_type == "pdf": - # Extract with a very large char budget to get the full text. - text, _, _ = extract_pdf_text(data, max_chars=10_000_000, page_offset=0) - return text + raise ValueError("PDF files must use page-and-character cursor pagination.") try: text = data.decode("utf-8") @@ -166,10 +136,7 @@ def _full_text_for_file(url: str, data: bytes, file_type: str) -> str: text = data.decode("utf-8", errors="replace") if file_type in {"html", "htm"}: - # Reuse scrape_url's cleaning but request a large budget to get the - # full cleaned text. scrape_url itself now supports pagination, but - # here we want the full body for manual slicing. - text = scrape_url(url, max_chars=10_000_000) + text = extract_readable_text(data) return text @@ -223,19 +190,43 @@ def search_and_fetch( payload = parse_token(continue_from, _TOOL_NAME) except ValueError as exc: return f"Invalid continuation token: {exc}" - resumed_matched = payload.get("matched", []) - if not isinstance(resumed_matched, list) or not resumed_matched: + raw_matched = payload.get("matched", []) + if not isinstance(raw_matched, list) or not raw_matched: return ( get_prompt("fetch.no_results", fallback="") .replace("{file_type}", ft) .replace("{query}", query) ) - file_idx = int(payload.get("file_idx", 0)) - file_offset = int(payload.get("offset", 0)) + resumed_matched = [ + {"title": str(item.get("title", "")), "href": str(item.get("href", ""))} + for item in raw_matched + if isinstance(item, dict) + ] + if payload.get("v") == 2 and payload.get("file_type") != ft: + return "Invalid continuation token: file type does not match." + try: + file_idx = int(payload.get("file_idx", 0)) + file_offset = int(payload.get("file_offset", payload.get("offset", 0))) + pdf_page_index = int(payload.get("pdf_page_index", 0)) + pdf_page_char_offset = int(payload.get("pdf_page_char_offset", 0)) + except (TypeError, ValueError) as exc: + return f"Invalid continuation token: cursor is not an integer: {exc}" + legacy_pdf_offset = ( + file_offset + if payload.get("v") == 1 and "pdf_page_index" not in payload + else None + ) + if min(file_idx, file_offset, pdf_page_index, pdf_page_char_offset) < 0: + return "Invalid continuation token: cursor values must not be negative." + if not resumed_matched or file_idx >= len(resumed_matched): + return "Invalid continuation token: file index is out of range." return _render_file_page( matched=resumed_matched, file_idx=file_idx, file_offset=file_offset, + pdf_page_index=pdf_page_index, + pdf_page_char_offset=pdf_page_char_offset, + legacy_pdf_offset=legacy_pdf_offset, ft=ft, per_file_chars=per_file_chars, session_id=session_id, @@ -265,6 +256,9 @@ def search_and_fetch( matched=matched, file_idx=0, file_offset=0, + pdf_page_index=0, + pdf_page_char_offset=0, + legacy_pdf_offset=None, ft=ft, per_file_chars=per_file_chars, session_id=session_id, @@ -277,6 +271,9 @@ def _render_file_page( matched: list[dict[str, str]], file_idx: int, file_offset: int, + pdf_page_index: int, + pdf_page_char_offset: int, + legacy_pdf_offset: int | None, ft: str, per_file_chars: int, session_id: str | None, @@ -288,8 +285,7 @@ def _render_file_page( re-run the web search. """ if file_idx >= len(matched): - # All candidates exhausted. - return format_paged_result("", has_more=False, next_token=None) + return "" item = matched[file_idx] url = item.get("href", "") @@ -305,59 +301,157 @@ def _render_file_page( f"=== {p['candidate']} {file_idx + 1} ===\n" f"URL: {url}\n{p['title']}: {title}\n{p['status']}: {p['fetch_failed']} ({exc})\n" ) - next_token = make_token( - _TOOL_NAME, - url, - 0, - matched=_slim_matched(matched), - file_idx=file_idx + 1, - file_offset=0, + has_more = file_idx + 1 < len(matched) + cursor = { + "file_idx": file_idx + 1, + "file_offset": 0, + "pdf_page_index": 0, + "pdf_page_char_offset": 0, + } + page = PageResult(body, cursor, cursor, has_more) + return format_page_to_limit( + lambda _budget: page, + lambda next_cursor: _make_fetch_token(matched, next_cursor, ft), + requested_chars=per_file_chars, + result_limit=settings.max_tool_result_chars, ) - return format_paged_result(body, has_more=True, next_token=next_token) - - full_text = _full_text_for_file(url, data, ft) - page_text = full_text[file_offset : file_offset + per_file_chars] - has_more_in_file = len(full_text) > file_offset + 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: - store_fetched_file(session_id, name, data, mime, page_text, url) - - body = ( + prefix = ( f"=== {p['candidate']} {file_idx + 1} ===\n" f"URL: {url}\n" f"{p['title']}: {title}\n" f"{p['type']}: {mime}\n" f"{p['size']}: {len(data)} bytes\n" - f"{p['chars']}: {len(page_text)}\n" - f"{p['separator']}\n{page_text}\n" + f"{p['separator']}\n" ) - if has_more_in_file: - # Continue within the same file. - next_token = make_token( - _TOOL_NAME, - url, - file_offset + per_file_chars, - matched=_slim_matched(matched), - file_idx=file_idx, - file_offset=file_offset + per_file_chars, - ) - return format_paged_result(body, has_more=True, next_token=next_token) + if ft == "pdf": + if legacy_pdf_offset is not None: + try: + pdf_page_index, pdf_page_char_offset = legacy_pdf_offset_to_cursor( + data, legacy_pdf_offset + ) + except Exception as exc: # noqa: BLE001 + return f"Invalid continuation token: {exc}" + preview = extract_pdf_page(data, max_chars=min(per_file_chars, 1000)).text + + def page_factory(body_budget: int) -> PageResult: + source_page = extract_pdf_page( + data, + max_chars=body_budget, + page_index=pdf_page_index, + page_char_offset=pdf_page_char_offset, + ) + if source_page.has_more: + next_cursor = { + "file_idx": file_idx, + "file_offset": 0, + "pdf_page_index": source_page.next_cursor["page_index"], + "pdf_page_char_offset": source_page.next_cursor["page_char_offset"], + } + has_more = True + else: + next_cursor = { + "file_idx": file_idx + 1, + "file_offset": 0, + "pdf_page_index": 0, + "pdf_page_char_offset": 0, + } + has_more = file_idx + 1 < len(matched) + return PageResult( + source_page.text, + { + "file_idx": file_idx, + "file_offset": 0, + "pdf_page_index": pdf_page_index, + "pdf_page_char_offset": pdf_page_char_offset, + }, + next_cursor, + has_more, + source_page.warnings, + ) - # Move to the next candidate file. - if file_idx + 1 < len(matched): - next_token = make_token( - _TOOL_NAME, - url, - 0, - matched=_slim_matched(matched), - file_idx=file_idx + 1, - file_offset=0, + else: + full_text = _full_text_for_file(url, data, ft) + preview = full_text[: min(per_file_chars, 1000)] + + def page_factory(body_budget: int) -> PageResult: + source_page = text_page(full_text, file_offset, body_budget) + if source_page.has_more: + next_cursor = { + "file_idx": file_idx, + "file_offset": source_page.next_cursor["offset"], + "pdf_page_index": 0, + "pdf_page_char_offset": 0, + } + has_more = True + else: + next_cursor = { + "file_idx": file_idx + 1, + "file_offset": 0, + "pdf_page_index": 0, + "pdf_page_char_offset": 0, + } + has_more = file_idx + 1 < len(matched) + return PageResult( + source_page.text, + { + "file_idx": file_idx, + "file_offset": file_offset, + "pdf_page_index": 0, + "pdf_page_char_offset": 0, + }, + next_cursor, + has_more, + ) + + if session_id: + store_fetched_file(session_id, name, data, mime, preview, url) + + try: + result = format_page_to_limit( + page_factory, + lambda cursor: _make_fetch_token(matched, cursor, ft), + requested_chars=per_file_chars, + result_limit=settings.max_tool_result_chars, + prefix=prefix, + suffix="\n", + ) + logger.info( + ( + "fetch page type=%s file_idx=%d file_offset=%d pdf_page_index=%d " + "pdf_page_char_offset=%d result_chars=%d has_more=%s" + ), + ft, + file_idx, + file_offset, + pdf_page_index, + pdf_page_char_offset, + len(result), + "---page-info---" in result, ) - return format_paged_result(body, has_more=True, next_token=next_token) + return result + except (PageSizeError, ValueError) as exc: + return f"File pagination failed: {exc}" + - return format_paged_result(body, has_more=False, next_token=None)
\ No newline at end of file +def _make_fetch_token( + matched: list[dict[str, str]], cursor: dict[str, int], file_type: str +) -> str: + file_idx = cursor["file_idx"] + url = matched[file_idx].get("href", "") if file_idx < len(matched) else "completed" + return make_token( + _TOOL_NAME, + url, + cursor["file_offset"], + matched=_slim_matched(matched), + file_idx=file_idx, + file_offset=cursor["file_offset"], + pdf_page_index=cursor["pdf_page_index"], + pdf_page_char_offset=cursor["pdf_page_char_offset"], + file_type=file_type, + ) diff --git a/src/ollama_web/tools/helper/__init__.py b/src/ollama_web/tools/helper/__init__.py index f201101..198a300 100644 --- a/src/ollama_web/tools/helper/__init__.py +++ b/src/ollama_web/tools/helper/__init__.py @@ -1,12 +1,13 @@ """Internal helper functions for tools.""" from .image import resize_image -from .pdf import extract_pdf_text +from .pdf import extract_pdf_page, extract_pdf_text from .safe_http import UnsafeURL, fetch_public_bytes, validate_public_http_url __all__ = [ "UnsafeURL", "extract_pdf_text", + "extract_pdf_page", "fetch_public_bytes", "resize_image", "validate_public_http_url", diff --git a/src/ollama_web/tools/helper/pagination.py b/src/ollama_web/tools/helper/pagination.py index 748338b..9d17823 100644 --- a/src/ollama_web/tools/helper/pagination.py +++ b/src/ollama_web/tools/helper/pagination.py @@ -1,32 +1,17 @@ -"""Pagination helpers for paged tool results. - -This module provides an opaque continuation-token mechanism so that LLM-driven -tools can return large content in pages without requiring the model to remember -numeric offsets. The token encodes the resumption state (URL + offset) as -base64url(JSON) and is meant to be passed back verbatim by the model via the -``continue_from`` parameter. - -Design rationale: -- The token is opaque: the LLM should copy the ``next_token`` value from a - previous result and pass it to ``continue_from`` without interpreting it. -- The state is encoded in the token itself, so the server stays stateless and - does not need to cache previous results. -- A ``tool`` field guards against a token issued by one tool being replayed - against a different tool. -""" +"""Lossless pagination helpers for tool results.""" from __future__ import annotations import base64 import json +from collections.abc import Callable +from dataclasses import dataclass, field from typing import Any -# Token format version. Bump when the payload structure changes in a -# backwards-incompatible way; ``parse_token`` rejects unknown versions. -_TOKEN_VERSION = 1 +_TOKEN_VERSION = 2 +_SUPPORTED_TOKEN_VERSIONS = {1, 2} +MAX_TOKEN_CHARS = 8192 -# Marker that separates the page body from the page-info footer in tool -# results. Kept in sync with ``llm._trim_tool_result``. PAGE_INFO_MARKER = "\n---page-info---" @@ -34,29 +19,32 @@ class InvalidTokenError(ValueError): """Raised when a continuation token is malformed or invalid.""" +class PageSizeError(ValueError): + """Raised when page metadata alone cannot fit in the result budget.""" + + +@dataclass(frozen=True) +class PageResult: + """One lossless source slice and the cursor immediately after that slice.""" + + text: str + start_cursor: dict[str, int] + next_cursor: dict[str, int] + has_more: bool + warnings: tuple[str, ...] = field(default_factory=tuple) + + def _b64url_encode(data: bytes) -> str: return base64.urlsafe_b64encode(data).decode("ascii").rstrip("=") def _b64url_decode(token: str) -> bytes: - # Re-add any stripped padding before decoding. padding = "=" * (-len(token) % 4) return base64.urlsafe_b64decode(token + padding) def make_token(tool: str, url: str, offset: int, **extra: Any) -> str: - """Create an opaque continuation token. - - Args: - tool: Name of the issuing tool (e.g. ``"scrape_url"``). - url: The URL the content was fetched from. - offset: The resumption offset (character or page index, tool-specific). - **extra: Additional tool-specific fields to embed in the token payload - (e.g. ``matched`` candidate list for ``search_and_fetch``). - - Returns: - A base64url-encoded JSON string. - """ + """Create a version-2 opaque continuation token.""" payload: dict[str, Any] = { "v": _TOKEN_VERSION, "tool": tool, @@ -64,36 +52,23 @@ def make_token(tool: str, url: str, offset: int, **extra: Any) -> str: "offset": int(offset), } payload.update(extra) - # ensure_ascii=True keeps the token pure ASCII so truncation of the - # token by the model or tool-result trimming cannot split a multibyte - # UTF-8 sequence (which would raise UnicodeDecodeError on decode). raw = json.dumps(payload, ensure_ascii=True, separators=(",", ":")).encode("ascii") - return _b64url_encode(raw) + token = _b64url_encode(raw) + if len(token) > MAX_TOKEN_CHARS: + raise PageSizeError(f"Continuation token exceeds {MAX_TOKEN_CHARS} characters.") + return token def parse_token(token: str | None, expected_tool: str) -> dict[str, Any]: - """Validate and decode a continuation token. - - Args: - token: The opaque token string (typically from ``continue_from``). - ``None`` or an empty string raises :class:`InvalidTokenError`. - expected_tool: The tool name that should appear in the token's ``tool`` - field. Mismatches raise :class:`InvalidTokenError`. - - Returns: - The decoded payload dict (includes ``url``, ``offset`` and any extra - fields the issuing tool embedded). - - Raises: - InvalidTokenError: If the token is empty, not valid base64, not valid - JSON, missing required fields, has an unsupported version, or names a - different tool than ``expected_tool``. - """ + """Validate and decode a version-1 or version-2 continuation token.""" if not token or not token.strip(): raise InvalidTokenError("Empty continuation token.") + token = token.strip() + if len(token) > MAX_TOKEN_CHARS: + raise InvalidTokenError(f"Continuation token exceeds {MAX_TOKEN_CHARS} characters.") try: - raw = _b64url_decode(token.strip()) + raw = _b64url_decode(token) except (ValueError, base64.binascii.Error) as exc: # type: ignore[attr-defined] raise InvalidTokenError(f"Token is not valid base64url: {exc}") from exc @@ -104,77 +79,86 @@ def parse_token(token: str | None, expected_tool: str) -> dict[str, Any]: if not isinstance(payload, dict): raise InvalidTokenError("Token payload is not a JSON object.") - - version = payload.get("v") - if version != _TOKEN_VERSION: - raise InvalidTokenError(f"Unsupported token version: {version!r}") - - tool = payload.get("tool") - if tool != expected_tool: - raise InvalidTokenError(f"Token was issued for tool {tool!r}, expected {expected_tool!r}.") - + if payload.get("v") not in _SUPPORTED_TOKEN_VERSIONS: + raise InvalidTokenError(f"Unsupported token version: {payload.get('v')!r}") + if payload.get("tool") != expected_tool: + raise InvalidTokenError( + f"Token was issued for tool {payload.get('tool')!r}, expected {expected_tool!r}." + ) if "url" not in payload or "offset" not in payload: raise InvalidTokenError("Token payload missing 'url' or 'offset'.") - + if not isinstance(payload["url"], str) or not payload["url"]: + raise InvalidTokenError("Token 'url' is not a non-empty string.") try: payload["offset"] = int(payload["offset"]) except (TypeError, ValueError) as exc: raise InvalidTokenError(f"Token 'offset' is not an integer: {exc}") from exc - + if payload["offset"] < 0: + raise InvalidTokenError("Token 'offset' must not be negative.") return payload def format_paged_result(text: str, *, has_more: bool, next_token: str | None) -> str: - """Wrap page body with a page-info footer. - - The footer is appended after :data:`PAGE_INFO_MARKER` so that - ``llm._trim_tool_result`` can protect it from truncation. When - ``has_more`` is False the ``next_token`` line is omitted. - - Args: - text: The page body text. - has_more: Whether more content is available. - next_token: The continuation token to include when ``has_more`` is True. - - Returns: - The formatted result string with the page-info footer. - """ + """Append the protected continuation footer when unread content remains.""" if not has_more: - # No footer needed when there is nothing more to fetch. Keeping the - # body unchanged preserves backward compatibility for short results. return text - if not next_token: - # Defensive: has_more without a token is a programming error; fall back - # to no footer to avoid confusing the model. - return text - - footer = f"{PAGE_INFO_MARKER}\nhas_more: true\nnext_token: {next_token}" - return f"{text}{footer}" + raise PageSizeError("A continuation token is required when has_more is true.") + return f"{text}{PAGE_INFO_MARKER}\nhas_more: true\nnext_token: {next_token}" + + +def format_page_to_limit( + page_factory: Callable[[int], PageResult], + token_factory: Callable[[dict[str, int]], str], + *, + requested_chars: int, + result_limit: int, + prefix: str = "", + suffix: str = "", +) -> str: + """Render the largest lossless page whose complete envelope fits the limit.""" + body_budget = max(1, min(requested_chars, result_limit)) + for _ in range(8): + page = page_factory(body_budget) + warning_text = "".join(f"[warning] {warning}\n" for warning in page.warnings) + body = f"{prefix}{warning_text}{page.text}{suffix}" + token = token_factory(page.next_cursor) if page.has_more else None + rendered = format_paged_result(body, has_more=page.has_more, next_token=token) + if len(rendered) <= result_limit: + return rendered + overflow = len(rendered) - result_limit + new_budget = body_budget - overflow - 16 + if new_budget >= body_budget: + new_budget = body_budget - 1 + if new_budget <= 0: + break + body_budget = new_budget + raise PageSizeError("Paged result metadata exceeds the configured result limit.") + + +def text_page(text: str, offset: int, max_chars: int) -> PageResult: + """Return a lossless character-offset page from ``text``.""" + if offset < 0 or offset > len(text): + raise ValueError("Text pagination offset is out of range.") + end = min(offset + max_chars, len(text)) + return PageResult( + text=text[offset:end], + start_cursor={"offset": offset}, + next_cursor={"offset": end}, + has_more=end < len(text), + ) def parse_paged_result(result: str) -> tuple[str, bool, str | None]: - """Split a paged result into (body, has_more, next_token). - - This is the inverse of :func:`format_paged_result` and is mainly useful in - tests and for chaining tool results. - - Args: - result: The formatted result string from :func:`format_paged_result`. - - Returns: - A tuple of (body_text, has_more, next_token_or_None). - """ + """Split a formatted result into body, has-more flag, and token.""" if PAGE_INFO_MARKER not in result: return result, False, None - body, _, footer = result.partition(PAGE_INFO_MARKER) has_more = "has_more: true" in footer next_token: str | None = None if has_more: for line in footer.splitlines(): - line = line.strip() - if line.startswith("next_token:"): + if line.strip().startswith("next_token:"): next_token = line.split(":", 1)[1].strip() break return body, has_more, next_token diff --git a/src/ollama_web/tools/helper/pdf.py b/src/ollama_web/tools/helper/pdf.py index 7b9fdac..243a837 100644 --- a/src/ollama_web/tools/helper/pdf.py +++ b/src/ollama_web/tools/helper/pdf.py @@ -1,4 +1,4 @@ -"""PDF text extraction helpers.""" +"""Lossless PDF text extraction helpers.""" from __future__ import annotations @@ -7,78 +7,117 @@ import io from pypdf import PdfReader from ...config import settings +from .pagination import PageResult -def extract_pdf_text( +def extract_pdf_page( data: bytes, - max_chars: int | None = None, - page_offset: int = 0, -) -> tuple[str, bool, int]: - """Extract text from PDF bytes, starting at ``page_offset``. - - Args: - data: Raw PDF bytes. - max_chars: Maximum characters to return. Defaults to - ``settings.fetch_max_chars``. - page_offset: Zero-based index of the first page to extract. Pages before - this index are skipped. Used by the pagination layer to resume. - - Returns: - A tuple ``(text, has_more, next_page_offset)`` where: - - ``text`` is the extracted text for this page (truncated to - ``max_chars``). - - ``has_more`` is True when there are remaining pages after the - extracted slice. - - ``next_page_offset`` is the page index to resume from on the next - call (only meaningful when ``has_more`` is True). - """ + *, + max_chars: int, + page_index: int = 0, + page_char_offset: int = 0, + max_size_mb: int | None = None, +) -> PageResult: + """Extract a PDF slice using a page-and-character continuation cursor.""" + start = {"page_index": page_index, "page_char_offset": page_char_offset} if not data: - return "Empty PDF data.", False, 0 - - 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.", False, 0 + return PageResult("Empty PDF data.", start, start, False) + size_limit_mb = settings.max_fetch_mb if max_size_mb is None else max_size_mb + if len(data) > size_limit_mb * 1024 * 1024: + return PageResult("PDF exceeds the configured size limit.", start, start, False) + if page_index < 0 or page_char_offset < 0: + return PageResult("Invalid PDF continuation cursor.", start, start, False) try: reader = PdfReader(stream=io.BytesIO(data)) except Exception as exc: # noqa: BLE001 - return f"Failed to parse PDF: {exc}", False, 0 + return PageResult(f"Failed to parse PDF: {exc}", start, start, False) total_pages = len(reader.pages) - if page_offset < 0: - page_offset = 0 - if page_offset >= total_pages: - return "No text could be extracted from the PDF.", False, 0 - - # Cap the number of pages processed per call to max_pdf_pages. - end = min(page_offset + settings.max_pdf_pages, total_pages) + if page_index > total_pages or (page_index == total_pages and page_char_offset): + return PageResult("Invalid PDF continuation cursor.", start, start, False) + if page_index == total_pages: + return PageResult("", start, start, False) parts: list[str] = [] - collected = 0 - last_extracted_page = page_offset - for page_index in range(page_offset, end): + warnings: list[str] = [] + remaining = max(1, max_chars) + current_page = page_index + current_char = page_char_offset + pages_processed = 0 + + while current_page < total_pages and pages_processed < settings.max_pdf_pages: try: - text = reader.pages[page_index].extract_text() - except Exception: # noqa: BLE001 + page_text = reader.pages[current_page].extract_text() or "" + except Exception as exc: # noqa: BLE001 + detail = str(exc).replace("\n", " ")[:120] + warnings.append(f"PDF page {current_page + 1} could not be extracted: {detail}") + if current_page + 1 < total_pages: + parts.append("\n\n") + current_page += 1 + current_char = 0 + pages_processed += 1 continue - if text: - parts.append(text) - collected += len(text) - last_extracted_page = page_index - # Stop as soon as we reach the per-call character budget so the page - # boundary for the next token is accurate. - if collected >= limit: + + if current_char > len(page_text): + return PageResult("Invalid PDF continuation cursor.", start, start, False) + + available = page_text[current_char:] + take = min(remaining, len(available)) + if take: + parts.append(available[:take]) + current_char += take + remaining -= take + + if current_char < len(page_text): + break + + if current_page + 1 < total_pages: + parts.append("\n\n") + current_page += 1 + current_char = 0 + pages_processed += 1 + if remaining <= 0: break - full = "\n\n".join(parts) - if len(full) > limit: - full = full[:limit].rsplit(" ", 1)[0] + "\n…[truncated]" + next_cursor = {"page_index": current_page, "page_char_offset": current_char} + has_more = current_page < total_pages + text = "".join(parts) + if not text and not warnings and not has_more: + text = "No text could be extracted from the PDF." + return PageResult(text, start, next_cursor, has_more, tuple(warnings)) + + +def legacy_pdf_offset_to_cursor(data: bytes, flat_offset: int) -> tuple[int, int]: + """Translate a v1 ``\n\n``-joined character offset to the v2 PDF cursor.""" + if flat_offset < 0: + raise ValueError("Legacy PDF offset must not be negative.") + reader = PdfReader(stream=io.BytesIO(data)) + remaining = flat_offset + for page_index, page in enumerate(reader.pages): + try: + page_text = page.extract_text() or "" + except Exception: # noqa: BLE001 + page_text = "" + if remaining <= len(page_text): + return page_index, remaining + remaining -= len(page_text) + if page_index + 1 < len(reader.pages): + if remaining <= 2: + return page_index + 1, 0 + remaining -= 2 + if remaining == 0: + return len(reader.pages), 0 + raise ValueError("Legacy PDF offset is out of range.") - # has_more when there are pages beyond the last one we extracted. - has_more = (last_extracted_page + 1) < total_pages - next_page_offset = last_extracted_page + 1 if has_more else 0 - if not full: - return "No text could be extracted from the PDF.", has_more, next_page_offset - return full, has_more, next_page_offset +def extract_pdf_text( + data: bytes, + max_chars: int | None = None, + page_offset: int = 0, +) -> tuple[str, bool, int]: + """Backward-compatible wrapper for page-only internal callers.""" + limit = max_chars if max_chars is not None and max_chars > 0 else settings.fetch_max_chars + page = extract_pdf_page(data, max_chars=limit, page_index=page_offset) + text = "".join(f"[warning] {warning}\n" for warning in page.warnings) + page.text + return text, page.has_more, page.next_cursor["page_index"] if page.has_more else 0 diff --git a/src/ollama_web/tools/pdf.py b/src/ollama_web/tools/pdf.py index 9fdb497..4faf8ac 100644 --- a/src/ollama_web/tools/pdf.py +++ b/src/ollama_web/tools/pdf.py @@ -2,67 +2,148 @@ from __future__ import annotations +import logging +from collections.abc import Callable +from typing import Any + from ..config import settings -from .helper.pagination import format_paged_result, make_token, parse_token -from .helper.pdf import extract_pdf_text +from .helper.pagination import ( + PageResult, + PageSizeError, + format_page_to_limit, + make_token, + parse_token, +) +from .helper.pdf import extract_pdf_page from .helper.safe_http import UnsafeURL, fetch_public_bytes _TOOL_NAME = "pdf_to_text" +_ATTACHMENT_PREFIX = "attachment:" +logger = logging.getLogger("ollama_web.tools.pdf") +AttachmentLoader = Callable[[str], bytes | None] -def pdf_to_text( - url_or_path: str, max_chars: int | None = None, continue_from: str | None = None -) -> str: - """Extract text from a local PDF path or remote URL. - The text is returned one page at a time. When more pages remain, a - ``next_token`` is appended that can be passed back via ``continue_from`` - to fetch the next batch of pages. +def _cursor_value(payload: dict[str, Any], key: str, fallback: int = 0) -> int: + raw = payload.get(key, fallback) + if isinstance(raw, bool) or not isinstance(raw, (int, str)): + raise ValueError(f"Token {key!r} is not an integer.") + try: + value = int(raw) + except (TypeError, ValueError) as exc: + raise ValueError(f"Token {key!r} is not an integer.") from exc + if value < 0: + raise ValueError(f"Token {key!r} must not be negative.") + return value - Args: - url_or_path: File path or URL pointing to a PDF. - max_chars: Maximum characters to return per page. - continue_from: Opaque continuation token returned in a previous call's - ``next_token``. When provided, extraction resumes from the encoded - page offset. - - Returns: - Extracted plain text for this page, followed by a ``---page-info---`` - footer when more pages are available, or an error message. - """ + +def _pdf_to_text( + url_or_path: str, + max_chars: int | None, + continue_from: str | None, + attachment_loader: AttachmentLoader | None, +) -> str: if not url_or_path or not url_or_path.strip(): return "Invalid PDF source." src = url_or_path.strip() - - # Resolve the resumption page offset from the continuation token, if any. - page_offset = 0 + page_index = 0 + page_char_offset = 0 if continue_from: try: payload = parse_token(continue_from, _TOOL_NAME) + src = str(payload.get("url") or src) + # Version 1 used offset as the PDF page index. + page_index = _cursor_value(payload, "page_index", int(payload["offset"])) + page_char_offset = _cursor_value(payload, "page_char_offset") except ValueError as exc: return f"Invalid continuation token: {exc}" - src = str(payload.get("url") or src) - page_offset = int(payload.get("offset", 0)) - if page_offset < 0: - page_offset = 0 - data: bytes | None = None - if src.startswith("http://") or src.startswith("https://"): + if src.startswith(_ATTACHMENT_PREFIX): + file_id = src.removeprefix(_ATTACHMENT_PREFIX) + data = attachment_loader(file_id) if attachment_loader is not None else None + if data is None: + return "Attachment PDF is unavailable in this session." + source_kind = "attachment" + elif src.startswith("http://") or src.startswith("https://"): try: - data, _, _ = fetch_public_bytes(src, timeout=settings.scrape_timeout) + data, final_url, _ = fetch_public_bytes(src, timeout=settings.scrape_timeout) + src = final_url or src except UnsafeURL as exc: return f"Blocked URL: {exc}" except Exception as exc: # noqa: BLE001 return f"Failed to fetch PDF: {exc}" + source_kind = "remote" else: return "Only http(s) PDF URLs are supported." - if data is None: - return "No PDF data loaded." + requested = max_chars if max_chars is not None and max_chars > 0 else settings.fetch_max_chars + + def page_factory(body_budget: int) -> PageResult: + return extract_pdf_page( + data, + max_chars=body_budget, + page_index=page_index, + page_char_offset=page_char_offset, + max_size_mb=(settings.max_upload_mb if source_kind == "attachment" else None), + ) + + def token_factory(cursor: dict[str, int]) -> str: + return make_token( + _TOOL_NAME, + src, + cursor["page_index"], + page_index=cursor["page_index"], + page_char_offset=cursor["page_char_offset"], + ) + + try: + result = format_page_to_limit( + page_factory, + token_factory, + requested_chars=requested, + result_limit=settings.max_tool_result_chars, + ) + except PageSizeError as exc: + return f"PDF pagination failed: {exc}" - text, has_more, next_page_offset = extract_pdf_text( - data, max_chars=max_chars, page_offset=page_offset + logger.info( + "pdf page source=%s page_index=%d page_char_offset=%d result_chars=%d has_more=%s", + source_kind, + page_index, + page_char_offset, + len(result), + "---page-info---" in result, ) - next_token = make_token(_TOOL_NAME, src, next_page_offset) if has_more else None - return format_paged_result(text, has_more=has_more, next_token=next_token) + return result + + +def pdf_to_text( + url_or_path: str, max_chars: int | None = None, continue_from: str | None = None +) -> str: + """Extract a losslessly paged PDF from an HTTP(S) URL. + + Args: + url_or_path: HTTP(S) PDF URL. In chat, a session attachment URI may also + be supplied by the attachment context. Arbitrary filesystem paths are + never accepted. + max_chars: Requested maximum source characters for this page. The tool + may return fewer characters so metadata and its token remain intact. + continue_from: Opaque continuation token from the preceding result. + """ + return _pdf_to_text(url_or_path, max_chars, continue_from, None) + + +def make_session_pdf_tool(attachment_loader: AttachmentLoader) -> Callable[..., str]: + """Create a session-scoped ``pdf_to_text`` supporting attachment IDs.""" + + def session_pdf_to_text( + url_or_path: str, + max_chars: int | None = None, + continue_from: str | None = None, + ) -> str: + return _pdf_to_text(url_or_path, max_chars, continue_from, attachment_loader) + + session_pdf_to_text.__name__ = _TOOL_NAME + session_pdf_to_text.__doc__ = pdf_to_text.__doc__ + return session_pdf_to_text diff --git a/src/ollama_web/tools/scrape.py b/src/ollama_web/tools/scrape.py index 39dab83..2f85c38 100644 --- a/src/ollama_web/tools/scrape.py +++ b/src/ollama_web/tools/scrape.py @@ -2,13 +2,37 @@ from __future__ import annotations +import logging + from bs4 import BeautifulSoup from ..config import settings -from .helper.pagination import format_paged_result, make_token, parse_token +from .helper.pagination import ( + PageSizeError, + format_page_to_limit, + make_token, + parse_token, + text_page, +) from .helper.safe_http import UnsafeURL, fetch_public_bytes _TOOL_NAME = "scrape_url" +logger = logging.getLogger("ollama_web.tools.scrape") + + +def extract_readable_text(data: bytes, content_type: str = "") -> str: + """Decode and normalize the complete readable body of an HTML payload.""" + 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") + soup = BeautifulSoup(html, "html.parser") + for tag in soup(["script", "style", "noscript", "template", "head", "nav", "footer", "aside"]): + tag.decompose() + container = soup.find("article") or soup.find("main") or soup + text = container.get_text(separator="\n") + lines = [line.strip() for line in text.splitlines()] + return "\n".join(line for line in lines if line) def scrape_url(url: str, max_chars: int | None = None, continue_from: str | None = None) -> str: @@ -47,43 +71,33 @@ def scrape_url(url: str, max_chars: int | None = None, continue_from: str | None # The token's URL takes precedence; it identifies the page being paged. url = str(payload.get("url") or url) offset = int(payload.get("offset", 0)) - if offset < 0: - offset = 0 try: 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}" - soup = BeautifulSoup(html, "html.parser") - - # Remove non-content tags. - for tag in soup(["script", "style", "noscript", "template", "head", "nav", "footer", "aside"]): - tag.decompose() - - # Prefer article/main content when available. - container = soup.find("article") or soup.find("main") or soup - text = container.get_text(separator="\n") - # Collapse blank lines and trailing whitespace per line. - lines = [ln.strip() for ln in text.splitlines()] - cleaned = "\n".join(ln for ln in lines if ln) + cleaned = extract_readable_text(data, content_type) if not cleaned: return "No readable text content found." - # Slice the requested page from the full cleaned text. - page = cleaned[offset : offset + limit] - if not page: - # Offset is past the end; nothing left to return. - return format_paged_result("", has_more=False, next_token=None) - - has_more = len(cleaned) > offset + limit - next_token = make_token(_TOOL_NAME, url, offset + limit) if has_more else None - return format_paged_result(page, has_more=has_more, next_token=next_token) + try: + result = format_page_to_limit( + lambda body_budget: text_page(cleaned, offset, body_budget), + lambda cursor: make_token(_TOOL_NAME, url, cursor["offset"]), + requested_chars=limit, + result_limit=settings.max_tool_result_chars, + ) + logger.info( + "scrape page offset=%d result_chars=%d has_more=%s", + offset, + len(result), + "---page-info---" in result, + ) + return result + except (PageSizeError, ValueError) as exc: + return f"Scrape pagination failed: {exc}" 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(): |