From 515f96a7779f117a7e37f854de2ded99160d0c76 Mon Sep 17 00:00:00 2001 From: tanaka takashi Date: Sun, 2 Aug 2026 15:50:41 +0900 Subject: scrape_url / pdf_to_text / search_and_fetch はページング（継続トークン方式）に対応 MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit mypy 6エラー修正 --- src/ollama_web/__init__.py | 2 +- src/ollama_web/config.py | 15 +-- src/ollama_web/i18n/__init__.py | 6 +- src/ollama_web/llm.py | 55 +++++--- src/ollama_web/mcp.py | 15 +-- src/ollama_web/prompts/__init__.py | 6 +- src/ollama_web/prompts/en.json | 4 +- src/ollama_web/prompts/ja.json | 4 +- src/ollama_web/routes/chat.py | 19 ++- src/ollama_web/routes/mcp.py | 4 +- src/ollama_web/routes/models.py | 2 +- src/ollama_web/routes/sessions.py | 5 +- src/ollama_web/sessions.py | 2 +- src/ollama_web/settings_store.py | 4 +- src/ollama_web/tools/fetch.py | 212 ++++++++++++++++++++++++------ src/ollama_web/tools/helper/pagination.py | 177 +++++++++++++++++++++++++ src/ollama_web/tools/helper/pdf.py | 65 +++++++-- src/ollama_web/tools/pdf.py | 39 +++++- src/ollama_web/tools/scrape.py | 50 +++++-- src/ollama_web/tools/search.py | 4 +- 20 files changed, 563 insertions(+), 127 deletions(-) create mode 100644 src/ollama_web/tools/helper/pagination.py (limited to 'src') 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\n{body}\n" - 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"\n{text}\n" + f'\n{text}\n' ) 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]] = [] -- cgit v1.2.3