Build a minimal WeKnora-inspired document Q&A engine in Python
Welcome to another official UVF IT build. Today we are creating a local, private Retrieval-Augmented Generation (RAG) assistant inspired by Tencent's WeKnora. We will ingest text documents, chunk them, and answer questions using local embeddings and a simple vector search, all without sending data to external paid APIs.
This guide focuses on the core 'document to answer' pipeline. We will use Python, SQLite for metadata, and sentence-transformers for local vector embeddings. By the end, you will have a working CLI tool that can search your local knowledge base and generate context-aware responses.
First, we need a clean environment to ensure our local model and dependencies don't conflict with your system's Python packages. We will create a dedicated directory and a virtual environment to keep our project isolated.
Inside this environment, we install sentence-transformers which handles the heavy lifting of converting text into vector embeddings. We also install numpy for the mathematical operations required for cosine similarity search later on.
mkdir weknora-clone
cd weknora-clone
python -m venv venv
source venv/bin/activate # On Windows use: venv\Scripts\activate
pip install sentence-transformers numpyNow we create the core engine in rag.py. We start by defining our SQLite database structure. We need a table to store our document chunks, the original text, and the binary representation of the vector embeddings.
The ingest function takes a file path, reads the content, and splits it into manageable 500-character chunks. Each chunk is then passed through the all-MiniLM-L6-v2 model to generate a vector, which is stored in the database alongside the text.
rag.py
import sqlite3
import numpy as np
from sentence_transformers import SentenceTransformer
import os
# Initialize the embedding model (downloads on first run)
model = SentenceTransformer('all-MiniLM-L6-v2')
DB_NAME = 'knowledge.db'
def init_db():
"""Initialize the SQLite database schema."""
conn = sqlite3.connect(DB_NAME)
cursor = conn.cursor()
cursor.execute('''
CREATE TABLE IF NOT EXISTS documents (
id INTEGER PRIMARY KEY AUTOINCREMENT,
filename TEXT NOT NULL,
chunk_index INTEGER NOT NULL,
content TEXT NOT NULL,
embedding BLOB NOT NULL
)
''')
conn.commit()
conn.close()
def chunk_text(text, chunk_size=500):
"""Split text into fixed-size chunks."""
chunks = []
for i in range(0, len(text), chunk_size):
chunks.append(text[i:i + chunk_size])
return chunks
def ingest(file_path):
"""Read a file, chunk it, embed it, and store in DB."""
if not os.path.exists(file_path):
print(f"Error: File {file_path} not found.")
return
with open(file_path, 'r', encoding='utf-8') as f:
content = f.read()
chunks = chunk_text(content)
print(f"Ingesting {len(chunks)} chunks from {file_path}...")
# Generate embeddings for all chunks
embeddings = model.encode(chunks, convert_to_numpy=True)
conn = sqlite3.connect(DB_NAME)
cursor = conn.cursor()
for i, (chunk, emb) in enumerate(zip(chunks, embeddings)):
# Store vector as a BLOB
emb_blob = emb.tobytes()
cursor.execute(
"INSERT INTO documents (filename, chunk_index, content, embedding) VALUES (?, ?, ?, ?)",
(file_path, i, chunk, emb_blob)
)
conn.commit()
conn.close()
print(f"Successfully ingested {file_path}")
if __name__ == "__main__":
init_db()With data stored, we need a way to find the most relevant parts of our documents. We add a search function that takes a user's question, converts it into a vector using the same model, and compares it against all stored vectors.
Since SQLite doesn't have native vector search capabilities, we load the embeddings into memory as NumPy arrays. We calculate the cosine similarity between the query vector and every document vector, then return the top 3 most similar chunks.
rag.py
import sqlite3
import numpy as np
from sentence_transformers import SentenceTransformer
import os
# Initialize the embedding model (downloads on first run)
model = SentenceTransformer('all-MiniLM-L6-v2')
DB_NAME = 'knowledge.db'
def init_db():
"""Initialize the SQLite database schema."""
conn = sqlite3.connect(DB_NAME)
cursor = conn.cursor()
cursor.execute('''
CREATE TABLE IF NOT EXISTS documents (
id INTEGER PRIMARY KEY AUTOINCREMENT,
filename TEXT NOT NULL,
chunk_index INTEGER NOT NULL,
content TEXT NOT NULL,
embedding BLOB NOT NULL
)
''')
conn.commit()
conn.close()
def chunk_text(text, chunk_size=500):
"""Split text into fixed-size chunks."""
chunks = []
for i in range(0, len(text), chunk_size):
chunks.append(text[i:i + chunk_size])
return chunks
def ingest(file_path):
"""Read a file, chunk it, embed it, and store in DB."""
if not os.path.exists(file_path):
print(f"Error: File {file_path} not found.")
return
with open(file_path, 'r', encoding='utf-8') as f:
content = f.read()
chunks = chunk_text(content)
print(f"Ingesting {len(chunks)} chunks from {file_path}...")
# Generate embeddings for all chunks
embeddings = model.encode(chunks, convert_to_numpy=True)
conn = sqlite3.connect(DB_NAME)
cursor = conn.cursor()
for i, (chunk, emb) in enumerate(zip(chunks, embeddings)):
# Store vector as a BLOB
emb_blob = emb.tobytes()
cursor.execute(
"INSERT INTO documents (filename, chunk_index, content, embedding) VALUES (?, ?, ?, ?)",
(file_path, i, chunk, emb_blob)
)
conn.commit()
conn.close()
print(f"Successfully ingested {file_path}")
def search(query, top_k=3):
"""Search for the most similar chunks to the query."""
# Embed the query
query_embedding = model.encode(query, convert_to_numpy=True)
# Load all embeddings from DB
conn = sqlite3.connect(DB_NAME)
cursor = conn.cursor()
cursor.execute("SELECT id, content, embedding FROM documents")
rows = cursor.fetchall()
conn.close()
if not rows:
return []
# Prepare data for NumPy calculation
ids = []
contents = []
embeddings = []
for row in rows:
ids.append(row[0])
contents.append(row[1])
# Convert BLOB back to numpy array
emb = np.frombuffer(row[2], dtype=np.float32)
embeddings.append(emb)
embeddings_array = np.array(embeddings)
# Calculate cosine similarity
# Normalize vectors for cosine similarity
query_norm = query_embedding / (np.linalg.norm(query_embedding) + 1e-9)
embeddings_norm = embeddings_array / (np.linalg.norm(embeddings_array, axis=1, keepdims=True) + 1e-9)
similarities = np.dot(embeddings_norm, query_norm)
# Get top k indices
top_indices = np.argsort(similarities)[::-1][:top_k]
results = []
for idx in top_indices:
results.append({
'id': ids[idx],
'content': contents[idx],
'score': float(similarities[idx])
})
return results
if __name__ == "__main__":
init_db()WeKnora uses an LLM to generate answers. For this local clone, we will create a generate_answer function that acts as a 'mock LLM'.
It will take the retrieved context and the original question, and return a structured response. In a real scenario, you would swap this with a local LLM like llama-cpp-python or transformers.
rag.py
def generate_answer(question, context_chunks):
"""
Mock LLM that formats retrieved context into a structured answer.
In production, replace this with a call to a local LLM.
"""
print("\n--- Generating Answer ---")
print(f"Question: {question}")
print("\nContext Retrieved:")
if not context_chunks:
return "I could not find any relevant information in the knowledge base to answer your question."
answer_parts = []
for i, chunk in enumerate(context_chunks, 1):
# Truncate long chunks for readability
display_text = chunk[:200] + "..." if len(chunk) > 200 else chunk
answer_parts.append(f"{i}. {display_text}")
final_answer = "\n".join(answer_parts)
return f"Based on the retrieved context, here is the relevant information:\n{final_answer}"Create a simple command-line interface using argparse to allow users to ingest files or ask questions.
Run the script to test the full pipeline: ingest a sample text file, then ask a question related to its content.
rag.py
import argparse
import sys
# Load model once at startup to avoid reloading for every command
print("Loading embedding model...")
model = SentenceTransformer('all-MiniLM-L6-v2')
print("Model loaded.")
def main():
parser = argparse.ArgumentParser(description="Local RAG Knowledge Base")
subparsers = parser.add_subparsers(dest='command', help='Available commands')
# Ingest command
ingest_parser = subparsers.add_parser('ingest', help='Ingest a text file into the knowledge base')
ingest_parser.add_argument('file', type=str, help='Path to the text file to ingest')
# Ask command
ask_parser = subparsers.add_parser('ask', help='Ask a question to the knowledge base')
ask_parser.add_argument('question', type=str, help='The question to ask')
args = parser.parse_args()
if args.command == 'ingest':
ingest_document(args.file)
elif args.command == 'ask':
# 1. Search for relevant chunks
results = search(args.question, top_k=3)
# 2. Generate answer using mock LLM
answer = generate_answer(args.question, results)
print(answer)
else:
parser.print_help()
if __name__ == '__main__':
main()| Check | Command / Why it matters |
|---|---|
| Environment is active and dependencies installed | Run python -c "import sentence_transformers" to ensure no import errors. |
| Database is created and populated | Run python rag.py ingest sample.txt and verify knowledge.db file size increases. |
| Search returns relevant chunks | Run python rag.py ask "What is the main topic?" and ensure the output contains text from sample.txt. |
| Mock LLM generates structured response | Verify the output is formatted as a numbered list of context chunks. |
You now have a working local RAG pipeline! To extend this, you can replace the mock LLM with a local model like llama-3-8b via ollama, or add support for PDF ingestion using pypdf. Remember to keep your dataset small for the in-memory vector search to remain fast.