pageindex 0.3.0.dev2__tar.gz → 0.3.0.dev3__tar.gz

This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
Files changed (36) hide show
  1. {pageindex-0.3.0.dev2 → pageindex-0.3.0.dev3}/PKG-INFO +7 -4
  2. {pageindex-0.3.0.dev2 → pageindex-0.3.0.dev3}/README.md +6 -3
  3. {pageindex-0.3.0.dev2 → pageindex-0.3.0.dev3}/pageindex/__init__.py +7 -0
  4. {pageindex-0.3.0.dev2 → pageindex-0.3.0.dev3}/pageindex/agent.py +14 -5
  5. {pageindex-0.3.0.dev2 → pageindex-0.3.0.dev3}/pageindex/backend/cloud.py +27 -4
  6. {pageindex-0.3.0.dev2 → pageindex-0.3.0.dev3}/pageindex/backend/local.py +77 -32
  7. {pageindex-0.3.0.dev2 → pageindex-0.3.0.dev3}/pageindex/client.py +20 -3
  8. {pageindex-0.3.0.dev2 → pageindex-0.3.0.dev3}/pageindex/cloud_api.py +22 -3
  9. {pageindex-0.3.0.dev2 → pageindex-0.3.0.dev3}/pageindex/config.py +114 -10
  10. {pageindex-0.3.0.dev2 → pageindex-0.3.0.dev3}/pageindex/errors.py +15 -2
  11. {pageindex-0.3.0.dev2 → pageindex-0.3.0.dev3}/pageindex/index/page_index.py +45 -12
  12. {pageindex-0.3.0.dev2 → pageindex-0.3.0.dev3}/pageindex/index/page_index_md.py +22 -2
  13. {pageindex-0.3.0.dev2 → pageindex-0.3.0.dev3}/pageindex/index/pipeline.py +38 -15
  14. {pageindex-0.3.0.dev2 → pageindex-0.3.0.dev3}/pageindex/index/utils.py +195 -47
  15. pageindex-0.3.0.dev3/pageindex/page_index.py +40 -0
  16. pageindex-0.3.0.dev3/pageindex/page_index_md.py +19 -0
  17. {pageindex-0.3.0.dev2 → pageindex-0.3.0.dev3}/pageindex/parser/markdown.py +22 -6
  18. {pageindex-0.3.0.dev2 → pageindex-0.3.0.dev3}/pageindex/parser/pdf.py +7 -1
  19. {pageindex-0.3.0.dev2 → pageindex-0.3.0.dev3}/pageindex/retrieve.py +18 -48
  20. {pageindex-0.3.0.dev2 → pageindex-0.3.0.dev3}/pageindex/storage/sqlite.py +53 -4
  21. {pageindex-0.3.0.dev2 → pageindex-0.3.0.dev3}/pageindex/types.py +13 -3
  22. {pageindex-0.3.0.dev2 → pageindex-0.3.0.dev3}/pyproject.toml +4 -1
  23. pageindex-0.3.0.dev2/pageindex/page_index.py +0 -15
  24. pageindex-0.3.0.dev2/pageindex/page_index_md.py +0 -38
  25. {pageindex-0.3.0.dev2 → pageindex-0.3.0.dev3}/LICENSE +0 -0
  26. {pageindex-0.3.0.dev2 → pageindex-0.3.0.dev3}/pageindex/backend/__init__.py +0 -0
  27. {pageindex-0.3.0.dev2 → pageindex-0.3.0.dev3}/pageindex/backend/protocol.py +0 -0
  28. {pageindex-0.3.0.dev2 → pageindex-0.3.0.dev3}/pageindex/collection.py +0 -0
  29. {pageindex-0.3.0.dev2 → pageindex-0.3.0.dev3}/pageindex/events.py +0 -0
  30. {pageindex-0.3.0.dev2 → pageindex-0.3.0.dev3}/pageindex/index/__init__.py +0 -0
  31. {pageindex-0.3.0.dev2 → pageindex-0.3.0.dev3}/pageindex/parser/__init__.py +0 -0
  32. {pageindex-0.3.0.dev2 → pageindex-0.3.0.dev3}/pageindex/parser/protocol.py +0 -0
  33. {pageindex-0.3.0.dev2 → pageindex-0.3.0.dev3}/pageindex/storage/__init__.py +0 -0
  34. {pageindex-0.3.0.dev2 → pageindex-0.3.0.dev3}/pageindex/storage/protocol.py +0 -0
  35. {pageindex-0.3.0.dev2 → pageindex-0.3.0.dev3}/pageindex/tokens.py +0 -0
  36. {pageindex-0.3.0.dev2 → pageindex-0.3.0.dev3}/pageindex/utils.py +0 -0
@@ -1,6 +1,6 @@
1
1
  Metadata-Version: 2.4
2
2
  Name: pageindex
3
- Version: 0.3.0.dev2
3
+ Version: 0.3.0.dev3
4
4
  Summary: Python SDK for PageIndex
5
5
  License: MIT
6
6
  License-File: LICENSE
@@ -275,10 +275,13 @@ You can customize the processing with additional optional arguments:
275
275
  --toc-check-pages Pages to check for table of contents (default: 20)
276
276
  --max-pages-per-node Max pages per node (default: 10)
277
277
  --max-tokens-per-node Max tokens per node (default: 20000)
278
- --if-add-node-id Add node ID (yes/no, default: yes)
279
- --if-add-node-summary Add node summary (yes/no, default: yes)
280
- --if-add-doc-description Add doc description (yes/no, default: yes)
278
+ --if-add-node-id Add node IDs (on by default; disable with: --if-add-node-id no)
279
+ --if-add-node-summary Add node summaries (on by default; disable with: --if-add-node-summary no)
280
+ --if-add-doc-description Add a document description (on by default; disable with: --if-add-doc-description no)
281
+ --if-add-node-text Add raw text to nodes (off by default; enable with: --if-add-node-text)
281
282
  ```
283
+ These flags take no value by default (a bare `--if-add-node-id` turns it on); the
284
+ legacy `--if-add-node-id no` form still works for turning an option off.
282
285
  </details>
283
286
 
284
287
  <details>
@@ -238,10 +238,13 @@ You can customize the processing with additional optional arguments:
238
238
  --toc-check-pages Pages to check for table of contents (default: 20)
239
239
  --max-pages-per-node Max pages per node (default: 10)
240
240
  --max-tokens-per-node Max tokens per node (default: 20000)
241
- --if-add-node-id Add node ID (yes/no, default: yes)
242
- --if-add-node-summary Add node summary (yes/no, default: yes)
243
- --if-add-doc-description Add doc description (yes/no, default: yes)
241
+ --if-add-node-id Add node IDs (on by default; disable with: --if-add-node-id no)
242
+ --if-add-node-summary Add node summaries (on by default; disable with: --if-add-node-summary no)
243
+ --if-add-doc-description Add a document description (on by default; disable with: --if-add-doc-description no)
244
+ --if-add-node-text Add raw text to nodes (off by default; enable with: --if-add-node-text)
244
245
  ```
246
+ These flags take no value by default (a bare `--if-add-node-id` turns it on); the
247
+ legacy `--if-add-node-id no` form still works for turning an option off.
245
248
  </details>
246
249
 
247
250
  <details>
@@ -61,4 +61,11 @@ __all__ = [
61
61
  "IndexingError",
62
62
  "CloudAPIError",
63
63
  "FileTypeError",
64
+ # Legacy top-level exports (pre-SDK API), kept so `from pageindex import *`
65
+ # still binds them.
66
+ "page_index",
67
+ "md_to_tree",
68
+ "get_document",
69
+ "get_document_structure",
70
+ "get_page_content",
64
71
  ]
@@ -20,7 +20,7 @@ OPEN_SYSTEM_PROMPT = """
20
20
  You are PageIndex, a document QA assistant.
21
21
  TOOL USE:
22
22
  - Call list_documents() to see available documents; use doc_name and doc_description to pick which doc(s) are relevant.
23
- - Call get_document(doc_id) to confirm status and page/line count.
23
+ - Call get_document(doc_id) to confirm the document's name and type.
24
24
  - Call get_document_structure(doc_id) to identify relevant page ranges.
25
25
  - Call get_page_content(doc_id, pages="5-7") with tight ranges; never fetch the whole document.
26
26
  - Before each tool call, output one short sentence explaining the reason.
@@ -33,7 +33,7 @@ Answer based only on tool output. Be concise.
33
33
  SCOPED_SYSTEM_PROMPT = """
34
34
  You are PageIndex, a document QA assistant.
35
35
  TOOL USE:
36
- - Call get_document(doc_id) to confirm status and page/line count.
36
+ - Call get_document(doc_id) to confirm the document's name and type.
37
37
  - Call get_document_structure(doc_id) to identify relevant page ranges.
38
38
  - Call get_page_content(doc_id, pages="5-7") with tight ranges; never fetch the whole document.
39
39
  - Before each tool call, output one short sentence explaining the reason.
@@ -46,20 +46,29 @@ Answer based only on tool output. Be concise.
46
46
  """
47
47
 
48
48
 
49
+ def _defang_delimiters(text: str) -> str:
50
+ """Strip '<'/'>' so untrusted text can never form a literal <docs>/</docs>
51
+ (or any other tag-shaped string) that would prematurely close the
52
+ wrap_with_doc_context() delimiter and escape the untrusted-data boundary."""
53
+ return text.replace("<", "").replace(">", "")
54
+
55
+
49
56
  def wrap_with_doc_context(docs: list[dict], question: str) -> str:
50
57
  """Prepend a doc-context block to the user question for scoped queries.
51
58
 
52
59
  Document fields (especially doc_description, which is LLM-generated at
53
60
  index time) are untrusted text that may contain adversarial instructions.
54
61
  We wrap them in a <docs>...</docs> delimiter and tell the agent in the
55
- system prompt to treat the block as data only.
62
+ system prompt to treat the block as data only. '<'/'>' are stripped from
63
+ the untrusted fields first so embedded content can never form a literal
64
+ </docs> (or any other tag) that closes the delimiter early.
56
65
  """
57
66
  lines = []
58
67
  for d in docs:
59
- line = f"- {d['doc_id']}: {d.get('doc_name', '')}"
68
+ line = f"- {d['doc_id']}: {_defang_delimiters(d.get('doc_name', ''))}"
60
69
  desc = d.get("doc_description") or ""
61
70
  if desc:
62
- line += f" — {desc}"
71
+ line += f" — {_defang_delimiters(desc)}"
63
72
  lines.append(line)
64
73
  label = "document" if len(docs) == 1 else "documents"
65
74
  return (
@@ -13,13 +13,12 @@ import urllib.parse
13
13
  import requests
14
14
  from typing import AsyncIterator
15
15
 
16
+ from ..cloud_api import API_BASE # single source of truth for the cloud base URL
16
17
  from ..errors import CloudAPIError, DocumentNotFoundError, PageIndexError
17
18
  from ..events import QueryEvent
18
19
 
19
20
  logger = logging.getLogger(__name__)
20
21
 
21
- API_BASE = "https://api.pageindex.ai"
22
-
23
22
  _INTERNAL_TOOLS = frozenset({"ToolSearch", "Read", "Grep", "Glob", "Bash", "Edit", "Write"})
24
23
 
25
24
 
@@ -89,7 +88,9 @@ class CloudBackend:
89
88
 
90
89
  @staticmethod
91
90
  def _validate_collection_name(name: str) -> None:
92
- if not re.match(r'^[a-zA-Z0-9_-]{1,128}$', name):
91
+ # .fullmatch() (not .match()): a $-anchored .match() would accept a
92
+ # trailing newline ("papers\n") because $ matches just before a final \n.
93
+ if not re.fullmatch(r'[a-zA-Z0-9_-]{1,128}', name):
93
94
  raise PageIndexError(
94
95
  f"Invalid collection name: {name!r}. "
95
96
  "Must be 1-128 chars of [a-zA-Z0-9_-]."
@@ -162,6 +163,12 @@ class CloudBackend:
162
163
  folder_id = self._get_folder_id(name)
163
164
  if folder_id:
164
165
  self._request("DELETE", f"/folder/{self._enc(folder_id)}/")
166
+ # Drop the cached id so a later same-name op re-resolves instead of
167
+ # reusing the now-deleted folder_id. Only when it was a REAL id —
168
+ # if folder_id was falsy, the cache holds the "folders unavailable
169
+ # on this plan" None sentinel, which must survive so we don't
170
+ # re-issue a doomed GET /folders/ on the next call.
171
+ self._folder_id_cache.pop(name, None)
165
172
 
166
173
  # ── Document management ───────────────────────────────────────────────
167
174
 
@@ -239,7 +246,12 @@ class CloudBackend:
239
246
  if isinstance(all_pages, list):
240
247
  return [
241
248
  {"page": p.get("page", p.get("page_index")),
242
- "content": p.get("content", p.get("markdown", ""))}
249
+ "content": p.get("content", p.get("markdown", "")),
250
+ # Cloud OCR pages carry an `images` list (empty on text-only
251
+ # pages). Preserve it — omitting when empty, mirroring the local
252
+ # backend — so cloud callers get the same PageContent shape and
253
+ # the SDK-prompted UI can render figures.
254
+ **({"images": p["images"]} if p.get("images") else {})}
243
255
  for p in all_pages
244
256
  if p.get("page", p.get("page_index")) in page_nums
245
257
  ]
@@ -307,6 +319,8 @@ class CloudBackend:
307
319
  "doc_ids cannot be empty; pass None to query the whole collection"
308
320
  )
309
321
  doc_id = doc_ids if doc_ids else self._get_all_doc_ids(collection)
322
+ if not doc_id:
323
+ raise ValueError("collection has no documents to query")
310
324
  # A non-streaming completion returns nothing until generation
311
325
  # finishes, so it needs far more than the default 30s. retries=1:
312
326
  # retrying this non-idempotent call would redo the full server-side
@@ -340,6 +354,8 @@ class CloudBackend:
340
354
  "doc_ids cannot be empty; pass None to query the whole collection"
341
355
  )
342
356
  doc_id = doc_ids if doc_ids else self._get_all_doc_ids(collection)
357
+ if not doc_id:
358
+ raise ValueError("collection has no documents to query")
343
359
  headers = self._headers
344
360
  # Queue carries QueryEvent, an Exception to re-raise, or None (end).
345
361
  queue: asyncio.Queue[QueryEvent | Exception | None] = asyncio.Queue()
@@ -380,6 +396,13 @@ class CloudBackend:
380
396
  timeout=120,
381
397
  )
382
398
  resp_holder["resp"] = resp
399
+ # The consumer may have abandoned the stream while we were still
400
+ # blocked in requests.post() (its connect phase, before resp
401
+ # existed to close). Now that resp exists, bail immediately
402
+ # rather than reading/draining a stream nobody is listening to;
403
+ # the finally block closes resp and pushes the sentinel.
404
+ if stop.is_set():
405
+ return
383
406
  if resp.status_code != 200:
384
407
  body = resp.text[:500] if resp.text else ""
385
408
  raise CloudAPIError(
@@ -12,12 +12,14 @@ from ..parser.pdf import PdfParser
12
12
  from ..parser.markdown import MarkdownParser
13
13
  from ..storage.protocol import StorageEngine
14
14
  from ..index.pipeline import build_index
15
- from ..index.utils import parse_pages, get_pdf_page_content, get_md_page_content, remove_fields
15
+ from ..index.utils import parse_pages, get_pdf_page_content, remove_fields
16
16
  from ..backend.protocol import AgentTools
17
17
  from ..errors import (FileTypeError, DocumentNotFoundError, CollectionNotFoundError,
18
18
  IndexingError, PageIndexError)
19
19
 
20
- _COLLECTION_NAME_RE = re.compile(r'^[a-zA-Z0-9_-]{1,128}$')
20
+ # Matched with .fullmatch() (not .match()): a $-anchored .match() would accept a
21
+ # trailing newline ("papers\n") because $ matches just before a final \n.
22
+ _COLLECTION_NAME_RE = re.compile(r'[a-zA-Z0-9_-]{1,128}')
21
23
 
22
24
 
23
25
  class LocalBackend:
@@ -45,7 +47,7 @@ class LocalBackend:
45
47
 
46
48
  # Collection management
47
49
  def _validate_collection_name(self, name: str) -> None:
48
- if not _COLLECTION_NAME_RE.match(name):
50
+ if not _COLLECTION_NAME_RE.fullmatch(name):
49
51
  raise PageIndexError(f"Invalid collection name: {name!r}. Must be 1-128 chars of [a-zA-Z0-9_-].")
50
52
 
51
53
  def create_collection(self, name: str) -> None:
@@ -116,26 +118,23 @@ class LocalBackend:
116
118
  parsed = parser.parse(file_path, model=self._model, images_dir=images_dir)
117
119
  result = build_index(parsed, model=self._model, opt=self._index_config)
118
120
 
119
- # Cache page text for fast retrieval (avoids re-reading files)
121
+ # Cache page text for fast retrieval (avoids re-reading files) and to
122
+ # reconstruct node text on demand (get_document(include_text=True),
123
+ # get_page_content fallback) independent of whether IndexConfig kept
124
+ # text in the stored structure. build_index() already applies
125
+ # if_add_node_text to result["structure"] for every strategy, so no
126
+ # extra stripping is needed here.
120
127
  pages = [{"page": n.index, "content": n.content,
121
128
  **({"images": n.images} if n.images else {})}
122
129
  for n in parsed.nodes if n.content]
123
130
 
124
- # Strip text from structure to save storage space (PDF only;
125
- # markdown needs text in structure for fallback retrieval)
126
- doc_type = ext.lstrip(".")
127
- if doc_type == "pdf":
128
- clean_structure = remove_fields(result["structure"], fields=["text"])
129
- else:
130
- clean_structure = result["structure"]
131
-
132
131
  self._storage.save_document(collection, doc_id, {
133
132
  "doc_name": parsed.doc_name,
134
133
  "doc_description": result.get("doc_description", ""),
135
134
  "file_path": str(managed_path),
136
135
  "file_hash": file_hash,
137
- "doc_type": doc_type,
138
- "structure": clean_structure,
136
+ "doc_type": ext.lstrip("."),
137
+ "structure": result["structure"],
139
138
  "pages": pages,
140
139
  })
141
140
  except sqlite3.IntegrityError:
@@ -158,6 +157,19 @@ class LocalBackend:
158
157
 
159
158
  return doc_id
160
159
 
160
+ def _require_document(self, collection: str, doc_id: str) -> dict:
161
+ """Return the document's storage row, or raise DocumentNotFoundError.
162
+
163
+ Single source of truth for "does this doc exist" — every public method
164
+ and agent tool below goes through this, so a missing doc always
165
+ surfaces the same way instead of each caller re-implementing its own
166
+ (and potentially inconsistent) existence check.
167
+ """
168
+ doc = self._storage.get_document(collection, doc_id)
169
+ if not doc:
170
+ raise DocumentNotFoundError(f"Document {doc_id} not found")
171
+ return doc
172
+
161
173
  def get_document(self, collection: str, doc_id: str, include_text: bool = False) -> dict:
162
174
  """Get document metadata with structure.
163
175
 
@@ -166,9 +178,7 @@ class LocalBackend:
166
178
  from cached page content. WARNING: may be very large — do NOT
167
179
  use in agent/LLM contexts as it can exhaust the context window.
168
180
  """
169
- doc = self._storage.get_document(collection, doc_id)
170
- if not doc:
171
- raise DocumentNotFoundError(f"Document {doc_id} not found")
181
+ doc = self._require_document(collection, doc_id)
172
182
  doc["structure"] = self._storage.get_document_structure(collection, doc_id)
173
183
  if include_text:
174
184
  pages = self._storage.get_pages(collection, doc_id) or []
@@ -178,7 +188,13 @@ class LocalBackend:
178
188
 
179
189
  @staticmethod
180
190
  def _fill_node_text(nodes: list, page_map: dict) -> None:
181
- """Recursively fill 'text' on structure nodes from cached page content."""
191
+ """Recursively fill 'text' on structure nodes from cached page content.
192
+
193
+ Two node conventions, one per indexing strategy: content_based (PDF)
194
+ nodes span a start_index..end_index page range; level_based (Markdown)
195
+ nodes map 1:1 to a single page keyed by line_num. Handling only the
196
+ first would silently leave Markdown nodes with no text.
197
+ """
182
198
  for node in nodes:
183
199
  start = node.get("start_index")
184
200
  end = node.get("end_index")
@@ -186,16 +202,17 @@ class LocalBackend:
186
202
  node["text"] = "\n".join(
187
203
  page_map.get(p, "") for p in range(start, end + 1)
188
204
  )
205
+ elif "line_num" in node:
206
+ node["text"] = page_map.get(node["line_num"], "")
189
207
  if "nodes" in node:
190
208
  LocalBackend._fill_node_text(node["nodes"], page_map)
191
209
 
192
210
  def get_document_structure(self, collection: str, doc_id: str) -> list:
211
+ self._require_document(collection, doc_id)
193
212
  return self._storage.get_document_structure(collection, doc_id)
194
213
 
195
214
  def get_page_content(self, collection: str, doc_id: str, pages: str) -> list:
196
- doc = self._storage.get_document(collection, doc_id)
197
- if not doc:
198
- raise DocumentNotFoundError(f"Document {doc_id} not found")
215
+ doc = self._require_document(collection, doc_id)
199
216
  page_nums = parse_pages(pages)
200
217
 
201
218
  # Try cached pages first (fast, no file I/O)
@@ -203,22 +220,24 @@ class LocalBackend:
203
220
  if cached_pages:
204
221
  return [p for p in cached_pages if p["page"] in page_nums]
205
222
 
206
- # Fallback to reading from file
223
+ # Fallback: re-derive from the source file, same as the PDF path below
224
+ # — never from the stored structure, whose 'text' field may have been
225
+ # stripped (if_add_node_text=False, the default). Reachable only for a
226
+ # custom StorageEngine that doesn't cache pages (the built-in
227
+ # SQLiteStorage always does).
207
228
  if doc["doc_type"] == "pdf":
208
229
  return get_pdf_page_content(doc["file_path"], page_nums)
209
230
  else:
210
- structure = self._storage.get_document_structure(collection, doc_id)
211
- return get_md_page_content(structure, page_nums)
231
+ parser = self._resolve_parser(doc["file_path"])
232
+ parsed = parser.parse(doc["file_path"], model=self._model)
233
+ page_map = {n.index: n.content for n in parsed.nodes}
234
+ return [{"page": p, "content": page_map[p]} for p in page_nums if p in page_map]
212
235
 
213
236
  def list_documents(self, collection: str) -> list[dict]:
214
237
  return self._storage.list_documents(collection)
215
238
 
216
239
  def delete_document(self, collection: str, doc_id: str) -> None:
217
- doc = self._storage.get_document(collection, doc_id)
218
- if not doc:
219
- # Parity with the cloud backend, which surfaces HTTP 404 as
220
- # DocumentNotFoundError — a typo'd doc_id should not pass silently.
221
- raise DocumentNotFoundError(f"Document {doc_id} not found")
240
+ doc = self._require_document(collection, doc_id)
222
241
  if doc.get("file_path"):
223
242
  Path(doc["file_path"]).unlink(missing_ok=True)
224
243
  # Clean up images directory: files/{collection}/{doc_id}/
@@ -233,13 +252,17 @@ class LocalBackend:
233
252
  - doc_ids=None (open mode): includes ``list_documents``; agent picks docs itself.
234
253
  - doc_ids=[...] (scoped mode): no ``list_documents``; the other tools
235
254
  hard-enforce the whitelist and reject out-of-scope doc_ids.
255
+
256
+ Note ``is not None``: an empty list is a scope of *nothing* (reject every
257
+ doc), NOT open mode. Using truthiness would let ``doc_ids=[]`` collapse to
258
+ ``None`` and silently grant access to the whole collection.
236
259
  """
237
260
  from agents import function_tool
238
261
  import json
239
262
  storage = self._storage
240
263
  col_name = collection
241
264
  backend = self
242
- scope = set(doc_ids) if doc_ids else None
265
+ scope = set(doc_ids) if doc_ids is not None else None
243
266
 
244
267
  def _reject(doc_id: str) -> str | None:
245
268
  if scope is not None and doc_id not in scope:
@@ -255,7 +278,14 @@ class LocalBackend:
255
278
  rejection = _reject(doc_id)
256
279
  if rejection:
257
280
  return rejection
258
- return json.dumps(storage.get_document(col_name, doc_id))
281
+ try:
282
+ # _require_document (not backend.get_document) deliberately:
283
+ # the metadata-only row, no 'structure' — keeps this tool's
284
+ # output small for the agent's context window.
285
+ doc = backend._require_document(col_name, doc_id)
286
+ except DocumentNotFoundError:
287
+ return json.dumps({"error": f"doc_id '{doc_id}' not found."})
288
+ return json.dumps(doc)
259
289
 
260
290
  @function_tool
261
291
  def get_document_structure(doc_id: str) -> str:
@@ -263,6 +293,10 @@ class LocalBackend:
263
293
  rejection = _reject(doc_id)
264
294
  if rejection:
265
295
  return rejection
296
+ try:
297
+ backend._require_document(col_name, doc_id)
298
+ except DocumentNotFoundError:
299
+ return json.dumps({"error": f"doc_id '{doc_id}' not found."})
266
300
  structure = storage.get_document_structure(col_name, doc_id)
267
301
  return json.dumps(remove_fields(structure, fields=["text"]), ensure_ascii=False)
268
302
 
@@ -272,7 +306,18 @@ class LocalBackend:
272
306
  rejection = _reject(doc_id)
273
307
  if rejection:
274
308
  return rejection
275
- result = backend.get_page_content(col_name, doc_id, pages)
309
+ try:
310
+ result = backend.get_page_content(col_name, doc_id, pages)
311
+ except DocumentNotFoundError:
312
+ return json.dumps({"error": f"doc_id '{doc_id}' not found."})
313
+ except (ValueError, AttributeError) as e:
314
+ # A malformed page spec ("all", "5-") is a recoverable bad tool
315
+ # argument: hand the model an actionable error it can correct
316
+ # (mirroring the legacy retrieval tool) rather than letting the
317
+ # ValueError surface as the agent SDK's generic tool-failure text.
318
+ return json.dumps({
319
+ "error": f"Invalid pages format: {pages!r}. Use '5-7', '3,8', or '12'. Error: {e}"
320
+ })
276
321
  return json.dumps(result, ensure_ascii=False)
277
322
 
278
323
  tools = [get_document, get_document_structure, get_page_content]
@@ -5,6 +5,7 @@ from typing import Any, Iterator
5
5
 
6
6
  from typing_extensions import deprecated
7
7
 
8
+ from .cloud_api import API_BASE
8
9
  from .collection import Collection
9
10
  from .config import IndexConfig
10
11
  from .errors import PageIndexAPIError
@@ -50,7 +51,7 @@ class PageIndexClient:
50
51
  # Or use LocalClient / CloudClient for explicit mode selection
51
52
  """
52
53
 
53
- BASE_URL = "https://api.pageindex.ai"
54
+ BASE_URL = API_BASE # single source of truth lives in cloud_api
54
55
 
55
56
  def __init__(self, api_key: str | None = None, model: str = None,
56
57
  retrieve_model: str = None, storage_path: str = None,
@@ -113,14 +114,30 @@ class PageIndexClient:
113
114
 
114
115
  @staticmethod
115
116
  def _validate_llm_provider(model: str) -> None:
116
- """Validate model and check API key via litellm. Warns if key seems missing."""
117
+ """Validate the model string and require an API key for providers that
118
+ need one. Local / keyless providers (ollama, lm_studio, …) are skipped so
119
+ a keyless LiteLLM model isn't rejected at construction time."""
117
120
  try:
118
121
  import litellm
119
- litellm.model_cost_map_url = ""
120
122
  _, provider, _, _ = litellm.get_llm_provider(model=model)
121
123
  except Exception:
122
124
  return
123
125
 
126
+ # LiteLLM providers that run locally / self-hosted and need no API key
127
+ # by default (litellm itself falls back to a placeholder key for these
128
+ # rather than erroring — see e.g. hosted_vllm's transformation.py).
129
+ # This list is necessarily a manual allowlist (litellm.validate_environment
130
+ # isn't reliable enough to derive it from); extend it as litellm adds
131
+ # more local-inference providers.
132
+ keyless = {
133
+ "ollama", "ollama_chat", "lm_studio", "hosted_vllm", "vllm",
134
+ "xinference", "llamafile", "triton", "oobabooga",
135
+ "openai_like", "custom_openai", "custom", "docker_model_runner",
136
+ "petals",
137
+ }
138
+ if provider in keyless:
139
+ return
140
+
124
141
  key = litellm.get_api_key(llm_provider=provider, dynamic_api_key=None)
125
142
  if not key:
126
143
  import os
@@ -8,11 +8,16 @@ import requests
8
8
 
9
9
  from .errors import PageIndexAPIError
10
10
 
11
+ # Single source of truth for the cloud API base URL — imported by the modern
12
+ # CloudBackend (as API_BASE) and PageIndexClient so a staging/migration change
13
+ # only has to happen here.
14
+ API_BASE = "https://api.pageindex.ai"
15
+
11
16
 
12
17
  class LegacyCloudAPI:
13
18
  """Compatibility layer for the pageindex 0.2.x cloud SDK API."""
14
19
 
15
- BASE_URL = "https://api.pageindex.ai"
20
+ BASE_URL = API_BASE
16
21
 
17
22
  def __init__(self, api_key: str, base_url: str | None = None):
18
23
  self.api_key = api_key
@@ -85,7 +90,11 @@ class LegacyCloudAPI:
85
90
  def get_tree(self, doc_id: str, node_summary: bool = False) -> dict[str, Any]:
86
91
  response = self._request(
87
92
  "GET",
88
- f"/doc/{self._enc(doc_id)}/?type=tree&summary={node_summary}",
93
+ # Lowercase the bool: a Python f-string renders True/False with a
94
+ # capital letter, but the API expects summary=true/false (the modern
95
+ # CloudBackend sends lowercase). A case-sensitive server would
96
+ # otherwise silently drop node summaries.
97
+ f"/doc/{self._enc(doc_id)}/?type=tree&summary={'true' if node_summary else 'false'}",
89
98
  "Failed to get tree result",
90
99
  )
91
100
  return response.json()
@@ -147,6 +156,12 @@ class LegacyCloudAPI:
147
156
  payload["temperature"] = temperature
148
157
  if enable_citations:
149
158
  payload["enable_citations"] = enable_citations
159
+ # Forward stream_metadata so the wire request matches the caller's intent
160
+ # (and stays correct if the server ever gates metadata chunks behind it),
161
+ # mirroring the modern CloudBackend which always sends it. It only affects
162
+ # streaming responses, where it selects the raw dict-chunk parser below.
163
+ if stream_metadata:
164
+ payload["stream_metadata"] = stream_metadata
150
165
 
151
166
  response = self._request(
152
167
  "POST",
@@ -224,7 +239,11 @@ class LegacyCloudAPI:
224
239
  f"/doc/{self._enc(doc_id)}/",
225
240
  "Failed to delete document",
226
241
  )
227
- return response.json()
242
+ # A successful DELETE may come back with an empty body (the documented
243
+ # examples don't consume one, and REST APIs commonly return no content
244
+ # for deletes). Don't let json() raise JSONDecodeError on success —
245
+ # the document is already gone; return an empty dict.
246
+ return response.json() if response.content else {}
228
247
 
229
248
  def list_documents(
230
249
  self,