ChromaDB is a lightweight vector database that stores documents as vectors and finds close matches. In a retrieval-augmented generation (RAG) workflow, it helps retrieve relevant information for an AI response. section 16.1.3
import ollama
import chromadb
# Step 1: Create a vector database
client = chromadb.Client()
collection = client.create_collection("my_docs")
# Step 2: Add documents to the database
documents = [ #1
"Ollama is a tool for running LLMs locally on your computer.",
"Streamlit is a Python framework for building web applications.",
"RAG combines retrieval with generation for better AI responses.",
"Virtual environments isolate Python packages per project.",
"The temperature parameter controls how creative the AI's responses are.",
"System prompts define the AI's personality and behavior rules.",
"ChromaDB is a lightweight vector database, a database optimized for storing vectors and finding the closest matches, for similarity search.",
"LangChain is a framework for building complex LLM applications.",
]
collection.add( #2
documents=documents,
ids=[f"doc{i}" for i in range(len(documents))]
)
print(f"Added {len(documents)} documents to the database.\n")
# Step 3: Query the database
query = "What is Ollama?"
results = collection.query(
query_texts=[query],
n_results=2 #3
)
print(f"Query: {query}")
print(f"Most relevant documents:")
for i, doc in enumerate(results["documents"][0]):
print(f" {i+1}. {doc}")
# Step 4: Use retrieved documents as context for the LLM
context = "\n".join(results["documents"][0])
prompt = f"""Based on the following context, answer the question.
If the context does not contain enough information, say so.
Context:
{context}
Question: {query}
Answer:"""
response = ollama.chat(
model="gemma3:4b",
messages=[{"role": "user", "content": prompt}]
)
print(f"\nAI Response:\n{response['message']['content']}")
The example creates a ChromaDB collection, adds sample documents, and uses the database in a local RAG workflow. section 16.1.3
Finding a relevant document
For the query “What is Ollama?”, the example retrieves a document describing Ollama as a tool for running language models locally, then uses the retrieved context to answer. section 16.1.3
