Skip to content
Categories:

Retrieval-Augmented Generation

Post date:
Author:

System Architecture Overview

Your RAG system will consist of three core components:

  1. Document Processing Pipeline – Ingests and chunks your security documents
  2. Vector Database – Stores embeddings using your nomic-embed model
  3. 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

  1. Periodic Re-indexing: Update embeddings monthly
  2. Cache Management: Clear query cache weekly
  3. 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*