Semantic Search Implementation for Agent Knowledge Bases: Building Vector Embeddings

Keyword search breaks the moment your users ask questions your documents don’t literally contain. An agent that searches for “employee offboarding process” will miss the document titled “Termination Checklist and IT Deprovisioning Steps” — even though it’s exactly what they need. Semantic search embeddings solve this by converting both queries and documents into vectors in a shared meaning space, where proximity equals relevance regardless of exact wording. This article shows you how to build that retrieval layer from scratch: embedding models, chunking strategies, vector storage, and ranking — with working Python code throughout.

Why Keyword Search Fails Agent Knowledge Bases

Agents retrieving context from a knowledge base using BM25 or simple string matching hit a predictable ceiling. The failure mode isn’t edge cases — it’s the majority of real-world queries. Users paraphrase. Documents use jargon. Questions are conceptual and answers are procedural. The vocabulary mismatch between “how do I cancel a subscription” and your internal doc titled “Membership Termination Policy” is enough to return zero results or rank the wrong chunk first.

Embedding-based retrieval encodes meaning as a high-dimensional vector. You embed everything once at index time, then embed the query at retrieval time, and compare them with cosine similarity or dot product. The query and document never need to share a single word to match — they just need to occupy similar regions in vector space.

The tradeoff: it costs compute to generate embeddings, you need a vector store, and retrieval quality depends heavily on your chunking strategy and embedding model choice. None of this is hard, but all of it matters.

Choosing Your Embedding Model

The model you pick determines your vector quality ceiling. You can’t fix a bad embedding model downstream with better retrieval logic — garbage in, garbage out at query time.

OpenAI text-embedding-3-small and 3-large

text-embedding-3-small is my default recommendation for most production setups. It produces 1536-dimensional vectors, costs $0.02 per million tokens (roughly $0.000002 per average 100-token chunk), and performs well on standard retrieval benchmarks. For 10,000 document chunks at 200 tokens each, you’re spending about $0.04 to build the index. That’s a rounding error.

text-embedding-3-large (3072 dimensions) performs better on multilingual content and complex reasoning tasks but costs 13x more. I’d only upgrade if you’re embedding technical documentation with heavy domain-specific vocabulary or multilingual corpora.

Sentence Transformers (Open Source)

all-MiniLM-L6-v2 runs locally, is fast, and is good enough for many use cases. If you’re embedding on the fly in a self-hosted setup or have privacy constraints preventing API calls, start here. The quality gap vs. OpenAI’s models is real but not catastrophic for English-language retrieval tasks.

BAAI/bge-large-en-v1.5 is the open-source model I’d reach for when quality matters more than inference speed. It consistently benchmarks close to OpenAI’s 3-small on MTEB while running fully local.

Cohere Embed v3

Cohere’s embed-english-v3.0 is worth considering if you’re already in their ecosystem. It supports an input_type parameter that lets you distinguish between query and document embeddings — which actually improves retrieval quality by optimizing the vector space for asymmetric search rather than treating both sides identically. At $0.10 per million tokens, it’s 5x more expensive than OpenAI’s 3-small, but the asymmetric embedding support is a genuine differentiator.

Chunking Strategy: This Is Where Most Implementations Break

Bad chunking kills retrieval quality more often than bad embedding models. The goal is chunks that are semantically coherent — each chunk should contain one complete idea, not half a sentence from one paragraph and the beginning of the next.

Fixed-Size Chunking with Overlap

The naive approach: split every N tokens with a K-token overlap between consecutive chunks. It’s fast to implement and works well enough for homogeneous content like FAQ documents or dense technical reference material.

from tiktoken import encoding_for_model

def chunk_text(text: str, chunk_size: int = 200, overlap: int = 40) -> list[str]:
    """
    Split text into overlapping token chunks.
    overlap prevents context loss at chunk boundaries.
    """
    enc = encoding_for_model("text-embedding-3-small")
    tokens = enc.encode(text)
    chunks = []
    
    start = 0
    while start < len(tokens):
        end = min(start + chunk_size, len(tokens))
        chunk_tokens = tokens[start:end]
        chunks.append(enc.decode(chunk_tokens))
        # Move forward by chunk_size minus overlap
        start += chunk_size - overlap
    
    return chunks

The overlap parameter is load-bearing. Without it, a concept that spans a chunk boundary gets cut in half, and neither chunk retrieves correctly. 15–20% overlap is a reasonable starting point.

Semantic Chunking

Better but slower: split on sentence boundaries and merge sentences until adding the next one would exceed your target size. This keeps sentences intact and produces chunks that feel more natural to the embedding model.

import spacy

nlp = spacy.load("en_core_web_sm")

def semantic_chunk(text: str, max_tokens: int = 200) -> list[str]:
    """
    Chunk by sentence boundaries using spaCy sentence detection.
    Merges sentences until max_tokens would be exceeded.
    """
    from tiktoken import encoding_for_model
    enc = encoding_for_model("text-embedding-3-small")
    
    doc = nlp(text)
    sentences = [sent.text.strip() for sent in doc.sents]
    
    chunks = []
    current_chunk = []
    current_tokens = 0
    
    for sentence in sentences:
        sent_tokens = len(enc.encode(sentence))
        
        if current_tokens + sent_tokens > max_tokens and current_chunk:
            chunks.append(" ".join(current_chunk))
            current_chunk = [sentence]
            current_tokens = sent_tokens
        else:
            current_chunk.append(sentence)
            current_tokens += sent_tokens
    
    if current_chunk:
        chunks.append(" ".join(current_chunk))
    
    return chunks

For knowledge bases with structured documents (policies, SOPs, product docs), I prefer semantic chunking. For conversational transcripts or freeform text, fixed-size with overlap is usually fine.

Building the Embedding Pipeline

Here’s a minimal but production-ready embedding pipeline using OpenAI’s API with rate limiting and batching handled correctly:

import openai
import time
from typing import list

client = openai.OpenAI()  # uses OPENAI_API_KEY env var

def embed_chunks(
    chunks: list[str],
    model: str = "text-embedding-3-small",
    batch_size: int = 100
) -> list[list[float]]:
    """
    Embed a list of text chunks in batches.
    OpenAI allows up to 2048 inputs per request but batching at 100
    keeps request sizes manageable and avoids timeout issues.
    """
    all_embeddings = []
    
    for i in range(0, len(chunks), batch_size):
        batch = chunks[i:i + batch_size]
        
        try:
            response = client.embeddings.create(
                input=batch,
                model=model
            )
            # Response embeddings are ordered to match input
            batch_embeddings = [item.embedding for item in response.data]
            all_embeddings.extend(batch_embeddings)
            
        except openai.RateLimitError:
            time.sleep(60)  # Back off for a minute, then retry
            response = client.embeddings.create(input=batch, model=model)
            all_embeddings.extend([item.embedding for item in response.data])
    
    return all_embeddings

Vector Storage Options

You have three meaningful choices depending on your scale and infrastructure preferences.

ChromaDB (Local, Zero Setup)

For prototyping and single-server deployments under ~1M vectors, ChromaDB is the fastest path. It runs in-process with no external dependencies, persists to disk, and has a clean Python API.

import chromadb

# Persistent client — survives restarts
chroma_client = chromadb.PersistentClient(path="./chroma_db")
collection = chroma_client.get_or_create_collection(
    name="knowledge_base",
    metadata={"hnsw:space": "cosine"}  # Use cosine similarity
)

def index_documents(chunks: list[str], embeddings: list[list[float]], doc_ids: list[str]):
    collection.add(
        documents=chunks,
        embeddings=embeddings,
        ids=doc_ids
    )

def retrieve(query_embedding: list[float], top_k: int = 5) -> list[dict]:
    results = collection.query(
        query_embeddings=[query_embedding],
        n_results=top_k,
        include=["documents", "distances", "metadatas"]
    )
    return results

Pinecone (Managed, Production Scale)

When you need multi-tenant isolation, horizontal scaling, or you’re embedding across microservices, Pinecone removes infrastructure overhead. Their serverless tier starts at $0.04 per million vectors stored per month and $2 per million query units — reasonable for most agent deployments. The main annoyance: index creation is asynchronous and the free tier has only one project namespace.

pgvector (Postgres Extension)

If you’re already running Postgres, pgvector lets you add a vector column to existing tables and query with cosine similarity using standard SQL. This is genuinely underrated for teams that want semantic search without adding a new database to manage. Performance degrades past ~500K vectors without careful HNSW index tuning, but for most knowledge base sizes it’s more than sufficient.

Retrieval Ranking: Don’t Just Return the Top Result

Raw cosine similarity ranking has a known weakness: it scores each chunk independently, so a document with three mediocre chunks can collectively beat one document with the single best chunk. For agent knowledge retrieval, you care about the best matching passages, not average document relevance.

Reranking with Cross-Encoders

A two-stage retrieval pipeline — retrieve top-K with embeddings, then rerank with a cross-encoder — consistently outperforms single-stage embedding retrieval. Cohere’s Rerank API (rerank-english-v3.0) takes your query and candidate passages and returns relevance scores without you managing the model. Cost is $2 per 1,000 searches, which adds up at scale but is worth it for knowledge-critical retrieval.

import cohere

co = cohere.Client()  # uses COHERE_API_KEY env var

def rerank_results(query: str, candidate_docs: list[str], top_n: int = 3) -> list[dict]:
    """
    Rerank retrieved candidates using Cohere's cross-encoder.
    Call this after initial vector retrieval, not instead of it.
    """
    response = co.rerank(
        query=query,
        documents=candidate_docs,
        top_n=top_n,
        model="rerank-english-v3.0"
    )
    
    return [
        {
            "document": result.document["text"],
            "relevance_score": result.relevance_score,
            "original_index": result.index
        }
        for result in response.results
    ]

In my testing, this two-stage pipeline typically improves retrieval precision by 15–30% on knowledge bases with mixed-quality content — meaning your agent gets the right context chunk rather than a plausible-sounding but wrong one.

Metadata Filtering: The Underused Power Feature

Pure semantic search is globally promiscuous — it will retrieve the most relevant chunk from any document. For agent knowledge bases, you usually want scoped retrieval: only search documents relevant to this user’s role, this product version, or this department. Store metadata alongside your embeddings and filter before or during vector search.

# When indexing, attach metadata to each chunk
collection.add(
    documents=chunks,
    embeddings=embeddings,
    ids=doc_ids,
    metadatas=[
        {"source": "hr-policy", "department": "engineering", "version": "2024-Q1"}
        for _ in chunks
    ]
)

# At query time, filter by metadata before ranking
results = collection.query(
    query_embeddings=[query_embedding],
    n_results=10,
    where={"department": "engineering"},  # Pre-filter
    include=["documents", "distances"]
)

This pattern is essential when your knowledge base spans multiple tenants or domains. Without it, semantic search will happily return a perfectly matched chunk from the wrong organizational context.

When to Use This and Who Should Build It

Solo founders and small teams: Start with ChromaDB + OpenAI text-embedding-3-small + semantic chunking. You can have this running in an afternoon. Don’t add reranking until you’ve confirmed retrieval quality is actually your bottleneck — it usually isn’t in the early stages.

Teams building multi-user products: Add metadata filtering from day one, even if you only have one tenant now. Retrofitting it later requires reindexing everything. Use Pinecone or pgvector depending on whether you want managed infrastructure or Postgres consolidation.

High-stakes retrieval (legal, medical, compliance): The two-stage embedding + reranking pipeline is non-negotiable. A missed retrieval that causes an agent to give wrong compliance guidance is a real business risk. Pay the reranking cost.

The core insight that makes semantic search embeddings worth building: your agent is only as good as the context it retrieves. A perfectly prompted agent with bad retrieval will consistently fail. A mediocre prompt with precise retrieval will consistently succeed. Invest in the retrieval layer first — it’s the foundation everything else runs on.