ChatIndex: Tree indexing and retrieval for long conversational memory
ChatIndex - Tree Indexing for Long Conversations
ChatIndex is a context management system that enables LLMs to efficiently navigate and utilize long conversation histories through hierarchical tree-based indexing and intelligent reasoning-based retrieval.
Table of Contents
- Inspiration & Comparisons - Context Tree Specification - Installation - Complete Workflow - Phase 1: Building the Tree - Phase 2: Querying the TreeMotivation
Current AI chat assistants face a fundamental challenge: context management in long conversations. While current LLM apps use multiple separate conversations to bypass context limits, a truly human-like AI assistant should maintain a single, coherent conversation thread, making efficient context management critical.
Although modern LLMs have longer contexts, they still suffer from the long-context problem (e.g. context rot problem) - reasoning ability decreases as context grows longer. Memory-based systems (e.g. Dynamic Cheatsheet, mem0) have been invented to alleviate the context rot problem; however, memory-based representations are inherently lossy and inevitably lose information from the original conversation. In principle, no lossy representation is universally perfect for all downstream tasks. This leads to two key requirements for defining a flexible in-context management system:
- Preserve raw data: An index system that can retrieve the original conversation when necessary
- Multi-resolution access: Ability to retrieve information at different levels of detail on demand
ChatIndex Introduction
ChatIndex is designed to meet these requirements by constructing a hierarchical tree index โ which we call a Context Tree (CTree) โ that captures the structure and semantic organization of a long conversation. Unlike memory-based architectures that store only compressed, lossy summaries, ChatIndex preserves the complete raw conversation and layers a topic hierarchy on top: * Leaf nodes store raw conversational segments. * Internal nodes store topic summaries that abstract and represent their child nodes.This forms a multi-level topic hierarchy in which higher nodes represent broader themes and lower nodes convey increasingly specific details. See the figure below for an illustration.
When a query arrives, ChatIndex performs a top-down search through the topic tree. At each node, it evaluates whether the summary provides enough information for the query. If it does, the traversal stops and the system returns that higher-level summary; if not, it continues downward until more detailed informationโor the raw conversationโis accessed. This design offers:
- Dynamic retrieval resolution โ the system returns only as much detail as needed.
- Lossless fallback โ the raw conversation is always accessible when required.
- Efficient reasoning โ large contexts are reduced to the minimally sufficient subset.
Inspiration & Comparisons
ChatIndex is an extension of PageIndex, a tree-based index system for long documents, but adapted for conversational contexts with two key differences:
- Dynamic vs. Static:
- Structured vs. Unstructured:
Connection to B+-Tree
ChatIndex's Context Tree is structurally similar to a B+ tree in databases, but instead of indexing keys on disk, it indexes conversational context:
- Leaf nodes: (raw records)
- Internal nodes: (routing keys)
- Bounded fan-out (
max_children)
Key difference:
B+-trees use numeric/lexicographic key comparisons, while ChatIndex uses contextual relevance judged by an LLM to decide which branch to follow.
In short, ChatIndex borrows the efficient hierarchical structure of a B+-tree, but replaces key-based lookup with reasoning-based navigation over topics.
Context Tree Specification
A Context Tree consists of two types of nodes:
1. TopicNode
- Represents a conversation topic/subtopic
- Contains:
topic_name: Descriptive name (2-5 words)
- summary: Brief summary of the topic content
- startindex, endindex: Range of messages in the conversation
- children: List of child TopicNodes or MessageNodes
- subnodecount: Number of direct children
2. MessageNode
- Represents a leaf node containing an actual conversation exchange
- Contains:
system_message: Optional system message
- user_message: User's message
- assistant_message: Assistant's response
- message_index: Position in the conversation history
Tree Structure Properties
- Temporal ordering: New topic nodes can only be children of the current node or its ancestors
- Tree width control: The maximum width of a tree layer is controlled by
max_children, which guarantees that when conducting layer-wise tree search for relevant conversations, the context length will be controlled
Quick Start
Installation
- Clone the repository:
git clone https://github.com/yourusername/ChatIndex.git
cd ChatIndex
- Install dependencies:
pip install -r requirements.txt
- Set up API keys:
# For building trees (Phase 1)
export OPENAIAPIKEY="your-openai-key"
For querying trees (Phase 2)
export ANTHROPICAPIKEY="your-anthropic-key"
Or use a .env file:
echo "OPENAIAPIKEY=your-openai-key" > .env
echo "ANTHROPICAPIKEY=your-anthropic-key" >> .env
Complete Workflow
Here's the full pipeline from conversation to intelligent retrieval:
from ctree import CTree
from retrieval.llmtools import queryctree
import os
============================================
Phase 1: Build the conversation tree
============================================
tree = CTree(max_children=10)
Add your conversation messages
messages = [
{"role": "system", "content": "You are a helpful assistant."},
{"role": "user", "content": "What is Python?"},
{"role": "assistant", "content": "Python is a high-level programming language..."},
# ... more messages
]
tree.add(messages)
Save for later use
tree.save('my_conversation.json')
tree.print_tree()
============================================
Phase 2: Query the conversation tree
============================================
Load the tree (can be done in a separate session)
tree = CTree.load('my_conversation.json')
Ask questions about the conversation
result = query_ctree(
apikey=os.getenv("ANTHROPICAPI_KEY"),
ctree=tree,
user_query="What programming concepts were discussed?",
max_turns=50
)
print(result["final_response"]) print(f"Retrieved answer using {result['turns_used']} turns")
Phase 1: Building the Tree
Build a hierarchical index of your conversation:
from ctree import CTree
Initialize tree
tree = CTree(max_children=10)
Add conversation exchanges
messages = [
{"role": "system", "content": "You are a helpful programming tutor."},
{"role": "user", "content": "What is Python?"},
{"role": "assistant", "content": "Python is a high-level programming language..."}
]
tree.add(messages)
Save and visualize
tree.save('conversation_tree.json')
tree.print_tree()
Tree Generation Example
To build a Context Tree from an existing conversation history, run the included demo:
python demo.py
See demo.py for the complete implementation details.
Phase 2: Querying the Tree
Query your indexed conversation efficiently:
from ctree import CTree
from retrieval.llmtools import queryctree
import os
Load indexed conversation
tree = CTree.load('conversation_tree.json')
Ask questions
result = query_ctree(
apikey=os.getenv("ANTHROPICAPI_KEY"),
ctree=tree,
user_query="What topics were discussed about network protocols?",
max_turns=50 # default is 50
)
print(result["final_response"])
Key benefits:
- Cost reduction - Only retrieves relevant conversation segments
- Reasoning-based navigation - LLM explores the tree autonomously
- Multi-resolution - Gets summaries or full messages as needed
Example Tree Output
Here's what a Context Tree structure looks like (from the demo):
ROOT
โโโ Network Protocols and Testing
โ โโโ ICMP Message Types Comparison
โ โ โโโ [Message 0-2]
โ โโโ Ping Command and Network Testing
โ โ โโโ Ping Command and Basic Network Testing
โ โ โ โโโ [Message 2-4]
โ โ โ โโโ [Message 4-6]
โ โ โ โโโ [Message 6-10]
โ โ โโโ Advanced Network Diagnostics
โ โ โโโ [Message 10-12]
โ โโโ Network Protocol Analysis
โ โโโ [Message 12-16]
โโโ ... more topics
See ./save/conversation_tree.json for a complete tree visualization
Advanced Usage
Streaming Responses
Get real-time responses while querying:
from retrieval.llmtools import queryctree_streaming
def on_text(chunk): print(chunk, end='', flush=True)
def ontooluse(toolname, toolinput): print(f"\n[Using: {tool_name}]", flush=True)
result = queryctreestreaming( apikey=os.getenv("ANTHROPICAPI_KEY"), ctree=tree, user_query="What are the main topics?", ontextchunk=on_text, ontooluse=ontooluse )
Direct Tool Access
Use the tools directly without the LLM wrapper:
from retrieval.llm_tools import ChatIndexTools
tools = ChatIndexTools(tree)
Navigate the tree
root = tools.viewnodeand_children([]) # View root
topic = tools.viewnodeand_children([0]) # View first topic
Get messages
messages = tools.getnodemessages(0, 10) # Get messages 0-10
Roadmap
- [x] Hierarchical tree indexing - Build topic-based conversation trees
- [x] LLM-guided retrieval - Intelligent navigation with tools
- [x] Streaming support - Real-time responses
- [ ] Offline tree optimization - Post-processing for better structure
- [ ] Multi-LLM support - Support for different LLMs in retrieval
- [ ] Incremental updates - Efficiently update trees with new messages
- [ ] Vector search integration - Hybrid retrieval combining tree + embeddings
Contributing
This project is currently under active development. Any contributions are welcome! Please feel free to:
- Submit issues for bugs or feature requests
- Open pull requests with improvements
- Share your use cases and feedback