RAG Without Vectors: How PageIndex Retrieves by Reasoning

pip set up pageindex openai requests

from pageindex import PageIndexClient
import pageindex.utils as utils
import os
from getpass import getpass

PAGEINDEX_API_KEY = getpass('Enter PageIndex API Key: ')
pi_client = PageIndexClient(api_key=PAGEINDEX_API_KEY)

import openai

OPENAI_API_KEY = getpass('Enter OpenAI API Key: ')

async def call_llm(immediate, mannequin="gpt-5.4", temperature=0):
    shopper = openai.AsyncOpenAI(api_key=OPENAI_API_KEY)
    response = await shopper.chat.completions.create(
        mannequin=mannequin,
        messages=[{"role": "user", "content": prompt}],
        temperature=temperature
    )
    return response.decisions[0].message.content material.strip()

# ─────────────────────────────────────────────
# Step 1: Build the PageIndex Tree
# ─────────────────────────────────────────────

# 1.1 Download the Transformer paper and submit it
import os, requests

pdf_url = "https://arxiv.org/pdf/1706.03762.pdf"
pdf_path = os.path.be part of("knowledge", pdf_url.break up("/")[-1])
os.makedirs("knowledge", exist_ok=True)

print("Downloading 'Attention Is All You Need'...")
response = requests.get(pdf_url)
with open(pdf_path, "wb") as f:
    f.write(response.content material)
print(f" Saved to {pdf_path}")

doc_id = pi_client.submit_document(pdf_path)["doc_id"]
print(f" Document submitted. doc_id: {doc_id}")

# 1.2 Retrieve the tree (ballot till prepared)
import time

print("nWaiting for PageIndex tree to be prepared", finish="")
whereas not pi_client.is_retrieval_ready(doc_id):
    print(".", finish="", flush=True)
    time.sleep(5)

tree = pi_client.get_tree(doc_id, node_summary=True)["result"]
print("nn Document Tree Structure:")
utils.print_tree(tree)

# ─────────────────────────────────────────────
# Step 2: Reasoning-Based Retrieval
# ─────────────────────────────────────────────

# 2.1 Define a question that requires navigating throughout sections
import json

# This question is deliberately cross-cutting -- it could actually't be answered
# by a single part, which is the place tree search shines over top-k.
question = "Why did the authors select self-attention over recurrence, and what are the complexity trade-offs they in contrast?"

tree_without_text = utils.remove_fields(tree.copy(), fields=["text"])

search_prompt = f"""
You are given a query and a hierarchical tree construction of a analysis paper.
Each node has a node_id, title, and a abstract of its content material.

Your job: establish ALL nodes which might be more likely to include info related to answering the query.
Think rigorously -- the reply could also be unfold throughout a number of sections.

Question: {question}

Document tree:
{json.dumps(tree_without_text, indent=2)}

Reply ONLY on this JSON format, no preamble:
{{
    "considering": "<step-by-step reasoning about which nodes are related and why>",
    "node_list": ["node_id_1", "node_id_2", ...]
}}
"""

print(f' Query: "{question}"n')
print("Running tree search with GPT-5.4...")
tree_search_result = await call_llm(search_prompt)

# 2.2 Inspect the retrieval reasoning and matched nodes
node_map = utils.create_node_mapping(tree)
result_json = json.hundreds(tree_search_result)

print("n LLM Reasoning:")
utils.print_wrapped(result_json["thinking"])

print("n Retrieved Nodes:")
for node_id in result_json["node_list"]:
    node = node_map[node_id]
    print(f"  • [{node['node_id']}] Page {node['page_index']:>2} -- {node['title']}")

# ─────────────────────────────────────────────
# Step 3: Answer Generation
# ─────────────────────────────────────────────

# 3.1 Stitch collectively context from all retrieved nodes
node_list = result_json["node_list"]
relevant_content = "nn---nn".be part of(
    f"[Section: {node_map[nid]['title']}]n{node_map[nid]['text']}"
    for nid in node_list
)

print(f"n Retrieved Context Preview (first 1200 chars):n")
utils.print_wrapped(relevant_content[:1200] + "...n")

# 3.2 Generate a structured reply grounded within the retrieved sections
answer_prompt = f"""
You are a technical assistant. Answer the query under utilizing ONLY the offered context.
Be particular -- reference precise design decisions, numbers, and trade-offs talked about within the textual content.

Question: {question}

Context:
{relevant_content}

Structure your reply as:
1. The core motivation for selecting self-attention
2. The particular complexity comparisons made (embody any tables or numbers)
3. Any caveats or limitations the authors acknowledged
"""

print(" Generating reply...n")
reply = await call_llm(answer_prompt)

print("─" * 60)
print(" Final Answer:n")
utils.print_wrapped(reply)
print("─" * 60)

query2 = "How does the multi-head consideration mechanism work, and what's the function of scaling in dot-product consideration?"

search_prompt2 = f"""
You are given a query and a hierarchical tree construction of a analysis paper.
Identify all nodes more likely to include the reply.

Question: {query2}

Document tree:
{json.dumps(tree_without_text, indent=2)}

Reply ONLY on this JSON format:
{{
    "considering": "<reasoning>",
    "node_list": ["node_id_1", ...]
}}
"""

print(f'nn Second Query: "{query2}"n')
result2_raw = await call_llm(search_prompt2)
result2 = json.hundreds(result2_raw)

print(" Reasoning:")
utils.print_wrapped(result2["thinking"])

relevant_content2 = "nn---nn".be part of(
    f"[Section: {node_map[nid]['title']}]n{node_map[nid]['text']}"
    for nid in result2["node_list"]
)

answer_prompt2 = f"""
Answer the next query utilizing ONLY the offered context.
Explain the mechanism clearly, as if for a technical weblog put up.

Question: {query2}
Context: {relevant_content2}
"""

answer2 = await call_llm(answer_prompt2)
print("n Answer:n")
utils.print_wrapped(answer2)