mirror of
https://github.com/VectifyAI/PageIndex.git
synced 2026-07-15 21:11:05 +02:00
fix: address xhigh code-review findings on 2d46d68..8f536cb
Verified 12 findings from an xhigh-effort review of the prior review-fix
batch; all confirmed real. Most trace back to one root cause: the
build_index() text-stripping fix (8f536cb) correctly stopped Markdown from
leaking full text by default, but broke every path that assumed text could
be re-read later.
Correctness:
- LocalBackend._fill_node_text (get_document(include_text=True)) only handled
PDF's start_index/end_index convention; Markdown nodes use line_num and got
silently empty text. Now handles both.
- get_page_content's Markdown fallback (triggered when a StorageEngine
legitimately returns None from get_pages()) read from the now-text-stripped
structure. It now re-derives from the source file, mirroring the PDF
fallback, so it no longer depends on structure text at all.
- add_document's PDF-only text-stripping branch (with the stale "markdown
needs text in structure for fallback retrieval" comment) is now dead/wrong
since build_index() already applies if_add_node_text uniformly — removed.
- _validate_llm_provider's keyless-provider allowlist was missing several
local LiteLLM providers (xinference, llamafile, triton, oobabooga,
openai_like, docker_model_runner, custom, custom_openai, petals) that need
no API key just like ollama/lm_studio; expanded.
- The three agent-tool closures (get_document, get_document_structure,
get_page_content) had three different not-found patterns; two bypassed the
backend's DocumentNotFoundError entirely. Extracted LocalBackend.
_require_document as the single existence check every method/tool now uses.
- examples/agentic_vectorless_rag_demo.py's hand-rolled Agent() didn't apply
the litellm/ prefix normalization the SDK does internally, so its own
documented "any LiteLLM provider" claim broke for non-openai models.
- cloud delete_collection's cache eviction removed the "folders unavailable"
None sentinel too, forcing a wasted re-fetch; now only pops on a real id.
Cleanup / altitude:
- build_index() skips the remove_structure_text walk entirely when text was
never added (content_based + if_add_node_summary=False + if_add_node_text=
False) instead of a guaranteed no-op tree walk.
- page_index()'s locals()-capture-as-kwargs (fragile by construction) replaced
with an explicit dict of the named parameters.
- run_pageindex.py's _cli_bool and the page_index_md.py legacy shim's
_coerce_bool were duplicate, diverging implementations; both now bind
directly to the canonical pageindex.index.page_index_md._coerce_bool.
- retrieve.py's _get_md_page_content delegated its own traversal instead of
calling the canonical get_md_page_content; now a one-line delegation.
- FileTypeError's docstring now calls out the except-ordering gotcha from
also subclassing ValueError.
17 new regression tests (tests/test_review_fixes_2.py) plus 2 updated in
tests/test_legacy_shims.py for the simplified md_to_tree shim. Full suite:
210 passed, 2 skipped.
Claude-Session: https://claude.ai/code/session_01Kx5DgKbhK1N8autqXH8SmS
This commit is contained in:
parent
8f536cb8d7
commit
04cb9cb02d
12 changed files with 333 additions and 124 deletions
|
|
@ -58,6 +58,16 @@ Answer based only on tool output. Be concise.
|
|||
"""
|
||||
|
||||
|
||||
def _normalize_model_for_agents_sdk(model: str) -> str:
|
||||
"""The OpenAI Agents SDK only recognizes 'openai/' and 'litellm/' model
|
||||
prefixes; route any other LiteLLM-style provider path (e.g. 'anthropic/...')
|
||||
through litellm explicitly, mirroring what PageIndex itself does internally
|
||||
for its built-in agent."""
|
||||
if model and "/" in model and not model.startswith(("litellm/", "openai/")):
|
||||
return f"litellm/{model}"
|
||||
return model
|
||||
|
||||
|
||||
def query_agent(col, doc_id: str, prompt: str, model: str, verbose: bool = False) -> str:
|
||||
"""Run a document QA agent using the OpenAI Agents SDK.
|
||||
|
||||
|
|
@ -90,7 +100,7 @@ def query_agent(col, doc_id: str, prompt: str, model: str, verbose: bool = False
|
|||
name="PageIndex",
|
||||
instructions=AGENT_SYSTEM_PROMPT,
|
||||
tools=[get_document, get_document_structure, get_page_content],
|
||||
model=model,
|
||||
model=_normalize_model_for_agents_sdk(model),
|
||||
# model_settings=ModelSettings(reasoning={"effort": "low", "summary": "auto"}), # Uncomment to enable reasoning
|
||||
)
|
||||
|
||||
|
|
|
|||
|
|
@ -162,9 +162,12 @@ class CloudBackend:
|
|||
folder_id = self._get_folder_id(name)
|
||||
if folder_id:
|
||||
self._request("DELETE", f"/folder/{self._enc(folder_id)}/")
|
||||
# Drop the cached id so a later same-name op re-resolves instead of
|
||||
# reusing the now-deleted folder_id.
|
||||
self._folder_id_cache.pop(name, None)
|
||||
# Drop the cached id so a later same-name op re-resolves instead of
|
||||
# reusing the now-deleted folder_id. Only when it was a REAL id —
|
||||
# if folder_id was falsy, the cache holds the "folders unavailable
|
||||
# on this plan" None sentinel, which must survive so we don't
|
||||
# re-issue a doomed GET /folders/ on the next call.
|
||||
self._folder_id_cache.pop(name, None)
|
||||
|
||||
# ── Document management ───────────────────────────────────────────────
|
||||
|
||||
|
|
|
|||
|
|
@ -12,7 +12,7 @@ from ..parser.pdf import PdfParser
|
|||
from ..parser.markdown import MarkdownParser
|
||||
from ..storage.protocol import StorageEngine
|
||||
from ..index.pipeline import build_index
|
||||
from ..index.utils import parse_pages, get_pdf_page_content, get_md_page_content, remove_fields
|
||||
from ..index.utils import parse_pages, get_pdf_page_content, remove_fields
|
||||
from ..backend.protocol import AgentTools
|
||||
from ..errors import (FileTypeError, DocumentNotFoundError, CollectionNotFoundError,
|
||||
IndexingError, PageIndexError)
|
||||
|
|
@ -116,26 +116,23 @@ class LocalBackend:
|
|||
parsed = parser.parse(file_path, model=self._model, images_dir=images_dir)
|
||||
result = build_index(parsed, model=self._model, opt=self._index_config)
|
||||
|
||||
# Cache page text for fast retrieval (avoids re-reading files)
|
||||
# Cache page text for fast retrieval (avoids re-reading files) and to
|
||||
# reconstruct node text on demand (get_document(include_text=True),
|
||||
# get_page_content fallback) independent of whether IndexConfig kept
|
||||
# text in the stored structure. build_index() already applies
|
||||
# if_add_node_text to result["structure"] for every strategy, so no
|
||||
# extra stripping is needed here.
|
||||
pages = [{"page": n.index, "content": n.content,
|
||||
**({"images": n.images} if n.images else {})}
|
||||
for n in parsed.nodes if n.content]
|
||||
|
||||
# Strip text from structure to save storage space (PDF only;
|
||||
# markdown needs text in structure for fallback retrieval)
|
||||
doc_type = ext.lstrip(".")
|
||||
if doc_type == "pdf":
|
||||
clean_structure = remove_fields(result["structure"], fields=["text"])
|
||||
else:
|
||||
clean_structure = result["structure"]
|
||||
|
||||
self._storage.save_document(collection, doc_id, {
|
||||
"doc_name": parsed.doc_name,
|
||||
"doc_description": result.get("doc_description", ""),
|
||||
"file_path": str(managed_path),
|
||||
"file_hash": file_hash,
|
||||
"doc_type": doc_type,
|
||||
"structure": clean_structure,
|
||||
"doc_type": ext.lstrip("."),
|
||||
"structure": result["structure"],
|
||||
"pages": pages,
|
||||
})
|
||||
except sqlite3.IntegrityError:
|
||||
|
|
@ -158,6 +155,19 @@ class LocalBackend:
|
|||
|
||||
return doc_id
|
||||
|
||||
def _require_document(self, collection: str, doc_id: str) -> dict:
|
||||
"""Return the document's storage row, or raise DocumentNotFoundError.
|
||||
|
||||
Single source of truth for "does this doc exist" — every public method
|
||||
and agent tool below goes through this, so a missing doc always
|
||||
surfaces the same way instead of each caller re-implementing its own
|
||||
(and potentially inconsistent) existence check.
|
||||
"""
|
||||
doc = self._storage.get_document(collection, doc_id)
|
||||
if not doc:
|
||||
raise DocumentNotFoundError(f"Document {doc_id} not found")
|
||||
return doc
|
||||
|
||||
def get_document(self, collection: str, doc_id: str, include_text: bool = False) -> dict:
|
||||
"""Get document metadata with structure.
|
||||
|
||||
|
|
@ -166,9 +176,7 @@ class LocalBackend:
|
|||
from cached page content. WARNING: may be very large — do NOT
|
||||
use in agent/LLM contexts as it can exhaust the context window.
|
||||
"""
|
||||
doc = self._storage.get_document(collection, doc_id)
|
||||
if not doc:
|
||||
raise DocumentNotFoundError(f"Document {doc_id} not found")
|
||||
doc = self._require_document(collection, doc_id)
|
||||
doc["structure"] = self._storage.get_document_structure(collection, doc_id)
|
||||
if include_text:
|
||||
pages = self._storage.get_pages(collection, doc_id) or []
|
||||
|
|
@ -178,7 +186,13 @@ class LocalBackend:
|
|||
|
||||
@staticmethod
|
||||
def _fill_node_text(nodes: list, page_map: dict) -> None:
|
||||
"""Recursively fill 'text' on structure nodes from cached page content."""
|
||||
"""Recursively fill 'text' on structure nodes from cached page content.
|
||||
|
||||
Two node conventions, one per indexing strategy: content_based (PDF)
|
||||
nodes span a start_index..end_index page range; level_based (Markdown)
|
||||
nodes map 1:1 to a single page keyed by line_num. Handling only the
|
||||
first would silently leave Markdown nodes with no text.
|
||||
"""
|
||||
for node in nodes:
|
||||
start = node.get("start_index")
|
||||
end = node.get("end_index")
|
||||
|
|
@ -186,20 +200,17 @@ class LocalBackend:
|
|||
node["text"] = "\n".join(
|
||||
page_map.get(p, "") for p in range(start, end + 1)
|
||||
)
|
||||
elif "line_num" in node:
|
||||
node["text"] = page_map.get(node["line_num"], "")
|
||||
if "nodes" in node:
|
||||
LocalBackend._fill_node_text(node["nodes"], page_map)
|
||||
|
||||
def get_document_structure(self, collection: str, doc_id: str) -> list:
|
||||
# Parity with get_document / the cloud backend: a missing doc must raise,
|
||||
# not masquerade as an empty structure.
|
||||
if not self._storage.get_document(collection, doc_id):
|
||||
raise DocumentNotFoundError(f"Document {doc_id} not found")
|
||||
self._require_document(collection, doc_id)
|
||||
return self._storage.get_document_structure(collection, doc_id)
|
||||
|
||||
def get_page_content(self, collection: str, doc_id: str, pages: str) -> list:
|
||||
doc = self._storage.get_document(collection, doc_id)
|
||||
if not doc:
|
||||
raise DocumentNotFoundError(f"Document {doc_id} not found")
|
||||
doc = self._require_document(collection, doc_id)
|
||||
page_nums = parse_pages(pages)
|
||||
|
||||
# Try cached pages first (fast, no file I/O)
|
||||
|
|
@ -207,22 +218,24 @@ class LocalBackend:
|
|||
if cached_pages:
|
||||
return [p for p in cached_pages if p["page"] in page_nums]
|
||||
|
||||
# Fallback to reading from file
|
||||
# Fallback: re-derive from the source file, same as the PDF path below
|
||||
# — never from the stored structure, whose 'text' field may have been
|
||||
# stripped (if_add_node_text=False, the default). Reachable only for a
|
||||
# custom StorageEngine that doesn't cache pages (the built-in
|
||||
# SQLiteStorage always does).
|
||||
if doc["doc_type"] == "pdf":
|
||||
return get_pdf_page_content(doc["file_path"], page_nums)
|
||||
else:
|
||||
structure = self._storage.get_document_structure(collection, doc_id)
|
||||
return get_md_page_content(structure, page_nums)
|
||||
parser = self._resolve_parser(doc["file_path"])
|
||||
parsed = parser.parse(doc["file_path"], model=self._model)
|
||||
page_map = {n.index: n.content for n in parsed.nodes}
|
||||
return [{"page": p, "content": page_map[p]} for p in page_nums if p in page_map]
|
||||
|
||||
def list_documents(self, collection: str) -> list[dict]:
|
||||
return self._storage.list_documents(collection)
|
||||
|
||||
def delete_document(self, collection: str, doc_id: str) -> None:
|
||||
doc = self._storage.get_document(collection, doc_id)
|
||||
if not doc:
|
||||
# Parity with the cloud backend, which surfaces HTTP 404 as
|
||||
# DocumentNotFoundError — a typo'd doc_id should not pass silently.
|
||||
raise DocumentNotFoundError(f"Document {doc_id} not found")
|
||||
doc = self._require_document(collection, doc_id)
|
||||
if doc.get("file_path"):
|
||||
Path(doc["file_path"]).unlink(missing_ok=True)
|
||||
# Clean up images directory: files/{collection}/{doc_id}/
|
||||
|
|
@ -259,8 +272,12 @@ class LocalBackend:
|
|||
rejection = _reject(doc_id)
|
||||
if rejection:
|
||||
return rejection
|
||||
doc = storage.get_document(col_name, doc_id)
|
||||
if not doc:
|
||||
try:
|
||||
# _require_document (not backend.get_document) deliberately:
|
||||
# the metadata-only row, no 'structure' — keeps this tool's
|
||||
# output small for the agent's context window.
|
||||
doc = backend._require_document(col_name, doc_id)
|
||||
except DocumentNotFoundError:
|
||||
return json.dumps({"error": f"doc_id '{doc_id}' not found."})
|
||||
return json.dumps(doc)
|
||||
|
||||
|
|
@ -270,7 +287,9 @@ class LocalBackend:
|
|||
rejection = _reject(doc_id)
|
||||
if rejection:
|
||||
return rejection
|
||||
if not storage.get_document(col_name, doc_id):
|
||||
try:
|
||||
backend._require_document(col_name, doc_id)
|
||||
except DocumentNotFoundError:
|
||||
return json.dumps({"error": f"doc_id '{doc_id}' not found."})
|
||||
structure = storage.get_document_structure(col_name, doc_id)
|
||||
return json.dumps(remove_fields(structure, fields=["text"]), ensure_ascii=False)
|
||||
|
|
|
|||
|
|
@ -122,8 +122,18 @@ class PageIndexClient:
|
|||
except Exception:
|
||||
return
|
||||
|
||||
# LiteLLM providers that run locally and need no API key.
|
||||
keyless = {"ollama", "ollama_chat", "lm_studio", "hosted_vllm", "vllm"}
|
||||
# LiteLLM providers that run locally / self-hosted and need no API key
|
||||
# by default (litellm itself falls back to a placeholder key for these
|
||||
# rather than erroring — see e.g. hosted_vllm's transformation.py).
|
||||
# This list is necessarily a manual allowlist (litellm.validate_environment
|
||||
# isn't reliable enough to derive it from); extend it as litellm adds
|
||||
# more local-inference providers.
|
||||
keyless = {
|
||||
"ollama", "ollama_chat", "lm_studio", "hosted_vllm", "vllm",
|
||||
"xinference", "llamafile", "triton", "oobabooga",
|
||||
"openai_like", "custom_openai", "custom", "docker_model_runner",
|
||||
"petals",
|
||||
}
|
||||
if provider in keyless:
|
||||
return
|
||||
|
||||
|
|
|
|||
|
|
@ -44,5 +44,9 @@ class FileTypeError(PageIndexError, ValueError):
|
|||
|
||||
Also subclasses ValueError so pre-SDK ``except ValueError`` around indexing
|
||||
(0.2.x raised ValueError for an unsupported file format) still catches it.
|
||||
Note: because of this, an ``except ValueError`` clause ahead of an
|
||||
``except FileTypeError`` clause in the same try block will catch it first —
|
||||
if you need FileTypeError-specific handling, put that except before (or
|
||||
instead of) a bare ValueError one.
|
||||
"""
|
||||
pass
|
||||
|
|
|
|||
|
|
@ -1112,15 +1112,24 @@ def page_index_main(doc, opt=None):
|
|||
|
||||
def page_index(doc, model=None, toc_check_page_num=None, max_page_num_each_node=None, max_token_num_each_node=None,
|
||||
if_add_node_id=None, if_add_node_summary=None, if_add_doc_description=None, if_add_node_text=None):
|
||||
|
||||
# Snapshot the call args BEFORE importing IndexConfig — otherwise the
|
||||
# imported class would be captured by locals() and rejected by
|
||||
# IndexConfig(extra="forbid").
|
||||
user_opt = {
|
||||
arg: value for arg, value in locals().items()
|
||||
if arg != "doc" and value is not None
|
||||
}
|
||||
from ..config import IndexConfig
|
||||
|
||||
# Explicit dict of the named kwargs — NOT locals(), which would also
|
||||
# capture any local variable defined above this line (e.g. the IndexConfig
|
||||
# import itself) and get rejected by IndexConfig(extra="forbid"). Unlike a
|
||||
# locals() snapshot, this stays correct regardless of what gets added to
|
||||
# the function body later.
|
||||
user_opt = {
|
||||
"model": model,
|
||||
"toc_check_page_num": toc_check_page_num,
|
||||
"max_page_num_each_node": max_page_num_each_node,
|
||||
"max_token_num_each_node": max_token_num_each_node,
|
||||
"if_add_node_id": if_add_node_id,
|
||||
"if_add_node_summary": if_add_node_summary,
|
||||
"if_add_doc_description": if_add_doc_description,
|
||||
"if_add_node_text": if_add_node_text,
|
||||
}
|
||||
user_opt = {k: v for k, v in user_opt.items() if v is not None}
|
||||
opt = IndexConfig(**user_opt)
|
||||
return page_index_main(doc, opt)
|
||||
|
||||
|
|
|
|||
|
|
@ -110,13 +110,17 @@ def build_index(parsed: ParsedDocument, model: str = None, opt=None) -> dict:
|
|||
clean_structure, model=opt.model
|
||||
)
|
||||
|
||||
# 'text' may have been populated for summary/description generation, or
|
||||
# by build_tree_from_levels for the level_based (Markdown) path. Strip it
|
||||
# LAST, for BOTH strategies, unless explicitly requested — otherwise a
|
||||
# default index leaks each node's full text into get_document_structure /
|
||||
# storage, inconsistent with if_add_node_text=False, the README, and the
|
||||
# legacy md_to_tree.
|
||||
if not opt.if_add_node_text:
|
||||
# 'text' is populated for level_based (Markdown, always) or for
|
||||
# content_based when if_add_node_text/if_add_node_summary requested it.
|
||||
# Strip it LAST, for BOTH strategies, unless explicitly requested —
|
||||
# otherwise a default index leaks each node's full text into
|
||||
# get_document_structure / storage, inconsistent with
|
||||
# if_add_node_text=False, the README, and the legacy md_to_tree. Skip
|
||||
# the walk entirely when text was never added in the first place
|
||||
# (content_based with if_add_node_text=if_add_node_summary=False) —
|
||||
# there's nothing to strip.
|
||||
text_present = strategy == "level_based" or opt.if_add_node_text or opt.if_add_node_summary
|
||||
if text_present and not opt.if_add_node_text:
|
||||
remove_structure_text(structure)
|
||||
|
||||
return result
|
||||
|
|
|
|||
|
|
@ -3,8 +3,9 @@
|
|||
# pageindex/index/page_index_md.py (the single source of truth). This module
|
||||
# re-exports it so legacy imports keep working.
|
||||
#
|
||||
# The canonical md_to_tree takes booleans; legacy callers passed 'yes'/'no'
|
||||
# strings, so the wrapper below coerces them (a bare 'no' is otherwise truthy).
|
||||
# The canonical md_to_tree coerces legacy 'yes'/'no' string flags itself (a
|
||||
# bare 'no' would otherwise be truthy) — this shim used to duplicate that
|
||||
# coercion in its own wrapper; now it just re-exports the canonical function.
|
||||
import warnings
|
||||
|
||||
warnings.warn(
|
||||
|
|
@ -16,23 +17,3 @@ warnings.warn(
|
|||
)
|
||||
|
||||
from .index.page_index_md import * # noqa: F401,F403,E402
|
||||
from .index.page_index_md import md_to_tree as _md_to_tree # noqa: E402
|
||||
|
||||
_BOOL_PARAMS = (
|
||||
"if_thinning", "if_add_node_summary", "if_add_doc_description",
|
||||
"if_add_node_text", "if_add_node_id",
|
||||
)
|
||||
|
||||
|
||||
def _coerce_bool(value):
|
||||
if isinstance(value, str):
|
||||
return value.strip().lower() in ("yes", "true", "1", "y", "on")
|
||||
return bool(value)
|
||||
|
||||
|
||||
async def md_to_tree(*args, **kwargs):
|
||||
"""Legacy wrapper: coerce 'yes'/'no' string flags to bool, then delegate."""
|
||||
for key in _BOOL_PARAMS:
|
||||
if key in kwargs:
|
||||
kwargs[key] = _coerce_bool(kwargs[key])
|
||||
return await _md_to_tree(*args, **kwargs)
|
||||
|
|
|
|||
|
|
@ -2,9 +2,9 @@ import json
|
|||
import PyPDF2
|
||||
|
||||
try:
|
||||
from .index.utils import get_number_of_pages, remove_fields
|
||||
from .index.utils import get_number_of_pages, remove_fields, get_md_page_content
|
||||
except ImportError:
|
||||
from index.utils import get_number_of_pages, remove_fields
|
||||
from index.utils import get_number_of_pages, remove_fields, get_md_page_content
|
||||
|
||||
|
||||
# ── Helpers ──────────────────────────────────────────────────────────────────
|
||||
|
|
@ -54,29 +54,9 @@ def _get_pdf_page_content(doc_info: dict, page_nums: list[int]) -> list[dict]:
|
|||
|
||||
|
||||
def _get_md_page_content(doc_info: dict, page_nums: list[int]) -> list[dict]:
|
||||
"""
|
||||
For Markdown documents, 'pages' are line numbers.
|
||||
Return only the nodes whose line_num is one of ``page_nums`` (exact match),
|
||||
not the whole [min(page_nums), max(page_nums)] range.
|
||||
"""
|
||||
if not page_nums:
|
||||
return []
|
||||
wanted = set(page_nums)
|
||||
results = []
|
||||
seen = set()
|
||||
|
||||
def _traverse(nodes):
|
||||
for node in nodes:
|
||||
ln = node.get('line_num')
|
||||
if ln in wanted and ln not in seen:
|
||||
seen.add(ln)
|
||||
results.append({'page': ln, 'content': node.get('text', '')})
|
||||
if node.get('nodes'):
|
||||
_traverse(node['nodes'])
|
||||
|
||||
_traverse(doc_info.get('structure', []))
|
||||
results.sort(key=lambda x: x['page'])
|
||||
return results
|
||||
"""For Markdown documents, 'pages' are line numbers. Delegates to the
|
||||
canonical implementation so the two never drift again."""
|
||||
return get_md_page_content(doc_info.get('structure', []), page_nums)
|
||||
|
||||
|
||||
# ── Tool functions ────────────────────────────────────────────────────────────
|
||||
|
|
|
|||
|
|
@ -2,19 +2,14 @@ import argparse
|
|||
import os
|
||||
import json
|
||||
from pageindex.index.page_index import *
|
||||
from pageindex.index.page_index_md import md_to_tree
|
||||
# Reuse the canonical yes/no coercion (as _cli_bool) instead of a second copy —
|
||||
# a bare ``--flag`` (no value) resolves to True via argparse's ``const``; an
|
||||
# explicit value keeps the legacy yes/no style working, so ``--flag no`` turns
|
||||
# it off. argparse only ever passes a str here (const/default bypass type=).
|
||||
from pageindex.index.page_index_md import md_to_tree, _coerce_bool as _cli_bool
|
||||
from pageindex.config import IndexConfig
|
||||
|
||||
|
||||
def _cli_bool(value):
|
||||
"""Parse a CLI boolean flag value.
|
||||
|
||||
A bare ``--flag`` (no value) resolves to True via ``const``; an explicit
|
||||
value keeps the legacy yes/no style working, so ``--flag no`` turns it off.
|
||||
"""
|
||||
return str(value).strip().lower() in ("yes", "true", "1", "y", "on")
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
# Set up argument parser
|
||||
parser = argparse.ArgumentParser(description='Process PDF or Markdown document and generate structure')
|
||||
|
|
|
|||
|
|
@ -67,17 +67,37 @@ def test_configloader_no_longer_needs_config_yaml():
|
|||
ConfigLoader().load({"nope": 1})
|
||||
|
||||
|
||||
def test_md_to_tree_shim_coerces_yes_no_strings(monkeypatch):
|
||||
"""Canonical md_to_tree takes booleans; the shim must coerce legacy
|
||||
'yes'/'no' strings so a bare 'no' doesn't read as truthy True."""
|
||||
import pageindex.page_index_md as shim
|
||||
captured = {}
|
||||
def test_md_to_tree_shim_is_the_canonical_function():
|
||||
"""The shim no longer wraps md_to_tree with its own coercion — the
|
||||
canonical implementation coerces internally, so the shim is a pure
|
||||
re-export (single source of truth, can't diverge from the canonical
|
||||
behavior)."""
|
||||
with warnings.catch_warnings():
|
||||
warnings.simplefilter("ignore")
|
||||
import pageindex.page_index_md as shim
|
||||
import pageindex.index.page_index_md as canonical
|
||||
assert shim.md_to_tree is canonical.md_to_tree
|
||||
|
||||
async def fake(*args, **kwargs):
|
||||
captured.update(kwargs)
|
||||
return {"ok": True}
|
||||
|
||||
monkeypatch.setattr(shim, "_md_to_tree", fake)
|
||||
asyncio.run(shim.md_to_tree(md_path="x.md", if_add_node_summary="no", if_add_node_id="yes"))
|
||||
assert captured["if_add_node_summary"] is False
|
||||
assert captured["if_add_node_id"] is True
|
||||
def test_md_to_tree_coerces_legacy_yes_no_strings(tmp_path):
|
||||
"""A bare 'no' must not read as truthy True — exercised end-to-end (no
|
||||
LLM calls needed with summary/description disabled)."""
|
||||
from pageindex.index.page_index_md import md_to_tree
|
||||
|
||||
md_path = tmp_path / "doc.md"
|
||||
md_path.write_text("# Title\nbody\n\n## Sub\nmore body\n")
|
||||
|
||||
result = asyncio.run(md_to_tree(
|
||||
md_path=str(md_path),
|
||||
if_add_node_summary="no",
|
||||
if_add_node_id="yes",
|
||||
if_add_doc_description="no",
|
||||
))
|
||||
assert "doc_description" not in result
|
||||
|
||||
def _has_summary(nodes):
|
||||
return any("summary" in n or (n.get("nodes") and _has_summary(n["nodes"]))
|
||||
for n in nodes)
|
||||
|
||||
assert not _has_summary(result["structure"])
|
||||
assert all("node_id" in n for n in result["structure"])
|
||||
|
|
|
|||
174
tests/test_review_fixes_2.py
Normal file
174
tests/test_review_fixes_2.py
Normal file
|
|
@ -0,0 +1,174 @@
|
|||
"""Regression tests for the second review pass (xhigh code-review of
|
||||
2d46d68..8f536cb): the Markdown text-stripping fix's fallout, plus the other
|
||||
directly-fixable findings from that pass."""
|
||||
import asyncio
|
||||
|
||||
import pytest
|
||||
|
||||
from pageindex.config import IndexConfig
|
||||
|
||||
|
||||
def _md_backend(tmp_path):
|
||||
from pageindex.backend.local import LocalBackend
|
||||
from pageindex.storage.sqlite import SQLiteStorage
|
||||
|
||||
backend = LocalBackend(
|
||||
storage=SQLiteStorage(str(tmp_path / "t.db")),
|
||||
files_dir=str(tmp_path / "f"), model="gpt-4o",
|
||||
index_config=IndexConfig(if_add_node_summary=False, if_add_doc_description=False),
|
||||
)
|
||||
backend.get_or_create_collection("c")
|
||||
return backend
|
||||
|
||||
|
||||
def _write_md(tmp_path, name="doc.md"):
|
||||
path = tmp_path / name
|
||||
path.write_text("# Title\nfirst section body\n\n## Sub\nsecond section body\n")
|
||||
return str(path)
|
||||
|
||||
|
||||
# ── #1: get_document(include_text=True) must fill text for Markdown nodes ────
|
||||
def test_get_document_include_text_fills_markdown_nodes(tmp_path):
|
||||
backend = _md_backend(tmp_path)
|
||||
doc_id = backend.add_document("c", _write_md(tmp_path))
|
||||
|
||||
def _texts(nodes):
|
||||
for n in nodes:
|
||||
yield n.get("text")
|
||||
if n.get("nodes"):
|
||||
yield from _texts(n["nodes"])
|
||||
|
||||
without = backend.get_document("c", doc_id, include_text=False)
|
||||
assert not any(_texts(without["structure"]))
|
||||
|
||||
with_text = backend.get_document("c", doc_id, include_text=True)
|
||||
texts = list(_texts(with_text["structure"]))
|
||||
assert texts, "expected at least one node"
|
||||
assert any(t for t in texts), "Markdown nodes must get real text, not all empty"
|
||||
assert any("first section body" in t or "second section body" in t for t in texts if t)
|
||||
|
||||
|
||||
# ── #2: get_page_content's Markdown fallback re-derives from the source file ──
|
||||
def test_get_page_content_markdown_fallback_reads_from_file(tmp_path):
|
||||
backend = _md_backend(tmp_path)
|
||||
md_path = _write_md(tmp_path)
|
||||
doc_id = backend.add_document("c", md_path)
|
||||
|
||||
# Simulate a StorageEngine that doesn't cache pages (protocol explicitly
|
||||
# allows get_pages() to return None) by clearing the cached pages column.
|
||||
conn = backend._storage._get_conn()
|
||||
conn.execute("UPDATE documents SET pages = NULL WHERE doc_id = ?", (doc_id,))
|
||||
|
||||
result = backend.get_page_content("c", doc_id, "1")
|
||||
assert result and result[0]["content"], "fallback must return real text, not empty"
|
||||
assert "first section body" in result[0]["content"]
|
||||
|
||||
|
||||
# ── #3: keyless provider allowlist covers other local LiteLLM providers ──────
|
||||
@pytest.mark.parametrize("model", [
|
||||
"ollama/llama3", "lm_studio/x", "xinference/llama2", "llamafile/x",
|
||||
"triton/x", "oobabooga/x", "openai_like/x", "docker_model_runner/x",
|
||||
])
|
||||
def test_validate_llm_provider_accepts_more_keyless_providers(model):
|
||||
from pageindex.client import LocalClient
|
||||
LocalClient._validate_llm_provider(model) # must not raise
|
||||
|
||||
|
||||
# ── #4: agent-tool closures consistently raise/error on a missing doc ────────
|
||||
def test_agent_tools_consistently_report_missing_doc(tmp_path):
|
||||
import json
|
||||
import asyncio as _asyncio
|
||||
from agents.tool_context import ToolContext
|
||||
|
||||
backend = _md_backend(tmp_path)
|
||||
# Open-mode tools (doc_ids=None) so we probe not-found handling directly,
|
||||
# not the separate out-of-scope rejection path.
|
||||
tools = backend.get_agent_tools("c", doc_ids=None)
|
||||
by_name = {t.name: t for t in tools.function_tools}
|
||||
|
||||
for name in ("get_document", "get_document_structure", "get_page_content"):
|
||||
tool = by_name[name]
|
||||
kwargs = {"doc_id": "ghost"}
|
||||
if name == "get_page_content":
|
||||
kwargs["pages"] = "1"
|
||||
raw_args = json.dumps(kwargs)
|
||||
ctx = ToolContext(context=None, tool_name=name, tool_call_id="1", tool_arguments=raw_args)
|
||||
out = _asyncio.run(tool.on_invoke_tool(ctx, raw_args))
|
||||
parsed = json.loads(out)
|
||||
assert "error" in parsed and "ghost" in parsed["error"], f"{name} did not report not-found consistently: {parsed}"
|
||||
|
||||
|
||||
# ── #6: cloud delete_collection preserves the "folders unavailable" sentinel ──
|
||||
def test_cloud_delete_collection_preserves_unavailable_sentinel(monkeypatch):
|
||||
from pageindex.backend.cloud import CloudBackend
|
||||
|
||||
backend = CloudBackend(api_key="pi-test")
|
||||
backend._folder_id_cache["papers"] = None # folders-unavailable sentinel
|
||||
called = []
|
||||
monkeypatch.setattr(backend, "_request", lambda *a, **k: called.append(a) or {})
|
||||
backend.delete_collection("papers")
|
||||
assert not called, "no DELETE should fire when folder_id is the unavailable sentinel"
|
||||
assert "papers" in backend._folder_id_cache and backend._folder_id_cache["papers"] is None
|
||||
|
||||
|
||||
def test_cloud_delete_collection_still_clears_real_folder_id(monkeypatch):
|
||||
from pageindex.backend.cloud import CloudBackend
|
||||
|
||||
backend = CloudBackend(api_key="pi-test")
|
||||
backend._folder_id_cache["papers"] = "folder-123"
|
||||
monkeypatch.setattr(backend, "_request", lambda *a, **k: {})
|
||||
backend.delete_collection("papers")
|
||||
assert "papers" not in backend._folder_id_cache
|
||||
|
||||
|
||||
# ── #7: remove_structure_text is skipped when text was never added ───────────
|
||||
def test_build_index_skips_text_strip_when_no_text_was_added(monkeypatch):
|
||||
from pageindex.index import pipeline
|
||||
from pageindex.parser.protocol import ContentNode, ParsedDocument
|
||||
|
||||
calls = []
|
||||
# build_index() imports remove_structure_text locally (`from .utils import
|
||||
# ...` inside the function body), so patch it on the utils module itself.
|
||||
import pageindex.index.utils as utils_mod
|
||||
monkeypatch.setattr(utils_mod, "remove_structure_text", lambda s: calls.append(s) or s)
|
||||
|
||||
nodes = [ContentNode(content="page one text", tokens=5, index=1)]
|
||||
parsed = ParsedDocument(doc_name="d", nodes=nodes)
|
||||
opt = IndexConfig(if_add_node_summary=False, if_add_doc_description=False,
|
||||
if_add_node_text=False)
|
||||
pipeline.build_index(parsed, opt=opt)
|
||||
assert calls == [], "remove_structure_text must not run when no text was ever added"
|
||||
|
||||
|
||||
def test_build_index_still_strips_text_when_summary_added_it(monkeypatch):
|
||||
from pageindex.index import pipeline
|
||||
from pageindex.parser.protocol import ContentNode, ParsedDocument
|
||||
|
||||
calls = []
|
||||
import pageindex.index.utils as utils_mod
|
||||
monkeypatch.setattr(utils_mod, "remove_structure_text", lambda s: calls.append(s) or s)
|
||||
|
||||
nodes = [ContentNode(content="page one text", tokens=5, index=1)]
|
||||
parsed = ParsedDocument(doc_name="d", nodes=nodes)
|
||||
opt = IndexConfig(if_add_node_summary=True, if_add_doc_description=False,
|
||||
if_add_node_text=False)
|
||||
pipeline.build_index(parsed, opt=opt)
|
||||
assert len(calls) == 1, "text WAS added for summary generation, so it must still be stripped"
|
||||
|
||||
|
||||
# ── #9/#10: run_pageindex._cli_bool and the shim's md_to_tree are the same
|
||||
# object as the canonical implementation (no drift possible) ──────
|
||||
def test_cli_bool_is_the_canonical_coerce_bool():
|
||||
import run_pageindex
|
||||
from pageindex.index.page_index_md import _coerce_bool
|
||||
assert run_pageindex._cli_bool is _coerce_bool
|
||||
|
||||
|
||||
# ── #11: retrieve._get_md_page_content delegates to the canonical function ───
|
||||
def test_retrieve_md_page_content_delegates_to_canonical():
|
||||
from pageindex import retrieve
|
||||
structure = [{"line_num": 5, "text": "five", "nodes": [
|
||||
{"line_num": 40, "text": "forty", "nodes": []},
|
||||
]}]
|
||||
out = retrieve._get_md_page_content({"structure": structure}, [5])
|
||||
assert [r["page"] for r in out] == [5]
|
||||
Loading…
Add table
Add a link
Reference in a new issue