Add Table of Contents (TOC) RAG system with new modules, database enhancements, and UI integration.

main.py

@@ -16,6 +16,11 @@ from src.vector_rag.retriever import retrieve as vector_retrieve
 from src.vector_rag.llm_engine import chat_with_rag as vector_chat_with_rag, chat_no_context as vector_chat_no_context
 from src.vector_rag.indexer import index_all_documents as vector_index_all
 
+# TOC RAG imports
+from src.toc_rag.retriever import retrieve as toc_retrieve
+from src.toc_rag.llm_engine import chat_with_rag as toc_chat_with_rag, chat_no_context as toc_chat_no_context
+from src.toc_rag.indexer import index_all_documents as toc_index_all
+
 from src.cost_tracker import add_chat_tokens, get_token_summary, clear_stats
 
 app = FastAPI(title="Local RAG Comparator")
@@ -118,10 +123,31 @@ def get_document_chunks(filename: str):
     except Exception as e:
         print(f"Error fetching vector chunks: {e}")
 
+    # ── TOC Index chunks ──
+    toc_chunks = []
+    try:
+        toc_collection = get_collection("toc_index")
+        toc_results = toc_collection.get(where={"source": filename})
+        
+        if toc_results and toc_results["ids"]:
+            for i in range(len(toc_results["ids"])):
+                toc_chunks.append({
+                    "id": toc_results["ids"][i],
+                    "text": toc_results["documents"][i] if toc_results["documents"] else "",
+                    "node_id": toc_results["metadatas"][i].get("node_id", "") if toc_results["metadatas"] else "",
+                    "title": toc_results["metadatas"][i].get("title", "") if toc_results["metadatas"] else "",
+                    "tokens": toc_results["metadatas"][i].get("tokens", 0) if toc_results["metadatas"] else 0
+                })
+            toc_chunks.sort(key=lambda x: x["node_id"])
+            
+    except Exception as e:
+        print(f"Error fetching toc chunks: {e}")
+
     return {
         "filename": filename,
         "page_chunks": page_chunks,
-        "vector_chunks": vector_chunks
+        "vector_chunks": vector_chunks,
+        "toc_chunks": toc_chunks
     }
 
 @app.get("/api/index/status")
@@ -193,6 +219,36 @@ def chat_vector(req: ChatRequest):
             
     return StreamingResponse(generate(), media_type="text/event-stream")
 
+@app.post("/api/chat/toc")
+def chat_toc(req: ChatRequest):
+    def generate():
+        try:
+            nodes = toc_retrieve(req.query, top_k=req.top_k)
+            source_data = [
+                {
+                    "source": n.source,
+                    "node_id": n.node_id,
+                    "title": n.title,
+                    "score": n.score,
+                    "text": n.text,
+                    "tokens": n.tokens,
+                }
+                for n in nodes
+            ]
+            yield f"data: {json.dumps({'type': 'sources', 'sources': source_data})}\n\n"
+            
+            generator = toc_chat_with_rag(req.query, nodes) if nodes else toc_chat_no_context(req.query)
+                
+            for chunk in generator:
+                if chunk.get("type") == "stats":
+                    add_chat_tokens(chunk.get("prompt_eval_count", 0), chunk.get("eval_count", 0))
+                yield f"data: {json.dumps(chunk)}\n\n"
+                
+        except Exception as e:
+            yield f"data: {json.dumps({'type': 'error', 'content': str(e)})}\n\n"
+            
+    return StreamingResponse(generate(), media_type="text/event-stream")
+
 @app.post("/api/upload")
 async def upload_files(background_tasks: BackgroundTasks, files: list[UploadFile] = File(...)):
     global is_indexing
@@ -217,6 +273,8 @@ async def upload_files(background_tasks: BackgroundTasks, files: list[UploadFile
             page_index_all(progress_callback=progress_cb)
             progress_cb("Auto-indexing: Starting Vector RAG...")
             vector_index_all(progress_callback=progress_cb)
+            progress_cb("Auto-indexing: Starting TOC RAG...")
+            toc_index_all(progress_callback=progress_cb)
             progress_cb("Successfully indexed all documents.")
         except Exception as e:
             progress_cb(f"Error during indexing: {str(e)}")
@@ -247,6 +305,8 @@ def index_documents(background_tasks: BackgroundTasks):
             page_index_all(progress_callback=progress_cb)
             progress_cb("Starting Vector Indexing...")
             vector_index_all(progress_callback=progress_cb)
+            progress_cb("Starting TOC Indexing...")
+            toc_index_all(progress_callback=progress_cb)
             progress_cb("Indexing completed.")
         except Exception as e:
             progress_cb(f"Error: {str(e)}")
@@ -259,16 +319,9 @@ def index_documents(background_tasks: BackgroundTasks):
 
 @app.post("/api/index/clear")
 def clear_index():
-    # remove both collections and underlying storage directory
-    from src.shared_db import get_chroma_client, clear_chroma_storage
-
-    # Clear ChromaDB collections via client API
-    client = get_chroma_client()
-    collections = client.list_collections()
-    for coll in collections:
-        client.delete_collection(name=coll.name)
+    from src.shared_db import clear_chroma_storage
 
-    # Purge on‑disk database files so we start fresh next time
+    # Purge on-disk database files and re-initialize a fresh client
     clear_chroma_storage()
 
     # Clear any app-specific stats

src/cost_tracker.py

@@ -17,6 +17,7 @@ def _default_stats():
         "embedding_tokens": 0,
         "prompt_tokens": 0,
         "completion_tokens": 0,
+        "toc_analysis_tokens": 0,
     }
 
 
@@ -55,6 +56,15 @@ def add_chat_tokens(prompt_count: int, completion_count: int):
         save_stats(stats)
 
 
+def add_toc_analysis_tokens(count: int):
+    if count <= 0:
+        return
+    with _lock:
+        stats = load_stats()
+        stats["toc_analysis_tokens"] = stats.get("toc_analysis_tokens", 0) + count
+        save_stats(stats)
+
+
 def get_token_summary():
     """Return raw token counts only – no monetary costs."""
     with _lock:
@@ -63,6 +73,7 @@ def get_token_summary():
             stats.get("embedding_tokens", 0)
             + stats.get("prompt_tokens", 0)
             + stats.get("completion_tokens", 0)
+            + stats.get("toc_analysis_tokens", 0)
         )
         return {
             "tokens": stats,

src/shared_db.py

@@ -42,12 +42,21 @@ def get_collection(name: str):
 def clear_chroma_storage():
     """Remove all persistent chroma data on disk and reset the client."""
     global _client
-    # If client exists, try closing it cleanly (PersistentClient doesn't provide close) by deleting reference
-    _client = None
-    try:
-        shutil.rmtree(CHROMA_PATH)
-        Path(CHROMA_PATH).mkdir(exist_ok=True, parents=True)
-        print(f"Cleared chroma storage directory: {CHROMA_PATH}")
-    except Exception as e:
-        print(f"Error clearing chroma storage: {e}")
-        raise
+    with _lock:
+        # Drop the old client reference so SQLite file handles are released
+        _client = None
+
+        try:
+            shutil.rmtree(CHROMA_PATH)
+            Path(CHROMA_PATH).mkdir(exist_ok=True, parents=True)
+            print(f"Cleared chroma storage directory: {CHROMA_PATH}")
+        except Exception as e:
+            print(f"Error clearing chroma storage: {e}")
+            raise
+
+        # Eagerly create a fresh client so subsequent calls don't hit a stale db
+        try:
+            _client = chromadb.PersistentClient(path=CHROMA_PATH)
+            print("Fresh ChromaDB Client re-initialized after clear.")
+        except Exception as e:
+            print(f"Warning: could not re-init ChromaDB Client: {e}")

src/toc_rag/__init__.py

@@ -0,0 +1 @@
+# TOC-based Index RAG module

src/toc_rag/indexer.py

@@ -0,0 +1,361 @@
+"""
+indexer.py
+----------
+Main indexer for the TOC-based Index RAG.
+Parses documents into hierarchical tree structures,
+embeds each node's text, and stores in ChromaDB.
+
+Supports PDF (via tree_builder), MD/TXT (via md_parser).
+"""
+
+import os
+import asyncio
+from pathlib import Path
+from typing import Generator
+
+import fitz  # PyMuPDF
+import ollama
+from docx import Document
+from rich.console import Console
+from rich.progress import track
+from src.shared_db import get_collection as get_shared_collection
+from src.cost_tracker import add_embedding_tokens
+
+import tiktoken
+_tokenizer = tiktoken.get_encoding("cl100k_base")
+
+from .toc_utils import (
+    get_page_tokens, JsonLogger, ConfigLoader,
+    write_node_id, add_node_text, remove_structure_text,
+    structure_to_list, format_structure,
+    generate_summaries_for_structure,
+    create_clean_structure_for_description, generate_doc_description,
+    get_pdf_name, count_tokens, LLM_MODEL,
+)
+from .tree_builder import tree_parser
+from .md_parser import md_to_tree
+
+console = Console()
+
+# ─── Config ────────────────────────────────────────────────────────────────────
+EMBED_MODEL = "nomic-embed-text"
+COLLECTION = "toc_index"
+DOCS_FOLDER = "./documents"
+# ───────────────────────────────────────────────────────────────────────────────
+
+
+
+def get_chroma_collection():
+    return get_shared_collection(COLLECTION)
+
+
+CHUNK_CHAR_LIMIT = 6000    # chars per sub-chunk (~1500 tokens, well within 8192 model limit)
+CHUNK_OVERLAP_CHARS = 1000  # overlap between consecutive sub-chunks
+
+
+def _chunk_text_for_embedding(text: str, header: str = "") -> list[str]:
+    """Split long text into overlapping sub-chunks for embedding.
+
+    Each chunk is prefixed with *header* (e.g. section title + summary) so
+    the embedding captures the node context even for chunks deep in the text.
+
+    Short texts (≤ CHUNK_CHAR_LIMIT) are returned as a single-item list.
+    """
+    # Header is always prepended; account for its length
+    available = CHUNK_CHAR_LIMIT - len(header)
+    if available < 500:
+        available = 500  # safety floor
+
+    # If text fits in one chunk, return as-is
+    if len(text) <= available:
+        return [f"{header}{text}".strip()] if header else [text]
+
+    chunks: list[str] = []
+    start = 0
+    while start < len(text):
+        end = start + available
+        chunk_body = text[start:end]
+        full_chunk = f"{header}{chunk_body}".strip() if header else chunk_body
+        chunks.append(full_chunk)
+        # Advance with overlap
+        start += available - CHUNK_OVERLAP_CHARS
+        if start >= len(text):
+            break
+
+    return chunks
+
+
+def embed_text(text: str) -> tuple[list[float], int]:
+    """Generate embedding locally using Ollama.
+    Text MUST already be chunked to fit the model context window."""
+    try:
+        response = ollama.embeddings(model=EMBED_MODEL, prompt=text)
+    except Exception as e:
+        print(f"[TOC-RAG] Embed error (len={len(text)} chars): {e}")
+        raise
+    exact_tokens = len(_tokenizer.encode(text))
+    add_embedding_tokens(exact_tokens)
+    return response["embedding"], exact_tokens
+
+
+def flatten_tree(structure, parent_title=""):
+    """Flatten tree structure into a list of indexable nodes."""
+    nodes = []
+    for node in structure:
+        title = node.get('title', '')
+        full_title = f"{parent_title} > {title}" if parent_title else title
+        text = node.get('text', '')
+        summary = node.get('summary', node.get('prefix_summary', ''))
+
+        # Build the indexable text: title + summary + text content
+        indexable_text = f"Section: {full_title}\n"
+        if summary:
+            indexable_text += f"Summary: {summary}\n"
+        if text:
+            indexable_text += f"Content:\n{text}"
+
+        nodes.append({
+            'node_id': node.get('node_id', ''),
+            'title': title,
+            'full_title': full_title,
+            'text': indexable_text.strip(),
+            'summary': summary,
+            'start_index': node.get('start_index'),
+            'end_index': node.get('end_index'),
+            'line_num': node.get('line_num'),
+        })
+
+        if 'nodes' in node and node['nodes']:
+            nodes.extend(flatten_tree(node['nodes'], full_title))
+
+    return nodes
+
+
+def _build_tree_for_pdf(filepath: str, progress_callback=None):
+    """Build tree structure for a PDF using the full TOC pipeline."""
+    logger = JsonLogger(filepath)
+    opt = ConfigLoader().load({
+        'model': LLM_MODEL,
+        'if_add_node_id': 'yes',
+        'if_add_node_summary': 'yes',
+        'if_add_doc_description': 'no',
+        'if_add_node_text': 'yes',
+    })
+
+    if progress_callback:
+        progress_callback("TOC RAG: Parsing PDF pages...")
+
+    page_list = get_page_tokens(filepath)
+    logger.info({'total_pages': len(page_list), 'total_tokens': sum(p[1] for p in page_list)})
+
+    async def build():
+        if progress_callback:
+            progress_callback("TOC RAG: Building document tree structure...")
+        structure = await tree_parser(page_list, opt, doc=filepath, logger=logger)
+
+        if progress_callback:
+            progress_callback("TOC RAG: Assigning node IDs...")
+        write_node_id(structure)
+
+        if progress_callback:
+            progress_callback("TOC RAG: Extracting node text...")
+        add_node_text(structure, page_list)
+
+        if progress_callback:
+            progress_callback("TOC RAG: Generating summaries...")
+        await generate_summaries_for_structure(structure, model=opt.model)
+
+        return structure
+
+    # Run async pipeline
+    try:
+        loop = asyncio.get_event_loop()
+        if loop.is_running():
+            # If we're already in an async context, create a new loop in a thread
+            import concurrent.futures
+            with concurrent.futures.ThreadPoolExecutor() as pool:
+                future = pool.submit(asyncio.run, build())
+                structure = future.result()
+        else:
+            structure = asyncio.run(build())
+    except RuntimeError:
+        structure = asyncio.run(build())
+
+    return structure
+
+
+def _build_tree_for_text(filepath: str, progress_callback=None):
+    """Build tree structure for text/markdown files."""
+    if progress_callback:
+        progress_callback(f"TOC RAG: Parsing {Path(filepath).suffix} file...")
+
+    async def build():
+        return await md_to_tree(
+            filepath,
+            if_thinning=True,
+            min_token_threshold=5000,
+            if_add_node_summary='yes',
+            summary_token_threshold=200,
+            model=LLM_MODEL,
+            if_add_node_text='yes',
+            if_add_node_id='yes',
+        )
+
+    try:
+        loop = asyncio.get_event_loop()
+        if loop.is_running():
+            import concurrent.futures
+            with concurrent.futures.ThreadPoolExecutor() as pool:
+                future = pool.submit(asyncio.run, build())
+                result = future.result()
+        else:
+            result = asyncio.run(build())
+    except RuntimeError:
+        result = asyncio.run(build())
+
+    return result.get('structure', [])
+
+
+def index_document(filepath: str, collection, progress_callback=None) -> int:
+    """Index all nodes of a single document. Returns node count indexed."""
+    ext = Path(filepath).suffix.lower()
+
+    if ext == '.pdf':
+        structure = _build_tree_for_pdf(filepath, progress_callback)
+    elif ext in ('.md', '.txt'):
+        structure = _build_tree_for_text(filepath, progress_callback)
+    elif ext == '.docx':
+        # Convert DOCX to text, then treat as text
+        doc = Document(filepath)
+        temp_txt = Path(filepath).with_suffix('.tmp.txt')
+        with open(temp_txt, 'w', encoding='utf-8') as f:
+            f.write('\n'.join(p.text for p in doc.paragraphs))
+        structure = _build_tree_for_text(str(temp_txt), progress_callback)
+        temp_txt.unlink(missing_ok=True)
+    else:
+        console.print(f"[yellow]Skipping unsupported file for TOC indexing: {filepath}[/yellow]")
+        return 0
+
+    if not structure:
+        console.print(f"[red]No structure extracted from {filepath}[/red]")
+        return 0
+
+    # Flatten tree to indexable nodes
+    flat_nodes = flatten_tree(structure)
+    if not flat_nodes:
+        console.print(f"[red]No nodes to index from {filepath}[/red]")
+        return 0
+
+    filename = Path(filepath).name
+    ids, embeddings, documents, metadatas = [], [], [], []
+    total = len(flat_nodes)
+    chunks_total = 0
+
+    for i, node in enumerate(track(flat_nodes, description=f"  Embedding TOC nodes for {filename}"), 1):
+        if progress_callback:
+            progress_callback(f"TOC RAG: Embedding {filename} node {i}/{total}")
+
+        text = node['text']
+        if not text.strip():
+            continue
+
+        # Build a short header that gets prepended to every sub-chunk
+        header = f"Section: {node['full_title']}\n"
+        if node.get('summary'):
+            header += f"Summary: {node['summary']}\n"
+        header += "Content:\n"
+
+        # Split into sub-chunks (single chunk if short enough)
+        raw_content = node.get('text', '')
+        # Remove the header we already built from raw_content if flatten_tree already added it
+        # (flatten_tree builds the full indexable text with header)  
+        content_only = raw_content
+        for prefix in [f"Section: {node['full_title']}\n", f"Summary: {node.get('summary','')}\n", "Content:\n"]:
+            if content_only.startswith(prefix):
+                content_only = content_only[len(prefix):]
+
+        sub_chunks = _chunk_text_for_embedding(content_only, header=header)
+
+        for chunk_idx, chunk_text in enumerate(sub_chunks):
+            if len(sub_chunks) == 1:
+                doc_id = f"{Path(filepath).stem}__toc_node_{node['node_id']}"
+            else:
+                doc_id = f"{Path(filepath).stem}__toc_node_{node['node_id']}_chunk_{chunk_idx}"
+
+            # Skip if already indexed
+            existing = collection.get(ids=[doc_id])
+            if existing["ids"]:
+                continue
+
+            try:
+                emb, tokens = embed_text(chunk_text)
+            except Exception as e:
+                console.print(f"  [red]Skipping node {node['node_id']} chunk {chunk_idx}: {e}[/red]")
+                continue
+
+            ids.append(doc_id)
+            embeddings.append(emb)
+            documents.append(chunk_text)
+            metadatas.append({
+                "source": filename,
+                "node_id": node['node_id'],
+                "title": node.get('title', ''),
+                "full_title": node.get('full_title', ''),
+                "doc_type": Path(filepath).suffix.lower().replace('.', ''),
+                "tokens": tokens,
+                "summary": node.get('summary', ''),
+                "chunk_index": chunk_idx,
+                "total_chunks": len(sub_chunks),
+            })
+            chunks_total += 1
+
+    if ids:
+        collection.add(
+            ids=ids,
+            embeddings=embeddings,
+            documents=documents,
+            metadatas=metadatas,
+        )
+
+    console.print(f"  [dim]{chunks_total} embedding(s) from {total} node(s)[/dim]")
+    return chunks_total
+
+
+def index_all_documents(docs_folder: str = DOCS_FOLDER, progress_callback=None) -> None:
+    """Scan the documents folder and index everything with TOC structure."""
+    collection = get_chroma_collection()
+    folder = Path(docs_folder)
+    supported = [".pdf", ".txt", ".docx", ".md"]
+    files = [f for f in folder.iterdir() if f.suffix.lower() in supported]
+
+    if not files:
+        if progress_callback:
+            progress_callback("No documents to index for TOC RAG.")
+        console.print(f"[yellow]No supported files found in {docs_folder}[/yellow]")
+        return
+
+    console.print(f"\n[bold cyan]🌳 TOC RAG: Found {len(files)} document(s) to index[/bold cyan]\n")
+
+    total_nodes = 0
+    for filepath in files:
+        console.print(f"[green]→ TOC Processing:[/green] {filepath.name}")
+        try:
+            count = index_document(str(filepath), collection, progress_callback)
+            total_nodes += count
+            console.print(f"  [dim]Indexed {count} node(s)[/dim]\n")
+        except Exception as e:
+            console.print(f"  [red]Error indexing {filepath.name}: {e}[/red]\n")
+
+    if progress_callback:
+        progress_callback("TOC Indexing Complete.")
+
+    console.print(f"[bold green]✅ TOC RAG Done! Total nodes indexed: {total_nodes}[/bold green]")
+    console.print(f"[dim]ChromaDB collection size: {collection.count()} nodes[/dim]")
+
+
+if __name__ == "__main__":
+    console.print("[bold]🌳 TOC Indexer — Local RAG[/bold]")
+    console.print(f"Embed model : {EMBED_MODEL}")
+    console.print(f"LLM model   : {LLM_MODEL}")
+    console.print(f"Docs folder : {DOCS_FOLDER}\n")
+    index_all_documents()

src/toc_rag/llm_engine.py

@@ -0,0 +1,93 @@
+"""
+llm_engine.py
+-------------
+Wraps Ollama to call your local LLM model.
+Builds a RAG prompt using retrieved TOC nodes as context.
+Streams response with token tracking.
+"""
+
+from typing import Generator
+import ollama
+from .retriever import RetrievedNode, build_context
+
+# ─── Config ────────────────────────────────────────────────────────────────────
+LLM_MODEL = "llama3.2:1b"
+# ───────────────────────────────────────────────────────────────────────────────
+
+SYSTEM_PROMPT = """You are a helpful document assistant using tree-structured document indexing.
+Answer the user's question using ONLY the provided context sections.
+Always cite your sources by mentioning the filename and section title.
+If the answer is not found in the context, say so honestly — do not make up information.
+"""
+
+
+def build_rag_prompt(query: str, nodes: list[RetrievedNode]) -> str:
+    context = build_context(nodes)
+    return f"""Use the following document sections to answer the question.
+
+=== CONTEXT (Tree-Indexed Sections) ===
+{context}
+
+=== QUESTION ===
+{query}
+
+=== ANSWER ==="""
+
+
+def chat_with_rag(
+    query: str,
+    nodes: list[RetrievedNode],
+    chat_history: list[dict] | None = None,
+) -> Generator[dict, None, None]:
+    """
+    Stream a response from the local LLM using retrieved tree nodes as context.
+    Yields text chunks as they arrive (streaming).
+    """
+    messages = [{"role": "system", "content": SYSTEM_PROMPT}]
+
+    if chat_history:
+        messages.extend(chat_history)
+
+    rag_prompt = build_rag_prompt(query, nodes)
+    messages.append({"role": "user", "content": rag_prompt})
+
+    stream = ollama.chat(
+        model=LLM_MODEL,
+        messages=messages,
+        stream=True,
+    )
+
+    for chunk in stream:
+        token = chunk.get("message", {}).get("content", "")
+        if token:
+            yield {"type": "token", "content": token}
+
+        if chunk.get("done"):
+            yield {
+                "type": "stats",
+                "prompt_eval_count": chunk.get("prompt_eval_count", 0),
+                "eval_count": chunk.get("eval_count", 0)
+            }
+
+
+def chat_no_context(query: str, chat_history: list[dict] | None = None) -> Generator[dict, None, None]:
+    """Fallback: LLM response when no relevant nodes are found."""
+    messages = [{"role": "system", "content": SYSTEM_PROMPT}]
+    if chat_history:
+        messages.extend(chat_history)
+    messages.append({
+        "role": "user",
+        "content": f"{query}\n\n(No relevant document sections were found for this query.)"
+    })
+    stream = ollama.chat(model=LLM_MODEL, messages=messages, stream=True)
+    for chunk in stream:
+        token = chunk.get("message", {}).get("content", "")
+        if token:
+            yield {"type": "token", "content": token}
+
+        if chunk.get("done"):
+            yield {
+                "type": "stats",
+                "prompt_eval_count": chunk.get("prompt_eval_count", 0),
+                "eval_count": chunk.get("eval_count", 0)
+            }

src/toc_rag/md_parser.py

@@ -0,0 +1,297 @@
+"""
+md_parser.py
+------------
+Markdown-to-tree parsing for .md and .txt files.
+Extracts headers, builds hierarchical tree, supports thinning
+and summary generation.
+
+Replica of the reference page-indexing repository.
+"""
+
+import os
+import re
+import json
+import asyncio
+
+from .toc_utils import (
+    count_tokens, write_node_id, format_structure,
+    structure_to_list, create_clean_structure_for_description,
+    generate_doc_description, generate_node_summary,
+    ChatGPT_API_async,
+)
+
+
+# ─── Summary Helpers ──────────────────────────────────────────────────────────
+
+async def get_node_summary(node, summary_token_threshold=200, model=None):
+    node_text = node.get('text')
+    num_tokens = count_tokens(node_text, model=model)
+    if num_tokens < summary_token_threshold:
+        return node_text
+    else:
+        return await generate_node_summary(node, model=model)
+
+
+async def generate_summaries_for_structure_md(structure, summary_token_threshold, model=None):
+    nodes = structure_to_list(structure)
+    tasks = [get_node_summary(node, summary_token_threshold=summary_token_threshold, model=model) for node in nodes]
+    summaries = await asyncio.gather(*tasks)
+
+    for node, summary in zip(nodes, summaries):
+        if not node.get('nodes'):
+            node['summary'] = summary
+        else:
+            node['prefix_summary'] = summary
+    return structure
+
+
+# ─── Markdown Extraction ──────────────────────────────────────────────────────
+
+def extract_nodes_from_markdown(markdown_content):
+    """Find headers in markdown, skipping code blocks."""
+    header_pattern = r'^(#{1,6})\s+(.+)$'
+    code_block_pattern = r'^```'
+    node_list = []
+
+    lines = markdown_content.split('\n')
+    in_code_block = False
+
+    for line_num, line in enumerate(lines, 1):
+        stripped_line = line.strip()
+
+        if re.match(code_block_pattern, stripped_line):
+            in_code_block = not in_code_block
+            continue
+
+        if not stripped_line:
+            continue
+
+        if not in_code_block:
+            match = re.match(header_pattern, stripped_line)
+            if match:
+                title = match.group(2).strip()
+                node_list.append({'node_title': title, 'line_num': line_num})
+
+    return node_list, lines
+
+
+def extract_node_text_content(node_list, markdown_lines):
+    """Extract text content between headers."""
+    all_nodes = []
+    for node in node_list:
+        line_content = markdown_lines[node['line_num'] - 1]
+        header_match = re.match(r'^(#{1,6})', line_content)
+
+        if header_match is None:
+            print(f"Warning: Line {node['line_num']} does not contain a valid header: '{line_content}'")
+            continue
+
+        processed_node = {
+            'title': node['node_title'],
+            'line_num': node['line_num'],
+            'level': len(header_match.group(1))
+        }
+        all_nodes.append(processed_node)
+
+    for i, node in enumerate(all_nodes):
+        start_line = node['line_num'] - 1
+        if i + 1 < len(all_nodes):
+            end_line = all_nodes[i + 1]['line_num'] - 1
+        else:
+            end_line = len(markdown_lines)
+
+        node['text'] = '\n'.join(markdown_lines[start_line:end_line]).strip()
+    return all_nodes
+
+
+# ─── Token Counting / Thinning ─────────────────────────────────────────────────
+
+def update_node_list_with_text_token_count(node_list, model=None):
+    """Calculate cumulative token counts for nodes including children."""
+    def find_all_children(parent_index, parent_level, node_list):
+        children_indices = []
+        for i in range(parent_index + 1, len(node_list)):
+            current_level = node_list[i]['level']
+            if current_level <= parent_level:
+                break
+            children_indices.append(i)
+        return children_indices
+
+    result_list = node_list.copy()
+
+    for i in range(len(result_list) - 1, -1, -1):
+        current_node = result_list[i]
+        current_level = current_node['level']
+        children_indices = find_all_children(i, current_level, result_list)
+
+        node_text = current_node.get('text', '')
+        total_text = node_text
+
+        for child_index in children_indices:
+            child_text = result_list[child_index].get('text', '')
+            if child_text:
+                total_text += '\n' + child_text
+
+        result_list[i]['text_token_count'] = count_tokens(total_text, model=model)
+
+    return result_list
+
+
+def tree_thinning_for_index(node_list, min_node_token=None, model=None):
+    """Merge small nodes into their parents."""
+    def find_all_children(parent_index, parent_level, node_list):
+        children_indices = []
+        for i in range(parent_index + 1, len(node_list)):
+            current_level = node_list[i]['level']
+            if current_level <= parent_level:
+                break
+            children_indices.append(i)
+        return children_indices
+
+    result_list = node_list.copy()
+    nodes_to_remove = set()
+
+    for i in range(len(result_list) - 1, -1, -1):
+        if i in nodes_to_remove:
+            continue
+
+        current_node = result_list[i]
+        current_level = current_node['level']
+        total_tokens = current_node.get('text_token_count', 0)
+
+        if total_tokens < min_node_token:
+            children_indices = find_all_children(i, current_level, result_list)
+
+            children_texts = []
+            for child_index in sorted(children_indices):
+                if child_index not in nodes_to_remove:
+                    child_text = result_list[child_index].get('text', '')
+                    if child_text.strip():
+                        children_texts.append(child_text)
+                    nodes_to_remove.add(child_index)
+
+            if children_texts:
+                parent_text = current_node.get('text', '')
+                merged_text = parent_text
+                for child_text in children_texts:
+                    if merged_text and not merged_text.endswith('\n'):
+                        merged_text += '\n\n'
+                    merged_text += child_text
+
+                result_list[i]['text'] = merged_text
+                result_list[i]['text_token_count'] = count_tokens(merged_text, model=model)
+
+    for index in sorted(nodes_to_remove, reverse=True):
+        result_list.pop(index)
+
+    return result_list
+
+
+# ─── Tree Building ─────────────────────────────────────────────────────────────
+
+def build_tree_from_nodes(node_list):
+    """Convert flat node list to nested tree structure."""
+    if not node_list:
+        return []
+
+    stack = []
+    root_nodes = []
+    node_counter = 1
+
+    for node in node_list:
+        current_level = node['level']
+
+        tree_node = {
+            'title': node['title'],
+            'node_id': str(node_counter).zfill(4),
+            'text': node['text'],
+            'line_num': node['line_num'],
+            'nodes': []
+        }
+        node_counter += 1
+
+        while stack and stack[-1][1] >= current_level:
+            stack.pop()
+
+        if not stack:
+            root_nodes.append(tree_node)
+        else:
+            parent_node, _ = stack[-1]
+            parent_node['nodes'].append(tree_node)
+
+        stack.append((tree_node, current_level))
+
+    return root_nodes
+
+
+# ─── Main Entry ────────────────────────────────────────────────────────────────
+
+async def md_to_tree(md_path, if_thinning=False, min_token_threshold=None,
+                     if_add_node_summary='no', summary_token_threshold=None,
+                     model=None, if_add_doc_description='no',
+                     if_add_node_text='no', if_add_node_id='yes'):
+    """Top-level async entry for markdown-to-tree conversion."""
+    with open(md_path, 'r', encoding='utf-8') as f:
+        markdown_content = f.read()
+
+    print("Extracting nodes from markdown...")
+    node_list, markdown_lines = extract_nodes_from_markdown(markdown_content)
+
+    print("Extracting text content from nodes...")
+    nodes_with_content = extract_node_text_content(node_list, markdown_lines)
+
+    if if_thinning:
+        nodes_with_content = update_node_list_with_text_token_count(nodes_with_content, model=model)
+        print("Thinning nodes...")
+        nodes_with_content = tree_thinning_for_index(nodes_with_content, min_token_threshold, model=model)
+
+    print("Building tree from nodes...")
+    tree_structure = build_tree_from_nodes(nodes_with_content)
+
+    if if_add_node_id == 'yes':
+        write_node_id(tree_structure)
+
+    print("Formatting tree structure...")
+
+    if if_add_node_summary == 'yes':
+        tree_structure = format_structure(
+            tree_structure,
+            order=['title', 'node_id', 'summary', 'prefix_summary', 'text', 'line_num', 'nodes']
+        )
+
+        print("Generating summaries for each node...")
+        tree_structure = await generate_summaries_for_structure_md(
+            tree_structure, summary_token_threshold=summary_token_threshold, model=model
+        )
+
+        if if_add_node_text == 'no':
+            tree_structure = format_structure(
+                tree_structure,
+                order=['title', 'node_id', 'summary', 'prefix_summary', 'line_num', 'nodes']
+            )
+
+        if if_add_doc_description == 'yes':
+            print("Generating document description...")
+            clean_structure = create_clean_structure_for_description(tree_structure)
+            doc_description = generate_doc_description(clean_structure, model=model)
+            return {
+                'doc_name': os.path.splitext(os.path.basename(md_path))[0],
+                'doc_description': doc_description,
+                'structure': tree_structure,
+            }
+    else:
+        if if_add_node_text == 'yes':
+            tree_structure = format_structure(
+                tree_structure,
+                order=['title', 'node_id', 'summary', 'prefix_summary', 'text', 'line_num', 'nodes']
+            )
+        else:
+            tree_structure = format_structure(
+                tree_structure,
+                order=['title', 'node_id', 'summary', 'prefix_summary', 'line_num', 'nodes']
+            )
+
+    return {
+        'doc_name': os.path.splitext(os.path.basename(md_path))[0],
+        'structure': tree_structure,
+    }

src/toc_rag/retriever.py

@@ -0,0 +1,93 @@
+"""
+retriever.py
+------------
+Node-level retrieval for the TOC-based Index RAG.
+Embeds query with Ollama, searches ChromaDB toc_index collection,
+returns top-K nodes with metadata.
+"""
+
+import ollama
+from dataclasses import dataclass
+from src.shared_db import get_collection as get_shared_collection
+from src.cost_tracker import add_embedding_tokens
+
+import tiktoken
+_tokenizer = tiktoken.get_encoding("cl100k_base")
+
+# ─── Config ────────────────────────────────────────────────────────────────────
+EMBED_MODEL = "nomic-embed-text"
+COLLECTION = "toc_index"
+TOP_K = 5
+# ───────────────────────────────────────────────────────────────────────────────
+
+
+@dataclass
+class RetrievedNode:
+    text:       str
+    source:     str
+    node_id:    str
+    title:      str
+    score:      float   # cosine distance (lower = more similar)
+    tokens:     int = 0
+    summary:    str = ""
+    full_title: str = ""
+
+
+def get_collection():
+    return get_shared_collection(COLLECTION)
+
+
+def embed_query(query: str) -> list[float]:
+    response = ollama.embeddings(model=EMBED_MODEL, prompt=query)
+    exact_tokens = len(_tokenizer.encode(query))
+    add_embedding_tokens(exact_tokens)
+    return response["embedding"]
+
+
+def retrieve(query: str, top_k: int = TOP_K) -> list[RetrievedNode]:
+    """
+    Embed the user query and return the top-k most relevant tree nodes
+    from the ChromaDB toc_index collection.
+    """
+    collection = get_collection()
+
+    if collection.count() == 0:
+        return []
+
+    query_emb = embed_query(query)
+
+    results = collection.query(
+        query_embeddings=[query_emb],
+        n_results=min(top_k, collection.count()),
+        include=["documents", "metadatas", "distances"],
+    )
+
+    nodes = []
+    for doc, meta, dist in zip(
+        results["documents"][0],
+        results["metadatas"][0],
+        results["distances"][0],
+    ):
+        nodes.append(RetrievedNode(
+            text       = doc,
+            source     = meta.get("source", "unknown"),
+            node_id    = meta.get("node_id", ""),
+            title      = meta.get("title", ""),
+            score      = round(float(dist), 4),
+            tokens     = int(meta.get("tokens", 0)),
+            summary    = meta.get("summary", ""),
+            full_title = meta.get("full_title", ""),
+        ))
+
+    return nodes
+
+
+def build_context(nodes: list[RetrievedNode]) -> str:
+    """Format retrieved nodes into a context block for the LLM prompt."""
+    parts = []
+    for i, node in enumerate(nodes, start=1):
+        label = node.full_title or node.title
+        parts.append(
+            f"[Source {i}: {node.source} — Section: {label}]\n{node.text}"
+        )
+    return "\n\n---\n\n".join(parts)

src/toc_rag/toc_detector.py

@@ -0,0 +1,219 @@
+"""
+toc_detector.py
+---------------
+TOC detection and extraction functions.
+Scans document pages to find Table of Contents,
+extracts TOC content, and checks for page numbers.
+
+Replica of the reference page-indexing repository,
+adapted for local Ollama LLM.
+"""
+
+import re
+from .toc_utils import (
+    ChatGPT_API, ChatGPT_API_with_finish_reason,
+    extract_json, get_json_content,
+)
+
+
+def toc_detector_single_page(content, model=None):
+    """Detect if a single page contains a table of contents."""
+    prompt = f"""
+    Your job is to detect if there is a table of content provided in the given text.
+
+    Given text: {content}
+
+    return the following JSON format:
+    {{
+        "thinking": <why do you think there is a table of content in the given text>
+        "toc_detected": "<yes or no>",
+    }}
+
+    Directly return the final JSON structure. Do not output anything else.
+    Please note: abstract,summary, notation list, figure list, table list, etc. are not table of contents."""
+
+    response = ChatGPT_API(model=model, prompt=prompt)
+    json_content = extract_json(response)
+    return json_content.get('toc_detected', 'no')
+
+
+def detect_page_index(toc_content, model=None):
+    """Detect if TOC has page numbers/indices."""
+    print('start detect_page_index')
+    prompt = f"""
+    You will be given a table of contents.
+
+    Your job is to detect if there are page numbers/indices given within the table of contents.
+
+    Given text: {toc_content}
+
+    Reply format:
+    {{
+        "thinking": <why do you think there are page numbers/indices given within the table of contents>
+        "page_index_given_in_toc": "<yes or no>"
+    }}
+    Directly return the final JSON structure. Do not output anything else."""
+
+    response = ChatGPT_API(model=model, prompt=prompt)
+    json_content = extract_json(response)
+    return json_content.get('page_index_given_in_toc', 'no')
+
+
+def extract_toc_content(content, model=None):
+    """Extract and clean raw TOC text from page content."""
+    prompt = f"""
+    Your job is to extract the full table of contents from the given text, replace ... with :
+
+    Given text: {content}
+
+    Directly return the full table of contents content. Do not output anything else."""
+
+    response, finish_reason = ChatGPT_API_with_finish_reason(model=model, prompt=prompt)
+
+    if_complete = check_if_toc_transformation_is_complete(content, response, model)
+    if if_complete == "yes" and finish_reason == "finished":
+        return response
+
+    chat_history = [
+        {"role": "user", "content": prompt},
+        {"role": "assistant", "content": response},
+    ]
+    prompt = "please continue the generation of table of contents, directly output the remaining part of the structure"
+    new_response, finish_reason = ChatGPT_API_with_finish_reason(model=model, prompt=prompt, chat_history=chat_history)
+    response = response + new_response
+    if_complete = check_if_toc_transformation_is_complete(content, response, model)
+
+    retry_count = 0
+    while not (if_complete == "yes" and finish_reason == "finished"):
+        chat_history = [
+            {"role": "user", "content": prompt},
+            {"role": "assistant", "content": response},
+        ]
+        prompt = "please continue the generation of table of contents, directly output the remaining part of the structure"
+        new_response, finish_reason = ChatGPT_API_with_finish_reason(model=model, prompt=prompt, chat_history=chat_history)
+        response = response + new_response
+        if_complete = check_if_toc_transformation_is_complete(content, response, model)
+
+        retry_count += 1
+        if retry_count > 5:
+            raise Exception('Failed to complete table of contents after maximum retries')
+
+    return response
+
+
+def toc_extractor(page_list, toc_page_list, model):
+    """Concatenate TOC pages, clean up, and detect if page indices exist."""
+    def transform_dots_to_colon(text):
+        text = re.sub(r'\.{5,}', ': ', text)
+        text = re.sub(r'(?:\. ){5,}\.?', ': ', text)
+        return text
+
+    toc_content = ""
+    for page_index in toc_page_list:
+        toc_content += page_list[page_index][0]
+    toc_content = transform_dots_to_colon(toc_content)
+    has_page_index = detect_page_index(toc_content, model=model)
+
+    return {
+        "toc_content": toc_content,
+        "page_index_given_in_toc": has_page_index
+    }
+
+
+def find_toc_pages(start_page_index, page_list, opt, logger=None):
+    """Scan pages to find TOC pages."""
+    print('start find_toc_pages')
+    last_page_is_yes = False
+    toc_page_list = []
+    i = start_page_index
+
+    while i < len(page_list):
+        if i >= opt.toc_check_page_num and not last_page_is_yes:
+            break
+        detected_result = toc_detector_single_page(page_list[i][0], model=opt.model)
+        if detected_result == 'yes':
+            if logger:
+                logger.info(f'Page {i} has toc')
+            toc_page_list.append(i)
+            last_page_is_yes = True
+        elif detected_result == 'no' and last_page_is_yes:
+            if logger:
+                logger.info(f'Found the last page with toc: {i-1}')
+            break
+        i += 1
+
+    if not toc_page_list and logger:
+        logger.info('No toc found')
+
+    return toc_page_list
+
+
+def check_toc(page_list, opt=None):
+    """Main entry: find TOC pages, extract content, check for page numbers."""
+    toc_page_list = find_toc_pages(start_page_index=0, page_list=page_list, opt=opt)
+    if len(toc_page_list) == 0:
+        print('no toc found')
+        return {'toc_content': None, 'toc_page_list': [], 'page_index_given_in_toc': 'no'}
+    else:
+        print('toc found')
+        toc_json = toc_extractor(page_list, toc_page_list, opt.model)
+
+        if toc_json['page_index_given_in_toc'] == 'yes':
+            print('index found')
+            return {
+                'toc_content': toc_json['toc_content'],
+                'toc_page_list': toc_page_list,
+                'page_index_given_in_toc': 'yes'
+            }
+        else:
+            current_start_index = toc_page_list[-1] + 1
+
+            while (toc_json['page_index_given_in_toc'] == 'no' and
+                   current_start_index < len(page_list) and
+                   current_start_index < opt.toc_check_page_num):
+
+                additional_toc_pages = find_toc_pages(
+                    start_page_index=current_start_index,
+                    page_list=page_list,
+                    opt=opt
+                )
+
+                if len(additional_toc_pages) == 0:
+                    break
+
+                additional_toc_json = toc_extractor(page_list, additional_toc_pages, opt.model)
+                if additional_toc_json['page_index_given_in_toc'] == 'yes':
+                    print('index found')
+                    return {
+                        'toc_content': additional_toc_json['toc_content'],
+                        'toc_page_list': additional_toc_pages,
+                        'page_index_given_in_toc': 'yes'
+                    }
+                else:
+                    current_start_index = additional_toc_pages[-1] + 1
+
+            print('index not found')
+            return {
+                'toc_content': toc_json['toc_content'],
+                'toc_page_list': toc_page_list,
+                'page_index_given_in_toc': 'no'
+            }
+
+
+def check_if_toc_transformation_is_complete(content, toc, model=None):
+    """Check if the TOC transformation is complete."""
+    prompt = f"""
+    You are given a raw table of contents and a table of contents.
+    Your job is to check if the table of contents is complete.
+
+    Reply format:
+    {{
+        "thinking": <why do you think the cleaned table of contents is complete or not>
+        "completed": "yes" or "no"
+    }}
+    Directly return the final JSON structure. Do not output anything else."""
+
+    prompt = prompt + '\n Raw Table of contents:\n' + content + '\n Cleaned Table of contents:\n' + toc
+    response = ChatGPT_API(model=model, prompt=prompt)
+    json_content = extract_json(response)
+    return json_content.get('completed', 'no')

src/toc_rag/toc_transformer.py

@@ -0,0 +1,312 @@
+"""
+toc_transformer.py
+------------------
+TOC-to-JSON transformation and page offset calculation.
+Converts raw TOC text to structured JSON, extracts physical indices,
+and calculates page offsets.
+
+Replica of the reference page-indexing repository.
+"""
+
+import json
+import copy
+import re
+from .toc_utils import (
+    ChatGPT_API, ChatGPT_API_with_finish_reason,
+    extract_json, get_json_content,
+    convert_page_to_int, convert_physical_index_to_int,
+    count_tokens,
+)
+
+
+def check_if_toc_transformation_is_complete(content, toc, model=None):
+    """Check if the TOC transformation is complete."""
+    prompt = f"""
+    You are given a raw table of contents and a table of contents.
+    Your job is to check if the table of contents is complete.
+
+    Reply format:
+    {{
+        "thinking": <why do you think the cleaned table of contents is complete or not>
+        "completed": "yes" or "no"
+    }}
+    Directly return the final JSON structure. Do not output anything else."""
+
+    prompt = prompt + '\n Raw Table of contents:\n' + str(content) + '\n Cleaned Table of contents:\n' + str(toc)
+    response = ChatGPT_API(model=model, prompt=prompt)
+    json_content = extract_json(response)
+    return json_content.get('completed', 'no')
+
+
+def toc_transformer(toc_content, model=None):
+    """Convert raw TOC text to structured JSON."""
+    print('start toc_transformer')
+    init_prompt = """
+    You are given a table of contents, You job is to transform the whole table of content into a JSON format included table_of_contents.
+
+    structure is the numeric system which represents the index of the hierarchy section in the table of contents. For example, the first section has structure index 1, the first subsection has structure index 1.1, the second subsection has structure index 1.2, etc.
+
+    The response should be in the following JSON format: 
+    {
+    table_of_contents: [
+        {
+            "structure": <structure index, "x.x.x" or None> (string),
+            "title": <title of the section>,
+            "page": <page number or None>,
+        },
+        ...
+        ],
+    }
+    You should transform the full table of contents in one go.
+    Directly return the final JSON structure, do not output anything else. """
+
+    prompt = init_prompt + '\n Given table of contents\n:' + toc_content
+    last_complete, finish_reason = ChatGPT_API_with_finish_reason(model=model, prompt=prompt)
+    if_complete = check_if_toc_transformation_is_complete(toc_content, last_complete, model)
+    if if_complete == "yes" and finish_reason == "finished":
+        last_complete = extract_json(last_complete)
+        if isinstance(last_complete, dict) and 'table_of_contents' in last_complete:
+            cleaned_response = convert_page_to_int(last_complete['table_of_contents'])
+        elif isinstance(last_complete, list):
+            cleaned_response = convert_page_to_int(last_complete)
+        else:
+            cleaned_response = []
+        return cleaned_response
+
+    last_complete = get_json_content(last_complete)
+    retry_count = 0
+    while not (if_complete == "yes" and finish_reason == "finished"):
+        position = last_complete.rfind('}')
+        if position != -1:
+            last_complete = last_complete[:position+2]
+        prompt = f"""
+        Your task is to continue the table of contents json structure, directly output the remaining part of the json structure.
+        The response should be in the following JSON format: 
+
+        The raw table of contents json structure is:
+        {toc_content}
+
+        The incomplete transformed table of contents json structure is:
+        {last_complete}
+
+        Please continue the json structure, directly output the remaining part of the json structure."""
+
+        new_complete, finish_reason = ChatGPT_API_with_finish_reason(model=model, prompt=prompt)
+
+        if new_complete.startswith('```json'):
+            new_complete = get_json_content(new_complete)
+            last_complete = last_complete + new_complete
+
+        if_complete = check_if_toc_transformation_is_complete(toc_content, last_complete, model)
+
+        retry_count += 1
+        if retry_count > 5:
+            break
+
+    try:
+        last_complete = json.loads(last_complete)
+    except json.JSONDecodeError:
+        last_complete = extract_json(last_complete)
+
+    if isinstance(last_complete, dict) and 'table_of_contents' in last_complete:
+        cleaned_response = convert_page_to_int(last_complete['table_of_contents'])
+    elif isinstance(last_complete, list):
+        cleaned_response = convert_page_to_int(last_complete)
+    else:
+        cleaned_response = []
+    return cleaned_response
+
+
+def toc_index_extractor(toc, content, model=None):
+    """Add physical_index to TOC items from document pages."""
+    print('start toc_index_extractor')
+    toc_extractor_prompt = """
+    You are given a table of contents in a json format and several pages of a document, your job is to add the physical_index to the table of contents in the json format.
+
+    The provided pages contains tags like <physical_index_X> and <physical_index_X> to indicate the physical location of the page X.
+
+    The structure variable is the numeric system which represents the index of the hierarchy section in the table of contents. For example, the first section has structure index 1, the first subsection has structure index 1.1, the second subsection has structure index 1.2, etc.
+
+    The response should be in the following JSON format: 
+    [
+        {
+            "structure": <structure index, "x.x.x" or None> (string),
+            "title": <title of the section>,
+            "physical_index": "<physical_index_X>" (keep the format)
+        },
+        ...
+    ]
+
+    Only add the physical_index to the sections that are in the provided pages.
+    If the section is not in the provided pages, do not add the physical_index to it.
+    Directly return the final JSON structure. Do not output anything else."""
+
+    prompt = toc_extractor_prompt + '\nTable of contents:\n' + str(toc) + '\nDocument pages:\n' + content
+    response = ChatGPT_API(model=model, prompt=prompt)
+    json_content = extract_json(response)
+    return json_content
+
+
+def add_page_number_to_toc(part, structure, model=None):
+    """Fill page numbers from document pages into TOC structure."""
+    fill_prompt_seq = """
+    You are given an JSON structure of a document and a partial part of the document. Your task is to check if the title that is described in the structure is started in the partial given document.
+
+    The provided text contains tags like <physical_index_X> and <physical_index_X> to indicate the physical location of the page X. 
+
+    If the full target section starts in the partial given document, insert the given JSON structure with the "start": "yes", and "start_index": "<physical_index_X>".
+
+    If the full target section does not start in the partial given document, insert "start": "no",  "start_index": None.
+
+    The response should be in the following format. 
+        [
+            {
+                "structure": <structure index, "x.x.x" or None> (string),
+                "title": <title of the section>,
+                "start": "<yes or no>",
+                "physical_index": "<physical_index_X> (keep the format)" or None
+            },
+            ...
+        ]    
+    The given structure contains the result of the previous part, you need to fill the result of the current part, do not change the previous result.
+    Directly return the final JSON structure. Do not output anything else."""
+
+    if isinstance(part, list):
+        part = ''.join(part)
+
+    prompt = fill_prompt_seq + f"\n\nCurrent Partial Document:\n{part}\n\nGiven Structure\n{json.dumps(structure, indent=2)}\n"
+    current_json_raw = ChatGPT_API(model=model, prompt=prompt)
+    json_result = extract_json(current_json_raw)
+
+    if isinstance(json_result, list):
+        for item in json_result:
+            if 'start' in item:
+                del item['start']
+    return json_result
+
+
+def remove_page_number(data):
+    """Remove page_number field from data recursively."""
+    if isinstance(data, dict):
+        data.pop('page_number', None)
+        data.pop('page', None)
+        for key in list(data.keys()):
+            if 'nodes' in key:
+                remove_page_number(data[key])
+    elif isinstance(data, list):
+        for item in data:
+            remove_page_number(item)
+    return data
+
+
+def extract_matching_page_pairs(toc_page, toc_physical_index, start_page_index):
+    """Match TOC items by title and extract page/physical_index pairs."""
+    pairs = []
+    for phy_item in toc_physical_index:
+        if not isinstance(phy_item, dict):
+            continue
+        for page_item in toc_page:
+            if not isinstance(page_item, dict):
+                continue
+            if phy_item.get('title') == page_item.get('title'):
+                physical_index = phy_item.get('physical_index')
+                if physical_index is not None and int(physical_index) >= start_page_index:
+                    pairs.append({
+                        'title': phy_item.get('title'),
+                        'page': page_item.get('page'),
+                        'physical_index': physical_index
+                    })
+    return pairs
+
+
+def calculate_page_offset(pairs):
+    """Calculate the most common page offset from matching pairs."""
+    differences = []
+    for pair in pairs:
+        try:
+            physical_index = pair['physical_index']
+            page_number = pair['page']
+            difference = physical_index - page_number
+            differences.append(difference)
+        except (KeyError, TypeError):
+            continue
+
+    if not differences:
+        return 0
+
+    difference_counts = {}
+    for diff in differences:
+        difference_counts[diff] = difference_counts.get(diff, 0) + 1
+
+    most_common = max(difference_counts.items(), key=lambda x: x[1])[0]
+    return most_common
+
+
+def add_page_offset_to_toc_json(data, offset):
+    """Apply page offset to convert logical page numbers to physical indices."""
+    for i in range(len(data)):
+        if data[i].get('page') is not None and isinstance(data[i]['page'], int):
+            data[i]['physical_index'] = data[i]['page'] + offset
+            del data[i]['page']
+    return data
+
+
+def process_none_page_numbers(toc_items, page_list, start_index=1, model=None):
+    """Fix items that are missing physical_index by searching nearby pages."""
+    for i, item in enumerate(toc_items):
+        if "physical_index" not in item:
+            # Find previous physical_index
+            prev_physical_index = 0
+            for j in range(i - 1, -1, -1):
+                if toc_items[j].get('physical_index') is not None:
+                    prev_physical_index = toc_items[j]['physical_index']
+                    break
+
+            # Find next physical_index
+            next_physical_index = -1
+            for j in range(i + 1, len(toc_items)):
+                if toc_items[j].get('physical_index') is not None:
+                    next_physical_index = toc_items[j]['physical_index']
+                    break
+
+            if next_physical_index == -1:
+                next_physical_index = len(page_list) + start_index - 1
+
+            page_contents = []
+            for page_index in range(prev_physical_index, next_physical_index + 1):
+                list_index = page_index - start_index
+                if 0 <= list_index < len(page_list):
+                    page_text = f"<physical_index_{page_index}>\n{page_list[list_index][0]}\n<physical_index_{page_index}>\n\n"
+                    page_contents.append(page_text)
+
+            item_copy = copy.deepcopy(item)
+            item_copy.pop('page', None)
+            result = add_page_number_to_toc(page_contents, item_copy, model)
+
+            if isinstance(result, list) and len(result) > 0:
+                pi = result[0].get('physical_index')
+                if isinstance(pi, str) and pi.startswith('<physical_index'):
+                    item['physical_index'] = int(pi.split('_')[-1].rstrip('>').strip())
+                    item.pop('page', None)
+
+    return toc_items
+
+
+def single_toc_item_index_fixer(section_title, content, model=None):
+    """Fix a single TOC item's physical index using LLM."""
+    toc_extractor_prompt = """
+    You are given a section title and several pages of a document, your job is to find the physical index of the start page of the section in the partial document.
+
+    The provided pages contains tags like <physical_index_X> and <physical_index_X> to indicate the physical location of the page X.
+
+    Reply in a JSON format:
+    {
+        "thinking": <explain which page, started and closed by <physical_index_X>, contains the start of this section>,
+        "physical_index": "<physical_index_X>" (keep the format)
+    }
+    Directly return the final JSON structure. Do not output anything else."""
+
+    prompt = toc_extractor_prompt + '\nSection Title:\n' + str(section_title) + '\nDocument pages:\n' + content
+    response = ChatGPT_API(model=model, prompt=prompt)
+    json_content = extract_json(response)
+    return convert_physical_index_to_int(json_content.get('physical_index', ''))

src/toc_rag/toc_utils.py

@@ -0,0 +1,459 @@
+"""
+toc_utils.py
+------------
+Utility functions for the TOC-based Index RAG.
+Handles token counting, JSON extraction, LLM wrappers (Ollama),
+structure helpers, logging, and config management.
+
+Adapted from the reference page-indexing repository to use
+local Ollama instead of OpenAI.
+"""
+
+import os
+import re
+import json
+import copy
+import asyncio
+import tiktoken
+from pathlib import Path
+from io import BytesIO
+from collections import OrderedDict
+
+import fitz  # PyMuPDF
+import ollama
+
+from src.cost_tracker import add_toc_analysis_tokens
+
+# ─── Token Counting ───────────────────────────────────────────────────────────
+
+_tokenizer = tiktoken.get_encoding("cl100k_base")
+
+LLM_MODEL = "gpt-oss:20b-cloud"
+
+
+def count_tokens(text, model=None):
+    """Count tokens in text using tiktoken."""
+    if not text:
+        return 0
+    return len(_tokenizer.encode(str(text)))
+
+
+# ─── JSON Extraction ──────────────────────────────────────────────────────────
+
+def extract_json(text):
+    """Extract JSON from LLM response text, handling markdown code blocks."""
+    if not text:
+        return {}
+    text = text.strip()
+    # Remove markdown code block wrappers
+    if text.startswith("```json"):
+        text = text[7:]
+    elif text.startswith("```"):
+        text = text[3:]
+    if text.endswith("```"):
+        text = text[:-3]
+    text = text.strip()
+
+    try:
+        return json.loads(text)
+    except json.JSONDecodeError:
+        # Try to find JSON in the text
+        # Look for array
+        match = re.search(r'\[.*\]', text, re.DOTALL)
+        if match:
+            try:
+                return json.loads(match.group())
+            except json.JSONDecodeError:
+                pass
+        # Look for object
+        match = re.search(r'\{.*\}', text, re.DOTALL)
+        if match:
+            try:
+                return json.loads(match.group())
+            except json.JSONDecodeError:
+                pass
+    return {}
+
+
+def get_json_content(text):
+    """Extract raw JSON content string from markdown-wrapped response."""
+    if not text:
+        return text
+    text = text.strip()
+    if text.startswith("```json"):
+        text = text[7:]
+    elif text.startswith("```"):
+        text = text[3:]
+    if text.endswith("```"):
+        text = text[:-3]
+    return text.strip()
+
+
+# ─── LLM Wrappers (Ollama) ────────────────────────────────────────────────────
+
+def ChatGPT_API(model=None, prompt=None, chat_history=None):
+    """Synchronous LLM call via Ollama. Returns response text."""
+    model = model or LLM_MODEL
+    messages = []
+    if chat_history:
+        messages.extend(chat_history)
+    messages.append({"role": "user", "content": prompt})
+
+    response = ollama.chat(model=model, messages=messages)
+    content = response.get("message", {}).get("content", "")
+
+    # Track tokens
+    prompt_tokens = response.get("prompt_eval_count", 0)
+    completion_tokens = response.get("eval_count", 0)
+    add_toc_analysis_tokens(prompt_tokens + completion_tokens)
+
+    return content
+
+
+def ChatGPT_API_with_finish_reason(model=None, prompt=None, chat_history=None):
+    """Synchronous LLM call that also returns finish reason."""
+    model = model or LLM_MODEL
+    messages = []
+    if chat_history:
+        messages.extend(chat_history)
+    messages.append({"role": "user", "content": prompt})
+
+    response = ollama.chat(model=model, messages=messages)
+    content = response.get("message", {}).get("content", "")
+
+    # Track tokens
+    prompt_tokens = response.get("prompt_eval_count", 0)
+    completion_tokens = response.get("eval_count", 0)
+    add_toc_analysis_tokens(prompt_tokens + completion_tokens)
+
+    # Ollama always finishes (no length limit truncation like OpenAI)
+    finish_reason = "finished"
+    if response.get("done_reason") == "length":
+        finish_reason = "length"
+
+    return content, finish_reason
+
+
+async def ChatGPT_API_async(model=None, prompt=None):
+    """Async LLM call via Ollama (runs sync call in executor)."""
+    loop = asyncio.get_event_loop()
+    return await loop.run_in_executor(None, ChatGPT_API, model, prompt, None)
+
+
+# ─── Type Converters ──────────────────────────────────────────────────────────
+
+def convert_page_to_int(data):
+    """Convert 'page' fields to int in TOC data."""
+    if isinstance(data, list):
+        for item in data:
+            if isinstance(item, dict) and 'page' in item:
+                try:
+                    if item['page'] is not None:
+                        item['page'] = int(item['page'])
+                except (ValueError, TypeError):
+                    item['page'] = None
+    return data
+
+
+def convert_physical_index_to_int(data):
+    """Convert physical_index from '<physical_index_X>' format to int."""
+    if isinstance(data, str):
+        # Single value
+        match = re.search(r'physical_index_(\d+)', data)
+        if match:
+            return int(match.group(1))
+        try:
+            return int(data)
+        except (ValueError, TypeError):
+            return None
+
+    if isinstance(data, list):
+        for item in data:
+            if isinstance(item, dict) and 'physical_index' in item:
+                pi = item['physical_index']
+                if isinstance(pi, str):
+                    match = re.search(r'physical_index_(\d+)', pi)
+                    if match:
+                        item['physical_index'] = int(match.group(1))
+                    else:
+                        try:
+                            item['physical_index'] = int(pi)
+                        except (ValueError, TypeError):
+                            item['physical_index'] = None
+                elif isinstance(pi, (int, float)):
+                    item['physical_index'] = int(pi)
+    return data
+
+
+# ─── Structure Helpers ─────────────────────────────────────────────────────────
+
+def post_processing(toc_items, end_index):
+    """Convert flat TOC list into tree with start_index/end_index."""
+    if not toc_items:
+        return []
+
+    result = []
+    for i, item in enumerate(toc_items):
+        node = {
+            'title': item.get('title', ''),
+            'start_index': item.get('physical_index'),
+        }
+        if item.get('structure'):
+            node['structure'] = item['structure']
+        if item.get('appear_start'):
+            node['appear_start'] = item['appear_start']
+
+        # End index is the start of the next section, or the document end
+        if i + 1 < len(toc_items):
+            node['end_index'] = toc_items[i + 1].get('physical_index', end_index)
+        else:
+            node['end_index'] = end_index
+
+        result.append(node)
+
+    return result
+
+
+def add_preface_if_needed(toc_items):
+    """Insert a preface node if first section doesn't start at page 1."""
+    if not toc_items:
+        return toc_items
+
+    first_pi = toc_items[0].get('physical_index')
+    if first_pi is not None and first_pi > 1:
+        preface = {
+            'structure': None,
+            'title': 'Preface',
+            'physical_index': 1,
+            'appear_start': 'yes',
+        }
+        toc_items.insert(0, preface)
+
+    return toc_items
+
+
+def write_node_id(structure, counter=None):
+    """Assign sequential node IDs to tree nodes."""
+    if counter is None:
+        counter = [1]
+
+    for node in structure:
+        node['node_id'] = str(counter[0]).zfill(4)
+        counter[0] += 1
+        if 'nodes' in node and node['nodes']:
+            write_node_id(node['nodes'], counter)
+
+
+def add_node_text(structure, page_list):
+    """Populate 'text' field from page_list for each node."""
+    for node in structure:
+        start = node.get('start_index', 1)
+        end = node.get('end_index', start)
+        texts = []
+        for pi in range(start, end + 1):
+            idx = pi - 1
+            if 0 <= idx < len(page_list):
+                texts.append(page_list[idx][0])
+        node['text'] = '\n'.join(texts)
+
+        if 'nodes' in node and node['nodes']:
+            add_node_text(node['nodes'], page_list)
+
+
+def remove_structure_text(structure):
+    """Remove 'text' field from structure."""
+    for node in structure:
+        node.pop('text', None)
+        if 'nodes' in node and node['nodes']:
+            remove_structure_text(node['nodes'])
+
+
+def format_structure(structure, order=None):
+    """Reorder dict keys in structure according to given order."""
+    if order is None:
+        return structure
+
+    result = []
+    for node in structure:
+        ordered = OrderedDict()
+        for key in order:
+            if key == 'nodes':
+                if 'nodes' in node and node['nodes']:
+                    ordered['nodes'] = format_structure(node['nodes'], order)
+            elif key in node:
+                ordered[key] = node[key]
+        result.append(dict(ordered))
+    return result
+
+
+def structure_to_list(structure):
+    """Flatten tree structure to a flat list of nodes."""
+    result = []
+    for node in structure:
+        result.append(node)
+        if 'nodes' in node and node['nodes']:
+            result.extend(structure_to_list(node['nodes']))
+    return result
+
+
+def create_clean_structure_for_description(structure):
+    """Create a minimal structure for document description generation."""
+    clean = []
+    for node in structure:
+        item = {'title': node.get('title', '')}
+        if 'summary' in node:
+            item['summary'] = node['summary']
+        if 'prefix_summary' in node:
+            item['prefix_summary'] = node['prefix_summary']
+        if 'nodes' in node and node['nodes']:
+            item['nodes'] = create_clean_structure_for_description(node['nodes'])
+        clean.append(item)
+    return clean
+
+
+# ─── Summary Generation ───────────────────────────────────────────────────────
+
+async def generate_node_summary(node, model=None):
+    """Generate a summary for a single node using LLM."""
+    title = node.get('title', '')
+    text = node.get('text', '')
+
+    prompt = f"""Summarize the following section concisely (2-3 sentences max).
+
+Section title: {title}
+Section content:
+{text[:8000]}
+
+Provide only the summary, no extra formatting."""
+
+    return await ChatGPT_API_async(model=model, prompt=prompt)
+
+
+async def generate_summaries_for_structure(structure, model=None):
+    """Generate summaries for all nodes in the structure."""
+    nodes = structure_to_list(structure)
+    nodes_with_text = [n for n in nodes if n.get('text')]
+
+    if not nodes_with_text:
+        return structure
+
+    summary_tasks = [generate_node_summary(n, model=model) for n in nodes_with_text]
+    summaries = await asyncio.gather(*summary_tasks, return_exceptions=True)
+
+    for node, summary in zip(nodes_with_text, summaries):
+        if isinstance(summary, Exception):
+            print(f"[TOC-RAG] Summary error for '{node.get('title', '')}': {summary}")
+            node['summary'] = ''
+        else:
+            if not node.get('nodes'):
+                node['summary'] = summary
+            else:
+                node['prefix_summary'] = summary
+
+    return structure
+
+
+def generate_doc_description(structure, model=None):
+    """Generate a document-level description from the structure."""
+    prompt = f"""Based on the following document structure, provide a brief description of
+what this document is about (2-3 sentences).
+
+Document structure:
+{json.dumps(structure, indent=2)[:6000]}
+
+Provide only the description."""
+
+    return ChatGPT_API(model=model, prompt=prompt)
+
+
+# ─── PDF Helpers ───────────────────────────────────────────────────────────────
+
+def get_page_tokens(doc):
+    """Extract pages from PDF. Returns list of (text, token_count) tuples."""
+    if isinstance(doc, str):
+        pdf_doc = fitz.open(doc)
+    elif isinstance(doc, BytesIO):
+        pdf_doc = fitz.open(stream=doc, filetype="pdf")
+    else:
+        raise ValueError("Expected file path or BytesIO")
+
+    page_list = []
+    for page in pdf_doc:
+        text = page.get_text("text").strip()
+        tokens = count_tokens(text)
+        page_list.append((text, tokens))
+
+    return page_list
+
+
+def get_pdf_name(doc):
+    """Get the filename from a PDF path or BytesIO."""
+    if isinstance(doc, str):
+        return Path(doc).stem
+    return "document"
+
+
+# ─── Logger ────────────────────────────────────────────────────────────────────
+
+class JsonLogger:
+    """Simple logger that prints structured JSON logs."""
+
+    def __init__(self, doc=None):
+        self.doc_name = get_pdf_name(doc) if doc else "unknown"
+
+    def info(self, msg):
+        if isinstance(msg, (dict, list)):
+            print(f"[TOC-RAG][{self.doc_name}] {json.dumps(msg, default=str)[:500]}")
+        else:
+            print(f"[TOC-RAG][{self.doc_name}] {str(msg)[:500]}")
+
+    def error(self, msg):
+        print(f"[TOC-RAG][ERROR][{self.doc_name}] {str(msg)[:500]}")
+
+
+# ─── Config ───────────────────────────────────────────────────────────────────
+
+class _Config:
+    """Simple config object with attribute access."""
+    def __init__(self, **kwargs):
+        for k, v in kwargs.items():
+            setattr(self, k, v)
+
+
+class ConfigLoader:
+    """Load config from defaults + user overrides."""
+
+    DEFAULTS = {
+        "model": LLM_MODEL,
+        "toc_check_page_num": 20,
+        "max_page_num_each_node": 10,
+        "max_token_num_each_node": 20000,
+        "if_add_node_id": "yes",
+        "if_add_node_summary": "yes",
+        "if_add_doc_description": "no",
+        "if_add_node_text": "no",
+    }
+
+    def load(self, user_opt=None):
+        merged = dict(self.DEFAULTS)
+        if user_opt:
+            merged.update({k: v for k, v in user_opt.items() if v is not None})
+        return _Config(**merged)
+
+
+# ─── Debug Printing ───────────────────────────────────────────────────────────
+
+def print_json(data):
+    """Pretty print JSON data."""
+    print(json.dumps(data, indent=2, ensure_ascii=False, default=str))
+
+
+def print_toc(structure, indent=0):
+    """Print tree table of contents."""
+    for node in structure:
+        prefix = "  " * indent
+        title = node.get('title', '')
+        node_id = node.get('node_id', '')
+        print(f"{prefix}{node_id} {title}")
+        if 'nodes' in node and node['nodes']:
+            print_toc(node['nodes'], indent + 1)

src/toc_rag/tree_builder.py

@@ -0,0 +1,716 @@
+"""
+tree_builder.py
+---------------
+Tree construction, async verification, fix loops, and recursive splitting.
+Main async pipeline for building document structure trees from PDFs.
+
+Replica of the reference page-indexing repository,
+adapted for local Ollama LLM.
+"""
+
+import json
+import copy
+import math
+import random
+import asyncio
+
+from .toc_utils import (
+    ChatGPT_API, ChatGPT_API_async, ChatGPT_API_with_finish_reason,
+    extract_json, count_tokens, convert_physical_index_to_int,
+    post_processing, add_preface_if_needed,
+    write_node_id, add_node_text, remove_structure_text,
+    format_structure, structure_to_list,
+    create_clean_structure_for_description,
+    generate_summaries_for_structure, generate_doc_description,
+    get_page_tokens, get_pdf_name,
+    JsonLogger, ConfigLoader,
+)
+from .toc_detector import check_toc
+from .toc_transformer import (
+    toc_transformer, toc_index_extractor, add_page_number_to_toc,
+    remove_page_number, extract_matching_page_pairs,
+    calculate_page_offset, add_page_offset_to_toc_json,
+    process_none_page_numbers, single_toc_item_index_fixer,
+)
+
+
+# ─── Title Verification (async) ───────────────────────────────────────────────
+
+async def check_title_appearance(item, page_list, start_index=1, model=None):
+    """Check if a section title appears on its indicated page."""
+    title = item['title']
+    if 'physical_index' not in item or item['physical_index'] is None:
+        return {'list_index': item.get('list_index'), 'answer': 'no', 'title': title, 'page_number': None}
+
+    page_number = item['physical_index']
+    idx = page_number - start_index
+    if idx < 0 or idx >= len(page_list):
+        return {'list_index': item.get('list_index'), 'answer': 'no', 'title': title, 'page_number': page_number}
+
+    page_text = page_list[idx][0]
+
+    prompt = f"""
+    Your job is to check if the given section appears or starts in the given page_text.
+
+    Note: do fuzzy matching, ignore any space inconsistency in the page_text.
+
+    The given section title is {title}.
+    The given page_text is {page_text}.
+    
+    Reply format:
+    {{
+        "thinking": <why do you think the section appears or starts in the page_text>
+        "answer": "yes or no" (yes if the section appears or starts in the page_text, no otherwise)
+    }}
+    Directly return the final JSON structure. Do not output anything else."""
+
+    response = await ChatGPT_API_async(model=model, prompt=prompt)
+    response = extract_json(response)
+    answer = response.get('answer', 'no')
+    return {'list_index': item.get('list_index'), 'answer': answer, 'title': title, 'page_number': page_number}
+
+
+async def check_title_appearance_in_start(title, page_text, model=None, logger=None):
+    """Check if a section starts at the beginning of a page."""
+    prompt = f"""
+    You will be given the current section title and the current page_text.
+    Your job is to check if the current section starts in the beginning of the given page_text.
+    If there are other contents before the current section title, then the current section does not start in the beginning of the given page_text.
+    If the current section title is the first content in the given page_text, then the current section starts in the beginning of the given page_text.
+
+    Note: do fuzzy matching, ignore any space inconsistency in the page_text.
+
+    The given section title is {title}.
+    The given page_text is {page_text}.
+    
+    reply format:
+    {{
+        "thinking": <why do you think the section appears or starts in the page_text>
+        "start_begin": "yes or no" (yes if the section starts in the beginning of the page_text, no otherwise)
+    }}
+    Directly return the final JSON structure. Do not output anything else."""
+
+    response = await ChatGPT_API_async(model=model, prompt=prompt)
+    response = extract_json(response)
+    if logger:
+        logger.info(f"Response: {response}")
+    return response.get("start_begin", "no")
+
+
+async def check_title_appearance_in_start_concurrent(structure, page_list, model=None, logger=None):
+    """Check all titles for start-of-page appearance concurrently."""
+    if logger:
+        logger.info("Checking title appearance in start concurrently")
+
+    for item in structure:
+        if item.get('physical_index') is None:
+            item['appear_start'] = 'no'
+
+    tasks = []
+    valid_items = []
+    for item in structure:
+        if item.get('physical_index') is not None:
+            idx = item['physical_index'] - 1
+            if 0 <= idx < len(page_list):
+                page_text = page_list[idx][0]
+                tasks.append(check_title_appearance_in_start(item['title'], page_text, model=model, logger=logger))
+                valid_items.append(item)
+            else:
+                item['appear_start'] = 'no'
+
+    if tasks:
+        results = await asyncio.gather(*tasks, return_exceptions=True)
+        for item, result in zip(valid_items, results):
+            if isinstance(result, Exception):
+                if logger:
+                    logger.error(f"Error checking start for {item['title']}: {result}")
+                item['appear_start'] = 'no'
+            else:
+                item['appear_start'] = result
+
+    return structure
+
+
+# ─── Page Grouping ─────────────────────────────────────────────────────────────
+
+def page_list_to_group_text(page_contents, token_lengths, max_tokens=20000, overlap_page=1):
+    """Group pages into text chunks respecting token limits."""
+    num_tokens = sum(token_lengths)
+
+    if num_tokens <= max_tokens:
+        page_text = "".join(page_contents)
+        return [page_text]
+
+    subsets = []
+    current_subset = []
+    current_token_count = 0
+
+    expected_parts_num = math.ceil(num_tokens / max_tokens)
+    average_tokens_per_part = math.ceil(((num_tokens / expected_parts_num) + max_tokens) / 2)
+
+    for i, (page_content, page_tokens) in enumerate(zip(page_contents, token_lengths)):
+        if current_token_count + page_tokens > average_tokens_per_part:
+            subsets.append(''.join(current_subset))
+            overlap_start = max(i - overlap_page, 0)
+            current_subset = list(page_contents[overlap_start:i])
+            current_token_count = sum(token_lengths[overlap_start:i])
+
+        current_subset.append(page_content)
+        current_token_count += page_tokens
+
+    if current_subset:
+        subsets.append(''.join(current_subset))
+
+    print('divide page_list to groups', len(subsets))
+    return subsets
+
+
+# ─── TOC Generation (no-TOC documents) ────────────────────────────────────────
+
+def generate_toc_init(part, model=None):
+    """Generate initial TOC structure from document text using LLM."""
+    print('start generate_toc_init')
+    prompt = """
+    You are an expert in extracting hierarchical tree structure, your task is to generate the tree structure of the document.
+
+    The structure variable is the numeric system which represents the index of the hierarchy section in the table of contents. For example, the first section has structure index 1, the first subsection has structure index 1.1, the second subsection has structure index 1.2, etc.
+
+    For the title, you need to extract the original title from the text, only fix the space inconsistency.
+
+    The provided text contains tags like <physical_index_X> and <physical_index_X> to indicate the start and end of page X. 
+
+    For the physical_index, you need to extract the physical index of the start of the section from the text. Keep the <physical_index_X> format.
+
+    The response should be in the following format. 
+        [
+            {
+                "structure": <structure index, "x.x.x"> (string),
+                "title": <title of the section, keep the original title>,
+                "physical_index": "<physical_index_X> (keep the format)"
+            },
+        ]
+
+    Directly return the final JSON structure. Do not output anything else."""
+
+    prompt = prompt + '\nGiven text\n:' + part
+    response, finish_reason = ChatGPT_API_with_finish_reason(model=model, prompt=prompt)
+
+    if finish_reason == 'finished':
+        return extract_json(response)
+    else:
+        raise Exception(f'finish reason: {finish_reason}')
+
+
+def generate_toc_continue(toc_content, part, model=None):
+    """Continue TOC generation for additional document parts."""
+    print('start generate_toc_continue')
+    prompt = """
+    You are an expert in extracting hierarchical tree structure.
+    You are given a tree structure of the previous part and the text of the current part.
+    Your task is to continue the tree structure from the previous part to include the current part.
+
+    The structure variable is the numeric system which represents the index of the hierarchy section in the table of contents. For example, the first section has structure index 1, the first subsection has structure index 1.1, the second subsection has structure index 1.2, etc.
+
+    For the title, you need to extract the original title from the text, only fix the space inconsistency.
+
+    The provided text contains tags like <physical_index_X> and <physical_index_X> to indicate the start and end of page X.
+    
+    For the physical_index, you need to extract the physical index of the start of the section from the text. Keep the <physical_index_X> format.
+
+    The response should be in the following format. 
+        [
+            {
+                "structure": <structure index, "x.x.x"> (string),
+                "title": <title of the section, keep the original title>,
+                "physical_index": "<physical_index_X> (keep the format)"
+            },
+            ...
+        ]    
+
+    Directly return the additional part of the final JSON structure. Do not output anything else."""
+
+    prompt = prompt + '\nGiven text\n:' + part + '\nPrevious tree structure\n:' + json.dumps(toc_content, indent=2)
+    response, finish_reason = ChatGPT_API_with_finish_reason(model=model, prompt=prompt)
+    if finish_reason == 'finished':
+        return extract_json(response)
+    else:
+        raise Exception(f'finish reason: {finish_reason}')
+
+
+# ─── Processing Pipelines ─────────────────────────────────────────────────────
+
+def process_no_toc(page_list, start_index=1, model=None, logger=None):
+    """Full pipeline for documents without a TOC."""
+    page_contents = []
+    token_lengths = []
+    for page_index in range(start_index, start_index + len(page_list)):
+        page_text = f"<physical_index_{page_index}>\n{page_list[page_index-start_index][0]}\n<physical_index_{page_index}>\n\n"
+        page_contents.append(page_text)
+        token_lengths.append(count_tokens(page_text, model))
+    group_texts = page_list_to_group_text(page_contents, token_lengths)
+    if logger:
+        logger.info(f'len(group_texts): {len(group_texts)}')
+
+    toc_with_page_number = generate_toc_init(group_texts[0], model)
+    for group_text in group_texts[1:]:
+        toc_with_page_number_additional = generate_toc_continue(toc_with_page_number, group_text, model)
+        toc_with_page_number.extend(toc_with_page_number_additional)
+    if logger:
+        logger.info(f'generate_toc: {toc_with_page_number}')
+
+    toc_with_page_number = convert_physical_index_to_int(toc_with_page_number)
+    if logger:
+        logger.info(f'convert_physical_index_to_int: {toc_with_page_number}')
+
+    return toc_with_page_number
+
+
+def process_toc_no_page_numbers(toc_content, toc_page_list, page_list, start_index=1, model=None, logger=None):
+    """Pipeline for documents with TOC but no page numbers."""
+    page_contents = []
+    token_lengths = []
+    toc_content = toc_transformer(toc_content, model)
+    if logger:
+        logger.info(f'toc_transformer: {toc_content}')
+    for page_index in range(start_index, start_index + len(page_list)):
+        page_text = f"<physical_index_{page_index}>\n{page_list[page_index-start_index][0]}\n<physical_index_{page_index}>\n\n"
+        page_contents.append(page_text)
+        token_lengths.append(count_tokens(page_text, model))
+
+    group_texts = page_list_to_group_text(page_contents, token_lengths)
+    if logger:
+        logger.info(f'len(group_texts): {len(group_texts)}')
+
+    toc_with_page_number = copy.deepcopy(toc_content)
+    for group_text in group_texts:
+        toc_with_page_number = add_page_number_to_toc(group_text, toc_with_page_number, model)
+    if logger:
+        logger.info(f'add_page_number_to_toc: {toc_with_page_number}')
+
+    toc_with_page_number = convert_physical_index_to_int(toc_with_page_number)
+    if logger:
+        logger.info(f'convert_physical_index_to_int: {toc_with_page_number}')
+
+    return toc_with_page_number
+
+
+def process_toc_with_page_numbers(toc_content, toc_page_list, page_list, toc_check_page_num=None, model=None, logger=None):
+    """Pipeline for documents with TOC and page numbers."""
+    toc_with_page_number = toc_transformer(toc_content, model)
+    if logger:
+        logger.info(f'toc_with_page_number: {toc_with_page_number}')
+
+    toc_no_page_number = remove_page_number(copy.deepcopy(toc_with_page_number))
+
+    start_page_index = toc_page_list[-1] + 1
+    main_content = ""
+    for page_index in range(start_page_index, min(start_page_index + (toc_check_page_num or 20), len(page_list))):
+        main_content += f"<physical_index_{page_index+1}>\n{page_list[page_index][0]}\n<physical_index_{page_index+1}>\n\n"
+
+    toc_with_physical_index = toc_index_extractor(toc_no_page_number, main_content, model)
+    if logger:
+        logger.info(f'toc_with_physical_index: {toc_with_physical_index}')
+
+    toc_with_physical_index = convert_physical_index_to_int(toc_with_physical_index)
+    if logger:
+        logger.info(f'toc_with_physical_index: {toc_with_physical_index}')
+
+    matching_pairs = extract_matching_page_pairs(toc_with_page_number, toc_with_physical_index, start_page_index)
+    if logger:
+        logger.info(f'matching_pairs: {matching_pairs}')
+
+    offset = calculate_page_offset(matching_pairs)
+    if logger:
+        logger.info(f'offset: {offset}')
+
+    toc_with_page_number = add_page_offset_to_toc_json(toc_with_page_number, offset)
+    if logger:
+        logger.info(f'toc_with_page_number: {toc_with_page_number}')
+
+    toc_with_page_number = process_none_page_numbers(toc_with_page_number, page_list, model=model)
+    if logger:
+        logger.info(f'toc_with_page_number: {toc_with_page_number}')
+
+    return toc_with_page_number
+
+
+# ─── Validation ─────────────────────────────────────────────────────────��──────
+
+def validate_and_truncate_physical_indices(toc_with_page_number, page_list_length, start_index=1, logger=None):
+    """Validate and truncate physical indices exceeding document length."""
+    if not toc_with_page_number:
+        return toc_with_page_number
+
+    max_allowed_page = page_list_length + start_index - 1
+    truncated_items = []
+
+    for i, item in enumerate(toc_with_page_number):
+        if item.get('physical_index') is not None:
+            original_index = item['physical_index']
+            if original_index > max_allowed_page:
+                item['physical_index'] = None
+                truncated_items.append({
+                    'title': item.get('title', 'Unknown'),
+                    'original_index': original_index
+                })
+                if logger:
+                    logger.info(f"Removed physical_index for '{item.get('title', 'Unknown')}' (was {original_index})")
+
+    if truncated_items and logger:
+        logger.info(f"Total removed items: {len(truncated_items)}")
+
+    print(f"Document validation: {page_list_length} pages, max allowed index: {max_allowed_page}")
+    if truncated_items:
+        print(f"Truncated {len(truncated_items)} TOC items that exceeded document length")
+
+    return toc_with_page_number
+
+
+# ─── Verification ─────────────────────────────────────────────────────────────
+
+async def verify_toc(page_list, list_result, start_index=1, N=None, model=None):
+    """Verify TOC accuracy by checking title appearances on pages."""
+    print('start verify_toc')
+    last_physical_index = None
+    for item in reversed(list_result):
+        if item.get('physical_index') is not None:
+            last_physical_index = item['physical_index']
+            break
+
+    if last_physical_index is None or last_physical_index < len(page_list) / 2:
+        return 0, []
+
+    if N is None:
+        sample_indices = range(0, len(list_result))
+    else:
+        N = min(N, len(list_result))
+        sample_indices = random.sample(range(0, len(list_result)), N)
+
+    indexed_sample_list = []
+    for idx in sample_indices:
+        item = list_result[idx]
+        if item.get('physical_index') is not None:
+            item_with_index = item.copy()
+            item_with_index['list_index'] = idx
+            indexed_sample_list.append(item_with_index)
+
+    tasks = [
+        check_title_appearance(item, page_list, start_index, model)
+        for item in indexed_sample_list
+    ]
+    results = await asyncio.gather(*tasks)
+
+    correct_count = 0
+    incorrect_results = []
+    for result in results:
+        if result['answer'] == 'yes':
+            correct_count += 1
+        else:
+            incorrect_results.append(result)
+
+    checked_count = len(results)
+    accuracy = correct_count / checked_count if checked_count > 0 else 0
+    print(f"accuracy: {accuracy*100:.2f}%")
+    return accuracy, incorrect_results
+
+
+# ─── Fix Loops ─────────────────────────────────────────────────────────────────
+
+async def fix_incorrect_toc(toc_with_page_number, page_list, incorrect_results, start_index=1, model=None, logger=None):
+    """Fix incorrect TOC entries by searching within nearby pages."""
+    print(f'start fix_incorrect_toc with {len(incorrect_results)} incorrect results')
+    incorrect_indices = {result['list_index'] for result in incorrect_results}
+    end_index = len(page_list) + start_index - 1
+
+    async def process_and_check_item(incorrect_item):
+        list_index = incorrect_item['list_index']
+
+        if list_index < 0 or list_index >= len(toc_with_page_number):
+            return {
+                'list_index': list_index,
+                'title': incorrect_item['title'],
+                'physical_index': incorrect_item.get('physical_index'),
+                'is_valid': False
+            }
+
+        prev_correct = None
+        for i in range(list_index - 1, -1, -1):
+            if i not in incorrect_indices and 0 <= i < len(toc_with_page_number):
+                pi = toc_with_page_number[i].get('physical_index')
+                if pi is not None:
+                    prev_correct = pi
+                    break
+        if prev_correct is None:
+            prev_correct = start_index - 1
+
+        next_correct = None
+        for i in range(list_index + 1, len(toc_with_page_number)):
+            if i not in incorrect_indices and 0 <= i < len(toc_with_page_number):
+                pi = toc_with_page_number[i].get('physical_index')
+                if pi is not None:
+                    next_correct = pi
+                    break
+        if next_correct is None:
+            next_correct = end_index
+
+        page_contents = []
+        for page_index in range(prev_correct, next_correct + 1):
+            li = page_index - start_index
+            if 0 <= li < len(page_list):
+                page_text = f"<physical_index_{page_index}>\n{page_list[li][0]}\n<physical_index_{page_index}>\n\n"
+                page_contents.append(page_text)
+        content_range = ''.join(page_contents)
+
+        physical_index_int = single_toc_item_index_fixer(incorrect_item['title'], content_range, model)
+
+        check_item = incorrect_item.copy()
+        check_item['physical_index'] = physical_index_int
+        check_result = await check_title_appearance(check_item, page_list, start_index, model)
+
+        return {
+            'list_index': list_index,
+            'title': incorrect_item['title'],
+            'physical_index': physical_index_int,
+            'is_valid': check_result['answer'] == 'yes'
+        }
+
+    tasks = [process_and_check_item(item) for item in incorrect_results]
+    results = await asyncio.gather(*tasks, return_exceptions=True)
+    results = [r for r in results if not isinstance(r, Exception)]
+
+    invalid_results = []
+    for result in results:
+        if result['is_valid']:
+            list_idx = result['list_index']
+            if 0 <= list_idx < len(toc_with_page_number):
+                toc_with_page_number[list_idx]['physical_index'] = result['physical_index']
+            else:
+                invalid_results.append(result)
+        else:
+            invalid_results.append(result)
+
+    if logger:
+        logger.info(f'invalid_results: {invalid_results}')
+
+    return toc_with_page_number, invalid_results
+
+
+async def fix_incorrect_toc_with_retries(toc_with_page_number, page_list, incorrect_results, start_index=1, max_attempts=3, model=None, logger=None):
+    """Fix incorrect TOC with retry loop."""
+    print('start fix_incorrect_toc')
+    fix_attempt = 0
+    current_toc = toc_with_page_number
+    current_incorrect = incorrect_results
+
+    while current_incorrect:
+        print(f"Fixing {len(current_incorrect)} incorrect results")
+        current_toc, current_incorrect = await fix_incorrect_toc(
+            current_toc, page_list, current_incorrect, start_index, model, logger
+        )
+        fix_attempt += 1
+        if fix_attempt >= max_attempts:
+            if logger:
+                logger.info("Maximum fix attempts reached")
+            break
+
+    return current_toc, current_incorrect
+
+
+# ─── Main Async Dispatcher ────────────────────────────────────────────────────
+
+async def meta_processor(page_list, mode=None, toc_content=None, toc_page_list=None, start_index=1, opt=None, logger=None):
+    """Main async dispatcher with fallback chain."""
+    print(mode)
+    print(f'start_index: {start_index}')
+
+    if mode == 'process_toc_with_page_numbers':
+        toc_with_page_number = process_toc_with_page_numbers(
+            toc_content, toc_page_list, page_list,
+            toc_check_page_num=opt.toc_check_page_num, model=opt.model, logger=logger
+        )
+    elif mode == 'process_toc_no_page_numbers':
+        toc_with_page_number = process_toc_no_page_numbers(
+            toc_content, toc_page_list, page_list, model=opt.model, logger=logger
+        )
+    else:
+        toc_with_page_number = process_no_toc(page_list, start_index=start_index, model=opt.model, logger=logger)
+
+    toc_with_page_number = [item for item in toc_with_page_number if item.get('physical_index') is not None]
+
+    toc_with_page_number = validate_and_truncate_physical_indices(
+        toc_with_page_number, len(page_list), start_index=start_index, logger=logger
+    )
+
+    accuracy, incorrect_results = await verify_toc(
+        page_list, toc_with_page_number, start_index=start_index, model=opt.model
+    )
+
+    if logger:
+        logger.info({
+            'mode': mode,
+            'accuracy': accuracy,
+            'incorrect_count': len(incorrect_results)
+        })
+
+    if accuracy == 1.0 and len(incorrect_results) == 0:
+        return toc_with_page_number
+    if accuracy > 0.6 and len(incorrect_results) > 0:
+        toc_with_page_number, _ = await fix_incorrect_toc_with_retries(
+            toc_with_page_number, page_list, incorrect_results,
+            start_index=start_index, max_attempts=3, model=opt.model, logger=logger
+        )
+        return toc_with_page_number
+    else:
+        if mode == 'process_toc_with_page_numbers':
+            return await meta_processor(
+                page_list, mode='process_toc_no_page_numbers',
+                toc_content=toc_content, toc_page_list=toc_page_list,
+                start_index=start_index, opt=opt, logger=logger
+            )
+        elif mode == 'process_toc_no_page_numbers':
+            return await meta_processor(
+                page_list, mode='process_no_toc',
+                start_index=start_index, opt=opt, logger=logger
+            )
+        else:
+            raise Exception('Processing failed')
+
+
+# ─── Recursive Large Node Splitting ────────────────���──────────────────────────
+
+async def process_large_node_recursively(node, page_list, opt=None, logger=None):
+    """Recursively split large nodes into sub-nodes."""
+    start_idx = node.get('start_index', 1)
+    end_idx = node.get('end_index', start_idx)
+    node_page_list = page_list[start_idx - 1:end_idx]
+    token_num = sum([page[1] for page in node_page_list])
+
+    if (end_idx - start_idx > opt.max_page_num_each_node and
+            token_num >= opt.max_token_num_each_node):
+        print(f'large node: {node["title"]} start: {start_idx} end: {end_idx} tokens: {token_num}')
+
+        node_toc_tree = await meta_processor(
+            node_page_list, mode='process_no_toc',
+            start_index=start_idx, opt=opt, logger=logger
+        )
+        node_toc_tree = await check_title_appearance_in_start_concurrent(
+            node_toc_tree, page_list, model=opt.model, logger=logger
+        )
+
+        valid_items = [item for item in node_toc_tree if item.get('physical_index') is not None]
+
+        if valid_items and node['title'].strip() == valid_items[0]['title'].strip():
+            node['nodes'] = post_processing(valid_items[1:], end_idx)
+            node['end_index'] = valid_items[1]['start_index'] if len(valid_items) > 1 else end_idx
+        else:
+            node['nodes'] = post_processing(valid_items, end_idx)
+            node['end_index'] = valid_items[0]['start_index'] if valid_items else end_idx
+
+    if 'nodes' in node and node['nodes']:
+        tasks = [
+            process_large_node_recursively(child, page_list, opt, logger=logger)
+            for child in node['nodes']
+        ]
+        await asyncio.gather(*tasks)
+
+    return node
+
+
+# ─── Top-Level Entry Points ───────────────────────────────────────────────────
+
+async def tree_parser(page_list, opt, doc=None, logger=None):
+    """Top-level async entry point for building the document tree."""
+    check_toc_result = check_toc(page_list, opt)
+    if logger:
+        logger.info(check_toc_result)
+
+    if (check_toc_result.get("toc_content") and
+            check_toc_result["toc_content"].strip() and
+            check_toc_result["page_index_given_in_toc"] == "yes"):
+        toc_with_page_number = await meta_processor(
+            page_list, mode='process_toc_with_page_numbers',
+            start_index=1,
+            toc_content=check_toc_result['toc_content'],
+            toc_page_list=check_toc_result['toc_page_list'],
+            opt=opt, logger=logger
+        )
+    else:
+        toc_with_page_number = await meta_processor(
+            page_list, mode='process_no_toc',
+            start_index=1, opt=opt, logger=logger
+        )
+
+    toc_with_page_number = add_preface_if_needed(toc_with_page_number)
+    toc_with_page_number = await check_title_appearance_in_start_concurrent(
+        toc_with_page_number, page_list, model=opt.model, logger=logger
+    )
+
+    valid_toc_items = [item for item in toc_with_page_number if item.get('physical_index') is not None]
+
+    toc_tree = post_processing(valid_toc_items, len(page_list))
+    tasks = [
+        process_large_node_recursively(node, page_list, opt, logger=logger)
+        for node in toc_tree
+    ]
+    await asyncio.gather(*tasks)
+
+    return toc_tree
+
+
+def page_index_main(doc, opt=None):
+    """Main synchronous entry point for PDF page indexing."""
+    from io import BytesIO
+    logger = JsonLogger(doc)
+
+    is_valid_pdf = (
+        (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)
+
+    logger.info({'total_page_number': len(page_list)})
+    logger.info({'total_token': sum([page[1] for page in page_list])})
+
+    async def page_index_builder():
+        structure = await tree_parser(page_list, opt, doc=doc, logger=logger)
+        if opt.if_add_node_id == 'yes':
+            write_node_id(structure)
+        if opt.if_add_node_text == 'yes':
+            add_node_text(structure, page_list)
+        if opt.if_add_node_summary == 'yes':
+            if opt.if_add_node_text == 'no':
+                add_node_text(structure, page_list)
+            await generate_summaries_for_structure(structure, model=opt.model)
+            if opt.if_add_node_text == 'no':
+                remove_structure_text(structure)
+            if opt.if_add_doc_description == 'yes':
+                clean_structure = create_clean_structure_for_description(structure)
+                doc_description = generate_doc_description(clean_structure, model=opt.model)
+                return {
+                    'doc_name': get_pdf_name(doc),
+                    'doc_description': doc_description,
+                    'structure': structure,
+                }
+        return {
+            'doc_name': get_pdf_name(doc),
+            'structure': structure,
+        }
+
+    return asyncio.run(page_index_builder())
+
+
+def page_index(doc, model=None, toc_check_page_num=None, max_page_num_each_node=None,
+               max_token_num_each_node=None, if_add_node_id=None, if_add_node_summary=None,
+               if_add_doc_description=None, if_add_node_text=None):
+    """Public API: index a document and return its tree structure."""
+    import os
+    user_opt = {
+        arg: value for arg, value in locals().items()
+        if arg != "doc" and value is not None
+    }
+    opt = ConfigLoader().load(user_opt)
+    return page_index_main(doc, opt)

static/index.html

@@ -36,6 +36,8 @@
                     </div>
                     <div><span class="cost-label">Response:</span> <span id="stat-response" class="cost-value">0</span>
                     </div>
+                    <div><span class="cost-label">TOC Analysis:</span> <span id="stat-toc" class="cost-value">0</span>
+                    </div>
                     <div class="cost-total-row">
                         <span class="cost-label" style="font-weight: 600;">Total Tokens:</span>
                         <span id="stat-total" class="cost-value cost-total">0</span>
@@ -62,7 +64,7 @@
         <main class="main-content">
             <header class="top-nav">
                 <h1>RAG Comparator</h1>
-                <p>Vector RAG vs Page Index RAG</p>
+                <p>Vector RAG vs Page Index RAG vs TOC Index RAG</p>
             </header>
 
             <!-- Loading Overlay -->
@@ -113,12 +115,16 @@
                         <div class="column-header">Vector RAG (Sliding Window)</div>
                         <div id="vector-chunks-list" class="chunks-list"></div>
                     </div>
+                    <div class="chunk-column">
+                        <div class="column-header">TOC Index RAG (Tree Nodes)</div>
+                        <div id="toc-chunks-list" class="chunks-list"></div>
+                    </div>
                 </div>
             </div>
         </div>
     </div>
 
-    <script src="/static/script.js?v=8"></script>
+    <script src="/static/script.js?v=9"></script>
 </body>
 
 </html>

static/script.js

@@ -209,6 +209,7 @@ document.addEventListener('DOMContentLoaded', () => {
             document.getElementById('stat-embed').textContent = data.tokens.embedding_tokens.toLocaleString();
             document.getElementById('stat-prompt').textContent = data.tokens.prompt_tokens.toLocaleString();
             document.getElementById('stat-response').textContent = data.tokens.completion_tokens.toLocaleString();
+            document.getElementById('stat-toc').textContent = (data.tokens.toc_analysis_tokens || 0).toLocaleString();
             document.getElementById('stat-total').textContent = data.total_tokens.toLocaleString();
         } catch (error) {
             console.error('Error fetching stats:', error);
@@ -269,9 +270,11 @@ document.addEventListener('DOMContentLoaded', () => {
         // Create Cards
         const vectorCard = createResponseCard('Vector RAG');
         const pageCard = createResponseCard('Page Index RAG');
+        const tocCard = createResponseCard('TOC Index RAG');
 
         responsesGrid.appendChild(vectorCard);
         responsesGrid.appendChild(pageCard);
+        responsesGrid.appendChild(tocCard);
         sysMsg.appendChild(responsesGrid);
         chatContainer.appendChild(sysMsg);
 
@@ -279,9 +282,10 @@ document.addEventListener('DOMContentLoaded', () => {
 
         const topK = parseInt(topKInput.value) || 5;
 
-        // Trigger both SSE calls
+        // Trigger all three SSE calls
         fetchSSE('/api/chat/vector', { query, top_k: topK }, vectorCard);
         fetchSSE('/api/chat/page', { query, top_k: topK }, pageCard);
+        fetchSSE('/api/chat/toc', { query, top_k: topK }, tocCard);
     }
 
     async function fetchSSE(url, payload, cardElement) {
@@ -334,7 +338,7 @@ document.addEventListener('DOMContentLoaded', () => {
                                 // Render sources
                                 let sourcesHtml = '';
                                 data.sources.forEach((s, idx) => {
-                                    let label = s.chunk_index !== undefined ? `Chunk #${s.chunk_index}` : `Page ${s.page_num}`;
+                                    let label = s.chunk_index !== undefined ? `Chunk #${s.chunk_index}` : s.node_id !== undefined ? `Node ${s.node_id}: ${s.title || ''}` : `Page ${s.page_num}`;
                                     sourcesHtml += `
                                         <div class="source-item">
                                             <span class="source-meta">[${idx + 1}] ${s.source} — ${label} (score: ${s.score})</span>
@@ -382,6 +386,7 @@ document.addEventListener('DOMContentLoaded', () => {
     const modalDocTitle = document.getElementById('modal-doc-title');
     const pageChunksList = document.getElementById('page-chunks-list');
     const vectorChunksList = document.getElementById('vector-chunks-list');
+    const tocChunksList = document.getElementById('toc-chunks-list');
 
     closeModalBtn.addEventListener('click', () => {
         chunksModal.style.display = 'none';
@@ -398,6 +403,7 @@ document.addEventListener('DOMContentLoaded', () => {
         modalDocTitle.textContent = filename;
         pageChunksList.innerHTML = '<div class="spinner large-spinner"></div>';
         vectorChunksList.innerHTML = '<div class="spinner large-spinner"></div>';
+        tocChunksList.innerHTML = '<div class="spinner large-spinner"></div>';
         chunksModal.style.display = 'flex';
 
         try {
@@ -470,10 +476,44 @@ document.addEventListener('DOMContentLoaded', () => {
                 vectorChunksList.innerHTML = '<div style="padding:16px;">No vector chunks found.</div>';
             }
 
+            // Render TOC Chunks
+            tocChunksList.innerHTML = '';
+            let totalTocTokens = 0;
+            if (data.toc_chunks && data.toc_chunks.length > 0) {
+                data.toc_chunks.forEach(chunk => {
+                    const el = document.createElement('div');
+                    el.className = 'chunk-item';
+                    el.innerHTML = `
+                        <span class="chunk-meta" style="display:flex; justify-content:space-between;">
+                            <span>Node ${chunk.node_id}: ${chunk.title || ''}</span>
+                            <span style="color:var(--text-secondary); font-size:0.85em;">Tokens: ${chunk.tokens || 0}</span>
+                        </span>
+                        <div class="chunk-text">${chunk.text}</div>
+                    `;
+                    tocChunksList.appendChild(el);
+                    totalTocTokens += (chunk.tokens || 0);
+                });
+
+                const totalEl = document.createElement('div');
+                totalEl.style.marginTop = '16px';
+                totalEl.style.padding = '12px';
+                totalEl.style.backgroundColor = 'var(--bg-secondary)';
+                totalEl.style.borderRadius = '6px';
+                totalEl.style.fontWeight = 'bold';
+                totalEl.style.display = 'flex';
+                totalEl.style.justifyContent = 'space-between';
+                totalEl.style.border = '1px solid var(--border-color)';
+                totalEl.innerHTML = `<span>Total TOC RAG Tokens:</span> <span>${totalTocTokens.toLocaleString()}</span>`;
+                tocChunksList.appendChild(totalEl);
+            } else {
+                tocChunksList.innerHTML = '<div style="padding:16px;">No TOC nodes found.</div>';
+            }
+
         } catch (error) {
             console.error(error);
             pageChunksList.innerHTML = '<div style="color:red; padding:16px;">Error loading data.</div>';
             vectorChunksList.innerHTML = '<div style="color:red; padding:16px;">Error loading data.</div>';
+            tocChunksList.innerHTML = '<div style="color:red; padding:16px;">Error loading data.</div>';
         }
     }
 });

static/style.css

@@ -188,7 +188,7 @@ body {
 
 .system-responses {
     display: grid;
-    grid-template-columns: 1fr 1fr;
+    grid-template-columns: 1fr 1fr 1fr;
     gap: 24px;
     width: 100%;
 }