TeselaDB
DocumentationDownload

RAG over your documents

TeselaDB

TeselaDB

TeselaDB team

40 minIntermediate

Split your documents into chunks, store them with their vector and recall the closest ones to each question to give them to the model.

A language model answers better about your documents if you give it, along with the question, the chunks that talk about it. Here you store each chunk in TeselaDB with its vector and recall the closest ones with $vectorSearch to build the text you pass to the model.

What you will build

Two Python scripts. ingest.py splits each Markdown file in docs/ into paragraphs and stores them in kb.chunks with their vector and their source file. ask.py recalls the three paragraphs closest to a question and prints the prompt with those paragraphs and their source.

Requirements

  • Python 3.9 or later and a docs/ folder with your documents in Markdown.
  • A TeselaDB cluster, a user with readWrite on the kb database and its string in the TESELADB_URI variable.
  • PyMongo and sentence-transformers. The examples use the all-MiniLM-L6-v2 model, with 384 dimensions:
Shell
pip install pymongo sentence-transformers

1. Store the chunks

Save this as ingest.py. One paragraph per chunk is a simple starting point; with long documents, split by sections or by a fixed number of words:

Python
import os
from pathlib import Path

from pymongo import MongoClient
from pymongo.operations import SearchIndexModel
from sentence_transformers import SentenceTransformer

model = SentenceTransformer("all-MiniLM-L6-v2")
chunks = MongoClient(os.environ["TESELADB_URI"]).kb.chunks

# One chunk per paragraph, with the file it comes from and its vector.
for path in Path("docs").glob("*.md"):
    paragraphs = [p.strip() for p in path.read_text().split("\n\n") if p.strip()]
    chunks.insert_many(
        {"source": path.name, "text": text, "embedding": model.encode(text).tolist()} for text in paragraphs
    )

chunks.create_search_index(SearchIndexModel(
    name="chunks",
    type="vectorSearch",
    definition={"fields": [{"type": "vector", "path": "embedding", "numDimensions": 384, "similarity": "cosine"}]},
))
print(chunks.count_documents({}), "chunks")

2. Recall the context of a question

Save this as ask.py. The prompt asks the model to answer only from the context and to cite the file of each fact:

Python
import os
import sys

from pymongo import MongoClient
from sentence_transformers import SentenceTransformer

model = SentenceTransformer("all-MiniLM-L6-v2")
chunks = MongoClient(os.environ["TESELADB_URI"]).kb.chunks


def context_for(question: str, limit: int = 3) -> list[dict]:
    return list(chunks.aggregate([
        {"$vectorSearch": {"index": "chunks", "path": "embedding", "queryVector": model.encode(question).tolist(),
                           "numCandidates": 100, "limit": limit}},
        {"$project": {"_id": 0, "source": 1, "text": 1}},
    ]))


question = sys.argv[1]
context = context_for(question)
prompt = "Answer only with this context and cite the source of each fact.\n\n"
prompt += "\n\n".join(f"[{c['source']}] {c['text']}" for c in context)
prompt += f"\n\nQuestion: {question}"
print(prompt)

Try it

With two files in docs/, one about adding nodes and one about draining them, store the chunks and ask, waiting a few seconds between the two commands so the vector index that ingest.py creates finishes building before ask.py queries it:

Shell
python ingest.py
python ask.py "how do I take a machine out of the cluster?"

The prompt starts with the paragraph of drain.md, even though the question, which asks how to take a machine out of the cluster, does not say "drain". Pass it to the language model you use as the user message.

Clean up

JavaScript
db.getSiblingDB("kb").chunks.drop();

Run it from mongosh when you are done. Dropping kb.chunks also drops its vector index and leaves nothing of the test in the kb database.

Next steps