fix: code review fixes, cloud streaming, naming alignment

- Critical: preserve text in markdown structure for fallback retrieval
- Cloud: SSE response close, folder cache dict, truncate error body
- Cloud: filter internal tools, async-safe streaming via to_thread
- SQLite: multi-thread connection tracking, context manager
- Security: collection name validation, parse_pages range cap
- Polish: use count_tokens wrapper, _EXAMPLES_DIR naming, QueryStream public
- Backend protocol: add @runtime_checkable

Author: KylinMountain

.gitignore

@@ -5,5 +5,6 @@ __pycache__
 .venv/
 logs/
 pageindex.egg-info/
+*.db
 venv/
 uv.lock

examples/cloud_demo.py

@@ -1,20 +1,26 @@
 """
-PageIndex Cloud Demo
+Agentic Vectorless RAG with PageIndex SDK - Cloud Demo
 
-Usage:
+Uses CloudClient for fully-managed document indexing and QA.
+No LLM API key needed — the cloud service handles everything.
+
+Steps:
+  1 — Upload and index a PDF via PageIndex cloud
+  2 — Stream a question with tool call visibility
+
+Requirements:
     pip install pageindex
     export PAGEINDEX_API_KEY=your-api-key
-    python examples/cloud_demo.py
 """
 import asyncio
 import os
 from pathlib import Path
 import requests
 from pageindex import CloudClient
 
-_DIR = Path(__file__).parent
+_EXAMPLES_DIR = Path(__file__).parent
 PDF_URL = "https://arxiv.org/pdf/1706.03762.pdf"
-PDF_PATH = _DIR / "documents" / "attention.pdf"
+PDF_PATH = _EXAMPLES_DIR / "documents" / "attention.pdf"
 
 # Download PDF if needed
 if not PDF_PATH.exists():

examples/local_demo.py

@@ -1,19 +1,25 @@
 """
-PageIndex Local Demo
+Agentic Vectorless RAG with PageIndex SDK - Local Demo
 
-Usage:
-    pip install pageindex
-    python examples/local_demo.py
+A simple example of using LocalClient for self-hosted document indexing
+and agent-based QA. The agent uses OpenAI Agents SDK to reason over
+the document's tree structure index.
+
+Steps:
+  1 — Download and index a PDF
+  2 — Stream a question with tool call visibility
+
+Requirements: pip install pageindex openai-agents
 """
 import asyncio
 from pathlib import Path
 import requests
 from pageindex import LocalClient
 
-_DIR = Path(__file__).parent
+_EXAMPLES_DIR = Path(__file__).parent
 PDF_URL = "https://arxiv.org/pdf/1706.03762.pdf"
-PDF_PATH = _DIR / "documents" / "attention.pdf"
-WORKSPACE = _DIR / "workspace"
+PDF_PATH = _EXAMPLES_DIR / "documents" / "attention.pdf"
+WORKSPACE = _EXAMPLES_DIR / "workspace"
 
 # Download PDF if needed
 if not PDF_PATH.exists():

pageindex/__init__.py

@@ -6,6 +6,7 @@
 
 # SDK exports
 from .client import PageIndexClient, LocalClient, CloudClient
+from .config import IndexConfig
 from .collection import Collection
 from .parser.protocol import ContentNode, ParsedDocument, DocumentParser
 from .storage.protocol import StorageEngine
@@ -23,6 +24,7 @@
     "PageIndexClient",
     "LocalClient",
     "CloudClient",
+    "IndexConfig",
     "Collection",
     "ContentNode",
     "ParsedDocument",

pageindex/backend/cloud.py

@@ -20,11 +20,14 @@
 
 API_BASE = "https://api.pageindex.ai"
 
+_INTERNAL_TOOLS = frozenset({"ToolSearch", "Read", "Grep", "Glob", "Bash", "Edit", "Write"})
+
 
 class CloudBackend:
     def __init__(self, api_key: str):
         self._api_key = api_key
         self._headers = {"api_key": api_key}
+        self._folder_id_cache: dict[str, str | None] = {}
 
     # ── HTTP helpers ──────────────────────────────────────────────────────
 
@@ -38,7 +41,8 @@ def _request(self, method: str, path: str, **kwargs) -> dict:
                     time.sleep(2 ** attempt)
                     continue
                 if resp.status_code != 200:
-                    raise CloudAPIError(f"Cloud API error {resp.status_code}: {resp.text}")
+                    body = resp.text[:500] if resp.text else ""
+                    raise CloudAPIError(f"Cloud API error {resp.status_code}: {body}")
                 return resp.json() if resp.content else {}
             except requests.RequestException as e:
                 if attempt == 2:
@@ -63,12 +67,14 @@ def _enc(value: str) -> str:
     def create_collection(self, name: str) -> None:
         self._validate_collection_name(name)
         try:
-            self._request("POST", "/folder/", json={"name": name})
+            resp = self._request("POST", "/folder/", json={"name": name})
+            self._folder_id_cache[name] = resp.get("folder", {}).get("id")
         except CloudAPIError as e:
             if "403" in str(e):
                 logger.warning(
                     "Folders require a Max plan. Upgrade at https://dash.pageindex.ai/subscription"
                 )
+                self._folder_id_cache[name] = None
             else:
                 raise
 
@@ -78,31 +84,33 @@ def get_or_create_collection(self, name: str) -> None:
             data = self._request("GET", "/folders/")
             for folder in data.get("folders", []):
                 if folder.get("name") == name:
-                    self._folder_id_cache = folder["id"]
+                    self._folder_id_cache[name] = folder["id"]
                     return
             resp = self._request("POST", "/folder/", json={"name": name})
-            self._folder_id_cache = resp.get("folder", {}).get("id")
+            self._folder_id_cache[name] = resp.get("folder", {}).get("id")
         except CloudAPIError as e:
             if "403" in str(e):
                 logger.warning(
                     "Folders require a Max plan. Documents will be stored without folder organization. "
                     "Upgrade at https://dash.pageindex.ai/subscription"
                 )
-                self._folder_id_cache = None
+                self._folder_id_cache[name] = None
             else:
                 raise
 
     def _get_folder_id(self, name: str) -> str | None:
         """Resolve collection name to folder ID. Returns None if folders not available."""
-        if hasattr(self, '_folder_id_cache'):
-            return self._folder_id_cache
+        if name in self._folder_id_cache:
+            return self._folder_id_cache.get(name)
         try:
             data = self._request("GET", "/folders/")
             for folder in data.get("folders", []):
                 if folder.get("name") == name:
+                    self._folder_id_cache[name] = folder["id"]
                     return folder["id"]
         except CloudAPIError:
             pass
+        self._folder_id_cache[name] = None
         return None
 
     def list_collections(self) -> list[str]:
@@ -204,67 +212,90 @@ async def query_stream(self, collection: str, question: str,
         - mcp_tool_use_start: tool call started (has tool_name, server_name)
         - tool_use: tool call argument delta
         - tool_use_stop: tool call ended
+
+        Note: Uses synchronous ``requests`` under the hood.  The blocking
+        HTTP call and line iteration are offloaded to a thread via
+        ``asyncio.to_thread`` so the caller's event loop is not blocked.
+        For full async streaming consider migrating to ``httpx.AsyncClient``.
         """
+        import asyncio
+
         doc_id = doc_ids if doc_ids else self._get_all_doc_ids(collection)
-        resp = requests.post(
-            f"{API_BASE}/chat/completions/",
-            headers=self._headers,
-            json={
-                "messages": [{"role": "user", "content": question}],
-                "doc_id": doc_id,
-                "stream": True,
-                "stream_metadata": True,
-            },
-            stream=True,
-            timeout=120,
-        )
-        if resp.status_code != 200:
-            raise CloudAPIError(f"Cloud streaming error {resp.status_code}: {resp.text}")
-
-        current_tool_name = None
-        current_tool_args = []
-
-        for line in resp.iter_lines(decode_unicode=True):
-            if not line or not line.startswith("data: "):
-                continue
-            data_str = line[6:]
-            if data_str.strip() == "[DONE]":
-                break
-            try:
-                chunk = json.loads(data_str)
-            except json.JSONDecodeError:
-                continue
-
-            meta = chunk.get("block_metadata", {})
-            block_type = meta.get("type", "")
-            choices = chunk.get("choices", [])
-            delta = choices[0].get("delta", {}) if choices else {}
-            content = delta.get("content", "")
-
-            if block_type == "mcp_tool_use_start":
-                current_tool_name = meta.get("tool_name", "")
-                current_tool_args = []
-
-            elif block_type == "tool_use":
-                if content:
-                    current_tool_args.append(content)
-
-            elif block_type == "tool_use_stop":
-                # Skip internal tools (ToolSearch, Read, Grep, etc.)
-                _INTERNAL_TOOLS = {"ToolSearch", "Read", "Grep", "Glob", "Bash", "Edit", "Write"}
-                if current_tool_name and current_tool_name not in _INTERNAL_TOOLS:
-                    args_str = "".join(current_tool_args)
-                    yield QueryEvent(type="tool_call", data={
-                        "name": current_tool_name,
-                        "args": args_str,
-                    })
-                current_tool_name = None
-                current_tool_args = []
 
-            elif block_type == "text" and content:
-                yield QueryEvent(type="answer_delta", data=content)
+        headers = self._headers
+
+        def _stream() -> list[tuple[str, object]]:
+            """Execute the blocking SSE request in a worker thread."""
+            resp = requests.post(
+                f"{API_BASE}/chat/completions/",
+                headers=headers,
+                json={
+                    "messages": [{"role": "user", "content": question}],
+                    "doc_id": doc_id,
+                    "stream": True,
+                    "stream_metadata": True,
+                },
+                stream=True,
+                timeout=120,
+            )
+            try:
+                if resp.status_code != 200:
+                    body = resp.text[:500] if resp.text else ""
+                    raise CloudAPIError(
+                        f"Cloud streaming error {resp.status_code}: {body}"
+                    )
 
-        yield QueryEvent(type="answer_done", data="")
+                events: list[tuple[str, object]] = []
+                current_tool_name = None
+                current_tool_args: list[str] = []
+
+                for line in resp.iter_lines(decode_unicode=True):
+                    if not line or not line.startswith("data: "):
+                        continue
+                    data_str = line[6:]
+                    if data_str.strip() == "[DONE]":
+                        break
+                    try:
+                        chunk = json.loads(data_str)
+                    except json.JSONDecodeError:
+                        continue
+
+                    meta = chunk.get("block_metadata", {})
+                    block_type = meta.get("type", "")
+                    choices = chunk.get("choices", [])
+                    delta = choices[0].get("delta", {}) if choices else {}
+                    content = delta.get("content", "")
+
+                    if block_type == "mcp_tool_use_start":
+                        current_tool_name = meta.get("tool_name", "")
+                        current_tool_args = []
+
+                    elif block_type == "tool_use":
+                        if content:
+                            current_tool_args.append(content)
+
+                    elif block_type == "tool_use_stop":
+                        # Skip internal tools (ToolSearch, Read, Grep, etc.)
+                        if current_tool_name and current_tool_name not in _INTERNAL_TOOLS:
+                            args_str = "".join(current_tool_args)
+                            events.append(("tool_call", {
+                                "name": current_tool_name,
+                                "args": args_str,
+                            }))
+                        current_tool_name = None
+                        current_tool_args = []
+
+                    elif block_type == "text" and content:
+                        events.append(("answer_delta", content))
+
+                events.append(("answer_done", ""))
+                return events
+            finally:
+                resp.close()
+
+        events = await asyncio.to_thread(_stream)
+        for event_type, event_data in events:
+            yield QueryEvent(type=event_type, data=event_data)
 
     def _get_all_doc_ids(self, collection: str) -> list[str]:
         """Get all document IDs in a collection."""

pageindex/backend/local.py

@@ -1,6 +1,7 @@
 # pageindex/backend/local.py
 import hashlib
 import os
+import re
 import uuid
 import shutil
 from pathlib import Path
@@ -12,17 +13,19 @@
 from ..index.pipeline import build_index
 from ..index.utils import parse_pages, get_pdf_page_content, get_md_page_content, remove_fields
 from ..backend.protocol import AgentTools
-from ..errors import FileTypeError, DocumentNotFoundError, IndexingError
+from ..errors import FileTypeError, DocumentNotFoundError, IndexingError, PageIndexError
+
+_COLLECTION_NAME_RE = re.compile(r'^[a-zA-Z0-9_-]{1,128}$')
 
 
 class LocalBackend:
     def __init__(self, storage: StorageEngine, files_dir: str, model: str = None,
-                 retrieve_model: str = None, api_key: str = None):
+                 retrieve_model: str = None, index_config=None):
         self._storage = storage
         self._files_dir = Path(files_dir)
         self._model = model
         self._retrieve_model = retrieve_model or model
-        self._api_key = api_key
+        self._index_config = index_config
         self._parsers: list[DocumentParser] = [PdfParser(), MarkdownParser()]
 
     def register_parser(self, parser: DocumentParser) -> None:
@@ -39,10 +42,16 @@ def _resolve_parser(self, file_path: str) -> DocumentParser:
         raise FileTypeError(f"No parser for extension: {ext}")
 
     # Collection management
+    def _validate_collection_name(self, name: str) -> None:
+        if not _COLLECTION_NAME_RE.match(name):
+            raise PageIndexError(f"Invalid collection name: {name!r}. Must be 1-128 chars of [a-zA-Z0-9_-].")
+
     def create_collection(self, name: str) -> None:
+        self._validate_collection_name(name)
         self._storage.create_collection(name)
 
     def get_or_create_collection(self, name: str) -> None:
+        self._validate_collection_name(name)
         self._storage.get_or_create_collection(name)
 
     def list_collections(self) -> list[str]:
@@ -87,21 +96,26 @@ def add_document(self, collection: str, file_path: str) -> str:
 
         try:
             parsed = parser.parse(file_path, model=self._model)
-            result = build_index(parsed, model=self._model)
+            result = build_index(parsed, model=self._model, opt=self._index_config)
 
             # Cache page text for fast retrieval (avoids re-reading files)
             pages = [{"page": n.index, "content": n.content}
                      for n in parsed.nodes if n.content]
 
-            # Strip text from structure to save storage space
-            clean_structure = remove_fields(result["structure"], fields=["text"])
+            # 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": ext.lstrip("."),
+                "doc_type": doc_type,
                 "structure": clean_structure,
                 "pages": pages,
             })