Building Production RAG Pipelines: Chunking, Embeddings, and Retrieval at Scale

// domain.go — core types flowing through the pipeline
package rag
 
import (
	"time"
	"github.com/google/uuid"
)
 
type Document struct {
	ID        uuid.UUID
	Title     string
	Content   string
	Source    string
	Metadata  map[string]string
	UpdatedAt time.Time
}
 
type Chunk struct {
	ID               uuid.UUID
	DocumentID       uuid.UUID
	Content          string
	Index            int
	Metadata         map[string]string
	Embedding        []float32
	EmbeddingVersion string  // Critical: track model version
	TokenCount       int
}
 
type SearchRequest struct {
	Query          string
	TopK           int
	MinScore       float64
	MetadataFilter map[string]string
	UseMRR         bool
	MRRLambda      float64  // 0 = max diversity, 1 = max relevance
}
 
type SearchResult struct {
	Chunk     Chunk
	Score     float64
	RankScore float64
}

// chunker.go — semantic chunking
package rag
 
import (
	"fmt"
	"strings"
	"github.com/google/uuid"
	"github.com/pkoukk/tiktoken-go"
)
 
type SemanticChunker struct {
	targetSize int // tokens (e.g., 512)
	overlap    int // tokens (e.g., 50 for 10% overlap)
	enc        *tiktoken.Tiktoken
}
 
func NewSemanticChunker(targetSize, overlap int) (*SemanticChunker, error) {
	enc, err := tiktoken.EncodingForModel("text-embedding-3-small")
	if err != nil {
		return nil, fmt.Errorf("init tokenizer: %w", err)
	}
	return &SemanticChunker{
		targetSize: targetSize,
		overlap:    overlap,
		enc:        enc,
	}, nil
}
 
// Chunk splits on paragraph boundaries, then merges into target-sized chunks.
func (c *SemanticChunker) Chunk(doc Document) []Chunk {
	// Split by double newline (paragraph boundary)
	paras := strings.Split(doc.Content, "\n\n")
 
	var chunks []Chunk
	var buffer strings.Builder
	bufferTokens := 0
 
	for _, para := range paras {
		para = strings.TrimSpace(para)
		if para == "" {
			continue
		}
 
		paraTokens := len(c.enc.Encode(para, nil, nil))
 
		// If adding this paragraph exceeds target size, flush buffer.
		if bufferTokens+paraTokens > c.targetSize && buffer.Len() > 0 {
			chunks = append(chunks, Chunk{
				ID:         uuid.New(),
				DocumentID: doc.ID,
				Content:    buffer.String(),
				Index:      len(chunks),
				Metadata:   copyMetadata(doc.Metadata),
				TokenCount: bufferTokens,
			})
			buffer.Reset()
			bufferTokens = 0
		}
 
		// Add paragraph to buffer
		if buffer.Len() > 0 {
			buffer.WriteString("\n\n")
		}
		buffer.WriteString(para)
		bufferTokens += paraTokens + 2 // +2 for separator
	}
 
	// Flush final buffer
	if buffer.Len() > 0 {
		chunks = append(chunks, Chunk{
			ID:         uuid.New(),
			DocumentID: doc.ID,
			Content:    buffer.String(),
			Index:      len(chunks),
			Metadata:   copyMetadata(doc.Metadata),
			TokenCount: bufferTokens,
		})
	}
 
	return chunks
}
 
func copyMetadata(src map[string]string) map[string]string {
	dst := make(map[string]string, len(src))
	for k, v := range src {
		dst[k] = v
	}
	return dst
}

// embedder.go — embedding with versioning and caching
package rag
 
import (
	"context"
	"fmt"
	"os"
	"sync"
	"time"
 
	"github.com/sashabaranov/go-openai"
)
 
const EmbeddingVersion = "text-embedding-3-small::20240415"
 
type EmbeddingService struct {
	client   *openai.Client
	version  string
	cache    map[string][]float32
	cacheMu  sync.RWMutex
	hitRate  float64
}
 
func NewEmbeddingService() *EmbeddingService {
	return &EmbeddingService{
		client:  openai.NewClient(os.Getenv("OPENAI_API_KEY")),
		version: EmbeddingVersion,
		cache:   make(map[string][]float32),
	}
}
 
// Embed calls OpenAI with in-memory caching and version tracking.
// Returns embeddings aligned to the input order of texts, so callers can
// assign results by index. Cache hits and freshly fetched vectors are both
// written back to their original position — never appended in arrival order.
func (es *EmbeddingService) Embed(ctx context.Context, texts []string) ([][]float32, error) {
	embeddings := make([][]float32, len(texts))
	var toEmbed []string
	var toEmbedIdx []int // original index in texts for each uncached entry
	hitCount := 0
 
	// Check cache, recording the input position of every miss.
	es.cacheMu.RLock()
	for i, text := range texts {
		if emb, ok := es.cache[text]; ok {
			embeddings[i] = emb
			hitCount++
		} else {
			toEmbed = append(toEmbed, text)
			toEmbedIdx = append(toEmbedIdx, i)
		}
	}
	es.cacheMu.RUnlock()
 
	// Call API for uncached texts (batch up to 2048 per OpenAI limits)
	if len(toEmbed) > 0 {
		resp, err := es.client.CreateEmbeddings(ctx, openai.EmbeddingRequest{
			Input: toEmbed,
			Model: openai.SmallEmbedding3,
		})
		if err != nil {
			return nil, fmt.Errorf("embedding API: %w", err)
		}
		if len(resp.Data) != len(toEmbed) {
			return nil, fmt.Errorf("embedding API returned %d vectors for %d inputs", len(resp.Data), len(toEmbed))
		}
 
		// Cache each vector and place it at its ORIGINAL input index.
		// OpenAI does not guarantee resp.Data is returned in request order — each
		// item carries an Index into the request Input — so key by data.Index.
		es.cacheMu.Lock()
		for _, data := range resp.Data {
			embedding := make([]float32, len(data.Embedding))
			for j, v := range data.Embedding {
				embedding[j] = float32(v)
			}
			es.cache[toEmbed[data.Index]] = embedding
			embeddings[toEmbedIdx[data.Index]] = embedding
		}
		es.cacheMu.Unlock()
	}
 
	// Track cache hit rate for monitoring
	if len(texts) > 0 {
		es.hitRate = float64(hitCount) / float64(len(texts))
	}
 
	return embeddings, nil
}
 
// Version returns the embedding model version (stored in chunk metadata).
func (es *EmbeddingService) Version() string {
	return es.version
}

-- chunks table with pgvector
CREATE TABLE chunks (
    id uuid PRIMARY KEY,
    document_id uuid NOT NULL,
    content text NOT NULL,
    chunk_index int NOT NULL,
    embedding vector(1536),  -- OpenAI text-embedding-3-small
    embedding_version varchar(50) NOT NULL,  -- e.g., "text-embedding-3-small::20240415"
    token_count int,
    metadata jsonb,  -- {source, category, author, chunk_section}
    created_at timestamp DEFAULT now(),
    updated_at timestamp DEFAULT now()
);
 
-- HNSW index for vector search (m=16, ef_construction=200)
-- This trades insertion cost (~5% slower writes) for query speed (2-3x faster)
CREATE INDEX idx_chunks_embedding ON chunks USING hnsw (embedding vector_cosine_ops)
WITH (m = 16, ef_construction = 200);
 
-- Metadata index for filtering
CREATE INDEX idx_chunks_metadata ON chunks USING gin(metadata);
 
-- Partial index for current embedding version (ignore deprecated)
CREATE INDEX idx_chunks_current_version ON chunks (embedding_version)
WHERE embedding_version = 'text-embedding-3-small::20240415';
 
-- DocumentID index for deletion (idempotent re-ingestion)
CREATE INDEX idx_chunks_document_id ON chunks(document_id);

// retriever.go
package rag
 
import (
	"context"
	"database/sql"
	"fmt"
	"log/slog"
	pgvector "github.com/pgvector/pgvector-go"
)
 
type ChunkRepository struct {
	db     *sql.DB
	logger *slog.Logger
}
 
// Search retrieves top-k chunks by cosine similarity, with optional metadata filtering.
// Falls back to keyword search if vector search returns no results (edge case).
func (r *ChunkRepository) Search(
	ctx context.Context,
	queryEmbedding []float32,
	topK int,
	minScore float64,
	filters map[string]string,
) ([]SearchResult, error) {
	vec := pgvector.NewVector(queryEmbedding)
 
	// Vector similarity search (cosine <=> in pgvector)
	query := `
		SELECT id, document_id, content, metadata, token_count,
		       1 - (embedding <=> $1::vector) as similarity
		FROM chunks
		WHERE embedding_version = $2
		AND 1 - (embedding <=> $1::vector) >= $3
	`
	args := []interface{}{vec, "text-embedding-3-small::20240415", minScore}
 
	// Add metadata filters (exact match)
	filterIdx := 4
	for key, val := range filters {
		query += fmt.Sprintf(" AND metadata->>'%s' = $%d", key, filterIdx)
		args = append(args, val)
		filterIdx++
	}
 
	query += fmt.Sprintf(" ORDER BY embedding <=> $1 LIMIT $%d", filterIdx)
	args = append(args, topK)
 
	rows, err := r.db.QueryContext(ctx, query, args...)
	if err != nil {
		r.logger.Error("vector search failed", "error", err)
		return nil, fmt.Errorf("query error: %w", err)
	}
	defer rows.Close()
 
	var results []SearchResult
	for rows.Next() {
		var sr SearchResult
		var tokCount int
		if err := rows.Scan(&sr.Chunk.ID, &sr.Chunk.DocumentID, &sr.Chunk.Content,
			&sr.Chunk.Metadata, &tokCount, &sr.Score); err != nil {
			return nil, err
		}
		sr.Chunk.TokenCount = tokCount
		results = append(results, sr)
	}
 
	return results, rows.Err()
}

type EmbeddingCache struct {
	mu       sync.RWMutex
	cache    map[string][]float32
	lru      []string
	maxSize  int
}
 
func (ec *EmbeddingCache) GetOrCompute(text string, fn func() ([]float32, error)) ([]float32, error) {
	ec.mu.RLock()
	if emb, ok := ec.cache[text]; ok {
		ec.mu.RUnlock()
		return emb, nil
	}
	ec.mu.RUnlock()
 
	emb, err := fn()
	if err != nil {
		return nil, err
	}
 
	ec.mu.Lock()
	ec.cache[text] = emb
	ec.lru = append(ec.lru, text)
 
	// Evict oldest 20% when cache exceeds maxSize
	if len(ec.cache) > ec.maxSize {
		evict := ec.maxSize / 5
		for i := 0; i < evict && len(ec.lru) > 0; i++ {
			old := ec.lru[0]
			ec.lru = ec.lru[1:]
			delete(ec.cache, old)
		}
	}
	ec.mu.Unlock()
 
	return emb, nil
}

func (r *ChunkRepository) SearchWithFallback(
	ctx context.Context,
	queryEmbedding []float32,
	topK int,
) ([]SearchResult, error) {
	results, err := r.Search(ctx, queryEmbedding, topK, 0.5, nil)
	if err != nil && isConnectionError(err) {
		// Vector DB is down; return cached top-100 popular chunks
		r.logger.Warn("vector DB unavailable, using fallback cache")
		return r.cachedPopularResults(), nil
	}
	if err != nil {
		return nil, err
	}
	return results, nil
}
 
func isConnectionError(err error) bool {
	s := err.Error()
	return strings.Contains(s, "connection refused") ||
		strings.Contains(s, "i/o timeout") ||
		strings.Contains(s, "connection reset")
}

// mmr.go
import "math"
 
func ComputeMMR(results []SearchResult, lambda float64, topK int) []SearchResult {
	selected := []SearchResult{}
	remaining := results
 
	for len(selected) < topK && len(remaining) > 0 {
		bestIdx := 0
		bestScore := lambda*remaining[0].Score - (1-lambda)*maxSimilarityToSelected(remaining[0], selected)
 
		for i := 1; i < len(remaining); i++ {
			score := lambda*remaining[i].Score - (1-lambda)*maxSimilarityToSelected(remaining[i], selected)
			if score > bestScore {
				bestScore = score
				bestIdx = i
			}
		}
 
		selected = append(selected, remaining[bestIdx])
		remaining = append(remaining[:bestIdx], remaining[bestIdx+1:]...)
	}
 
	return selected
}
 
func cosineSimilarity(a, b []float32) float64 {
	var dot float64
	var normA, normB float64
	for i := range a {
		dot += float64(a[i]) * float64(b[i])
		normA += float64(a[i]) * float64(a[i])
		normB += float64(b[i]) * float64(b[i])
	}
	if normA == 0 || normB == 0 {
		return 0
	}
	return dot / (math.Sqrt(normA) * math.Sqrt(normB))
}
 
func maxSimilarityToSelected(r SearchResult, selected []SearchResult) float64 {
	maxSim := 0.0
	for _, s := range selected {
		sim := cosineSimilarity(r.Chunk.Embedding, s.Chunk.Embedding)
		if sim > maxSim {
			maxSim = sim
		}
	}
	return maxSim
}

// eval_test.go
package rag
 
import (
	"context"
	"testing"
	"time"
)
 
type EvalCase struct {
	Query            string
	ExpectedChunkIDs []string
	Category         string
}
 
func TestRetrievalQuality(t *testing.T) {
	goldenSet := loadGoldenTestSet("testdata/golden_evals.json")
 
	var hitRates, precisions, recalls []float64
	var latencies []time.Duration
 
	for _, tc := range goldenSet {
		emb, err := embedder.Embed(context.Background(), []string{tc.Query})
		if err != nil {
			t.Fatalf("embed: %v", err)
		}
 
		start := time.Now()
		results, err := repo.Search(context.Background(), emb[0], 10, 0.5, nil)
		latencies = append(latencies, time.Since(start))
		if err != nil {
			t.Fatalf("search: %v", err)
		}
 
		// Compute metrics
		retrieved := extractIDs(results)
		hit := contains(retrieved, tc.ExpectedChunkIDs[0])
		hitRates = append(hitRates, boolToFloat(hit))
 
		prec := precision(retrieved, tc.ExpectedChunkIDs)
		recall := recall(retrieved, tc.ExpectedChunkIDs)
		precisions = append(precisions, prec)
		recalls = append(recalls, recall)
	}
 
	// Assert thresholds (fail if regression detected)
	avgHit := avg(hitRates)
	if avgHit < 0.75 {
		t.Errorf("hit rate %.3f below threshold 0.75", avgHit)
	}
 
	avgPrec := avg(precisions)
	if avgPrec < 0.65 {
		t.Errorf("precision %.3f below threshold 0.65", avgPrec)
	}
 
	p99Latency := percentile(latencies, 99)
	if p99Latency > 500*time.Millisecond {
		t.Errorf("p99 latency %v exceeds 500ms threshold", p99Latency)
	}
 
	t.Logf("Hit Rate: %.3f | Precision: %.3f | Recall: %.3f | p99 Latency: %v",
		avgHit, avgPrec, avg(recalls), p99Latency)
}

// ingest.go
package rag
 
import (
	"context"
	"fmt"
	"log/slog"
	"time"
)
 
type Ingester struct {
	chunker  *SemanticChunker
	embedder *EmbeddingService
	repo     *ChunkRepository
	logger   *slog.Logger
}
 
func NewIngester(
	chunker *SemanticChunker,
	embedder *EmbeddingService,
	repo *ChunkRepository,
	logger *slog.Logger,
) *Ingester {
	return &Ingester{
		chunker:  chunker,
		embedder: embedder,
		repo:     repo,
		logger:   logger,
	}
}
 
// Ingest processes a document: chunk → embed → store (idempotent).
// Replaces all existing chunks for the document.
func (ing *Ingester) Ingest(ctx context.Context, doc Document) error {
	start := time.Now()
 
	// Step 1: Delete existing chunks (idempotent update)
	if err := ing.repo.DeleteByDocument(ctx, doc.ID); err != nil {
		return fmt.Errorf("delete existing chunks: %w", err)
	}
 
	// Step 2: Chunk the document
	chunks := ing.chunker.Chunk(doc)
	if len(chunks) == 0 {
		ing.logger.Warn("document produced no chunks",
			"document_id", doc.ID,
			"title", doc.Title)
		return nil
	}
 
	// Step 3: Embed chunks in batches (OpenAI limit: 2048 per request)
	const embedBatchSize = 256
	for i := 0; i < len(chunks); i += embedBatchSize {
		end := i + embedBatchSize
		if end > len(chunks) {
			end = len(chunks)
		}
 
		texts := make([]string, end-i)
		for j := i; j < end; j++ {
			texts[j-i] = chunks[j].Content
		}
 
		embeddings, err := ing.embedder.Embed(ctx, texts)
		if err != nil {
			return fmt.Errorf("embed batch %d-%d: %w", i, end, err)
		}
 
		for j, emb := range embeddings {
			chunks[i+j].Embedding = emb
			chunks[i+j].EmbeddingVersion = ing.embedder.Version()
		}
	}
 
	// Step 4: Store chunks
	if err := ing.repo.BatchUpsert(ctx, chunks); err != nil {
		return fmt.Errorf("store chunks: %w", err)
	}
 
	ing.logger.Info("document ingested",
		"document_id", doc.ID,
		"title", doc.Title,
		"chunks", len(chunks),
		"duration", time.Since(start),
	)
 
	return nil
}
 
// IngestBatch processes multiple documents concurrently with bounded parallelism.
func (ing *Ingester) IngestBatch(
	ctx context.Context, docs []Document, concurrency int,
) error {
	sem := make(chan struct{}, concurrency)
	errCh := make(chan error, len(docs))
 
	for _, doc := range docs {
		sem <- struct{}{}
		go func(d Document) {
			defer func() { <-sem }()
			if err := ing.Ingest(ctx, d); err != nil {
				errCh <- fmt.Errorf("ingest %s: %w", d.ID, err)
			}
		}(doc)
	}
 
	// Wait for all goroutines to finish
	for i := 0; i < cap(sem); i++ {
		sem <- struct{}{}
	}
	close(errCh)
 
	var errs []error
	for err := range errCh {
		errs = append(errs, err)
	}
 
	if len(errs) > 0 {
		return fmt.Errorf("%d documents failed ingestion: %w", len(errs), errs[0])
	}
	return nil
}