Categories: docs
Retrieval-Augmented Generation
Post date:
Author: g.t.n
System Architecture Overview
Your RAG system will consist of three core components:
- Document Processing Pipeline – Ingests and chunks your security documents
- Vector Database – Stores embeddings using your nomic-embed model
- Retrieval & Generation – Queries relevant context and generates responses with Gemma-3
Required Software Stack
Core Dependencies
pip install langchain-community
pip install chromadb
pip install sentence-transformers
pip install requests # For Ollama API calls
pip install pypdf2 pymupdf # PDF processing
pip install python-docx # Word documents
Ollama Setup & Configuration
# Download and install Ollama from https://ollama.com/download
# Then pull your model (you'll need to convert your GGUF to Ollama format)
ollama create gemma-3-4b-security -f ./Modelfile
ollama run gemma-3-4b-security
Modelfile for your Gemma-3 GGUF
FROM ./unsloth_gemma-3-4b-it-GGUF_gemma-3-4b-it-Q4_K_M.gguf
TEMPLATE """{{ if .System }}<|system|>
{{ .System }}<|end|>
{{ end }}{{ if .Prompt }}<|user|>
{{ .Prompt }}<|end|>
{{ end }}<|assistant|>
{{ .Response }}<|end|>
"""
PARAMETER temperature 0.3
PARAMETER num_ctx 4096
PARAMETER num_gpu 99 # Use all available GPU layers
Implementation Steps
Step 1: Document Processor
import os
from pathlib import Path
from langchain.text_splitter import RecursiveCharacterTextSplitter
from langchain_community.document_loaders import (
PyPDFLoader, TextLoader, DirectoryLoader
)
class SecurityDocumentProcessor:
def __init__(self, chunk_size=1000, chunk_overlap=200):
self.text_splitter = RecursiveCharacterTextSplitter(
chunk_size=chunk_size,
chunk_overlap=chunk_overlap,
separators=["\n\n", "\n", ".", "!", "?", ",", " ", ""]
)
def load_documents(self, directory_path):
"""Load various document types from directory"""
loaders = {
"*.pdf": PyPDFLoader,
"*.txt": TextLoader,
"*.log": TextLoader,
"*.md": TextLoader
}
documents = []
for pattern, loader_class in loaders.items():
loader = DirectoryLoader(
directory_path,
glob=pattern,
loader_cls=loader_class
)
documents.extend(loader.load())
return self.text_splitter.split_documents(documents)
Step 2: Embedding Model Setup (Ollama-Based)
import requests
import json
import numpy as np
from typing import List, Union
class NomicEmbedder:
def __init__(self, model_name="nomic-embed-text:latest", base_url="http://localhost:11434"):
self.model_name = model_name
self.base_url = base_url
self.api_url = f"{base_url}/api/embeddings"
# Verify model is available
if not self._check_model_available():
raise ValueError(f"Model {model_name} not found in Ollama. Run: ollama pull {model_name}")
def _check_model_available(self):
"""Check if the embedding model is available in Ollama"""
try:
response = requests.get(f"{self.base_url}/api/tags", timeout=5)
if response.status_code == 200:
models = response.json().get('models', [])
return any(model['name'] == self.model_name for model in models)
except:
pass
return False
def _get_embeddings(self, texts: Union[str, List[str]]) -> np.ndarray:
"""Get embeddings from Ollama API"""
if isinstance(texts, str):
texts = [texts]
payload = {
"model": self.model_name,
"prompt": texts[0] if len(texts) == 1 else texts
}
try:
response = requests.post(self.api_url, json=payload, timeout=30)
response.raise_for_status()
result = response.json()
# Handle single vs batch responses
if 'embedding' in result:
return np.array([result['embedding']])
elif 'embeddings' in result:
return np.array(result['embeddings'])
else:
raise ValueError("Unexpected response format from Ollama")
except requests.exceptions.RequestException as e:
raise ConnectionError(f"Failed to get embeddings from Ollama: {str(e)}")
def embed_documents(self, texts: List[str], batch_size: int = 10) -> np.ndarray:
"""Generate embeddings for document chunks with batching"""
all_embeddings = []
# Process in batches to avoid overwhelming Ollama
for i in range(0, len(texts), batch_size):
batch = texts[i:i + batch_size]
print(f"Processing embedding batch {i//batch_size + 1}/{(len(texts)-1)//batch_size + 1}")
for text in batch: # Ollama embeddings API processes one at a time
embedding = self._get_embeddings(text)
all_embeddings.extend(embedding)
return np.array(all_embeddings)
def embed_query(self, query: str) -> np.ndarray:
"""Generate embedding for search query"""
return self._get_embeddings(query)[0]
# Alternative: Hybrid approach using sentence-transformers as fallback
class HybridNomicEmbedder:
def __init__(self, model_identifier="nomic-embed-text:latest", base_url="http://localhost:11434"):
self.use_ollama = True
# Try Ollama first
try:
self.ollama_embedder = NomicEmbedder(model_identifier, base_url)
except (ValueError, ConnectionError):
print("Ollama not available, falling back to sentence-transformers")
self.use_ollama = False
# Fallback to sentence-transformers with local GGUF
from sentence_transformers import SentenceTransformer
if model_identifier.endswith('.gguf'):
self.model = SentenceTransformer(model_identifier, device='cuda')
else:
# Try to use sentence-transformers with model name
self.model = SentenceTransformer('nomic-ai/nomic-embed-text-v1.5', device='cuda')
def embed_documents(self, texts: List[str]) -> np.ndarray:
if self.use_ollama:
return self.ollama_embedder.embed_documents(texts)
else:
return self.model.encode(texts, show_progress_bar=True)
def embed_query(self, query: str) -> np.ndarray:
if self.use_ollama:
return self.ollama_embedder.embed_query(query)
else:
return self.model.encode([query])[0]
Step 3: Vector Database with ChromaDB
import chromadb
from chromadb.config import Settings
class VectorStore:
def __init__(self, collection_name="security_docs", persist_path="./chroma_db"):
self.client = chromadb.PersistentClient(
path=persist_path,
settings=Settings(anonymized_telemetry=False)
)
self.collection = self.client.get_or_create_collection(
name=collection_name,
metadata={"description": "Security research documents"}
)
def add_documents(self, documents, embeddings, metadatas):
"""Add documents to vector store"""
ids = [f"doc_{i}" for i in range(len(documents))]
self.collection.add(
documents=documents,
embeddings=embeddings.tolist(),
metadatas=metadatas,
ids=ids
)
def similarity_search(self, query_embedding, k=5):
"""Retrieve most relevant documents"""
results = self.collection.query(
query_embeddings=[query_embedding.tolist()],
n_results=k
)
return results
Step 4: LLM Integration with Ollama
import requests
import json
class OllamaLLM:
def __init__(self, model_name="gemma-3-4b-security", base_url="http://localhost:11434"):
self.model_name = model_name
self.base_url = base_url
self.api_url = f"{base_url}/api/generate"
def generate_response(self, context, query, max_tokens=512):
"""Generate response using retrieved context"""
prompt = f"""You are a cybersecurity research assistant. Based on the following security research context, provide a detailed and accurate answer to the question.
Context:
{context}
Question: {query}
Answer:"""
payload = {
"model": self.model_name,
"prompt": prompt,
"stream": False,
"options": {
"temperature": 0.3,
"num_predict": max_tokens,
"stop": ["Question:", "Context:"]
}
}
try:
response = requests.post(self.api_url, json=payload, timeout=60)
response.raise_for_status()
result = response.json()
return result.get('response', '').strip()
except requests.exceptions.RequestException as e:
return f"Error communicating with Ollama: {str(e)}"
def health_check(self):
"""Check if Ollama service is running"""
try:
response = requests.get(f"{self.base_url}/api/tags", timeout=5)
return response.status_code == 200
except:
return False
Step 5: Complete RAG System with Ollama Embeddings
class SecurityRAGSystem:
def __init__(self, ollama_model_name, nomic_model_identifier="nomic-embed-text:latest"):
# Use Ollama-based embeddings
self.embedder = NomicEmbedder(nomic_model_identifier)
self.vector_store = VectorStore()
self.llm = OllamaLLM(ollama_model_name)
# Verify both services are running
if not self.llm.health_check():
raise ConnectionError("Ollama service not running. Start with: ollama serve")
def ingest_documents(self, document_directory):
"""Process and store documents"""
processor = SecurityDocumentProcessor()
documents = processor.load_documents(document_directory)
# Extract text and metadata
texts = [doc.page_content for doc in documents]
metadatas = [doc.metadata for doc in documents]
# Generate embeddings using Ollama
print("Generating embeddings via Ollama...")
embeddings = self.embedder.embed_documents(texts)
# Store in vector database
self.vector_store.add_documents(texts, embeddings, metadatas)
print(f"Ingested {len(documents)} document chunks")
def query(self, question, k=3):
"""Query the RAG system"""
# Embed the question using Ollama
query_embedding = self.embedder.embed_query(question)
# Retrieve relevant documents
results = self.vector_store.similarity_search(query_embedding, k=k)
# Combine retrieved context
context = "\n\n".join(results['documents'][0])
# Generate response
print("Generating response...")
response = self.llm.generate_response(context, question)
return {
'answer': response,
'sources': results['metadatas'][0],
'relevance_scores': results.get('distances', [[]])[0]
}
Step 6: Usage Example with Full Ollama Integration
# Initialize the RAG system - now fully Ollama-based
rag = SecurityRAGSystem(
ollama_model_name="gemma-3-4b-security",
nomic_model_identifier="nomic-embed-text:latest" # Your existing Ollama model
)
# Ingest your security documents
rag.ingest_documents("./security_documents/")
# Query the system
result = rag.query("What are common SQL injection attack patterns in web applications?")
print("Answer:", result['answer'])
print("Sources:", result['sources'])
print("Relevance Scores:", result['relevance_scores'])
# For hybrid approach (fallback to local GGUF if Ollama unavailable)
hybrid_rag = SecurityRAGSystem(
ollama_model_name="gemma-3-4b-security",
nomic_model_identifier="path/to/nomic-embed-text-v1.5.f16.gguf" # Fallback path
)
Verifying Your Ollama Setup
def verify_ollama_setup():
"""Verify both models are available in Ollama"""
try:
# Check embedding model
embedder = NomicEmbedder("nomic-embed-text:latest")
test_embedding = embedder.embed_query("test query")
print(f"✓ Embedding model working - dimension: {len(test_embedding)}")
# Check generation model
llm = OllamaLLM("gemma-3-4b-security")
if llm.health_check():
test_response = llm.generate_response("Test context", "What is this?", max_tokens=50)
print(f"✓ Generation model working - sample: {test_response[:50]}...")
return True
except Exception as e:
print(f"✗ Setup issue: {e}")
return False
# Run verification
verify_ollama_setup()
Ollama Service Management
import subprocess
import time
class OllamaManager:
@staticmethod
def start_service():
"""Start Ollama service if not running"""
try:
subprocess.Popen(["ollama", "serve"], shell=True)
time.sleep(3) # Give service time to start
return True
except:
return False
@staticmethod
def list_models():
"""List available Ollama models"""
try:
result = subprocess.run(["ollama", "list"],
capture_output=True, text=True, shell=True)
return result.stdout
except:
return "Error listing models"
@staticmethod
def pull_model(model_name):
"""Pull a model from Ollama registry"""
try:
subprocess.run(["ollama", "pull", model_name], shell=True)
return True
except:
return False
Quick Start Guide
1. Install Ollama
# Download from https://ollama.com/download
# Or via PowerShell:
winget install Ollama.Ollama
2. Create Your Custom Model
# Create Modelfile in your project directory
# Copy the Modelfile content above, then:
ollama create gemma-3-4b-security -f Modelfile
# Verify creation
ollama list
3. Install Python Dependencies
pip install langchain-community chromadb sentence-transformers requests pypdf2 pymupdf python-docx
4. Start the System
# Start Ollama service
ollama serve
# Run your RAG system
python security_rag.py
Troubleshooting Common Issues
Ollama won’t start:
- Check Windows Defender/Firewall settings
- Ensure port 11434 is available
- Run as administrator if needed
Model loading errors:
- Verify GGUF file path in Modelfile
- Check available disk space (models need 2-3x their size)
- Monitor GPU memory usage
Slow performance:
- Reduce context window:
num_ctx 2048 - Adjust batch size in embedding generation
- Use SSD storage for vector database
Query Optimization
# Implement semantic caching for repeated queries
from functools import lru_cache
@lru_cache(maxsize=100)
def cached_query(self, question_hash):
return self.query(question)
Document Preprocessing
# Security-specific text preprocessing
def preprocess_security_text(text):
# Preserve technical terms, IPs, hashes
import re
# Normalize but preserve security indicators
text = re.sub(r'\s+', ' ', text) # Normalize whitespace
text = re.sub(r'[^\w\s\.\-\:\/]', ' ', text) # Keep relevant chars
return text.strip()
Monitoring & Maintenance
Performance Metrics
import time
import psutil
def monitor_system_resources():
cpu_usage = psutil.cpu_percent()
memory = psutil.virtual_memory()
gpu_memory = # Use nvidia-ml-py for GPU monitoring
return {
'cpu_usage': cpu_usage,
'ram_usage': memory.percent,
'available_ram': memory.available / 1024**3, # GB
}
Regular Maintenance Tasks
- Periodic Re-indexing: Update embeddings monthly
- Cache Management: Clear query cache weekly
- Performance Tuning: Monitor response times and adjust parameters
Security Considerations
- Data Privacy: All processing remains local
- Access Control: Implement file-based permissions
- Audit Trail: Log all queries and document access
- Backup Strategy: Regular vector database backups
Expected Performance
- Initial Setup: 2-4 hours for 10GB of documents
- Query Speed: 3-8 seconds per query
- Memory Usage: ~8-12GB RAM during operation
- Storage: ~20-30% additional space for embeddings
research notes (click to expand):
model, system, operator: qwen3_4b, acer-nitro-5-4gb-vram, gtn |
notes: failed* |
