mirror of
https://github.com/VectifyAI/PageIndex.git
synced 2026-10-02 07:44:37 +08:00
fix: indexing failures surface loud, chat envelopes append verbatim (#421)
Indexing: dead credentials or a missing model fail the run instead of storing a document with blank summaries; a 400 (context_length_exceeded) skips the retry ladder — the prompt will not shrink — and stays a per-prompt failure the run absorbs; all-empty model replies can no longer store a retrieval-ready document; the one-sentence doc description absorbs its own context overflow instead of discarding a fully indexed document; the heading-less flash refusal points at mode='standard'. Chat: messages() output is append-verbatim clean — unset response-only defaults are dropped (no "caller": null the request schema rejects); Claude cache marks follow the wire routing; model_settings and name are openai_agent_config parameters; one Anthropic client per backend; lifted thinking defaults are clamped to the model's output ceiling from LiteLLM's capability map. Store and inputs: lone surrogates are scrubbed from page text and the stored basename, so the returned name is byte-for-byte the stored name and the rename warning fires; NaN/Infinity metadata is rejected at the gate; every cloud error now carries its HTTP status. CLI: the flash lane resolves the summary model through ConfigLoader like the standard and markdown lanes; an empty flash structure errors like the SDK instead of writing "structure": [] with exit 0; --summary-model reaches the markdown lane; the SDK page-spec surface keeps 0.2.10's whitespace tolerance while the tool layer stays strict. pypdfium2 stays on the 5.x line for every install; the 4.x code paths are tested compatibility insurance with their own CI leg; process-pool construction failure falls back to the sequential parse; a py3.10 GC flake in text extraction is fixed. Port of feat/local-chat 0667e3b..1993740 (28 commits); README and assets untouched.
This commit is contained in:
Ray
@@ -33,7 +33,13 @@ jobs:
|
||||
fail-fast: false
|
||||
matrix:
|
||||
python-version: ["3.10", "3.13"]
|
||||
name: test py${{ matrix.python-version }}
|
||||
agent-frameworks: [without, with]
|
||||
pdfium: ["5"]
|
||||
include:
|
||||
- python-version: "3.10"
|
||||
agent-frameworks: with
|
||||
pdfium: "4"
|
||||
name: test py${{ matrix.python-version }} (${{ matrix.agent-frameworks }} frameworks, pdfium ${{ matrix.pdfium }})
|
||||
timeout-minutes: 15
|
||||
steps:
|
||||
- uses: actions/checkout@9c091bb21b7c1c1d1991bb908d89e4e9dddfe3e0 # v7.0.0
|
||||
@@ -42,6 +48,14 @@ jobs:
|
||||
python-version: ${{ matrix.python-version }}
|
||||
cache: pip
|
||||
- run: pip install -r requirements.txt pytest
|
||||
# mirror tests.yml's install shapes: the gate must cover the framework
|
||||
# tests, the no-framework import paths, and the pdfium 4.x insurance alike
|
||||
- if: matrix.agent-frameworks == 'without'
|
||||
run: pip uninstall -y openai-agents
|
||||
- if: matrix.agent-frameworks == 'with'
|
||||
run: pip install openai-agents claude-agent-sdk anthropic
|
||||
- if: matrix.pdfium == '4'
|
||||
run: pip install "pypdfium2<5"
|
||||
- run: python -m pytest -q
|
||||
|
||||
publish:
|
||||
|
||||
@@ -21,7 +21,13 @@ jobs:
|
||||
matrix:
|
||||
python-version: ["3.10", "3.13"]
|
||||
agent-frameworks: [without, with]
|
||||
name: py${{ matrix.python-version }} (${{ matrix.agent-frameworks }} frameworks)
|
||||
pdfium: ["5"]
|
||||
include:
|
||||
# one leg holds the 4.x insurance line
|
||||
- python-version: "3.10"
|
||||
agent-frameworks: with
|
||||
pdfium: "4"
|
||||
name: py${{ matrix.python-version }} (${{ matrix.agent-frameworks }} frameworks, pdfium ${{ matrix.pdfium }})
|
||||
timeout-minutes: 15
|
||||
steps:
|
||||
- uses: actions/checkout@9c091bb21b7c1c1d1991bb908d89e4e9dddfe3e0 # v7.0.0
|
||||
@@ -35,4 +41,6 @@ jobs:
|
||||
run: pip uninstall -y openai-agents
|
||||
- if: matrix.agent-frameworks == 'with'
|
||||
run: pip install openai-agents claude-agent-sdk anthropic
|
||||
- if: matrix.pdfium == '4'
|
||||
run: pip install "pypdfium2<5"
|
||||
- run: python -m pytest -q
|
||||
|
||||
@@ -49,8 +49,10 @@ def query_agent(client: PageIndexLocalClient, doc_id: str, prompt: str, verbose:
|
||||
Tool calls are always printed; verbose=True also prints arguments and output previews.
|
||||
"""
|
||||
agent = Agent(
|
||||
**client.openai_agent_config(doc_id=doc_id),
|
||||
# model_settings=ModelSettings(reasoning={"effort": "low", "summary": "auto"}), # from agents.model_settings import ModelSettings
|
||||
**client.openai_agent_config(
|
||||
doc_id=doc_id,
|
||||
# model_settings=ModelSettings(reasoning={"effort": "low", "summary": "auto"}), # from agents.model_settings import ModelSettings
|
||||
),
|
||||
)
|
||||
|
||||
async def _run():
|
||||
|
||||
+31
-34
@@ -12,7 +12,10 @@ exist on the cloud.
|
||||
Tools never raise: every outcome, including errors, is returned as the
|
||||
same JSON envelope the cloud emits ({"success": true, ...} /
|
||||
{"error": ...}) — arguments outside a pruned local signature come back as
|
||||
that envelope too, on the direct and the call_tool path alike.
|
||||
that envelope too, on the direct and the call_tool path alike, except
|
||||
browse_documents' ``recursive``: call_tool honors it, because the flat
|
||||
no-folders shape it asks for is trivially true here. One exception to
|
||||
never-raise: a cloud 401/403 re-raises PageIndexAPIError.
|
||||
"""
|
||||
from __future__ import annotations
|
||||
|
||||
@@ -429,7 +432,9 @@ def _refetch_entry(client, doc_id: str) -> Optional[dict[str, Any]]:
|
||||
|
||||
|
||||
def _await_completion(client, entry: dict[str, Any], wait: bool) -> dict[str, Any]:
|
||||
"""Re-poll a processing document for up to 3 minutes when wait is set."""
|
||||
"""Re-poll a processing document for up to 3 minutes when wait is set.
|
||||
Local documents are stored already terminal, so the wait never engages
|
||||
there."""
|
||||
doc_id = entry.get("id")
|
||||
if not wait or not doc_id or entry.get("status") in ("completed", "failed"):
|
||||
return entry
|
||||
@@ -517,13 +522,23 @@ class _PageSpecError(ValueError):
|
||||
self.code = code
|
||||
|
||||
|
||||
# re.ASCII mirrors the ECMA regex semantics the contract's pattern carries
|
||||
_PAGE_SPEC_RE = re.compile(
|
||||
TOOL_CONTRACT["get_page_content"]["schema"]["properties"]["pages"]["pattern"],
|
||||
re.ASCII)
|
||||
|
||||
|
||||
def _expand_pages(pages) -> list[int]:
|
||||
"""Expand '1-3,7' into sorted distinct pages — the one parser for the
|
||||
SDK surface and the tool layer. Raises _PageSpecError with code
|
||||
'invalid', 'too_many', or 'nonpositive'."""
|
||||
SDK surface and the tool layer, holding both to the contract's pages
|
||||
pattern. Raises _PageSpecError with code 'invalid', 'too_many', or
|
||||
'nonpositive'."""
|
||||
if not isinstance(pages, str):
|
||||
raise _PageSpecError("invalid",
|
||||
f"Invalid page specification: {pages!r}")
|
||||
if not _PAGE_SPEC_RE.fullmatch(pages):
|
||||
raise _PageSpecError(
|
||||
"invalid", f"Invalid page specification '{pages}'")
|
||||
too_many = (f"Page specification '{pages}' spans more than "
|
||||
f"{_MAX_REQUESTED_PAGES} pages; request a narrower range")
|
||||
expanded: set[int] = set()
|
||||
@@ -1461,31 +1476,19 @@ def _cloud_bridge(client, gated: bool = False):
|
||||
return bridge
|
||||
|
||||
|
||||
def _read_only_tools(tools_meta: list[dict]) -> list[dict]:
|
||||
"""The management gate for consumers without a framework permission
|
||||
layer: only tools the server marks read-only, guarded against a server
|
||||
annotation regression silently disabling every tool."""
|
||||
filtered = [meta for meta in tools_meta
|
||||
if (meta.get("annotations") or {}).get("readOnlyHint") is True]
|
||||
if tools_meta and not filtered:
|
||||
raise PageIndexAPIError(
|
||||
"The MCP server returned tools but none are annotated "
|
||||
"read-only — a server annotation regression would otherwise "
|
||||
"silently disable every tool. Pass include_management=True "
|
||||
"to expose the unfiltered list."
|
||||
)
|
||||
return filtered
|
||||
|
||||
|
||||
def _require_local_scope(client, doc_ids) -> None:
|
||||
"""The allowlist is enforced in-process; cloud lookups run server-side,
|
||||
so accepting doc_ids there would be advisory-only — refuse loudly.
|
||||
An empty allowlist is refused too: it would scope the agent to
|
||||
nothing, with no signal to the caller."""
|
||||
def _require_doc_selection(doc_ids) -> None:
|
||||
"""An empty selection fails loud: washed to None it would mean
|
||||
"everything", while the tool-layer allowlist would mean "nothing"."""
|
||||
if doc_ids is not None and not doc_ids:
|
||||
raise PageIndexAPIError(
|
||||
"doc_id is empty. Pass one or more document IDs, or omit "
|
||||
"doc_id to give the agent the whole library.")
|
||||
|
||||
|
||||
def _require_local_scope(client, doc_ids) -> None:
|
||||
"""The allowlist is enforced in-process; cloud lookups run server-side,
|
||||
so accepting doc_ids there would be advisory-only — refuse loudly."""
|
||||
_require_doc_selection(doc_ids)
|
||||
if doc_ids is not None and getattr(client, "api_key", None):
|
||||
raise PageIndexAPIError(
|
||||
"doc_ids scoping applies to local tools only — cloud calls "
|
||||
@@ -1503,8 +1506,6 @@ def _tool_specs(client, include_management: bool = False, doc_ids=None,
|
||||
if getattr(client, "api_key", None):
|
||||
bridge = _cloud_bridge(client, gated=not include_management)
|
||||
tools_meta = bridge.list_tools()
|
||||
if not include_management:
|
||||
tools_meta = _read_only_tools(tools_meta)
|
||||
return [(str(meta.get("name") or "tool"),
|
||||
meta.get("description") or "",
|
||||
copy.deepcopy(meta.get("inputSchema"))
|
||||
@@ -1531,7 +1532,8 @@ def build_agent_tools(client, include_management: bool = False,
|
||||
synthesized from the server's schemas, calls proxied over MCP. Local:
|
||||
the built-in contract tools over the local store. Every function returns
|
||||
the JSON envelope as a string and never raises for arguments its
|
||||
signature accepts (cloud-only parameters are absent from the local
|
||||
signature accepts — except a cloud 401/403, which re-raises
|
||||
PageIndexAPIError (cloud-only parameters are absent from the local
|
||||
signatures; the call_tool path answers them with the guided envelope).
|
||||
``doc_ids`` is the local allowlist, as in ``_tool_specs``.
|
||||
"""
|
||||
@@ -1619,12 +1621,7 @@ def doc_targeting_block(client, doc_id, scoped: bool = False) -> Optional[str]:
|
||||
if doc_id is None:
|
||||
return None
|
||||
doc_ids = [doc_id] if isinstance(doc_id, str) else list(doc_id)
|
||||
if not doc_ids:
|
||||
# An empty selection must fail loud: washing it to None would mean
|
||||
# "everything", and the tool-layer allowlist would mean "nothing".
|
||||
raise PageIndexAPIError(
|
||||
"doc_id is empty. Pass one or more document IDs, or omit "
|
||||
"doc_id to give the agent the whole library.")
|
||||
_require_doc_selection(doc_ids)
|
||||
details = []
|
||||
missing = []
|
||||
for one_id in doc_ids:
|
||||
|
||||
+63
-33
@@ -2,6 +2,7 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import os
|
||||
import re
|
||||
import threading
|
||||
import time
|
||||
import warnings
|
||||
@@ -37,6 +38,10 @@ def _preload_litellm() -> None:
|
||||
|
||||
def _parse_pages(pages: str) -> list[int]:
|
||||
from .agent_tools import _PageSpecError, _expand_pages
|
||||
if isinstance(pages, str):
|
||||
# 0.2.10 tolerated whitespace on this surface; the tool layer stays
|
||||
# on the strict contract pattern.
|
||||
pages = re.sub(r"\s*([,-])\s*", r"\1", pages.strip())
|
||||
try:
|
||||
return _expand_pages(pages)
|
||||
except _PageSpecError as exc:
|
||||
@@ -253,6 +258,8 @@ class PageIndexClient:
|
||||
status = self.get_document(doc_id).get("status")
|
||||
poll_failures = 0
|
||||
except (PageIndexAPIError, requests.RequestException) as exc:
|
||||
if getattr(exc, "status_code", None) in (401, 403, 404):
|
||||
raise # a definite answer, not a poll failure
|
||||
# Tolerate transient poll failures; a 30-minute wait should
|
||||
# not die on one 502 or dropped connection.
|
||||
poll_failures += 1
|
||||
@@ -477,8 +484,9 @@ class PageIndexClient:
|
||||
prompt prefix cache-marked automatically. The non-stream
|
||||
response carries the final answer only; streaming yields the
|
||||
agent's visible text as it is produced, including narration before
|
||||
tool calls. ``finish_reason`` reports loop completion ("stop") —
|
||||
the engine does not surface per-turn backend finish reasons. For
|
||||
tool calls. ``finish_reason`` carries the final turn's native
|
||||
finish reason — "stop", or the backend's "length" /
|
||||
"content_filter" when the last turn was cut short. For
|
||||
the tool-use process and prompt-cache round-trip use
|
||||
``responses()`` or ``messages()``.
|
||||
|
||||
@@ -811,13 +819,15 @@ class PageIndexClient:
|
||||
``get_document_structure``, ``get_page_content``).
|
||||
|
||||
Each function takes JSON-serializable arguments, returns a JSON
|
||||
string, and reports failures inside that JSON instead of raising.
|
||||
string, and reports failures inside that JSON instead of raising —
|
||||
except a cloud 401/403, which raises PageIndexAPIError.
|
||||
|
||||
Args:
|
||||
include_management (bool): Also expose tools that modify the
|
||||
library. Local: adds ``remove_document``. Cloud: by default
|
||||
only tools the server marks read-only are exposed; True
|
||||
exposes the server's complete list (upload, delete, ...).
|
||||
library. Local: adds ``remove_document``. Cloud: the URL
|
||||
is the gate — the default serves what the read-only
|
||||
endpoint (``?tools=read``) registers; True connects to
|
||||
the full ``/mcp`` list (upload, delete, ...).
|
||||
doc_id: Local only — restrict the tools to this document ID
|
||||
(or list of IDs), enforced at the tool layer: out-of-scope
|
||||
lookups return NOT_FOUND. Raises on cloud, where scoping
|
||||
@@ -856,10 +866,10 @@ class PageIndexClient:
|
||||
|
||||
Args:
|
||||
include_management (bool): Also expose tools that modify the
|
||||
library (delete, upload). Default off: the in-process
|
||||
cloud default serves only server-annotated read-only
|
||||
tools, and ``hosted=True`` connects OpenAI to the
|
||||
read-only endpoint (``/mcp?tools=read``) instead.
|
||||
library (delete, upload). Default off: on cloud the URL
|
||||
is the gate — in-process and ``hosted=True`` alike
|
||||
connect to the read-only endpoint (``/mcp?tools=read``);
|
||||
True switches to the full ``/mcp`` list.
|
||||
hosted (bool): Cloud only — hand the MCP connection to OpenAI
|
||||
for server-side tool execution (OpenAI models only).
|
||||
doc_id: Local only — restrict the tools to this document ID
|
||||
@@ -875,12 +885,8 @@ class PageIndexClient:
|
||||
"""doc_id for the tool layer: passed through locally (structural
|
||||
allowlist), dropped on cloud where scoping is server-side and the
|
||||
config helpers keep prompt-level targeting."""
|
||||
if doc_id is not None and not doc_id:
|
||||
# An empty scope means "nothing" locally (empty allowlist) and
|
||||
# cannot be represented on cloud; both refuse it loudly.
|
||||
raise PageIndexAPIError(
|
||||
"doc_id is empty. Pass one or more document IDs, or omit "
|
||||
"doc_id to give the agent the whole library.")
|
||||
from .agent_tools import _require_doc_selection
|
||||
_require_doc_selection(doc_id)
|
||||
if not getattr(self, "api_key", None):
|
||||
return doc_id
|
||||
return None
|
||||
@@ -890,6 +896,8 @@ class PageIndexClient:
|
||||
doc_id: Optional[Union[str, list[str]]] = None,
|
||||
include_management: bool = False,
|
||||
model: Optional[str] = None,
|
||||
model_settings: Optional[Any] = None,
|
||||
name: str = "PageIndex",
|
||||
) -> dict[str, Any]:
|
||||
"""
|
||||
Document QA ``Agent`` kwargs for the OpenAI Agents SDK in one
|
||||
@@ -906,15 +914,12 @@ class PageIndexClient:
|
||||
environment, so its model auth comes from there —
|
||||
``chat_backend`` does not travel with it.
|
||||
|
||||
Prompt caching configures itself for most destinations (OpenAI
|
||||
server-side; Anthropic- and Bedrock-hosted Claude via LiteLLM's
|
||||
defaults). Vertex-hosted Claude is the exception — pass the
|
||||
injection points yourself::
|
||||
|
||||
Agent(**config, model_settings=ModelSettings(extra_args={
|
||||
"cache_control_injection_points": [
|
||||
{"location": "message", "role": "system"},
|
||||
{"location": "message", "index": -1}]}))
|
||||
Prompt caching: OpenAI models cache server-side on their own;
|
||||
LiteLLM-routed Claude (Anthropic, Bedrock, Vertex) gets its
|
||||
cache marks from the bundled ``model_settings``. Pass
|
||||
``model_settings`` here to layer your own on top — your fields
|
||||
win and ``extra_args`` merge. Replacing the returned key
|
||||
wholesale drops the marks instead.
|
||||
|
||||
Args:
|
||||
doc_id: Document ID or list of IDs to target, as in
|
||||
@@ -926,11 +931,16 @@ class PageIndexClient:
|
||||
model: Backend model name; overrides the local default. Same
|
||||
grammar as ``chat_model`` (LiteLLM names; bare names are
|
||||
OpenAI-compatible shorthand).
|
||||
model_settings: Your own ``ModelSettings``, merged on top of
|
||||
the bundled cache marks; included verbatim when no marks
|
||||
apply.
|
||||
name (str): Agent display name; in composition it also seeds
|
||||
the SDK-derived handoff and ``as_tool`` names.
|
||||
"""
|
||||
from .agent_tools import build_agent_instructions
|
||||
scope = self._local_doc_scope(doc_id)
|
||||
config: dict[str, Any] = {
|
||||
"name": "PageIndex",
|
||||
"name": name,
|
||||
"instructions": build_agent_instructions(
|
||||
self, doc_id, scoped=scope is not None,
|
||||
include_management=include_management),
|
||||
@@ -944,6 +954,20 @@ class PageIndexClient:
|
||||
# caller's process, outside our completion helpers.
|
||||
from .utils import _repair_litellm_types
|
||||
_repair_litellm_types()
|
||||
# Marks follow this lane's routing: the SDK strips litellm/
|
||||
# and LiteLLM resolves the rest (bare claude-* → Anthropic);
|
||||
# names without the prefix ride the SDK's OpenAI provider.
|
||||
from .local_chat import _litellm_claude_marks
|
||||
extra_args = _litellm_claude_marks(
|
||||
config["model"].removeprefix("litellm/"))
|
||||
if extra_args:
|
||||
from agents import ModelSettings
|
||||
config["model_settings"] = ModelSettings(
|
||||
extra_args=extra_args)
|
||||
if model_settings is not None:
|
||||
marks = config.get("model_settings")
|
||||
config["model_settings"] = (marks.resolve(model_settings)
|
||||
if marks else model_settings)
|
||||
return config
|
||||
|
||||
def as_anthropic_tools(self, include_management: bool = False,
|
||||
@@ -979,9 +1003,10 @@ class PageIndexClient:
|
||||
|
||||
Args:
|
||||
include_management (bool): Also expose tools that modify the
|
||||
library. Local: adds ``remove_document``. Cloud: by default
|
||||
only tools the server marks read-only are exposed; True
|
||||
exposes the server's complete list (upload, delete, ...).
|
||||
library. Local: adds ``remove_document``. Cloud: the URL
|
||||
is the gate — the default serves what the read-only
|
||||
endpoint (``?tools=read``) registers; True connects to
|
||||
the full ``/mcp`` list (upload, delete, ...).
|
||||
asynchronous (bool): Build ``beta_async_tool`` runnables for
|
||||
``AsyncAnthropic`` (each tool call runs in a worker
|
||||
thread, keeping blocking I/O off your event loop). The
|
||||
@@ -1063,7 +1088,8 @@ class PageIndexClient:
|
||||
}
|
||||
|
||||
def as_claude_mcp(self, include_management: bool = False,
|
||||
doc_id: Optional[Union[str, list[str]]] = None):
|
||||
doc_id: Optional[Union[str, list[str]]] = None,
|
||||
server_name: str = "pageindex"):
|
||||
"""
|
||||
``mcp_servers`` entry for the Claude Agent SDK.
|
||||
|
||||
@@ -1077,6 +1103,8 @@ class PageIndexClient:
|
||||
``pip install 'pageindex[claude]'``). ``doc_id`` (local only)
|
||||
restricts those tools to that document ID (or list), enforced at
|
||||
the tool layer; it raises on cloud, where scoping is server-side.
|
||||
``server_name`` names the in-process server — match it to the key
|
||||
you register the entry under (cloud entries carry no name).
|
||||
|
||||
Cloud hosts that surface MCP server instructions receive the same
|
||||
guidance ``agent_instructions()`` returns natively — passing both
|
||||
@@ -1095,7 +1123,8 @@ class PageIndexClient:
|
||||
)
|
||||
"""
|
||||
from .integrations.claude_agent_sdk import build_claude_mcp
|
||||
return build_claude_mcp(self, include_management, doc_ids=doc_id)
|
||||
return build_claude_mcp(self, include_management, doc_ids=doc_id,
|
||||
server_name=server_name)
|
||||
|
||||
def claude_agent_config(
|
||||
self,
|
||||
@@ -1122,7 +1151,8 @@ class PageIndexClient:
|
||||
(tool scoping is server-side).
|
||||
include_management (bool): Also allow tools that modify the
|
||||
library.
|
||||
server_name (str): Key the server is registered under.
|
||||
server_name (str): Key the server is registered under;
|
||||
locally also the name the SDK server declares.
|
||||
"""
|
||||
from .agent_tools import build_agent_instructions
|
||||
scope = self._local_doc_scope(doc_id)
|
||||
@@ -1131,7 +1161,7 @@ class PageIndexClient:
|
||||
self, doc_id, scoped=scope is not None,
|
||||
include_management=include_management),
|
||||
"mcp_servers": {server_name: self.as_claude_mcp(
|
||||
include_management, doc_id=scope)},
|
||||
include_management, doc_id=scope, server_name=server_name)},
|
||||
# Pre-approval only — the server itself is already gated (the
|
||||
# read-only endpoint on cloud, the registered set locally).
|
||||
"allowed_tools": [f"mcp__{server_name}"],
|
||||
|
||||
+30
-10
@@ -78,7 +78,9 @@ class CloudAPI:
|
||||
)
|
||||
|
||||
if response.status_code != 200:
|
||||
raise PageIndexAPIError(f"Failed to submit document: {response.text}")
|
||||
raise PageIndexAPIError(
|
||||
f"Failed to submit document: {response.text}",
|
||||
status_code=response.status_code)
|
||||
return response.json()
|
||||
|
||||
# ---------- OCR FUNCTIONALITY ----------
|
||||
@@ -103,7 +105,9 @@ class CloudAPI:
|
||||
timeout=30
|
||||
)
|
||||
if response.status_code != 200:
|
||||
raise PageIndexAPIError(f"Failed to get OCR result: {response.text}")
|
||||
raise PageIndexAPIError(
|
||||
f"Failed to get OCR result: {response.text}",
|
||||
status_code=response.status_code)
|
||||
return response.json()
|
||||
|
||||
# ---------- TREE GENERATION ----------
|
||||
@@ -127,7 +131,9 @@ class CloudAPI:
|
||||
timeout=30
|
||||
)
|
||||
if response.status_code != 200:
|
||||
raise PageIndexAPIError(f"Failed to get tree result: {response.text}")
|
||||
raise PageIndexAPIError(
|
||||
f"Failed to get tree result: {response.text}",
|
||||
status_code=response.status_code)
|
||||
return response.json()
|
||||
|
||||
# ---------- RETRIEVAL ----------
|
||||
@@ -156,7 +162,9 @@ class CloudAPI:
|
||||
timeout=30
|
||||
)
|
||||
if response.status_code != 200:
|
||||
raise PageIndexAPIError(f"Failed to submit retrieval: {response.text}")
|
||||
raise PageIndexAPIError(
|
||||
f"Failed to submit retrieval: {response.text}",
|
||||
status_code=response.status_code)
|
||||
return response.json()
|
||||
|
||||
def get_retrieval(self, retrieval_id: str) -> Dict[str, Any]:
|
||||
@@ -175,7 +183,9 @@ class CloudAPI:
|
||||
timeout=30
|
||||
)
|
||||
if response.status_code != 200:
|
||||
raise PageIndexAPIError(f"Failed to get retrieval result: {response.text}")
|
||||
raise PageIndexAPIError(
|
||||
f"Failed to get retrieval result: {response.text}",
|
||||
status_code=response.status_code)
|
||||
return response.json()
|
||||
|
||||
# ---------- CHAT COMPLETIONS ----------
|
||||
@@ -229,7 +239,9 @@ class CloudAPI:
|
||||
)
|
||||
|
||||
if response.status_code != 200:
|
||||
raise PageIndexAPIError(f"Failed to get chat completion: {response.text}")
|
||||
raise PageIndexAPIError(
|
||||
f"Failed to get chat completion: {response.text}",
|
||||
status_code=response.status_code)
|
||||
|
||||
if stream:
|
||||
if stream_metadata:
|
||||
@@ -333,7 +345,9 @@ class CloudAPI:
|
||||
timeout=30
|
||||
)
|
||||
if response.status_code != 200:
|
||||
raise PageIndexAPIError(f"Failed to delete document: {response.text}")
|
||||
raise PageIndexAPIError(
|
||||
f"Failed to delete document: {response.text}",
|
||||
status_code=response.status_code)
|
||||
return response.json() if response.content else {}
|
||||
|
||||
def list_documents(self, limit: int = 50, offset: int = 0, folder_id: Optional[str] = None) -> Dict[str, Any]:
|
||||
@@ -369,7 +383,9 @@ class CloudAPI:
|
||||
timeout=30
|
||||
)
|
||||
if response.status_code != 200:
|
||||
raise PageIndexAPIError(f"Failed to list documents: {response.text}")
|
||||
raise PageIndexAPIError(
|
||||
f"Failed to list documents: {response.text}",
|
||||
status_code=response.status_code)
|
||||
return response.json()
|
||||
|
||||
# ---------- FOLDER MANAGEMENT ----------
|
||||
@@ -406,7 +422,9 @@ class CloudAPI:
|
||||
timeout=30
|
||||
)
|
||||
if response.status_code != 200:
|
||||
raise PageIndexAPIError(f"Failed to create folder: {response.text}")
|
||||
raise PageIndexAPIError(
|
||||
f"Failed to create folder: {response.text}",
|
||||
status_code=response.status_code)
|
||||
return response.json()
|
||||
|
||||
def list_folders(self, parent_folder_id: Optional[str] = None) -> Dict[str, Any]:
|
||||
@@ -433,5 +451,7 @@ class CloudAPI:
|
||||
timeout=30
|
||||
)
|
||||
if response.status_code != 200:
|
||||
raise PageIndexAPIError(f"Failed to list folders: {response.text}")
|
||||
raise PageIndexAPIError(
|
||||
f"Failed to list folders: {response.text}",
|
||||
status_code=response.status_code)
|
||||
return response.json()
|
||||
|
||||
+2
-2
@@ -1,6 +1,6 @@
|
||||
class PageIndexAPIError(Exception):
|
||||
"""status_code carries the HTTP status when the failure came from a
|
||||
non-200 cloud response; None otherwise (local mode, client-side)."""
|
||||
"""status_code carries the HTTP status when the raising site passes it;
|
||||
None does not imply local/client-side."""
|
||||
|
||||
def __init__(self, *args: object, status_code: int | None = None) -> None:
|
||||
super().__init__(*args)
|
||||
|
||||
+10
-14
@@ -99,12 +99,15 @@ def page_index_flash(pdf, summary=True, summary_model=None,
|
||||
optimize: str | bool | None = None, optimize_expand=None,
|
||||
optimize_model=None, summary_concurrency=None,
|
||||
use_embedded_toc=True) -> dict:
|
||||
"""Build a PageIndex tree structure from a PDF using layout statistics. The tree extraction itself uses no LLM; by default an LLM writes node summaries and expands the tree (``summary=False, optimize=False`` runs fully LLM-free). Args: pdf: path to a PDF file (``str`` or ``pathlib.Path``) or an in-memory binary stream (``io.BytesIO``). summary: if True, generate LLM summaries for each node (requires ``summary_model``). summary_model: the LLM model identifier to use for summary generation. optimize: ``"full"`` for merge + LLM expand (fails fast with ``PageIndexAPIError`` when no LLM key is configured), ``"merge"`` for deterministic merge only, ``False`` to disable. ``True`` is accepted as ``"full"`` for backward compatibility; defaults to ``"full"``. optimize_expand: deprecated — use ``optimize``. Honored only when ``optimize`` is not passed (or is the legacy ``True``): ``False`` maps to ``"merge"``, ``True`` to ``"full"``. optimize_model: the LLM model for expand (defaults to the summary model). summary_concurrency: maximum simultaneous summary model calls; None uses the library default. use_embedded_toc: if True, consume the PDF's embedded bookmarks when trustworthy: deep bookmarks become the frame and the detected sections they lack are grafted back in after noise filtering, coarse ones become the chapter frame with detected nodes re-hung under them (deeper sparse entries are filled in when the page text confirms them, and garbled extracted titles are repaired from the bookmark strings), garbage ones are ignored; adds a ``toc_source`` key to the result. On by default; pass False for the pure detected structure. Returns: dict with keys ``doc_name``, ``doc_title``, ``structure`` (a list of nested ``{"title", "start_index", "end_index", "nodes"}`` dicts; page indexes are 1-based) and ``has_abstract_or_references_section`` (True when a top-level entry is an abstract or references heading). With ``optimize`` an ``optimize`` key reports merge/expand counts and before/after search-cost metrics. """
|
||||
"""Build a PageIndex tree structure from a PDF using layout statistics. The tree extraction itself uses no LLM; by default an LLM writes node summaries and expands the tree (``summary=False, optimize=False`` runs fully LLM-free). Args: pdf: path to a PDF file (``str`` or ``pathlib.Path``) or an in-memory binary stream (``io.BytesIO``). summary: if True, generate LLM summaries for each node (requires ``summary_model``). summary_model: the LLM model identifier to use for summary generation. optimize: ``"full"`` for merge + LLM expand (a model unreachable after the retry ladder — a missing credential included — fails the run loudly from expand itself; a per-prompt rejection leaves just that node collapsed), ``"merge"`` for deterministic merge only, ``False`` to disable. ``True`` is accepted as ``"full"`` for backward compatibility; defaults to ``"full"``. Expand needs readable page text, so a bookmark-only or scanned PDF runs the merge half only (``expands`` reports 0). optimize_expand: deprecated — use ``optimize``. Honored only when ``optimize`` is not passed (or is the legacy ``True``): ``False`` maps to ``"merge"``, ``True`` to ``"full"``. optimize_model: the LLM model for expand (defaults to the summary model). summary_concurrency: maximum simultaneous summary model calls; None uses the library default. use_embedded_toc: if True, consume the PDF's embedded bookmarks when trustworthy: deep bookmarks become the frame and the detected sections they lack are grafted back in after noise filtering, coarse ones become the chapter frame with detected nodes re-hung under them (deeper sparse entries are filled in when the page text confirms them, and garbled extracted titles are repaired from the bookmark strings), garbage ones are ignored; adds a ``toc_source`` key to the result. On by default; pass False for the pure detected structure. Returns: dict with keys ``doc_name``, ``doc_title``, ``structure`` (a list of nested ``{"title", "start_index", "end_index", "nodes"}`` dicts; page indexes are 1-based) and ``has_abstract_or_references_section`` (True when a top-level entry is an abstract or references heading). With ``optimize`` an ``optimize`` key reports merge/expand counts and before/after search-cost metrics. """
|
||||
if optimize_expand is not None:
|
||||
import warnings
|
||||
warnings.warn(
|
||||
"optimize_expand is deprecated: pass optimize='full', 'merge', "
|
||||
"or False.", DeprecationWarning, stacklevel=2)
|
||||
"or False. When optimize is not passed it maps onto it (False "
|
||||
"-> 'merge', True -> 'full'), so the optimize pass now runs "
|
||||
"where the old optimize=False default ran nothing.",
|
||||
DeprecationWarning, stacklevel=2)
|
||||
if optimize is None or optimize is True:
|
||||
# legacy spellings only — an explicit 'full'/'merge' wins
|
||||
optimize = "merge" if optimize_expand is False else "full"
|
||||
@@ -113,21 +116,14 @@ def page_index_flash(pdf, summary=True, summary_model=None,
|
||||
elif optimize not in ("full", "merge"):
|
||||
raise ValueError(
|
||||
f"optimize must be 'full', 'merge', or False, got {optimize!r}")
|
||||
if optimize == "full":
|
||||
from ..errors import PageIndexAPIError
|
||||
from ..utils import ConfigLoader, _llm_backend, _openai_missing_keys
|
||||
model = (optimize_model or summary_model
|
||||
or ConfigLoader().load().summary_model)
|
||||
if not _llm_backend.get() and _openai_missing_keys(model):
|
||||
raise PageIndexAPIError(
|
||||
"optimize='full' runs LLM expand and no LLM key is "
|
||||
"configured — set OPENAI_API_KEY, or pass optimize='merge' "
|
||||
"or optimize=False for the LLM-free tree.")
|
||||
result = extract_toc(_validate_pdf(pdf), use_embedded_toc=use_embedded_toc)
|
||||
structure = result.get("structure", [])
|
||||
if optimize and structure:
|
||||
result["optimize"] = _optimize(structure, result.get("page_texts") or [],
|
||||
optimize == "full",
|
||||
# bookmark-only extractions carry no page_texts and scanned ones
|
||||
# only empty strings; expand needs text
|
||||
pages = result.get("page_texts") or []
|
||||
result["optimize"] = _optimize(structure, pages,
|
||||
optimize == "full" and any(pages),
|
||||
optimize_model or summary_model)
|
||||
if summary and structure:
|
||||
import asyncio
|
||||
|
||||
@@ -171,16 +171,17 @@ def extract_toc(
|
||||
"doc_title": None,
|
||||
"structure": [],
|
||||
"has_abstract_or_references_section": False,
|
||||
# the summary/expand passes read these like on the normal path
|
||||
"page_texts": ["\n".join(block_text(block)
|
||||
for block in (page.secondary_slot or []))
|
||||
for page in pages],
|
||||
}
|
||||
# Bookmarks need no extracted text, so they can still structure a
|
||||
# document this gate wrote off as unreadable.
|
||||
if use_embedded_toc:
|
||||
from .embedded_toc import apply_embedded_toc
|
||||
result["structure"], result["toc_source"] = apply_embedded_toc(
|
||||
[], doc_handle, len(pages),
|
||||
page_texts=["\n".join(block_text(block)
|
||||
for block in (page.secondary_slot or []))
|
||||
for page in pages],
|
||||
[], doc_handle, len(pages), page_texts=result["page_texts"],
|
||||
)
|
||||
return result
|
||||
|
||||
|
||||
@@ -71,6 +71,8 @@ def _extract_raw_chars(page, text_page) -> tuple[list[dict], list[dict]]:
|
||||
if 0xDC00 <= low <= 0xDFFF:
|
||||
codepoint = ((codepoint & 0x3FF) << 10) + (low & 0x3FF) + 0x10000
|
||||
skip_next = True
|
||||
if 0xD800 <= codepoint <= 0xDFFF:
|
||||
codepoint = 0xFFFD # unpaired surrogate: not utf-8 encodable
|
||||
# u == 0 (PDFium found no unicode for the glyph) is KEPT as '\x00':
|
||||
# text extraction emits the raw charcode for unmapped codes, so its items
|
||||
# really contain chr(0) for extension-font pieces at code 0, and the
|
||||
@@ -79,9 +81,9 @@ def _extract_raw_chars(page, text_page) -> tuple[list[dict], list[dict]]:
|
||||
ch_str = chr(codepoint)
|
||||
is_ws = js_is_ws(codepoint)
|
||||
# FPDFText_IsGenerated returns a c_int: 1 generated, 0 real, -1 error.
|
||||
# Only a POSITIVE 1 may mark a char generated -- the -1 has to read the
|
||||
# same way here as it does in the page-mode unicode walk, or the two
|
||||
# char sets disagree and that walk desyncs.
|
||||
# Only a POSITIVE 1 may mark a char generated. This is the package's
|
||||
# only read: the page-mode unicode walk consumes this flag rather
|
||||
# than re-reading PDFium.
|
||||
is_gen = is_generated(text_page, index_value) == 1
|
||||
# PDFium inserts is_generated chars as layout placeholders for
|
||||
# Td/Tm jumps with no literal content-stream char (typically
|
||||
|
||||
@@ -4,6 +4,7 @@ from __future__ import annotations
|
||||
|
||||
import bisect
|
||||
import difflib
|
||||
import re
|
||||
from collections import Counter
|
||||
|
||||
from .text_normalize import _is_whitespace
|
||||
@@ -13,6 +14,8 @@ from .code_walk import (
|
||||
_walk_codes,
|
||||
)
|
||||
|
||||
_SURROGATES = re.compile("[\ud800-\udfff]")
|
||||
|
||||
|
||||
def _apply_font_unicode(
|
||||
raw_chars: list[dict],
|
||||
@@ -21,7 +24,7 @@ def _apply_font_unicode(
|
||||
pdf_doc,
|
||||
map_cache: dict,
|
||||
) -> None:
|
||||
"""Patch each char's unicode to span merger glyph Unicode (`map.get(code) or chr(code)`, content stream tokenizer glyph mapping) where PDFium's decode disagrees. Two granularities, both gated by _walk_codes' both-streams-exhaust rule: - object mode (when PDFium's text objects pair consistent with the page's show ops, the _assign_flush_ids precondition): each object's chars are walked against its own show op's codes. This is immune to PDFium's textpage segment reordering (e.g. math-heavy page margin labels emitted at a different page position than paint order) because chars keep stream order WITHIN an object; a desync rolls back only that object. - page mode (counts differ, e.g. PDFium splitting a TJ into several objects): all non-generated textpage chars are walked against all show ops' codes in paint order; any desync rolls back the whole page. """
|
||||
"""Patch each char's unicode to span merger glyph Unicode (`map.get(code) or chr(code)`, surrogate-band results replaced with U+FFFD, content stream tokenizer glyph mapping) where PDFium's decode disagrees. Two granularities, both gated by _walk_codes' both-streams-exhaust rule: - object mode (when PDFium's text objects pair consistent with the page's show ops, the _assign_flush_ids precondition): each object's chars are walked against its own show op's codes. This is immune to PDFium's textpage segment reordering (e.g. math-heavy page margin labels emitted at a different page position than paint order) because chars keep stream order WITHIN an object; a desync rolls back only that object. - page mode (counts differ, e.g. PDFium splitting a TJ into several objects): all non-generated chars from the extraction census (surrogate pairs already merged) are walked against all show ops' codes in paint order; any desync rolls back the whole page. """
|
||||
if not show_codes:
|
||||
return
|
||||
|
||||
@@ -37,9 +40,12 @@ def _apply_font_unicode(
|
||||
if entry is None:
|
||||
return None
|
||||
next_block, measure_item = entry
|
||||
# Broken font data (uniD83D glyph names, surrogate-band CIDs) yields
|
||||
# lone-surrogate targets; patched into chars they crash utf-8 saves.
|
||||
if next_block == 1:
|
||||
return [measure_item.get(code) or chr(code) for code in other_numbers]
|
||||
return [measure_item.get((other_numbers[key_value] << 8) | other_numbers[key_value + 1]) or chr((other_numbers[key_value] << 8) | other_numbers[key_value + 1])
|
||||
return [_SURROGATES.sub("\ufffd", measure_item.get(code) or chr(code))
|
||||
for code in other_numbers]
|
||||
return [_SURROGATES.sub("\ufffd", measure_item.get((other_numbers[key_value] << 8) | other_numbers[key_value + 1]) or chr((other_numbers[key_value] << 8) | other_numbers[key_value + 1]))
|
||||
for key_value in range(0, len(other_numbers) - 1, 2)]
|
||||
|
||||
def apply(patches: list[tuple[int, str]], drops: list[int],
|
||||
|
||||
@@ -12,7 +12,8 @@ Workers run pass 1 + pass 2 per page assuming the union stays empty and
|
||||
poison the run the moment any page accumulates an extent; the driver then
|
||||
discards the parallel attempt and reruns the document on the sequential
|
||||
path, which is the source of truth. Any other worker failure falls back the
|
||||
same way, so this entry can only ever return sequential-identical output.
|
||||
same way, so this entry returns sequential-identical output — except in a
|
||||
spawn child re-importing an unguarded __main__, where it re-raises.
|
||||
|
||||
Worker startup pays the full package import chain plus its own document
|
||||
open; ``min_pages`` routes documents too small to amortize that to the
|
||||
@@ -24,6 +25,7 @@ from __future__ import annotations
|
||||
import multiprocessing
|
||||
import os
|
||||
import sys
|
||||
import threading
|
||||
from concurrent.futures import ProcessPoolExecutor
|
||||
from contextlib import contextmanager
|
||||
from io import BytesIO
|
||||
@@ -56,27 +58,42 @@ _worker_pdf_doc = None
|
||||
_worker_font_maps: dict = {}
|
||||
|
||||
|
||||
_window_lock = threading.Lock()
|
||||
_window_depth = 0
|
||||
_window_saved: dict = {}
|
||||
|
||||
|
||||
@contextmanager
|
||||
def _anonymous_main():
|
||||
"""Hide __main__'s import identity while workers spawn: spawn re-executes
|
||||
the caller's script in every worker otherwise, which for an unguarded
|
||||
script means one duplicate full run per worker. Our workers import
|
||||
everything by module name and never need __main__.
|
||||
everything by module name and never need __main__. Depth-counted so
|
||||
overlapping windows restore the true originals, not a mid-window snapshot.
|
||||
|
||||
ponytail: window covers the whole map; a concurrent pool spawned from
|
||||
another thread whose tasks live in __main__ would break during it."""
|
||||
global _window_depth, _window_saved
|
||||
main = sys.modules.get("__main__")
|
||||
if main is None:
|
||||
yield
|
||||
return
|
||||
d = main.__dict__
|
||||
saved = {k: d.pop(k) for k in ("__file__", "__spec__") if k in d}
|
||||
d["__spec__"] = None # get_preparation_data reads it via attribute access
|
||||
with _window_lock:
|
||||
_window_depth += 1
|
||||
if _window_depth == 1:
|
||||
_window_saved = {k: d.pop(k) for k in ("__file__", "__spec__")
|
||||
if k in d}
|
||||
d["__spec__"] = None # get_preparation_data reads it via attribute access
|
||||
try:
|
||||
yield
|
||||
finally:
|
||||
d.pop("__spec__", None)
|
||||
d.update(saved)
|
||||
with _window_lock:
|
||||
_window_depth -= 1
|
||||
if _window_depth == 0:
|
||||
d.pop("__spec__", None)
|
||||
d.update(_window_saved)
|
||||
_window_saved = {}
|
||||
|
||||
|
||||
def _init_worker(kind: str, payload) -> None:
|
||||
@@ -140,12 +157,19 @@ def parse_charlevel_meta_parallel(
|
||||
if w <= 1 or n_pages < min_pages:
|
||||
return parse_charlevel_meta(doc_handle)
|
||||
|
||||
executor = ProcessPoolExecutor(
|
||||
max_workers=w,
|
||||
mp_context=multiprocessing.get_context("spawn"),
|
||||
initializer=_init_worker,
|
||||
initargs=src,
|
||||
)
|
||||
try:
|
||||
executor = ProcessPoolExecutor(
|
||||
max_workers=w,
|
||||
mp_context=multiprocessing.get_context("spawn"),
|
||||
initializer=_init_worker,
|
||||
initargs=src,
|
||||
)
|
||||
except Exception:
|
||||
# Restricted environments (no working POSIX semaphores) refuse the
|
||||
# pool at construction; the sequential path needs none of that.
|
||||
if getattr(multiprocessing.current_process(), "_inheriting", False):
|
||||
raise
|
||||
return parse_charlevel_meta(doc_handle)
|
||||
try:
|
||||
with _anonymous_main():
|
||||
results = list(executor.map(_run_page, range(n_pages)))
|
||||
|
||||
@@ -14,7 +14,8 @@ from .._version import sdk_version
|
||||
from ..errors import PageIndexAPIError
|
||||
|
||||
|
||||
def build_claude_mcp(client, include_management: bool = False, doc_ids=None):
|
||||
def build_claude_mcp(client, include_management: bool = False, doc_ids=None,
|
||||
server_name: str = "pageindex"):
|
||||
from ..agent_tools import _require_local_scope
|
||||
_require_local_scope(client, doc_ids)
|
||||
if getattr(client, "api_key", None):
|
||||
@@ -61,5 +62,5 @@ def build_claude_mcp(client, include_management: bool = False, doc_ids=None):
|
||||
for name, description, schema, invoke
|
||||
in _tool_specs(client, include_management, doc_ids)
|
||||
]
|
||||
return create_sdk_mcp_server(name="pageindex", version=sdk_version(),
|
||||
return create_sdk_mcp_server(name=server_name, version=sdk_version(),
|
||||
tools=tools)
|
||||
|
||||
+40
-5
@@ -5,6 +5,7 @@ import json
|
||||
import logging
|
||||
import multiprocessing
|
||||
import os
|
||||
import re
|
||||
import uuid
|
||||
from datetime import datetime, timezone
|
||||
from typing import Any
|
||||
@@ -15,6 +16,14 @@ from .utils import run_off_loop
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
_SURROGATES = re.compile("[\ud800-\udfff]")
|
||||
|
||||
|
||||
def _scrub_surrogates(text: str) -> str:
|
||||
"""Lone surrogates (surrogateescape'd names, PyPDF2's surrogatepass
|
||||
decodes) cannot encode to UTF-8; replace with U+FFFD."""
|
||||
return _SURROGATES.sub("\ufffd", text)
|
||||
|
||||
|
||||
def _now_iso() -> str:
|
||||
"""Naive UTC, millisecond precision."""
|
||||
@@ -76,7 +85,7 @@ class LocalAPI:
|
||||
"Failed to submit document: metadata must be a dict."
|
||||
)
|
||||
try:
|
||||
json.dumps(metadata)
|
||||
json.dumps(metadata, allow_nan=False)
|
||||
except (TypeError, ValueError) as e:
|
||||
raise PageIndexAPIError(
|
||||
f"Failed to submit document: metadata must be valid JSON. {e}"
|
||||
@@ -108,7 +117,11 @@ class LocalAPI:
|
||||
raise PageIndexAPIError(
|
||||
"Failed to submit document: PDF has no content. All pages are blank."
|
||||
)
|
||||
self._unique_doc_name(os.path.basename(file_path))
|
||||
# Surrogates from a surrogateescape'd filesystem name would be
|
||||
# mangled by the store's errors="replace" write; scrub now so the
|
||||
# returned name is byte-for-byte the stored name.
|
||||
doc_name = _scrub_surrogates(os.path.basename(file_path))
|
||||
self._unique_doc_name(doc_name)
|
||||
|
||||
try:
|
||||
if mode == "flash":
|
||||
@@ -124,6 +137,7 @@ class LocalAPI:
|
||||
raise
|
||||
except Exception as e:
|
||||
raise PageIndexAPIError(f"Failed to submit document: {e}") from e
|
||||
self._check_page_bounds(structure, len(page_texts))
|
||||
|
||||
doc_id = "pi-" + uuid.uuid4().hex
|
||||
pages = [{"page_index": i + 1, "markdown": text}
|
||||
@@ -134,7 +148,7 @@ class LocalAPI:
|
||||
with self._store.lock():
|
||||
meta = {
|
||||
"id": doc_id,
|
||||
"name": self._unique_doc_name(os.path.basename(file_path)),
|
||||
"name": self._unique_doc_name(doc_name),
|
||||
"description": description,
|
||||
"status": "completed",
|
||||
"createdAt": _now_iso(),
|
||||
@@ -163,12 +177,32 @@ class LocalAPI:
|
||||
"Please use a different file name."
|
||||
)
|
||||
|
||||
@staticmethod
|
||||
def _check_page_bounds(structure: list, page_count: int) -> None:
|
||||
"""The tree (pdfium) and stored pages (PyPDF2) come from different
|
||||
parsers; a span outside 1..page_count IndexErrors every later read."""
|
||||
stack = list(structure)
|
||||
while stack:
|
||||
node = stack.pop()
|
||||
start, end = node.get("start_index"), node.get("end_index")
|
||||
if (start is not None and end is not None
|
||||
and not (1 <= start and end <= page_count)):
|
||||
raise PageIndexAPIError(
|
||||
f"Failed to submit document: the extracted structure "
|
||||
f"references pages {start}-{end} outside the PDF's "
|
||||
f"{page_count} readable pages."
|
||||
)
|
||||
stack.extend(node.get("nodes") or [])
|
||||
|
||||
@staticmethod
|
||||
def _extract_page_texts(file_path: str) -> list[str]:
|
||||
import PyPDF2
|
||||
with open(file_path, "rb") as f:
|
||||
reader = PyPDF2.PdfReader(f)
|
||||
return [page.extract_text() or "" for page in reader.pages]
|
||||
# PyPDF2 decodes broken ToUnicode maps with surrogatepass; lone
|
||||
# surrogates would crash every utf-8 JSON save downstream.
|
||||
return [_scrub_surrogates(page.extract_text() or "")
|
||||
for page in reader.pages]
|
||||
|
||||
def _index_standard(self, file_path: str, page_texts: list[str]) -> tuple[list, str | None]:
|
||||
from .page_index_classic import page_index_main
|
||||
@@ -203,7 +237,8 @@ class LocalAPI:
|
||||
if not structure:
|
||||
raise PageIndexAPIError(
|
||||
"Failed to submit document: PageIndex Flash could not extract "
|
||||
"a structure from this PDF."
|
||||
"a structure from this PDF. Try mode='standard', which builds "
|
||||
"the structure with the model."
|
||||
)
|
||||
write_node_id(structure)
|
||||
description = generate_doc_description(
|
||||
|
||||
+130
-79
@@ -4,7 +4,6 @@ from __future__ import annotations
|
||||
import asyncio
|
||||
import hashlib
|
||||
import json
|
||||
import os
|
||||
import queue
|
||||
import threading
|
||||
import time
|
||||
@@ -215,15 +214,7 @@ def _openai_model(protocol: str, model_name: str, backend=None):
|
||||
from .utils import _litellm_model, _repair_litellm_types
|
||||
_repair_litellm_types()
|
||||
try:
|
||||
wire = _litellm_model(model_name, backend)
|
||||
except litellm.AuthenticationError as exc:
|
||||
raise PageIndexAPIError(
|
||||
"The OpenAI backend is not configured: set the "
|
||||
"OPENAI_API_KEY environment variable, pass an api_key "
|
||||
"in chat_backend / backend (any value works for keyless "
|
||||
"OPENAI_BASE_URL servers), or point chat_model at "
|
||||
"another provider (e.g. 'anthropic/...')."
|
||||
) from exc
|
||||
wire = _litellm_model(model_name)
|
||||
except litellm.NotFoundError as exc:
|
||||
raise PageIndexAPIError(str(exc)) from exc
|
||||
return LitellmModel(wire, api_key=(backend or {}).get("api_key"),
|
||||
@@ -235,33 +226,42 @@ def _reported_model(model_name: str) -> str:
|
||||
return model_name.removeprefix("litellm/").removeprefix("openai/")
|
||||
|
||||
|
||||
def _cache_extra_args(model_name: str) -> Optional[dict]:
|
||||
def _litellm_claude_marks(wire: str) -> Optional[dict]:
|
||||
"""Claude's prompt caching is opt-in per request: on Claude models
|
||||
routed through LiteLLM (Anthropic direct, Bedrock, Vertex — each
|
||||
channel live-verified), mark the managed system prefix and the newest
|
||||
message via LiteLLM's injection param so the loop's later turns and a
|
||||
conversation's next calls read them instead of repaying full price.
|
||||
Provider resolution is LiteLLM's own, so this predicate can never
|
||||
disagree with where the request actually routes."""
|
||||
if "/" not in model_name or model_name.startswith("openai/"):
|
||||
return None
|
||||
``wire`` is the name LiteLLM itself resolves — each lane strips its
|
||||
own routing prefixes first, because the lanes normalize differently
|
||||
(the chat wire treats bare names as OpenAI shorthand; the Agents SDK
|
||||
hands bare names to LiteLLM's own resolution)."""
|
||||
try:
|
||||
from litellm import get_llm_provider
|
||||
model, provider, _, _ = get_llm_provider(
|
||||
model=model_name.removeprefix("litellm/"))
|
||||
model, provider, _, _ = get_llm_provider(model=wire)
|
||||
except Exception:
|
||||
return None
|
||||
if provider == "anthropic" or (provider in ("bedrock", "vertex_ai")
|
||||
and "claude" in model.lower()):
|
||||
# The pair LiteLLM itself seeds for Anthropic and Bedrock: the
|
||||
# stable prefix plus the newest message, so each turn re-reads
|
||||
# the turns before it. Passing it explicitly extends it to Vertex.
|
||||
# The stable prefix plus the newest message, so each turn re-reads
|
||||
# the turns before it. LiteLLM seeds nothing unprompted, so this
|
||||
# pair is the marks' sole source.
|
||||
return {"cache_control_injection_points": [
|
||||
{"location": "message", "role": "system"},
|
||||
{"location": "message", "index": -1}]}
|
||||
return None
|
||||
|
||||
|
||||
def _cache_extra_args(model_name: str) -> Optional[dict]:
|
||||
"""The chat lane's marks: normalized exactly as _litellm_model
|
||||
normalizes the wire (bare names get openai/), so this predicate
|
||||
cannot disagree with where chat_completions actually routes."""
|
||||
wire = model_name.removeprefix("litellm/")
|
||||
if "/" not in wire or wire.startswith("openai/"):
|
||||
return None
|
||||
return _litellm_claude_marks(wire)
|
||||
|
||||
|
||||
def _openai_protocol(model_name: str) -> bool:
|
||||
"""Destinations that speak the OpenAI protocol on the wire, where
|
||||
prompt_cache_key means something and extra_body lands in the request
|
||||
@@ -327,12 +327,9 @@ def _openai_agent(client, protocol: str, model_name: str, instructions: str,
|
||||
body = {**(body or {}), **extra_body}
|
||||
else:
|
||||
extra_args = {**(extra_args or {}), **extra_body}
|
||||
return Agent(
|
||||
name="PageIndex",
|
||||
instructions=instructions,
|
||||
tools=build_openai_tools(client, doc_ids=doc_ids),
|
||||
model=_openai_model(protocol, model_name, conn or None),
|
||||
model_settings=ModelSettings(
|
||||
from pydantic import ValidationError
|
||||
try:
|
||||
settings = ModelSettings(
|
||||
temperature=temperature, top_p=top_p, max_tokens=max_tokens,
|
||||
reasoning=reasoning,
|
||||
# Streamed runs otherwise carry no usage at all (agents forwards
|
||||
@@ -340,7 +337,15 @@ def _openai_agent(client, protocol: str, model_name: str, instructions: str,
|
||||
include_usage=True,
|
||||
extra_body=body,
|
||||
extra_headers=extra_headers,
|
||||
extra_args=extra_args),
|
||||
extra_args=extra_args)
|
||||
except ValidationError as exc:
|
||||
raise PageIndexAPIError(f"Invalid model settings: {exc}") from exc
|
||||
return Agent(
|
||||
name="PageIndex",
|
||||
instructions=instructions,
|
||||
tools=build_openai_tools(client, doc_ids=doc_ids),
|
||||
model=_openai_model(protocol, model_name, conn or None),
|
||||
model_settings=settings,
|
||||
)
|
||||
|
||||
|
||||
@@ -370,21 +375,37 @@ def _conversation_cache_key(model_name: str, instructions: str, doc_id,
|
||||
return "pageindex-" + hashlib.sha256(seed.encode()).hexdigest()[:16]
|
||||
|
||||
|
||||
def _model_backend_error(exc) -> PageIndexAPIError:
|
||||
def _model_backend_error(exc, lane: str) -> PageIndexAPIError:
|
||||
"""Wrap a provider failure; the sol-class refusal (chatcmpl rejects
|
||||
function tools while reasoning is on) gets its two documented exits
|
||||
appended, since the fix is a different lane, not a retry."""
|
||||
function tools while reasoning is on) gets its documented exits
|
||||
appended, since the fix is a different route, not a retry. The exits
|
||||
are per-lane: of the chat lane's three, two are dead ends for a
|
||||
responses() caller — it IS the other lane, and its reasoning knob is
|
||||
``reasoning``, not ``reasoning_effort``."""
|
||||
message = f"The model backend failed: {exc}"
|
||||
if "Function tools with reasoning_effort" in str(exc):
|
||||
message += (
|
||||
" — this model runs tools on the Responses lane: upgrade "
|
||||
"litellm (newer releases route it there automatically), pass "
|
||||
"reasoning_effort (older litellm routes explicit efforts), or "
|
||||
"call responses() instead."
|
||||
"litellm (newer releases route it there automatically)"
|
||||
)
|
||||
message += (
|
||||
", pass reasoning_effort (older litellm routes explicit "
|
||||
"efforts), or call responses() instead." if lane == "chat"
|
||||
else "."
|
||||
)
|
||||
return PageIndexAPIError(message)
|
||||
|
||||
|
||||
def _translate_run_error(exc, max_turns, lane) -> PageIndexAPIError:
|
||||
"""The uncaught-run ladder every agent door shares."""
|
||||
from agents.exceptions import AgentsException, MaxTurnsExceeded
|
||||
if isinstance(exc, MaxTurnsExceeded):
|
||||
return _wrap_max_turns(max_turns)
|
||||
if isinstance(exc, AgentsException):
|
||||
return PageIndexAPIError(f"The agent backend failed: {exc}")
|
||||
return _model_backend_error(exc, lane)
|
||||
|
||||
|
||||
def _run_kwargs(max_turns) -> dict:
|
||||
# No traces — the caller opted into QA, not telemetry.
|
||||
from agents import RunConfig
|
||||
@@ -583,13 +604,9 @@ def run_chat_completions(client, messages, stream: bool = False,
|
||||
try:
|
||||
result = _run_sync(_run_closing(agent,
|
||||
Runner.run(agent, input=items, **run_kwargs)))
|
||||
except MaxTurnsExceeded as exc:
|
||||
raise _wrap_max_turns(max_turns) from exc
|
||||
except AgentsException as exc:
|
||||
raise PageIndexAPIError(
|
||||
f"The agent backend failed: {exc}") from exc
|
||||
except openai.OpenAIError as exc:
|
||||
raise _model_backend_error(exc) from exc
|
||||
except (MaxTurnsExceeded, AgentsException,
|
||||
openai.OpenAIError) as exc:
|
||||
raise _translate_run_error(exc, max_turns, "chat") from exc
|
||||
return {
|
||||
"id": f"chatcmpl-{uuid.uuid4().hex}",
|
||||
"object": "chat.completion",
|
||||
@@ -628,13 +645,9 @@ def run_chat_completions(client, messages, stream: bool = False,
|
||||
and isinstance(event.data, ResponseTextDeltaEvent)):
|
||||
yield chunk({"content": event.data.delta})
|
||||
completed = True
|
||||
except MaxTurnsExceeded as exc:
|
||||
raise _wrap_max_turns(max_turns) from exc
|
||||
except AgentsException as exc:
|
||||
raise PageIndexAPIError(
|
||||
f"The agent backend failed: {exc}") from exc
|
||||
except openai.OpenAIError as exc:
|
||||
raise _model_backend_error(exc) from exc
|
||||
except (MaxTurnsExceeded, AgentsException,
|
||||
openai.OpenAIError) as exc:
|
||||
raise _translate_run_error(exc, max_turns, "chat") from exc
|
||||
finally:
|
||||
if not completed and hasattr(streamed, "cancel"):
|
||||
streamed.cancel() # abandoned/failed: stop the agent task
|
||||
@@ -736,14 +749,10 @@ def run_responses(client, input, model: Optional[str] = None,
|
||||
result = _run_sync(_run_closing(agent,
|
||||
Runner.run(agent, input=[dict(item) for item in items],
|
||||
**run_kwargs)))
|
||||
except MaxTurnsExceeded as exc:
|
||||
raise _wrap_max_turns(max_turns) from exc
|
||||
except AgentsException as exc:
|
||||
raise PageIndexAPIError(
|
||||
f"The agent backend failed: {exc}") from exc
|
||||
except openai.OpenAIError as exc:
|
||||
raise PageIndexAPIError(
|
||||
f"The model backend failed: {exc}") from exc
|
||||
except (MaxTurnsExceeded, AgentsException,
|
||||
openai.OpenAIError) as exc:
|
||||
raise _translate_run_error(exc, max_turns,
|
||||
"responses") from exc
|
||||
transcript = result.to_input_list()[len(items):]
|
||||
return envelope(transcript, result.raw_responses)
|
||||
|
||||
@@ -801,16 +810,17 @@ def run_responses(client, input, model: Optional[str] = None,
|
||||
data["sequence_number"] = sequence
|
||||
yield data
|
||||
completed = True
|
||||
except MaxTurnsExceeded as exc:
|
||||
raise _wrap_max_turns(max_turns) from exc
|
||||
except AgentsException as exc:
|
||||
if recorded.get("status") not in ("failed", "incomplete"):
|
||||
raise PageIndexAPIError(
|
||||
f"The agent backend failed: {exc}") from exc
|
||||
# a run the envelope already reports failed/incomplete is done —
|
||||
# except for max_turns, which always gets its guidance
|
||||
if (isinstance(exc, MaxTurnsExceeded)
|
||||
or recorded.get("status") not in ("failed", "incomplete")):
|
||||
raise _translate_run_error(exc, max_turns,
|
||||
"responses") from exc
|
||||
completed = True
|
||||
except openai.OpenAIError as exc:
|
||||
raise PageIndexAPIError(
|
||||
f"The model backend failed: {exc}") from exc
|
||||
raise _translate_run_error(exc, max_turns,
|
||||
"responses") from exc
|
||||
finally:
|
||||
if not completed and hasattr(streamed, "cancel"):
|
||||
streamed.cancel() # abandoned/failed: stop the agent task
|
||||
@@ -847,14 +857,35 @@ def _require_anthropic() -> None:
|
||||
) from exc
|
||||
|
||||
|
||||
_ANTHROPIC_CLIENTS: dict = {} # backend key -> client, kept open for reuse
|
||||
|
||||
|
||||
def _anthropic_client(backend=None):
|
||||
"""The backend client — the seam tests replace with a fake transport."""
|
||||
"""The backend client — the seam tests replace with a fake transport.
|
||||
One client per backend: each construction pays ~45 ms of SSL-context
|
||||
build and a cold connection pool. A backend whose values defeat
|
||||
hashing constructs per call, as before."""
|
||||
import anthropic
|
||||
kwargs = _sdk_backend(backend)
|
||||
try:
|
||||
return anthropic.Anthropic(**_sdk_backend(backend))
|
||||
key = tuple(sorted(
|
||||
(k, tuple(sorted(v.items())) if isinstance(v, dict) else v)
|
||||
for k, v in kwargs.items()))
|
||||
hash(key)
|
||||
except TypeError:
|
||||
key = None
|
||||
if key in _ANTHROPIC_CLIENTS:
|
||||
return _ANTHROPIC_CLIENTS[key]
|
||||
try:
|
||||
client = anthropic.Anthropic(**kwargs)
|
||||
except TypeError as exc:
|
||||
raise PageIndexAPIError(
|
||||
f"The Anthropic backend is not configured: {exc}") from exc
|
||||
if key is not None and len(_ANTHROPIC_CLIENTS) < 8:
|
||||
# ponytail: cache capped at 8 backends; the tail constructs per call.
|
||||
# setdefault: never evict a client another thread may already hold.
|
||||
client = _ANTHROPIC_CLIENTS.setdefault(key, client)
|
||||
return client
|
||||
|
||||
|
||||
def _anthropic_system(extra_system, block: Optional[str]) -> list[dict]:
|
||||
@@ -892,10 +923,12 @@ def _cache_marks(system_blocks, messages) -> int:
|
||||
|
||||
def _dump_block(block) -> Any:
|
||||
"""A content block as a plain JSON dict, minus SDK-internal fields the
|
||||
API rejects (ParsedBetaTextBlock.__api_exclude__, e.g. parsed_output)."""
|
||||
API rejects (ParsedBetaTextBlock.__api_exclude__, e.g. parsed_output)
|
||||
and unset response-only defaults (exclude_unset, like the SDK's own
|
||||
request serializer — an explicit null fails the request schema)."""
|
||||
if hasattr(block, "model_dump"):
|
||||
exclude = getattr(type(block), "__api_exclude__", None)
|
||||
return block.model_dump(mode="json",
|
||||
return block.model_dump(mode="json", exclude_unset=True,
|
||||
exclude=set(exclude) if exclude else None)
|
||||
return block
|
||||
|
||||
@@ -929,11 +962,19 @@ def _default_max_tokens(model: str, thinking=None) -> int:
|
||||
"""The wire-required per-turn budget when the caller sets none: 8192,
|
||||
except the claude-3 generation whose output ceiling is 4096. The wire
|
||||
also requires max_tokens > thinking.budget_tokens, so an enabled
|
||||
budget lifts the default above itself."""
|
||||
budget lifts the default above itself — clamped to the model's output
|
||||
ceiling where LiteLLM's capability map knows it."""
|
||||
budget = (thinking.get("budget_tokens")
|
||||
if isinstance(thinking, dict) else None)
|
||||
if isinstance(budget, int):
|
||||
return budget + 8192
|
||||
if isinstance(budget, int) and not isinstance(budget, bool):
|
||||
want = budget + 8192
|
||||
try:
|
||||
import litellm
|
||||
ceiling = (litellm.model_cost.get(model)
|
||||
or {}).get("max_output_tokens")
|
||||
except Exception:
|
||||
ceiling = None
|
||||
return min(want, ceiling) if ceiling else want
|
||||
return 4096 if model.startswith(_CLAUDE_4096_MODELS) else 8192
|
||||
|
||||
|
||||
@@ -976,17 +1017,11 @@ def run_messages(client, messages, model: str,
|
||||
{"cache_control": {"type": "ephemeral"}}
|
||||
if _cache_marks(system_blocks, prepared) < 4 else {})
|
||||
merged = _merged_backend(client, backend)
|
||||
# The SDK defers credential resolution to request time and raises a
|
||||
# bare TypeError there — pre-check for the contract's PageIndexAPIError.
|
||||
if not merged and not (os.environ.get("ANTHROPIC_API_KEY")
|
||||
or os.environ.get("ANTHROPIC_AUTH_TOKEN")):
|
||||
raise PageIndexAPIError(
|
||||
"The Anthropic backend is not configured: set the "
|
||||
"ANTHROPIC_API_KEY environment variable, or pass an api_key "
|
||||
"in chat_backend / backend.")
|
||||
backend_client = _anthropic_client(merged)
|
||||
# A caller-owned http_client must survive the per-call closes below.
|
||||
owns_transport = "http_client" not in (merged or {})
|
||||
# Close only a per-call construction: cached clients stay open for
|
||||
# reuse; a caller-owned http_client survives regardless.
|
||||
owns_transport = ("http_client" not in (merged or {})
|
||||
and backend_client not in _ANTHROPIC_CLIENTS.values())
|
||||
if max_tokens is None:
|
||||
max_tokens = _default_max_tokens(model, thinking)
|
||||
runner = backend_client.beta.messages.tool_runner(
|
||||
@@ -1011,6 +1046,14 @@ def run_messages(client, messages, model: str,
|
||||
except anthropic.AnthropicError as exc:
|
||||
raise PageIndexAPIError(
|
||||
f"The model backend failed: {exc}") from exc
|
||||
except TypeError as exc:
|
||||
# the SDK's request-time credential-resolution failure
|
||||
if "authentication" not in str(exc).lower():
|
||||
raise
|
||||
raise PageIndexAPIError(
|
||||
"The Anthropic backend is not configured: set the "
|
||||
"ANTHROPIC_API_KEY environment variable, or pass an "
|
||||
f"api_key in chat_backend / backend. ({exc})") from exc
|
||||
finally:
|
||||
# runs on exhaustion and abandonment (GeneratorExit) alike
|
||||
if owns_transport:
|
||||
@@ -1022,6 +1065,14 @@ def run_messages(client, messages, model: str,
|
||||
except anthropic.AnthropicError as exc:
|
||||
raise PageIndexAPIError(
|
||||
f"The model backend failed: {exc}") from exc
|
||||
except TypeError as exc:
|
||||
# the SDK's request-time credential-resolution failure
|
||||
if "authentication" not in str(exc).lower():
|
||||
raise
|
||||
raise PageIndexAPIError(
|
||||
"The Anthropic backend is not configured: set the "
|
||||
"ANTHROPIC_API_KEY environment variable, or pass an "
|
||||
f"api_key in chat_backend / backend. ({exc})") from exc
|
||||
finally:
|
||||
# safe here: the params read-back below does no HTTP
|
||||
if owns_transport:
|
||||
|
||||
@@ -15,7 +15,10 @@ logger = logging.getLogger(__name__)
|
||||
def _write_json_atomic(path: Path, data) -> None:
|
||||
tmp = path.with_name(path.name + f".{uuid.uuid4().hex}.tmp")
|
||||
try:
|
||||
with open(tmp, "w", encoding="utf-8") as f:
|
||||
# errors=: a lone surrogate (os.fsdecode'd path in metadata, an
|
||||
# LLM-written \ud83d escape) must not crash the store after a whole
|
||||
# indexing run — it is replaced instead.
|
||||
with open(tmp, "w", encoding="utf-8", errors="replace") as f:
|
||||
json.dump(data, f, ensure_ascii=False)
|
||||
f.flush()
|
||||
os.fsync(f.fileno())
|
||||
|
||||
@@ -91,7 +91,9 @@ class McpBridge:
|
||||
) from exc
|
||||
# Strict id correlation only — accepting any result-bearing message
|
||||
# would return a stale or mis-correlated reply as this call's.
|
||||
reply = next((m for m in messages if m.get("id") == request_id), None)
|
||||
reply = next((m for m in messages
|
||||
if isinstance(m, dict) and m.get("id") == request_id),
|
||||
None)
|
||||
if reply is None:
|
||||
raise PageIndexAPIError(
|
||||
"MCP server response contained no reply matching the request."
|
||||
|
||||
@@ -242,7 +242,7 @@ def clean_tree_for_output(tree_nodes):
|
||||
return cleaned_nodes
|
||||
|
||||
|
||||
async def md_to_tree(md_path, if_thinning=False, min_token_threshold=None, if_add_node_summary='no', summary_token_threshold=None, model=None, if_add_doc_description='no', if_add_node_text='no', if_add_node_id='yes'):
|
||||
async def md_to_tree(md_path, if_thinning=False, min_token_threshold=None, if_add_node_summary='no', summary_token_threshold=None, model=None, if_add_doc_description='no', if_add_node_text='no', if_add_node_id='yes', summary_model=None):
|
||||
with open(md_path, 'r', encoding='utf-8') as f:
|
||||
markdown_content = f.read()
|
||||
line_count = markdown_content.count('\n') + 1
|
||||
@@ -267,11 +267,12 @@ async def md_to_tree(md_path, if_thinning=False, min_token_threshold=None, if_ad
|
||||
print(f"Formatting tree structure...")
|
||||
|
||||
if if_add_node_summary == 'yes':
|
||||
summary_model = summary_model or model
|
||||
# Always include text for summary generation
|
||||
tree_structure = format_structure(tree_structure, order = ['title', 'node_id', 'line_num', 'summary', 'prefix_summary', 'text', 'nodes'])
|
||||
|
||||
|
||||
print(f"Generating summaries for each node...")
|
||||
tree_structure = await generate_summaries_for_structure_md(tree_structure, summary_token_threshold=summary_token_threshold, model=model)
|
||||
tree_structure = await generate_summaries_for_structure_md(tree_structure, summary_token_threshold=summary_token_threshold, model=summary_model)
|
||||
|
||||
if if_add_node_text == 'no':
|
||||
# Remove text after summary generation if not requested
|
||||
@@ -281,7 +282,7 @@ async def md_to_tree(md_path, if_thinning=False, min_token_threshold=None, if_ad
|
||||
print(f"Generating document description...")
|
||||
# Create a clean structure without unnecessary fields for description generation
|
||||
clean_structure = create_clean_structure_for_description(tree_structure)
|
||||
doc_description = generate_doc_description(clean_structure, model=model)
|
||||
doc_description = generate_doc_description(clean_structure, model=summary_model)
|
||||
return {
|
||||
'doc_name': os.path.splitext(os.path.basename(md_path))[0],
|
||||
'doc_description': doc_description,
|
||||
|
||||
@@ -61,8 +61,8 @@ import re
|
||||
import sys
|
||||
from types import SimpleNamespace
|
||||
|
||||
from .utils import (ConfigLoader, _is_unrecoverable, _openai_missing_keys,
|
||||
llm_acompletion, strip_internal_keys)
|
||||
from .utils import (ConfigLoader, _is_unrecoverable, llm_acompletion,
|
||||
strip_internal_keys)
|
||||
|
||||
TRIGGER_PAGES = 5 # only look ahead on nodes larger than this
|
||||
ROUTING_COST = 1 # R(v), in pages
|
||||
@@ -624,6 +624,9 @@ def children_from_cache(node, cache, kinds):
|
||||
async def propose_children(node, pages, args):
|
||||
"""Generate one temporary level of children via the model. Validated, not committed."""
|
||||
start, end = node["start_index"], subtree_end(node)
|
||||
end = min(end, len(pages)) # a tree from another parser may overrun pages
|
||||
if end < start:
|
||||
return [] # the whole span is beyond the loaded pages
|
||||
block = "\n".join(
|
||||
f"<page_{n}>\n{pages[n - 1][:PAGE_CHARS]}\n</page_{n}>" for n in range(start, end + 1))
|
||||
answer = await ask_model(args.model, EXPAND_PROMPT.format(
|
||||
@@ -872,12 +875,6 @@ async def main():
|
||||
args = parser.parse_args()
|
||||
|
||||
model = args.model or default_model()
|
||||
if args.expand and not args.plan:
|
||||
missing = _openai_missing_keys(model)
|
||||
if missing:
|
||||
sys.exit(f"{', '.join(missing)} is not set "
|
||||
f"(expand model: {model}).")
|
||||
|
||||
original = json.load(open(args.structure))
|
||||
structure = copy.deepcopy(original["structure"])
|
||||
pages, lines = load_pages(args.pdf)
|
||||
|
||||
+67
-70
@@ -77,32 +77,13 @@ def run_off_loop(func, *args):
|
||||
return pool.submit(func, *args).result()
|
||||
|
||||
|
||||
def _openai_missing_keys(model):
|
||||
"""Missing env keys for the pre-check, which covers only OpenAI-shaped
|
||||
names (bare or ``openai/``): other providers resolve credentials their
|
||||
own way at call time (IAM chains, ADC, Ollama's localhost default),
|
||||
invisible to env inspection — the chat lane draws the same line.
|
||||
``litellm/``-prefixed names are exempt: the prefix is an explicit
|
||||
routing choice, and litellm resolves credentials beyond the
|
||||
environment (litellm.api_key, a keyless OPENAI_BASE_URL server).
|
||||
Truthiness, not litellm's validate_environment, which reports a blank
|
||||
exported key as present."""
|
||||
if model.startswith("litellm/"):
|
||||
return []
|
||||
if "/" in model and not model.startswith("openai/"):
|
||||
return []
|
||||
return ([] if (os.getenv("OPENAI_API_KEY") or "").strip()
|
||||
else ["OPENAI_API_KEY"])
|
||||
|
||||
|
||||
def _litellm_model(model, backend):
|
||||
def _litellm_model(model):
|
||||
"""Normalize to LiteLLM's grammar (``litellm/`` strips, bare names get
|
||||
the ``openai/`` wire form — same as the chat lane) and fail fast on a
|
||||
missing key or unknown provider, with status codes the retry loop and
|
||||
the summary/optimize passes treat as unrecoverable."""
|
||||
the ``openai/`` wire form — same as the chat lane) and refuse an
|
||||
unknown provider with the 404 the retry loop treats as unrecoverable.
|
||||
Credentials are LiteLLM's own call, made at the first completion."""
|
||||
if not model:
|
||||
return model
|
||||
raw = model
|
||||
model = _strip_prefix(model, "litellm/")
|
||||
if "/" not in model:
|
||||
model = f"openai/{model}"
|
||||
@@ -119,69 +100,69 @@ def _litellm_model(model, backend):
|
||||
f"this model id, use 'openai/{model}' and point "
|
||||
f"OPENAI_BASE_URL at the server.",
|
||||
llm_provider=None, model=model)
|
||||
if not backend:
|
||||
missing = _openai_missing_keys(raw)
|
||||
if missing:
|
||||
raise litellm.AuthenticationError(
|
||||
f"missing API key for {model}: {', '.join(missing)}",
|
||||
llm_provider=None, model=model)
|
||||
return model
|
||||
|
||||
|
||||
# Misconfiguration: no retry can fix a rejected key or a model that does not
|
||||
# exist, and every later call fails the same way. Deliberately not 400, which
|
||||
# also carries context_length_exceeded, a per-prompt failure the caller absorbs
|
||||
# today. An unknown status is a transport failure and stays retryable.
|
||||
# exist, and every later call fails the same way. An unknown status is a
|
||||
# transport failure and stays retryable.
|
||||
_UNRECOVERABLE_STATUS = frozenset({401, 403, 404})
|
||||
|
||||
# A 400 (context_length_exceeded) is equally unfixable by retry — the prompt
|
||||
# will not shrink — but it is per-prompt: the ladder raises it immediately
|
||||
# and consumers absorb it instead of failing the run.
|
||||
_NO_RETRY_STATUS = _UNRECOVERABLE_STATUS | frozenset({400})
|
||||
|
||||
|
||||
class LLMRetriesExhausted(RuntimeError):
|
||||
"""The retry ladder gave up; carries the last error's status_code."""
|
||||
|
||||
def __init__(self, message, status_code=None):
|
||||
super().__init__(message)
|
||||
self.status_code = status_code
|
||||
|
||||
|
||||
def _is_unrecoverable(exc: Exception) -> bool:
|
||||
if isinstance(exc, LLMRetriesExhausted):
|
||||
# 400 carries context_length_exceeded, the per-prompt failure the
|
||||
# caller absorbs (see above); any other exhausted ladder is fatal.
|
||||
return exc.status_code != 400
|
||||
return getattr(exc, "status_code", None) in _UNRECOVERABLE_STATUS
|
||||
|
||||
|
||||
def _no_cache_seeding_kwargs(backend):
|
||||
"""litellm 1.97 auto-marks Claude requests for prompt caching (system +
|
||||
last message); indexing prompts are single-shot and unique, so every call
|
||||
would pay the cache-write premium with nothing ever read back. A
|
||||
system-role-only injection point matches no indexing message, and its
|
||||
presence stops litellm seeding its own defaults; backend keys still
|
||||
win."""
|
||||
return {"cache_control_injection_points":
|
||||
[{"location": "message", "role": "system"}],
|
||||
**(backend or {})}
|
||||
|
||||
|
||||
def llm_completion(model, prompt, chat_history=None, return_finish_reason=False):
|
||||
import litellm
|
||||
max_retries = 10
|
||||
messages = list(chat_history) + [{"role": "user", "content": prompt}] if chat_history else [{"role": "user", "content": prompt}]
|
||||
backend = _llm_backend.get()
|
||||
model = _litellm_model(model, backend)
|
||||
model = _litellm_model(model)
|
||||
_repair_litellm_types()
|
||||
for i in range(max_retries):
|
||||
try:
|
||||
response = litellm.completion(
|
||||
model=model,
|
||||
messages=messages,
|
||||
drop_params=True,
|
||||
response = litellm.completion(**{
|
||||
"model": model,
|
||||
"messages": messages,
|
||||
"drop_params": True,
|
||||
# the loop is the retry policy; the merge lets a backend override win
|
||||
**{"max_retries": 0, **_no_cache_seeding_kwargs(backend)},
|
||||
)
|
||||
"max_retries": 0,
|
||||
**(backend or {}),
|
||||
})
|
||||
content = response.choices[0].message.content
|
||||
if return_finish_reason:
|
||||
finish_reason = "max_output_reached" if response.choices[0].finish_reason == "length" else "finished"
|
||||
return content, finish_reason
|
||||
return content
|
||||
except Exception as e:
|
||||
if _is_unrecoverable(e):
|
||||
if getattr(e, "status_code", None) in _NO_RETRY_STATUS:
|
||||
raise
|
||||
print('************* Retrying *************')
|
||||
logging.error(f"Error: {e}")
|
||||
if i < max_retries - 1:
|
||||
time.sleep(1)
|
||||
else:
|
||||
raise RuntimeError(
|
||||
f"LLM completion failed after {max_retries} retries"
|
||||
raise LLMRetriesExhausted(
|
||||
f"LLM completion failed after {max_retries} retries: {e}",
|
||||
status_code=getattr(e, "status_code", None),
|
||||
) from e
|
||||
|
||||
|
||||
@@ -190,27 +171,29 @@ async def llm_acompletion(model, prompt):
|
||||
max_retries = 10
|
||||
messages = [{"role": "user", "content": prompt}]
|
||||
backend = _llm_backend.get()
|
||||
model = _litellm_model(model, backend)
|
||||
model = _litellm_model(model)
|
||||
_repair_litellm_types()
|
||||
for i in range(max_retries):
|
||||
try:
|
||||
response = await litellm.acompletion(
|
||||
model=model,
|
||||
messages=messages,
|
||||
drop_params=True,
|
||||
**{"max_retries": 0, **_no_cache_seeding_kwargs(backend)},
|
||||
)
|
||||
response = await litellm.acompletion(**{
|
||||
"model": model,
|
||||
"messages": messages,
|
||||
"drop_params": True,
|
||||
"max_retries": 0,
|
||||
**(backend or {}),
|
||||
})
|
||||
return response.choices[0].message.content
|
||||
except Exception as e:
|
||||
if _is_unrecoverable(e):
|
||||
if getattr(e, "status_code", None) in _NO_RETRY_STATUS:
|
||||
raise
|
||||
print('************* Retrying *************')
|
||||
logging.error(f"Error: {e}")
|
||||
if i < max_retries - 1:
|
||||
await asyncio.sleep(1)
|
||||
else:
|
||||
raise RuntimeError(
|
||||
f"LLM completion failed after {max_retries} retries"
|
||||
raise LLMRetriesExhausted(
|
||||
f"LLM completion failed after {max_retries} retries: {e}",
|
||||
status_code=getattr(e, "status_code", None),
|
||||
) from e
|
||||
|
||||
|
||||
@@ -738,7 +721,8 @@ async def generate_summaries_for_structure(structure, model=None):
|
||||
if nodes and not any(node['summary'] for node in nodes):
|
||||
raise RuntimeError(
|
||||
"Summary generation failed for all nodes "
|
||||
"(check LLM credentials and model availability)"
|
||||
"(every summary call failed or returned empty; "
|
||||
"check the model and its context limits)"
|
||||
)
|
||||
return structure
|
||||
|
||||
@@ -835,10 +819,16 @@ async def summarize_tree(structure, pdf_pages, model=None,
|
||||
summary are left untouched; leaves under `small_node_tokens` use their raw
|
||||
text as the summary without a model call."""
|
||||
semaphore = asyncio.Semaphore(concurrency or SUMMARY_CONCURRENCY)
|
||||
asked = answered = False
|
||||
|
||||
async def ask(prompt):
|
||||
nonlocal asked, answered
|
||||
asked = True
|
||||
async with semaphore:
|
||||
return await llm_acompletion(model, prompt)
|
||||
reply = await llm_acompletion(model, prompt)
|
||||
if reply:
|
||||
answered = True
|
||||
return reply
|
||||
|
||||
async def leaf_summary(node):
|
||||
text = get_text_of_pdf_pages(pdf_pages, node['start_index'], node['end_index'])
|
||||
@@ -927,13 +917,16 @@ async def summarize_tree(structure, pdf_pages, model=None,
|
||||
if isinstance(r, Exception) and _is_unrecoverable(r):
|
||||
raise r
|
||||
|
||||
# Raw-text leaves summarize without the model, so they cannot vouch for
|
||||
# it: a run whose every model call failed still fails loud.
|
||||
def _any_summary(nodes):
|
||||
return any(n.get('summary') or _any_summary(n.get('nodes') or [])
|
||||
for n in nodes)
|
||||
if not _any_summary(structure):
|
||||
if (asked and not answered) or not _any_summary(structure):
|
||||
raise RuntimeError(
|
||||
"Summary generation failed for all nodes "
|
||||
"(check LLM credentials and model availability)"
|
||||
"(every summary call failed or returned empty; "
|
||||
"check the model and its context limits)"
|
||||
)
|
||||
|
||||
strip_internal_keys(structure)
|
||||
@@ -973,8 +966,12 @@ def generate_doc_description(structure, model=None):
|
||||
"""
|
||||
try:
|
||||
return llm_completion(model, prompt)
|
||||
except RuntimeError:
|
||||
return ""
|
||||
except Exception as e:
|
||||
# Per-prompt 400: the unbounded whole-tree prompt overran the
|
||||
# context; the indexed document survives with no description.
|
||||
if getattr(e, "status_code", None) == 400:
|
||||
return ""
|
||||
raise
|
||||
|
||||
|
||||
def reorder_dict(data, key_order):
|
||||
|
||||
+12
-10
@@ -3,7 +3,7 @@ import os
|
||||
import json
|
||||
from pageindex import *
|
||||
from pageindex.page_index_md import md_to_tree
|
||||
from pageindex.utils import ConfigLoader, _openai_missing_keys
|
||||
from pageindex.utils import ConfigLoader
|
||||
|
||||
# Keep LiteLLM's import off the network (frozen bundled model-cost map);
|
||||
# an explicit user setting wins.
|
||||
@@ -32,7 +32,7 @@ if __name__ == "__main__":
|
||||
parser.add_argument('--model', type=str, default=None,
|
||||
help='(legacy) Same as --index-model')
|
||||
parser.add_argument('--summary-model', type=str, default=None,
|
||||
help='Model for node summaries (defaults to --index-model, then --model, then config.yaml)')
|
||||
help='Model for node summaries (falls back to config.yaml summary_model, then --index-model, then --model)')
|
||||
|
||||
parser.add_argument('--toc-check-pages', type=int, default=None,
|
||||
help='Number of pages to check for table of contents (PDF only)')
|
||||
@@ -94,15 +94,12 @@ if __name__ == "__main__":
|
||||
|
||||
if args.mode == 'flash':
|
||||
from pageindex.flash import page_index_flash
|
||||
summary_model = (args.summary_model or args.index_model
|
||||
or args.model
|
||||
or ConfigLoader().load().summary_model)
|
||||
summary_model = ConfigLoader().load({k: v for k, v in {
|
||||
'summary_model': args.summary_model,
|
||||
'index_model': args.index_model,
|
||||
'model': args.model,
|
||||
}.items() if v is not None}).summary_model
|
||||
will_summarize = args.summary if args.summary is not None else True
|
||||
if will_summarize or args.optimize == 'full':
|
||||
missing = _openai_missing_keys(summary_model)
|
||||
if missing:
|
||||
raise SystemExit(
|
||||
f"Missing API key for {summary_model}: {', '.join(missing)}")
|
||||
toc_with_page_number = page_index_flash(
|
||||
args.pdf_path,
|
||||
optimize=args.optimize if args.optimize != 'off' else False,
|
||||
@@ -111,6 +108,9 @@ if __name__ == "__main__":
|
||||
use_embedded_toc=args.embedded_toc if args.embedded_toc is not None else True,
|
||||
summary=will_summarize,
|
||||
)
|
||||
if not toc_with_page_number.get('structure'):
|
||||
raise ValueError("PageIndex Flash could not extract a structure from this PDF; "
|
||||
"try --mode standard, which builds the structure with the model")
|
||||
if 'optimize' in toc_with_page_number:
|
||||
o = toc_with_page_number['optimize']
|
||||
print(f"Optimize: merges={o['merges']} expands={o['expands']}, "
|
||||
@@ -169,6 +169,7 @@ if __name__ == "__main__":
|
||||
user_opt = {
|
||||
'index_model': args.index_model,
|
||||
'model': args.model,
|
||||
'summary_model': args.summary_model,
|
||||
}
|
||||
|
||||
# Load config with defaults from config.yaml
|
||||
@@ -184,6 +185,7 @@ if __name__ == "__main__":
|
||||
if_add_node_summary=args.if_add_node_summary,
|
||||
summary_token_threshold=args.summary_token_threshold,
|
||||
model=opt.model,
|
||||
summary_model=opt.summary_model,
|
||||
if_add_doc_description=args.if_add_doc_description,
|
||||
if_add_node_text=args.if_add_node_text,
|
||||
if_add_node_id=args.if_add_node_id
|
||||
|
||||
+124
-5
@@ -213,6 +213,17 @@ def test_browse_documents_relevance_unsupported(client, store_path):
|
||||
assert "local mode" in bad_sort["error"]
|
||||
|
||||
|
||||
def test_expand_pages_enforces_contract_pattern():
|
||||
"""int() alone is far laxer than the published pages pattern; an
|
||||
out-of-contract spelling must reject, never resolve to another page."""
|
||||
from pageindex.agent_tools import _PageSpecError, _expand_pages
|
||||
assert _expand_pages("1-3, 7") == [1, 2, 3, 7]
|
||||
for bad in ["1_0", "+5", "٥", "１", " 1", "1 - 3"]:
|
||||
with pytest.raises(_PageSpecError) as excinfo:
|
||||
_expand_pages(bad)
|
||||
assert excinfo.value.code == "invalid"
|
||||
|
||||
|
||||
def test_browse_documents_empty_and_folder_error(client):
|
||||
payload, is_error = run(client, "browse_documents")
|
||||
assert not is_error
|
||||
@@ -746,6 +757,11 @@ def test_claude_agent_config_local(client, store_path):
|
||||
config = client.claude_agent_config(doc_id="pi-a")
|
||||
assert "report.pdf" in config["system_prompt"]
|
||||
assert config["allowed_tools"] == ["mcp__pageindex"]
|
||||
assert config["mcp_servers"]["pageindex"]["name"] == "pageindex"
|
||||
# The SDK server's declared identity follows the registration key.
|
||||
renamed = client.claude_agent_config(server_name="docs")
|
||||
assert renamed["mcp_servers"]["docs"]["name"] == "docs"
|
||||
assert renamed["allowed_tools"] == ["mcp__docs"]
|
||||
|
||||
|
||||
def test_openai_agent_config_local(client, store_path):
|
||||
@@ -759,6 +775,7 @@ def test_openai_agent_config_local(client, store_path):
|
||||
assert config["model"] == client.retrieve_model
|
||||
assert client.openai_agent_config(model="gpt-x")["model"] == "gpt-x"
|
||||
assert Agent(**client.openai_agent_config()).name == "PageIndex"
|
||||
assert client.openai_agent_config(name="Researcher")["name"] == "Researcher"
|
||||
|
||||
|
||||
def test_openai_agent_config_model_speaks_the_agents_sdk_grammar(tmp_path):
|
||||
@@ -777,6 +794,59 @@ def test_openai_agent_config_model_speaks_the_agents_sdk_grammar(tmp_path):
|
||||
== "litellm/groq/llama-x")
|
||||
|
||||
|
||||
def test_openai_agent_config_carries_cache_marks_for_litellm_claude(tmp_path):
|
||||
"""LiteLLM-routed Claude gets the same cache marks the engine
|
||||
attaches in chat_completions(); OpenAI-bound models stay unmarked
|
||||
(their caching is server-side, and LiteLLM seeds nothing on its
|
||||
own)."""
|
||||
pytest.importorskip("agents")
|
||||
from pageindex.local_chat import _cache_extra_args
|
||||
client = PageIndexLocalClient(storage_path=str(tmp_path / "s"),
|
||||
chat_model="anthropic/claude-x")
|
||||
settings = client.openai_agent_config()["model_settings"]
|
||||
assert settings.extra_args == _cache_extra_args("anthropic/claude-x")
|
||||
assert "cache_control_injection_points" in settings.extra_args
|
||||
assert "model_settings" not in client.openai_agent_config(model="gpt-x")
|
||||
# The per-call override is marked by its own routing, not the default's.
|
||||
marked = client.openai_agent_config(model="bedrock/claude-y")
|
||||
assert "cache_control_injection_points" in marked["model_settings"].extra_args
|
||||
|
||||
|
||||
def test_openai_agent_config_marks_bare_claude_behind_litellm_prefix(tmp_path):
|
||||
"""In this lane litellm/<bare-claude> routes to Anthropic — the Agents
|
||||
SDK strips the prefix and LiteLLM resolves the bare name — unlike the
|
||||
chat lane, whose wire treats bare names as OpenAI shorthand. The marks
|
||||
follow this lane's routing, not the chat lane's."""
|
||||
pytest.importorskip("agents")
|
||||
pytest.importorskip("litellm")
|
||||
client = PageIndexLocalClient(storage_path=str(tmp_path / "s"))
|
||||
cfg = client.openai_agent_config(model="litellm/claude-sonnet-4-5")
|
||||
assert cfg["model"] == "litellm/claude-sonnet-4-5"
|
||||
assert "cache_control_injection_points" in cfg["model_settings"].extra_args
|
||||
# Without the prefix the SDK's default OpenAI provider serves the name.
|
||||
assert "model_settings" not in client.openai_agent_config(
|
||||
model="claude-sonnet-4-5")
|
||||
assert "model_settings" not in client.openai_agent_config(
|
||||
model="litellm/gpt-4o")
|
||||
|
||||
|
||||
def test_openai_agent_config_merges_caller_model_settings(tmp_path):
|
||||
"""Caller model_settings merge on top of the bundled cache marks
|
||||
(caller fields win, extra_args dict-merge); with no marks the
|
||||
caller's object rides through verbatim."""
|
||||
pytest.importorskip("agents")
|
||||
from agents import ModelSettings
|
||||
client = PageIndexLocalClient(storage_path=str(tmp_path / "s"),
|
||||
chat_model="anthropic/claude-x")
|
||||
mine = ModelSettings(temperature=0.2, extra_args={"top_k": 5})
|
||||
merged = client.openai_agent_config(model_settings=mine)["model_settings"]
|
||||
assert merged.temperature == 0.2
|
||||
assert merged.extra_args["top_k"] == 5
|
||||
assert "cache_control_injection_points" in merged.extra_args
|
||||
verbatim = client.openai_agent_config(model="gpt-x", model_settings=mine)
|
||||
assert verbatim["model_settings"] is mine
|
||||
|
||||
|
||||
def test_plain_functions_answer_bad_arguments_with_the_envelope(client,
|
||||
store_path):
|
||||
"""agent_tools() functions must not raise into a framework loop:
|
||||
@@ -1234,6 +1304,10 @@ class _FakeBridge:
|
||||
]
|
||||
|
||||
def list_tools(self):
|
||||
# mirrors the live server: the read endpoint serves the read subset
|
||||
if "tools=read" in self.url:
|
||||
return [tool for tool in self.tools
|
||||
if (tool.get("annotations") or {}).get("readOnlyHint")]
|
||||
return self.tools
|
||||
|
||||
def instructions(self):
|
||||
@@ -1266,8 +1340,7 @@ def test_cloud_agent_tools_discover_live_tool_set(cloud_with_fake_bridge):
|
||||
# instructions fetch and the hosted/MCP registrations.
|
||||
assert bridge.url == "https://api.pageindex.ai/mcp?tools=read"
|
||||
assert bridge.headers == {"Authorization": "Bearer pi-test-key"}
|
||||
# Default: only tools the server marks read-only; unannotated tools are
|
||||
# treated as non-read-only.
|
||||
# The endpoint is the only gate: whatever it serves is exposed verbatim.
|
||||
assert [t.__name__ for t in tools] == ["search_documents", "get_document"]
|
||||
assert "ESCALATION tool" in tools[0].__doc__
|
||||
|
||||
@@ -1759,7 +1832,9 @@ def test_annotation_for_both_nullable_encodings():
|
||||
== Optional[str])
|
||||
|
||||
|
||||
def test_cloud_agent_tools_empty_filter_raises(monkeypatch):
|
||||
def test_cloud_agent_tools_trust_the_gated_endpoint(monkeypatch):
|
||||
"""The ?tools=read endpoint is the only gate: what it serves is exposed
|
||||
verbatim, with no client-side annotation second-guessing."""
|
||||
import pageindex.mcp_bridge as mcp_bridge
|
||||
|
||||
class _AllWriteBridge:
|
||||
@@ -1773,11 +1848,37 @@ def test_cloud_agent_tools_empty_filter_raises(monkeypatch):
|
||||
|
||||
monkeypatch.setattr(mcp_bridge, "McpBridge", _AllWriteBridge)
|
||||
cloud = PageIndexCloudClient(api_key="pi-test-key")
|
||||
with pytest.raises(PageIndexAPIError, match="annotation"):
|
||||
cloud.agent_tools()
|
||||
assert len(cloud.agent_tools()) == 1
|
||||
assert len(cloud.agent_tools(include_management=True)) == 1
|
||||
|
||||
|
||||
def test_extract_result_skips_non_dict_messages():
|
||||
"""A 200 body of null, a batched array, or an SSE string frame must
|
||||
surface as the contract's PageIndexAPIError, not an AttributeError."""
|
||||
from pageindex.mcp_bridge import McpBridge
|
||||
|
||||
bridge = McpBridge("http://unused", {})
|
||||
|
||||
class _Resp:
|
||||
def __init__(self, text, content_type="application/json"):
|
||||
self.headers = {"Content-Type": content_type}
|
||||
self.status_code = 200
|
||||
self.content = text.encode("utf-8")
|
||||
|
||||
def json(self):
|
||||
return json.loads(self.content)
|
||||
|
||||
for body in ("null", '[{"jsonrpc": "2.0", "id": 1, "result": {}}]',
|
||||
'"hello"'):
|
||||
with pytest.raises(PageIndexAPIError, match="no reply matching"):
|
||||
bridge._extract_result(_Resp(body), 1)
|
||||
|
||||
sse = ('data: "noise"\n\n'
|
||||
'data: {"jsonrpc": "2.0", "id": 1, "result": {"ok": true}}\n\n')
|
||||
result = bridge._extract_result(_Resp(sse, "text/event-stream"), 1)
|
||||
assert result == {"ok": True}
|
||||
|
||||
|
||||
def test_bridge_call_tool_surfaces_iserror(monkeypatch):
|
||||
import requests as requests_mod
|
||||
import pageindex.mcp_bridge as mcp_bridge
|
||||
@@ -2363,6 +2464,24 @@ def test_submit_wait_poll_error_carries_doc_id(fake_cloud_client, monkeypatch):
|
||||
cloud.submit_document("whatever.pdf", wait=True)
|
||||
|
||||
|
||||
def test_submit_wait_reraises_definite_poll_answers(fake_cloud_client,
|
||||
monkeypatch):
|
||||
"""A 401/403/404 poll answer is final: re-raised untouched, no retries, no keep-polling advice."""
|
||||
cloud = fake_cloud_client(["processing"])
|
||||
polls = {"n": 0}
|
||||
|
||||
def denied(doc_id):
|
||||
polls["n"] += 1
|
||||
raise PageIndexAPIError("Failed to get document metadata: 401",
|
||||
status_code=401)
|
||||
|
||||
monkeypatch.setattr(cloud, "get_document", denied)
|
||||
with pytest.raises(PageIndexAPIError, match="401") as err:
|
||||
cloud.submit_document("whatever.pdf", wait=True)
|
||||
assert polls["n"] == 1
|
||||
assert "Processing continues" not in str(err.value)
|
||||
|
||||
|
||||
def test_config_helpers_reject_empty_doc_id_on_cloud():
|
||||
"""An explicitly empty scope must not silently widen to the whole
|
||||
library — cloud has no tool-layer allowlist to enforce it."""
|
||||
|
||||
+384
-45
@@ -282,30 +282,47 @@ def test_submit_defaults_to_flash(local_client, sample_pdf, monkeypatch):
|
||||
assert local_client._api._store.get_meta(doc_id)["mode"] == "flash"
|
||||
|
||||
|
||||
def test_submit_rejects_structure_beyond_stored_pages(local_client, sample_pdf,
|
||||
monkeypatch):
|
||||
"""A tree spanning pages the store lacks fails submit instead of saving a doc whose reads IndexError."""
|
||||
monkeypatch.setattr(
|
||||
pageindex.flash, "page_index_flash",
|
||||
lambda pdf, **kwargs: {
|
||||
"doc_name": "sample.pdf",
|
||||
"structure": [{"title": "Root", "start_index": 1,
|
||||
"end_index": 3, "summary": "s", "nodes": []}]})
|
||||
monkeypatch.setattr(pageindex.utils, "llm_completion",
|
||||
lambda model, prompt, **kw: "d.")
|
||||
with pytest.raises(PageIndexAPIError, match="pages 1-3 outside"):
|
||||
local_client.submit_document(sample_pdf)
|
||||
assert local_client._api._store.list_metas() == []
|
||||
|
||||
|
||||
def test_submit_survives_pypdf2_lone_surrogates(local_client, sample_pdf,
|
||||
monkeypatch):
|
||||
"""PyPDF2 decodes broken ToUnicode with surrogatepass; the store gets U+FFFD, not a utf-8-fatal str."""
|
||||
import PyPDF2
|
||||
monkeypatch.setattr(PyPDF2.PageObject, "extract_text",
|
||||
lambda self: "\ud83dello broken")
|
||||
monkeypatch.setattr(
|
||||
pageindex.flash, "page_index_flash",
|
||||
lambda p, **kwargs: {
|
||||
"structure": [{"title": "T", "start_index": 1,
|
||||
"end_index": 1, "summary": "s", "nodes": []}]})
|
||||
monkeypatch.setattr(pageindex.utils, "llm_completion",
|
||||
lambda model, prompt, **kw: "d.")
|
||||
doc_id = local_client.submit_document(sample_pdf)["doc_id"]
|
||||
markdown = local_client.get_ocr(doc_id)["result"][0]["markdown"]
|
||||
assert "\ud83d" not in markdown
|
||||
assert markdown.startswith("�ello")
|
||||
|
||||
|
||||
def test_page_index_flash_rejects_unknown_optimize():
|
||||
from pageindex.flash import page_index_flash
|
||||
with pytest.raises(ValueError, match="optimize must be"):
|
||||
page_index_flash("never-opened.pdf", optimize="off")
|
||||
|
||||
|
||||
def test_llm_completion_missing_key_raises_immediately(monkeypatch):
|
||||
import openai
|
||||
import litellm # first import may load a .env; delenv after it
|
||||
monkeypatch.delenv("OPENAI_API_KEY", raising=False)
|
||||
with pytest.raises(openai.OpenAIError, match="OPENAI_API_KEY"):
|
||||
pageindex.utils.llm_completion("gpt-4o", "probe")
|
||||
with pytest.raises(openai.OpenAIError, match="OPENAI_API_KEY"):
|
||||
asyncio.run(pageindex.utils.llm_acompletion("gpt-4o", "probe"))
|
||||
# unknown bare names are OpenAI shorthand, so the same check applies
|
||||
with pytest.raises(openai.OpenAIError, match="OPENAI_API_KEY"):
|
||||
pageindex.utils.llm_completion("my-finetune-v2", "probe")
|
||||
# a blank exported key is as missing as no key (litellm's
|
||||
# validate_environment reports it present)
|
||||
monkeypatch.setenv("OPENAI_API_KEY", " ")
|
||||
with pytest.raises(openai.OpenAIError, match="OPENAI_API_KEY"):
|
||||
pageindex.utils.llm_completion("gpt-4o", "probe")
|
||||
|
||||
|
||||
def test_llm_completion_refuses_unknown_provider(monkeypatch):
|
||||
"""A first segment LiteLLM does not know (a HuggingFace repo id like
|
||||
Qwen/...) is refused with the openai/ escape before the retry loop,
|
||||
@@ -320,11 +337,12 @@ def test_llm_completion_refuses_unknown_provider(monkeypatch):
|
||||
def test_submit_missing_llm_key_fails_loud(local_client, sample_pdf, monkeypatch):
|
||||
import litellm # first import may load a .env; delenv after it
|
||||
monkeypatch.delenv("OPENAI_API_KEY", raising=False)
|
||||
monkeypatch.setattr("pageindex.utils.time.sleep", lambda s: None)
|
||||
def first_llm_call(*args, **kwargs):
|
||||
return pageindex.utils.llm_completion("gpt-4o", "probe")
|
||||
monkeypatch.setattr(page_index_module, "page_index_main", first_llm_call)
|
||||
monkeypatch.setattr(pageindex.flash, "page_index_flash", first_llm_call)
|
||||
for kwargs in ({}, {"mode": "flash"}):
|
||||
for kwargs in ({"mode": "standard"}, {"mode": "flash"}):
|
||||
with pytest.raises(PageIndexAPIError, match="OPENAI_API_KEY"):
|
||||
local_client.submit_document(sample_pdf, **kwargs)
|
||||
assert local_client.list_documents()["total"] == 0
|
||||
@@ -344,6 +362,17 @@ def test_submit_rejections(local_client, sample_pdf, tmp_path):
|
||||
local_client.submit_document(sample_pdf, beta_headers=["block_reference"])
|
||||
|
||||
|
||||
def test_write_json_atomic_replaces_lone_surrogates(tmp_path):
|
||||
"""A lone surrogate (an os.fsdecode'd path in metadata, an LLM-written
|
||||
\\ud83d escape) must not crash the store's writer after a whole
|
||||
indexing run — it lands as the encoder's replacement character."""
|
||||
from pageindex.local_store import _read_json, _write_json_atomic
|
||||
|
||||
path = tmp_path / "doc.json"
|
||||
_write_json_atomic(path, {"src": "bad-\udcff-path"})
|
||||
assert _read_json(path)["src"] == "bad-?-path"
|
||||
|
||||
|
||||
def test_corrupt_pdf_raises_api_error(local_client, tmp_path):
|
||||
bad = tmp_path / "bad.pdf"
|
||||
bad.write_bytes(b"%PDF-1.4 garbage with no xref or trailer")
|
||||
@@ -637,14 +666,16 @@ def test_list_documents_skips_unsafe_directory_names(
|
||||
assert [d["id"] for d in listing["documents"]] == [indexed_doc]
|
||||
|
||||
|
||||
def test_generate_doc_description_error_boundary(monkeypatch):
|
||||
def test_generate_doc_description_propagates_failures(monkeypatch):
|
||||
"""No swallow: a dead model fails the run instead of storing ''."""
|
||||
def raiser(exc):
|
||||
def _f(*args, **kwargs):
|
||||
raise exc
|
||||
return _f
|
||||
monkeypatch.setattr(pageindex.utils, "llm_completion",
|
||||
raiser(RuntimeError("retries exhausted")))
|
||||
assert pageindex.utils.generate_doc_description([]) == ""
|
||||
with pytest.raises(RuntimeError):
|
||||
pageindex.utils.generate_doc_description([])
|
||||
monkeypatch.setattr(pageindex.utils, "llm_completion",
|
||||
raiser(ValueError("provider rejected the model")))
|
||||
with pytest.raises(ValueError):
|
||||
@@ -712,10 +743,107 @@ def test_summarize_tree_child_unrecoverable_raises(monkeypatch):
|
||||
structure, pdf_pages, small_node_tokens=0))
|
||||
|
||||
|
||||
def test_llm_completion_suppresses_litellm_cache_seeding(monkeypatch):
|
||||
"""Indexing prompts are single-shot: without an explicit injection
|
||||
point litellm 1.97 seeds its own cache marks and every call pays the
|
||||
write premium for nothing. Backend keys still override ours."""
|
||||
def test_summarize_tree_fails_loud_when_every_model_call_fails(monkeypatch):
|
||||
"""A failure foreign to the retry ladder (not LLMRetriesExhausted)
|
||||
blanks per node; the asked-and-never-answered backstop must still fail
|
||||
the run loud — a raw-text short leaf cannot vouch for it."""
|
||||
async def exhausted(model, prompt):
|
||||
raise RuntimeError("LLM call failed after 10 attempts")
|
||||
monkeypatch.setattr(pageindex.utils, "llm_acompletion", exhausted)
|
||||
pdf_pages = [("tiny", 1), ("beta " * 300, 300)]
|
||||
structure = [{"title": "R", "start_index": 1, "end_index": 2,
|
||||
"nodes": [
|
||||
{"title": "A", "start_index": 1, "end_index": 1},
|
||||
{"title": "B", "start_index": 2, "end_index": 2}]}]
|
||||
with pytest.raises(RuntimeError, match="every summary call failed"):
|
||||
asyncio.run(pageindex.utils.summarize_tree(structure, pdf_pages))
|
||||
|
||||
|
||||
def test_summarize_tree_all_short_leaves_need_no_model(monkeypatch):
|
||||
"""A tree whose every node summarizes from raw text makes zero model
|
||||
calls and must not be mistaken for a failed run."""
|
||||
async def unexpected(model, prompt):
|
||||
raise AssertionError("no model call expected")
|
||||
monkeypatch.setattr(pageindex.utils, "llm_acompletion", unexpected)
|
||||
structure = [{"title": "A", "start_index": 1, "end_index": 1}]
|
||||
out = asyncio.run(pageindex.utils.summarize_tree(structure, [("tiny", 1)]))
|
||||
assert out[0]["summary"] == "tiny"
|
||||
|
||||
|
||||
def test_summarize_tree_partial_exhaustion_fails_loud(monkeypatch):
|
||||
"""One lucky call must not vouch for a model that then went away: a
|
||||
ladder-exhausted node raises instead of silently blanking."""
|
||||
calls = {"n": 0}
|
||||
|
||||
async def flaky(model, prompt):
|
||||
calls["n"] += 1
|
||||
if calls["n"] == 1:
|
||||
return '{"points": ["p"], "summary": "ok"}'
|
||||
raise pageindex.utils.LLMRetriesExhausted(
|
||||
"LLM completion failed after 10 retries", status_code=500)
|
||||
monkeypatch.setattr(pageindex.utils, "llm_acompletion", flaky)
|
||||
pdf_pages = [("alpha " * 300, 300), ("beta " * 300, 300)]
|
||||
structure = [{"title": "A", "start_index": 1, "end_index": 1},
|
||||
{"title": "B", "start_index": 2, "end_index": 2}]
|
||||
with pytest.raises(pageindex.utils.LLMRetriesExhausted):
|
||||
asyncio.run(pageindex.utils.summarize_tree(structure, pdf_pages))
|
||||
|
||||
|
||||
def test_generate_summaries_partial_exhaustion_fails_loud(monkeypatch):
|
||||
calls = {"n": 0}
|
||||
|
||||
async def flaky(model, prompt):
|
||||
calls["n"] += 1
|
||||
if calls["n"] == 1:
|
||||
return "fine"
|
||||
raise pageindex.utils.LLMRetriesExhausted(
|
||||
"LLM completion failed after 10 retries", status_code=500)
|
||||
monkeypatch.setattr(pageindex.utils, "llm_acompletion", flaky)
|
||||
structure = [{"title": "A", "text": "t1",
|
||||
"nodes": [{"title": "B", "text": "t2"}]}]
|
||||
with pytest.raises(pageindex.utils.LLMRetriesExhausted):
|
||||
asyncio.run(
|
||||
pageindex.utils.generate_summaries_for_structure(structure))
|
||||
|
||||
|
||||
def test_expand_exhausted_ladder_fails_loud(monkeypatch):
|
||||
"""Keyless/broken-model expand must kill the run, not degrade to
|
||||
no_children after burning the retry ladder on every node."""
|
||||
import pageindex.tree_optimize as tree_optimize
|
||||
|
||||
async def exhausted(model, prompt):
|
||||
raise pageindex.utils.LLMRetriesExhausted(
|
||||
"LLM completion failed after 10 retries", status_code=500)
|
||||
monkeypatch.setattr(tree_optimize, "llm_acompletion", exhausted)
|
||||
structure = [{"title": "T", "start_index": 1, "end_index": 8,
|
||||
"node_id": "0001", "nodes": []}]
|
||||
pages = ["heading\nbody text"] * 8
|
||||
lines = [["heading", "body text"]] * 8
|
||||
with pytest.raises(pageindex.utils.LLMRetriesExhausted):
|
||||
asyncio.run(tree_optimize.optimize(structure, pages, lines,
|
||||
model="m", do_expand=True))
|
||||
|
||||
|
||||
def test_expand_absorbs_per_prompt_rejection(monkeypatch):
|
||||
"""A 400-exhausted node (context_length_exceeded) stays collapsed and
|
||||
the run survives — the documented per-prompt absorption."""
|
||||
import pageindex.tree_optimize as tree_optimize
|
||||
|
||||
async def rejected(model, prompt):
|
||||
raise pageindex.utils.LLMRetriesExhausted(
|
||||
"LLM completion failed after 10 retries", status_code=400)
|
||||
monkeypatch.setattr(tree_optimize, "llm_acompletion", rejected)
|
||||
structure = [{"title": "T", "start_index": 1, "end_index": 8,
|
||||
"node_id": "0001", "nodes": []}]
|
||||
pages = ["heading\nbody text"] * 8
|
||||
lines = [["heading", "body text"]] * 8
|
||||
outcome = asyncio.run(tree_optimize.optimize(structure, pages, lines,
|
||||
model="m", do_expand=True))
|
||||
assert outcome["expands"] == 0
|
||||
|
||||
|
||||
def test_llm_completion_backend_reaches_litellm(monkeypatch):
|
||||
"""No cache params of our own; backend keys reach litellm and win the merge."""
|
||||
import litellm
|
||||
captured = {}
|
||||
|
||||
@@ -728,15 +856,41 @@ def test_llm_completion_suppresses_litellm_cache_seeding(monkeypatch):
|
||||
monkeypatch.setattr(litellm, "completion", fake_completion)
|
||||
monkeypatch.setenv("OPENAI_API_KEY", "k")
|
||||
assert pageindex.utils.llm_completion("gpt-4o", "probe") == "ok"
|
||||
assert captured["cache_control_injection_points"] == [
|
||||
{"location": "message", "role": "system"}]
|
||||
assert "cache_control_injection_points" not in captured
|
||||
token = pageindex.utils._llm_backend.set(
|
||||
{"api_key": "x", "cache_control_injection_points": []})
|
||||
{"api_key": "x", "max_retries": 3})
|
||||
try:
|
||||
pageindex.utils.llm_completion("gpt-4o", "probe")
|
||||
finally:
|
||||
pageindex.utils._llm_backend.reset(token)
|
||||
assert captured["cache_control_injection_points"] == []
|
||||
assert captured["api_key"] == "x"
|
||||
assert captured["max_retries"] == 3
|
||||
|
||||
|
||||
def test_backend_overrides_reserved_kwargs_without_retry(monkeypatch):
|
||||
"""A backend key colliding with our own kwargs wins the merge instead
|
||||
of raising TypeError through the retry ladder."""
|
||||
import litellm
|
||||
calls = {"n": 0}
|
||||
captured = {}
|
||||
|
||||
def fake_completion(**kwargs):
|
||||
calls["n"] += 1
|
||||
captured.clear()
|
||||
captured.update(kwargs)
|
||||
message = types.SimpleNamespace(content="ok")
|
||||
choice = types.SimpleNamespace(message=message, finish_reason="stop")
|
||||
return types.SimpleNamespace(choices=[choice])
|
||||
monkeypatch.setattr(litellm, "completion", fake_completion)
|
||||
monkeypatch.setattr(pageindex.utils.time, "sleep", lambda s: None)
|
||||
monkeypatch.setenv("OPENAI_API_KEY", "k")
|
||||
token = pageindex.utils._llm_backend.set({"drop_params": False})
|
||||
try:
|
||||
assert pageindex.utils.llm_completion("gpt-4o", "probe") == "ok"
|
||||
finally:
|
||||
pageindex.utils._llm_backend.reset(token)
|
||||
assert captured["drop_params"] is False
|
||||
assert calls["n"] == 1
|
||||
|
||||
|
||||
def test_delete_survives_marker_tamper(local_client, tmp_path):
|
||||
@@ -885,6 +1039,32 @@ def test_cloud_error_and_empty_delete(cloud, monkeypatch):
|
||||
assert client.delete_document("pi-1") == {}
|
||||
|
||||
|
||||
def test_cloud_errors_carry_status_code(cloud, monkeypatch, sample_pdf):
|
||||
"""Every non-200 raise carries the HTTP status, so callers can branch
|
||||
on 429-vs-401 instead of parsing message text."""
|
||||
client, calls, fake = cloud
|
||||
_patch_requests(monkeypatch,
|
||||
lambda m, url, kw: FakeResponse(status_code=418, text="no"))
|
||||
attempts = [
|
||||
lambda: client.submit_document(sample_pdf),
|
||||
lambda: client.get_ocr("pi-1"),
|
||||
lambda: client.get_tree("pi-1"),
|
||||
lambda: client.submit_query("pi-1", "q"),
|
||||
lambda: client.get_retrieval("r-1"),
|
||||
lambda: client.chat_completions(
|
||||
messages=[{"role": "user", "content": "q"}]),
|
||||
lambda: client.get_document("pi-1"),
|
||||
lambda: client.delete_document("pi-1"),
|
||||
lambda: client.list_documents(),
|
||||
lambda: client.create_folder("f"),
|
||||
lambda: client.list_folders(),
|
||||
]
|
||||
for attempt in attempts:
|
||||
with pytest.raises(PageIndexAPIError) as err:
|
||||
attempt()
|
||||
assert err.value.status_code == 418
|
||||
|
||||
|
||||
def test_cloud_chat_stream_parsing(cloud, monkeypatch):
|
||||
client, calls, fake = cloud
|
||||
lines = [
|
||||
@@ -964,40 +1144,55 @@ def test_backend_scopes_the_index_lane(tmp_path, monkeypatch):
|
||||
assert captured["api_base"] == "http://b"
|
||||
|
||||
|
||||
def test_index_precheck_covers_only_openai_shaped(monkeypatch):
|
||||
"""The missing-key pre-check fires only for OpenAI-shaped names — other
|
||||
providers resolve credentials at call time (IAM chains, ADC, Ollama's
|
||||
localhost default), invisible to env inspection, so the lane must not
|
||||
block them up front."""
|
||||
def test_index_lane_makes_no_key_prejudgment(monkeypatch):
|
||||
"""Credentials are LiteLLM's call at completion time: keyless
|
||||
environments reach the wire untouched for every provider shape."""
|
||||
pytest.importorskip("litellm")
|
||||
import litellm
|
||||
from types import SimpleNamespace
|
||||
from pageindex.utils import llm_completion
|
||||
|
||||
monkeypatch.delenv("OPENAI_API_KEY", raising=False)
|
||||
with pytest.raises(litellm.AuthenticationError, match="missing API key"):
|
||||
llm_completion("my-finetune-v2", "p")
|
||||
|
||||
reply = SimpleNamespace(choices=[SimpleNamespace(
|
||||
message=SimpleNamespace(content="ok"), finish_reason="stop")])
|
||||
monkeypatch.setattr(litellm, "completion", lambda **kw: reply)
|
||||
monkeypatch.setattr(litellm, "validate_environment",
|
||||
lambda *a, **k: pytest.fail("env pre-check ran"))
|
||||
assert llm_completion("my-finetune-v2", "p") == "ok"
|
||||
assert llm_completion("ollama/llama3", "p") == "ok"
|
||||
assert llm_completion("bedrock/anthropic.claude-sonnet", "p") == "ok"
|
||||
|
||||
|
||||
def test_litellm_routing_prefix_skips_key_precheck(monkeypatch):
|
||||
"""litellm/-prefixed names are an explicit routing choice: litellm
|
||||
resolves credentials beyond the environment (litellm.api_key, a
|
||||
keyless OPENAI_BASE_URL server), so the env pre-check stands aside."""
|
||||
def test_llm_completion_surfaces_litellms_credential_verdict(monkeypatch):
|
||||
"""LiteLLM's own missing-credentials error (a retryable-shaped 500)
|
||||
rides the retry loop and lands verbatim in the terminal error."""
|
||||
pytest.importorskip("litellm")
|
||||
import litellm # noqa: F401 — first import may load a .env; delenv after
|
||||
from pageindex.utils import _litellm_model, _openai_missing_keys
|
||||
from pageindex.utils import llm_completion
|
||||
|
||||
monkeypatch.delenv("OPENAI_API_KEY", raising=False)
|
||||
assert _openai_missing_keys("litellm/gpt-4o") == []
|
||||
assert _litellm_model("litellm/gpt-4o", None) == "openai/gpt-4o"
|
||||
monkeypatch.setattr("pageindex.utils.time.sleep", lambda s: None)
|
||||
with pytest.raises(RuntimeError, match="OPENAI_API_KEY"):
|
||||
llm_completion("gpt-4o", "probe")
|
||||
|
||||
async def _nosleep(s):
|
||||
pass
|
||||
|
||||
monkeypatch.setattr("pageindex.utils.asyncio.sleep", _nosleep)
|
||||
with pytest.raises(RuntimeError, match="OPENAI_API_KEY"):
|
||||
asyncio.run(pageindex.utils.llm_acompletion("gpt-4o", "probe"))
|
||||
|
||||
|
||||
def test_litellm_model_normalizes_without_key_prejudgment(monkeypatch):
|
||||
"""_litellm_model only normalizes and provider-checks — a keyless
|
||||
environment changes nothing for any spelling."""
|
||||
pytest.importorskip("litellm")
|
||||
import litellm # noqa: F401 — first import may load a .env; delenv after
|
||||
from pageindex.utils import _litellm_model
|
||||
|
||||
monkeypatch.delenv("OPENAI_API_KEY", raising=False)
|
||||
assert _litellm_model("litellm/gpt-4o") == "openai/gpt-4o"
|
||||
assert _litellm_model("gpt-4o") == "openai/gpt-4o"
|
||||
|
||||
|
||||
def test_custom_provider_map_passes_provider_precheck(monkeypatch):
|
||||
@@ -1009,7 +1204,7 @@ def test_custom_provider_map_passes_provider_precheck(monkeypatch):
|
||||
|
||||
monkeypatch.setattr(litellm, "custom_provider_map",
|
||||
[{"provider": "my-llm", "custom_handler": object()}])
|
||||
assert _litellm_model("my-llm/model-a", None) == "my-llm/model-a"
|
||||
assert _litellm_model("my-llm/model-a") == "my-llm/model-a"
|
||||
|
||||
|
||||
def test_backend_args_are_local_only():
|
||||
@@ -1079,3 +1274,147 @@ def test_format_tree_node_keeps_key_items():
|
||||
assert out["key_items"] == ["1.1 Alpha", "1.2 Beta", "1.3 Gamma"]
|
||||
assert "key_items" not in _format_tree_node(
|
||||
{"title": "t", "node_id": "0001", "start_index": 1}, False)
|
||||
|
||||
|
||||
# ── retry-ladder and summary fail-loud edges (twelfth review) ──
|
||||
|
||||
def test_summarize_tree_all_empty_replies_fail_loud(monkeypatch):
|
||||
"""Empty-content replies (content filter, spent output cap) must not
|
||||
vouch for the model: a raw-text short leaf cannot carry the run when
|
||||
every model reply comes back blank."""
|
||||
async def blank(model, prompt):
|
||||
return ""
|
||||
monkeypatch.setattr(pageindex.utils, "llm_acompletion", blank)
|
||||
pdf_pages = [("tiny", 1), ("beta " * 300, 300)]
|
||||
structure = [{"title": "R", "start_index": 1, "end_index": 2,
|
||||
"nodes": [
|
||||
{"title": "A", "start_index": 1, "end_index": 1},
|
||||
{"title": "B", "start_index": 2, "end_index": 2}]}]
|
||||
with pytest.raises(RuntimeError, match="returned empty"):
|
||||
asyncio.run(pageindex.utils.summarize_tree(structure, pdf_pages))
|
||||
|
||||
|
||||
def test_summarize_tree_partial_empty_reply_absorbed(monkeypatch):
|
||||
"""One blank reply among good ones stays the documented per-node
|
||||
absorption: blank summary, run survives."""
|
||||
async def flaky(model, prompt):
|
||||
if "alpha" in prompt:
|
||||
return ""
|
||||
return '{"points": [], "summary": "ok"}'
|
||||
monkeypatch.setattr(pageindex.utils, "llm_acompletion", flaky)
|
||||
pdf_pages = [("alpha " * 300, 300), ("beta " * 300, 300)]
|
||||
structure = [{"title": "A", "start_index": 1, "end_index": 1},
|
||||
{"title": "B", "start_index": 2, "end_index": 2}]
|
||||
out = asyncio.run(pageindex.utils.summarize_tree(structure, pdf_pages))
|
||||
assert [n["summary"] for n in out] == ["", "ok"]
|
||||
|
||||
|
||||
def test_generate_doc_description_absorbs_context_overflow(monkeypatch):
|
||||
"""A per-prompt 400 (the whole-tree prompt overran the context) keeps
|
||||
the indexed document: empty description instead of a lost run. Anything
|
||||
else keeps propagating (see the sibling no-swallow test)."""
|
||||
class Rejected(Exception):
|
||||
status_code = 400
|
||||
|
||||
def boom(model, prompt):
|
||||
raise Rejected("context_length_exceeded")
|
||||
monkeypatch.setattr(pageindex.utils, "llm_completion", boom)
|
||||
assert pageindex.utils.generate_doc_description([]) == ""
|
||||
|
||||
|
||||
def test_llm_completion_400_skips_the_retry_ladder(monkeypatch):
|
||||
"""A 400 rejects this prompt permanently — retrying cannot shrink it:
|
||||
one wire call, raised raw so consumers can absorb it per policy."""
|
||||
pytest.importorskip("litellm")
|
||||
import litellm
|
||||
|
||||
class Rejected(Exception):
|
||||
status_code = 400
|
||||
|
||||
calls = {"n": 0}
|
||||
|
||||
def reject(**kw):
|
||||
calls["n"] += 1
|
||||
raise Rejected("context_length_exceeded")
|
||||
monkeypatch.setattr(litellm, "completion", reject)
|
||||
monkeypatch.setattr("pageindex.utils.time.sleep", lambda s: None)
|
||||
with pytest.raises(Rejected):
|
||||
pageindex.utils.llm_completion("gpt-4o", "p")
|
||||
assert calls["n"] == 1
|
||||
|
||||
|
||||
def test_llm_acompletion_400_skips_the_retry_ladder(monkeypatch):
|
||||
pytest.importorskip("litellm")
|
||||
import litellm
|
||||
|
||||
class Rejected(Exception):
|
||||
status_code = 400
|
||||
|
||||
calls = {"n": 0}
|
||||
|
||||
async def reject(**kw):
|
||||
calls["n"] += 1
|
||||
raise Rejected("context_length_exceeded")
|
||||
monkeypatch.setattr(litellm, "acompletion", reject)
|
||||
|
||||
async def _nosleep(s):
|
||||
pass
|
||||
monkeypatch.setattr("pageindex.utils.asyncio.sleep", _nosleep)
|
||||
with pytest.raises(Rejected):
|
||||
asyncio.run(pageindex.utils.llm_acompletion("gpt-4o", "p"))
|
||||
assert calls["n"] == 1
|
||||
|
||||
|
||||
def test_parse_pages_keeps_whitespace_tolerance():
|
||||
"""0.2.10 accepted whitespace around parts (int() tolerance); the SDK
|
||||
surface keeps that while the tool layer stays on the strict pattern."""
|
||||
from pageindex.client import _parse_pages
|
||||
assert _parse_pages(" 1-3") == [1, 2, 3]
|
||||
assert _parse_pages("5 - 7") == [5, 6, 7]
|
||||
assert _parse_pages("3 ,8") == [3, 8]
|
||||
assert _parse_pages("1-3\n") == [1, 2, 3]
|
||||
assert _parse_pages("\t2") == [2]
|
||||
with pytest.raises(PageIndexAPIError):
|
||||
_parse_pages("1 2")
|
||||
|
||||
|
||||
def test_submit_scrubs_surrogates_from_the_stored_name(
|
||||
local_client, sample_pdf, monkeypatch):
|
||||
"""The store scrubs at write time; scrubbing the basename at entry keeps
|
||||
the returned name identical to the stored name and lets the rename
|
||||
warning fire. (APFS refuses surrogate filenames, so the basename is
|
||||
patched instead of the filesystem.)"""
|
||||
import os
|
||||
|
||||
def fake_page_index_main(doc, opt=None, logger=None, page_list=None):
|
||||
return {"doc_name": "x", "doc_description": "d",
|
||||
"structure": json.loads(json.dumps(STRUCTURE))}
|
||||
monkeypatch.setattr(page_index_module, "page_index_main",
|
||||
fake_page_index_main)
|
||||
real = os.path.basename
|
||||
monkeypatch.setattr(os.path, "basename",
|
||||
lambda p: "re\udcffport.pdf"
|
||||
if real(str(p)) == "sample.pdf" else real(p))
|
||||
with pytest.warns(UserWarning, match="stored as"):
|
||||
result = local_client.submit_document(sample_pdf, mode="standard")
|
||||
assert result["name"] == "re\ufffdport.pdf"
|
||||
docs = local_client.list_documents()["documents"]
|
||||
assert [d["name"] for d in docs] == ["re\ufffdport.pdf"]
|
||||
|
||||
|
||||
def test_submit_rejects_nan_metadata(local_client):
|
||||
"""json.dumps' default (allow_nan=True) passes NaN/Infinity that no
|
||||
strict JSON parser accepts; the gate must reject them before they reach
|
||||
disk and every tool envelope."""
|
||||
with pytest.raises(PageIndexAPIError, match="valid JSON"):
|
||||
local_client.submit_document("/nonexistent.pdf",
|
||||
metadata={"score": float("nan")})
|
||||
|
||||
|
||||
def test_submit_flash_empty_structure_points_to_standard(
|
||||
local_client, sample_pdf, monkeypatch):
|
||||
"""The heading-less hard-fail names its way out: mode='standard'."""
|
||||
monkeypatch.setattr(pageindex.flash, "page_index_flash",
|
||||
lambda pdf, **kwargs: {"structure": []})
|
||||
with pytest.raises(PageIndexAPIError, match="mode='standard'"):
|
||||
local_client.submit_document(sample_pdf)
|
||||
|
||||
+318
-19
@@ -13,6 +13,21 @@ import pytest
|
||||
PDF = Path(__file__).parent.parent / "examples" / "documents" / "earthmover.pdf"
|
||||
|
||||
|
||||
def test_read_bookmarks_same_on_pdfium_4_and_5():
|
||||
"""Users may install pypdfium2 4.x or 5.x (floor >=4.30) — the bookmark
|
||||
reader has one branch per major and both must yield the same entries.
|
||||
The CI pdfium-4 leg runs this against the 4.x branch; everywhere else
|
||||
it pins the 5.x branch to the same values."""
|
||||
from pageindex.flash.embedded_toc import read_bookmarks
|
||||
|
||||
pdf = Path(__file__).parent.parent / "examples" / "documents" / "attention-residuals.pdf"
|
||||
bookmarks = read_bookmarks(str(pdf))
|
||||
assert len(bookmarks) == 22
|
||||
assert bookmarks[0] == {"title": "Introduction", "level": 1, "page": 2}
|
||||
assert bookmarks[2] == {"title": "Training Deep Networks via Residuals",
|
||||
"level": 2, "page": 3}
|
||||
|
||||
|
||||
@pytest.mark.skipif(int(version("pypdfium2").split(".")[0]) < 5,
|
||||
reason="extraction is pinned to pdfium 5.x font-name semantics")
|
||||
def test_page_text_pins_pdfium5_semantics():
|
||||
@@ -45,20 +60,17 @@ def test_page_mode_walk_uses_merged_surrogate_census():
|
||||
assert unmapped["ch"] == "β"
|
||||
|
||||
|
||||
def test_optimize_full_fails_fast_without_a_key(tmp_path, monkeypatch):
|
||||
"""optimize='full' runs LLM expand: with no key configured it must be
|
||||
an instant, guided PageIndexAPIError — raised before any PDF work (a
|
||||
bogus path proves the ordering) — while the LLM-free spellings and a
|
||||
backend-carrying indexing scope stay untouched."""
|
||||
def test_optimize_full_keyless_reports_file_errors_first(tmp_path, monkeypatch):
|
||||
"""No credential pre-check: a bad path is a FileNotFoundError even
|
||||
keyless (validation runs first), and the LLM-free spellings still run
|
||||
end to end."""
|
||||
from conftest import build_pdf
|
||||
from pageindex import PageIndexAPIError
|
||||
from pageindex.flash import page_index_flash
|
||||
from pageindex.utils import _llm_backend
|
||||
import litellm # noqa: F401 — first import may load a .env; delenv after it
|
||||
monkeypatch.delenv("OPENAI_API_KEY", raising=False)
|
||||
monkeypatch.delenv("CHATGPT_API_KEY", raising=False)
|
||||
|
||||
with pytest.raises(PageIndexAPIError, match="optimize='merge'"):
|
||||
with pytest.raises(FileNotFoundError):
|
||||
page_index_flash(str(tmp_path / "missing.pdf"), summary=False)
|
||||
|
||||
pdf = tmp_path / "doc.pdf"
|
||||
@@ -68,12 +80,46 @@ def test_optimize_full_fails_fast_without_a_key(tmp_path, monkeypatch):
|
||||
result = page_index_flash(str(pdf), summary=False, optimize=False)
|
||||
assert "structure" in result
|
||||
|
||||
token = _llm_backend.set({"api_key": "k"})
|
||||
try:
|
||||
with pytest.raises(FileNotFoundError):
|
||||
page_index_flash(str(tmp_path / "missing.pdf"), summary=False)
|
||||
finally:
|
||||
_llm_backend.reset(token)
|
||||
|
||||
def test_empty_outline_gate_carries_page_texts(tmp_path):
|
||||
"""The gate's bookmark-built trees feed the same summary/expand passes
|
||||
as detected ones, so its result must carry the per-page text too."""
|
||||
from conftest import build_pdf
|
||||
from pageindex.flash.main import extract_toc
|
||||
|
||||
pdf = tmp_path / "doc.pdf"
|
||||
pdf.write_bytes(build_pdf(["Alpha body", "Beta body"]))
|
||||
result = extract_toc(str(pdf))
|
||||
assert result["structure"] == [] # the short-document gate fired
|
||||
assert len(result["page_texts"]) == 2
|
||||
assert "Alpha" in result["page_texts"][0]
|
||||
|
||||
|
||||
def test_propose_children_clamps_to_loaded_pages(monkeypatch):
|
||||
"""A tree from another parser may overrun the loaded pages; the span is
|
||||
clamped instead of IndexErroring into a silently frozen node."""
|
||||
import asyncio
|
||||
from types import SimpleNamespace
|
||||
import pageindex.tree_optimize as tree_optimize
|
||||
|
||||
seen = {}
|
||||
|
||||
async def fake_ask(model, prompt):
|
||||
seen["prompt"] = prompt
|
||||
return {"subsections": []}
|
||||
|
||||
monkeypatch.setattr(tree_optimize, "ask_model", fake_ask)
|
||||
node = {"title": "T", "start_index": 1, "end_index": 3, "node_id": "n1"}
|
||||
out = asyncio.run(tree_optimize.propose_children(
|
||||
node, ["page one", "page two"], SimpleNamespace(model="m")))
|
||||
assert out == []
|
||||
assert "<page_2>" in seen["prompt"] and "<page_3>" not in seen["prompt"]
|
||||
|
||||
seen.clear()
|
||||
node = {"title": "T", "start_index": 3, "end_index": 4, "node_id": "n2"}
|
||||
out = asyncio.run(tree_optimize.propose_children(
|
||||
node, ["page one", "page two"], SimpleNamespace(model="m")))
|
||||
assert out == [] and "prompt" not in seen # fully beyond: no model call
|
||||
|
||||
|
||||
def test_bootstrap_reimport_is_not_swallowed(monkeypatch):
|
||||
@@ -107,6 +153,30 @@ def test_bootstrap_reimport_is_not_swallowed(monkeypatch):
|
||||
assert len(out) == len(meta) > 0 # normal failures still fall back sequentially
|
||||
|
||||
|
||||
def test_pool_construction_failure_falls_back_sequential(monkeypatch):
|
||||
"""Restricted environments (no working POSIX semaphores) refuse the pool
|
||||
at construction, before any work is mapped; indexing must take the
|
||||
sequential path, not die — except in a bootstrapping spawn child, where
|
||||
a sequential rerun would duplicate the whole run per worker."""
|
||||
import multiprocessing
|
||||
|
||||
from pageindex.flash import parser_pdfium_parallel as mod
|
||||
|
||||
class RefusedExecutor:
|
||||
def __init__(self, *a, **k):
|
||||
raise OSError("Function not implemented")
|
||||
|
||||
monkeypatch.setattr(mod, "ProcessPoolExecutor", RefusedExecutor)
|
||||
out, meta = mod.parse_charlevel_meta_parallel(str(PDF), workers=2,
|
||||
min_pages=1)
|
||||
assert len(out) == len(meta) > 0
|
||||
|
||||
cur = multiprocessing.current_process()
|
||||
monkeypatch.setattr(cur, "_inheriting", True, raising=False)
|
||||
with pytest.raises(OSError):
|
||||
mod.parse_charlevel_meta_parallel(str(PDF), workers=2, min_pages=1)
|
||||
|
||||
|
||||
def test_submit_document_refuses_during_bootstrap(tmp_path, monkeypatch):
|
||||
import multiprocessing
|
||||
|
||||
@@ -154,7 +224,7 @@ def test_optimize_wins_over_deprecated_optimize_expand(tmp_path, monkeypatch):
|
||||
|
||||
pdf = tmp_path / "doc.pdf"
|
||||
pdf.write_bytes(build_pdf(["1 Introduction", "Body text"]))
|
||||
# resolved to "full" before the precedence fix (keyless → fail-fast)
|
||||
# explicit "merge" wins even when the deprecated flag says expand
|
||||
with pytest.warns(DeprecationWarning):
|
||||
result = page_index_flash(str(pdf), summary=False,
|
||||
optimize="merge", optimize_expand=True)
|
||||
@@ -164,7 +234,236 @@ def test_optimize_wins_over_deprecated_optimize_expand(tmp_path, monkeypatch):
|
||||
optimize=True, optimize_expand=False)
|
||||
assert "structure" in result
|
||||
# optimize=None means unset ("full"), not off
|
||||
from pageindex import PageIndexAPIError
|
||||
with pytest.raises(PageIndexAPIError, match="optimize='merge'"):
|
||||
page_index_flash(str(tmp_path / "missing.pdf"), summary=False,
|
||||
optimize=None)
|
||||
from pageindex.flash import api as flash_api
|
||||
seen = {}
|
||||
|
||||
def fake_optimize(structure, pages, do_expand, model):
|
||||
seen["do_expand"] = do_expand
|
||||
return {"merges": 0}
|
||||
|
||||
monkeypatch.setattr(flash_api, "_optimize", fake_optimize)
|
||||
monkeypatch.setattr(flash_api, "extract_toc",
|
||||
lambda pdf, use_embedded_toc=True: {
|
||||
"structure": [{"title": "T", "start_index": 1,
|
||||
"end_index": 1, "nodes": []}],
|
||||
"page_texts": ["body"]})
|
||||
page_index_flash(str(pdf), summary=False, optimize=None)
|
||||
assert seen["do_expand"] is True
|
||||
|
||||
|
||||
def test_lone_surrogate_from_broken_tounicode_is_replaced(monkeypatch):
|
||||
"""An unpaired UTF-16 surrogate leaves as U+FFFD, not a str that crashes utf-8 save."""
|
||||
import json
|
||||
from io import BytesIO
|
||||
|
||||
import pypdfium2 as pdfium
|
||||
import pypdfium2.raw as pdfium_c
|
||||
from conftest import build_pdf
|
||||
from pageindex.flash.parser_pdfium_charlevel.char_extract import (
|
||||
_extract_raw_chars)
|
||||
|
||||
orig = pdfium_c.FPDFText_GetUnicode
|
||||
monkeypatch.setattr(pdfium_c, "FPDFText_GetUnicode",
|
||||
lambda tp, i: 0xD83D if i == 0 else orig(tp, i))
|
||||
pdf = pdfium.PdfDocument(BytesIO(build_pdf(["Hello broken cmap"])))
|
||||
page = pdf[0]
|
||||
# hold the textpage: GC finalizes an unreferenced one mid-extraction,
|
||||
# closing the handle so every per-char call reads back 0
|
||||
text_page = page.get_textpage()
|
||||
raw_chars, _objects = _extract_raw_chars(page, text_page.raw)
|
||||
text = "".join(char["ch"] for char in raw_chars)
|
||||
assert "\ud83d" not in text
|
||||
assert text.startswith("�ello")
|
||||
json.dumps(text) # the save-time crash this guards against
|
||||
|
||||
|
||||
def test_lone_surrogate_targets_never_patched_into_chars():
|
||||
"""A surrogate-band code with no cmap entry (chr fallback) must not patch a lone surrogate back in."""
|
||||
from pageindex.flash.parser_pdfium_charlevel.unicode_apply import (
|
||||
_apply_font_unicode)
|
||||
|
||||
char = {"i": 0, "ch": "X", "is_gen": False}
|
||||
show_codes = [(7, (0xD8, 0x3D), 100.0)]
|
||||
map_cache = {7: (2, {})} # Identity map, no ToUnicode: target = chr(0xD83D)
|
||||
|
||||
_apply_font_unicode([char], [], show_codes, None, map_cache)
|
||||
|
||||
assert char["ch"] == "�"
|
||||
|
||||
|
||||
def test_lone_surrogate_from_single_byte_map_is_replaced():
|
||||
"""The single-byte branch scrubs mapped lone surrogates like the
|
||||
two-byte branch does."""
|
||||
from pageindex.flash.parser_pdfium_charlevel.unicode_apply import (
|
||||
_apply_font_unicode)
|
||||
|
||||
char = {"i": 0, "ch": "X", "is_gen": False}
|
||||
show_codes = [(7, (0x41,), 100.0)]
|
||||
map_cache = {7: (1, {0x41: "\ud83d"})}
|
||||
|
||||
_apply_font_unicode([char], [], show_codes, None, map_cache)
|
||||
|
||||
assert char["ch"] == "�"
|
||||
|
||||
|
||||
def test_anonymous_main_overlapping_windows_restore(monkeypatch):
|
||||
"""The last window out must restore the true originals, not a mid-window snapshot."""
|
||||
import sys
|
||||
import threading
|
||||
|
||||
from pageindex.flash.parser_pdfium_parallel import _anonymous_main
|
||||
|
||||
main = sys.modules["__main__"]
|
||||
spec = object()
|
||||
monkeypatch.setattr(main, "__file__", "sentinel-file", raising=False)
|
||||
monkeypatch.setattr(main, "__spec__", spec, raising=False)
|
||||
a_in, b_in, a_out = (threading.Event() for _ in range(3))
|
||||
errors = []
|
||||
|
||||
def first():
|
||||
try:
|
||||
with _anonymous_main():
|
||||
a_in.set()
|
||||
assert b_in.wait(5)
|
||||
a_out.set()
|
||||
except BaseException as exc: # in-thread failures only warn in pytest
|
||||
errors.append(exc)
|
||||
|
||||
def second():
|
||||
try:
|
||||
assert a_in.wait(5)
|
||||
with _anonymous_main():
|
||||
b_in.set()
|
||||
assert a_out.wait(5)
|
||||
except BaseException as exc:
|
||||
errors.append(exc)
|
||||
|
||||
threads = [threading.Thread(target=first), threading.Thread(target=second)]
|
||||
for thread in threads:
|
||||
thread.start()
|
||||
for thread in threads:
|
||||
thread.join(10)
|
||||
assert not errors
|
||||
assert main.__spec__ is spec
|
||||
assert main.__file__ == "sentinel-file"
|
||||
|
||||
|
||||
def test_optimize_full_skips_expand_without_page_texts(tmp_path, monkeypatch):
|
||||
"""A bookmark-only extraction (no page_texts) skips expand; merge still runs."""
|
||||
from conftest import build_pdf
|
||||
from pageindex.flash import api as flash_api
|
||||
|
||||
monkeypatch.setenv("OPENAI_API_KEY", "k")
|
||||
calls = {}
|
||||
|
||||
def fake_optimize(structure, pages, do_expand, model):
|
||||
calls["pages"] = pages
|
||||
calls["do_expand"] = do_expand
|
||||
return {"merges": 0}
|
||||
|
||||
monkeypatch.setattr(flash_api, "_optimize", fake_optimize)
|
||||
monkeypatch.setattr(flash_api, "extract_toc",
|
||||
lambda pdf, use_embedded_toc=True: {
|
||||
"structure": [{"title": "T", "start_index": 1,
|
||||
"end_index": 1, "nodes": []}]})
|
||||
pdf = tmp_path / "doc.pdf"
|
||||
pdf.write_bytes(build_pdf(["x"]))
|
||||
result = flash_api.page_index_flash(str(pdf), summary=False)
|
||||
assert calls == {"pages": [], "do_expand": False}
|
||||
assert result["optimize"] == {"merges": 0}
|
||||
|
||||
|
||||
def test_optimize_full_skips_expand_on_textless_pages(tmp_path, monkeypatch):
|
||||
"""Scanned PDFs yield page_texts of empty strings; expand still skips —
|
||||
proposals against empty text are all rejected, so the calls are waste."""
|
||||
from conftest import build_pdf
|
||||
from pageindex.flash import api as flash_api
|
||||
|
||||
monkeypatch.setenv("OPENAI_API_KEY", "k")
|
||||
calls = {}
|
||||
|
||||
def fake_optimize(structure, pages, do_expand, model):
|
||||
calls["do_expand"] = do_expand
|
||||
return {"merges": 0}
|
||||
|
||||
monkeypatch.setattr(flash_api, "_optimize", fake_optimize)
|
||||
monkeypatch.setattr(flash_api, "extract_toc",
|
||||
lambda pdf, use_embedded_toc=True: {
|
||||
"structure": [{"title": "T", "start_index": 1,
|
||||
"end_index": 2, "nodes": []}],
|
||||
"page_texts": ["", ""]})
|
||||
pdf = tmp_path / "doc.pdf"
|
||||
pdf.write_bytes(build_pdf(["x"]))
|
||||
result = flash_api.page_index_flash(str(pdf), summary=False)
|
||||
assert calls == {"do_expand": False}
|
||||
assert result["optimize"] == {"merges": 0}
|
||||
|
||||
|
||||
def test_optimize_expand_warning_names_the_behavior_change(tmp_path,
|
||||
monkeypatch):
|
||||
"""The deprecation must say the optimize pass now runs, not just that
|
||||
the parameter was renamed."""
|
||||
from conftest import build_pdf
|
||||
from pageindex.flash import api as flash_api
|
||||
|
||||
monkeypatch.setattr(flash_api, "extract_toc",
|
||||
lambda pdf, use_embedded_toc=True: {"structure": []})
|
||||
pdf = tmp_path / "doc.pdf"
|
||||
pdf.write_bytes(build_pdf(["x"]))
|
||||
with pytest.warns(DeprecationWarning, match="now runs"):
|
||||
flash_api.page_index_flash(str(pdf), summary=False,
|
||||
optimize_expand=False)
|
||||
|
||||
|
||||
# ── run_pageindex.py flash branch (twelfth review) ──
|
||||
|
||||
SCRIPT = Path(__file__).resolve().parent.parent / "run_pageindex.py"
|
||||
|
||||
|
||||
def _run_flash_cli(monkeypatch, tmp_path, argv, structure):
|
||||
"""Drive run_pageindex.py in-process with a stubbed flash indexer."""
|
||||
import runpy
|
||||
import sys
|
||||
|
||||
import pageindex.flash
|
||||
|
||||
pdf = tmp_path / "t.pdf"
|
||||
pdf.write_bytes(b"%PDF-1.4 stub")
|
||||
captured = {}
|
||||
|
||||
def fake_flash(path, **kw):
|
||||
captured.update(kw)
|
||||
return {"structure": structure}
|
||||
monkeypatch.setattr(pageindex.flash, "page_index_flash", fake_flash)
|
||||
monkeypatch.setattr(sys, "argv",
|
||||
["run_pageindex.py", "--pdf_path", str(pdf), *argv])
|
||||
monkeypatch.chdir(tmp_path)
|
||||
runpy.run_path(str(SCRIPT), run_name="__main__")
|
||||
return captured
|
||||
|
||||
|
||||
def test_flash_cli_summary_model_follows_config_chain(monkeypatch, tmp_path):
|
||||
"""--model must not outrank a file-supplied summary_model: the flash
|
||||
branch resolves through ConfigLoader's chain like the standard and
|
||||
markdown branches, and like the --summary-model help promises."""
|
||||
import pageindex.utils as U
|
||||
|
||||
cfg = tmp_path / "config.yaml"
|
||||
cfg.write_text((Path(U.__file__).parent / "config.yaml").read_text()
|
||||
+ "\nsummary_model: yaml-summary\n")
|
||||
orig = U.ConfigLoader.__init__
|
||||
monkeypatch.setattr(U.ConfigLoader, "__init__",
|
||||
lambda self, default_path=None: orig(self, str(cfg)))
|
||||
captured = _run_flash_cli(
|
||||
monkeypatch, tmp_path, ["--model", "cli-model"],
|
||||
[{"title": "T", "start_index": 1, "end_index": 1}])
|
||||
assert captured["summary_model"] == "yaml-summary"
|
||||
assert captured["optimize_model"] == "yaml-summary"
|
||||
|
||||
|
||||
def test_flash_cli_rejects_empty_structure(monkeypatch, tmp_path):
|
||||
"""A PDF flash cannot structure must error like the SDK does, not write
|
||||
"structure": [] and exit 0 with a success message."""
|
||||
with pytest.raises(ValueError, match="try --mode standard"):
|
||||
_run_flash_cli(monkeypatch, tmp_path, [], [])
|
||||
assert not (tmp_path / "results").exists()
|
||||
|
||||
+186
-42
@@ -8,6 +8,11 @@ import types
|
||||
import httpx # via the hard `openai` dependency
|
||||
import pytest
|
||||
|
||||
try: # anthropic >= 1.0 validates http_client against httpx2
|
||||
import httpx2 as anthropic_httpx
|
||||
except ImportError: # older anthropic rides classic httpx
|
||||
anthropic_httpx = httpx
|
||||
|
||||
import pageindex.local_chat as local_chat
|
||||
from pageindex import (PageIndexAPIError, PageIndexCloudClient,
|
||||
PageIndexLocalClient)
|
||||
@@ -688,14 +693,15 @@ def fake_anthropic(monkeypatch):
|
||||
state["calls"].append(json.loads(request.content))
|
||||
body = responses[len(state["calls"]) - 1]
|
||||
if isinstance(body, str): # pre-rendered SSE
|
||||
return httpx.Response(
|
||||
return anthropic_httpx.Response(
|
||||
200, content=body.encode(),
|
||||
headers={"content-type": "text/event-stream"})
|
||||
return httpx.Response(200, json=body)
|
||||
return anthropic_httpx.Response(200, json=body)
|
||||
|
||||
fake = anthropic.Anthropic(
|
||||
api_key="test",
|
||||
http_client=httpx.Client(transport=httpx.MockTransport(handler)))
|
||||
http_client=anthropic_httpx.Client(
|
||||
transport=anthropic_httpx.MockTransport(handler)))
|
||||
monkeypatch.setattr(local_chat, "_anthropic_client",
|
||||
lambda backend=None: fake)
|
||||
return state["calls"]
|
||||
@@ -924,13 +930,20 @@ def test_responses_envelope_fields_and_cache_group(client, store_path,
|
||||
|
||||
def test_sol_class_refusal_names_its_exits():
|
||||
"""The chatcmpl+tools-while-reasoning 400 is a lane problem, not a
|
||||
retry problem — the wrapped error must name every exit."""
|
||||
err = local_chat._model_backend_error(Exception(
|
||||
retry problem — the wrapped error must name every exit that is real
|
||||
for the caller's lane. Chat has three; a responses() caller gets only
|
||||
the litellm upgrade (it IS the other lane, and its reasoning knob is
|
||||
``reasoning``, not ``reasoning_effort``)."""
|
||||
refusal = Exception(
|
||||
"Error code: 400 - Function tools with reasoning_effort are not "
|
||||
"supported for gpt-5.6-sol in /v1/chat/completions."))
|
||||
assert "responses()" in str(err) and "litellm" in str(err)
|
||||
assert "pass reasoning_effort" in str(err)
|
||||
plain = local_chat._model_backend_error(Exception("rate limited"))
|
||||
"supported for gpt-5.6-sol in /v1/chat/completions.")
|
||||
chat = str(local_chat._model_backend_error(refusal, "chat"))
|
||||
assert "responses()" in chat and "litellm" in chat
|
||||
assert "pass reasoning_effort" in chat
|
||||
resp = str(local_chat._model_backend_error(refusal, "responses"))
|
||||
assert "upgrade litellm" in resp
|
||||
assert "responses()" not in resp and "pass reasoning_effort" not in resp
|
||||
plain = local_chat._model_backend_error(Exception("rate limited"), "chat")
|
||||
assert "responses()" not in str(plain)
|
||||
|
||||
|
||||
@@ -1224,14 +1237,15 @@ def test_envelope_model_strips_openai_routing_prefix(store_path, fake_model):
|
||||
|
||||
|
||||
@needs_agents
|
||||
def test_chat_missing_openai_key_fails_loud(monkeypatch):
|
||||
"""A missing backend credential surfaces as the SDK's own error type,
|
||||
like every other precondition on the chat surfaces."""
|
||||
def test_chat_model_builds_keyless_without_prejudgment(monkeypatch):
|
||||
"""No key pre-judgment at model build: credentials are LiteLLM's call
|
||||
at run time, so a keyless build succeeds for every spelling."""
|
||||
pytest.importorskip("litellm") # first import may load a .env; delenv after
|
||||
from agents.extensions.models.litellm_model import LitellmModel
|
||||
monkeypatch.delenv("OPENAI_API_KEY", raising=False)
|
||||
for name in ("gpt-4o", "openai/gpt-4o"):
|
||||
with pytest.raises(PageIndexAPIError, match="OPENAI_API_KEY"):
|
||||
local_chat._openai_model("chat", name)
|
||||
model = local_chat._openai_model("chat", name)
|
||||
assert isinstance(model, LitellmModel)
|
||||
|
||||
|
||||
@needs_agents
|
||||
@@ -1432,13 +1446,14 @@ def test_messages_provider_errors_wrap_as_sdk_errors(client, store_path,
|
||||
seed_doc(store_path, "pi-a", "report.pdf")
|
||||
|
||||
def handler(request):
|
||||
return httpx.Response(429, json={
|
||||
return anthropic_httpx.Response(429, json={
|
||||
"type": "error",
|
||||
"error": {"type": "rate_limit_error", "message": "slow down"}})
|
||||
|
||||
fake = anthropic.Anthropic(
|
||||
api_key="test", max_retries=0,
|
||||
http_client=httpx.Client(transport=httpx.MockTransport(handler)))
|
||||
http_client=anthropic_httpx.Client(
|
||||
transport=anthropic_httpx.MockTransport(handler)))
|
||||
monkeypatch.setattr(local_chat, "_anthropic_client",
|
||||
lambda backend=None: fake)
|
||||
with pytest.raises(PageIndexAPIError, match="model backend failed"):
|
||||
@@ -1824,11 +1839,12 @@ def test_messages_extra_headers_reach_the_wire(client, monkeypatch):
|
||||
|
||||
def handler(request):
|
||||
seen["beta"] = request.headers.get("anthropic-beta")
|
||||
return httpx.Response(200, json=_anthropic_message(
|
||||
return anthropic_httpx.Response(200, json=_anthropic_message(
|
||||
[{"type": "text", "text": "ok"}], "end_turn"))
|
||||
|
||||
fake = anthropic.Anthropic(api_key="t", http_client=httpx.Client(
|
||||
transport=httpx.MockTransport(handler)))
|
||||
fake = anthropic.Anthropic(
|
||||
api_key="t", http_client=anthropic_httpx.Client(
|
||||
transport=anthropic_httpx.MockTransport(handler)))
|
||||
monkeypatch.setattr(local_chat, "_anthropic_client",
|
||||
lambda backend=None: fake)
|
||||
client.messages("q", model="claude-sonnet-4-5",
|
||||
@@ -1865,6 +1881,56 @@ def test_messages_default_max_tokens_clears_thinking_budget(client, fake_anthrop
|
||||
assert calls[0]["max_tokens"] == 11000 # explicit value passes through
|
||||
|
||||
|
||||
@needs_anthropic
|
||||
def test_anthropic_client_cached_per_backend(monkeypatch):
|
||||
"""One real client per backend: construction pays ~45ms of SSL-context
|
||||
build and a cold connection pool each call otherwise."""
|
||||
monkeypatch.setattr(local_chat, "_ANTHROPIC_CLIENTS", {})
|
||||
a = local_chat._anthropic_client({"api_key": "k"})
|
||||
assert local_chat._anthropic_client({"api_key": "k"}) is a
|
||||
assert local_chat._anthropic_client({"api_key": "k2"}) is not a
|
||||
|
||||
|
||||
@needs_anthropic
|
||||
def test_anthropic_client_construction_race_keeps_first(monkeypatch):
|
||||
"""A constructor losing the store race must adopt the winner rather
|
||||
than evict a client other threads may already hold."""
|
||||
monkeypatch.setattr(local_chat, "_ANTHROPIC_CLIENTS", {})
|
||||
winner = object()
|
||||
real = anthropic.Anthropic
|
||||
|
||||
def racing(**kwargs):
|
||||
local_chat._ANTHROPIC_CLIENTS[(("api_key", "k"),)] = winner
|
||||
return real(**kwargs)
|
||||
monkeypatch.setattr(anthropic, "Anthropic", racing)
|
||||
assert local_chat._anthropic_client({"api_key": "k"}) is winner
|
||||
|
||||
|
||||
@needs_anthropic
|
||||
def test_messages_reuses_cached_client_across_runs(client, monkeypatch):
|
||||
"""A cached backend client survives the per-run close: two consecutive
|
||||
runs ride the same client (a closed one refuses the second request),
|
||||
and the cache hit constructs nothing."""
|
||||
def handler(request):
|
||||
return anthropic_httpx.Response(
|
||||
200, json=_anthropic_message([{"type": "text", "text": "ok"}],
|
||||
"end_turn"))
|
||||
cached = anthropic.Anthropic(
|
||||
api_key="test",
|
||||
http_client=anthropic_httpx.Client(
|
||||
transport=anthropic_httpx.MockTransport(handler)))
|
||||
monkeypatch.setattr(local_chat, "_ANTHROPIC_CLIENTS",
|
||||
{(("api_key", "test"),): cached})
|
||||
|
||||
def boom(**kwargs):
|
||||
raise AssertionError("cache hit expected — no new construction")
|
||||
monkeypatch.setattr(anthropic, "Anthropic", boom)
|
||||
for _ in range(2):
|
||||
result = client.messages("q", model="claude-test", max_tokens=64,
|
||||
backend={"api_key": "test"})
|
||||
assert result["stop_reason"] == "end_turn"
|
||||
|
||||
|
||||
def test_record_chat_finish_records_and_delegates():
|
||||
recorded = {}
|
||||
closed = {"n": 0}
|
||||
@@ -1939,10 +2005,10 @@ def test_chat_completions_reports_native_finish_reason(client, store_path,
|
||||
|
||||
|
||||
@needs_agents
|
||||
def test_chat_gate_honors_litellm_routing_and_custom_providers(monkeypatch):
|
||||
"""Mirrors the indexing lane: an explicit litellm/ prefix skips the env
|
||||
key pre-check, and custom_provider_map providers pass the allowlist;
|
||||
a name LiteLLM cannot route is still refused up front."""
|
||||
def test_chat_model_honors_litellm_routing_and_custom_providers(monkeypatch):
|
||||
"""litellm/ spellings and custom_provider_map providers pass the
|
||||
provider allowlist; a name LiteLLM cannot route is still refused up
|
||||
front."""
|
||||
pytest.importorskip("litellm")
|
||||
import litellm # first import may load a .env; delenv after it
|
||||
from agents.extensions.models.litellm_model import LitellmModel
|
||||
@@ -1984,19 +2050,6 @@ def test_openai_protocol_predicate_follows_litellm_routing():
|
||||
assert not local_chat._openai_protocol(name), name
|
||||
|
||||
|
||||
@needs_agents
|
||||
def test_chat_backend_without_key_stands_aside_like_index_lane(monkeypatch):
|
||||
"""Any non-empty backend dict suppresses the key pre-check (utils rule)."""
|
||||
pytest.importorskip("litellm")
|
||||
import litellm # noqa: F401 — first import may load a .env; delenv after it
|
||||
from agents.extensions.models.litellm_model import LitellmModel
|
||||
|
||||
monkeypatch.delenv("OPENAI_API_KEY", raising=False)
|
||||
model = local_chat._openai_model(
|
||||
"chat", "gpt-test", {"base_url": "http://localhost:9"})
|
||||
assert isinstance(model, LitellmModel)
|
||||
|
||||
|
||||
@needs_agents
|
||||
def test_responses_model_marks_caller_owned_transport(monkeypatch):
|
||||
monkeypatch.setenv("OPENAI_API_KEY", "sk-test")
|
||||
@@ -2019,8 +2072,8 @@ def test_responses_model_marks_caller_owned_transport(monkeypatch):
|
||||
@needs_anthropic
|
||||
def test_messages_keeps_caller_owned_http_client_open(client):
|
||||
body = _anthropic_message([{"type": "text", "text": "a"}], "end_turn")
|
||||
shared = httpx.Client(transport=httpx.MockTransport(
|
||||
lambda request: httpx.Response(200, json=body)))
|
||||
shared = anthropic_httpx.Client(transport=anthropic_httpx.MockTransport(
|
||||
lambda request: anthropic_httpx.Response(200, json=body)))
|
||||
out = client.messages("q", model="claude-test",
|
||||
backend={"api_key": "t", "http_client": shared})
|
||||
assert out["content"][0]["text"] == "a"
|
||||
@@ -2032,9 +2085,100 @@ def test_messages_keeps_caller_owned_http_client_open(client):
|
||||
|
||||
@needs_anthropic
|
||||
def test_messages_without_credentials_raises_contract_error(client,
|
||||
monkeypatch):
|
||||
monkeypatch,
|
||||
tmp_path):
|
||||
"""No pre-check: the SDK's own request-time credential-resolution
|
||||
failure is translated into the contract's PageIndexAPIError — for a
|
||||
bare call, a credential-less backend dict, and the unset-env-var
|
||||
shape ({"api_key": None}) alike."""
|
||||
monkeypatch.delenv("ANTHROPIC_API_KEY", raising=False)
|
||||
monkeypatch.delenv("ANTHROPIC_AUTH_TOKEN", raising=False)
|
||||
with pytest.raises(PageIndexAPIError,
|
||||
match="Anthropic backend is not configured"):
|
||||
client.messages("q", model="claude-test")
|
||||
monkeypatch.delenv("ANTHROPIC_PROFILE", raising=False)
|
||||
monkeypatch.setenv("HOME", str(tmp_path)) # no ant-auth profile fallback
|
||||
for backend in (None, {"timeout": 30}, {"api_key": None}):
|
||||
with pytest.raises(PageIndexAPIError,
|
||||
match="Anthropic backend is not configured"):
|
||||
client.messages("q", model="claude-test", backend=backend)
|
||||
|
||||
|
||||
def test_cache_extra_args_follow_wire_normalization():
|
||||
"""Chat lane only: litellm/<bare-name> rides the OpenAI protocol on
|
||||
this wire (bare names get openai/), so it must carry no Anthropic
|
||||
cache marks; explicit anthropic routes keep them. The agents lane
|
||||
routes the same spelling to Anthropic and marks it — see
|
||||
test_openai_agent_config_marks_bare_claude_behind_litellm_prefix."""
|
||||
pytest.importorskip("litellm")
|
||||
assert local_chat._cache_extra_args("litellm/claude-sonnet-4-5") is None
|
||||
assert local_chat._cache_extra_args("anthropic/claude-x") is not None
|
||||
assert local_chat._cache_extra_args("litellm/anthropic/claude-x") is not None
|
||||
|
||||
|
||||
@needs_agents
|
||||
def test_bad_model_settings_wrap_as_contract_error(client):
|
||||
"""A mistyped sampling param is the SDK's verdict (pydantic),
|
||||
translated into the contract's PageIndexAPIError like every other
|
||||
door failure."""
|
||||
with pytest.raises(PageIndexAPIError, match="Invalid model settings"):
|
||||
client.chat_completions("q", temperature="hot")
|
||||
|
||||
|
||||
@needs_agents
|
||||
def test_translate_run_error_routes_all_three_kinds():
|
||||
"""The shared ladder every agent door delegates to: max_turns guidance
|
||||
first (a MaxTurnsExceeded is also an AgentsException), then the
|
||||
agents-framework wrap, then the model-backend wrap."""
|
||||
import openai
|
||||
from agents.exceptions import AgentsException, MaxTurnsExceeded
|
||||
|
||||
assert "max_turns (3)" in str(
|
||||
local_chat._translate_run_error(MaxTurnsExceeded("over"), 3, "chat"))
|
||||
assert "agent backend failed" in str(
|
||||
local_chat._translate_run_error(AgentsException("boom"), None,
|
||||
"responses"))
|
||||
assert "model backend failed" in str(
|
||||
local_chat._translate_run_error(openai.OpenAIError("down"), None,
|
||||
"chat"))
|
||||
|
||||
|
||||
def test_cache_marks_counts_system_and_message_blocks():
|
||||
"""The counter guards the API's 4-breakpoint limit; marks live on
|
||||
system blocks and on message content blocks, never on plain strings."""
|
||||
system = [{"type": "text", "text": "s",
|
||||
"cache_control": {"type": "ephemeral"}},
|
||||
{"type": "text", "text": "t"}]
|
||||
messages = [
|
||||
{"role": "user", "content": "plain strings carry no marks"},
|
||||
{"role": "user", "content": [
|
||||
{"type": "text", "text": "a",
|
||||
"cache_control": {"type": "ephemeral"}},
|
||||
{"type": "text", "text": "b"}]},
|
||||
]
|
||||
assert local_chat._cache_marks(system, messages) == 2
|
||||
assert local_chat._cache_marks([], []) == 0
|
||||
|
||||
|
||||
def test_dump_block_omits_unset_response_defaults():
|
||||
"""messages() tells the caller to append result["messages"] verbatim;
|
||||
a response-only default like tool_use's caller must not surface as an
|
||||
explicit null the request schema has no variant for."""
|
||||
pytest.importorskip("anthropic")
|
||||
from anthropic.types.beta import BetaToolUseBlock
|
||||
block = BetaToolUseBlock(id="tu_1", input={}, name="t", type="tool_use")
|
||||
assert local_chat._dump_block(block) == {
|
||||
"id": "tu_1", "input": {}, "name": "t", "type": "tool_use"}
|
||||
|
||||
|
||||
def test_default_max_tokens_respects_output_ceilings():
|
||||
"""A lifted thinking default must not overshoot the model's output
|
||||
ceiling — the wire rejects max_tokens above it; bool is not a budget."""
|
||||
lift = local_chat._default_max_tokens
|
||||
enabled = {"type": "enabled", "budget_tokens": 30000}
|
||||
assert lift("claude-opus-4-1", enabled) == 32000
|
||||
assert lift("claude-sonnet-4-5-20250929",
|
||||
{"type": "enabled", "budget_tokens": 60000}) == 64000
|
||||
assert lift("claude-opus-4-1",
|
||||
{"type": "enabled", "budget_tokens": 10000}) == 18192
|
||||
assert lift("claude-test",
|
||||
{"type": "enabled", "budget_tokens": 10000}) == 18192
|
||||
assert lift("claude-sonnet-4-5",
|
||||
{"type": "enabled", "budget_tokens": True}) == 8192
|
||||
|
||||
@@ -35,8 +35,11 @@ class MarkdownCliTest(unittest.TestCase):
|
||||
with tempfile.TemporaryDirectory() as tmp:
|
||||
md = Path(tmp) / "notes.md"
|
||||
md.write_text("# Title\n\nIntro.\n\n## Section\n\nBody.\n")
|
||||
env = {k: v for k, v in os.environ.items()
|
||||
if k not in ("OPENAI_API_KEY", "CHATGPT_API_KEY")}
|
||||
env = dict(os.environ)
|
||||
# present-but-empty beats deletion: utils' load_dotenv() does
|
||||
# not override existing vars, so the repo .env key stays out
|
||||
env["OPENAI_API_KEY"] = ""
|
||||
env["CHATGPT_API_KEY"] = ""
|
||||
env["PYTHONPATH"] = str(script.parent)
|
||||
res = subprocess.run(
|
||||
[sys.executable, str(script), "--md_path", str(md)],
|
||||
@@ -46,6 +49,57 @@ class MarkdownCliTest(unittest.TestCase):
|
||||
self.assertTrue(out.exists(), res.stdout.decode())
|
||||
json.loads(out.read_text())
|
||||
|
||||
def test_md_cli_summary_model_drives_summary_calls(self):
|
||||
"""--summary-model owns the markdown summary lane: node summaries
|
||||
and the doc description bill it, never the index model given
|
||||
alongside — the same chain the flag's help promises on PDFs."""
|
||||
import json
|
||||
import os
|
||||
import subprocess
|
||||
import sys
|
||||
import tempfile
|
||||
from pathlib import Path
|
||||
|
||||
script = Path(__file__).resolve().parent.parent / "run_pageindex.py"
|
||||
driver = (
|
||||
"import json, runpy, sys\n"
|
||||
"import pageindex.utils as U\n"
|
||||
"seen = []\n"
|
||||
"async def fake_acompletion(model, prompt, **kw):\n"
|
||||
" seen.append(model)\n"
|
||||
" return 'node summary'\n"
|
||||
"def fake_completion(model, prompt, **kw):\n"
|
||||
" seen.append(model)\n"
|
||||
" return 'doc description'\n"
|
||||
"U.llm_acompletion = fake_acompletion\n"
|
||||
"U.llm_completion = fake_completion\n"
|
||||
"target = sys.argv[1]\n"
|
||||
"sys.argv = [target] + sys.argv[2:]\n"
|
||||
"runpy.run_path(target, run_name='__main__')\n"
|
||||
"print('MODELS_SEEN=' + json.dumps(sorted(set(seen))))\n"
|
||||
)
|
||||
with tempfile.TemporaryDirectory() as tmp:
|
||||
md = Path(tmp) / "notes.md"
|
||||
md.write_text("# Title\n\nIntro.\n\n## Section\n\nBody.\n")
|
||||
drv = Path(tmp) / "driver.py"
|
||||
drv.write_text(driver)
|
||||
env = dict(os.environ)
|
||||
env["PYTHONPATH"] = str(script.parent)
|
||||
res = subprocess.run(
|
||||
[sys.executable, str(drv), str(script),
|
||||
"--md_path", str(md),
|
||||
"--if-add-node-summary", "yes",
|
||||
"--if-add-doc-description", "yes",
|
||||
"--summary-token-threshold", "1",
|
||||
"--summary-model", "SUMMARY-SENTINEL",
|
||||
"--index-model", "INDEX-DECOY"],
|
||||
capture_output=True, cwd=tmp, env=env, timeout=180)
|
||||
self.assertEqual(res.returncode, 0, res.stderr.decode())
|
||||
line = next(l for l in res.stdout.decode().splitlines()
|
||||
if l.startswith("MODELS_SEEN="))
|
||||
self.assertEqual(json.loads(line[len("MODELS_SEEN="):]),
|
||||
["SUMMARY-SENTINEL"], res.stdout.decode())
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
|
||||
Reference in New Issue
Block a user