RAG for DevOps Runbooks: Build a Searchable Knowledge Base with ChromaDB and Anthropic SDK

pip install anthropic chromadb sentence-transformers python-frontmatter

import chromadb
import anthropic
import frontmatter
import os
import hashlib
from pathlib import Path
from typing import Optional
 
 
def chunk_markdown(content: str, chunk_size: int = 800, overlap: int = 100) -> list[str]:
    """
    Split markdown into overlapping chunks that preserve section context.
    Prefers splitting at headers/paragraph boundaries.
    """
    lines = content.split('\n')
    chunks = []
    current_chunk_lines = []
    current_section_header = ""
    current_length = 0
 
    for line in lines:
        # Track current section for context injection
        if line.startswith('#'):
            current_section_header = line.strip()
 
        current_chunk_lines.append(line)
        current_length += len(line) + 1
 
        if current_length >= chunk_size:
            chunk_text = '\n'.join(current_chunk_lines).strip()
            if chunk_text:
                chunks.append(chunk_text)
 
            # Overlap: keep last N characters worth of lines
            overlap_lines = []
            overlap_len = 0
            for l in reversed(current_chunk_lines):
                if overlap_len + len(l) > overlap:
                    break
                overlap_lines.insert(0, l)
                overlap_len += len(l) + 1
 
            # Start new chunk with section header for context
            current_chunk_lines = []
            if current_section_header and current_section_header not in ''.join(overlap_lines):
                current_chunk_lines.append(current_section_header)
            current_chunk_lines.extend(overlap_lines)
            current_length = sum(len(l) + 1 for l in current_chunk_lines)
 
    # Don't lose the last chunk
    if current_chunk_lines:
        chunk_text = '\n'.join(current_chunk_lines).strip()
        if chunk_text:
            chunks.append(chunk_text)
 
    return chunks
 
 
def load_runbooks(runbooks_dir: str) -> list[dict]:
    """
    Load all markdown files from a directory as documents.
    Returns list of {content, title, source, category, doc_id}.
    """
    documents = []
    runbooks_path = Path(runbooks_dir)
 
    for md_file in runbooks_path.glob("**/*.md"):
        try:
            post = frontmatter.load(str(md_file))
            title = post.get("title", md_file.stem.replace("-", " ").title())
            category = post.get("category", md_file.parent.name)
            content = post.content
 
            chunks = chunk_markdown(content)
            for i, chunk in enumerate(chunks):
                doc_id = hashlib.md5(f"{md_file}:{i}".encode()).hexdigest()
                documents.append({
                    "id": doc_id,
                    "content": chunk,
                    "title": title,
                    "source": str(md_file.relative_to(runbooks_path)),
                    "category": category,
                    "chunk_index": i,
                    "total_chunks": len(chunks)
                })
 
        except Exception as e:
            print(f"Error loading {md_file}: {e}")
 
    return documents

from sentence_transformers import SentenceTransformer
 
embedding_model = SentenceTransformer("all-MiniLM-L6-v2")  # Small, fast, good quality
 
 
def build_runbook_index(
    runbooks_dir: str,
    chroma_path: str = "./chroma_runbooks"
) -> chromadb.Collection:
    """
    Build or update the ChromaDB index from runbook documents.
    """
    chroma_client = chromadb.PersistentClient(path=chroma_path)
 
    # Get or create collection
    collection = chroma_client.get_or_create_collection(
        name="runbooks",
        metadata={"hnsw:space": "cosine"}
    )
 
    print("Loading runbooks...")
    documents = load_runbooks(runbooks_dir)
    print(f"Loaded {len(documents)} chunks from {runbooks_dir}")
 
    # Check which docs are already indexed
    existing_ids = set(collection.get()["ids"])
    new_docs = [d for d in documents if d["id"] not in existing_ids]
 
    if not new_docs:
        print("Index is up to date.")
        return collection
 
    print(f"Indexing {len(new_docs)} new chunks...")
 
    # Batch embed and insert
    batch_size = 64
    for i in range(0, len(new_docs), batch_size):
        batch = new_docs[i:i + batch_size]
        texts = [d["content"] for d in batch]
        embeddings = embedding_model.encode(texts).tolist()
 
        collection.add(
            ids=[d["id"] for d in batch],
            embeddings=embeddings,
            documents=[d["content"] for d in batch],
            metadatas=[{
                "title": d["title"],
                "source": d["source"],
                "category": d["category"],
                "chunk_index": d["chunk_index"]
            } for d in batch]
        )
 
    print(f"Indexed {len(new_docs)} new chunks. Total: {collection.count()}")
    return collection

claude_client = anthropic.Anthropic(api_key=os.getenv("ANTHROPIC_API_KEY"))
 
 
def retrieve_relevant_chunks(
    query: str,
    collection: chromadb.Collection,
    n_results: int = 5,
    category_filter: Optional[str] = None
) -> list[dict]:
    """Embed query and retrieve top-k most relevant document chunks."""
    query_embedding = embedding_model.encode([query]).tolist()[0]
 
    where_filter = {}
    if category_filter:
        where_filter = {"category": {"$eq": category_filter}}
 
    results = collection.query(
        query_embeddings=[query_embedding],
        n_results=n_results,
        where=where_filter if where_filter else None,
        include=["documents", "metadatas", "distances"]
    )
 
    chunks = []
    for doc, meta, distance in zip(
        results["documents"][0],
        results["metadatas"][0],
        results["distances"][0]
    ):
        chunks.append({
            "content": doc,
            "title": meta["title"],
            "source": meta["source"],
            "category": meta["category"],
            "relevance_score": 1 - distance  # Convert distance to similarity
        })
 
    return chunks
 
 
def query_runbooks(
    query: str,
    collection: chromadb.Collection,
    n_results: int = 5,
    category_filter: Optional[str] = None
) -> dict:
    """
    Full RAG pipeline: retrieve context, generate grounded answer.
    """
    # Retrieve relevant chunks
    chunks = retrieve_relevant_chunks(query, collection, n_results, category_filter)
 
    if not chunks:
        return {
            "answer": "No relevant runbooks found for this query. Please check your documentation or ask a team member.",
            "sources": [],
            "confidence": "low"
        }
 
    # Build context from retrieved chunks
    context_parts = []
    for i, chunk in enumerate(chunks):
        context_parts.append(
            f"[Source {i+1}: {chunk['title']} ({chunk['source']})]\n{chunk['content']}"
        )
 
    context = "\n\n---\n\n".join(context_parts)
 
    # Generate answer with Claude
    prompt = f"""You are a DevOps knowledge assistant. Answer the question using ONLY the provided runbook context.
 
## Question
{query}
 
## Relevant Runbook Content
{context}
 
## Instructions
- Answer directly and specifically based on the runbook content provided
- If the runbooks contain step-by-step commands, include them exactly
- If the question cannot be answered from the provided context, say so clearly
- Cite which source (Source 1, Source 2, etc.) you used for each part of your answer
- Do not make up information not in the provided context
- Format the answer clearly with sections and code blocks where appropriate"""
 
    response = claude_client.messages.create(
        model="claude-sonnet-5",
        max_tokens=2000,
        messages=[{"role": "user", "content": prompt}]
    )
 
    answer = response.content[0].text
 
    # Determine confidence based on top relevance score
    top_score = chunks[0]["relevance_score"] if chunks else 0
    confidence = "high" if top_score > 0.8 else "medium" if top_score > 0.6 else "low"
 
    return {
        "answer": answer,
        "sources": [{"title": c["title"], "source": c["source"], "score": round(c["relevance_score"], 3)} for c in chunks],
        "confidence": confidence,
        "top_relevance_score": top_score
    }

from fastapi import FastAPI
from pydantic import BaseModel
 
app = FastAPI(title="DevOps Runbook RAG")
collection = None
 
 
@app.on_event("startup")
async def startup():
    global collection
    runbooks_dir = os.getenv("RUNBOOKS_DIR", "./runbooks")
    collection = build_runbook_index(runbooks_dir)
 
 
class QueryRequest(BaseModel):
    question: str
    category: Optional[str] = None
    n_results: int = 5
 
 
@app.post("/query")
async def query_runbooks_api(request: QueryRequest):
    if collection is None:
        return {"error": "Index not ready"}
    return query_runbooks(request.question, collection, request.n_results, request.category)
 
 
@app.post("/index/rebuild")
async def rebuild_index():
    global collection
    runbooks_dir = os.getenv("RUNBOOKS_DIR", "./runbooks")
    collection = build_runbook_index(runbooks_dir)
    return {"status": "rebuilt", "total_chunks": collection.count()}

# Query
curl -X POST http://localhost:8000/query \
  -H "Content-Type: application/json" \
  -d '{"question": "How do we restore the payments database from the latest backup?"}'
 
# Response:
# {
#   "answer": "Based on [Source 1: Database Recovery Runbook], restore the payments database using:\n\n```bash\n# 1. Stop the application\nkubectl scale deployment payment-api --replicas=0 -n production\n\n# 2. Restore from latest snapshot\naws rds restore-db-instance-to-point-in-time ...\n```\n",
#   "sources": [{"title": "Database Recovery Runbook", "source": "runbooks/database-recovery.md", "score": 0.91}],
#   "confidence": "high"
# }