treerag 0.1.0

Lightweight structural RAG library — index documents as trees, query without a vector DB.

treerag 🌳

A lightweight structural RAG library — index documents as hierarchical trees, query them without a vector database.

Works with any LangChain-compatible LLM: OpenAI, Anthropic, Gemini, Ollama.

---

Install

# Base
pip install treerag

# With OpenAI
pip install "treerag[openai]"

# With Anthropic
pip install "treerag[anthropic]"

# Everything
pip install "treerag[all]"

---

Quick Start

from treerag import index_document, make_summarizer, ask, make_retriever
from langchain_openai import ChatOpenAI

llm = ChatOpenAI(model="gpt-4o")

# 1. Index
doc = index_document("my_doc.md", summarizer=make_summarizer(llm))

# 2. Ask — returns enriched LangChain AIMessage
result = ask("What does this cover?", doc, make_retriever(llm))
print(result.content)            # answer text
print(result.references)         # sections used
print(result.response_metadata)  # token usage, model, etc

---

Supported Input Formats

# Files
doc = index_document("my_doc.md")
doc = index_document("report.pdf")
doc = index_document("document.docx")
doc = index_document("notes.txt")

# URLs
doc = index_document("https://docs.example.com")

---

Response Formats

# Default — raw LangChain AIMessage with .references attached
result = ask("What is this?", doc, retriever)
print(result.content)      # answer text
print(result.references)   # [{"node_id", "title", "path", "file_name"}, ...]

# Plain dict
result = ask("What is this?", doc, retriever, return_raw=False)
print(result["answer"])
print(result["references"])

# Streaming
for chunk in ask("What is this?", doc, retriever, stream=True):
    if isinstance(chunk, dict):
        print("\nRefs:", chunk["__references__"])
    else:
        print(chunk, end="", flush=True)

---

Async Support

from treerag import aask, make_async_retriever

retriever = make_async_retriever(llm)
result = await aask("What is this?", doc, retriever)
print(result.content)

# Async streaming
async for chunk in await aask("What is this?", doc, retriever, stream=True):
    if isinstance(chunk, dict):
        print(chunk["__references__"])
    else:
        print(chunk, end="", flush=True)

---

Multi-Document Q&A

from treerag import ask_multi

doc1 = get_document_by_id("uuid-1")
doc2 = get_document_by_id("uuid-2")

result = ask_multi("What is the budget cap?", [doc1, doc2], retriever)
print(result.content)
print(result.references)  # shows which file each section came from

---

Registry Management

from treerag import list_documents, get_document_by_id, delete_document

# List all indexed docs
for doc in list_documents():
    print(doc["name"], doc["doc_id"], doc["total_sections"])

# Load by UUID (fast, O(1))
doc = get_document_by_id("3f7a1c2d-9b4e-4f8a-b2d1-6e5c3a9f0e12")

# Load by file name
from treerag import get_document
doc = get_document("my_doc.md")

# Delete by UUID
delete_document("3f7a1c2d-9b4e-4f8a-b2d1-6e5c3a9f0e12")

---

Custom Prompts

# Domain-specific summarizer
summarizer = make_summarizer(
    llm,
    system_prompt="You are a legal expert. Summarize clauses, obligations, and key terms precisely."
)

# Domain-specific retriever
retriever = make_retriever(
    llm,
    answer_system_prompt=(
        "You are a legal assistant. "
        "Provide precise answers using ONLY the provided document context. "
        "Avoid assumptions."
    )
)

# Per-question extra context
result = ask(
    "What are the key obligations and clauses defined in this document?",
    doc,
    retriever,
    extra_context="This document is a legal agreement outlining terms, obligations, and conditions."
)

---

Provider Examples

from langchain_openai import ChatOpenAI
from langchain_anthropic import ChatAnthropic
from langchain_google_genai import ChatGoogleGenerativeAI
from langchain_ollama import ChatOllama

# OpenAI — cost optimized: cheap model for search, smart model for answers
retriever = make_retriever(
    llm=ChatOpenAI(model="gpt-4o"),
    search_llm=ChatOpenAI(model="gpt-4o-mini"),
)

# Anthropic
retriever = make_retriever(ChatAnthropic(model="claude-haiku-4-5-20251001"))

# Gemini
retriever = make_retriever(ChatGoogleGenerativeAI(model="gemini-2.0-flash"))

# Ollama (local, free)
retriever = make_retriever(ChatOllama(model="llama3"))

---

Production Usage (no local storage)

# persist=False — returns dict, you handle storage
doc = index_document("my_doc.md", persist=False)
db.save(doc["doc_id"], doc)         # PostgreSQL
redis.set(doc["doc_id"], json.dumps(doc))  # Redis

# Load from your storage and pass directly to ask()
doc = db.get("some-uuid")
result = ask("Your question?", doc, retriever)

---

How It Works

Indexing:

File / URL
    ↓  read_file()           — read content
    ↓  parse_sections()      — split by markdown headers
    ↓  make_summarizer(llm)  — AI summaries per section
    ↓  build_hierarchy()     — parent → child tree
    ↓  flatten_tree()        — {node_id: content} lookup
    ↓  save_registry()       — indexed_docs.json (UUID key)

Q&A:

Query
    ↓  tree search   — LLM reads outline → picks node IDs
    ↓  fetch content — full text from flat_nodes
    ↓  answer gen    — LLM reads sections → generates answer
    ↓  AIMessage with .references attached

---

License

MIT