diff options
29 files changed, 1322 insertions, 143 deletions
diff --git a/docs/tools-schema.md b/docs/tools-schema.md new file mode 100644 index 0000000..ea438d3 --- /dev/null +++ b/docs/tools-schema.md @@ -0,0 +1,372 @@ +# ollama-web 標準搭載ツール スキーマ定義 + +本ドキュメントは `src/ollama_web/tools/` 以下に実装されている、LLM に公開される標準搭載ツールのスキーマ定義と設計状況を整理したものです。実装ベース（コード）を正として記載しています。 + +## 目次 + +- [概要](#概要) +- [ツール一覧](#ツール一覧) +- [1. web_search](#1-web_search) +- [2. scrape_url](#2-scrape_url) +- [3. pdf_to_text](#3-pdf_to_text) +- [4. search_and_fetch](#4-search_and_fetch) +- [ヘルパーモジュール](#ヘルパーモジュール) +- [レジストリと実行フロー](#レジストリと実行フロー) +- [ツール結果の切り詰め（truncation）ポリシー](#ツール結果の切り詰めtruncationポリシー) +- [ページング（継続トークン）仕様](#ページング継続トークン仕様) + +--- + +## 概要 + +ollama-web は LLM（Ollama）にツール呼び出し（function calling）機能を提供するため、4つの標準ツールを組み込んでいます。各ツールは Python の callable として実装され、docstring から Ollama 互換のスキーマが推論されます。MCP 経由で外部ツールを追加することも可能ですが、本ドキュメントでは標準搭載ツールのみを扱います。 + +ツールの登録・実行は [`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 呼び出し回数とキャッシュ要件の懸念からページング対象外です。 + +## ツール一覧 + +| ツール名 | 実装ファイル | 概要 | ページング対応 | +|----------|--------------|------|----------------| +| `web_search` | [`search.py`](src/ollama_web/tools/search.py:1) | DuckDuckGo による Web 検索 | ❌ 対象外 | +| `scrape_url` | [`scrape.py`](src/ollama_web/tools/scrape.py:1) | URL のページ本文抽出 | ✅ 文字オフセット | +| `pdf_to_text` | [`pdf.py`](src/ollama_web/tools/pdf.py:1) | リモート PDF のテキスト抽出 | ✅ ページオフセット | +| `search_and_fetch` | [`fetch.py`](src/ollama_web/tools/fetch.py:1) | ファイル種別検索＋取得 | ✅ ファイル単位 | + +--- + +## 1. web_search + +- **実装**: [`web_search()`](src/ollama_web/tools/search.py:10) +- **バックエンド**: [`ddgs`](https://pypi.org/project/ddgs/) (DuckDuckGo Search) + +### スキーマ + +| パラメータ | 型 | 必須 | デフォルト | 説明 | +|------------|----|------|------------|------| +| `query` | `str` | ✅ | — | 検索クエリ文字列。空文字列は空リストを返す | +| `max_results` | `int \| None` | ❌ | [`settings.search_max_results`](src/ollama_web/config.py:50) (=5) | 返却件数上限。0 以下の場合はデフォルト値を使用 | + +### 戻り値 + +`list[dict[str, str]]` — 各要素は以下のキーを持つ: + +| キー | 型 | 説明 | +|------|----|------| +| `title` | `str` | 検索結果のタイトル | +| `href` | `str` | 結果URL（`href` が無い場合は `url` をフォールバック） | +| `body` | `str` | スニペット（`body` が無い場合は `snippet` をフォールバック） | + +### 設計上の注意点 + +- 検索失敗時は例外を握りつぶし空リスト `[]` を返す（チャットループ継続のため）。 +- **ページング対象外**。検索結果キャッシュが必要になること、DuckDuckGo API 呼び出し回数増でアクセス制限に引っかかりやすくなる懸念から除外。 + +--- + +## 2. scrape_url + +- **実装**: [`scrape_url()`](src/ollama_web/tools/scrape.py:11) +- **依存**: [`BeautifulSoup`](https://www.crummy.com/software/BeautifulSoup/)、[`fetch_public_bytes()`](src/ollama_web/tools/helper/safe_http.py:72) + +### スキーマ + +| パラメータ | 型 | 必須 | デフォルト | 説明 | +|------------|----|------|------------|------| +| `url` | `str` | ✅ | — | 取得対象の絶対URL。空文字列は `"Invalid URL."` を返す | +| `max_chars` | `int \| None` | ❌ | [`settings.scrape_max_chars`](src/ollama_web/config.py:49) (=4000) | 1ページあたりの返却文字数上限。0 以下の場合はデフォルト値を使用 | +| `continue_from` | `str \| None` | ❌ | `None` | 前回戻り値の `next_token`。指定時は続きのページを取得 | + +### 戻り値 + +`str` — 1ページ分の本文テキスト。続きがある場合は末尾にページ情報フッターが付与される: + +``` +{本文1ページ分} + +---page-info--- +has_more: true +next_token: {token} +``` + +`has_more: false` のときは `next_token` 行が省略され、フッター自体も付与されない（短文の場合）。 + +### 処理フロー + +1. `continue_from` があれば [`parse_token()`](src/ollama_web/tools/helper/pagination.py:1) で `(url, offset)` を復元（トークンの `url` が引数より優先） +2. [`fetch_public_bytes()`](src/ollama_web/tools/helper/safe_http.py:72) で HTML バイト列を取得（SSRF 保護・リダイレクト追跡・サイズ制限付き） +3. `Content-Type` の `charset` からエンコーディングを判定（無ければ `utf-8`） +4. BeautifulSoup で `script`/`style`/`noscript`/`template`/`head`/`nav`/`footer`/`aside` を除去 +5. `article` > `main` > ルート の優先順位で本文コンテナを特定 +6. 改行区切りでテキスト化し、空行を圧縮（**全文をメモリに保持**） +7. `cleaned[offset : offset + max_chars]` を切り出して返却 +8. `len(cleaned) > offset + max_chars` なら `has_more: true` + `next_token`（offset=`offset + max_chars`）を付与 + +### 設計上の注意点 + +- 不正URL/ブロックURLは `Blocked URL: ...`、取得失敗は `Failed to fetch ...` のエラー文字列を返す。 +- 不正な `continue_from` トークンは `Invalid continuation token: ...` エラーを返す。 +- 全文を一度解析してメモリに保持する（`max_fetch_mb`=10MB 上限があるため実用上問題ない範囲）。 + +--- + +## 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/) 使用） + +### スキーマ + +| パラメータ | 型 | 必須 | デフォルト | 説明 | +|------------|----|------|------------|------| +| `url_or_path` | `str` | ✅ | — | PDF の URL またはローカルパス。空文字列は `"Invalid PDF source."` を返す | +| `max_chars` | `int \| None` | ❌ | [`settings.fetch_max_chars`](src/ollama_web/config.py:52) (=8000) | 1ページあたりの返却文字数上限 | +| `continue_from` | `str \| None` | ❌ | `None` | 前回戻り値の `next_token`。指定時は続きページを取得 | + +### 戻り値 + +`str` — 1ページ分のテキスト、またはエラーメッセージ。続きがある場合は `scrape_url` と同様のページ情報フッターが付与される。 + +### 制約 + +- **`http://` / `https://` のみサポート**。ローカルパスは `"Only http(s) PDF URLs are supported."` を返す。 +- [`extract_pdf_text()`](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. `continue_from` があればトークンから `(url, page_offset)` を復元 +2. PDF を全件取得・解析 +3. `page_offset` ページ目から順にテキスト抽出し `max_chars` に達したら停止 +4. 残りページがあれば `has_more: true` + `next_token`（`page_offset = 最後に抽出したページ + 1`）を付与 + +--- + +## 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) + +### スキーマ + +| パラメータ | 型 | 必須 | デフォルト | 説明 | +|------------|----|------|------------|------| +| `query` | `str` | ✅ | — | 検索キーワード。空文字列は `fetch.empty_query` プロンプトを返す | +| `file_type` | `str` | ✅ | — | ファイル拡張子（`pdf`/`csv`/`txt`/`md`/`json`/`html`/`htm`/`xml`/`py`/`js`/`ts` 等）。空文字列は `fetch.invalid_type` プロンプトを返す | +| `max_results` | `int \| None` | ❌ | [`settings.fetch_max_results`](src/ollama_web/config.py:51) (=3) | 取得候補URL数上限。`fetch_max_results` でキャップ | +| `max_chars` | `int \| None` | ❌ | [`settings.fetch_max_chars`](src/ollama_web/config.py:52) (=8000) | ファイルごとの1ページあたりの返却文字数上限 | +| `session_id` | `str \| None` | ❌ | `None` | 指定時、取得ファイルをセッションに永続化するためのキューに格納 | +| `continue_from` | `str \| None` | ❌ | `None` | 前回戻り値の `next_token`。指定時は候補リストを再検索せず続きを取得 | + +> **注**: `session_id` は LLM から直接指定されることを想定していない内部パラメータ。チャットルート（[`routes/chat.py`](src/ollama_web/routes/chat.py:72)）が `functools.wraps` でラップし、現在のセッションID を自動注入する。 + +### 戻り値 + +`str` — 現在のファイルページのメタデータと抽出テキスト。各ファイルは以下の形式: + +``` +=== 候補 {idx} === +URL: {url} +タイトル: {title} +形式: {mime} +サイズ: {size} bytes +文字数: {chars} +--- +{text} +``` + +続きがある場合（同一ファイル内の続き、または次候補ファイル）はページ情報フッターが付与される。 + +### 処理フロー + +1. `continue_from` があればトークンから `(候補リスト, file_idx, file_offset)` を復元（候補リストはトークンに埋め込み済み、再検索不要） +2. トークンが無い場合は従来通り検索→候補抽出 +3. 現在のファイルインデックスから順に処理: + - ファイル内で `max_chars` に達したら、そのファイルの `next_token` を生成して**そこで停止**（1回1ファイルのページ継続を優先） + - ファイルが完了したら次のファイルへ +4. 残り候補 or 残りファイル内容があれば `has_more: true` + `next_token` を付与 + +> トークンペイロードは `search_and_fetch` 用に拡張され、`matched`（候補URLリスト）を埋め込むことで再検索なしで続きを取得（API呼び出し増加を防止）。 + +### 設計上の注意点 + +- 取得失敗時は次候補へ進む `next_token` を発行し、ページングチェーンを継続。 +- 不正な `continue_from` トークンは `Invalid continuation token: ...` エラーを返す。 + +--- + +## ヘルパーモジュール + +### safe_http.py + +- **実装**: [`safe_http.py`](src/ollama_web/tools/helper/safe_http.py:1) +- [`validate_public_http_url()`](src/ollama_web/tools/helper/safe_http.py:34): SSRF 保護。`http(s)` のみ許可、localhost/プライベートIP/メタデータIP（`169.254.169.254`）を拒否、DNS 解決後の IP も検証 +- [`fetch_public_bytes()`](src/ollama_web/tools/helper/safe_http.py:72): ストリーミング取得。最大6回のリダイレクト追跡、`Content-Length` および実サイズで [`settings.max_fetch_mb`](src/ollama_web/config.py:53) (=10MB) を超える場合は拒否 + +### helper/pdf.py + +- **実装**: [`extract_pdf_text()`](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)` のタプル。 + +### helper/image.py + +- **実装**: [`resize_image()`](src/ollama_web/tools/helper/image.py:32) +- ビジョンモデル用の画像前処理。`max_dimension`=1024 でリサイズ、JPEG/WEBP は quality 指定。[`settings.max_image_pixels`](src/ollama_web/config.py:57) (=16M pixels) 超過時は元データを返す。LLM ツールとしては公開されていない内部ヘルパー。 + +### helper/pagination.py + +- **実装**: [`pagination.py`](src/ollama_web/tools/helper/pagination.py:1) +- [`make_token()`](src/ollama_web/tools/helper/pagination.py:1): 継続トークン生成（`base64url(JSON)`） +- [`parse_token()`](src/ollama_web/tools/helper/pagination.py:1): トークン検証・デコード。`tool` 一致・`v` バージョン・必須フィールドを検証し、不正時は `InvalidTokenError` +- [`format_paged_result()`](src/ollama_web/tools/helper/pagination.py:1): 戻り値ラッピング。`---page-info---` マーカーで本文とフッターを分離 +- [`parse_paged_result()`](src/ollama_web/tools/helper/pagination.py:1): 戻り値分割（テスト・チェーニング用） + +--- + +## レジストリと実行フロー + +```mermaid +flowchart LR + A[chat route] --> B[_build_registry] + B --> C[default_registry] + C --> D[ToolRegistry] + D -->|register| E[web_search] + D -->|register| F[scrape_url] + D -->|register| G[search_and_fetch] + D -->|register| H[pdf_to_text] + B -->|MCP| I[collect_mcp_tools] + I --> D + A --> J[llm.astream_chat_with_tools] + J -->|tool call| K[registry.execute] + K --> E + K --> F + K --> G + K --> H +``` + +- [`default_registry()`](src/ollama_web/tools/registry.py:117) が4つの標準ツールを登録した [`ToolRegistry`](src/ollama_web/tools/registry.py:17) を返す +- [`ToolRegistry.ollama_tools`](src/ollama_web/tools/registry.py:63) はビルトイン callable と MCP ツール定義 dict のリストを返す（Ollama の `tools` 引数用） +- [`ToolRegistry.execute()`](src/ollama_web/tools/registry.py:78) がツール名と引数（dict または JSON 文字列）を受け取り、JSON 文字列を返す。例外時は `{"error": ...}` の JSON + +--- + +## ツール結果の切り詰め（truncation）ポリシー + +ツール結果は LLM コンテキストに戻される前に**多段階**で切り詰められる。ページング対応ツールの `---page-info---` フッターは保護され、本文側が優先的に削られる。 + +```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 で連続ツール呼び出しを打ち切り] +``` + +| ステージ | 場所 | 関数/定数 | 制限値 | +|----------|------|-----------|--------| +| 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文字**（累積超過時に各結果を再削り） | +| 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]` を付与 + +--- + +## ページング（継続トークン）仕様 + +### 設計方針 + +- **不透明な継続トークン**（`base64url(JSON)`）を採用。LLMは値を解釈せずコピペするだけ +- 数値 `offset`/`page` は使わない → 「0始まりか1始まりか」「前回値の記憶」の推論負荷を排除 +- 状態はトークン自体にエンコード → **ステートレス**。seed変動でも「直前のツール結果の `next_token` を渡す」だけで動く +- サーバー側キャッシュ不要（トークン = 状態） + +### トークンペイロード構造 + +```json +{ + "v": 1, + "tool": "scrape_url", + "url": "https://example.com/page", + "offset": 4000 +} +``` + +`search_and_fetch` は追加フィールドを埋め込む: + +```json +{ + "v": 1, + "tool": "search_and_fetch", + "url": "https://example.com/file.pdf", + "offset": 8000, + "matched": [{"title": "...", "href": "...", "body": "..."}], + "file_idx": 1, + "file_offset": 8000 +} +``` + +### 戻り値形式 + +``` +{本文} + +---page-info--- +has_more: true +next_token: {token} +``` + +- `has_more: false` のときは `next_token` 行を省略（フッター自体も付与しない） +- `---page-info---` は [`_trim_tool_result()`](src/ollama_web/llm.py:1) の切り詰め対象外 + +### LLM誤用対策 + +- **不透明トークン**: 推測補完が不要 → seed 変更でもばらつきにくい +- **直感的なオプション名**: `continue_from`（「続きを取得する」という意図が明確） +- **システムプロンプトで明示**: [`prompts/ja.json`](src/ollama_web/prompts/ja.json) / [`en.json`](src/ollama_web/prompts/en.json) の `tool_system` に使い方を記載 +- **捏造対策**: `parse_token()` で JSON 構造・`tool` 一致・`v` 検証。捏造トークンは `InvalidTokenError` → エラーメッセージ返却 +- **退化しない**: `continue_from` 未指定時は従来動作（先頭から取得） + +### 設定項目 + +| 項目 | デフォルト | 説明 | +|------|------------|------| +| [`max_tool_result_chars`](src/ollama_web/config.py:58) | 6000 | page-info セクション分の余裕確保のため 3000→6000 に引き上げ | +| [`pagination_token_ttl`](src/ollama_web/config.py:59) | 1800 (秒) | トークン有効期限（将来のキャッシュ無効化用、現状は検証のみ） | + +### ページング全体フロー + +```mermaid +sequenceDiagram + participant LLM + participant Tool + participant Helper + LLM->>Tool: call(continue_from=None) + Tool->>Helper: fetch + parse全文 + Helper-->>Tool: 全文データ + Tool->>Tool: slice offset:offset+max_chars + Tool->>Helper: make_token(url, offset+max_chars) + Helper-->>Tool: next_token + Tool-->>LLM: text + has_more:true + next_token + Note over LLM: 続きが必要なら next_token を使用 + LLM->>Tool: call(continue_from=next_token) + Tool->>Helper: parse_token(next_token) + Helper-->>Tool: url, offset + Tool->>Helper: fetch + parse全文 + Helper-->>Tool: 全文データ + 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/__init__.py b/src/ollama_web/__init__.py index 10040d0..04b06c6 100644 --- a/src/ollama_web/__init__.py +++ b/src/ollama_web/__init__.py @@ -1,3 +1,3 @@ """Ollama web project.""" -__version__ = "0.1.0"
\ No newline at end of file +__version__ = "0.1.0" diff --git a/src/ollama_web/config.py b/src/ollama_web/config.py index 273e8bb..fa07384 100644 --- a/src/ollama_web/config.py +++ b/src/ollama_web/config.py @@ -16,15 +16,15 @@ def _csv_env(name: str) -> list[str]: class Settings: """Runtime settings resolved from environment variables.""" - ollama_host: str = field(default_factory=lambda: os.environ.get("OLLAMA_HOST", "http://127.0.0.1:11434")) + ollama_host: str = field( + default_factory=lambda: os.environ.get("OLLAMA_HOST", "http://127.0.0.1:11434") + ) default_model: str = field( default_factory=lambda: os.environ.get("OLLAMA_WEB_MODEL", "llama3.2") ) host: str = field(default_factory=lambda: os.environ.get("OLLAMA_WEB_HOST", "127.0.0.1")) port: int = field(default_factory=lambda: int(os.environ.get("OLLAMA_WEB_PORT", "8000"))) - data_dir: str = field( - default_factory=lambda: os.environ.get("OLLAMA_WEB_DATA_DIR", "data") - ) + data_dir: str = field(default_factory=lambda: os.environ.get("OLLAMA_WEB_DATA_DIR", "data")) max_upload_mb: int = field( default_factory=lambda: int(os.environ.get("OLLAMA_WEB_MAX_UPLOAD_MB", "20")) ) @@ -55,10 +55,9 @@ class Settings: max_message_chars: int = 20000 max_attachment_text_chars: int = 20000 max_image_pixels: int = 16_000_000 - max_tool_result_chars: int = 3000 - language: str = field( - default_factory=lambda: os.environ.get("OLLAMA_WEB_LANGUAGE", "ja") - ) + max_tool_result_chars: int = 6000 + pagination_token_ttl: int = 1800 + language: str = field(default_factory=lambda: os.environ.get("OLLAMA_WEB_LANGUAGE", "ja")) settings = Settings() diff --git a/src/ollama_web/i18n/__init__.py b/src/ollama_web/i18n/__init__.py index 56dd217..6ab540c 100644 --- a/src/ollama_web/i18n/__init__.py +++ b/src/ollama_web/i18n/__init__.py @@ -9,7 +9,7 @@ from __future__ import annotations import json from pathlib import Path -from typing import Any +from typing import Any, cast from ollama_web.config import settings @@ -22,7 +22,7 @@ def _load_messages(lang: str) -> dict[str, Any]: path = _MESSAGES_DIR / f"{lang}.json" try: with path.open("r", encoding="utf-8") as f: - return json.load(f) + return cast(dict[str, Any], json.load(f)) except (FileNotFoundError, json.JSONDecodeError): return {} @@ -81,4 +81,4 @@ def t(key: str, lang: str | None = None, fallback: str | None = None) -> str: value = _resolve(key, get_messages(resolve_language(lang))) if isinstance(value, str): return value - return fallback if fallback is not None else key + return fallback if fallback is not None else key
\ No newline at end of file diff --git a/src/ollama_web/llm.py b/src/ollama_web/llm.py index 76a61a5..2839051 100644 --- a/src/ollama_web/llm.py +++ b/src/ollama_web/llm.py @@ -30,6 +30,10 @@ MAX_TOOL_CONTEXT_CHARS = 24000 # Type alias matching ollama's expected tools signature. ToolCallable = Callable[..., Any] +# Marker that separates a paged tool result's body from its page-info footer. +# Must match ``ollama_web.tools.helper.pagination.PAGE_INFO_MARKER``. +_PAGE_INFO_MARKER = "\n---page-info---" + def _safe_tool_arguments(arguments: Any) -> Any: return redact_secrets(arguments) @@ -40,8 +44,6 @@ def _safe_tool_result(result: str) -> str: return str(redacted) - - def get_client(host: str | None = None) -> ollama.Client: """Return an ``ollama.Client`` configured for the configured host.""" import httpx @@ -340,17 +342,40 @@ def chat_with_tools( yield {"type": "thinking", "content": thinking} if content: yield {"type": "delta", "content": content} - done_event = {"type": "done"} + done_event_final: dict[str, Any] = {"type": "done"} if tokens: - done_event.update(tokens) - yield done_event + done_event_final.update(tokens) + yield done_event_final def _trim_tool_result(result: str, limit: int) -> str: - """Truncate ``result`` to ``limit`` chars while keeping it readable.""" + """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. + """ if len(result) <= limit: return result - # Prefer cutting at a paragraph boundary. + + # Protect the page-info footer from truncation. ``split`` keeps the + # marker on the footer side so the ``---page-info---`` line is preserved. + 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 + + # No footer: legacy truncation behaviour. cutoff = result[:limit].rsplit("\n\n", 1)[0] if len(cutoff) < limit * 0.5: # No good paragraph boundary; cut at a space instead. @@ -457,9 +482,7 @@ def stream_chat_with_tools( # stream with almost no meaningful content. Log it and retry a few # times before giving up. is_empty = ( - not full_message.get("content") - and not full_message.get("thinking") - and not tool_calls + not full_message.get("content") and not full_message.get("thinking") and not tool_calls ) if is_empty and empty_retries < MAX_EMPTY_RETRIES: empty_retries += 1 @@ -497,7 +520,7 @@ def stream_chat_with_tools( "type": "warning", "message": t("status.empty_response", lang=language), } - done_event = {"type": "done"} + done_event: dict[str, Any] = {"type": "done"} tokens = _token_counts(last_chunk) if tokens: done_event.update(tokens) @@ -594,11 +617,11 @@ def stream_chat_with_tools( if content: yield {"type": "delta", "content": content} logger.info("final answer stream finished elapsed=%.2fs", time.time() - t0) - done_event = {"type": "done"} + done_event_final: dict[str, Any] = {"type": "done"} tokens = _token_counts(last_chunk) if tokens: - done_event.update(tokens) - yield done_event + done_event_final.update(tokens) + yield done_event_final async def astream_chat_with_tools( @@ -690,9 +713,7 @@ def _assistant_message_chunk(full_message: dict[str, Any]) -> Message | None: name = c.get("name", "") args = c.get("arguments", {}) tc_list.append( - Message.ToolCall( - function=Message.ToolCall.Function(name=name, arguments=args) - ) + Message.ToolCall(function=Message.ToolCall.Function(name=name, arguments=args)) ) kwargs["tool_calls"] = tc_list return Message(**kwargs)
\ No newline at end of file diff --git a/src/ollama_web/mcp.py b/src/ollama_web/mcp.py index e52860b..835a160 100644 --- a/src/ollama_web/mcp.py +++ b/src/ollama_web/mcp.py @@ -104,18 +104,14 @@ def validate_mcp_server_name(name: Any) -> str: """Return a safe MCP server name or raise ``ValueError``.""" value = str(name) if not _is_safe_name(value): - raise ValueError( - "server name must match ^[A-Za-z0-9_-]{1,64}$ and must not contain '__'" - ) + raise ValueError("server name must match ^[A-Za-z0-9_-]{1,64}$ and must not contain '__'") return value def _validate_mcp_tool_name(name: Any) -> str: value = str(name) if not _is_safe_name(value): - raise ValueError( - "tool name must match ^[A-Za-z0-9_-]{1,64}$ and must not contain '__'" - ) + raise ValueError("tool name must match ^[A-Za-z0-9_-]{1,64}$ and must not contain '__'") return value @@ -249,7 +245,6 @@ def wrap_untrusted_tool_output(text: str) -> str: return f"{_UNTRUSTED_PREFIX}\n<tool-output>\n{body}\n</tool-output>" - def _to_ollama_tool(server_name: str, tool: Any) -> dict[str, Any]: """Convert an MCP ``Tool`` object into an ollama-compatible tool dict. @@ -437,7 +432,9 @@ async def collect_mcp_tools( results = await asyncio.gather(*tasks, return_exceptions=True) all_tools: list[dict[str, Any]] = [] for name, result in zip(servers.keys(), results, strict=True): - if isinstance(result, Exception): + # asyncio.gather(return_exceptions=True) can surface BaseException + # (e.g. CancelledError), so check BaseException here. + if isinstance(result, BaseException): logger.warning("failed to list tools from MCP server %s: %s", name, result) continue all_tools.extend(result) @@ -663,4 +660,4 @@ def make_mcp_executor( ) -> str: return call_mcp_tool_sync(config, full_name, arguments) - return executor + return executor
\ No newline at end of file diff --git a/src/ollama_web/prompts/__init__.py b/src/ollama_web/prompts/__init__.py index 59bdb48..5b6b53e 100644 --- a/src/ollama_web/prompts/__init__.py +++ b/src/ollama_web/prompts/__init__.py @@ -9,7 +9,7 @@ from __future__ import annotations import json from pathlib import Path -from typing import Any +from typing import Any, cast from ollama_web.config import settings @@ -23,7 +23,7 @@ def _load_prompts(lang: str) -> dict[str, Any]: path = _DIR / f"{lang}.json" try: with path.open("r", encoding="utf-8") as f: - return json.load(f) + return cast(dict[str, Any], json.load(f)) except (FileNotFoundError, json.JSONDecodeError): return {} @@ -54,4 +54,4 @@ def get_prompt(key: str, lang: str | None = None, fallback: str | None = None) - current = current[part] if isinstance(current, str): return current - return fallback if fallback is not None else key + return fallback if fallback is not None else key
\ No newline at end of file diff --git a/src/ollama_web/prompts/en.json b/src/ollama_web/prompts/en.json index 72a0b20..42e2e68 100644 --- a/src/ollama_web/prompts/en.json +++ b/src/ollama_web/prompts/en.json @@ -1,5 +1,5 @@ { - "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.", + "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.", "attachment_prefix": "The following is the content of files attached by the user:", @@ -16,4 +16,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 00698c2..88af710 100644 --- a/src/ollama_web/prompts/ja.json +++ b/src/ollama_web/prompts/ja.json @@ -1,5 +1,5 @@ { - "tool_system": "あなたはWeb検索・スクレイピング・ファイル取得ツールを利用できます。一度の回答で必要以上に多くの検索を同時に行わず、結果を確認してから次の調査を行ってください。各ツール結果は要約されて提供されるため、重要なポイントを把握して活用してください。", + "tool_system": "あなたはWeb検索・スクレイピング・ファイル取得ツールを利用できます。一度の回答で必要以上に多くの検索を同時に行わず、結果を確認してから次の調査を行ってください。各ツール結果は要約されて提供されるため、重要なポイントを把握して活用してください。ツール結果が途切れている場合、結果末尾の next_token を使って続きを取得できます。next_token の値は解釈せず、そのまま continue_from パラメータに渡してください。続きが必要ない場合は continue_from を省略してください。", "current_time": "現在時刻: {time}", "tool_limit_instruction": "ツール呼び出しの上限に達しました。これまでの検索結果と会話を元に、ユーザーが求めた回答を最終的にまとめてください。", "attachment_prefix": "以下はユーザーが添付したファイルの内容です：", @@ -16,4 +16,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 d69fa36..e5b90d3 100644 --- a/src/ollama_web/routes/chat.py +++ b/src/ollama_web/routes/chat.py @@ -113,8 +113,8 @@ def _attachment_context(store: SessionStore, session_id: str, file_ids: list[str continue text = _limit_text(chat_file.text or "", settings.max_attachment_text_chars) parts.append( - f"<attached-file name=\"{html.escape(chat_file.name)}\" id=\"{chat_file.id}\" " - f"source=\"{chat_file.source}\">\n{text}\n</attached-file>" + f'<attached-file name="{html.escape(chat_file.name)}" id="{chat_file.id}" ' + f'source="{chat_file.source}">\n{text}\n</attached-file>' ) if not parts: return "" @@ -276,8 +276,14 @@ async def _chat_event_stream( logger.info("enabled tools capability fallback for cloud model %s", model) async for event in llm.astream_chat_with_tools( - messages, model, think=think, host=host, registry=registry, - capabilities=capabilities, options=ollama_options, language=language, + messages, + model, + think=think, + host=host, + registry=registry, + capabilities=capabilities, + options=ollama_options, + language=language, ): if event.get("type") == "delta": assistant_text = _limit_text( @@ -296,8 +302,7 @@ async def _chat_event_stream( # Persist the error as the assistant message so it survives reloads # and re-renders in the chat pane. assistant_text = ( - f"[ERROR] {error_msg}\n\n" - f"{t('errors.vision_model_required', lang=language)}" + f"[ERROR] {error_msg}\n\n{t('errors.vision_model_required', lang=language)}" ) finally: @@ -375,4 +380,4 @@ async def chat(request: Request) -> StreamingResponse: async def _single_event(data: str) -> AsyncIterator[bytes]: """Yield a single SSE event and finish.""" - yield f"data: {data}\n\n".encode()
\ No newline at end of file + yield f"data: {data}\n\n".encode() diff --git a/src/ollama_web/routes/mcp.py b/src/ollama_web/routes/mcp.py index ee8ba65..ba3628c 100644 --- a/src/ollama_web/routes/mcp.py +++ b/src/ollama_web/routes/mcp.py @@ -90,9 +90,7 @@ async def put_servers(request: Request) -> JSONResponse: continue cleaned["mcpServers"][name] = cleaned_server else: - errors.append( - f"Server '{name}' must define either 'command' (stdio) or 'url' (http)" - ) + errors.append(f"Server '{name}' must define either 'command' (stdio) or 'url' (http)") if errors: return JSONResponse({"error": errors}, status_code=400) diff --git a/src/ollama_web/routes/models.py b/src/ollama_web/routes/models.py index 406679a..0cc9941 100644 --- a/src/ollama_web/routes/models.py +++ b/src/ollama_web/routes/models.py @@ -33,4 +33,4 @@ async def health(request: Request) -> JSONResponse: host = request.app.state.settings.ollama_host models = llm.list_models(host=host) reachable = models is not None - return JSONResponse({"reachable": reachable, "models": models})
\ No newline at end of file + return JSONResponse({"reachable": reachable, "models": models}) diff --git a/src/ollama_web/routes/sessions.py b/src/ollama_web/routes/sessions.py index 64e1bc9..e1e669c 100644 --- a/src/ollama_web/routes/sessions.py +++ b/src/ollama_web/routes/sessions.py @@ -127,7 +127,8 @@ 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"): - text = extract_pdf_text(data) + # extract_pdf_text now returns (text, has_more, next_page_offset). + text = extract_pdf_text(data)[0] else: # Try decode as text for text-like uploads. try: @@ -171,4 +172,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}) + return JSONResponse({"cleared": True, "session": session})
\ No newline at end of file diff --git a/src/ollama_web/sessions.py b/src/ollama_web/sessions.py index 388066b..4075649 100644 --- a/src/ollama_web/sessions.py +++ b/src/ollama_web/sessions.py @@ -92,7 +92,7 @@ class ChatMessage: token_count: dict[str, int] | None = None def to_dict(self) -> dict[str, Any]: - out = { + out: dict[str, Any] = { "role": self.role, "content": self.content, "attachments": list(self.attachments), diff --git a/src/ollama_web/settings_store.py b/src/ollama_web/settings_store.py index 6cdd022..d565ec6 100644 --- a/src/ollama_web/settings_store.py +++ b/src/ollama_web/settings_store.py @@ -61,9 +61,7 @@ def settings_file(data_dir: str | Path) -> Path: def _require_keys(value: dict[str, Any], allowed: set[str], where: str) -> None: unknown = set(value) - allowed if unknown: - raise SettingsValidationError( - f"Unknown {where} field(s): {', '.join(sorted(unknown))}" - ) + raise SettingsValidationError(f"Unknown {where} field(s): {', '.join(sorted(unknown))}") def _validate_options(raw: Any) -> dict[str, Any]: diff --git a/src/ollama_web/tools/fetch.py b/src/ollama_web/tools/fetch.py index 20031f1..a9a92f8 100644 --- a/src/ollama_web/tools/fetch.py +++ b/src/ollama_web/tools/fetch.py @@ -9,9 +9,11 @@ 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.safe_http import fetch_public_bytes from .scrape import scrape_url +from .search import web_search _FetchedFile = dict[str, Any] @@ -34,6 +36,8 @@ _FILETYPE_SUFFIXES: dict[str, str] = { "ts": "ts", } +_TOOL_NAME = "search_and_fetch" + def set_session_context(session_id: str | None) -> None: """Set the active session id so fetched files can be persisted later.""" @@ -108,7 +112,8 @@ def _text_from_bytes(url: str, data: bytes, file_type: str, max_chars: int) -> s name = f"file.{file_type}" if file_type == "pdf": - return extract_pdf_text(data, max_chars=max_chars) + text, _, _ = extract_pdf_text(data, max_chars=max_chars, page_offset=0) + return text # Decode text-like formats. try: @@ -127,25 +132,61 @@ def _text_from_bytes(url: str, data: bytes, file_type: str, max_chars: int) -> s 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. + """ + 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 + + 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"}: + # 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) + return text + + def search_and_fetch( query: str, file_type: str, max_results: int | None = None, max_chars: int | None = None, session_id: str | None = None, + continue_from: str | None = None, ) -> str: """Search the web for files of a given type and return their text content. + Results are returned one file at a time. When a file is truncated or more + candidate files remain, a ``next_token`` is appended that can be passed + back via ``continue_from`` to resume. + Args: query: Search keywords. file_type: Desired file extension (e.g. "pdf", "csv", "txt", "md", "json"). max_results: Number of candidate URLs to fetch. max_chars: Max characters per file to return. session_id: If provided, fetched files are queued for session persistence. + continue_from: Opaque continuation token returned in a previous call's + ``next_token``. When provided, resumes from the encoded file index and + in-file offset without re-running the search. Returns: - A formatted string containing metadata and extracted text for each - successfully fetched file. + A formatted string containing metadata and extracted text for the + current file page, followed by a ``---page-info---`` footer when more + content is available. """ if not query or not query.strip(): return get_prompt("fetch.empty_query") @@ -155,17 +196,39 @@ def search_and_fetch( return get_prompt("fetch.invalid_type") limit = min( - max_results - if max_results is not None and max_results > 0 - else settings.fetch_max_results, + max_results if max_results is not None and max_results > 0 else settings.fetch_max_results, settings.fetch_max_results, ) per_file_chars = ( max_chars if max_chars is not None and max_chars > 0 else settings.fetch_max_chars ) - from .search import web_search + # --- Pagination resume path --- + if continue_from: + try: + 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: + 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)) + return _render_file_page( + matched=resumed_matched, + file_idx=file_idx, + file_offset=file_offset, + ft=ft, + per_file_chars=per_file_chars, + session_id=session_id, + query=query, + ) + # --- Initial search path --- search_query = _filetype_query(query, ft) results = web_search(search_query, max_results=limit * 2) @@ -178,38 +241,109 @@ def search_and_fetch( break if not matched: - return get_prompt("fetch.no_results", fallback="").replace("{file_type}", ft).replace("{query}", query) + return ( + get_prompt("fetch.no_results", fallback="") + .replace("{file_type}", ft) + .replace("{query}", query) + ) - outputs: list[str] = [] - for idx, item in enumerate(matched, 1): - url = item.get("href", "") - title = item.get("title", "") - try: - data = _fetch_bytes(url) - except Exception as exc: # noqa: BLE001 - p = get_prompts() - outputs.append( - f"=== {p['candidate']} {idx} ===\n" - f"URL: {url}\n{p['title']}: {title}\n{p['status']}: {p['fetch_failed']} ({exc})\n" - ) - continue - - text = _text_from_bytes(url, data, ft, 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, text, url) - p = get_prompts() - outputs.append( - f"=== {p['candidate']} {idx} ===\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(text)}\n" - f"{p['separator']}\n{text}\n" + return _render_file_page( + matched=matched, + file_idx=0, + file_offset=0, + ft=ft, + per_file_chars=per_file_chars, + session_id=session_id, + query=query, + ) + + +def _render_file_page( + *, + matched: list[dict[str, str]], + file_idx: int, + file_offset: int, + ft: str, + per_file_chars: int, + session_id: str | None, + query: str, +) -> str: + """Render a single file page and emit a continuation token if needed. + + The token encodes the candidate list so subsequent pages do not need to + re-run the web search. + """ + if file_idx >= len(matched): + # All candidates exhausted. + return format_paged_result("", has_more=False, next_token=None) + + item = matched[file_idx] + url = item.get("href", "") + title = item.get("title", "") + p = get_prompts().get("fetch", {}) + + try: + data = _fetch_bytes(url) + except Exception as exc: # noqa: BLE001 + # On fetch failure, advance to the next candidate rather than aborting + # the whole pagination chain. + body = ( + 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=matched, + file_idx=file_idx + 1, + file_offset=0, + ) + 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 = ( + 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" + ) + + if has_more_in_file: + # Continue within the same file. + next_token = make_token( + _TOOL_NAME, + url, + file_offset + per_file_chars, + 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) + + # Move to the next candidate file. + if file_idx + 1 < len(matched): + next_token = make_token( + _TOOL_NAME, + url, + 0, + matched=matched, + file_idx=file_idx + 1, + file_offset=0, ) + return format_paged_result(body, has_more=True, next_token=next_token) - return "\n".join(outputs) + return format_paged_result(body, has_more=False, next_token=None)
\ No newline at end of file diff --git a/src/ollama_web/tools/helper/pagination.py b/src/ollama_web/tools/helper/pagination.py new file mode 100644 index 0000000..7f18ac8 --- /dev/null +++ b/src/ollama_web/tools/helper/pagination.py @@ -0,0 +1,177 @@ +"""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. +""" + +from __future__ import annotations + +import base64 +import json +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 + +# 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---" + + +class InvalidTokenError(ValueError): + """Raised when a continuation token is malformed or invalid.""" + + +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. + """ + payload: dict[str, Any] = { + "v": _TOKEN_VERSION, + "tool": tool, + "url": url, + "offset": int(offset), + } + payload.update(extra) + raw = json.dumps(payload, ensure_ascii=False, separators=(",", ":")).encode("utf-8") + return _b64url_encode(raw) + + +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``. + """ + if not token or not token.strip(): + raise InvalidTokenError("Empty continuation token.") + + try: + raw = _b64url_decode(token.strip()) + except (ValueError, base64.binascii.Error) as exc: # type: ignore[attr-defined] + raise InvalidTokenError(f"Token is not valid base64url: {exc}") from exc + + try: + payload = json.loads(raw.decode("utf-8")) + except (UnicodeDecodeError, json.JSONDecodeError) as exc: + raise InvalidTokenError(f"Token payload is not valid JSON: {exc}") from exc + + 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 "url" not in payload or "offset" not in payload: + raise InvalidTokenError("Token payload missing 'url' or 'offset'.") + + try: + payload["offset"] = int(payload["offset"]) + except (TypeError, ValueError) as exc: + raise InvalidTokenError(f"Token 'offset' is not an integer: {exc}") from exc + + 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. + """ + 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}" + + +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). + """ + 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:"): + 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 51d3b3a..7b9fdac 100644 --- a/src/ollama_web/tools/helper/pdf.py +++ b/src/ollama_web/tools/helper/pdf.py @@ -9,31 +9,76 @@ from pypdf import PdfReader from ...config import settings -def extract_pdf_text(data: bytes, max_chars: int | None = None) -> str: - """Extract text from PDF bytes.""" +def extract_pdf_text( + 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). + """ if not data: - return "Empty PDF 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." + return "PDF exceeds the configured size limit.", False, 0 try: reader = PdfReader(stream=io.BytesIO(data)) except Exception as exc: # noqa: BLE001 - return f"Failed to parse PDF: {exc}" + return f"Failed to parse PDF: {exc}", False, 0 + + 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) parts: list[str] = [] - for page in reader.pages[: settings.max_pdf_pages]: + collected = 0 + last_extracted_page = page_offset + for page_index in range(page_offset, end): try: - text = page.extract_text() - if text: - parts.append(text) + text = reader.pages[page_index].extract_text() except Exception: # noqa: BLE001 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: + break full = "\n\n".join(parts) if len(full) > limit: full = full[:limit].rsplit(" ", 1)[0] + "\n…[truncated]" - return full or "No text could be extracted from the PDF." + + # 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 diff --git a/src/ollama_web/tools/pdf.py b/src/ollama_web/tools/pdf.py index 373506c..9fdb497 100644 --- a/src/ollama_web/tools/pdf.py +++ b/src/ollama_web/tools/pdf.py @@ -3,24 +3,50 @@ from __future__ import annotations from ..config import settings +from .helper.pagination import format_paged_result, make_token, parse_token from .helper.pdf import extract_pdf_text from .helper.safe_http import UnsafeURL, fetch_public_bytes +_TOOL_NAME = "pdf_to_text" -def pdf_to_text(url_or_path: str, max_chars: int | None = None) -> str: + +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. + Args: url_or_path: File path or URL pointing to a PDF. - max_chars: Maximum characters to return. + 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 or an error message. + Extracted plain text for this page, followed by a ``---page-info---`` + footer when more pages are available, or an error message. """ 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 + if continue_from: + try: + payload = parse_token(continue_from, _TOOL_NAME) + 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://"): try: @@ -34,4 +60,9 @@ def pdf_to_text(url_or_path: str, max_chars: int | None = None) -> str: if data is None: return "No PDF data loaded." - return extract_pdf_text(data, max_chars=max_chars) + + text, has_more, next_page_offset = extract_pdf_text( + data, max_chars=max_chars, page_offset=page_offset + ) + 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) diff --git a/src/ollama_web/tools/scrape.py b/src/ollama_web/tools/scrape.py index d5fa2f4..39dab83 100644 --- a/src/ollama_web/tools/scrape.py +++ b/src/ollama_web/tools/scrape.py @@ -5,28 +5,51 @@ from __future__ import annotations from bs4 import BeautifulSoup from ..config import settings +from .helper.pagination import format_paged_result, make_token, parse_token from .helper.safe_http import UnsafeURL, fetch_public_bytes +_TOOL_NAME = "scrape_url" -def scrape_url(url: str, max_chars: int | None = None) -> str: + +def scrape_url(url: str, max_chars: int | None = None, continue_from: str | None = None) -> str: """Fetch a web page and return its main text content. Scripts, styles and other non-content elements are removed before - collapsing whitespace. The returned text is truncated to ``max_chars``. + collapsing whitespace. The returned text is a page of at most ``max_chars`` + characters; when more content is available, a ``next_token`` is appended + that can be passed back via ``continue_from`` to fetch the next page. Args: url: The absolute URL of the page to scrape. - max_chars: Maximum number of characters to return. Defaults to the value - in settings. + max_chars: Maximum number of characters to return per page. Defaults to + the value in settings. + continue_from: Opaque continuation token returned in a previous call's + ``next_token``. When provided, the page is resumed from the encoded + offset instead of the beginning. Returns: - The extracted plain text of the page, truncated to ``max_chars``. If the - request fails a short error description is returned instead. + The extracted plain text of the page (one page of ``max_chars``), + followed by a ``---page-info---`` footer when more content is available. + If the request fails a short error description is returned instead. """ if not url or not url.strip(): return "Invalid URL." limit = max_chars if max_chars is not None and max_chars > 0 else settings.scrape_max_chars + + # Resolve the resumption offset from the continuation token, if any. + offset = 0 + if continue_from: + try: + payload = parse_token(continue_from, _TOOL_NAME) + except ValueError as exc: + return f"Invalid continuation token: {exc}" + # 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" @@ -52,6 +75,15 @@ def scrape_url(url: str, max_chars: int | None = None) -> str: lines = [ln.strip() for ln in text.splitlines()] cleaned = "\n".join(ln for ln in lines if ln) - if len(cleaned) > limit: - cleaned = cleaned[:limit].rsplit(" ", 1)[0] + "\n…[truncated]" - return cleaned or "No readable text content found." + 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) diff --git a/src/ollama_web/tools/search.py b/src/ollama_web/tools/search.py index 0af2a3e..761f182 100644 --- a/src/ollama_web/tools/search.py +++ b/src/ollama_web/tools/search.py @@ -22,9 +22,7 @@ def web_search(query: str, max_results: int | None = None) -> list[dict[str, str return [] limit = ( - max_results - if max_results is not None and max_results > 0 - else settings.search_max_results + max_results if max_results is not None and max_results > 0 else settings.search_max_results ) results: list[dict[str, str]] = [] diff --git a/tests/__init__.py b/tests/__init__.py index 2b1df6a..ce993ff 100644 --- a/tests/__init__.py +++ b/tests/__init__.py @@ -1 +1 @@ -"""Test package for ollama-web."""
\ No newline at end of file +"""Test package for ollama-web.""" diff --git a/tests/test_clear_messages.py b/tests/test_clear_messages.py index a0653ca..675d5b8 100644 --- a/tests/test_clear_messages.py +++ b/tests/test_clear_messages.py @@ -91,4 +91,4 @@ def test_add_message_stores_token_count_and_clears_it(): assert "token_count" not in cleared reloaded_after = store.get(session_id) assert reloaded_after is not None - assert "token_count" not in reloaded_after
\ No newline at end of file + assert "token_count" not in reloaded_after diff --git a/tests/test_i18n.py b/tests/test_i18n.py index a2d0ffd..c8e4f33 100644 --- a/tests/test_i18n.py +++ b/tests/test_i18n.py @@ -2,7 +2,6 @@ from __future__ import annotations -import json import re from pathlib import Path @@ -10,7 +9,6 @@ import pytest from ollama_web import i18n, prompts - _SRC = Path(__file__).resolve().parent.parent / "src" / "ollama_web" @@ -106,6 +104,8 @@ def test_get_prompt_returns_expected_strings(): """Basic sanity checks for the get_prompt() helper.""" assert prompts.get_prompt("tool_system") == prompts.get_prompts("ja")["tool_system"] assert prompts.get_prompt("tool_system", lang="en") == prompts.get_prompts("en")["tool_system"] - assert prompts.get_prompt("current_time", lang="ja") == prompts.get_prompts("ja")["current_time"] + assert ( + prompts.get_prompt("current_time", lang="ja") == prompts.get_prompts("ja")["current_time"] + ) assert prompts.get_prompt("missing.prompt", fallback="fallback") == "fallback" assert prompts.get_prompt("missing.prompt") == "missing.prompt" diff --git a/tests/test_llm.py b/tests/test_llm.py index 3002a77..b4e68ad 100644 --- a/tests/test_llm.py +++ b/tests/test_llm.py @@ -104,4 +104,4 @@ def test_token_counts_extracts_from_object() -> None: def test_token_counts_returns_none_when_missing() -> None: assert _token_counts({}) is None - assert _token_counts(None) is None
\ No newline at end of file + assert _token_counts(None) is None diff --git a/tests/test_mcp_security_diagnostics.py b/tests/test_mcp_security_diagnostics.py index be73f63..f0e8afa 100644 --- a/tests/test_mcp_security_diagnostics.py +++ b/tests/test_mcp_security_diagnostics.py @@ -330,15 +330,13 @@ def test_diagnostic_llm_auto_executes_dangerous_mcp_tool_without_confirmation( ) ) - assert marker.exists() assert any( event.get("type") == "tool_start" and event["name"] == dangerous_name for event in events ) assert any( - event.get("type") == "tool_end" and event["name"] == dangerous_name - for event in events + event.get("type") == "tool_end" and event["name"] == dangerous_name for event in events ) assert any( json.loads(marker.read_text(encoding="utf-8"))["name"] == dangerous_name diff --git a/tests/test_models.py b/tests/test_models.py index 177b86b..6c648bf 100644 --- a/tests/test_models.py +++ b/tests/test_models.py @@ -66,4 +66,4 @@ def test_model_capabilities_route_still_works_without_slashes() -> None: assert response.status_code == 200 payload = response.json() assert "capabilities" in payload - assert "completion" in [c.lower() for c in payload["capabilities"]]
\ No newline at end of file + assert "completion" in [c.lower() for c in payload["capabilities"]] diff --git a/tests/test_settings.py b/tests/test_settings.py index 29ec2c1..0f334f4 100644 --- a/tests/test_settings.py +++ b/tests/test_settings.py @@ -123,8 +123,7 @@ def test_settings_api_and_page_elements(tmp_path: Path) -> None: assert 'data-settings-screen="ui"' in html assert 'data-settings-screen="ollama"' in html num_keep = ( - 'data-option="num_keep" data-type="int" type="number" ' - 'step="1" placeholder="24"' + 'data-option="num_keep" data-type="int" type="number" step="1" placeholder="24"' ) assert num_keep in html assert "<|eot_id|>" in html @@ -159,9 +158,7 @@ def test_options_are_forwarded_to_tool_and_final_rounds(monkeypatch: pytest.Monk "message": { "role": "assistant", "content": "", - "tool_calls": [ - {"function": {"name": "missing_tool", "arguments": {}}} - ], + "tool_calls": [{"function": {"name": "missing_tool", "arguments": {}}}], } } return {"message": {"role": "assistant", "content": "done"}} diff --git a/tests/test_tools.py b/tests/test_tools.py index 897c597..67c66ef 100644 --- a/tests/test_tools.py +++ b/tests/test_tools.py @@ -5,8 +5,21 @@ from __future__ import annotations import json from unittest.mock import MagicMock, patch +import pytest + +from ollama_web.tools.helper.pagination import ( + InvalidTokenError, + format_paged_result, + make_token, + parse_paged_result, + parse_token, +) from ollama_web.tools.registry import ToolRegistry, default_registry +# --------------------------------------------------------------------------- +# Registry tests (existing) +# --------------------------------------------------------------------------- + def test_registry_register_and_execute(): reg = ToolRegistry() @@ -61,6 +74,13 @@ def test_default_registry_has_tools(): reg = default_registry() assert "web_search" in reg.names assert "scrape_url" in reg.names + assert "pdf_to_text" in reg.names + assert "search_and_fetch" in reg.names + + +# --------------------------------------------------------------------------- +# web_search tests (existing) +# --------------------------------------------------------------------------- def test_web_search_empty_query(): @@ -94,7 +114,363 @@ def test_web_search_mocked(): assert result[0]["title"] == "T" +# --------------------------------------------------------------------------- +# scrape_url tests (existing + pagination) +# --------------------------------------------------------------------------- + + def test_scrape_url_empty(): from ollama_web.tools.scrape import scrape_url - assert scrape_url("") == "Invalid URL."
\ No newline at end of file + assert scrape_url("") == "Invalid URL." + + +def _fake_fetch_public_bytes(html: bytes, content_type: str = "text/html; charset=utf-8"): + """Return a fake fetch_public_bytes that yields the given HTML.""" + return MagicMock(return_value=(html, "https://example.com/page", content_type)) + + +def test_scrape_url_no_continue_from_returns_first_page(): + """continue_from=None returns the first page with has_more + next_token.""" + from ollama_web.tools import scrape as scrape_mod + + # Build HTML longer than max_chars so pagination kicks in. + body = "<article>" + ("A" * 100) + "</article>" + html = f"<html><body>{body}</body></html>" + with patch.object( + scrape_mod, "fetch_public_bytes", _fake_fetch_public_bytes(html.encode("utf-8")) + ): + result = scrape_mod.scrape_url("https://example.com/page", max_chars=40) + body_text, has_more, next_token = parse_paged_result(result) + assert len(body_text) == 40 + assert has_more is True + assert next_token is not None + + +def test_scrape_url_continue_from_returns_next_page(): + """Passing next_token returns the continuation page.""" + from ollama_web.tools import scrape as scrape_mod + + body = "<article>" + ("A" * 100) + "</article>" + html = f"<html><body>{body}</body></html>" + with patch.object( + scrape_mod, "fetch_public_bytes", _fake_fetch_public_bytes(html.encode("utf-8")) + ): + first = scrape_mod.scrape_url("https://example.com/page", max_chars=40) + _, _, next_token = parse_paged_result(first) + + with patch.object( + scrape_mod, "fetch_public_bytes", _fake_fetch_public_bytes(html.encode("utf-8")) + ): + second = scrape_mod.scrape_url( + "https://example.com/page", max_chars=40, continue_from=next_token + ) + body_text, has_more, _ = parse_paged_result(second) + # Second page should start at offset 40. + assert body_text == "A" * 40 + assert has_more is True + + +def test_scrape_url_pagination_covers_full_text(): + """Iterating all pages reconstructs the full cleaned text.""" + from ollama_web.tools import scrape as scrape_mod + + body = "<article>" + ("A" * 100) + "</article>" + html = f"<html><body>{body}</body></html>" + collected = "" + token = None + for _ in range(10): + with patch.object( + scrape_mod, "fetch_public_bytes", _fake_fetch_public_bytes(html.encode("utf-8")) + ): + result = scrape_mod.scrape_url( + "https://example.com/page", max_chars=40, continue_from=token + ) + body_text, has_more, token = parse_paged_result(result) + collected += body_text + if not has_more: + break + assert collected == "A" * 100 + + +def test_scrape_url_continue_from_invalid_token(): + """An invalid continuation token returns an error message.""" + from ollama_web.tools.scrape import scrape_url + + result = scrape_url("https://example.com/page", continue_from="not-a-valid-token") + assert "Invalid continuation token" in result + + +def test_scrape_url_short_text_no_next_token(): + """Text shorter than max_chars has no next_token.""" + from ollama_web.tools import scrape as scrape_mod + + html = b"<html><body><article>short</article></body></html>" + with patch.object(scrape_mod, "fetch_public_bytes", _fake_fetch_public_bytes(html)): + result = scrape_mod.scrape_url("https://example.com/page", max_chars=4000) + body_text, has_more, next_token = parse_paged_result(result) + assert body_text == "short" + assert has_more is False + assert next_token is None + + +# --------------------------------------------------------------------------- +# pdf_to_text tests (pagination) +# --------------------------------------------------------------------------- + + +def _make_multipage_pdf(num_pages: int, text_per_page: str = "Page content. ") -> bytes: + """Build an in-memory multi-page PDF using pypdf's writer.""" + from pypdf import PdfWriter + + 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() + writer.write(buf) + return buf.getvalue() + + +def test_pdf_to_text_only_https_supported(): + """Local paths are rejected (regression).""" + from ollama_web.tools.pdf import pdf_to_text + + result = pdf_to_text("/local/path.pdf") + assert "Only http(s) PDF URLs are supported." in result + + +def test_pdf_to_text_continue_from_invalid_token(): + """Invalid continuation token returns an error.""" + from ollama_web.tools.pdf import pdf_to_text + + result = pdf_to_text("https://example.com/doc.pdf", continue_from="bad-token") + assert "Invalid continuation token" in result + + +def test_pdf_to_text_pagination_roundtrip(monkeypatch): + """Pagination across pages reconstructs all page text.""" + 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) + + # 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"), + ) + + collected = "" + token = None + for _ in range(10): + result = pdf_mod.pdf_to_text( + "https://example.com/doc.pdf", max_chars=100, continue_from=token + ) + body_text, has_more, token = parse_paged_result(result) + collected += body_text + if not has_more: + break + assert collected == "".join(pages) + + +# --------------------------------------------------------------------------- +# search_and_fetch tests (pagination) +# --------------------------------------------------------------------------- + + +def test_search_and_fetch_no_results(monkeypatch): + """No matching candidates returns the no_results prompt (regression).""" + from ollama_web.tools import fetch as fetch_mod + + monkeypatch.setattr(fetch_mod, "web_search", lambda q, max_results=None: []) + result = fetch_mod.search_and_fetch("query", "pdf") + assert "No pdf files found" in result + + +def test_search_and_fetch_continue_from_invalid_token(): + """Invalid continuation token returns an error.""" + from ollama_web.tools.fetch import search_and_fetch + + result = search_and_fetch("query", "pdf", continue_from="bad-token") + assert "Invalid continuation token" 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 + + matched = [ + {"title": "File1", "href": "https://example.com/f1.txt", "body": ""}, + {"title": "File2", "href": "https://example.com/f2.txt", "body": ""}, + ] + monkeypatch.setattr(fetch_mod, "web_search", lambda q, max_results=None: matched) + # fetch_public_bytes returns distinct content per URL. + contents = { + "https://example.com/f1.txt": b"FILE1 " + b"A" * 200, + "https://example.com/f2.txt": b"FILE2 " + b"B" * 200, + } + monkeypatch.setattr( + fetch_mod, + "fetch_public_bytes", + lambda url, timeout=None: (contents[url], "", "text/plain"), + ) + + seen_titles = [] + token = None + for _ in range(20): + result = fetch_mod.search_and_fetch("query", "txt", max_chars=50, continue_from=token) + body_text, has_more, token = parse_paged_result(result) + # Record the first time each file title appears, preserving order. + if "File1" in body_text and "File1" not in seen_titles: + seen_titles.append("File1") + if "File2" in body_text and "File2" not in seen_titles: + seen_titles.append("File2") + if not has_more: + break + # Both files must be visited, in order. + assert seen_titles == ["File1", "File2"] + + +# --------------------------------------------------------------------------- +# pagination helper tests +# --------------------------------------------------------------------------- + + +def test_make_token_and_parse_token_roundtrip(): + token = make_token("scrape_url", "https://example.com/page", 4000) + payload = parse_token(token, "scrape_url") + assert payload["url"] == "https://example.com/page" + assert payload["offset"] == 4000 + assert payload["tool"] == "scrape_url" + assert payload["v"] == 1 + + +def test_parse_token_invalid_base64(): + with pytest.raises(InvalidTokenError): + parse_token("!!!not-base64!!!", "scrape_url") + + +def test_parse_token_invalid_json(): + import base64 + + token = base64.urlsafe_b64encode(b"{not json").decode("ascii").rstrip("=") + with pytest.raises(InvalidTokenError): + parse_token(token, "scrape_url") + + +def test_parse_token_wrong_tool(): + token = make_token("scrape_url", "https://example.com/page", 0) + with pytest.raises(InvalidTokenError): + parse_token(token, "pdf_to_text") + + +def test_parse_token_missing_fields(): + import base64 + + payload = json.dumps({"v": 1, "tool": "scrape_url"}).encode("utf-8") + token = base64.urlsafe_b64encode(payload).decode("ascii").rstrip("=") + with pytest.raises(InvalidTokenError): + parse_token(token, "scrape_url") + + +def test_parse_token_version_mismatch(): + import base64 + + payload = json.dumps({"v": 999, "tool": "scrape_url", "url": "https://x", "offset": 0}).encode( + "utf-8" + ) + token = base64.urlsafe_b64encode(payload).decode("ascii").rstrip("=") + with pytest.raises(InvalidTokenError): + parse_token(token, "scrape_url") + + +def test_parse_token_empty_or_none(): + with pytest.raises(InvalidTokenError): + parse_token("", "scrape_url") + with pytest.raises(InvalidTokenError): + parse_token(None, "scrape_url") + + +# --------------------------------------------------------------------------- +# format_paged_result tests +# --------------------------------------------------------------------------- + + +def test_format_paged_result_has_more_true(): + result = format_paged_result("body", has_more=True, next_token="tok") + assert "---page-info---" in result + assert "has_more: true" in result + assert "next_token: tok" in result + + +def test_format_paged_result_has_more_false(): + result = format_paged_result("body", has_more=False, next_token=None) + assert "---page-info---" not in result + assert "next_token" not in result + + +def test_format_paged_result_marker_at_end(): + result = format_paged_result("body", has_more=True, next_token="tok") + # Marker must come after the body. + assert result.index("body") < result.index("---page-info---") + + +# --------------------------------------------------------------------------- +# _trim_tool_result page-info protection tests +# --------------------------------------------------------------------------- + + +def test_trim_tool_result_preserves_page_info(): + 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) + + +def test_trim_tool_result_no_page_info_unchanged(): + """Without a footer, legacy truncation applies (regression).""" + from ollama_web.llm import _trim_tool_result + + result = "A" * 5000 + trimmed = _trim_tool_result(result, 200) + assert "---page-info---" not in trimmed + assert "…[tool result truncated]" in trimmed + + +def test_trim_tool_result_page_info_fits_limit(): + """When the whole result fits the limit, nothing is truncated.""" + from ollama_web.llm import _trim_tool_result + + result = "short body\n---page-info---\nhas_more: true\nnext_token: tok" + trimmed = _trim_tool_result(result, 200) + assert trimmed == result |