diff --git a/.github/workflows/publish.yml b/.github/workflows/publish.yml new file mode 100644 index 000000000..e96dc4cca --- /dev/null +++ b/.github/workflows/publish.yml @@ -0,0 +1,85 @@ +name: Publish to PyPI + +# Release flow (the git tag IS the version — nothing to bump in the repo): +# 1. git tag -a v0.2.9 -m "Release 0.2.9" +# 2. git push origin v0.2.9 +# 3. This workflow derives the version from the tag, injects it into +# pyproject.toml, builds, publishes to PyPI via OIDC trusted publishing +# (no stored secret), and creates a GitHub Release with generated notes. +# +# The tag must be a PEP 440 version with a leading `v`: +# v0.2.9 v0.2.9rc1 v0.2.9.dev1 +# PyPI rejects duplicate version uploads, so each tag must be a new version. +# Plain `pip install pageindex` skips dev/rc pre-releases — install one +# explicitly with `pip install pageindex==0.2.9.dev1`. +# +# One-time setup this workflow depends on: +# - PyPI: add a Trusted Publisher on the `pageindex` project pointing at +# repo VectifyAI/PageIndex, workflow `publish.yml`, environment `pypi`. +# - GitHub: create an Environment named `pypi` (Settings -> Environments). + +on: + push: + tags: + - "v*" + +jobs: + test: + runs-on: ubuntu-latest + permissions: + contents: read + strategy: + fail-fast: false + matrix: + python-version: ["3.10", "3.13"] + name: test py${{ matrix.python-version }} + timeout-minutes: 15 + steps: + - uses: actions/checkout@9c091bb21b7c1c1d1991bb908d89e4e9dddfe3e0 # v7.0.0 + - uses: actions/setup-python@ece7cb06caefa5fff74198d8649806c4678c61a1 # v6.3.0 + with: + python-version: ${{ matrix.python-version }} + cache: pip + - run: pip install -r requirements.txt pytest + - run: python -m pytest -q + + publish: + needs: test + runs-on: ubuntu-latest + environment: pypi + permissions: + id-token: write # OIDC trusted publishing to PyPI + contents: write # create the GitHub Release + steps: + - uses: actions/checkout@9c091bb21b7c1c1d1991bb908d89e4e9dddfe3e0 # v7.0.0 + + - uses: actions/setup-python@ece7cb06caefa5fff74198d8649806c4678c61a1 # v6.3.0 + with: + python-version: "3.12" + + - name: Set version from tag and build + run: | + set -euo pipefail + python -m pip install --upgrade build packaging + VERSION="${GITHUB_REF_NAME#v}" + echo "Publishing version: $VERSION" + # Fail early on a malformed tag instead of publishing a junk version. + export VERSION + python -c "import os; from packaging.version import Version; Version(os.environ['VERSION'])" + # The git tag is the single source of truth; overwrite the static + # placeholder in [tool.poetry] so the built artifacts carry $VERSION. + sed -i "s/^version = .*/version = \"$VERSION\"/" pyproject.toml + grep '^version = ' pyproject.toml + python -m build + + - name: Publish to PyPI + uses: pypa/gh-action-pypi-publish@cef221092ed1bacb1cc03d23a2d87d1d172e277b # release/v1.14.0 + + - name: Create GitHub Release + uses: softprops/action-gh-release@718ea10b132b3b2eba29c1007bb80653f286566b # v3.0.1 + with: + tag_name: ${{ github.ref_name }} + name: ${{ github.ref_name }} + prerelease: ${{ contains(github.ref_name, 'rc') || contains(github.ref_name, 'dev') || contains(github.ref_name, 'alpha') || contains(github.ref_name, 'beta') }} + generate_release_notes: true + files: dist/* diff --git a/.github/workflows/tests.yml b/.github/workflows/tests.yml index fc25d6c45..d8d9dbb38 100644 --- a/.github/workflows/tests.yml +++ b/.github/workflows/tests.yml @@ -20,22 +20,16 @@ jobs: fail-fast: false matrix: python-version: ["3.10", "3.13"] - # "without" proves the package works with no agent framework - # installed; "with" covers the framework integration paths. agent-frameworks: [without, with] name: py${{ matrix.python-version }} (${{ matrix.agent-frameworks }} frameworks) timeout-minutes: 15 steps: - - uses: actions/checkout@v4 - - uses: actions/setup-python@v5 + - uses: actions/checkout@9c091bb21b7c1c1d1991bb908d89e4e9dddfe3e0 # v7.0.0 + - uses: actions/setup-python@ece7cb06caefa5fff74198d8649806c4678c61a1 # v6.3.0 with: python-version: ${{ matrix.python-version }} cache: pip - # requirements.txt works on every branch; the packaging metadata - # (pyproject) does not exist on all of them. - run: pip install -r requirements.txt pytest - if: matrix.agent-frameworks == 'with' run: pip install openai-agents claude-agent-sdk - # python -m pytest puts the repo root on sys.path, so the in-repo - # `pageindex` package is imported without an install step. - run: python -m pytest -q diff --git a/.gitignore b/.gitignore index 23d6b5655..5193735ca 100644 --- a/.gitignore +++ b/.gitignore @@ -4,3 +4,5 @@ __pycache__ .env* .venv/ logs/ +.pageindex/ +dist/ diff --git a/examples/agentic_vectorless_rag_demo.py b/examples/agentic_vectorless_rag_demo.py index b4ed9c2f8..4fe5f179f 100644 --- a/examples/agentic_vectorless_rag_demo.py +++ b/examples/agentic_vectorless_rag_demo.py @@ -1,10 +1,10 @@ """ Agentic Vectorless RAG with PageIndex - Demo -A simple example of building a document QA agent with self-hosted PageIndex -and the OpenAI Agents SDK. Instead of vector similarity search and chunking, -PageIndex builds a hierarchical tree index and uses agentic LLM reasoning for -human-like, context-aware retrieval. +A simple example of building a document QA agent with the PageIndex SDK in +local mode and the OpenAI Agents SDK. Instead of vector similarity search and +chunking, PageIndex builds a hierarchical tree index and uses agentic LLM +reasoning for human-like, context-aware retrieval. Agent tools: - get_document() — document metadata (status, page count, etc.) @@ -12,11 +12,11 @@ - get_page_content() — retrieve text content of specific pages Steps: - 1 — Index a PDF and view its tree structure index + 1 — Index a PDF locally and view its tree structure index 2 — View document metadata 3 — Ask a question (agent reasons over the index and auto-calls tools) -Requirements: pip install openai-agents +Requirements: pip install openai-agents; OPENAI_API_KEY in the environment. """ import sys import json @@ -39,12 +39,12 @@ _EXAMPLES_DIR = Path(__file__).parent PDF_PATH = _EXAMPLES_DIR / "documents" / "attention-residuals.pdf" -WORKSPACE = _EXAMPLES_DIR / "workspace" +STORAGE_PATH = _EXAMPLES_DIR / ".pageindex" AGENT_SYSTEM_PROMPT = """ You are PageIndex, a document QA assistant. TOOL USE: -- Call get_document() first to confirm status and page/line count. +- Call get_document() first to confirm status and page count. - Call get_document_structure() to identify relevant page ranges. - Call get_page_content(pages="5-7") with tight ranges; never fetch the whole document. - Before each tool call, output one short sentence explaining the reason. @@ -62,27 +62,26 @@ def query_agent(client: PageIndexClient, doc_id: str, prompt: str, verbose: bool @function_tool def get_document() -> str: """Get document metadata: status, page count, name, and description.""" - return client.get_document(doc_id) + return json.dumps(client.get_document(doc_id)) @function_tool def get_document_structure() -> str: """Get the document's full tree structure (without text) to find relevant sections.""" - return client.get_document_structure(doc_id) + return json.dumps(client.get_document_structure(doc_id), ensure_ascii=False) @function_tool def get_page_content(pages: str) -> str: """ - Get the text content of specific pages or line numbers. + Get the text content of specific pages. Use tight ranges: e.g. '5-7' for pages 5 to 7, '3,8' for pages 3 and 8, '12' for page 12. - For Markdown documents, use line numbers from the structure's line_num field. """ - return client.get_page_content(doc_id, pages) + return json.dumps(client.get_page_content(doc_id, pages), ensure_ascii=False) agent = Agent( name="PageIndex", instructions=AGENT_SYSTEM_PROMPT, tools=[get_document, get_document_structure, get_page_content], - model=client.retrieve_model, + model=getattr(client, "retrieve_model", None), # model_settings=ModelSettings(reasoning={"effort": "low", "summary": "auto"}), # Uncomment to enable reasoning ) @@ -152,24 +151,25 @@ async def _run(): f.write(chunk) print("Download complete.\n") - # Setup - client = PageIndexClient(workspace=WORKSPACE) + # Setup: local mode — no PageIndex API key needed, your LLM key does the work + client = PageIndexClient(storage_path=str(STORAGE_PATH)) # Step 1: Index PDF and view tree structure print("=" * 60) print("Step 1: Index PDF and view tree structure") print("=" * 60) doc_id = next( - (did for did, doc in client.documents.items() if doc.get('doc_name') == PDF_PATH.name), + (doc["id"] for doc in client.list_documents(limit=100)["documents"] + if doc["name"] == PDF_PATH.name), None, ) if doc_id: print(f"\nLoaded cached doc_id: {doc_id}") else: - doc_id = client.index(PDF_PATH) + doc_id = client.submit_document(str(PDF_PATH))["doc_id"] print(f"\nIndexed. doc_id: {doc_id}") print("\nTree Structure (top-level sections):") - structure = json.loads(client.get_document_structure(doc_id)) + structure = client.get_tree(doc_id, node_summary=True)["result"] utils.print_tree(structure) # Step 2: View document metadata diff --git a/examples/workspace/12345678-abcd-4321-abcd-123456789abc.json b/examples/workspace/12345678-abcd-4321-abcd-123456789abc.json deleted file mode 100644 index 23351d5c5..000000000 --- a/examples/workspace/12345678-abcd-4321-abcd-123456789abc.json +++ /dev/null @@ -1,274 +0,0 @@ -{ - "id": "12345678-abcd-4321-abcd-123456789abc", - "type": "pdf", - "path": "../documents/attention-residuals.pdf", - "doc_name": "attention-residuals.pdf", - "doc_description": "This document introduces \"Attention Residuals\" (AttnRes) and its scalable variant \"Block AttnRes,\" novel mechanisms for replacing fixed residual accumulation in neural networks with learned, input-dependent depth-wise attention, addressing limitations of standard residual connections while optimizing memory, computation, and scalability for large-scale training and inference.", - "page_count": 21, - "structure": [ - { - "title": "Preface", - "node_id": "0000", - "start_index": 1, - "end_index": 2, - "summary": "The partial document introduces \"Attention Residuals\" (AttnRes), a novel approach to replace fixed residual accumulation in large language models (LLMs) with learned, input-dependent softmax attention over preceding layer outputs. This method addresses issues like uncontrolled hidden-state growth and dilution of layer contributions caused by standard residual connections with PreNorm. To enhance scalability, the document proposes \"Block AttnRes,\" which partitions layers into blocks and applies attention at the block level, reducing memory and communication overhead while maintaining performance gains. The document highlights system optimizations, such as cross-stage caching and a two-phase computation strategy, to make Block AttnRes efficient for large-scale training. Experiments confirm consistent improvements across model sizes, with AttnRes mitigating PreNorm dilution, leading to more uniform output magnitudes, gradient distributions, and better downstream task performance. Key contributions include the introduction of AttnRes and Block AttnRes, scalable infrastructure optimizations, and comprehensive evaluations demonstrating their effectiveness." - }, - { - "title": "Introduction", - "node_id": "0001", - "start_index": 2, - "end_index": 3, - "summary": "The partial document introduces \"Attention Residuals\" (AttnRes), a novel mechanism that replaces fixed residual accumulation in deep networks with learned softmax attention over depth. It highlights the limitations of standard residual connections, such as uniform layer contributions, irreversible information loss, and output growth, and draws parallels between depth-wise accumulation and sequence modeling in RNNs. AttnRes enables selective, content-aware aggregation of information across layers using attention weights, addressing these limitations. The document also proposes a scalable variant, Block AttnRes, which reduces memory and communication overhead for large-scale training. Key contributions include the development of AttnRes and Block AttnRes, system optimizations for scalability, and comprehensive evaluations demonstrating improved training dynamics, bounded hidden-state magnitudes, and better gradient distribution. The approach is validated through scaling law experiments, ablations, and downstream benchmarks, showing consistent performance improvements over standard residual connections." - }, - { - "title": "Motivation", - "node_id": "0002", - "start_index": 3, - "end_index": 3, - "summary": "The partial document discusses the concept of Attention Residuals in the context of deep learning models, particularly Transformers. It begins by introducing the notation and structure of input sequences and layers in a Transformer model. The document then explains residual learning, highlighting its importance in training deep networks by enabling gradients to bypass transformations through identity mapping. It expands on the limitations of traditional residual connections and highway networks, such as lack of selective access to earlier layer outputs, irreversible information loss, and output growth issues that destabilize training. To address these limitations, the document proposes Attention Residuals (AttnRes), a mechanism inspired by the duality of time and depth in sequence modeling. This approach introduces layer-specific attention weights to selectively aggregate information from all preceding layers, offering a unified view of time and depth while maintaining computational feasibility." - }, - { - "title": "Attention Residuals: A Unified View of Time and Depth", - "node_id": "0003", - "start_index": 3, - "end_index": 4, - "summary": "The partial document discusses the concept of \"Attention Residuals\" as a mechanism to address limitations in training deep networks with residual connections. It begins by explaining residual learning, its benefits in gradient flow, and its limitations, such as lack of selective access, irreversible information loss, and output growth. The document introduces \"Attention Residuals\" (AttnRes), which generalizes residual connections by allowing layers to selectively aggregate information from all preceding layers using attention mechanisms. It describes \"Full Attention Residuals,\" which compute attention weights over depth with softmax normalization, and highlights their computational and memory overhead. To address scalability challenges, the document proposes \"Block Attention Residuals,\" which partition layers into blocks, reducing memory and communication overhead by applying attention at the block level. The text also outlines the intra-block accumulation process and its efficiency in distributed training setups.", - "nodes": [ - { - "title": "Full Attention Residuals", - "node_id": "0004", - "start_index": 4, - "end_index": 4, - "summary": "The partial document discusses \"Attention Residuals\" in neural networks, focusing on two main approaches: Full Attention Residuals and Block Attention Residuals. \n\n1. **Full Attention Residuals**: This method computes attention weights using a kernel function with RMS normalization to prevent large-magnitude outputs from dominating. It introduces no additional memory overhead during vanilla training but incurs communication and memory overhead in distributed training due to the need to retain and transmit layer outputs across stages. A blockwise optimization strategy is proposed to reduce memory I/O by batching attention computation within groups of layers.\n\n2. **Block Attention Residuals**: This approach partitions layers into blocks, reducing memory and communication overhead by summing layer outputs within each block and applying attention only to block-level representations. This reduces the complexity from O(Ld) to O(Nd), where N is the number of blocks. The method ensures normalization to avoid biases from magnitude differences between blocks.\n\nThe document highlights the trade-offs between memory, computation, and communication overheads in these methods and introduces strategies to optimize their efficiency in distributed training setups." - }, - { - "title": "Block Attention Residuals", - "node_id": "0005", - "start_index": 4, - "end_index": 5, - "summary": "The partial document discusses \"Attention Residuals,\" focusing on two main variants: Full Attention Residuals (Full AttnRes) and Block Attention Residuals (Block AttnRes). \n\n1. **Full Attention Residuals (Full AttnRes):**\n - Defines attention weights using a kernel function with RMS normalization to prevent large-magnitude outputs from dominating.\n - Requires O(L²d) arithmetic and O(Ld) memory, with no additional memory overhead during vanilla training.\n - Highlights challenges in large-scale training, such as memory and communication overhead under pipeline parallelism.\n - Introduces blockwise optimization to reduce memory I/O but notes that cross-stage communication remains a bottleneck.\n\n2. **Block Attention Residuals (Block AttnRes):**\n - Partitions layers into blocks, reducing memory and communication overhead from O(Ld) to O(Nd) by summing layer outputs within blocks and applying attention over block-level representations.\n - Provides PyTorch-style pseudocode for implementation, detailing intra-block accumulation and inter-block attention mechanisms.\n - Improves efficiency by reducing memory and computation requirements, with block count N interpolating between Full AttnRes (N=L) and standard residual connections (N=1).\n - Enhances inference latency and bounds KV cache size through blockwise optimization.\n\nThe document also addresses infrastructure challenges for large-scale training, emphasizing the need to manage communication overhead and optimize system design for block-based attention mechanisms." - } - ] - }, - { - "title": "Infrastructure Design", - "node_id": "0006", - "start_index": 5, - "end_index": 6, - "summary": "The partial document describes the concept and implementation of Block Attention Residuals (Block AttnRes), a mechanism designed to improve memory and computational efficiency in attention-based models. It introduces inter-block attention, where attention is computed over block representations and partial sums, reducing memory and computation from O(L) and O(L²) to O(N) and O(N²), respectively. The document provides PyTorch-style pseudocode for the implementation, detailing how block representations and partial sums are managed across layers. It highlights the efficiency benefits of using block representations instead of individual outputs, with empirical findings suggesting that a block count of N≈8 balances performance and resource usage. \n\nThe document also addresses infrastructure challenges in large-scale training and inference. It discusses pipeline communication optimizations, such as cross-stage caching, to reduce redundant data transmission and improve efficiency during distributed training. For inference, it proposes a two-phase computation strategy and memory-efficient prefilling to handle long-context scenarios. An example of cache-based pipeline communication is provided, illustrating how caching minimizes communication overhead in distributed systems.", - "nodes": [ - { - "title": "Training", - "node_id": "0007", - "start_index": 6, - "end_index": 7, - "summary": "The partial document discusses the optimization of Attention Residuals (AttnRes) in training and inference for large-scale distributed systems. It introduces cross-stage caching to address communication and memory overheads in pipeline parallelism, reducing redundant data transmission and improving efficiency. The document details a two-phase computation strategy for Block AttnRes, which includes parallel inter-block attention and sequential intra-block attention with online softmax merging. This approach minimizes memory access and I/O overhead while maintaining a low training overhead. Additionally, it highlights the memory-efficient prefilling scheme for long-context inputs and explains how Block AttnRes compresses representations to reduce storage requirements. The document also provides algorithmic details and performance improvements in both training and inference scenarios." - }, - { - "title": "Inference", - "node_id": "0008", - "start_index": 7, - "end_index": 8, - "summary": "The partial document describes the technical details and implementation of Attention Residuals (AttnRes) in neural network architectures. It introduces a two-phase computation strategy for block-based attention, optimizing memory and computational efficiency. Phase 1 handles parallel inter-block attention, while Phase 2 processes sequential intra-block attention with an online softmax merge. The document highlights memory overhead reduction through cross-stage caching, sequence-sharded prefilling, and kernel fusion, achieving minimal training and inference latency overhead. It compares memory access costs across different residual mechanisms and demonstrates the efficiency of AttnRes, particularly in Block AttnRes, which compresses block representations. Experimental results show that AttnRes improves scaling behavior and validation loss compared to baseline models, with negligible parameter overhead and consistent performance gains across compute ranges." - } - ] - }, - { - "title": "Experiments", - "node_id": "0009", - "start_index": 8, - "end_index": 8, - "summary": "The partial document discusses the technical details and performance of the Attention Residuals (AttnRes) mechanism in transformer architectures. It highlights the memory efficiency and reduced inference latency of AttnRes compared to prior residual mechanisms like mHC. The document provides a breakdown of memory access costs for different schemes, emphasizing the two-phase inference schedule of AttnRes and its memory-efficient prefilling strategy, which significantly reduces memory overhead through sharding and chunked prefill techniques. It also describes the integration of AttnRes into a Mixture-of-Experts (MoE) Transformer architecture, detailing its minimal parameter addition and initialization strategy to ensure stable training. Additionally, the document presents experimental results, including scaling laws and validation loss comparisons across model variants, demonstrating that AttnRes achieves consistently lower loss while maintaining similar scaling behavior to the baseline.", - "nodes": [ - { - "title": "Scaling Laws", - "node_id": "0010", - "start_index": 8, - "end_index": 9, - "summary": "The partial document discusses the implementation and evaluation of Attention Residuals (AttnRes) in transformer architectures. It highlights the memory efficiency and reduced inference latency of AttnRes compared to prior residual mechanisms like mHC. The document introduces a two-phase inference schedule for AttnRes, optimizing memory access costs and reducing per-device memory usage through sharding and chunked prefill techniques. It describes the integration of AttnRes into a Mixture-of-Experts (MoE) Transformer architecture, maintaining minimal parameter overhead and ensuring stable training through specific initialization strategies. Experiments compare scaling laws and validation loss across model sizes, showing that both Full and Block AttnRes outperform baselines and mHC in terms of loss and compute efficiency. The main results include training recipes for large-scale models, leveraging hybrid attention mechanisms and progressive sequence length extension without additional modifications." - }, - { - "title": "Main Results", - "node_id": "0011", - "start_index": 9, - "end_index": 11, - "summary": "The partial document discusses the concept of Attention Residuals (AttnRes) in transformer models, comparing its performance and efficiency against baseline models and other methods. Key points include:\n\n1. **Model Configurations and Validation Loss**: Comparison of Baseline, Block AttnRes, Full AttnRes, and mHC(-lite) models across various configurations, showing that AttnRes consistently achieves lower validation loss, with Block AttnRes closely tracking Full AttnRes.\n\n2. **Scaling Laws**: Analysis of scaling behavior, demonstrating that Block AttnRes achieves significant compute efficiency and narrows the performance gap with Full AttnRes at larger scales.\n\n3. **Training Recipe**: Description of the training process for large models, including pre-training and mid-training phases, use of hybrid attention mechanisms, and progressive sequence length extension.\n\n4. **Training Dynamics**: Examination of validation loss, output magnitude, and gradient magnitude during training, highlighting how Block AttnRes mitigates issues like PreNorm dilution and uneven gradient flow.\n\n5. **Downstream Performance**: Evaluation of AttnRes on various benchmarks for language understanding, reasoning, and code/math tasks, showing consistent improvements over the baseline, particularly in multi-step reasoning and compositional tasks.\n\n6. **Ablation Study**: Validation of key design choices in AttnRes, comparing it with prior methods like DenseFormer and mHC. Full AttnRes achieves the best performance, while Block AttnRes offers a memory-efficient trade-off with competitive results.\n\n7. **Cross-Layer Access**: Exploration of different granularities of cross-layer access, with Block AttnRes providing an effective balance between performance and memory efficiency, and Full AttnRes offering the best results at higher memory costs." - }, - { - "title": "Ablation Study", - "node_id": "0012", - "start_index": 11, - "end_index": 12, - "summary": "The partial document focuses on the development and evaluation of Attention Residuals (AttnRes), a novel mechanism for improving Transformer models. Key points include:\n\n1. **Ablation Studies**: The document evaluates the impact of various design choices in AttnRes, such as input-dependent queries, input-independent mixing, softmax vs. sigmoid, multihead attention, and RMSNorm. Results show that input-dependent queries and RMSNorm improve performance, while softmax outperforms sigmoid due to sharper selection.\n\n2. **Comparison with Prior Methods**: AttnRes is compared against baseline PreNorm, DenseFormer, and mHC. AttnRes achieves superior performance, with Full AttnRes and Block AttnRes showing significant improvements in validation loss.\n\n3. **Cross-Layer Access**: Different granularities of cross-layer access are analyzed. Full AttnRes achieves the best performance, while Block AttnRes offers a memory-efficient trade-off. Sliding-window aggregation (SWA) is less effective, highlighting the importance of selectively accessing distant layers.\n\n4. **Performance on Benchmarks**: AttnRes outperforms the baseline on various benchmarks, particularly in multi-step reasoning tasks, code generation, and knowledge-oriented tasks, demonstrating its effectiveness in compositional tasks.\n\n5. **Optimal Architecture Analysis**: The study explores how AttnRes reshapes architectural scaling under fixed compute and parameter budgets. AttnRes favors deeper models with a shift in the optimal depth–width–attention trade-off, achieving consistently lower losses across configurations compared to the baseline.\n\n6. **Validation Loss Trends**: The document provides detailed validation loss trends across different configurations and block sizes, showing graceful degradation with increasing block size and highlighting the efficiency of finer-grained configurations." - }, - { - "title": "Analysis", - "node_id": "0013", - "start_index": 12, - "end_index": 12, - "summary": "The partial document discusses the evaluation and analysis of Attention Residuals (AttnRes) in Transformer architectures. Key points include:\n\n1. **Architecture Sweep**: A study under fixed compute and parameter budgets to analyze validation loss across different configurations of model depth (dmodel/Lb) and attention heads (H/Lb). AttnRes consistently outperforms the baseline in all configurations, with a notable shift in optimal depth from dmodel/Lb ≈ 60 (baseline) to dmodel/Lb ≈ 45 (AttnRes).\n\n2. **Component Design Ablations**:\n - **Input-dependent query**: Improves performance but adds computational complexity.\n - **Input-independent mixing**: Degrades performance compared to learned queries.\n - **Softmax vs. Sigmoid**: Softmax performs better due to sharper selection among sources.\n - **Multihead Attention**: Depth aggregation across heads reduces performance, indicating uniform depth-wise mixtures are optimal.\n - **RMSNorm on Keys**: Removing RMSNorm negatively impacts performance, especially for block-level representations, by preventing bias in attention weights.\n\n3. **Optimal Architecture Analysis**: Investigates how AttnRes influences depth–width–attention trade-offs under fixed compute and parameter constraints. AttnRes favors deeper models and achieves lower loss compared to conventional Transformer designs.", - "nodes": [ - { - "title": "Optimal Architecture", - "node_id": "0014", - "start_index": 12, - "end_index": 13, - "summary": "The partial document discusses the concept of Attention Residuals (AttnRes) in Transformer architectures, focusing on their design, performance, and analysis. Key points include:\n\n1. **Component Design Ablations**: The document evaluates various modifications to the attention mechanism, such as input-dependent queries, input-independent mixing, softmax vs. sigmoid, multihead attention, and RMSNorm on keys. These experiments highlight the impact of each component on performance, with findings such as the importance of softmax for competitive normalization and RMSNorm for preventing bias in attention weights.\n\n2. **Optimal Architecture Analysis**: A controlled study under fixed compute and parameter budgets examines how AttnRes reshapes architectural scaling preferences. Results show that AttnRes favors deeper, narrower networks compared to baseline Transformers, achieving lower validation loss across configurations. The optimal configuration shifts to a lower dmodel/Lb ratio, indicating better exploitation of depth.\n\n3. **Learned AttnRes Patterns**: Visualization of learned attention weights reveals key insights:\n - Preserved locality with layers attending strongly to immediate predecessors while forming selective skip connections.\n - Layer specialization, with embeddings retaining weight and distinct patterns in pre-attention and pre-MLP layers.\n - Block AttnRes effectively preserves structural patterns while acting as implicit regularization.\n\n4. **Performance Trends**: AttnRes consistently outperforms the baseline across configurations, with lower validation loss and sharper, more decisive weight distributions in block attention settings." - }, - { - "title": "Analyzing Learned AttnRes Patterns", - "node_id": "0015", - "start_index": 13, - "end_index": 14, - "summary": "The partial document discusses Attention Residuals (AttnRes) in deep learning models, focusing on their structure, behavior, and benefits. Key points include:\n\n1. **Depth-wise Attention Weight Distributions**: Analysis of weight distributions in a 16-head model with full and block Attention Residuals, highlighting diagonal dominance (locality), learned skip connections, and sharper weight distributions in block settings.\n\n2. **Learned AttnRes Patterns**: Observations include preserved locality, layer specialization, and the ability of block AttnRes to maintain essential information pathways while acting as implicit regularization.\n\n3. **Comparison of Residual Update Mechanisms**: A detailed comparison of various residual connection methods, including their update rules, weight types (fixed, learned, or dynamic), and source access.\n\n4. **Sequence-Depth Duality**: Exploration of the analogy between residual connections and RNNs, emphasizing how AttnRes replaces depth-wise recurrence with direct cross-layer attention for improved information propagation.\n\n5. **Residual Connections as Structured Matrices**: Formalization of residual connections as depth mixing matrices, comparing different methods based on weight generation and structural constraints.\n\nThe document emphasizes the advantages of AttnRes in leveraging depth, preserving structure, and enabling efficient information flow across layers." - } - ] - } - ] - }, - { - "title": "Discussions", - "node_id": "0016", - "start_index": 14, - "end_index": 14, - "summary": "The partial document discusses various residual update mechanisms in neural network architectures, comparing their update rules, weight types (fixed, learned-static, or input-dependent), and sources of earlier representations. It categorizes methods into single-state recurrence, multi-state recurrence, and cross-layer access, providing examples like Residual, ReZero, LayerScale, Highway, DeepNorm, KEEL, DenseNet, DenseFormer, MRLA, and AttnRes. The document explores the sequence-depth duality, drawing parallels between residual connections and recurrent neural networks (RNNs), and highlights how AttnRes replaces depth-wise recurrence with direct cross-layer attention. Additionally, it formalizes residual connections as structured matrices, introducing a depth mixing matrix to analyze how different methods aggregate outputs from previous layers, and discusses their weight generation and rank constraints.", - "nodes": [ - { - "title": "Sequence-Depth Duality", - "node_id": "0017", - "start_index": 14, - "end_index": 14, - "summary": "The partial document discusses various residual update mechanisms in neural network architectures, comparing their update rules, weight types (fixed, learned-static, or input-dependent), and sources of earlier representations. It categorizes methods into single-state recurrence, multi-state recurrence, and cross-layer access, providing examples like Residual, ReZero, LayerScale, Highway, DeepNorm, KEEL, DenseNet, DenseFormer, MRLA, and AttnRes. The document explores the sequence-depth duality, drawing parallels between residual connections and recurrent neural networks (RNNs), and highlights how AttnRes replaces depth-wise recurrence with cross-layer attention. Additionally, it formalizes residual connections as structured matrices, introducing a depth mixing matrix to analyze how different methods aggregate outputs from previous layers, and discusses their weight generation and rank constraints." - }, - { - "title": "Residual Connections as Structured Matrices", - "node_id": "0018", - "start_index": 14, - "end_index": 16, - "summary": "The partial document discusses various residual update mechanisms in neural networks, comparing their weight types (fixed, learned, or input-dependent) and source accessibility. It introduces AttnRes, a novel approach that replaces fixed residual accumulation with learned, input-dependent depth-wise attention, inspired by the sequence-depth duality. The document explores structured matrix perspectives, showing how residual variants can be viewed as depth-wise linear attention. It highlights the limitations of existing methods like single-state recurrence and multi-state recurrence, and contrasts them with AttnRes, which provides selective access to earlier-layer outputs. The paper also introduces Block AttnRes, a scalable variant that partitions layers into blocks to reduce memory and computational overhead while retaining performance gains. Empirical results validate the effectiveness of AttnRes and Block AttnRes, with discussions on normalization, scaling, depth stability, and cross-layer connectivity. The document concludes by emphasizing the practicality and scalability of Block AttnRes for large-scale models." - }, - { - "title": "Prior Residuals as Depth-Wise Linear Attention", - "node_id": "0019", - "start_index": 16, - "end_index": 16, - "summary": "The partial document discusses the concept of Attention Residuals (AttnRes), which replaces traditional residual accumulation with learned, input-dependent depth-wise attention. It explores the structured-matrix perspective, sequence-depth duality, and the role of state expansion in depth-wise linear attention. The document addresses challenges in normalization, scaling, and depth stability, comparing PreNorm and PostNorm approaches and introducing AttnRes as a solution to avoid cumulative magnitude growth and gradient vanishing. It highlights multi-state recurrence methods, cross-layer connectivity strategies, and the advantages of AttnRes in selectively accessing earlier-layer outputs. The introduction of Block AttnRes is proposed to address memory constraints in large-scale models by partitioning layers into blocks, reducing computational overhead while maintaining performance. Empirical studies validate the effectiveness of AttnRes and Block AttnRes, with scalability and efficiency improvements highlighted as key contributions." - } - ] - }, - { - "title": "Related Work", - "node_id": "0020", - "start_index": 16, - "end_index": 16, - "summary": "The partial document discusses the concept of Attention Residuals (AttnRes) and its application as depth-wise attention in neural networks. It explores the structured-matrix perspective, highlighting how existing residual variants can be interpreted as linear attention mechanisms over the depth axis. The document addresses challenges in normalization, scaling, and depth stability, comparing PreNorm and PostNorm approaches and introducing AttnRes as a solution to avoid cumulative magnitude growth and gradient vanishing. It also examines multi-state recurrence methods, cross-layer connectivity strategies, and their limitations, proposing AttnRes as a method that selectively aggregates earlier-layer outputs with softmax-normalized, input-dependent weights. The introduction of Block AttnRes is detailed as a scalable alternative to Full AttnRes, reducing memory and computational overhead by partitioning layers into blocks while maintaining performance gains. Empirical validation and practical implementation strategies, such as cross-stage caching and two-phase computation, are also discussed." - }, - { - "title": "Conclusion", - "node_id": "0021", - "start_index": 16, - "end_index": 20, - "summary": "The partial document discusses the concept of Attention Residuals (AttnRes), introducing a novel approach to residual accumulation in neural networks by leveraging depth-wise attention mechanisms. It explores the sequence-depth duality, interpreting residual variants as linear attention over the depth axis. The document highlights the challenges of normalization placement and gradient propagation in standard residual updates, comparing PreNorm and PostNorm methods, and presents AttnRes as a solution that avoids cumulative magnitude growth and gradient vanishing. It also examines multi-state recurrence and cross-layer connectivity, contrasting AttnRes with existing methods like Hyper-Connections, DenseNet, and MUDDFormer, emphasizing its selective access to earlier-layer outputs and efficient scaling. The introduction of Block AttnRes addresses memory constraints by partitioning layers into blocks, reducing computational overhead while maintaining performance. Empirical studies validate the scalability and efficiency of AttnRes, with future directions focusing on finer-grained blocking and hardware advancements." - }, - { - "title": "Contributions", - "node_id": "0022", - "start_index": 20, - "end_index": 21, - "summary": "The partial document discusses the concept of \"Attention Residuals\" and provides a technical explanation of optimized inference input/output (I/O) for Full Attention Residuals. It highlights the inefficiencies of a naïve implementation, where memory traffic scales linearly with depth, and introduces a two-phase scheduling approach to reduce I/O costs. The document explains the partitioning of layers into blocks and details the two phases: Phase 1 (batched inter-block attention) and Phase 2 (sequential intra-block attention). It provides mathematical formulations for read and write costs during these phases and demonstrates how batching inter-block reads reduces per-layer I/O complexity from O(L) to O(S+N). The approach maintains the model architecture while optimizing inference efficiency. Additionally, the document lists the contributors to the work, with equal contributions noted for some authors." - }, - { - "title": "Optimized Inference I/O for Full Attention Residuals", - "node_id": "0023", - "start_index": 21, - "end_index": 21, - "summary": "The partial document discusses an optimized inference I/O strategy for Full Attention Residuals (Full AttnRes) to reduce memory traffic, which scales linearly with model depth in a naïve implementation. It introduces a two-phase scheduling approach for inference, dividing the model into blocks to batch inter-block and intra-block computations. Phase 1 handles batched inter-block attention, reducing redundant memory reads by reusing key-value pairs across layers within a block. Phase 2 processes sequential intra-block dependencies. The document provides detailed calculations for read and write costs during both phases, showing that the proposed method reduces per-layer I/O complexity from O(L) to O(S+N), where S is the block size and N is the number of blocks. The approach maintains the model architecture while optimizing memory efficiency during inference." - } - ], - "pages": [ - { - "page": 1, - "content": "ATTENTIONRESIDUALS\nTECHNICALREPORT OFATTENTIONRESIDUALS\nKimi Team\n/gtbhttps://github.com/MoonshotAI/Attention-Residuals\nABSTRACT\nResidual connections [12] with PreNorm [60] are standard in modern LLMs, yet they accumulate\nall layer outputs with fixed unit weights. This uniform aggregation causes uncontrolled hidden-state\ngrowth with depth, progressively diluting each layer’s contribution [27]. We proposeAttention\nResiduals (AttnRes), which replaces this fixed accumulation with softmax attention over preceding\nlayer outputs, allowing each layer to selectively aggregate earlier representations with learned, input-\ndependent weights. To address the memory and communication overhead of attending over all\npreceding layer outputs for large-scale model training, we introduceBlock AttnRes, which partitions\nlayers into blocks and attends over block-level representations, reducing the memory footprint while\npreserving most of the gains of full AttnRes. Combined with cache-based pipeline communication\nand a two-phase computation strategy, Block AttnRes becomes a practical drop-in replacement for\nstandard residual connections with minimal overhead.\nScaling law experiments confirm that the improvement is consistent across model sizes, and ablations\nvalidate the benefit of content-dependent depth-wise selection. We further integrate AttnRes into\nthe Kimi Linear architecture [69] (48B total / 3B activated parameters) and pre-train on 1.4T tokens,\nwhere AttnRes mitigates PreNorm dilution, yielding more uniform output magnitudes and gradient\ndistribution across depth, and improves downstream performance across all evaluated tasks.\nEmbedding...AttentionMoEAttentionMoEOutput\n(a) Standard ResidualsEmbedding...αAttentionαMoEαAttentionαMoE\nwwwwOutput\nαw\nααααα\n(b) Full Attention ResidualsEmbedding···Blockn-2Blockn-1AttentionMoEAttentionMoEOutput\nα\nαααα\nααααα\nwwwww\nAttnRes Op(α)wQKV\n(c) Block Attention Residuals\nFigure 1: Overview of Attention Residuals.(a)Standard Residuals: standard residual connections with uniform additive accumulation.\n(b)Full AttnRes: each layer selectively aggregates all previous layer outputs via learned attention weights.(c)Block AttnRes: layers\nare grouped into blocks, reducing memory fromO(Ld)toO(Nd).arXiv:2603.15031v1 [cs.CL] 16 Mar 2026" - }, - { - "page": 2, - "content": "Attention ResidualsTECHNICALREPORT\n1 Introduction\nStandard residual connections [12] are thede factobuilding block of modern LLMs [35, 51, 9]. The update hl=\nhl−1+fl−1(hl−1)is widely understood as agradient highwaythat lets gradients bypass transformations via identity\nmappings, enabling stable training at depth. Yet residuals also play a second role that has received less attention.\nUnrolling the recurrence shows that every layer receives the same uniformly-weighted sum of all prior layer outputs;\nresiduals define how information aggregates across depth. Unlike sequence mixing and expert routing, which now\nemploy learnable input-dependent weighting [53, 20, 9], this depth-wise aggregation remains governed by fixed unit\nweights, with no mechanism to selectively emphasize or suppress individual layer contributions.\nIn practice, PreNorm [60] has become the dominant paradigm, yet its unweighted accumulation causes hidden-state\nmagnitudes to grow as O(L) with depth, progressively diluting each layer’s relative contribution [27]. Early-layer\ninformation is buried and cannot be selectively retrieved; empirically, a significant fraction of layers can be pruned with\nminimal loss [11]. Recent efforts such as scaled residual paths [54] and multi-stream recurrences [72] remain bound to\nthe additive recurrence, while methods that do introduce cross-layer access [36, 56] are difficult to scale. The situation\nparallels the challenges that recurrent neural networks (RNNs) faced over the sequence dimension before attention\nmechanism provided an alternative.\nWe observe a formal duality between depth-wise accumulation and the sequential recurrence in RNNs. Building\non this duality, we proposeAttention Residuals (AttnRes), which replaces the fixed accumulation hl=P\nivi\nwithhl=P\niαi→l·vi, where αi→laresoftmax attention weights computed from a single learned pseudo-query\nwl∈Rdper layer. This lightweight mechanism enables selective, content-aware retrieval across depth with only one\nd-dimensional vector per layer. Indeed, standard residual connections and prior recurrence-based variants can all be\nshown to perform depth-wiselinearattention; AttnRes generalizes them to depth-wise softmax attention, completing\nfor depth the same linear-to-softmaxtransition that proved transformative over sequences (§6.2, §6.1).\nIn standard training, Full AttnRes adds negligible overhead, since the layer outputs it requires are already retained for\nbackpropagation. At scale, however, activation recomputation and pipeline parallelism are routinely employed, and these\nactivations must now be explicitly preserved and communicated across pipeline stages. We introduceBlock AttnResto\nmaintain efficiency in this regime: layers are partitioned into Nblocks, each reduced to a single representation via\nstandard residuals, with cross-block attention applied only over the Nblock-level summaries. This brings both memory\nand communication down to O(Nd) , and together with infrastructure optimizations (§4), Block AttnRes serves as a\ndrop-in replacement for standard residual connections with marginal training cost and negligible inference latency\noverhead.\nScaling law experiments confirm that AttnRes consistently outperforms the baseline across compute budgets, with\nBlock AttnRes matching the loss of a baseline trained with 1.25× more compute. We further integrate AttnRes into\nthe Kimi Linear architecture [69] (48B total / 3B activated parameters) and pre-train on 1.4T tokens. Analysis of\nthe resulting training dynamics reveals that AttnRes mitigates PreNorm dilution, with output magnitudes remaining\nbounded across depth and gradient norms distributing more uniformly across layers. On downstream benchmarks, our\nfinal model improves over the baseline across all evaluated tasks.\nContributions\n•Attention Residuals.We propose AttnRes, which replaces fixed residual accumulation with learned softmax\nattention over depth, and its scalable variant Block AttnRes that reduces memory and communication from O(Ld) to\nO(Nd) . Through a unified structured-matrix analysis, we show that standard residuals and prior recurrence-based\nvariants correspond to depth-wiselinearattention, while AttnRes performs depth-wisesoftmaxattention.\n•Infrastructure for scale.We develop system optimizations that make Block AttnRes practical and efficient at scale,\nincluding cross-stage caching that eliminates redundant transfers under pipeline parallelism and a two-phase inference\nstrategy that amortizes cross-block attention via online softmax [31]. The resulting training overhead is marginal,\nand the inference latency overhead is less than 2% on typical inference workloads.\n•Comprehensive evaluation and analysis.We validate AttnRes through scaling law experiments, component\nablations, and downstream benchmarks on a 48B-parameter model pre-trained on 1.4T tokens, demonstrating\nconsistent improvements over standard residual connections. Training dynamics analysis further reveals that AttnRes\nmitigates PreNorm dilution, yielding bounded hidden-state magnitudes and more uniform gradient distribution across\ndepth.\n2" - }, - { - "page": 3, - "content": "Attention ResidualsTECHNICALREPORT\n2 Motivation\nNotation.Consider a batch of input sequences with shape B×T×d , where Bis the batch size, Tis the sequence\nlength, and dis the hidden dimension. For clarity, we write formulas for a single token: hl∈Rddenotes the hidden state\nentering layer l, where l∈ {1, . . . , L} is the layer index and Lis the total number of layers. The token embedding is h1.\nThe function flrepresents the transformation applied by layer l. In Transformer models, we treat each self-attention or\nMLP as an individuallayer.\n2.1 Training Deep Networks via Residuals\nResidual Learning.Residual learning [12] proves to be a critical technique in training deep networks as it allows\ngradients to bypass transformations. Specifically, each layer updates the hidden state as:\nhl=hl−1+fl−1(hl−1)\nExpanding this recurrence, the hidden state at layer lis the sum of the embedding and all preceding layer outputs:\nhl=h 1+Pl−1\ni=1fi(hi). The key insight behind residual connections isidentity mapping: each layer preserves a direct\npath for both information and gradients to flow unchanged. During back-propagation, the gradient with respect to an\nintermediate hidden state is:\n∂L\n∂hl=∂L\n∂hL·L−1Y\nj=l\u0012\nI+∂fj\n∂hj\u0013\nExpanding this product yields Iplus higher-order terms involving the layer Jacobians ∂fj/∂hj. The identity term is\nalways preserved, providing a direct gradient path from the loss to any layer regardless of depth.\nGeneralizing Residuals.While effective, the fixed unit coefficients in the residual update treat every layer’s con-\ntribution uniformly, offering no mechanism to adapt the mixing across depth. Highway networks [45] relax this by\nintroducing learned element-wise gates:\nhl= (1−g l)⊙h l−1+gl⊙fl−1(hl−1)\nwhere gl∈[0,1]dinterpolates between the transformation and the identity path. More generally, both are instances\nof a weighted recurrence hl=α l·hl−1+βl·fl−1(hl−1), with residual setting αl=βl=1and Highway setting\nαl=1−g l, βl=gl.\nLimitations.Whether fixed or gated, both approaches share a fundamental constraint: each layer can only access\nits immediate input hl−1, a single compressed state that conflates all earlier layer outputs, rather than the individual\noutputs themselves. This entails several limitations: (1)no selective access: different layer types (e.g., attention vs.\nMLP) receive the same aggregated state, despite potentially benefiting from different weightings; (2)irreversible loss:\ninformation lost through aggregation cannot be selectively recovered in deeper layers; and (3)output growth: later\nlayers learn increasingly larger outputs to gain influence over the accumulated residual, which can destabilize training.\nThese limitations motivate a mechanism that lets each layer selectively aggregate information from all preceding layers.\n3 Attention Residuals: A Unified View of Time and Depth\nThe limitations discussed above are reminiscent of similar bottlenecks in sequence modeling, suggesting that we seek\nsimilar solutions for the depth dimension.\nThe Duality of Time and Depth.Like RNNs over time, residual connections compress all prior information into a\nsingle state hlover depth. For sequence modeling, the Transformer improved upon RNNs by replacing recurrence with\nattention [3, 52], allowing each position to selectively access all previous positions with data-dependent weights. We\npropose the same methodology for depth:\nhl=α 0→l·h1+l−1X\ni=1αi→l·fi(hi)(1)\nwhere αi→lare layer-specific attention weights satisfyingPl−1\ni=0αi→l= 1. Unlike sequence length (which can reach\nmillions of tokens), network depth is typically modest ( L <1000 ), making O(L2)attention over depth computationally\nfeasible. We call this approachAttention Residuals, abbreviated asAttnRes.\n3" - }, - { - "page": 4, - "content": "Attention ResidualsTECHNICALREPORT\n3.1 Full Attention Residuals\nThe attention weights can be written as αi→l=ϕ(q l,ki)for a kernel function ϕ:Rd×Rd→R≥0, where qland\nkiare query and key vectors [23, 70]. Different choices of ϕrecover different residual variants (§6.2); we adopt\nϕ(q,k) = exp\u0000\nq⊤RMSNorm(k)\u0001\n[66] with normalization, yieldingsoftmaxattention over depth:\nαi→l=ϕ(ql,ki)\nPl−1\nj=0ϕ(ql,kj)(2)\nFor each layerl, we define:\nql=w l,k i=vi=\u001ah1 i= 0\nfi(hi) 1≤i≤l−1(3)\nwhere the query ql=w lis a layer-specific learnable vector in Rd. The RMSNorm inside ϕprevents layers with\nlarge-magnitude outputs from dominating the attention weights. The input to layerlis then:\nhl=l−1X\ni=0αi→l·vi (4)\nWe call this formfull attention residuals. For each token, Full AttnRes requires O(L2d)arithmetic and O(Ld) memory\nto store layer outputs. Since depth is far smaller than sequence length, the arithmetic cost is modest.\nOverhead.The O(Ld) memory overlaps entirely with the activations already retained for backpropagation, so Full\nAttnRes introduces no additional memory overhead in vanilla training. At scale, however, activation recomputation and\npipeline parallelism are widely adopted: layer outputs that would otherwise be freed and recomputed must now be kept\nalive for all subsequent layers, and under pipeline parallelism each must further be transmitted across stage boundaries.\nBoth the memory and communication overhead then grow asO(Ld).\nBlockwise optimization.A deliberate design choice in Full AttnRes is that thepseudo-query wlis a learned parameter\ndecoupled from the layer’s forward computation. This independence means that attention weights for any group of\nlayers can be computed in parallel without waiting for their sequential outputs, and in particular permits grouping the L\nlayers into Nblocks of Slayers each and batching the attention computation within each block, reducing per-layer\nmemory I/O from O(Ld) toO((S+N)d) (we defer the detailed two-phase strategy to §4). Under current distributed\ntraining regimes, however, the dominant cost is not local memory bandwidth but cross-stage communication under\npipeline parallelism: every layer output must still be transmitted between stages, and this O(Ld) communication\noverhead cannot be alleviated by local batching. This motivates the Block AttnRes variant introduced below, which\nreduces the number of cross-stage representations from LtoN. We anticipate that future interconnect improvements\nwill make the fullO(Ld)communication practical, fully realizing the potential of Full AttnRes.\n3.2 Block Attention Residuals\nWe proposeBlock Attention Residuals, which partitions the Llayers into Nblocks: within each block, the layer outputs\nare reduced to a single representation via summation, and across blocks, we apply full attention over only Nblock-level\nrepresentations and the token embedding. This reduces both memory and communication overhead from O(Ld) to\nO(Nd).\nIntra-Block Accumulation.Specifically, we divide the Llayers into Nblocks of S=L/N layers each, assuming\nLis divisible by N; otherwise, the last block contains the remaining LmodN layers. Let Bndenote the set of layer\nindices in blockn(n= 1, . . . , N). To form a block, we sum all of its layer outputs:\nbn=X\nj∈Bnfj(hj)(5)\nWe further denote bi\nnas the partial sum over the first ilayers in Bn, so that bn=bS\nn. When Lis not divisible by N,\nthe final partial sum is taken as the last block’s representation. As in Full AttnRes, the RMSNorm inside ϕprevents\nmagnitude differences between complete blocks and partial sums from biasing the attention weights.\n4" - }, - { - "page": 5, - "content": "Attention ResidualsTECHNICALREPORT\n1 def block_attn_res(blocks: list[Tensor], partial_block: Tensor, proj: Linear, norm: RMSNorm) -> Tensor:\n2 \"\"\"\n3 Inter-block attention: attend over block reps + partial sum.\n4 blocks:\n5 N tensors of shape [B, T, D]: completed block representations for each previous block\n6 partial_block:\n7 [B, T, D]: intra-block partial sum (b_n^i)\n8 \"\"\"\n9 V = torch.stack(blocks + [partial_block]) # [N+1, B, T, D]\n10 K = norm(V)\n11 logits = torch.einsum('d, n b t d -> n b t', proj.weight.squeeze(), K)\n12 h = torch.einsum('n b t, n b t d -> b t d', logits.softmax(0), V)\n13 return h\n14\n15 def forward(self, blocks: list[Tensor], hidden_states: Tensor) -> tuple[list[Tensor], Tensor]:\n16 partial_block = hidden_states\n17 # apply block attnres before attn\n18 # blocks already include token embedding\n19 h = block_attn_res(blocks, partial_block, self.attn_res_proj, self.attn_res_norm)\n20\n21 # if reaches block boundary, start new block\n22 # block_size counts ATTN + MLP; each transformer layer has 2\n23 if self.layer_number % (self.block_size // 2) == 0:\n24 blocks.append(partial_block)\n25 partial_block = None\n26\n27 # self-attention layer\n28 attn_out = self.attn(self.attn_norm(h))\n29 partial_block = partial_block + attn_out if partial_block is not None else attn_out\n30\n31 # apply block attnres before MLP\n32 h = block_attn_res(blocks, partial_block, self.mlp_res_proj, self.mlp_res_norm)\n33\n34 # MLP layer\n35 mlp_out = self.mlp(self.mlp_norm(h))\n36 partial_block = partial_block + mlp_out\n37\n38 return blocks, partial_block\nFigure 2: PyTorch-style pseudo code for Block Attention Residuals. block_attn_res computes softmax attention over block\nrepresentations using a learned pseudo-query wl;forward is a single-layer pass that maintains partial_block (bi\nn, intra-block\nresidual) andblocks([b 0, . . . ,b n−1], inter-block history).\nInter-Block Attention.In Full AttnRes, the input to layer lis computed by attending over all outputs up to fl−1(hl−1).\nThe block-wise variant replaces these individual outputs with block representations, defining b0=h 1so that the token\nembedding is always included as a source. For thei-th layer in blockn, the value matrix is:\nV=\u001a[b0,b1, . . . ,b n−1]⊤ifi= 1(first layer of blockn)\n[b0,b1, . . . ,b n−1,bi−1\nn]⊤ifi≥2(subsequent layers)(6)\nKeys and attention weights follow Eq. 3 and Eq. 2. The input of the very first layer of the network is the token\nembeddings, i.e. b0=h 1. In each block, the first layer receives the previous block representations and the token\nembeddings, and the subsequent layers additionally attend to the partial sum bi−1\nn. The final output layer aggregates all\nNblock representations. Fig. 2 provides PyTorch-style pseudocode for Block AttnRes.\nEfficiency.Since each layer now attends over Nblock representations rather than Lindividual outputs, memory\nreduces from O(L) toO(N) and computation from O(L2)toO(N2). The block count Ninterpolates between two\nextremes: N=L recovers Full AttnRes, while N= 1 reduces to standard residual connections with the embedding\nisolated as b0. Empirically, we find that N≈8 recovers most of the benefit across model scales, requiring only eight\nstored hidden states per token (see § 5).\nBeyond memory and computation, the block structure also benefits inference latency: block boundaries define the\ndispatch granularity for the blockwise optimization described in §3, and the fixed block count Nbounds the KV cache\nsize. The parallel inter-block results are merged with the sequential intra-block partial sums via online softmax [31],\npreserving exact equivalence (§4).\n4 Infrastructure Design\nBlock AttnRes introduces additional system challenges compared to standard residual connections. For large-scale\nmodel training, block representations must be propagated across pipeline stages, causing heavy communication in a\n5" - }, - { - "page": 6, - "content": "Attention ResidualsTECHNICALREPORT\nRANK0\nRANK1\nRANK2\nRANK3[b0] [ ]\n[b0] [b1]\n[b0,b1] [ ]\n[b0,b1] [b2]+ [b 1,b2][ ]\n+ [b 1,b2][b3]\n+ [b 2,b3][ ]\n+ [b 2,b3][b4]VIRTUALSTAGE0 VIRTUALSTAGE1\n1 2\n1 2\n1 2\n1 21 2\n1 2\n1 2\n1 2\nFigure 3: Cache-based pipeline communication example with 4 physical ranks and 2 virtual stages per rank, where hatched boxes\ndenote end of AttnRes blocks. Numbers indicate micro-batch indices. Each rank caches previously received blocks; stage transitions\nonly transmit incremental blocks (+[b 1,b2]) instead of the full history.\nnaïve implementation. During inference, repeated access to accumulated block representations increases latency, while\nlong-context prefilling amplifies the memory cost of caching block representations. We address these challenges with\ncross-stage caching in training, and with a two-phase computation strategy together with a memory-efficient prefilling\nscheme in inference.\n4.1 Training\nFor small-scale training, AttnRes adds a tiny computation overhead and no extra memory usage, as the activations\nneed to be saved for backpropagation regardless. Under large-scale distributed training, pipeline parallelism poses the\nprimary infrastructure challenge for AttnRes. Full AttnRes requires all Llayer outputs to be transmitted across stages;\nBlock AttnRes reduces this to Nblock representations, and the optimizations below further minimize the remaining\noverhead.\nPipeline communication.With standard residual connections, pipeline parallelism [18] transfers a fixed-size hidden\nstate between adjacent stages, independent of pipeline depth. Block AttnRes requires all accumulated block representa-\ntions at each stage for inter-block attention, and naïvely transmitting the full history at every transition incurs redundant\ncommunication.\nConsider an interleaved pipeline schedule [33] with Pphysical stages and Vvirtual stages per physical stage. For\nsimplicity, assume each physical stage produces on average Npblock representations of dimension dper token.1With\nC=PV total chunks (each physical stage in each virtual stage), the j-th chunk accumulates jNpblocks. Naïvely\ntransmitting all accumulated blocks at every transition incurs per-token communication cost:\nComm naïve=C−1X\nj=1jNp·d=C(C−1)\n2Npd.(7)\nCross-stage caching.Since each physical stage processes multiple virtual stages in succession, we can eliminate\nthis redundancy by caching blocks locally: blocks received during earlier virtual stages remain in local memory and\nneed not be re-transmitted. The first virtual stage ( v= 1 ) has no cache and accumulates normally; for v≥2 , each\ntransition conveys only the ∼PN pincremental blocks accumulated since the receiver’s corresponding chunk in the\nprevious virtual stage. Total communication reduces to:\nComm cached =P(P−1)\n2Npd\n|{z}\nfirst virtual stage+ (V−1)P2Npd|{z }\nsubsequent virtual stages.(8)\nCaching reduces peak per-transition cost from O(C) toO(P) , aV× improvement that enables full overlap with\ncomputation during steady-state 1F1B. The backward pass benefits from the same scheme. Fig. 3 illustrates this\noptimization withP=4andV=2: for the second virtual stage, caching eliminates 6 redundant block transmissions.\n1In practice, block boundaries need not align with physical stage boundaries. For example, in Fig. 3, each block spans two\nphysical stages, so only every other transition involves a newly completed block.\n6" - }, - { - "page": 7, - "content": "Attention ResidualsTECHNICALREPORT\nAlgorithm 1:Two-phase computation for blockn\nInput:Pseudo queries{w l}l∈Bn, block representations{b 0, . . . ,b n−1}\n/* Phase 1: Parallel inter-block attention */\n1Q←[w l]l∈Bn //[S, d]\n2K,V←[b 0;. . .;b n−1]//[n, d]\n3{o(1)\nl, m(1)\nl, ℓ(1)\nl}l∈Bn←ATTNWITHSTATS(Q,K,V)// Return LSE\n4\n/* Phase 2: Sequential intra-block attention + Onlinesoftmaxmerge */\n5i←0\n6forl∈ B ndo\n7ifi= 0then\n8h l←o(1)\nl/ℓ(1)\nl// Inter-block only\n9else\n10o(2)\nl, m(2)\nl, ℓ(2)\nl←ATTNWITHSTATS(w l,bi\nn,bi\nn)// Intra-block\n11m l←max(m(1)\nl, m(2)\nl)\n12h l←em(1)\nl−mlo(1)\nl+em(2)\nl−mlo(2)\nl\nem(1)\nl−mlℓ(1)\nl+em(2)\nl−mlℓ(2)\nl// Online softmax merge\n13i←i+ 1\n14bi\nn←bi−1\nn+fl(hl)// Update partial sum;b0\nn:=0\n15return{h l}l∈Bn\nMemory overhead.With cross-stage caching, each block is stored exactly once across all Vvirtual stages, which\nbecomes negligible relative to standard per-layer activation cache. Crucially, the per-layer activation footprint remains\nidentical to standard architectures, as activation checkpointing eliminates all inter-block attention intermediates, and the\ncheckpointed inputp lmatches the memory size of the hidden stateh lit replaces.\nIn terms of wall-clock time, Block AttnRes adds negligible training overhead when pipeline parallelism is not enabled;\nunder pipeline parallelism, the measured end-to-end overhead is less than 4%.\n4.2 Inference\nThe two-phase computation strategy described below applies to both Full and Block AttnRes: in either case, layers are\ngrouped into blocks of size S, with Phase 1 batching the inter-block queries and Phase 2 handling sequential intra-block\nlookback. For Full AttnRes, this reduces per-layer I/O from O(Ld) toO((S+N)d) (detailed derivation shown in\nAppendix B); Block AttnRes further reduces the stored representations from LtoN, since each block is compressed\ninto a single vector. In what follows, we focus on Block AttnRes and detail the two-phase computation strategy together\nwith a sequence-sharded prefilling scheme for long-context inputs.\nTwo-phase computation strategy.The layer-wise attention computation of Block AttnRes resembles autoregressive\ndecoding, where block representations serve as a shared KV cache reused across layers. A naïve implementation\ncomputes the attention residual at every layer, each requiring a full pass over all preceding blocks, resulting in O(L·N)\nmemory accesses. Since the pseudo-query vectors are decoupled from the forward computation (§3), all S=L/N\nqueries within a block can be batched into a single matrix multiplication, amortizing memory access from Sreads to 1.\nAlgorithm 1 instantiates a two-phase computation strategy exploiting this property.\n•Phase 1computes inter-block attention for all Slayers simultaneously via a single batched query against the cached\nblock representations, returning both outputs and softmax statistics (max and log-sum-exp). This amortizes the\nmemory access cost, reducing reads fromStimes to just once per block.\n•Phase 2computes intra-block attention sequentially for each layer using the evolving partial sum, then merges with\nPhase 1 outputs through online softmax [31]. Because the online- softmax merge is elementwise, this phase naturally\nadmits kernel fusion with surrounding operations, further reducing I/O overhead.\nWith the two-phase design, Phase 2 preserves an I/O footprint similar to that of standard residual connections, whereas\nthe main additional cost arises from Phase 1 inter-block attention. Because these inter-block reads are amortized across\n7" - }, - { - "page": 8, - "content": "Attention ResidualsTECHNICALREPORT\nall layers in a block through batching, the total per-layer memory access cost remains only (N\nS+ 3)d reads and 2d\nwrites (Table 1). This is substantially lower than the residual-stream I/O of prior residual generalizations such as (m)HC\nunder typical settings. In practice, Phase 1 can also partially overlap with the computation of the first layer in the block,\nfurther reducing its wall-clock impact. As a result, the end-to-end inference latency overhead is less than 2% on typical\ninference workloads.\nTable 1: Memory access cost per token per layer incurred by the residual mechanism under each scheme. The internal I/O of the layer\nfunction flis excluded. For AttnRes, both Full and Block variants use the two-phase inference schedule described in Appendix B;\namortized costs are averaged overNlayers within a block. Typical values:L=128,N=8,S=L/N=16,m=4.\nOperation Read WriteTotal I/O\nSymbolic Typical\nStandard Residuals Residual Merge2d d3d3d\nmHC (mstreams)Computeα l,βl,Al md m2+2m\n(8m+2)d+2m2+4m 34dApplyα l md+m d\nApplyβ l d+m md\nApplyA l md+m2md\nResidual Merge2md md\nAttnResFullPhase 1 (amortized)(N−1)d d(S+N)d24dPhase 2(S−1)d d\nBlockPhase 1 (amortized)N\nSd d \u0000N\nS+5\u0001\nd 5.5dPhase 23d d\nMemory-efficient prefilling.Storing block representations during prefilling requires N·T·d elements, which incurs\n15 GB of memory for a 128K-token sequence with 8 blocks. We mitigate this by sharding these representations along\nthe sequence dimension across Ptensor-parallel devices, allowing Phase 1 to execute independently on local sequence\nshards. The Phase 2 online- softmax merge then integrates into the standard TP all-reduce communication path: the\noutput is reduce-scattered, merged locally, and reconstructed via all-gather, naturally admitting kernel fusion with\noperations like RMSNorm . This reduces the per-device memory footprint to N·(T/P)·d —lowering the 128K-context\nexample from 15 GB to roughly 1.9 GB per device. Combined with chunked prefill (e.g., 16K chunk size), the overhead\nfurther reduces to under 0.3 GB per device.\n5 Experiments\nArchitecture Details.Our architecture is identical to Kimi Linear [69], a Mixture-of-Experts (MoE) Transformer\nfollowing the Moonlight [28] / DeepSeek-V3 [9] design, which interleaves Kimi Delta Attention (KDA) and Multi-Head\nLatent Attention (MLA) layers in a 3:1 ratio, each followed by an MoE feed-forward layer. The only modification is the\naddition of AttnRes to the residual connections; all other components (model depth, hidden dimensions, expert routing,\nand MLP structure) remain unchanged. AttnRes introduces only one RMSNorm and one pseudo-query vector wl∈Rd\nper layer, amounting to a negligible fraction of the total parameter count. Crucially, all pseudo-query vectors must be\ninitialized to zero. This ensures that the initial attention weights αi→lare uniform across source layers, which reduces\nAttnRes to an equal-weight average at the start of training and prevents training volatility, as we validated empirically.\n5.1 Scaling Laws\nWe sweep five model sizes (Table 2) and train three variants per size: a PreNorm baseline, Full AttnRes, and Block\nAttnRes with ≈8blocks. They are trained with an 8192-token context window and a cosine learning rate schedule.\nWithin each scaling law size group, all variants share identical hyperparameters selected under the baseline to ensure\nfair comparison; this setup intentionally favors the baseline and thus makes the comparison conservative. Following\nstandard practice, we fit power-law curves of the form L=A×C−α[22, 15], where Lis validation loss and Cis\ncompute measured in PFLOP/s-days.\nScaling Behavior.Fig. 4 presents the fitted scaling curves. The Baseline follows L= 1.891×C−0.057, while Block\nAttnRes fits L= 1.870×C−0.058, and Full AttnRes fits L= 1.865×C−0.057. All three variants exhibit a similar\nslope, but AttnRes consistently achieves lower loss across the entire compute range. Based on the fitted curves, at 5.6\n8" - }, - { - "page": 9, - "content": "Attention ResidualsTECHNICALREPORT\nTable 2: Baseline vs Block AttnRes ( N= 8 ) vs Full AttnRes vs mHC(-lite) [64]: Model configurations, Hyperparameters, and\nValidation Loss.\n# Act.\nParams†TokensL bH d model dff lr batch size‡ Val. Loss\nBaseline Block AttnRes Full AttnRes mHC(-lite)\n194M 038.7B 12 12 0896 4002.99×10−3192 1.931 1.9091.8991.906\n241M 045.4B 13 13 0960 4322.80×10−3256 1.895 1.875 1.8741.869\n296M 062.1B 14 14 1024 4642.50×10−3320 1.829 1.8091.8041.807\n436M 087.9B 16 16 1168 5282.20×10−3384 1.766 1.7461.7371.747\n528M 119.0B 17 17 1264 5602.02×10−3432 1.719 1.6931.6921.694\n†Denotes the number of activated parameters in our MoE models, excluding embeddings.\n‡All models were trained with a context length of 8192.\n⋆Lb=L/2denotes the number of Transformer blocks.\n0.5 1 2 51.71.81.9\n1.25×\nPFLOP/s-daysLossBaseline:1.891×C−0.057\nFull AttnRes:1.865×C−0.057\nBlock AttnRes:1.870×C−0.058\nFigure 4: Scaling law curves for Attention Residuals. Both Full and Block AttnRes consistently outperform the baseline across all\nscales. Block AttnRes closely tracks Full AttnRes, recovering most of the gain at the largest scale.\nPFLOP/s-days, Block AttnRes reaches 1.692 versus the Baseline’s 1.714, equivalent to a 1.25× compute advantage.\nThe gap between Full and Block AttnRes narrows with scale, shrinking to just 0.001 at the largest size. We also list\nmHC(-lite) [64] in Table 2 for reference. Full AttnRes outperforms mHC, while Block AttnRes matches it at lower\nmemory I/O per layer:5.5dversus34dfor mHC withm=4streams (Table 1).\n5.2 Main Results\nTraining recipe.The largest models we study are based on the full Kimi Linear 48B configuration: 27 Transformer\nblocks (54 layers) with 8 out of 256 routed experts plus 1 shared expert, yielding 48B total and 3B activated parameters.\nThis model applies Block AttnRes with 6 layers per block, producing 9 blocks plus the token embedding for a total of\n10 depth-wise sources.\nWe follow the same data and training recipe as the Kimi Linear 1.4T-token runs [69]: all models are pre-trained with a\n4096-token context window, the Muon optimizer [28], and a WSD (Warmup–Stable–Decay) learning rate schedule [16],\nwith a global batch size of 8M tokens. Training of the final model proceeds in two stages: (i) a WSD pre-training phase\non 1T tokens, followed by (ii) a mid-training phase on ≈400B high-quality tokens, following the annealing recipe of\nMoonlight [28].\nAfter mid-training, we continue training with progressively longer sequence length of 32K tokens. Since our architecture\nuses hybrid KDA/MLA attention [69], where MLA operates without positional encodings (NoPE) [61], context extension\nrequires no modifications such as YaRN [37] or attention temperature rescaling.\n9" - }, - { - "page": 10, - "content": "Attention ResidualsTECHNICALREPORT\n20k 40k 60k 80k 100k1.21.31.41.5\nStep(a) Validation Loss\nBaseline\nBlock AttnRes\n0 10 20051015\nTransformer Block Index(b) Output Magnitude\n0 10 200123\nTransformer Block Index(c) Gradient Magnitude (×10−5)\nFigure 5: Training dynamics of Baseline and Block AttnRes.(a)Validation loss during training.(b)Each transformer block’s output\nmagnitude at the end of training.(c)Each transformer block’s gradient magnitude.\nTraining dynamics.We compare the training dynamics of our final Baseline and Block AttnRes models over 1T\ntokens in Fig. 5.\n•Validation loss:AttnRes achieves consistently lower validation loss throughout training, with the gap widening\nduring the decay phase and resulting in a notably lower final loss.\n•Output magnitude:The Baseline suffers from the PreNorm dilution problem [60, 27]: as hidden-state magnitudes\ngrow monotonically with depth, deeper layers are compelled to learn increasingly large outputs from fixed-scale\nnormalized inputs to remain influential. Block AttnRes confines this growth within each block, as selective aggregation\nat block boundaries resets the accumulation, yielding a bounded periodic pattern.\n•Gradient magnitude:With all residual weights fixed to 1, the Baseline provides no means of regulating gradient\nflow across depth, leading to disproportionately large gradients in the earliest layers. The learnable softmax weights\nin Block AttnRes (Fig. 8) introduce competition among sources for probability mass, resulting in a substantially more\nuniform gradient distribution.\nTable 3: Performance comparison of AttnRes with the baseline, both after the same pre-training recipe. Best per-row results are\nbolded.\nBaseline AttnRes\nGeneralMMLU 73.574.6\nMMLU-Pro52.2 52.2\nGPQA-Diamond 36.944.4\nBBH 76.378.0\nARC-Challenge 64.665.7\nHellaSwag 83.283.4\nTriviaQA 69.971.8\nMath & CodeGSM8K 81.782.4\nMGSM 64.966.1\nMath 53.557.1\nCMath 84.785.1\nHumanEval 59.162.2\nMBPP 72.073.9\nChineseCMMLU 82.082.9\nC-Eval 79.682.5\nDownstream performance.Following the evaluation protocol of Kimi Linear [69], we assess both models across\nthree areas (Table 3):\n10" - }, - { - "page": 11, - "content": "Attention ResidualsTECHNICALREPORT\nTable 4: Ablation on key components of AttnRes (16-layer\nmodel).\nVariant Loss\nBaseline (PreNorm) 1.766\nDenseFormer [36] 1.767\nmHC [59] 1.747\nAttnRes Full 1.737\nw/ input-dependent query1.731\nw/ input-independent mixing1.749\nw/sigmoid1.741\nw/oRMSNorm1.743\nSWA (W= 1 + 8) 1.764\nBlock (S= 4) 1.746\nw/ multihead (H= 16)1.752\nw/oRMSNorm1.75032 16 8 4 21.7351.7401.7451.7501.7551.7601.7651.770\n1.757\n1.753\n1.748\n1.746 1.746Baseline (1.766)\nFull AttnRes i.e. S=1 (1.737)\nBlock size (S)Validation lossBaseline\nFull AttnRes\nBlock AttnRes\nFigure 6: Effect of block size on validation loss (16-layer model).\n•Language understanding and reasoning: MMLU [13], MMLU-Pro Hard [55], GPQA-Diamond [41], BBH [48],\nARC-Challenge [6], HellaSwag [65], and TriviaQA [21].\n•Reasoning (Code and Math): GSM8K [7], MGSM [44], Math [25], CMath [14], HumanEval [5], and MBPP [1].\n•Chinese language understanding: CMMLU [26] and C-Eval [19].\nAs shown in Table 3, Block AttnRes matches or outperforms the baseline on all benchmarks. The improvements are\nparticularly pronounced on multi-step reasoning tasks such as GPQA-Diamond (+7.5) and Minerva Math (+3.6), as\nwell as code generation such as HumanEval (+3.1), while knowledge-oriented benchmarks such as MMLU (+1.1)\nand TriviaQA (+1.9) also show solid gains. This pattern is consistent with the hypothesis that improved depth-wise\ninformation flow benefits compositional tasks, where later layers can selectively retrieve and build upon earlier\nrepresentations.\n5.3 Ablation Study\nWe conduct ablation studies on the 16-head model from Table 2 to validate key design choices in AttnRes (Table 4). All\nmodels share identical hyperparameters and compute budget.\nComparison with prior methods.We compare AttnRes against the PreNorm baseline (loss 1.766) and two rep-\nresentative methods that generalize residual connections. DenseFormer [36] grants each layer access to all previous\noutputs but combines them with fixed, input-independent scalar coefficients; it shows no gain over the baseline (1.767),\nhighlighting the importance of input-dependent weighting. mHC [59] introduces input dependence through mparallel\nstreams with learned mixing matrices, improving to 1.747. AttnRes takes this further with explicit content-dependent\nselection via softmax attention: Full AttnRes achieves 1.737 and Block AttnRes 1.746, outperforming both methods\nwith only a single query vector per layer.\nCross-layer access.We compare three granularities of cross-layer access. Full AttnRes follows directly from the\ntime–depth duality (§ 3), applying attention over all previous layers, and achieves the lowest loss (1.737). A simple\nway to reduce its memory cost is sliding-window aggregation (SWA), which retains only the most recent W=8 layer\noutputs plus the token embedding; it improves over baseline (1.764) but falls well short of both Full and Block AttnRes,\nsuggesting that selectively accessing distant layers matters more than attending to many nearby ones.\nBlock AttnRes offers a better trade-off: with block size S=4 it reaches 1.746 while keeping memory overhead constant\nper layer. Fig. 6 sweeps Sacross the full spectrum from S=1 (i.e. Full AttnRes) to increasingly coarse groupings. Loss\ndegrades gracefully as Sgrows, with S=2,4,8 all landing near 1.746 while larger blocks ( S=16,32 ) move toward\nbaseline. In practice, we fix the number of blocks to ≈8for infrastructure efficiency (§ 4). As future hardware alleviates\nmemory capacity constraints, adopting finer-grained block sizes or Full AttnRes represents a natural pathway to further\nimprove performance.\n11" - }, - { - "page": 12, - "content": "Attention ResidualsTECHNICALREPORT\n15 30 45 60 750.30.40.50.60.7 2.017 1.909 1.875 1.851 1.858\n1.990 1.902 1.862 1.852 1.862\n1.973 1.883 1.859 1.849 1.854\n1.952 1.868 1.850 1.849 1.857\n1.926 1.857 1.851 1.858 1.847\ndmodel/LbH/L b\n(a) Baseline15 30 45 60 751.954 1.890 1.843 1.828 1.824\n1.931 1.863 1.830 1.817 1.818\n1.917 1.841 1.819 1.812 1.817\n1.893 1.823 1.815 1.813 1.813\n1.877 1.816 1.820 1.806 1.802\ndmodel/Lb\n1.841.881.921.962\n(b) Attention Residuals\nFigure 7: Architecture sweep under fixed compute ( ≈6.5×1019FLOPs, ≈2.3×108active parameters). Each cell reports\nvalidation loss for a (dmodel/Lb, H/L b)configuration, where Lb=L/2 is the number of Transformer blocks; the star marks the\noptimum.\nComponent design.We further ablate individual components of the attention mechanism:\n•Input-dependent query.A natural extension is to make the query input-dependent by projecting it from the current\nhidden state. This further lowers loss to 1.731, but introduces a d×d projection per layer and requires sequential\nmemory access during decoding, so we default to the learned query.\n•Input-independent mixing.We removed the query and key and replaced them with learnable, input-independent\nscalars to weigh previous layers, which hurts performance (1.749 vs. 1.737).\n•softmax vs.sigmoid .Replacing softmax withsigmoid degrades performance (1.741). We attribute this to softmax ’s\ncompetitive normalization, which forces sharper selection among sources.\n•Multihead attention.We test per-head depth aggregation ( H=16 ) on Block AttnRes, allowing different channel\ngroups to attend to different source layers. This hurts performance (1.752 vs. 1.746), indicating that the optimal\ndepth-wise mixture is largely uniform across channels: when a layer’s output is relevant, it is relevant as a whole.\n•RMSNorm on keys.Removing RMSNorm degrades both Full AttnRes (1.743) and Block AttnRes (1.750). For\nFull AttnRes, it prevents individual layers with naturally larger outputs from dominating the softmax . This becomes\neven more critical for Block AttnRes, as block-level representations accumulate over more layers and can develop\nlarge magnitude differences;RMSNormprevents these from biasing the attention weights.\n5.4 Analysis\n5.4.1 Optimal Architecture\nTo understand how AttnRes reshapes optimal architectural scaling, we perform a controlled capacity reallocation\nstudy under a fixed compute and parameter budget. Our central question is whether AttnRes alters the preferred\ndepth–width–attention trade-off, and in particular, given its potential strength on the depth dimension, whether it favors\ndeeper models compared to conventional Transformer design heuristics. To isolate structural factors directly coupled\nto depth, we fix the per-expert MLP expansion ratio based on internal empirical observations ( dff/dmodel≈0.45 ).\nWe further fix total training compute (FLOPs ≈6.5×1019) and active parameters ( ≈2.3×108), ensuring that any\nperformance variation arises purely from architectural reallocation rather than overall capacity differences. Under\nthis constrained budget, we enumerate 25 configurations on a 5×5 grid over dmodel/Lb∈ {15,30,45,60,75} and\nH/L b∈ {0.3,0.4,0.5,0.6,0.7} , where Lb=L/2 is the number of Transformer blocks and Hthe number of attention\nheads. The results are shown in Fig. 7.\nBoth heatmaps exhibit a shared pattern: loss decreases with growing dmodel/Lband shrinking H/L b, and both methods\nreach their optima at H/L b≈0.3 . Despite this shared trend, AttnRes achieves a lower loss than the baseline in each of\nthe 25 configurations, by 0.019 –0.063 . The most apparent difference lies in the location of the optimum: the baseline\nachieves its lowest loss at dmodel/Lb≈60 (1.847 ), whereas AttnRes shifts it to dmodel/Lb≈45 (1.802 ). Under a fixed\n12" - }, - { - "page": 13, - "content": "Attention ResidualsTECHNICALREPORT\n0 5 10 15 20 25 301\n2\n3\n4\n5\n6\n7\n8\n9\n10\n11\n12\n13\n14\n15\n16\nSource IndexLayerFull AttnRes, Pre-Attn\n0 5 10 15 20 25 30\nSource IndexFull AttnRes, Pre-MLP\n0 1 2 3 4 5 6 7 81\n2\n3\n4\n5\n6\n7\n8\n9\n10\n11\n12\n13\n14\n15\n16\nBlock IndexLayerBlock AttnRes, Pre-Attn\n0 1 2 3 4 5 6 7 8\nBlock IndexBlock AttnRes, Pre-MLP\n00.20.40.60.8Weight\nFigure 8: Depth-wise attention weight distributions for a 16-head model with full (top) and block (bottom) Attention Residuals,\naveraged over tokens. The model has 16 attention and 16 MLP layers. Each row shows how the lth attention (left) or MLP (right)\nlayer distributes weight over previous sources. Diagonal dominance indicates locality remains the primary information pathway,\nwhile persistent weights on source 0 (embedding) and occasional off-diagonal concentrations reveal learned skip connections. Block\nattention (N= 8) recovers the essential structure with sharper, more decisive weight distributions.\nparameter budget, a lower dmodel/Lbcorresponds to a deeper, narrower network, suggesting that AttnRes can exploit\nadditional depth more effectively. We note that this preference for depth does not directly translate to a deployment\nrecommendation, as deeper models generally incur higher inference latency due to their sequential computation [39].\nRather, this sweep serves as a diagnostic that reveals where AttnRes benefits most, and this depth preference can be\nfactored into the architecture selection alongside inference cost.\n5.4.2 Analyzing Learned AttnRes Patterns\nWe visualize the learned weights αi→lin Fig. 8 for the 16-head model (from Table 2) with both full and block ( N=8 )\nAttnRes. Each heatmap shows how the lth attention or MLP layer (rows) allocates its attention over previous sources\n(columns), with pre-attention and pre-MLP layers shown separately. We highlight three key observations:\n•Preserved locality.Each layer attends most strongly to its immediate predecessor, yet selective off-diagonal\nconcentrations emerge (e.g., layer 4 attending to early sources, layers 15–16 reaching back under the block setting),\nindicating learned skip connections beyond the standard residual path.\n•Layer specialization.The embedding h1retains non-trivial weight throughout, especially in pre-attention layers.\nPre-MLP inputs show sharper diagonal reliance on recent representations, while pre-attention inputs maintain broader\nreceptive fields, consistent with attention routing information across layers and MLPs operating locally.\n•Block AttnRes preserves structure.Diagonal dominance, embedding persistence, and layer specialization all\ntransfer from the full to the block variant, suggesting that block-wise compression acts as implicit regularization\nwhile preserving the essential information pathways.\n13" - }, - { - "page": 14, - "content": "Attention ResidualsTECHNICALREPORT\nTable 5: Comparison of residual update mechanisms.Weight: whether the mixing coefficients are architecture-fixed, learned-static\n(fixed after training), or input-dependent (dynamic).Source: which earlier representations layer lcan access. Normalization is\nomitted from most formulas for clarity.\nMethod Update rule Weight Source\nSingle-state recurrence: layerlreceives onlyh l−1\nResidual [12]h l=hl−1+fl−1(hl−1)Fixedh l−1\nReZero [2]h l=hl−1+αl·fl−1(hl−1)Statich l−1\nLayerScale [50]h l=hl−1+ diag(λ l)·fl−1(hl−1)Statich l−1\nHighway [45]h l= (1−g l)⊙h l−1+gl⊙fl−1(hl−1)Dynamich l−1\nDeepNorm [54]h l= Norm(αh l−1+fl−1(hl−1))Fixedh l−1\nKEEL [4]h l= Norm(αh l−1+fl−1(Norm(h l−1)))Fixedh l−1\nMulti-state recurrence: layerlreceivesmstreams\nSiameseNorm [27]h1\nl=Norm(h1\nl−1+yl−1);h2\nl=h2\nl−1+yl−1 Fixed 2 streams\nHC/mHC [72, 59]H l=H l−1Al+fl−1(Hl−1αl−1)β⊤\nl−1 Dynamicmstreams\nDDL [67]H l= (I−β lklk⊤\nl)Hl−1+βlklv⊤\nl Dynamicd vstreams\nCross-layer access: layerlcan access individual earlier-layer outputs\nDenseNet [17]h l= ConvPool([h 1;f1(h1);. . .;f l−1(hl−1)])Static[h 1, . . . ,h l−1]\nDenseFormer [36]h l=α 0→lh1+Pl−1\ni=1αi→lfi(hi)Static[h 1, . . . ,h l−1]\nMRLA [10]1hl=Pl−1\ni=1σ\u0000\nConvPool(f l−1(hl−1))\u0001⊤σ\u0000\nConvPool(f i(hi))\u0001\nConv(f i(hi))Dynamic[h 1, . . . ,h l−1]\nFull2hl∝Pl−1\ni=0ϕ(w l,ki)vi Dynamic [h1, . . . ,h l−1]AttnRes (ours)Block3hl∝Pn−1\ni=0ϕ(w l,ki)vi+ϕ(w l,kj\nn)vj\nn Dynamic [b0, . . . ,b n−1,bj\nn]\n1ConvPool: pooling operation followed by convolution (channel projection).\n2ϕ(q,k) = exp\u0000\nq⊤RMSNorm(k)\u0001\n;ki=vi;v0=h 1,vi≥1=fi(hi).softmaxjointly normalized over all sources.\n3Sameϕand normalization as Full;v i=bi,vj\nn=bj\nn.\n6 Discussions\n6.1 Sequence-Depth Duality\nResidual connections propagate information over depth via a fixed recurrence hl=hl−1+fl−1(hl−1), much as RNNs\npropagate information over time. Test-Time Training (TTT) [46] formalizes the sequence side of this analogy (cf. Fast\nWeight Programmers [43, 32]), casting each recurrent step as gradient descent on a self-supervised loss:\nWt=W t−1−η∇ℓ(W t−1;xt),(9)\nwhere a slow network parameterizes ℓand the state Wis updated once per token. When fis linear, this reduces to\nvanilla linear attention St=St−1+ktv⊤\nt. The standard residual exhibits the same additive form along depth, with hl\nserving as the state and each layerf lacting as one “gradient step.”\nAs noted by [4], this duality extends to richer variants (Table 5). Data-dependent gates on the sequence side [47, 63]\ncorrespond to Highway networks [45] on the depth side; the delta rule [42, 62, 69] corresponds to DDL [67]; and\nMRLA [10] mirrors GLA’s [63] gated linear attention. These methods all refine the recurrent update while remaining\nwithin the recurrence paradigm. AttnRes goes a step further and replaces depth-wise recurrence with direct cross-layer\nattention, just as Transformers replaced temporal recurrence with self-attention. Since the number of layers in current\narchitectures remains well within the practical regime of softmax attention, we adopt vanilla depth-wise attention.\nIncorporating more expressive yet memory-efficient (e.g. linear-complexity) alternatives is a natural direction for future\nwork.\n6.2 Residual Connections as Structured Matrices\nThe residual variants discussed above can all be viewed as weighted aggregations over previous layer outputs. We\nformalize this with adepth mixing matrix M∈RL×L, where Mi→lis the weight that layer lassigns to the output of\nlayer i. The variants differ in how these weights arise (fixed, learned, or input-dependent) and whether Mis constrained\nto low rank or allowed to be dense. The semiseparable rank ofM[8] offers a unified lens for comparing them.\nConcretely, the input to layer lishl=Pl−1\ni=0Mi→lvi, where v0=h 1(embedding) and vi=fi(hi)fori≥1 . Fig. 9\nvisualizesMfor representative methods; we derive each below.\n14" - }, - { - "page": 15, - "content": "Attention ResidualsTECHNICALREPORT\nHighway\n\n1\nγ×\n1→2g2\nγ×\n1→3g2γ×\n2→3g3\nγ×\n1→4g2γ×\n2→4g3γ×\n3→4g4\n(m)HC\n\nβ⊤\n0α1\nβ⊤\n0A×\n1→2α2 β⊤\n1α2\nβ⊤\n0A×\n1→3α3β⊤\n1A×\n2→3α3 β⊤\n2α3\nβ⊤\n0A×\n1→4α4β⊤\n1A×\n2→4α4β⊤\n2A×\n3→4α4 β⊤\n3α4\n\nFull AttnRes\n\nϕ(w 1,k0)\nϕ(w 2,k0) ϕ(w 2,k1)\nϕ(w 3,k0) ϕ(w 3,k1) ϕ(w 3,k2)\nϕ(w 4,k0) ϕ(w 4,k1) ϕ(w 4,k2) ϕ(w 4,k3)\nBlock AttnRes\n\nϕ(w 1,k0)\nϕ(w 2,k0) ϕ(w 2,k1)\nϕ(w 3,k0)\nϕ(w 4,k0) ϕ(w 4,k3)ϕ(w 3,k1+k 2)\nϕ(w 4,k1+k 2)\n\nFigure 9: Depth mixing matrices Mfor four residual variants ( L=4 ; Block AttnRes uses block size S=2 ). Highway is shown with\nscalar gates for clarity. AttnRes panels show unnormalized ϕscores; background colors group entries that share the same source\n(Full AttnRes) or the same source block (Block AttnRes).\n•Standard residual [12], hl=hl−1+fl−1(hl−1). Expanding gives hl=Pl−1\ni=0vi, soMi→l= 1for all i < l andM\nis an all-ones lower-triangular matrix:\n\nh1\nh2\n...\nhL\n=\n1\n1 1\n.........\n1 1···1\n\nv0\nv1\n...\nvL−1\n\n•Highway [45], hl= (1−g l)hl−1+glfl−1(hl−1)(written here with scalar gates for clarity; the element-wise\nextension is straightforward). Defining the carry product γ×\ni→l:=Ql\nj=i+1(1−g j), the weights are M0→l=γ×\n1→l\nfor the embedding and Mi→l=gi+1γ×\ni+1→lfori≥1 . Since the cumulative products factor through scalar gates, M\nis 1-semiseparable [8], the same rank as the standard residual but with input-dependent weights. The weights sum to\none by construction, making Highway a softmax-free depth-wise instance of stick-breaking attention [49].\n• (m)HC [72, 59] maintainmparallel streamsH l∈Rd×m, updated via\nHl=H l−1Al+fl−1(Hl−1αl−1)β⊤\nl−1,\nwhere Al∈Rm×mis a learned transition matrix, αl−1∈Rmmixes streams into a single input for fl−1, and\nβl−1∈Rmdistributes the output back across streams. Unrolling the recurrence gives the effective weight\nMi→l=β⊤\niA×\ni+1→lαl,(10)\nwhereA×\ni→j:=Qj\nk=i+1Ak. The m×m transitions render Mm -semiseparable [8]. mHC [59, 64] further constrains\neachA lto be doubly stochastic, stabilizing the cumulative products across depth.\n•Full AttnRes computes Mi→l=α i→lviaϕ(w l,ki) = exp\u0000\nw⊤\nlRMSNorm(k i)\u0001\nwith normalization, where\nki=viare input-dependent layer outputs, yielding a dense, rank-LM.\n•Block AttnRes partitions layers into Nblocks B1, . . . ,B N. For sources iin a completed earlier block Bn, all share\nthe block-level key/value bn, soMi→l=αn→lfor every i∈ B n. Within the current block, each layer additionally\nattends over the evolving partial sum bi−1\nn, introducing one extra distinct source per intra-block position. The effective\nrank of Mtherefore lies between NandN+S (where Sis the block size), interpolating between standard residual\n(N=1) and Full AttnRes (N=L).\nPracticality.The structured-matrix perspective serves two purposes. First, it enables analytical insights that are not\napparent from the recurrence form alone. The input-dependent Mof AttnRes, for instance, reveals depth-wise attention\nsinks (§5.4.2), where certain layers consistently attract high weight regardless of input, mirroring the same phenomenon\nin sequence-wise attention [57]. Second, it informs new designs by exposing which properties of the kernel ϕmatter. For\nexample, when ϕdecomposes as ϕ(q,k) =φ(q)⊤φ(k) for some feature map φ[23], depth-wise attention collapses\ninto a recurrence—precisely the structure underlying the MRLA–GLA and DDL–DeltaNet correspondences noted\nabove.\n15" - }, - { - "page": 16, - "content": "Attention ResidualsTECHNICALREPORT\nPrior Residuals as Depth-Wise Linear AttentionThe structured-matrix perspective further relates to the sequence-\ndepth duality by showing that existing residual variants are, in effect, instances oflinearattention over the depth axis.\nFor example, the unrolled (m)HC weight Mi→l=β⊤\niA×\ni+1→lαl(Eq. 10) admits a natural attention interpretation in\nwhich αlplays the role of a query issued by layer l,βiserves as a key summarizing the contribution of layer i, and\nthe cumulative transition A×\ni+1→lacts as a depth-relative positional operator [69] governing the query–key interaction\nacross intervening layers. Notably, themparallel streams correspond to state expansion [40, 29] along the depth axis,\nexpanding the recurrent state from dtod×m and thereby increasing the semiseparable rank of M. [58] show that\nreplacing A×\ni+1→lwith the identity matrix still yields competitive performance, highlighting the role of state expansion.\nThrough this lens, methods like (m)HC thus act as depth-wiselinearattention with matrix-valued states, while AttnRes\nacts as depth-wisesoftmaxattention.\n7 Related Work\nNormalization, Scaling, and Depth Stability.The standard residual update hl+1=h l+fl(hl)[12] presents a\nfundamental tension betweennormalization placementandgradient propagation. PostNorm [52] maintains bounded\nmagnitudes but distorts gradients, as repeated normalization on the residual path compounds into gradient vanishing at\ndepth [60]. PreNorm [34, 60] restores a clean identity path yet introduces unbounded magnitude growth: since ∥hl∥\ngrows as O(L) , each layer’s relative contribution shrinks, compelling deeper layers to produce ever-larger outputs\nand limiting effective depth [27]. Subsequent work reconciles both desiderata via scaled residual paths [54], hybrid\nnormalization [73], amplified skip connections [4], or learned element-wise gates [45] (see Table 5). AttnRes sidesteps\nthis tension by replacing the additive recurrence with selective aggregation over individual earlier-layer outputs, avoiding\nboth the cumulative magnitude growth of PreNorm and the repeated scale contraction of PostNorm.\nMulti-State Recurrence.All single-state methods above condition layer lonly on hl−1, from which individual\nearlier-layer contributions cannot be selectively retrieved. Several methods address this by widening the recurrence\nto multiple parallel streams: Hyper-Connections [72] and its stabilized variant mHC [59] maintain mstreams with\nlearned mixing matrices; DDL [67] maintains a matrix state updated via a delta-rule erase-and-write mechanism;\nSiameseNorm [27] maintains two parameter-shared streams—one PreNorm and one PostNorm—to preserve identity\ngradients and bounded representations. While these methods alleviate information compression, they still condition\non the immediate predecessor’s state; AttnRes is orthogonal, providing selective access to individual earlier-layer\noutputs while remaining compatible with any normalization or gating scheme. We discuss the formal connection to\nHyper-Connections in § 6.2.\nCross-Layer Connectivity.A separate line of work bypasses the single-state bottleneck by giving each layer direct\naccess to individual earlier-layer outputs. The simplest approach uses static weights: DenseNet [17] concatenates all\npreceding feature maps; ELMo [38] computes a softmax -weighted sum of layer representations with learned scalar\nweights; DenseFormer [36] and ANCRe [68] assign learned per-layer scalar coefficients fixed after training. For\ninput-dependent aggregation, MUDDFormer [56] generates position-dependent weights via a small MLP across four\ndecoupled streams; MRLA [10] applies element-wise sigmoid gating over all previous layers, though its separable\nquery–key product is closer to linear attention than softmax -based retrieval. Other methods trade full cross-layer access\nfor more targeted designs: Value Residual Learning [71] accesses only a single earlier layer; LAuReL [30] augments\nthe residual with low-rank projections over the previous kactivations; Dreamer [24] combines sequence attention with\ndepth attention and sparse experts. AttnRes combines softmax -normalized, input-dependent weights with selective\naccess to all preceding layers through a single d-dimensional pseudo-query per layer, and introduces a block structure\nreducing cost from O(L2)toO(LN) . Cache-based pipeline communication and a two-phase computation strategy\n(§ 4) make Block AttnRes practical at scale with negligible overhead.\nConclusion\nInspired by the duality between sequence and depth, we introduce AttnRes, which replaces fixed, uniform residual\naccumulation with learned, input-dependent depth-wise attention. We validate the method through ablation studies and\nscaling law experiments, showing that its gains persist across scales. Because Full AttnRes must access all preceding\nlayer outputs at every layer, the memory footprint of cross-layer aggregation grows as O(Ld) , which is prohibitive\nfor large-scale models on current hardware. We therefore introduce Block AttnRes, which partitions layers into N\nblocks and attends over block-level representations. Empirically, using about 8 blocks recovers most of the gains of Full\nAttnRes, while finer-grained blocking remains a promising direction as future hardware constraints relax. Together with\ncross-stage caching and a two-phase computation strategy, Block AttnRes is practical at scale, incurring only marginal\ntraining overhead and minimal inference overhead.\n16" - }, - { - "page": 17, - "content": "Attention ResidualsTECHNICALREPORT\nReferences\n[1] Jacob Austin et al.Program Synthesis with Large Language Models. 2021. arXiv: 2108.07732 [cs.PL] .URL:\nhttps://arxiv.org/abs/2108.07732.\n[2] Thomas Bachlechner et al.ReZero is All You Need: Fast Convergence at Large Depth. 2020. arXiv: 2003.04887\n[cs.LG].URL:https://arxiv.org/abs/2003.04887.\n[3] Dzmitry Bahdanau, Kyunghyun Cho, and Yoshua Bengio.Neural Machine Translation by Jointly Learning to\nAlign and Translate. 2016. arXiv:1409.0473 [cs.CL].URL:https://arxiv.org/abs/1409.0473.\n[4] Chen Chen and Lai Wei.Post-LayerNorm Is Back: Stable, ExpressivE, and Deep. 2026. arXiv: 2601.19895\n[cs.LG].URL:https://arxiv.org/abs/2601.19895.\n[5] Mark Chen et al.Evaluating Large Language Models Trained on Code. 2021. arXiv: 2107.03374 [cs.LG] .\nURL:https://arxiv.org/abs/2107.03374.\n[6] Peter Clark et al. “Think you have Solved Question Answering? Try ARC, the AI2 Reasoning Challenge”. In:\narXiv:1803.05457v1(2018).\n[7] Karl Cobbe et al.Training Verifiers to Solve Math Word Problems. 2021. arXiv: 2110.14168 [cs.LG] .URL:\nhttps://arxiv.org/abs/2110.14168.\n[8] Tri Dao and Albert Gu. “Transformers are SSMs: Generalized Models and Efficient Algorithms Through\nStructured State Space Duality”. In:CoRRabs/2405.21060 (2024).DOI: 10.48550/ARXIV.2405.21060 . arXiv:\n2405.21060.URL:https://doi.org/10.48550/arXiv.2405.21060.\n[9] DeepSeek-AI et al.DeepSeek-V3 Technical Report. 2025. arXiv: 2412.19437 [cs.CL] .URL: https://arxiv.\norg/abs/2412.19437.\n[10] Yanwen Fang et al.Cross-Layer Retrospective Retrieving via Layer Attention. 2023. arXiv: 2302 . 03985\n[cs.CV].URL:https://arxiv.org/abs/2302.03985.\n[11] Andrey Gromov et al.The Unreasonable Ineffectiveness of the Deeper Layers. 2025. arXiv: 2403.17887\n[cs.CL].URL:https://arxiv.org/abs/2403.17887.\n[12] Kaiming He et al.Deep Residual Learning for Image Recognition. 2015. arXiv: 1512.03385 [cs.CV] .URL:\nhttps://arxiv.org/abs/1512.03385.\n[13] Dan Hendrycks et al.Measuring Massive Multitask Language Understanding. 2021. arXiv: 2009.03300\n[cs.CY].URL:https://arxiv.org/abs/2009.03300.\n[14] Dan Hendrycks et al.Measuring Mathematical Problem Solving With the MATH Dataset. 2021. arXiv: 2103.\n03874 [cs.LG].URL:https://arxiv.org/abs/2103.03874.\n[15] Jordan Hoffmann et al.Training Compute-Optimal Large Language Models. 2022. arXiv: 2203.15556 [cs.CL] .\nURL:https://arxiv.org/abs/2203.15556.\n[16] Shengding Hu et al.MiniCPM: Unveiling the Potential of Small Language Models with Scalable Training\nStrategies. 2024. arXiv:2404.06395 [cs.CL].URL:https://arxiv.org/abs/2404.06395.\n[17] Gao Huang et al.Densely Connected Convolutional Networks. 2018. arXiv: 1608.06993 [cs.CV] .URL:\nhttps://arxiv.org/abs/1608.06993.\n[18] Yanping Huang et al. “GPipe: Efficient Training of Giant Neural Networks using Pipeline Parallelism”. In:\nAdvances in NeurIPS. 2019.\n[19] Yuzhen Huang et al. “C-eval: A multi-level multi-discipline chinese evaluation suite for foundation models”. In:\nAdvances in NeurIPS36 (2023), pp. 62991–63010.\n[20] Robert A. Jacobs et al. “Adaptive Mixtures of Local Experts”. In:Neural Computation3.1 (1991), pp. 79–87.\nDOI:10.1162/neco.1991.3.1.79.\n[21] Mandar Joshi et al. “Triviaqa: A large scale distantly supervised challenge dataset for reading comprehension”.\nIn:arXiv preprint arXiv:1705.03551(2017).\n[22] Jared Kaplan et al.Scaling Laws for Neural Language Models. 2020. arXiv: 2001.08361 [cs.LG] .URL:\nhttps://arxiv.org/abs/2001.08361.\n[23] Angelos Katharopoulos et al. “Transformers are RNNs: Fast Autoregressive Transformers with Linear Attention”.\nIn:Proceedings of ICML. Ed. by Hal Daumé III and Aarti Singh. PMLR, 2020, pp. 5156–5165.URL: https:\n//proceedings.mlr.press/v119/katharopoulos20a.html.\n[24] Jonas Knupp et al.Depth-Recurrent Attention Mixtures: Giving Latent Reasoning the Attention it Deserves. 2026.\narXiv:2601.21582 [cs.AI].URL:https://arxiv.org/abs/2601.21582.\n[25] Aitor Lewkowycz et al.Solving Quantitative Reasoning Problems with Language Models. 2022. arXiv: 2206.\n14858 [cs.CL].URL:https://arxiv.org/abs/2206.14858.\n17" - }, - { - "page": 18, - "content": "Attention ResidualsTECHNICALREPORT\n[26] Haonan Li et al. “CMMLU: Measuring massive multitask language understanding in Chinese”. In:Findings\nof the Association for Computational Linguistics: ACL 2024. Ed. by Lun-Wei Ku, Andre Martins, and Vivek\nSrikumar. Bangkok, Thailand: Association for Computational Linguistics, Aug. 2024, pp. 11260–11285.DOI:\n10 . 18653 / v1 / 2024 . findings - acl . 671 .URL: https : / / aclanthology . org / 2024 . findings -\nacl.671/.\n[27] Tianyu Li et al.SiameseNorm: Breaking the Barrier to Reconciling Pre/Post-Norm. 2026. arXiv: 2602.08064\n[cs.LG].URL:https://arxiv.org/abs/2602.08064.\n[28] Jingyuan Liu et al.Muon is Scalable for LLM Training. 2025. arXiv: 2502.16982 [cs.LG] .URL: https:\n//arxiv.org/abs/2502.16982.\n[29] Brian Mak and Jeffrey Flanigan.Residual Matrix Transformers: Scaling the Size of the Residual Stream. 2025.\narXiv:2506.22696 [cs.LG].URL:https://arxiv.org/abs/2506.22696.\n[30] Gaurav Menghani, Ravi Kumar, and Sanjiv Kumar.LAuReL: Learned Augmented Residual Layer. 2025. arXiv:\n2411.07501 [cs.LG].URL:https://arxiv.org/abs/2411.07501.\n[31] Maxim Milakov and Natalia Gimelshein.Online normalizer calculation for softmax. 2018. arXiv: 1805.02867\n[cs.PF].URL:https://arxiv.org/abs/1805.02867.\n[32] Tsendsuren Munkhdalai et al. “Metalearned Neural Memory”. In:ArXivabs/1907.09720 (2019).URL: https:\n//api.semanticscholar.org/CorpusID:198179407.\n[33] Deepak Narayanan et al.Efficient Large-Scale Language Model Training on GPU Clusters Using Megatron-LM.\n2021. arXiv:2104.04473 [cs.CL].URL:https://arxiv.org/abs/2104.04473.\n[34] Toan Q. Nguyen and Julian Salazar. “Transformers without Tears: Improving the Normalization of Self-\nAttention”. In:Proceedings of IWSLT. Ed. by Jan Niehues et al. 2019.URL: https : / / aclanthology .\norg/2019.iwslt-1.17/.\n[35] OpenAI et al.GPT-4 Technical Report. 2024. arXiv: 2303.08774 [cs.CL] .URL: https://arxiv.org/abs/\n2303.08774.\n[36] Matteo Pagliardini et al.DenseFormer: Enhancing Information Flow in Transformers via Depth Weighted\nAveraging. 2024. arXiv:2402.02622 [cs.CL].URL:https://arxiv.org/abs/2402.02622.\n[37] Bowen Peng et al. “Yarn: Efficient context window extension of large language models”. In:arXiv preprint\narXiv:2309.00071(2023).\n[38] Matthew E. Peters et al. “Deep Contextualized Word Representations”. In:Proceedings of NAACL. 2018,\npp. 2227–2237.URL:https://aclanthology.org/N18-1202/.\n[39] Reiner Pope et al.Efficiently Scaling Transformer Inference. 2022. arXiv:2211.05102 [cs.LG].\n[40] Zhen Qin et al.HGRN2: Gated Linear RNNs with State Expansion. 2024. arXiv:2404.07904 [cs.CL].\n[41] David Rein et al. “Gpqa: A graduate-level google-proof q&a benchmark”. In:First Conference on Language\nModeling. 2024.\n[42] Imanol Schlag, Kazuki Irie, and Jürgen Schmidhuber. “Linear Transformers Are Secretly Fast Weight Program-\nmers”. In:Proceedings of ICML. Ed. by Marina Meila and Tong Zhang. PMLR, 2021, pp. 9355–9366.URL:\nhttps://proceedings.mlr.press/v139/schlag21a.html.\n[43] Jürgen Schmidhuber. “Learning to control fast-weight memories: An alternative to dynamic recurrent networks”.\nIn:Neural Computation4.1 (1992), pp. 131–139.\n[44] Freda Shi et al.Language Models are Multilingual Chain-of-Thought Reasoners. 2022. arXiv: 2210.03057\n[cs.CL].URL:https://arxiv.org/abs/2210.03057.\n[45] Rupesh Kumar Srivastava, Klaus Greff, and Jürgen Schmidhuber.Highway Networks. 2015. arXiv: 1505.00387\n[cs.LG].URL:https://arxiv.org/abs/1505.00387.\n[46] Yu Sun et al. “Learning to (Learn at Test Time): RNNs with Expressive Hidden States”. In:ArXivabs/2407.04620\n(2024).URL:https://api.semanticscholar.org/CorpusID:271039606.\n[47] Yutao Sun et al.Retentive Network: A Successor to Transformer for Large Language Models. 2023. arXiv:\n2307.08621 [cs.CL].\n[48] Mirac Suzgun et al. “Challenging big-bench tasks and whether chain-of-thought can solve them”. In:arXiv\npreprint arXiv:2210.09261(2022).\n[49] Shawn Tan et al. “Scaling Stick-Breaking Attention: An Efficient Implementation and In-depth Study”. In:\nProceedings of ICLR. 2025.\n[50] Hugo Touvron et al.Going deeper with Image Transformers. 2021. arXiv: 2103.17239 [cs.CV] .URL: https:\n//arxiv.org/abs/2103.17239.\n[51] Hugo Touvron et al.LLaMA: Open and Efficient Foundation Language Models. 2023. arXiv: 2302.13971\n[cs.CL].\n18" - }, - { - "page": 19, - "content": "Attention ResidualsTECHNICALREPORT\n[52] Ashish Vaswani et al. “Attention is All you Need”. In:Advances in NeurIPS. Ed. by I. Guyon et al. Curran\nAssociates, Inc., 2017.URL: https://proceedings.neurips.cc/paper_files/paper/2017/file/\n3f5ee243547dee91fbd053c1c4a845aa-Paper.pdf.\n[53] Ashish Vaswani et al. “Attention is All you Need”. In:Advances in NeurIPS. Ed. by I. Guyon et al. V ol. 30.\nCurran Associates, Inc., 2017.URL: https://proceedings.neurips.cc/paper_files/paper/2017/\nfile/3f5ee243547dee91fbd053c1c4a845aa-Paper.pdf.\n[54] Hongyu Wang et al.DeepNet: Scaling Transformers to 1,000 Layers. 2022. arXiv: 2203.00555 [cs.CL] .URL:\nhttps://arxiv.org/abs/2203.00555.\n[55] Yubo Wang et al. “Mmlu-pro: A more robust and challenging multi-task language understanding benchmark”. In:\nAdvances in NeurIPS37 (2024), pp. 95266–95290.\n[56] Da Xiao et al. “MUDDFormer: Breaking Residual Bottlenecks in Transformers via Multiway Dynamic Dense\nConnections”. In:Proceedings of ICML. 2025.\n[57] Guangxuan Xiao et al. “Efficient streaming language models with attention sinks”. In:arXiv preprint\narXiv:2309.17453(2023).\n[58] Tian Xie.Your DeepSeek mHC Might Not Need the “m”. Zhihu blog post. 2026.URL: https://zhuanlan.\nzhihu.com/p/2010852389670908320.\n[59] Zhenda Xie et al.mHC: Manifold-Constrained Hyper-Connections. 2026. arXiv: 2512.24880 [cs.CL] .URL:\nhttps://arxiv.org/abs/2512.24880.\n[60] Ruibin Xiong et al.On Layer Normalization in the Transformer Architecture. 2020. arXiv: 2002.04745 [cs.LG] .\nURL:https://arxiv.org/abs/2002.04745.\n[61] Bowen Yang et al.Rope to Nope and Back Again: A New Hybrid Attention Strategy. 2025. arXiv: 2501.18795\n[cs.CL].URL:https://arxiv.org/abs/2501.18795.\n[62] Songlin Yang, Jan Kautz, and Ali Hatamizadeh. “Gated Delta Networks: Improving Mamba2 with Delta Rule”.\nIn:Proceedings of ICLR. 2025.URL:https://openreview.net/forum?id=r8H7xhYPwz.\n[63] Songlin Yang et al. “Gated Linear Attention Transformers with Hardware-Efficient Training”. In:Proceedings of\nICML. PMLR, 2024.\n[64] Yongyi Yang and Jianyang Gao.mHC-lite: You Don’t Need 20 Sinkhorn-Knopp Iterations. 2026. arXiv: 2601.\n05732 [cs.LG].URL:https://arxiv.org/abs/2601.05732.\n[65] Rowan Zellers et al. “HellaSwag: Can a Machine Really Finish Your Sentence?” In:Proceedings of the 57th\nAnnual Meeting of the Association for Computational Linguistics. 2019.\n[66] Biao Zhang and Rico Sennrich. “Root mean square layer normalization”. In:Advances in NeurIPS32 (2019).\n[67] Yifan Zhang et al.Deep Delta Learning. 2026. arXiv: 2601.00417 [cs.LG] .URL: https://arxiv.org/\nabs/2601.00417.\n[68] Yilang Zhang et al.ANCRe: Adaptive Neural Connection Reassignment for Efficient Depth Scaling. 2026. arXiv:\n2602.09009 [cs.LG].URL:https://arxiv.org/abs/2602.09009.\n[69] Yu Zhang et al.Kimi Linear: An Expressive, Efficient Attention Architecture. 2025. arXiv: 2510.26692 [cs.CL] .\n[70] Shu Zhong et al.Understanding Transformer from the Perspective of Associative Memory. 2025. arXiv: 2505.\n19488 [cs.LG].URL:https://arxiv.org/abs/2505.19488.\n[71] Zhanchao Zhou et al. “Value Residual Learning”. In:Proceedings of ACL. Ed. by Wanxiang Che et al. Vienna,\nAustria, 2025, pp. 28341–28356.URL:https://aclanthology.org/2025.acl-long.1375/.\n[72] Defa Zhu et al.Hyper-Connections. 2025. arXiv: 2409.19606 [cs.LG] .URL: https://arxiv.org/abs/\n2409.19606.\n[73] Zhijian Zhuo et al.HybridNorm: Towards Stable and Efficient Transformer Training via Hybrid Normalization.\n2025. arXiv:2503.04598 [cs.CL].URL:https://arxiv.org/abs/2503.04598.\n19" - }, - { - "page": 20, - "content": "Attention ResidualsTECHNICALREPORT\nA Contributions\nThe authors are listed in order of the significance of their contributions, with those in project leadership roles appearing\nlast.\nGuangyu Chen∗\nYu Zhang∗\nJianlin Su∗\nWeixin Xu\nSiyuan Pan\nYaoyu Wang\nYucheng Wang\nGuanduo Chen\nBohong Yin\nYutian Chen\nJunjie Yan\nMing Wei\nY . Zhang\nFanqing Meng\nChao Hong\nXiaotong Xie\nShaowei Liu\nEnzhe Lu\nYunpeng TaiYanru Chen\nXin Men\nHaiqing Guo\nY . Charles\nHaoyu Lu\nLin Sui\nJinguo Zhu\nZaida Zhou\nWeiran He\nWeixiao Huang\nXinran Xu\nYuzhi Wang\nGuokun Lai\nYulun Du\nYuxin Wu\nZhilin Yang\nXinyu Zhou\n∗Equal contribution\n20" - }, - { - "page": 21, - "content": "Attention ResidualsTECHNICALREPORT\nB Optimized Inference I/O for Full Attention Residuals\nA naïve implementation of Full AttnRes scans all preceding layer outputs at every layer, so memory traffic scales\nlinearly with depth. As noted in §4.2, however, the pseudo-query wlis a learned parameter independent of both the\ninput and the hidden state. We can therefore batch inter-block accesses across layers in a two-phase schedule, bringing\ntotal I/O well below the naïve bound.\nNote that the block partition introduced below is purely an inference scheduling device. Unlike Block AttnRes, it leaves\nthe model architecture unchanged and does not replace per-layer sources with block summaries; it simply makes the\namortization argument concrete.\nSetupLet the model have Llayers and hidden dimension d, partitioned into Ncontiguous blocks of size S=L/N .\nInference proceeds one block at a time: Phase 1 jointly computes inter-block attention for all Slayers in the block\nagainst all preceding blocks, and Phase 2 walks through intra-block dependencies sequentially.\nPhase 1: Batched Inter-block Attention\nConsider block nwith its Slayers. The queries {wl}l∈Bnare all known before execution begins, so the (n−1)S\npreceding key–value pairs need only be read once from HBM and reused across all Squeries. The read cost for block n\nis therefore\nRead(n)\ninter= 2(n−1)Sd,(11)\nwhere the factor of2accounts for both keys and values. Summing over allNblocks and usingSN=L:\nRead inter=NX\nn=12(n−1)Sd= 2Sd·N(N−1)\n2=dL(N−1).(12)\nPhase 1 also writes oned-dimensional output per layer, givingWrite(n)\ninter=Sdper block and\nWrite inter=Ld(13)\nin total.\nPhase 2: Sequential Intra-block Attention\nPhase 1 covers all sources before the current block. Within the block, however, each layer depends on those before it,\nso these must be handled in order. Layer t(1≤t≤S ) reads t−1 intra-block key–value pairs at a cost of 2(t−1)d .\nSumming over one block:\nRead(n)\nintra=SX\nt=12(t−1)d=S(S−1)d.(14)\nPhase 2 also writes one output per layer, soWrite(n)\nintra=Sd.\nTotal Amortized I/O per Layer\nSumming both phases over allNblocks:\nRead total=dL(N−1) +N·S(S−1)d,Write total= 2Ld.(15)\nDividing byLand usingSN=L:\nRead per layer= (N−1)d+ (S−1)d= (S+N−2)d,Write per layer= 2d,(16)\nTotal I/O per layer= (S+N)d. (17)\nBatching inter-block reads thus brings per-layer I/O from O(L) down to O(S+N) . The schedule follows the same\ntwo-phase split as Block AttnRes: inter-block attention accounts for the bulk of the traffic, while sequential computation\nstays local within each block.\n21" - } - ] -} \ No newline at end of file diff --git a/examples/workspace/_meta.json b/examples/workspace/_meta.json deleted file mode 100644 index daf212c70..000000000 --- a/examples/workspace/_meta.json +++ /dev/null @@ -1,9 +0,0 @@ -{ - "12345678-abcd-4321-abcd-123456789abc": { - "type": "pdf", - "doc_name": "attention-residuals.pdf", - "doc_description": "This document introduces \"Attention Residuals\" (AttnRes) and its scalable variant \"Block AttnRes,\" novel mechanisms for replacing fixed residual accumulation in neural networks with learned, input-dependent depth-wise attention, addressing limitations of standard residual connections while optimizing memory, computation, and scalability for large-scale training and inference.", - "page_count": 21, - "path": "../documents/attention-residuals.pdf" - } -} \ No newline at end of file diff --git a/pageindex/__init__.py b/pageindex/__init__.py index 3f4df8d6d..3513668a2 100644 --- a/pageindex/__init__.py +++ b/pageindex/__init__.py @@ -1,5 +1,46 @@ -from .page_index import * -from .page_index_md import md_to_tree -from .retrieve import get_document, get_document_structure, get_page_content -from .client import PageIndexClient -from .tree_optimize import optimize_tree +"""PageIndex SDK.""" +from typing import TYPE_CHECKING as _TYPE_CHECKING + +from .client import PageIndexClient, PageIndexCloudClient, PageIndexLocalClient +from .errors import PageIndexAPIError + +if _TYPE_CHECKING: + from .flash import page_index_flash + from .page_index_classic import page_index, page_index_main + from .page_index_md import md_to_tree + from .tree_optimize import optimize_tree + +__all__ = [ + "PageIndexClient", "PageIndexCloudClient", "PageIndexLocalClient", + "PageIndexAPIError", + "page_index", "page_index_main", "page_index_flash", + "optimize_tree", "md_to_tree", +] + +_LAZY = { + "page_index_flash": ".flash", + "optimize_tree": ".tree_optimize", + "md_to_tree": ".page_index_md", +} +_SUBMODULES = {"client", "cloud_api", "errors", "flash", "local_api", + "local_store", "page_index_classic", "page_index_md", "tree_optimize", + "utils"} + + +def __getattr__(name): + if name.startswith("_"): + raise AttributeError(f"module {__name__!r} has no attribute {name!r}") + import importlib + if name in _SUBMODULES: + return importlib.import_module(f".{name}", __name__) + module = importlib.import_module(_LAZY.get(name, ".page_index_classic"), __name__) + try: + value = getattr(module, name) + except AttributeError: + raise AttributeError(f"module {__name__!r} has no attribute {name!r}") from None + globals()[name] = value + return value + + +def __dir__(): + return sorted(set(globals()) | set(__all__) | _SUBMODULES) diff --git a/pageindex/client.py b/pageindex/client.py index 894dab181..158c9b6f7 100644 --- a/pageindex/client.py +++ b/pageindex/client.py @@ -1,18 +1,23 @@ -import os -import uuid -import json -import asyncio -import concurrent.futures -from pathlib import Path +"""PageIndex SDK client: the 0.2.x cloud surface, now with a local mode.""" +from __future__ import annotations -import PyPDF2 +from typing import Any, Iterator, Optional, Union -from .page_index import page_index -from .page_index_md import md_to_tree -from .retrieve import get_document, get_document_structure, get_page_content -from .utils import ConfigLoader, remove_fields +from .errors import PageIndexAPIError -META_INDEX = "_meta.json" + +def _parse_pages(pages: str) -> list[int]: + result = [] + for part in pages.split(","): + part = part.strip() + if "-" in part: + start, end = (int(x) for x in part.split("-", 1)) + if start > end: + raise ValueError(f"Invalid range '{part}': start must be <= end") + result.extend(range(start, end + 1)) + else: + result.append(int(part)) + return sorted(set(result)) def _normalize_retrieve_model(model: str) -> str: @@ -27,208 +32,398 @@ def _normalize_retrieve_model(model: str) -> str: class PageIndexClient: """ - A client for indexing and retrieving document content. - Flow: index() -> get_document() / get_document_structure() / get_page_content() + Python SDK client for PageIndex. + + Cloud mode (an ``api_key`` is given) talks to the PageIndex API at + api.pageindex.ai, exactly like the 0.2.x SDK. Local mode (no ``api_key``) + runs the same operations on your machine: documents are indexed with the + open-source PageIndex pipeline using your own LLM provider key (e.g. + ``OPENAI_API_KEY`` in the environment) and stored under ``storage_path``. + + Args: + api_key (str, optional): PageIndex cloud API key + (https://dash.pageindex.ai/api-keys). Omit for local mode. + model (str, optional): Local mode only — LLM used to build document + trees. Defaults to the packaged config (see pageindex/config.yaml). + summary_model (str, optional): Local mode only — LLM used for node + summaries and document descriptions. + retrieve_model (str, optional): Local mode only — exposed as + ``client.retrieve_model`` (the agent demo reads it); the SDK + itself consumes it once agent-based local chat lands in a + later release. + storage_path (str, optional): Local mode only — directory where + indexed documents are stored. Defaults to ``./.pageindex``. - For agent-based QA, see examples/agentic_vectorless_rag_demo.py. + Usage: + client = PageIndexClient(api_key="...") # cloud + client = PageIndexClient() # local + + PageIndexCloudClient / PageIndexLocalClient pin the mode at construction + instead of inferring it from api_key. + + Local mode differences (all documented per method): indexing is + synchronous, only PDFs are supported, and ``chat_completions`` (until + agent-based local chat lands in a later release) / folders / + ``beta_headers`` / the deprecated retrieval API (``submit_query``, + ``get_retrieval``) are cloud-only. """ - def __init__(self, api_key: str = None, model: str = None, retrieve_model: str = None, workspace: str = None): - if api_key: - os.environ["OPENAI_API_KEY"] = api_key - elif not os.getenv("OPENAI_API_KEY") and os.getenv("CHATGPT_API_KEY"): - os.environ["OPENAI_API_KEY"] = os.getenv("CHATGPT_API_KEY") - self.workspace = Path(workspace).expanduser() if workspace else None - overrides = {} - if model: - overrides["model"] = model - if retrieve_model: - overrides["retrieve_model"] = retrieve_model - opt = ConfigLoader().load(overrides or None) - self.model = opt.model - self.retrieve_model = _normalize_retrieve_model(opt.retrieve_model or self.model) - if self.workspace: - self.workspace.mkdir(parents=True, exist_ok=True) - self.documents = {} - if self.workspace: - self._load_workspace() - - def index(self, file_path: str, mode: str = "auto") -> str: - """Index a document. Returns a document_id.""" - # Persist a canonical absolute path so workspace reloads do not - # reinterpret caller-relative paths against the workspace directory. - file_path = os.path.abspath(os.path.expanduser(file_path)) - if not os.path.exists(file_path): - raise FileNotFoundError(f"File not found: {file_path}") - - doc_id = str(uuid.uuid4()) - ext = os.path.splitext(file_path)[1].lower() - - is_pdf = ext == '.pdf' - is_md = ext in ['.md', '.markdown'] - - if mode == "pdf" or (mode == "auto" and is_pdf): - print(f"Indexing PDF: {file_path}") - result = page_index( - doc=file_path, - model=self.model, - if_add_node_summary='yes', - if_add_node_text='yes', - if_add_node_id='yes', - if_add_doc_description='yes' + + BASE_URL = "https://api.pageindex.ai" + + def __init__( + self, + api_key: Optional[str] = None, + *, + model: Optional[str] = None, + summary_model: Optional[str] = None, + retrieve_model: Optional[str] = None, + storage_path: Optional[str] = None, + ): + if api_key == "": + raise PageIndexAPIError( + "api_key is an empty string. Pass a real PageIndex API key for " + "cloud mode, or omit api_key entirely for local mode." ) - # Extract per-page text so queries don't need the original PDF - pages = [] - with open(file_path, 'rb') as f: - pdf_reader = PyPDF2.PdfReader(f) - for i, page in enumerate(pdf_reader.pages, 1): - pages.append({'page': i, 'content': page.extract_text() or ''}) - - self.documents[doc_id] = { - 'id': doc_id, - 'type': 'pdf', - 'path': file_path, - 'doc_name': result.get('doc_name', ''), - 'doc_description': result.get('doc_description', ''), - 'page_count': len(pages), - 'structure': result['structure'], - 'pages': pages, - } - - elif mode == "md" or (mode == "auto" and is_md): - print(f"Indexing Markdown: {file_path}") - coro = md_to_tree( - md_path=file_path, - if_thinning=False, - if_add_node_summary='yes', - summary_token_threshold=200, + if api_key is not None: + local_only = {"model": model, "summary_model": summary_model, + "retrieve_model": retrieve_model, "storage_path": storage_path} + passed = [name for name, value in local_only.items() if value is not None] + if passed: + raise PageIndexAPIError( + f"Local-mode arguments ({', '.join(passed)}) cannot be " + "combined with api_key — remove them, or omit api_key to " + "run locally." + ) + self.api_key = api_key + from .cloud_api import CloudAPI + self._api = CloudAPI(self) + else: + from .utils import ConfigLoader + overrides = {key: value for key, value in + {"model": model, "summary_model": summary_model, + "retrieve_model": retrieve_model}.items() + if value} + opt = ConfigLoader().load(overrides or None) + self.model = opt.model + self.summary_model = getattr(opt, "summary_model", None) or opt.model + self.retrieve_model = _normalize_retrieve_model( + getattr(opt, "retrieve_model", None) or opt.model) + self.storage_path = storage_path or ".pageindex" + from .local_api import LocalAPI + self._api = LocalAPI( + storage_path=self.storage_path, model=self.model, - if_add_doc_description='yes', - if_add_node_text='yes', - if_add_node_id='yes' + summary_model=self.summary_model, + retrieve_model=self.retrieve_model, ) - try: - asyncio.get_running_loop() - with concurrent.futures.ThreadPoolExecutor(max_workers=1) as pool: - result = pool.submit(asyncio.run, coro).result() - except RuntimeError: - result = asyncio.run(coro) - self.documents[doc_id] = { - 'id': doc_id, - 'type': 'md', - 'path': file_path, - 'doc_name': result.get('doc_name', ''), - 'doc_description': result.get('doc_description', ''), - 'line_count': result.get('line_count', 0), - 'structure': result['structure'], - } - else: - raise ValueError(f"Unsupported file format for: {file_path}") - - print(f"Indexing complete. Document ID: {doc_id}") - if self.workspace: - self._save_doc(doc_id) - return doc_id - - @staticmethod - def _make_meta_entry(doc: dict) -> dict: - """Build a lightweight meta entry from a document dict.""" - entry = { - 'type': doc.get('type', ''), - 'doc_name': doc.get('doc_name', ''), - 'doc_description': doc.get('doc_description', ''), - 'path': doc.get('path', ''), - } - if doc.get('type') == 'pdf': - entry['page_count'] = doc.get('page_count') - elif doc.get('type') == 'md': - entry['line_count'] = doc.get('line_count') - return entry - - @staticmethod - def _read_json(path) -> dict | None: - """Read a JSON file, returning None on any error.""" + + # ---------- DOCUMENT SUBMISSION ---------- + + def submit_document( + self, + file_path: str, + mode: Optional[str] = None, + beta_headers: Optional[list[str]] = None, + folder_id: Optional[str] = None, + metadata: Optional[dict] = None, + ) -> dict[str, Any]: + """ + Submit a PDF document for processing. Returns {'doc_id': ...}. + + Cloud: uploads the file; processing is asynchronous — poll + ``is_retrieval_ready(doc_id)`` before retrieving. + + Local: indexes the document in this call (it blocks while your LLM + builds the tree — minutes for a standard index of a long document), + then stores it under ``storage_path``. Pass ``mode="flash"`` to build + the tree with PageIndex Flash (layout-based extraction, no LLM calls + for the structure; node summaries and the document description still + use ``summary_model``). ``beta_headers`` and ``folder_id`` are + cloud-only. + + Args: + file_path (str): Path to the PDF file. + mode (str, optional): Processing mode. Local mode supports + "standard" and "flash"; omit it for standard indexing. Cloud + modes are passed through (e.g. "mcp"). + beta_headers (list[str], optional): Cloud-only beta feature headers. + folder_id (str, optional): Cloud-only folder (workspace) ID. + metadata (dict, optional): Your own JSON-serializable tags for the + document; returned in get_tree/get_ocr responses and + list_documents entries (both modes). + + Returns: + dict: {'doc_id': ...} + """ + return self._api.submit_document( + file_path=file_path, mode=mode, + beta_headers=beta_headers, folder_id=folder_id, metadata=metadata, + ) + + # ---------- OCR FUNCTIONALITY ---------- + + def get_ocr(self, doc_id: str, format: str = "page") -> dict[str, Any]: + """ + Get OCR status and results. + + Args: + doc_id (str): Document ID. + format (str): 'page' for page-based results, 'node' for node-based + results, or 'raw' for concatenated markdown. + + Returns: + dict: {'doc_id', 'status', 'retrieval_ready', 'result', ...}. + With 'page', result entries are {'page_index', 'markdown', ...}. + + Local: the "OCR" result is the text extracted from the PDF while + indexing (no OCR model runs locally, so scanned/image-only PDFs have + no local text). + """ + return self._api.get_ocr(doc_id=doc_id, format=format) + + def get_page_content(self, doc_id: str, pages: str) -> list[dict[str, Any]]: + """ + Get text content of specific pages. + + Args: + doc_id (str): Document ID. + pages (str): Page specifier — '5-7', '3,8', or '12'. + + Returns: + list: Matching entries from get_ocr (format='page'). + """ + wanted = set(_parse_pages(pages)) + result = self.get_ocr(doc_id, format="page") + all_pages = result["result"] + if all_pages is None: + raise PageIndexAPIError( + f"Document '{doc_id}' is not ready " + f"(status: {result.get('status', 'unknown')})" + ) + return [p for p in all_pages if p["page_index"] in wanted] + + # ---------- TREE GENERATION ---------- + + def get_tree(self, doc_id: str, node_summary: bool = False, + include_text: bool = True) -> dict[str, Any]: + """ + Get tree generation status and results. + + Args: + doc_id (str): Document ID. + node_summary (bool): Include node summaries in the tree. + include_text (bool): Include node text (default True). + False is useful for structure-only views (saves tokens). + + Returns: + dict: {'doc_id', 'status', 'retrieval_ready', 'result', ...} where + result nodes are {'title', 'node_id', 'page_index', ('summary' / + 'prefix_summary',) ('text',) 'nodes'}. + """ + tree = self._api.get_tree(doc_id=doc_id, node_summary=node_summary, + include_text=include_text) + if not include_text and tree.get("result"): + from .utils import remove_fields + tree["result"] = remove_fields(tree["result"], fields=["text"]) + return tree + + def get_document_structure(self, doc_id: str) -> list[dict[str, Any]]: + """ + Get the document's tree structure without text — summaries included. + + Returns: + list: Tree nodes with titles, page ranges, and summaries. + """ + return self.get_tree(doc_id, node_summary=True, include_text=False)["result"] + + def is_retrieval_ready(self, doc_id: str) -> bool: + """ + Check if a document is ready for retrieval. API errors (including a + missing document) are reported as False; transport errors (connection + failures, timeouts) propagate. + """ try: - with open(path, "r", encoding="utf-8") as f: - return json.load(f) - except (json.JSONDecodeError, OSError) as e: - print(f"Warning: corrupt {Path(path).name}: {e}") - return None - - def _save_doc(self, doc_id: str): - doc = self.documents[doc_id].copy() - # Strip text from structure nodes — redundant with pages (PDF only) - if doc.get('structure') and doc.get('type') == 'pdf': - doc['structure'] = remove_fields(doc['structure'], fields=['text']) - path = self.workspace / f"{doc_id}.json" - with open(path, "w", encoding="utf-8") as f: - json.dump(doc, f, ensure_ascii=False, indent=2) - self._save_meta(doc_id, self._make_meta_entry(doc)) - # Drop heavy fields; will lazy-load on demand - self.documents[doc_id].pop('structure', None) - self.documents[doc_id].pop('pages', None) - - def _rebuild_meta(self) -> dict: - """Scan individual doc JSON files and return a meta dict.""" - meta = {} - for path in self.workspace.glob("*.json"): - if path.name == META_INDEX: - continue - doc = self._read_json(path) - if doc and isinstance(doc, dict): - meta[path.stem] = self._make_meta_entry(doc) - return meta - - def _read_meta(self) -> dict | None: - """Read and validate _meta.json, returning None on any corruption.""" - meta = self._read_json(self.workspace / META_INDEX) - if meta is not None and not isinstance(meta, dict): - print(f"Warning: {META_INDEX} is not a JSON object, ignoring") - return None - return meta - - def _save_meta(self, doc_id: str, entry: dict): - meta = self._read_meta() or self._rebuild_meta() - meta[doc_id] = entry - meta_path = self.workspace / META_INDEX - with open(meta_path, "w", encoding="utf-8") as f: - json.dump(meta, f, ensure_ascii=False, indent=2) - - def _load_workspace(self): - meta = self._read_meta() - if meta is None: - meta = self._rebuild_meta() - if meta: - print(f"Loaded {len(meta)} document(s) from workspace (legacy mode).") - for doc_id, entry in meta.items(): - doc = dict(entry, id=doc_id) - if doc.get('path') and not os.path.isabs(doc['path']): - doc['path'] = str((self.workspace / doc['path']).resolve()) - self.documents[doc_id] = doc - - def _ensure_doc_loaded(self, doc_id: str): - """Load full document JSON on demand (structure, pages, etc.).""" - doc = self.documents.get(doc_id) - if not doc or doc.get('structure') is not None: - return - full = self._read_json(self.workspace / f"{doc_id}.json") - if not full: - return - doc['structure'] = full.get('structure', []) - if full.get('pages'): - doc['pages'] = full['pages'] - - def get_document(self, doc_id: str) -> str: - """Return document metadata JSON.""" - return get_document(self.documents, doc_id) - - def get_document_structure(self, doc_id: str) -> str: - """Return document tree structure JSON (without text fields).""" - if self.workspace: - self._ensure_doc_loaded(doc_id) - return get_document_structure(self.documents, doc_id) - - def get_page_content(self, doc_id: str, pages: str) -> str: - """Return page content for the given pages string (e.g. '5-7', '3,8', '12').""" - if self.workspace: - self._ensure_doc_loaded(doc_id) - return get_page_content(self.documents, doc_id, pages) + result = self.get_tree(doc_id) + return result.get("retrieval_ready", False) + except PageIndexAPIError: + return False + + # ---------- RETRIEVAL (cloud-only, deprecated) ---------- + + def submit_query(self, doc_id: str, query: str, thinking: bool = False) -> dict[str, Any]: + """ + Submit a retrieval query for a document. Returns {'retrieval_id': ...}. + + Cloud-only: the cloud API marks this endpoint deprecated in favor of + chat completions, so local mode does not implement it — raises + PageIndexAPIError. Use ``chat_completions`` (cloud) instead. + """ + return self._require_cloud( + "submit_query is cloud-only — the retrieval API is deprecated in " + "favor of chat completions; use chat_completions in cloud mode." + ).submit_query(doc_id=doc_id, query=query, thinking=thinking) + + def get_retrieval(self, retrieval_id: str) -> dict[str, Any]: + """ + Get retrieval status and results for a submitted query. + + Cloud-only: the cloud API marks this endpoint deprecated in favor of + chat completions, so local mode does not implement it — raises + PageIndexAPIError. Use ``chat_completions`` (cloud) instead. + """ + return self._require_cloud( + "get_retrieval is cloud-only — the retrieval API is deprecated in " + "favor of chat completions; use chat_completions in cloud mode." + ).get_retrieval(retrieval_id=retrieval_id) + + # ---------- CHAT COMPLETIONS ---------- + + def chat_completions( + self, + messages: list[dict[str, str]], + stream: bool = False, + doc_id: Optional[Union[str, list[str]]] = None, + temperature: Optional[float] = None, + stream_metadata: bool = False, + enable_citations: bool = False, + ) -> Union[dict[str, Any], Iterator[str], Iterator[dict[str, Any]]]: + """ + PageIndex Chat Completions, scoped to specific PageIndex documents. + + Args: + messages: Conversation messages with 'role' and 'content' keys. + stream: Enable streaming responses. + doc_id: Document ID or list of IDs to scope the conversation. + temperature: Sampling temperature (0.0-1.0). + stream_metadata: With stream=True, yield chunk dicts instead of + text pieces. + enable_citations: Enable citation instructions in responses. + + Returns: + - stream=False: complete response dict ({'id', 'object', 'created', + 'choices', 'usage'}) + - stream=True, stream_metadata=False: iterator of text chunks + - stream=True, stream_metadata=True: iterator of chunk dicts + + Local: not yet supported — raises PageIndexAPIError. Agent-based + local chat arrives in a later release. + """ + return self._require_cloud( + "chat_completions is not yet supported in local mode — it arrives " + "in a later release. Create the client with an api_key to use " + "cloud chat." + ).chat_completions( + messages=messages, stream=stream, doc_id=doc_id, + temperature=temperature, stream_metadata=stream_metadata, + enable_citations=enable_citations, + ) + + # ---------- DOCUMENT MANAGEMENT ---------- + + def get_document(self, doc_id: str) -> dict[str, Any]: + """ + Get document metadata: {'id', 'name', 'description', 'status', + 'createdAt', 'pageNum', 'folderId'}. Status is one of "queued", + "processing", "completed", "failed" (local documents are + always "completed"; local 'folderId' is always None). + + 'createdAt' is UTC with no timezone marker, in both modes. To show + it in the user's timezone:: + + from datetime import datetime, timezone + datetime.fromisoformat(doc["createdAt"]).replace( + tzinfo=timezone.utc).astimezone() + """ + return self._api.get_document(doc_id=doc_id) + + def delete_document(self, doc_id: str) -> dict[str, Any]: + """ + Delete a PageIndex document and all its associated data. + + Returns: + dict: {'message': 'Document deleted successfully.'}, or an empty + dict if the cloud API responds with no body. + """ + return self._api.delete_document(doc_id=doc_id) + + def list_documents( + self, + limit: int = 50, + offset: int = 0, + folder_id: Optional[str] = None, + ) -> dict[str, Any]: + """ + List documents with pagination, newest first. + + Args: + limit (int): Maximum documents to return (1-100). + offset (int): Number of documents to skip. + folder_id (str, optional): Cloud-only folder filter. + + Returns: + dict: {'documents': [...], 'total', 'limit', 'offset'}. + """ + return self._api.list_documents(limit=limit, offset=offset, folder_id=folder_id) + + # ---------- FOLDER MANAGEMENT ---------- + + def create_folder( + self, + name: str, + description: Optional[str] = None, + parent_folder_id: Optional[str] = None, + ) -> dict[str, Any]: + """ + Create a folder (workspace). Cloud-only: local mode raises + PageIndexAPIError. + """ + return self._require_cloud( + "create_folder is cloud-only — folders are not supported in local " + "mode. Create the client with an api_key to use folders." + ).create_folder( + name=name, description=description, parent_folder_id=parent_folder_id, + ) + + def list_folders(self, parent_folder_id: Optional[str] = None) -> dict[str, Any]: + """ + List folders. Cloud-only: local mode raises PageIndexAPIError. + """ + return self._require_cloud( + "list_folders is cloud-only — folders are not supported in local " + "mode. Create the client with an api_key to use folders." + ).list_folders( + parent_folder_id=parent_folder_id, + ) + + def _require_cloud(self, message: str): + from .cloud_api import CloudAPI + if not isinstance(self._api, CloudAPI): + raise PageIndexAPIError(message) + return self._api + + +class PageIndexCloudClient(PageIndexClient): + """Cloud mode — requires a real API key at construction.""" + + def __init__(self, api_key: str): + if not api_key: + raise PageIndexAPIError( + "PageIndexCloudClient requires a PageIndex API key — get one " + "at https://dash.pageindex.ai/api-keys." + ) + super().__init__(api_key) + + +class PageIndexLocalClient(PageIndexClient): + """Local mode — no api_key parameter, no cloud access.""" + + def __init__( + self, + *, + model: Optional[str] = None, + summary_model: Optional[str] = None, + retrieve_model: Optional[str] = None, + storage_path: Optional[str] = None, + ): + super().__init__(None, model=model, summary_model=summary_model, + retrieve_model=retrieve_model, storage_path=storage_path) diff --git a/pageindex/cloud_api.py b/pageindex/cloud_api.py new file mode 100644 index 000000000..b7597cc9a --- /dev/null +++ b/pageindex/cloud_api.py @@ -0,0 +1,433 @@ +"""Cloud mode of the PageIndex SDK, based on the 0.2.8 client.""" +import requests +from typing import Optional, Dict, Any, List, Union, Iterator +import json +import urllib.parse + +from .errors import PageIndexAPIError + + +def _enc(value: str) -> str: + """URL-encode a path segment (ids may contain / ? # or spaces).""" + return urllib.parse.quote(str(value), safe="") + + +class CloudAPI: + """ + Python SDK client for the PageIndex API. + """ + + def __init__(self, client): + self._client = client + + @property + def BASE_URL(self) -> str: + return self._client.BASE_URL + + @property + def api_key(self) -> str: + return self._client.api_key + + def _headers(self) -> Dict[str, str]: + return {"api_key": self.api_key} + + # ---------- DOCUMENT SUBMISSION ---------- + + def submit_document( + self, + file_path: str, + mode: Optional[str] = None, + beta_headers: Optional[List[str]] = None, + folder_id: Optional[str] = None, + metadata: Optional[Dict[str, Any]] = None, + ) -> Dict[str, Any]: + """ + Upload a PDF document for processing. The system will automatically process both tree generation and OCR. + Immediately returns a document identifier (`doc_id`) for subsequent operations. + + Args: + file_path (str): Path to the PDF file. + mode (str, optional): Processing mode (e.g., "mcp"). Defaults to None. + beta_headers (List[str], optional): Beta feature headers (e.g., ["block_reference"] + to enable block-level content with bounding boxes). Defaults to None. + folder_id (str, optional): Folder (workspace) ID to assign the document to. Defaults to None. + metadata (dict, optional): Your own JSON-serializable tags for the document; + returned in get_tree/get_ocr responses and list_documents entries. Defaults to None. + + Returns: + dict: {'doc_id': ...} + """ + data = {'if_retrieval': True} + if mode is not None: + data['mode'] = mode + if beta_headers is not None: + data['beta_headers'] = json.dumps(beta_headers) + if folder_id is not None: + data['folder_id'] = folder_id + if metadata is not None: + data['metadata'] = json.dumps(metadata) + + with open(file_path, "rb") as f: + response = requests.post( + f"{self.BASE_URL}/doc/", + headers=self._headers(), + files={'file': f}, + data=data + ) + + if response.status_code != 200: + raise PageIndexAPIError(f"Failed to submit document: {response.text}") + return response.json() + + # ---------- OCR FUNCTIONALITY ---------- + + def get_ocr(self, doc_id: str, format: str = "page") -> Dict[str, Any]: + """ + Get OCR processing status and results. + + Args: + doc_id (str): Document ID. + format (str): Result format. Use 'page' for page-based results, 'node' for node-based results, or 'raw' for concatenated markdown. Defaults to 'page'. + + Returns: + dict: API response with status and, if ready, OCR results. + """ + if format not in ["page", "node", "raw"]: + raise ValueError("Format parameter must be 'page', 'node', or 'raw'") + + response = requests.get( + f"{self.BASE_URL}/doc/{_enc(doc_id)}/?type=ocr&format={format}", + headers=self._headers(), + timeout=30 + ) + if response.status_code != 200: + raise PageIndexAPIError(f"Failed to get OCR result: {response.text}") + return response.json() + + # ---------- TREE GENERATION ---------- + + def get_tree(self, doc_id: str, node_summary: bool = False, + include_text: bool = True) -> Dict[str, Any]: + """ + Get tree generation status and results. + + Args: + doc_id (str): Document ID. + node_summary (bool): Include node summaries (default False). + + Returns: + dict: API response with status and, if ready, tree structure. + """ + response = requests.get( + f"{self.BASE_URL}/doc/{_enc(doc_id)}/?type=tree&summary={node_summary}" + f"&include_text={str(include_text).lower()}", + headers=self._headers(), + timeout=30 + ) + if response.status_code != 200: + raise PageIndexAPIError(f"Failed to get tree result: {response.text}") + return response.json() + + # ---------- RETRIEVAL ---------- + + def submit_query(self, doc_id: str, query: str, thinking: bool = False) -> Dict[str, Any]: + """ + Submit a retrieval query for a specific PageIndex document. + + Args: + doc_id (str): Document ID. + query (str): User question or information need. + thinking (bool, optional): If true, enables deeper retrieval. Default is False. + + Returns: + dict: {'retrieval_id': ...} + """ + payload = { + "doc_id": doc_id, + "query": query, + "thinking": thinking + } + response = requests.post( + f"{self.BASE_URL}/retrieval/", + headers=self._headers(), + json=payload, + timeout=30 + ) + if response.status_code != 200: + raise PageIndexAPIError(f"Failed to submit retrieval: {response.text}") + return response.json() + + def get_retrieval(self, retrieval_id: str) -> Dict[str, Any]: + """ + Get retrieval status and results. + + Args: + retrieval_id (str): Retrieval ID. + + Returns: + dict: Retrieval status and results. + """ + response = requests.get( + f"{self.BASE_URL}/retrieval/{_enc(retrieval_id)}/", + headers=self._headers(), + timeout=30 + ) + if response.status_code != 200: + raise PageIndexAPIError(f"Failed to get retrieval result: {response.text}") + return response.json() + + # ---------- CHAT COMPLETIONS ---------- + + def chat_completions( + self, + messages: List[Dict[str, str]], + stream: bool = False, + doc_id: Optional[Union[str, List[str]]] = None, + temperature: Optional[float] = None, + stream_metadata: bool = False, + enable_citations: bool = False + ) -> Union[Dict[str, Any], Iterator[str], Iterator[Dict[str, Any]]]: + """ + PageIndex Chat Completions. Optionally scoped to specific PageIndex documents. + + Args: + messages (List[Dict[str, str]]): Conversation messages with 'role' and 'content' keys. + stream (bool, optional): Enable streaming responses. Default is False. + doc_id (Optional[Union[str, List[str]]], optional): Document ID(s) to scope the conversation. Can be a single ID or a list of IDs. + temperature (Optional[float], optional): Sampling temperature. Default is None (uses API default). + stream_metadata (bool, optional): If True and stream=True, return raw chunks with metadata instead of just text. Default is False. + enable_citations (bool, optional): Enable citation instructions in responses. Default is False. + + Returns: + Union[Dict[str, Any], Iterator[str], Iterator[Dict[str, Any]]]: + - If stream=False: Complete response dictionary + - If stream=True and stream_metadata=False: Iterator of text content chunks + - If stream=True and stream_metadata=True: Iterator of raw response chunks with metadata + """ + payload = { + "messages": messages, + "stream": stream + } + + if doc_id is not None: + payload["doc_id"] = doc_id + + if temperature is not None: + payload["temperature"] = temperature + + if enable_citations: + payload["enable_citations"] = enable_citations + + response = requests.post( + f"{self.BASE_URL}/chat/completions/", + headers=self._headers(), + json=payload, + stream=stream, + timeout=120 if stream else 300 + ) + + if response.status_code != 200: + raise PageIndexAPIError(f"Failed to get chat completion: {response.text}") + + if stream: + if stream_metadata: + return self._stream_chat_response_raw(response) + else: + return self._stream_chat_response(response) + else: + return response.json() + + def _stream_chat_response(self, response: requests.Response) -> Iterator[str]: + """ + Parse streaming chat completion response. + + Args: + response: Streaming HTTP response + + Yields: + str: Content chunks from the streaming response + """ + try: + for line in response.iter_lines(): + if line: + line = line.decode('utf-8') + if line.startswith('data: '): + data = line[6:] + if data == '[DONE]': + break + + try: + chunk = json.loads(data) + choices = chunk.get("choices") or [{}] + content = choices[0].get("delta", {}).get("content", "") + if content: + yield content + except json.JSONDecodeError: + continue + finally: + response.close() + + def _stream_chat_response_raw(self, response: requests.Response) -> Iterator[Dict[str, Any]]: + """Streaming chat completion with full metadata, including citation events.""" + try: + for line in response.iter_lines(): + if line: + line = line.decode('utf-8') + if line.startswith('data: '): + data = line[6:] + if data == '[DONE]': + break + + try: + chunk = json.loads(data) + yield chunk + except json.JSONDecodeError: + continue + finally: + response.close() + + # ---------- DOCUMENT MANAGEMENT ---------- + + def get_document(self, doc_id: str) -> Dict[str, Any]: + """ + Get document metadata including id, name, description, status, createdAt, and pageNum. + + Args: + doc_id (str): Document ID. + + Returns: + dict: Document metadata containing: + - id (str): Document ID + - name (str): Document name + - description (str): Document description + - status (str): Processing status (e.g., "queued", "processing", "completed", "failed") + - createdAt (str): Creation timestamp in ISO format + - pageNum (int): Number of pages in the document + """ + response = requests.get( + f"{self.BASE_URL}/doc/{_enc(doc_id)}/metadata/", + headers=self._headers(), + timeout=30 + ) + if response.status_code != 200: + raise PageIndexAPIError(f"Failed to get document metadata: {response.text}") + return response.json() + + def delete_document(self, doc_id: str) -> Dict[str, Any]: + """ + Delete a PageIndex document and all its associated data. + + Args: + doc_id (str): Document ID. + + Returns: + dict: API response. + """ + response = requests.delete( + f"{self.BASE_URL}/doc/{_enc(doc_id)}/", + headers=self._headers(), + timeout=30 + ) + if response.status_code != 200: + raise PageIndexAPIError(f"Failed to delete document: {response.text}") + 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]: + """ + List all documents for the authenticated user with pagination. + + Args: + limit (int, optional): Maximum number of documents to return (1-100). Defaults to 50. + offset (int, optional): Number of documents to skip. Defaults to 0. + folder_id (str, optional): Filter by folder (workspace) ID. If provided, only documents + in the specified folder are returned. Defaults to None (all documents). + + Returns: + dict: API response containing: + - documents (List[Dict]): List of document metadata objects (each includes folderId) + - total (int): Total number of documents + - limit (int): Applied limit + - offset (int): Applied offset + """ + if limit < 1 or limit > 100: + raise ValueError("limit must be between 1 and 100") + if offset < 0: + raise ValueError("offset must be non-negative") + + params = {"limit": limit, "offset": offset} + if folder_id is not None: + params["folder_id"] = folder_id + + response = requests.get( + f"{self.BASE_URL}/docs/", + headers=self._headers(), + params=params, + timeout=30 + ) + if response.status_code != 200: + raise PageIndexAPIError(f"Failed to list documents: {response.text}") + return response.json() + + # ---------- FOLDER MANAGEMENT ---------- + + def create_folder( + self, + name: str, + description: Optional[str] = None, + parent_folder_id: Optional[str] = None, + ) -> Dict[str, Any]: + """ + Create a new folder (workspace). + + Args: + name (str): Folder name. + description (str, optional): Folder description. Defaults to None. + parent_folder_id (str, optional): Parent folder ID for nesting. Defaults to None (root level). + + Returns: + dict: Created folder metadata containing: + - folder (dict): Folder info with id, name, description, parent_folder_id, + created_at, file_count, children_count + """ + payload = {"name": name} + if description is not None: + payload["description"] = description + if parent_folder_id is not None: + payload["parent_folder_id"] = parent_folder_id + + response = requests.post( + f"{self.BASE_URL}/folder/", + headers=self._headers(), + json=payload, + timeout=30 + ) + if response.status_code != 200: + raise PageIndexAPIError(f"Failed to create folder: {response.text}") + return response.json() + + def list_folders(self, parent_folder_id: Optional[str] = None) -> Dict[str, Any]: + """ + List folders. + + Args: + parent_folder_id (str, optional): Use "root" for root-level folders only, + a folder ID for subfolders, or omit for all folders. + + Returns: + dict: API response containing: + - folders (List[Dict]): List of folder metadata objects + - total (int): Total number of folders + """ + params = {} + if parent_folder_id is not None: + params["parent_folder_id"] = parent_folder_id + + response = requests.get( + f"{self.BASE_URL}/folders/", + headers=self._headers(), + params=params, + timeout=30 + ) + if response.status_code != 200: + raise PageIndexAPIError(f"Failed to list folders: {response.text}") + return response.json() diff --git a/pageindex/errors.py b/pageindex/errors.py new file mode 100644 index 000000000..e460a956b --- /dev/null +++ b/pageindex/errors.py @@ -0,0 +1,2 @@ +class PageIndexAPIError(Exception): + pass diff --git a/pageindex/local_api.py b/pageindex/local_api.py new file mode 100644 index 000000000..0e9f682c8 --- /dev/null +++ b/pageindex/local_api.py @@ -0,0 +1,325 @@ +"""Local implementation of the PageIndex SDK surface.""" +from __future__ import annotations + +import asyncio +import json +import logging +import os +import uuid +from concurrent.futures import ThreadPoolExecutor +from datetime import datetime, timezone +from typing import Any + +from .errors import PageIndexAPIError +from .local_store import DocStore + +logger = logging.getLogger(__name__) + + +def _now_iso() -> str: + """Naive UTC, millisecond precision.""" + now = datetime.now(timezone.utc).replace(tzinfo=None) + return now.replace(microsecond=now.microsecond // 1000 * 1000).isoformat() + + +def _run_indexer(func, *args, **kwargs): + try: + asyncio.get_running_loop() + except RuntimeError: + return func(*args, **kwargs) + with ThreadPoolExecutor(max_workers=1) as pool: + return pool.submit(func, *args, **kwargs).result() + + +class LocalAPI: + """Backs PageIndexClient's local mode. One instance per client.""" + + def __init__(self, storage_path: str, model: str, summary_model: str, + retrieve_model: str): + self._store = DocStore(storage_path) + self._model = model + self._summary_model = summary_model + self._retrieve_model = retrieve_model + from .utils import ConfigLoader + self._config_loader = ConfigLoader() + + # ── indexing ── + + def submit_document( + self, + file_path: str, + mode: str | None = None, + beta_headers: list[str] | None = None, + folder_id: str | None = None, + metadata: dict | None = None, + ) -> dict[str, Any]: + if beta_headers is not None: + raise PageIndexAPIError( + "Failed to submit document: beta_headers is not supported in local mode." + ) + if folder_id is not None: + raise PageIndexAPIError( + "Failed to submit document: folders are not supported in local mode." + ) + if metadata is not None: + if not isinstance(metadata, dict): + raise PageIndexAPIError( + "Failed to submit document: metadata must be a dict." + ) + try: + json.dumps(metadata) + except (TypeError, ValueError) as e: + raise PageIndexAPIError( + f"Failed to submit document: metadata must be valid JSON. {e}" + ) from e + if mode not in (None, "standard", "flash"): + raise PageIndexAPIError( + f"Failed to submit document: unknown local processing mode {mode!r}. " + "Supported: None or 'standard' for standard indexing, or 'flash'." + ) + file_path = os.path.abspath(os.path.expanduser(str(file_path))) + if not os.path.isfile(file_path): + raise FileNotFoundError(f"No such file: {file_path}") + if not file_path.lower().endswith(".pdf"): + raise PageIndexAPIError( + "Failed to submit document: only PDF files are supported in local mode." + ) + + try: + page_texts = self._extract_page_texts(file_path) + except PageIndexAPIError: + raise + except Exception as e: + raise PageIndexAPIError( + f"Failed to submit document: could not read PDF: {e}" + ) from e + if not any(text.strip() for text in page_texts): + raise PageIndexAPIError( + "Failed to submit document: PDF has no content. All pages are blank." + ) + + try: + if mode == "flash": + structure, description = _run_indexer( + self._index_flash, file_path, page_texts + ) + else: + structure, description = _run_indexer( + self._index_standard, file_path, page_texts + ) + except PageIndexAPIError: + raise + except Exception as e: + raise PageIndexAPIError(f"Failed to submit document: {e}") from e + + doc_id = "pi-" + uuid.uuid4().hex + meta = { + "id": doc_id, + "name": os.path.basename(file_path), + "description": description, + "status": "completed", + "createdAt": _now_iso(), + "pageNum": len(page_texts), + "folderId": None, + "metadata": metadata, + "mode": mode or "standard", + } + pages = [{"page_index": i + 1, "markdown": text} + for i, text in enumerate(page_texts)] + from .utils import remove_fields + self._store.save_document( + doc_id, meta, remove_fields(structure, fields=["text"]), pages) + return {"doc_id": doc_id} + + @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] + + def _index_standard(self, file_path: str, page_texts: list[str]) -> tuple[list, str | None]: + from .page_index_classic import page_index_main + import litellm + page_list = [(text, litellm.token_counter(model=self._model, text=text)) + for text in page_texts] + opt = self._config_loader.load({ + "model": self._model, + "summary_model": self._summary_model, + "if_add_node_id": "yes", + "if_add_node_summary": "yes", + "if_add_node_text": "yes", + "if_add_doc_description": "yes", + }) + result = page_index_main(file_path, opt, logger=logger, page_list=page_list) + structure = result.get("structure") or [] + if not structure: + raise PageIndexAPIError( + "Failed to submit document: standard indexing produced no structure." + ) + return structure, result.get("doc_description") + + def _index_flash(self, file_path: str, page_texts: list[str]) -> tuple[list, str | None]: + from .flash import page_index_flash + from .utils import (add_node_text, create_clean_structure_for_description, + generate_doc_description, write_node_id) + result = page_index_flash(file_path, summary=True, + summary_model=self._summary_model) + structure = result.get("structure", []) + if not structure: + raise PageIndexAPIError( + "Failed to submit document: PageIndex Flash could not extract " + "a structure from this PDF." + ) + write_node_id(structure) + add_node_text(structure, [(text, 0) for text in page_texts]) + description = generate_doc_description( + create_clean_structure_for_description(structure), + model=self._summary_model, + ) + return structure, description + + # ── tree / ocr ── + + def _load_tree_with_text(self, doc_id: str, error_prefix: str) -> list: + from .utils import add_node_text + structure = self._require_data( + self._store.get_tree(doc_id), error_prefix) + pages = self._require_pages(doc_id, error_prefix) + pdf_pages = [(p.get("markdown", ""), 0) for p in pages] + add_node_text(structure, pdf_pages) + return structure + + def get_tree(self, doc_id: str, node_summary: bool = False, + include_text: bool = True) -> dict[str, Any]: + meta = self._require_doc(doc_id, "Failed to get tree result") + if include_text: + structure = self._load_tree_with_text(doc_id, "Failed to get tree result") + else: + structure = self._require_data( + self._store.get_tree(doc_id), "Failed to get tree result") + result = [_format_tree_node(node, node_summary) for node in structure] + return self._completed_envelope(doc_id, result, meta) + + def get_ocr(self, doc_id: str, format: str = "page") -> dict[str, Any]: + if format not in ["page", "node", "raw"]: + raise ValueError("Format parameter must be 'page', 'node', or 'raw'") + meta = self._require_doc(doc_id, "Failed to get OCR result") + if format == "node": + result: Any = [] + def _walk(nodes, level): + for node in nodes: + result.append({ + "title": node.get("title", ""), + "level": level, + "page_index": node.get("start_index"), + "text": node.get("text", ""), + }) + _walk(node.get("nodes") or [], level + 1) + _walk(self._load_tree_with_text(doc_id, + "Failed to get OCR result"), 1) + else: + pages = self._require_pages(doc_id, "Failed to get OCR result") + if format == "page": + result = pages + else: # raw + result = "\n\n".join(p.get("markdown", "") for p in pages) + return self._completed_envelope(doc_id, result, meta) + + @staticmethod + def _require_data(data, error_prefix: str): + if data is None: + raise PageIndexAPIError(f"{error_prefix}: stored document data is unreadable.") + return data + + def _require_pages(self, doc_id: str, error_prefix: str) -> list: + pages = self._require_data(self._store.get_pages(doc_id), error_prefix) + if not pages: + raise PageIndexAPIError( + f"{error_prefix}: stored document has no page content.") + return pages + + def _completed_envelope(self, doc_id: str, result, meta: dict) -> dict[str, Any]: + return { + "doc_id": doc_id, + "status": "completed", + "retrieval_ready": True, + "result": result, + "metadata": meta.get("metadata"), + "features": {}, + } + + # ── document management ── + + def _require_doc(self, doc_id: str, error_prefix: str) -> dict: + meta = self._store.get_meta(doc_id) + if meta is None: + raise PageIndexAPIError(f"{error_prefix}: Document not found.") + return meta + + def get_document(self, doc_id: str) -> dict[str, Any]: + meta = self._store.get_meta(doc_id) + if meta is None: + raise PageIndexAPIError("Failed to get document metadata: Document not found") + return {key: meta.get(key) for key in + ("id", "name", "description", "status", "createdAt", "pageNum", "folderId")} + + def delete_document(self, doc_id: str) -> dict[str, Any]: + if not self._store.delete_document(doc_id): + raise PageIndexAPIError("Failed to delete document: Document not found.") + return {"message": "Document deleted successfully."} + + def list_documents( + self, + limit: int = 50, + offset: int = 0, + folder_id: str | None = None, + ) -> dict[str, Any]: + if limit < 1 or limit > 100: + raise ValueError("limit must be between 1 and 100") + if offset < 0: + raise ValueError("offset must be non-negative") + if folder_id is not None: + raise PageIndexAPIError( + "Failed to list documents: folders are not supported in local mode." + ) + metas = sorted(self._store.list_metas(), key=lambda m: m.get("id") or "") + metas.sort(key=lambda m: m.get("createdAt") or "", reverse=True) + documents = [{ + "id": m.get("id"), + "name": m.get("name"), + "description": m.get("description"), + "status": m.get("status"), + "createdAt": m.get("createdAt"), + "pageNum": m.get("pageNum", 0), + "folderId": None, + "metadata": m.get("metadata"), + "features": {}, + } for m in metas[offset:offset + limit]] + return { + "documents": documents, + "total": len(metas), + "limit": limit, + "offset": offset, + } + + +def _format_tree_node(node: dict, node_summary: bool) -> dict: + children = node.get("nodes") or [] + out = { + "title": node.get("title", ""), + "node_id": node.get("node_id"), + "page_index": node.get("start_index"), + } + if node_summary: + summary = node.get("summary") + if summary is not None: + if children: + out["prefix_summary"] = summary + else: + out["summary"] = summary + if "text" in node: + out["text"] = node["text"] + if children: + out["nodes"] = [_format_tree_node(child, node_summary) for child in children] + return out diff --git a/pageindex/local_store.py b/pageindex/local_store.py new file mode 100644 index 000000000..37108f93c --- /dev/null +++ b/pageindex/local_store.py @@ -0,0 +1,164 @@ +"""On-disk document store behind PageIndexClient's local mode.""" +from __future__ import annotations + +import json +import logging +import os +import shutil +import uuid +from pathlib import Path + +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: + json.dump(data, f, ensure_ascii=False) + f.flush() + os.fsync(f.fileno()) + os.replace(tmp, path) + except BaseException: + tmp.unlink(missing_ok=True) + raise + + +def _read_json(path: Path): + try: + with open(path, "r", encoding="utf-8") as f: + return json.load(f) + except (FileNotFoundError, NotADirectoryError, IsADirectoryError, + PermissionError): + return None + except ValueError: + logger.warning("Unreadable JSON at %s; treating it as absent", path) + return None + + +def _is_safe_id(value: str) -> bool: + return ( + isinstance(value, str) + and value not in ("", ".", "..") + and os.path.basename(value) == value + and "\\" not in value + ) + + +def _is_valid_meta(meta, doc_id: str) -> bool: + if not isinstance(meta, dict) or meta.get("id") != doc_id: + return False + page_num = meta.get("pageNum") + return ( + isinstance(meta.get("name"), str) + and (meta.get("description") is None + or isinstance(meta.get("description"), str)) + and isinstance(meta.get("status"), str) + and isinstance(meta.get("createdAt"), str) + and isinstance(page_num, int) + and not isinstance(page_num, bool) + and page_num >= 0 + and (meta.get("folderId") is None + or isinstance(meta.get("folderId"), str)) + and (meta.get("metadata") is None + or isinstance(meta.get("metadata"), dict)) + and (meta.get("mode") is None or isinstance(meta.get("mode"), str)) + ) + + +class DocStore: + def __init__(self, storage_dir: str): + self._root = Path(storage_dir).expanduser() + self._docs = self._root / "docs" + self._manifest = self._root / "manifest.json" + + def _doc_dir(self, doc_id: str) -> Path | None: + if not _is_safe_id(doc_id): + return None + return self._docs / doc_id + + # ── manifest cache ── + def _read_manifest(self) -> dict: + data = _read_json(self._manifest) + docs = data.get("docs") if isinstance(data, dict) else None + return docs if isinstance(docs, dict) else {} + + def _write_manifest(self, docs: dict) -> None: + try: + _write_json_atomic(self._manifest, {"docs": docs}) + except OSError: + pass + + # ── documents ── + def save_document(self, doc_id: str, meta: dict, tree: list, pages: list) -> None: + doc_dir = self._doc_dir(doc_id) + if doc_dir is None: + raise ValueError(f"Invalid doc_id: {doc_id!r}") + doc_dir.mkdir(parents=True, exist_ok=True) + _write_json_atomic(doc_dir / "tree.json", tree) + _write_json_atomic(doc_dir / "pages.json", pages) + _write_json_atomic(doc_dir / "doc.json", meta) + manifest = self._read_manifest() + manifest[doc_id] = meta + self._write_manifest(manifest) + + def _read_doc_file(self, doc_id: str, name: str): + doc_dir = self._doc_dir(doc_id) + if doc_dir is None or not (doc_dir / "doc.json").is_file(): + return None + return _read_json(doc_dir / name) + + def get_meta(self, doc_id: str) -> dict | None: + doc_dir = self._doc_dir(doc_id) + if doc_dir is None or not (doc_dir / "doc.json").is_file(): + return None + meta = _read_json(doc_dir / "doc.json") + if not _is_valid_meta(meta, doc_id): + meta = self._read_manifest().get(doc_id) + return meta if _is_valid_meta(meta, doc_id) else None + + def get_tree(self, doc_id: str) -> list | None: + return self._read_doc_file(doc_id, "tree.json") + + def get_pages(self, doc_id: str) -> list | None: + return self._read_doc_file(doc_id, "pages.json") + + def list_metas(self) -> list[dict]: + if not self._docs.is_dir(): + return [] + with os.scandir(self._docs) as entries: + dir_names = {entry.name for entry in entries + if entry.is_dir() and _is_safe_id(entry.name)} + cached = self._read_manifest() + fresh = {} + for name in dir_names: + if not (self._docs / name / "doc.json").is_file(): + continue + meta = cached.get(name) + if not _is_valid_meta(meta, name): + meta = _read_json(self._docs / name / "doc.json") + if _is_valid_meta(meta, name): + fresh[name] = meta + if fresh != cached: + self._write_manifest(fresh) + return list(fresh.values()) + + def delete_document(self, doc_id: str) -> bool: + doc_dir = self._doc_dir(doc_id) + if doc_dir is None: + return False + try: + (doc_dir / "doc.json").unlink() + existed = True + except (FileNotFoundError, NotADirectoryError): + existed = False + except OSError: + if not (doc_dir / "doc.json").is_dir(): + raise + existed = False + if doc_dir.is_dir(): + shutil.rmtree(doc_dir, ignore_errors=True) + manifest = self._read_manifest() + if manifest.pop(doc_id, None) is not None: + self._write_manifest(manifest) + return existed diff --git a/pageindex/page_index.py b/pageindex/page_index_classic.py similarity index 99% rename from pageindex/page_index.py rename to pageindex/page_index_classic.py index c0b3ea935..5b846e6d0 100644 --- a/pageindex/page_index.py +++ b/pageindex/page_index_classic.py @@ -1100,19 +1100,21 @@ async def verify_toc(page_list, list_result, start_index=1, N=None, model=None): check_title_appearance(item, page_list, start_index, model) for item in indexed_sample_list ] - results = await asyncio.gather(*tasks) + results = await asyncio.gather(*tasks, return_exceptions=True) - # Process results + # Process results (skip exceptions from failed LLM calls) correct_count = 0 incorrect_results = [] for result in results: + if isinstance(result, BaseException): + continue if result['answer'] == 'yes': correct_count += 1 else: incorrect_results.append(result) - + # Calculate accuracy - checked_count = len(results) + checked_count = sum(1 for r in results if not isinstance(r, BaseException)) accuracy = correct_count / checked_count if checked_count > 0 else 0 print(f"accuracy: {accuracy*100:.2f}%") return accuracy, incorrect_results @@ -1229,18 +1231,19 @@ async def tree_parser(page_list, opt, doc=None, logger=None): return toc_tree -def page_index_main(doc, opt=None): - logger = JsonLogger(doc) - +def page_index_main(doc, opt=None, logger=None, page_list=None): + logger = logger or JsonLogger(doc) + is_valid_pdf = ( - (isinstance(doc, str) and os.path.isfile(doc) and doc.lower().endswith(".pdf")) or + (isinstance(doc, str) and os.path.isfile(doc) and doc.lower().endswith(".pdf")) or isinstance(doc, BytesIO) ) if not is_valid_pdf: raise ValueError("Unsupported input type. Expected a PDF file path or BytesIO object.") - print('Parsing PDF...') - page_list = get_page_tokens(doc, model=opt.model) + if page_list is None: + print('Parsing PDF...') + page_list = get_page_tokens(doc, model=opt.model) logger.info({'total_page_number': len(page_list)}) logger.info({'total_token': sum([page[1] for page in page_list])}) diff --git a/pageindex/retrieve.py b/pageindex/retrieve.py deleted file mode 100644 index 55c38509c..000000000 --- a/pageindex/retrieve.py +++ /dev/null @@ -1,137 +0,0 @@ -import json -import PyPDF2 - -try: - from .utils import get_number_of_pages, remove_fields -except ImportError: - from utils import get_number_of_pages, remove_fields - - -# ── Helpers ────────────────────────────────────────────────────────────────── - -def _parse_pages(pages: str) -> list[int]: - """Parse a pages string like '5-7', '3,8', or '12' into a sorted list of ints.""" - result = [] - for part in pages.split(','): - part = part.strip() - if '-' in part: - start, end = int(part.split('-', 1)[0].strip()), int(part.split('-', 1)[1].strip()) - if start > end: - raise ValueError(f"Invalid range '{part}': start must be <= end") - result.extend(range(start, end + 1)) - else: - result.append(int(part)) - return sorted(set(result)) - - -def _count_pages(doc_info: dict) -> int: - """Return total page count for a PDF document.""" - if doc_info.get('page_count'): - return doc_info['page_count'] - if doc_info.get('pages'): - return len(doc_info['pages']) - return get_number_of_pages(doc_info['path']) - - -def _get_pdf_page_content(doc_info: dict, page_nums: list[int]) -> list[dict]: - """Extract text for specific PDF pages (1-indexed). Prefer cached pages, fallback to PDF.""" - cached_pages = doc_info.get('pages') - if cached_pages: - page_map = {p['page']: p['content'] for p in cached_pages} - return [ - {'page': p, 'content': page_map[p]} - for p in page_nums if p in page_map - ] - path = doc_info['path'] - with open(path, 'rb') as f: - pdf_reader = PyPDF2.PdfReader(f) - total = len(pdf_reader.pages) - valid_pages = [p for p in page_nums if 1 <= p <= total] - return [ - {'page': p, 'content': pdf_reader.pages[p - 1].extract_text() or ''} - for p in valid_pages - ] - - -def _get_md_page_content(doc_info: dict, page_nums: list[int]) -> list[dict]: - """ - For Markdown documents, 'pages' are line numbers. - Find nodes whose line_num falls within [min(page_nums), max(page_nums)] and return their text. - """ - min_line, max_line = min(page_nums), max(page_nums) - results = [] - seen = set() - - def _traverse(nodes): - for node in nodes: - ln = node.get('line_num') - if ln and min_line <= ln <= max_line 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 - - -# ── Tool functions ──────────────────────────────────────────────────────────── - -def get_document(documents: dict, doc_id: str) -> str: - """Return JSON with document metadata: doc_id, doc_name, doc_description, type, status, page_count (PDF) or line_count (Markdown).""" - doc_info = documents.get(doc_id) - if not doc_info: - return json.dumps({'error': f'Document {doc_id} not found'}) - result = { - 'doc_id': doc_id, - 'doc_name': doc_info.get('doc_name', ''), - 'doc_description': doc_info.get('doc_description', ''), - 'type': doc_info.get('type', ''), - 'status': 'completed', - } - if doc_info.get('type') == 'pdf': - result['page_count'] = _count_pages(doc_info) - else: - result['line_count'] = doc_info.get('line_count', 0) - return json.dumps(result) - - -def get_document_structure(documents: dict, doc_id: str) -> str: - """Return tree structure JSON with text fields removed (saves tokens).""" - doc_info = documents.get(doc_id) - if not doc_info: - return json.dumps({'error': f'Document {doc_id} not found'}) - structure = doc_info.get('structure', []) - structure_no_text = remove_fields(structure, fields=['text']) - return json.dumps(structure_no_text, ensure_ascii=False) - - -def get_page_content(documents: dict, doc_id: str, pages: str) -> str: - """ - Retrieve page content for a document. - - pages format: '5-7', '3,8', or '12' - For PDF: pages are physical page numbers (1-indexed). - For Markdown: pages are line numbers corresponding to node headers. - - Returns JSON list of {'page': int, 'content': str}. - """ - doc_info = documents.get(doc_id) - if not doc_info: - return json.dumps({'error': f'Document {doc_id} not found'}) - - try: - page_nums = _parse_pages(pages) - except (ValueError, AttributeError) as e: - return json.dumps({'error': f'Invalid pages format: {pages!r}. Use "5-7", "3,8", or "12". Error: {e}'}) - - try: - if doc_info.get('type') == 'pdf': - content = _get_pdf_page_content(doc_info, page_nums) - else: - content = _get_md_page_content(doc_info, page_nums) - except Exception as e: - return json.dumps({'error': f'Failed to read page content: {e}'}) - - return json.dumps(content, ensure_ascii=False) diff --git a/pageindex/utils.py b/pageindex/utils.py index 92fc46d85..97f60a942 100644 --- a/pageindex/utils.py +++ b/pageindex/utils.py @@ -7,7 +7,6 @@ import PyPDF2 import copy import asyncio -import pymupdf from io import BytesIO from dotenv import load_dotenv load_dotenv() @@ -31,6 +30,12 @@ def count_tokens(text, model=None): return litellm.token_counter(model=model, text=text) +def _strip_prefix(s, prefix): + if s.startswith(prefix): + return s[len(prefix):] + return s + + def _is_openai_model(model): """Models without a provider prefix (no '/') use the openai SDK directly. For other providers, use 'provider/model' format (e.g. 'anthropic/claude-sonnet-4-6').""" @@ -57,18 +62,19 @@ def _is_unrecoverable(exc: Exception) -> bool: def llm_completion(model, prompt, chat_history=None, return_finish_reason=False): use_openai_sdk = _is_openai_model(model) if model: - model = model.removeprefix("litellm/") + model = _strip_prefix(model, "litellm/") if use_openai_sdk: - model = model.removeprefix("openai/") + model = _strip_prefix(model, "openai/") max_retries = 10 messages = list(chat_history) + [{"role": "user", "content": prompt}] if chat_history else [{"role": "user", "content": prompt}] + if use_openai_sdk: + global _openai_sync_client + if _openai_sync_client is None: + import openai + _openai_sync_client = openai.OpenAI(max_retries=0) for i in range(max_retries): try: if use_openai_sdk: - global _openai_sync_client - if _openai_sync_client is None: - import openai - _openai_sync_client = openai.OpenAI(max_retries=0) response = _openai_sync_client.chat.completions.create( model=model, messages=messages, @@ -94,27 +100,27 @@ def llm_completion(model, prompt, chat_history=None, return_finish_reason=False) if i < max_retries - 1: time.sleep(1) else: - logging.error('Max retries reached for prompt: ' + prompt) - if return_finish_reason: - return "", "error" - return "" + raise RuntimeError( + f"LLM completion failed after {max_retries} retries" + ) from e async def llm_acompletion(model, prompt): use_openai_sdk = _is_openai_model(model) if model: - model = model.removeprefix("litellm/") + model = _strip_prefix(model, "litellm/") if use_openai_sdk: - model = model.removeprefix("openai/") + model = _strip_prefix(model, "openai/") max_retries = 10 messages = [{"role": "user", "content": prompt}] + if use_openai_sdk: + global _openai_async_client + if _openai_async_client is None: + import openai + _openai_async_client = openai.AsyncOpenAI(max_retries=0) for i in range(max_retries): try: if use_openai_sdk: - global _openai_async_client - if _openai_async_client is None: - import openai - _openai_async_client = openai.AsyncOpenAI(max_retries=0) response = await _openai_async_client.chat.completions.create( model=model, messages=messages, @@ -136,10 +142,11 @@ async def llm_acompletion(model, prompt): if i < max_retries - 1: await asyncio.sleep(1) else: - logging.error('Max retries reached for prompt: ' + prompt) - return "" - - + raise RuntimeError( + f"LLM completion failed after {max_retries} retries" + ) from e + + def get_json_content(response): start_idx = response.find("```json") if start_idx != -1: @@ -180,7 +187,7 @@ def extract_json(content): # Remove any trailing commas before closing brackets/braces json_content = json_content.replace(',]', ']').replace(',}', '}') return json.loads(json_content) - except: + except Exception: logging.error("Failed to parse JSON even after cleanup") return {} except Exception as e: @@ -454,6 +461,7 @@ def get_page_tokens(pdf_path, model=None, pdf_parser="PyPDF2"): page_list.append((page_text, token_length)) return page_list elif pdf_parser == "PyMuPDF": + import pymupdf if isinstance(pdf_path, BytesIO): pdf_stream = pdf_path doc = pymupdf.open(stream=pdf_stream, filetype="pdf") @@ -471,12 +479,16 @@ def get_page_tokens(pdf_path, model=None, pdf_parser="PyPDF2"): def get_text_of_pdf_pages(pdf_pages, start_page, end_page): + if start_page is None or end_page is None: + return "" text = "" for page_num in range(start_page-1, end_page): text += pdf_pages[page_num][0] return text def get_text_of_pdf_pages_with_labels(pdf_pages, start_page, end_page): + if start_page is None or end_page is None: + return "" text = "" for page_num in range(start_page-1, end_page): text += f"\n{pdf_pages[page_num][0]}\n\n" @@ -522,12 +534,14 @@ def clean_structure_post(data): clean_structure_post(section) return data -def remove_fields(data, fields=['text']): +def remove_fields(data, fields=['text'], max_len=None): if isinstance(data, dict): - return {k: remove_fields(v, fields) + return {k: remove_fields(v, fields, max_len) for k, v in data.items() if k not in fields} elif isinstance(data, list): - return [remove_fields(item, fields) for item in data] + return [remove_fields(item, fields, max_len) for item in data] + elif isinstance(data, str): + return data[:max_len] + '...' if max_len is not None and len(data) > max_len else data return data def print_toc(tree, indent=0): @@ -648,10 +662,15 @@ async def generate_node_summary(node, model=None): async def generate_summaries_for_structure(structure, model=None): nodes = structure_to_list(structure) tasks = [generate_node_summary(node, model=model) for node in nodes] - summaries = await asyncio.gather(*tasks) + summaries = await asyncio.gather(*tasks, return_exceptions=True) for node, summary in zip(nodes, summaries): - node['summary'] = summary + node['summary'] = "" if isinstance(summary, BaseException) else summary + 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)" + ) return structure @@ -819,12 +838,27 @@ async def parent_summary(node): async def visit(node): children = node.get('nodes') or [] if children: - await asyncio.gather(*(visit(child) for child in children)) + await asyncio.gather(*(visit(child) for child in children), + return_exceptions=True) if node.get('summary'): return - node['summary'] = await (parent_summary(node) if children else leaf_summary(node)) + try: + node['summary'] = await (parent_summary(node) if children else leaf_summary(node)) + except Exception: + node['summary'] = "" + + await asyncio.gather(*(visit(root) for root in structure), + return_exceptions=True) + + 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): + raise RuntimeError( + "Summary generation failed for all nodes " + "(check LLM credentials and model availability)" + ) - await asyncio.gather(*(visit(root) for root in structure)) strip_internal_keys(structure) return structure @@ -860,8 +894,10 @@ def generate_doc_description(structure, model=None): Directly return the description, do not include any other text. """ - response = llm_completion(model, prompt) - return response + try: + return llm_completion(model, prompt) + except RuntimeError: + return "" def reorder_dict(data, key_order): @@ -951,25 +987,42 @@ def load(self, user_opt=None) -> config: merged = {**self._default_dict, **user_dict} return config(**merged) -def create_node_mapping(tree): - """Create a flat dict mapping node_id to node for quick lookup.""" +def create_node_mapping(tree, include_page_ranges=False, max_page=None): + """Map node_id to node; with include_page_ranges, to {"node", "start_index", + "end_index"} (end = next node's page_index, or max_page for the last node).""" + def get_all_nodes(tree): + if isinstance(tree, dict): + return [tree] + [node for child in tree.get('nodes', []) for node in get_all_nodes(child)] + elif isinstance(tree, list): + return [node for item in tree for node in get_all_nodes(item)] + return [] + + all_nodes = get_all_nodes(tree) + if not include_page_ranges: + return {node["node_id"]: node for node in all_nodes if node.get("node_id")} mapping = {} - def _traverse(nodes): - for node in nodes: - if node.get('node_id'): - mapping[node['node_id']] = node - if node.get('nodes'): - _traverse(node['nodes']) - _traverse(tree) + for i, node in enumerate(all_nodes): + if node.get("node_id"): + end_page = all_nodes[i + 1].get("page_index") if i + 1 < len(all_nodes) else max_page + mapping[node["node_id"]] = { + "node": node, + "start_index": node["page_index"], + "end_index": end_page, + } return mapping -def print_tree(tree, indent=0): +def print_tree(tree, exclude_fields=None, indent=0): + """Outline view; passing exclude_fields gives the 0.2.8 pprint view.""" + if exclude_fields is not None: + from pprint import pprint + pprint(remove_fields(tree, exclude_fields, max_len=40), sort_dicts=False, width=100) + return for node in tree: summary = node.get('summary') or node.get('prefix_summary', '') summary_str = f" — {summary[:60]}..." if summary else "" print(' ' * indent + f"[{node.get('node_id', '?')}] {node.get('title', '')}{summary_str}") if node.get('nodes'): - print_tree(node['nodes'], indent + 1) + print_tree(node['nodes'], indent=indent + 1) def print_wrapped(text, width=100): for line in text.splitlines(): diff --git a/pyproject.toml b/pyproject.toml new file mode 100644 index 000000000..deac66be3 --- /dev/null +++ b/pyproject.toml @@ -0,0 +1,53 @@ +[tool.poetry] +name = "pageindex" +version = "0.2.9" +description = "Python SDK for PageIndex — reasoning-based, vectorless document retrieval, cloud and local" +readme = "README.md" +license = "MIT" +authors = ["Ray "] +classifiers = [ + "Development Status :: 3 - Alpha", + "Intended Audience :: Developers", + "License :: OSI Approved :: MIT License", + "Programming Language :: Python :: 3", + "Programming Language :: Python :: 3.7", + "Programming Language :: Python :: 3.8", + "Programming Language :: Python :: 3.9", + "Programming Language :: Python :: 3.10", + "Programming Language :: Python :: 3.11", + "Programming Language :: Python :: 3.12", + "Programming Language :: Python :: 3.13", + "Topic :: Scientific/Engineering :: Artificial Intelligence", +] +keywords = ["rag", "document", "retrieval", "llm", "pageindex", "vectorless"] +packages = [{ include = "pageindex" }] +include = [ + { path = "pageindex/config.yaml", format = ["sdist", "wheel"] }, + { path = "pageindex/flash/data/*.json", format = ["sdist", "wheel"] }, +] +exclude = ["pageindex/flash/assets"] + +[tool.poetry.dependencies] +python = ">=3.7" +requests = ">=2.28.0" +openai = ">=1.70.0" +litellm = ">=1.84.0" +PyPDF2 = ">=3.0.0" +pypdfium2 = ">=4.30.0" +sortedcontainers = ">=2.4.0" +regex = ">=2024.0.0" +python-dotenv = ">=1.0.0" +pyyaml = ">=6.0" + +[tool.poetry.group.dev.dependencies] +pytest = ">=7.0" + +[tool.poetry.urls] +Repository = "https://github.com/VectifyAI/PageIndex" +Homepage = "https://pageindex.ai" +Documentation = "https://docs.pageindex.ai" +Issues = "https://github.com/VectifyAI/PageIndex/issues" + +[build-system] +requires = ["poetry-core"] +build-backend = "poetry.core.masonry.api" diff --git a/requirements.txt b/requirements.txt index c0b76deb1..5fd4f2e4e 100644 --- a/requirements.txt +++ b/requirements.txt @@ -1,6 +1,8 @@ litellm==1.84.0 -# openai-agents # optional: required for examples/agentic_vectorless_rag_demo.py -pymupdf==1.26.4 +openai>=1.70.0 +requests>=2.28.0 +# openai-agents # optional +# pymupdf # optional PyPDF2==3.0.1 pypdfium2==4.30.0 python-dotenv==1.2.2 diff --git a/tests/conftest.py b/tests/conftest.py new file mode 100644 index 000000000..40a4a9199 --- /dev/null +++ b/tests/conftest.py @@ -0,0 +1,52 @@ +import pytest + + +@pytest.fixture(autouse=True) +def _llm_key(monkeypatch): + """Deterministic key presence for every test; missing-key tests delenv.""" + monkeypatch.setenv("OPENAI_API_KEY", "test-key") + + +def build_pdf(page_texts): + """Build a minimal, uncompressed PDF (one Helvetica line per page) whose + text PyPDF2 can extract. Returns the PDF file bytes.""" + n = len(page_texts) + objects = [] + kids = " ".join(f"{3 + i} 0 R" for i in range(n)) + objects.append(b"<< /Type /Catalog /Pages 2 0 R >>") + objects.append(f"<< /Type /Pages /Kids [{kids}] /Count {n} >>".encode()) + font_obj = 3 + 2 * n + for i in range(n): + objects.append( + f"<< /Type /Page /Parent 2 0 R /MediaBox [0 0 612 792] " + f"/Resources << /Font << /F1 {font_obj} 0 R >> >> " + f"/Contents {3 + n + i} 0 R >>".encode() + ) + for text in page_texts: + safe = text.replace("\\", r"\\").replace("(", r"\(").replace(")", r"\)") + stream = f"BT /F1 12 Tf 72 720 Td ({safe}) Tj ET".encode() + objects.append(b"<< /Length %d >>\nstream\n%s\nendstream" % (len(stream), stream)) + objects.append(b"<< /Type /Font /Subtype /Type1 /BaseFont /Helvetica >>") + + out = bytearray(b"%PDF-1.4\n") + offsets = [] + for num, body in enumerate(objects, start=1): + offsets.append(len(out)) + out += b"%d 0 obj\n" % num + body + b"\nendobj\n" + xref_pos = len(out) + out += b"xref\n0 %d\n" % (len(objects) + 1) + out += b"0000000000 65535 f \n" + for offset in offsets: + out += b"%010d 00000 n \n" % offset + out += (b"trailer\n<< /Size %d /Root 1 0 R >>\nstartxref\n%d\n%%%%EOF\n" + % (len(objects) + 1, xref_pos)) + return bytes(out) + + +@pytest.fixture +def sample_pdf(tmp_path): + """A 2-page PDF with known, extractable text.""" + path = tmp_path / "sample.pdf" + path.write_bytes(build_pdf(["Hello page one about apples", + "Second page about bananas"])) + return str(path) diff --git a/tests/test_client.py b/tests/test_client.py new file mode 100644 index 000000000..50b7f5178 --- /dev/null +++ b/tests/test_client.py @@ -0,0 +1,712 @@ +"""SDK surface tests: PageIndexClient in local and cloud mode.""" +import asyncio +import importlib +import json +import re +import shutil +import types + +import pytest + +import pageindex.flash +import pageindex.utils +from pageindex import PageIndexClient, PageIndexAPIError + +page_index_module = importlib.import_module("pageindex.page_index_classic") + + +STRUCTURE = [ + { + "title": "Root Section", "node_id": "0000", + "start_index": 1, "end_index": 2, + "summary": "root summary", "text": "root text", + "nodes": [ + {"title": "Child Section", "node_id": "0001", + "start_index": 2, "end_index": 2, + "summary": "child summary", "text": "child text"}, + ], + }, +] + + +@pytest.fixture +def local_client(tmp_path): + return PageIndexClient(storage_path=str(tmp_path / "store")) + + +@pytest.fixture +def indexed_doc(local_client, sample_pdf, monkeypatch): + """A document indexed through a stubbed standard pipeline.""" + def fake_page_index_main(doc, opt=None, logger=None, page_list=None): + assert opt.if_add_node_summary == "yes" + assert opt.if_add_node_text == "yes" + assert logger is not None + assert page_list is not None + assert all(isinstance(t, tuple) and len(t) == 2 for t in page_list) + return {"doc_name": "sample.pdf", + "doc_description": "A test document.", + "structure": json.loads(json.dumps(STRUCTURE))} + monkeypatch.setattr(page_index_module, "page_index_main", fake_page_index_main) + return local_client.submit_document(sample_pdf)["doc_id"] + + +# ── constructor ── + +def test_empty_api_key_raises(): + with pytest.raises(PageIndexAPIError, match="empty string"): + PageIndexClient(api_key="") + + +def test_cloud_rejects_local_args(): + with pytest.raises(PageIndexAPIError, match="model, storage_path"): + PageIndexClient(api_key="k", model="m", storage_path="/tmp/x") + + +def test_local_client_does_not_touch_disk(tmp_path): + storage = tmp_path / "store" + PageIndexClient(storage_path=str(storage)) + assert not storage.exists() + + +def test_retrieve_model_carries_agents_sdk_prefix(tmp_path): + def resolved(retrieve_model): + return PageIndexClient(retrieve_model=retrieve_model, + storage_path=str(tmp_path / "s")).retrieve_model + + assert resolved("anthropic/claude-sonnet-4-6") == "litellm/anthropic/claude-sonnet-4-6" + for already_routable in ("gpt-4o", "openai/gpt-4o", "litellm/anthropic/claude-sonnet-4-6"): + assert resolved(already_routable) == already_routable + + +def test_explicit_mode_clients(tmp_path): + from pageindex import PageIndexCloudClient, PageIndexLocalClient + + for bad_key in (None, ""): + with pytest.raises(PageIndexAPIError, match="requires a PageIndex API key"): + PageIndexCloudClient(bad_key) + cloud = PageIndexCloudClient("k") + assert cloud.api_key == "k" and isinstance(cloud, PageIndexClient) + + local = PageIndexLocalClient(model="m", storage_path=str(tmp_path / "s")) + assert local.model == "m" and isinstance(local, PageIndexClient) + with pytest.raises(TypeError): + PageIndexLocalClient("k") + + +# ── local: indexing and reading ── + +def test_submit_and_get_tree(local_client, indexed_doc, tmp_path, monkeypatch): + tree = local_client.get_tree(indexed_doc, node_summary=True) + assert tree["status"] == "completed" + assert tree["retrieval_ready"] is True + root = tree["result"][0] + assert root["page_index"] == 1 + assert "start_index" not in root and "end_index" not in root + assert root["prefix_summary"] == "root summary" + assert "summary" not in root + child = root["nodes"][0] + assert child["summary"] == "child summary" + assert child["text"] == "Second page about bananas" + + no_summary = local_client.get_tree(indexed_doc)["result"][0] + assert "summary" not in no_summary and "prefix_summary" not in no_summary + + +def test_get_tree_include_text_false(local_client, indexed_doc): + tree = local_client.get_tree(indexed_doc, include_text=False) + root = tree["result"][0] + assert "text" not in root + assert "text" not in root["nodes"][0] + assert root["page_index"] == 1 + + with_text = local_client.get_tree(indexed_doc)["result"][0] + assert "text" in with_text + + +def test_get_document_structure(local_client, indexed_doc): + result = local_client.get_document_structure(indexed_doc) + assert isinstance(result, list) + root = result[0] + assert "text" not in root + assert "text" not in root["nodes"][0] + assert "prefix_summary" in root + assert root["nodes"][0]["summary"] == "child summary" + + +def test_get_page_content(local_client, indexed_doc): + pages = local_client.get_page_content(indexed_doc, "1") + assert len(pages) == 1 + assert pages[0]["page_index"] == 1 + assert "Hello page one" in pages[0]["markdown"] + + pages = local_client.get_page_content(indexed_doc, "1-2") + assert len(pages) == 2 + + pages = local_client.get_page_content(indexed_doc, "2,1") + assert [p["page_index"] for p in pages] == [1, 2] + + assert local_client.get_page_content(indexed_doc, "99") == [] + + with pytest.raises(ValueError): + local_client.get_page_content(indexed_doc, "abc") + + +def test_submit_does_not_create_cwd_logs(local_client, sample_pdf, tmp_path, monkeypatch): + monkeypatch.chdir(tmp_path) + def fake_page_index_main(doc, opt=None, logger=None, page_list=None): + logger.info({"probe": True}) + return {"doc_name": "sample.pdf", "doc_description": None, + "structure": json.loads(json.dumps(STRUCTURE))} + monkeypatch.setattr(page_index_module, "page_index_main", fake_page_index_main) + local_client.submit_document(sample_pdf) + assert not (tmp_path / "logs").exists() + + +def test_submit_flash(local_client, sample_pdf, monkeypatch): + calls = {} + def fake_flash(pdf, summary=True, summary_model=None, **kwargs): + calls["summary"] = summary + calls["summary_model"] = summary_model + return {"doc_name": "sample.pdf", + "structure": [{"title": "Flash Root", "start_index": 1, + "end_index": 2, "summary": "s", "nodes": []}]} + monkeypatch.setattr(pageindex.flash, "page_index_flash", fake_flash) + monkeypatch.setattr(pageindex.utils, "llm_completion", + lambda model, prompt, **kw: "Flash description.") + doc_id = local_client.submit_document(sample_pdf, mode="flash")["doc_id"] + assert calls == {"summary": True, "summary_model": local_client.summary_model} + root = local_client.get_tree(doc_id)["result"][0] + assert root["node_id"] == "0000" + assert "Hello page one" in root["text"] + assert local_client.get_document(doc_id)["description"] == "Flash description." + + +def test_llm_completion_missing_key_raises_immediately(monkeypatch): + import openai + monkeypatch.delenv("OPENAI_API_KEY", raising=False) + monkeypatch.setattr(pageindex.utils, "_openai_sync_client", None) + monkeypatch.setattr(pageindex.utils, "_openai_async_client", None) + with pytest.raises(openai.OpenAIError): + pageindex.utils.llm_completion("gpt-4o", "probe") + with pytest.raises(openai.OpenAIError): + asyncio.run(pageindex.utils.llm_acompletion("gpt-4o", "probe")) + + +def test_submit_missing_llm_key_fails_loud(local_client, sample_pdf, monkeypatch): + monkeypatch.delenv("OPENAI_API_KEY", raising=False) + monkeypatch.setattr(pageindex.utils, "_openai_sync_client", 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"}): + with pytest.raises(PageIndexAPIError, match="OPENAI_API_KEY"): + local_client.submit_document(sample_pdf, **kwargs) + assert local_client.list_documents()["total"] == 0 + + +def test_submit_rejections(local_client, sample_pdf, tmp_path): + with pytest.raises(FileNotFoundError): + local_client.submit_document(str(tmp_path / "missing.pdf")) + (tmp_path / "notes.txt").write_text("hi") + with pytest.raises(PageIndexAPIError, match="only PDF"): + local_client.submit_document(str(tmp_path / "notes.txt")) + with pytest.raises(PageIndexAPIError, match="unknown local processing mode"): + local_client.submit_document(sample_pdf, mode="mcp") + with pytest.raises(PageIndexAPIError, match="folders"): + local_client.submit_document(sample_pdf, folder_id="f1") + with pytest.raises(PageIndexAPIError, match="beta_headers"): + local_client.submit_document(sample_pdf, beta_headers=["block_reference"]) + + +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") + with pytest.raises(PageIndexAPIError, match="could not read PDF"): + local_client.submit_document(str(bad)) + + +def test_encrypted_pdf_raises_api_error(local_client, sample_pdf, tmp_path): + from PyPDF2 import PdfReader, PdfWriter + + writer = PdfWriter() + for page in PdfReader(sample_pdf).pages: + writer.add_page(page) + writer.encrypt("secret") + enc = tmp_path / "enc.pdf" + with open(enc, "wb") as f: + writer.write(f) + with pytest.raises(PageIndexAPIError, match="could not read PDF"): + local_client.submit_document(str(enc)) + + +def test_submit_explicit_standard_mode(local_client, sample_pdf, monkeypatch): + calls = [] + + def fake_page_index_main(doc, opt=None, logger=None, page_list=None): + calls.append(doc) + return { + "doc_name": "sample.pdf", + "doc_description": "A test document.", + "structure": json.loads(json.dumps(STRUCTURE)), + } + + monkeypatch.setattr(page_index_module, "page_index_main", fake_page_index_main) + doc_id = local_client.submit_document(sample_pdf, mode="standard")["doc_id"] + + assert calls == [sample_pdf] + assert local_client._api._store.get_meta(doc_id)["mode"] == "standard" + + +@pytest.mark.parametrize("mode", ["standard", "flash"]) +def test_submit_from_running_event_loop( + local_client, sample_pdf, monkeypatch, mode +): + def fake_index(*args): + asyncio.run(asyncio.sleep(0)) + return json.loads(json.dumps(STRUCTURE)), "A test document." + + monkeypatch.setattr(local_client._api, f"_index_{mode}", fake_index) + + async def submit(): + return local_client.submit_document(sample_pdf, mode=mode) + + doc_id = asyncio.run(submit())["doc_id"] + assert local_client.get_document(doc_id)["status"] == "completed" + + +def test_submit_with_metadata(local_client, sample_pdf, monkeypatch): + monkeypatch.setattr( + page_index_module, "page_index_main", + lambda doc, opt=None, logger=None, page_list=None: { + "doc_name": "sample.pdf", "doc_description": None, + "structure": json.loads(json.dumps(STRUCTURE))}) + tags = {"project": "alpha", "year": 2026} + doc_id = local_client.submit_document(sample_pdf, metadata=tags)["doc_id"] + assert local_client.get_tree(doc_id)["metadata"] == tags + assert local_client.get_ocr(doc_id)["metadata"] == tags + assert local_client.list_documents()["documents"][0]["metadata"] == tags + assert "metadata" not in local_client.get_document(doc_id) + + +def test_submit_metadata_validation(local_client, sample_pdf, monkeypatch): + indexed = [] + monkeypatch.setattr(page_index_module, "page_index_main", + lambda *args, **kwargs: indexed.append(1)) + with pytest.raises(PageIndexAPIError, match="metadata must be a dict"): + local_client.submit_document(sample_pdf, metadata=["not", "a", "dict"]) + with pytest.raises(PageIndexAPIError, match="valid JSON"): + local_client.submit_document(sample_pdf, metadata={"x": object()}) + assert indexed == [] + + +def test_blank_pdf_rejected(local_client, tmp_path): + from conftest import build_pdf + blank = tmp_path / "blank.pdf" + blank.write_bytes(build_pdf(["", ""])) + with pytest.raises(PageIndexAPIError, match="All pages are blank"): + local_client.submit_document(str(blank)) + + +def test_get_ocr(local_client, indexed_doc): + page = local_client.get_ocr(indexed_doc) + assert page["result"][0] == {"page_index": 1, + "markdown": "Hello page one about apples"} + raw = local_client.get_ocr(indexed_doc, format="raw") + assert raw["result"] == ("Hello page one about apples\n\n" + "Second page about bananas") + node = local_client.get_ocr(indexed_doc, format="node") + assert node["result"] == [ + {"title": "Root Section", "level": 1, "page_index": 1, + "text": "Hello page one about applesSecond page about bananas"}, + {"title": "Child Section", "level": 2, "page_index": 2, + "text": "Second page about bananas"}, + ] + with pytest.raises(ValueError): + local_client.get_ocr(indexed_doc, format="bogus") + + +def test_document_management(local_client, indexed_doc): + assert indexed_doc.startswith("pi-") + doc = local_client.get_document(indexed_doc) + assert doc["id"] == indexed_doc + assert re.fullmatch(r"\d{4}-\d{2}-\d{2}T\d{2}:\d{2}:\d{2}(\.\d{3}000)?", + doc["createdAt"]) + assert doc["name"] == "sample.pdf" + assert doc["description"] == "A test document." + assert doc["status"] == "completed" + assert doc["pageNum"] == 2 + assert doc["folderId"] is None + + listing = local_client.list_documents() + assert listing["total"] == 1 + assert listing["limit"] == 50 and listing["offset"] == 0 + assert listing["documents"][0]["id"] == indexed_doc + + assert local_client.is_retrieval_ready(indexed_doc) is True + + assert local_client.delete_document(indexed_doc) == { + "message": "Document deleted successfully."} + with pytest.raises(PageIndexAPIError, match="Document not found"): + local_client.delete_document(indexed_doc) + assert local_client.is_retrieval_ready(indexed_doc) is False + + +def test_manifest_write_through_and_self_heal(local_client, indexed_doc, tmp_path): + manifest_path = tmp_path / "store" / "manifest.json" + manifest = json.loads(manifest_path.read_text()) + assert manifest["docs"][indexed_doc]["name"] == "sample.pdf" + + # corrupt cache → listings rebuild it from the doc.json files + manifest_path.write_text("{broken") + assert local_client.list_documents()["total"] == 1 + assert indexed_doc in json.loads(manifest_path.read_text())["docs"] + + # missing cache → same + manifest_path.unlink() + assert local_client.list_documents()["total"] == 1 + + # doc dir removed behind the store's back → healed, not served stale + shutil.rmtree(tmp_path / "store" / "docs" / indexed_doc) + assert local_client.list_documents()["total"] == 0 + + local_client_meta = json.loads(manifest_path.read_text()) + assert local_client_meta == {"docs": {}} + + +@pytest.mark.parametrize("bad_entry", ["corrupt-entry", {"id": "wrong"}]) +def test_manifest_invalid_entry_self_heals( + local_client, indexed_doc, tmp_path, bad_entry +): + manifest_path = tmp_path / "store" / "manifest.json" + manifest_path.write_text(json.dumps({"docs": {indexed_doc: bad_entry}})) + + listing = local_client.list_documents() + + assert listing["total"] == 1 + assert listing["documents"][0]["id"] == indexed_doc + healed = json.loads(manifest_path.read_text()) + assert healed["docs"][indexed_doc]["name"] == "sample.pdf" + + +def test_manifest_updated_on_delete(local_client, indexed_doc, tmp_path): + local_client.delete_document(indexed_doc) + manifest = json.loads((tmp_path / "store" / "manifest.json").read_text()) + assert manifest == {"docs": {}} + + +def test_manifest_picks_up_external_doc(local_client, indexed_doc, tmp_path): + # a doc whose manifest update was lost (e.g. concurrent writer) still lists + docs_dir = tmp_path / "store" / "docs" + external_id = "11111111-1111-4111-8111-111111111111" + shutil.copytree(docs_dir / indexed_doc, docs_dir / external_id) + meta_path = docs_dir / external_id / "doc.json" + meta = json.loads(meta_path.read_text()) + meta["id"] = external_id + meta_path.write_text(json.dumps(meta)) + + ids = {d["id"] for d in local_client.list_documents()["documents"]} + assert ids == {indexed_doc, external_id} + + +def test_manifest_ignores_incomplete_dir(local_client, indexed_doc, tmp_path): + (tmp_path / "store" / "docs" / "crashed-save").mkdir() + listing = local_client.list_documents() + assert listing["total"] == 1 + assert listing["documents"][0]["id"] == indexed_doc + + +def test_torn_delete_never_lists_ghost(local_client, indexed_doc, tmp_path): + # crash mid-delete: doc.json gone, dir and manifest entry remain + doc_dir = tmp_path / "store" / "docs" / indexed_doc + (doc_dir / "doc.json").unlink() + + assert local_client.list_documents()["total"] == 0 + manifest = json.loads((tmp_path / "store" / "manifest.json").read_text()) + assert manifest == {"docs": {}} + with pytest.raises(PageIndexAPIError): + local_client.get_document(indexed_doc) + with pytest.raises(PageIndexAPIError, match="Document not found"): + local_client.delete_document(indexed_doc) + assert not doc_dir.exists() + + +def test_corrupt_doc_json_is_contained(local_client, indexed_doc, sample_pdf, tmp_path): + second = local_client.submit_document(sample_pdf)["doc_id"] + (tmp_path / "store" / "docs" / indexed_doc / "doc.json").write_text("{truncated") + + # manifest still holds a good copy of the meta — served consistently + assert local_client.get_document(indexed_doc)["id"] == indexed_doc + assert local_client.list_documents()["total"] == 2 + + # without the manifest copy, the doc is treated as absent, not a crash + (tmp_path / "store" / "manifest.json").unlink() + listing = local_client.list_documents() + assert listing["total"] == 1 + assert listing["documents"][0]["id"] == second + with pytest.raises(PageIndexAPIError): + local_client.get_document(indexed_doc) + assert local_client.is_retrieval_ready(indexed_doc) is False + + +def test_invalid_utf8_is_contained(local_client, indexed_doc, tmp_path): + doc_json = tmp_path / "store" / "docs" / indexed_doc / "doc.json" + doc_json.write_bytes(b'{"id": "\xff\xfe broken') + + # the manifest copy keeps serving, consistently across list and get + assert local_client.get_document(indexed_doc)["id"] == indexed_doc + assert local_client.list_documents()["total"] == 1 + + # even with the manifest corrupted the same way: no crash, self-heals + (tmp_path / "store" / "manifest.json").write_bytes(b"\xff\xfe") + assert local_client.list_documents()["total"] == 0 + with pytest.raises(PageIndexAPIError): + local_client.get_document(indexed_doc) + + +def test_corrupt_data_files_fail_loud(local_client, indexed_doc, tmp_path): + doc_dir = tmp_path / "store" / "docs" / indexed_doc + (doc_dir / "tree.json").write_bytes(b"\xff\xfe") + with pytest.raises(PageIndexAPIError, match="unreadable"): + local_client.get_tree(indexed_doc) + assert local_client.is_retrieval_ready(indexed_doc) is False + + (doc_dir / "pages.json").write_text("{broken") + with pytest.raises(PageIndexAPIError, match="unreadable"): + local_client.get_ocr(indexed_doc) + + # the metadata itself is intact, so listings stay honest + assert local_client.list_documents()["total"] == 1 + + +def test_get_tree_fails_loud_on_broken_pages(local_client, indexed_doc, tmp_path): + doc_dir = tmp_path / "store" / "docs" / indexed_doc + (doc_dir / "pages.json").write_text("{broken") + with pytest.raises(PageIndexAPIError, match="unreadable"): + local_client.get_tree(indexed_doc) + + +def test_get_tree_fails_loud_on_empty_pages(local_client, indexed_doc, tmp_path): + doc_dir = tmp_path / "store" / "docs" / indexed_doc + (doc_dir / "pages.json").write_text("[]") + with pytest.raises(PageIndexAPIError, match="no page content"): + local_client.get_tree(indexed_doc) + + +def test_data_file_as_directory_fails_loud(local_client, indexed_doc, tmp_path): + tree_path = tmp_path / "store" / "docs" / indexed_doc / "tree.json" + tree_path.unlink() + tree_path.mkdir() + with pytest.raises(PageIndexAPIError, match="unreadable"): + local_client.get_tree(indexed_doc) + + +def test_list_documents_skips_unsafe_directory_names( + local_client, indexed_doc, tmp_path +): + bad_dir = tmp_path / "store" / "docs" / "bad\\name" + bad_dir.mkdir() + (bad_dir / "doc.json").write_text("{}") + listing = local_client.list_documents() + assert [d["id"] for d in listing["documents"]] == [indexed_doc] + + +def test_generate_doc_description_error_boundary(monkeypatch): + 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([]) == "" + monkeypatch.setattr(pageindex.utils, "llm_completion", + raiser(ValueError("provider rejected the model"))) + with pytest.raises(ValueError): + pageindex.utils.generate_doc_description([]) + + +def test_generate_summaries_all_failed_raises(monkeypatch): + async def boom(model, prompt): + raise ValueError("bad key") + monkeypatch.setattr(pageindex.utils, "llm_acompletion", boom) + structure = [{"title": "A", "text": "t1", + "nodes": [{"title": "B", "text": "t2"}]}] + with pytest.raises(RuntimeError, match="all nodes"): + asyncio.run(pageindex.utils.generate_summaries_for_structure(structure)) + + +def test_generate_summaries_partial_failure_absorbed(monkeypatch): + async def flaky(model, prompt): + if "t1" in prompt: + raise ValueError("transient") + return "ok" + monkeypatch.setattr(pageindex.utils, "llm_acompletion", flaky) + structure = [{"title": "A", "text": "t1", + "nodes": [{"title": "B", "text": "t2"}]}] + result = asyncio.run(pageindex.utils.generate_summaries_for_structure(structure)) + summaries = {n["title"]: n["summary"] + for n in pageindex.utils.structure_to_list(result)} + assert summaries == {"A": "", "B": "ok"} + + +def test_delete_survives_marker_tamper(local_client, tmp_path): + tampered = tmp_path / "store" / "docs" / "tampered" / "doc.json" + tampered.mkdir(parents=True) + with pytest.raises(PageIndexAPIError, match="Document not found"): + local_client.delete_document("tampered") + assert not tampered.parent.exists() + + +def test_list_documents_validation(local_client): + with pytest.raises(ValueError): + local_client.list_documents(limit=0) + with pytest.raises(ValueError): + local_client.list_documents(offset=-1) + with pytest.raises(PageIndexAPIError, match="folders"): + local_client.list_documents(folder_id="f1") + + +def test_missing_document_errors(local_client): + with pytest.raises(PageIndexAPIError): + local_client.get_tree("nope") + with pytest.raises(PageIndexAPIError): + local_client.get_document("nope") + assert local_client.is_retrieval_ready("nope") is False + + +def test_traversal_ids_are_contained(local_client, indexed_doc, tmp_path): + store_root = tmp_path / "store" + with pytest.raises(PageIndexAPIError): + local_client.get_document("../../etc") + with pytest.raises(PageIndexAPIError): + local_client.delete_document("..") + assert (store_root / "docs").exists() + + +def test_folders_are_cloud_only(local_client): + with pytest.raises(PageIndexAPIError, match="cloud-only"): + local_client.create_folder("team") + with pytest.raises(PageIndexAPIError, match="cloud-only"): + local_client.list_folders() + + +# ── local: retrieval endpoints are cloud-only ── + +def test_retrieval_endpoints_cloud_only(local_client): + with pytest.raises(PageIndexAPIError, match="use chat_completions"): + local_client.submit_query("any", "q") + with pytest.raises(PageIndexAPIError, match="use chat_completions"): + local_client.get_retrieval("any") + + +def test_chat_completions_cloud_only(local_client): + with pytest.raises(PageIndexAPIError, match="not yet supported in local mode"): + local_client.chat_completions( + messages=[{"role": "user", "content": "q"}]) + + +# ── cloud mode: request wiring ── + +class FakeResponse: + def __init__(self, payload=None, status_code=200, text="", content=b"{}", + lines=None): + self._payload = payload if payload is not None else {} + self.status_code = status_code + self.text = text + self.content = content + self._lines = lines or [] + + def json(self): + return self._payload + + def iter_lines(self): + return iter(self._lines) + + def close(self): + pass + + +def _patch_requests(monkeypatch, handler): + """Replace cloud_api's requests module with per-verb fakes.""" + fake = types.SimpleNamespace( + post=lambda url, **kw: handler("POST", url, kw), + get=lambda url, **kw: handler("GET", url, kw), + delete=lambda url, **kw: handler("DELETE", url, kw), + Response=FakeResponse, + ) + monkeypatch.setattr("pageindex.cloud_api.requests", fake) + + +@pytest.fixture +def cloud(monkeypatch): + client = PageIndexClient(api_key="secret") + calls = [] + class Fake: + payload = {} + def handler(method, url, kw): + calls.append({"method": method, "url": url, **kw}) + return FakeResponse(Fake.payload) + _patch_requests(monkeypatch, handler) + return client, calls, Fake + + +def test_cloud_request_wiring(cloud, sample_pdf): + client, calls, fake = cloud + + fake.payload = {"doc_id": "pi-1"} + assert client.submit_document(sample_pdf) == {"doc_id": "pi-1"} + assert calls[-1]["url"] == "https://api.pageindex.ai/doc/" + assert calls[-1]["headers"] == {"api_key": "secret"} + assert calls[-1]["data"] == {"if_retrieval": True} + assert "timeout" not in calls[-1] + + client.submit_document(sample_pdf, metadata={"project": "alpha"}) + assert calls[-1]["data"]["metadata"] == json.dumps({"project": "alpha"}) + + fake.payload = {"status": "processing", "retrieval_ready": False} + client.get_tree("pi-1", node_summary=True) + assert calls[-1]["url"].endswith("/doc/pi-1/?type=tree&summary=True&include_text=true") + assert calls[-1]["timeout"] == 30 + assert client.is_retrieval_ready("pi-1") is False + + client.get_ocr("pi/../1") + assert "/doc/pi%2F..%2F1/" in calls[-1]["url"] + + client.BASE_URL = "https://staging.example" + client.api_key = "other" + client.get_document("pi-1") + assert calls[-1]["url"] == "https://staging.example/doc/pi-1/metadata/" + assert calls[-1]["headers"] == {"api_key": "other"} + + +def test_cloud_error_and_empty_delete(cloud, monkeypatch): + client, calls, fake = cloud + _patch_requests(monkeypatch, + lambda m, url, kw: FakeResponse(status_code=401, text="denied")) + with pytest.raises(PageIndexAPIError, + match="Failed to get document metadata: denied"): + client.get_document("pi-1") + + _patch_requests(monkeypatch, lambda m, url, kw: FakeResponse(content=b"")) + assert client.delete_document("pi-1") == {} + + +def test_cloud_chat_stream_parsing(cloud, monkeypatch): + client, calls, fake = cloud + lines = [ + b'data: {"choices": [{"delta": {"role": "assistant", "content": ""}}]}', + b'data: {"choices": [{"delta": {"content": "Hi"}}]}', + b"", + b'data: {"object": "chat.completion.citations", "citations": []}', + b'data: {"choices": [{"delta": {"content": " there"}}]}', + b"data: [DONE]", + ] + _patch_requests(monkeypatch, lambda m, url, kw: FakeResponse(lines=lines)) + pieces = list(client.chat_completions( + messages=[{"role": "user", "content": "q"}], stream=True)) + assert pieces == ["Hi", " there"] + + chunks = list(client.chat_completions( + messages=[{"role": "user", "content": "q"}], stream=True, + stream_metadata=True)) + assert {"object": "chat.completion.citations", "citations": []} in chunks diff --git a/tests/test_issue_163.py b/tests/test_issue_163.py index 7281cd14d..1e2999fe3 100644 --- a/tests/test_issue_163.py +++ b/tests/test_issue_163.py @@ -1,17 +1,11 @@ import pytest import sys import os -from importlib import import_module from unittest.mock import patch, MagicMock sys.path.insert(0, os.path.join(os.path.dirname(__file__), "..")) -# Patch attributes on the module object: the package exports a *function* -# named page_index that shadows the submodule, and Python 3.10's mock -# resolves the string "pageindex.page_index" to that function. -page_index_module = import_module("pageindex.page_index") - -from pageindex.page_index import ( +from pageindex.page_index_classic import ( check_if_toc_extraction_is_complete, check_if_toc_transformation_is_complete, toc_detector_single_page, @@ -22,50 +16,50 @@ class TestRobustKeyAccess: - @patch.object(page_index_module, "llm_completion", return_value="") + @patch("pageindex.page_index_classic.llm_completion", return_value="") def test_toc_detector_empty_response(self, mock_llm): result = toc_detector_single_page("some content", model="test") assert result == "no" - @patch.object(page_index_module, "llm_completion", return_value='{"toc_detected": "yes"}') + @patch("pageindex.page_index_classic.llm_completion", return_value='{"toc_detected": "yes"}') def test_toc_detector_valid_response(self, mock_llm): result = toc_detector_single_page("some content", model="test") assert result == "yes" - @patch.object(page_index_module, "llm_completion", return_value="not json at all") + @patch("pageindex.page_index_classic.llm_completion", return_value="not json at all") def test_toc_detector_malformed_response(self, mock_llm): result = toc_detector_single_page("some content", model="test") assert result == "no" - @patch.object(page_index_module, "llm_completion", return_value="") + @patch("pageindex.page_index_classic.llm_completion", return_value="") def test_extraction_complete_empty_response(self, mock_llm): result = check_if_toc_extraction_is_complete("doc", "toc", model="test") assert result == "no" - @patch.object(page_index_module, "llm_completion", return_value='{"completed": "yes"}') + @patch("pageindex.page_index_classic.llm_completion", return_value='{"completed": "yes"}') def test_extraction_complete_valid_response(self, mock_llm): result = check_if_toc_extraction_is_complete("doc", "toc", model="test") assert result == "yes" - @patch.object(page_index_module, "llm_completion", return_value="") + @patch("pageindex.page_index_classic.llm_completion", return_value="") def test_transformation_complete_empty_response(self, mock_llm): result = check_if_toc_transformation_is_complete("raw", "cleaned", model="test") assert result == "no" - @patch.object(page_index_module, "llm_completion", return_value='{"thinking": "looks fine", "completed": "yes"}') + @patch("pageindex.page_index_classic.llm_completion", return_value='{"thinking": "looks fine", "completed": "yes"}') def test_transformation_complete_valid_response(self, mock_llm): result = check_if_toc_transformation_is_complete("raw", "cleaned", model="test") assert result == "yes" - @patch.object(page_index_module, "llm_completion", return_value="") + @patch("pageindex.page_index_classic.llm_completion", return_value="") def test_detect_page_index_empty_response(self, mock_llm): result = detect_page_index("toc text", model="test") assert result == "no" class TestExtractTocContentRetryLoop: - @patch.object(page_index_module, "check_if_toc_transformation_is_complete") - @patch.object(page_index_module, "llm_completion") + @patch("pageindex.page_index_classic.check_if_toc_transformation_is_complete") + @patch("pageindex.page_index_classic.llm_completion") def test_completes_on_first_try(self, mock_llm, mock_check): mock_llm.return_value = ("full toc content", "finished") mock_check.return_value = "yes" @@ -73,8 +67,8 @@ def test_completes_on_first_try(self, mock_llm, mock_check): assert result == "full toc content" assert mock_llm.call_count == 1 - @patch.object(page_index_module, "check_if_toc_transformation_is_complete") - @patch.object(page_index_module, "llm_completion") + @patch("pageindex.page_index_classic.check_if_toc_transformation_is_complete") + @patch("pageindex.page_index_classic.llm_completion") def test_continues_on_incomplete(self, mock_llm, mock_check): mock_llm.side_effect = [ ("partial toc", "max_output_reached"), @@ -85,8 +79,8 @@ def test_continues_on_incomplete(self, mock_llm, mock_check): assert result == "partial toc continued toc" assert mock_llm.call_count == 2 - @patch.object(page_index_module, "check_if_toc_transformation_is_complete") - @patch.object(page_index_module, "llm_completion") + @patch("pageindex.page_index_classic.check_if_toc_transformation_is_complete") + @patch("pageindex.page_index_classic.llm_completion") def test_max_retries_raises_exception(self, mock_llm, mock_check): mock_llm.return_value = ("chunk", "max_output_reached") mock_check.return_value = "no" @@ -94,8 +88,8 @@ def test_max_retries_raises_exception(self, mock_llm, mock_check): extract_toc_content("raw content", model="test") assert mock_llm.call_count == 6 - @patch.object(page_index_module, "check_if_toc_transformation_is_complete") - @patch.object(page_index_module, "llm_completion") + @patch("pageindex.page_index_classic.check_if_toc_transformation_is_complete") + @patch("pageindex.page_index_classic.llm_completion") def test_chat_history_grows_incrementally(self, mock_llm, mock_check): call_count = [0] @@ -120,8 +114,8 @@ def side_effect(*args, **kwargs): class TestTocTransformerRetryLoop: - @patch.object(page_index_module, "check_if_toc_transformation_is_complete") - @patch.object(page_index_module, "llm_completion") + @patch("pageindex.page_index_classic.check_if_toc_transformation_is_complete") + @patch("pageindex.page_index_classic.llm_completion") def test_completes_on_first_try(self, mock_llm, mock_check): mock_llm.return_value = ( '{"table_of_contents": [{"structure": "1", "title": "Intro", "page": 1}]}', @@ -132,8 +126,8 @@ def test_completes_on_first_try(self, mock_llm, mock_check): assert len(result) == 1 assert result[0]["title"] == "Intro" - @patch.object(page_index_module, "check_if_toc_transformation_is_complete") - @patch.object(page_index_module, "llm_completion") + @patch("pageindex.page_index_classic.check_if_toc_transformation_is_complete") + @patch("pageindex.page_index_classic.llm_completion") def test_handles_missing_table_of_contents_key(self, mock_llm, mock_check): mock_llm.return_value = ('{"other_key": "value"}', "finished") mock_check.return_value = "yes" diff --git a/tests/test_package_surface.py b/tests/test_package_surface.py new file mode 100644 index 000000000..e10c3e8c5 --- /dev/null +++ b/tests/test_package_surface.py @@ -0,0 +1,62 @@ +"""What `pip install pageindex` exposes: 0.2.8 helper compat and import cost.""" +import subprocess +import sys + +from pageindex.utils import create_node_mapping, print_tree, remove_fields + +TREE = [ + {"title": "Root", "node_id": "0000", "page_index": 1, + "text": "root text", + "nodes": [ + {"title": "Child", "node_id": "0001", "page_index": 3, + "text": "child text"}, + ]}, + {"title": "Tail", "node_id": "0002", "page_index": 5, "text": "tail text"}, +] + + +# ── the published 0.2.8 pageindex.utils surface, as the cookbooks call it ── + +def test_remove_fields_max_len(): + out = remove_fields({"keep": "x" * 50, "text": "gone"}, max_len=10) + assert out == {"keep": "x" * 10 + "..."} + assert remove_fields({"keep": "short"}, max_len=10) == {"keep": "short"} + + +def test_create_node_mapping_flat(): + mapping = create_node_mapping(TREE) + assert set(mapping) == {"0000", "0001", "0002"} + assert mapping["0001"]["title"] == "Child" + + +def test_create_node_mapping_page_ranges(): + mapping = create_node_mapping(TREE, include_page_ranges=True, max_page=9) + assert mapping["0000"] == {"node": TREE[0], "start_index": 1, "end_index": 3} + assert mapping["0001"]["start_index"] == 3 + assert mapping["0001"]["end_index"] == 5 + assert mapping["0002"] == {"node": TREE[1], "start_index": 5, "end_index": 9} + + +def test_print_tree_exclude_fields(capsys): + print_tree(TREE, exclude_fields=["text"]) + out = capsys.readouterr().out + assert "Root" in out and "'text'" not in out + + print_tree(TREE) + assert "[0000] Root" in capsys.readouterr().out + + +# ── import cost: the SDK must not pay for the indexing stack ── + +def test_import_pageindex_is_lazy(): + probe = ( + "import sys; import pageindex; " + "heavy = [m for m in ('pageindex.page_index_classic', 'pageindex.flash', " + "'pageindex.utils', 'pageindex.tree_optimize', 'numpy', 'PyPDF2') " + "if m in sys.modules]; " + "print(','.join(heavy) or 'clean'); " + "print(type(pageindex.page_index_main).__name__)" + ) + out = subprocess.run([sys.executable, "-c", probe], + capture_output=True, text=True, check=True) + assert out.stdout.split() == ["clean", "function"] diff --git a/tests/test_page_index.py b/tests/test_page_index.py index 1ffd51581..da69a9a79 100644 --- a/tests/test_page_index.py +++ b/tests/test_page_index.py @@ -1,13 +1,7 @@ import unittest -from importlib import import_module from unittest.mock import Mock, patch -# Patch attributes on the module object: the package exports a *function* -# named page_index that shadows the submodule, and Python 3.10's mock -# resolves the string "pageindex.page_index" to that function. -page_index_module = import_module("pageindex.page_index") - -from pageindex.page_index import ( +from pageindex.page_index_classic import ( _secure_doc_text, process_no_toc, process_toc_no_page_numbers, @@ -25,10 +19,10 @@ def test_rejects_same_length_reordered_llm_toc(self): {"structure": "1", "title": "First", "physical_index": ""}, ] - with patch.object(page_index_module, "toc_transformer", return_value=toc), \ - patch.object(page_index_module, "count_tokens", return_value=1), \ - patch.object(page_index_module, "page_list_to_group_text", return_value=[" "]), \ - patch.object(page_index_module, "add_page_number_to_toc", return_value=reordered): + with patch("pageindex.page_index_classic.toc_transformer", return_value=toc), \ + patch("pageindex.page_index_classic.count_tokens", return_value=1), \ + patch("pageindex.page_index_classic.page_list_to_group_text", return_value=[" "]), \ + patch("pageindex.page_index_classic.add_page_number_to_toc", return_value=reordered): with self.assertRaises(ValueError): process_toc_no_page_numbers( "toc", @@ -38,17 +32,17 @@ def test_rejects_same_length_reordered_llm_toc(self): ) def test_process_no_toc_validates_continuation_chunks(self): - with patch.object(page_index_module, "count_tokens", return_value=1), \ - patch.object( - page_index_module, "page_list_to_group_text", + with patch("pageindex.page_index_classic.count_tokens", return_value=1), \ + patch( + "pageindex.page_index_classic.page_list_to_group_text", return_value=["", ""], ), \ - patch.object( - page_index_module, "generate_toc_init", + patch( + "pageindex.page_index_classic.generate_toc_init", return_value=[{"title": "First", "physical_index": ""}], ), \ - patch.object( - page_index_module, "generate_toc_continue", + patch( + "pageindex.page_index_classic.generate_toc_continue", return_value=[{"title": "Second", "physical_index": ""}], ): result = process_no_toc(