Semantic Document Search
Project 8: Build a Semantic Document Search Application
This project builds a search engine that ranks documents by meaning rather than keyword overlap, using the embeddings techniques from Unit 20. Unlike Project 2's File Search-based Q&A system, which answers questions using a managed retrieval pipeline, this project owns the entire retrieval stack directly — computing embeddings, storing them, and ranking by similarity — which is the right approach when the application needs control over ranking behavior, needs to combine semantic and keyword signals, or needs to search content that never goes through OpenAI's file ingestion pipeline at all (structured records, short text snippets, or ranked in combination with non-text signals).
Scope and Design Decisions
The application indexes a collection of text documents (support articles, blog posts, internal notes) and returns ranked results for a natural-language query, along with a similarity score. Three decisions shape the implementation:
- Documents are chunked before embedding, not embedded whole. A single embedding vector for an entire long document averages together many different topics, which weakens retrieval precision. Chunking preserves topical focus per vector.
- Cosine similarity is computed directly, not through a managed vector database, to keep the underlying mechanics visible. A production system at scale would use a dedicated vector database (pgvector, a hosted vector store) for the storage and nearest-neighbor search layer, but the ranking logic itself is the same regardless of storage backend, and implementing it directly here makes that logic inspectable.
- Search results include the matched chunk, not just the parent document. Returning the specific passage that matched is far more useful than returning "this document is relevant" and forcing the user to re-search within it.
Chunking and Embedding Documents
from openai import OpenAI
from dataclasses import dataclass
import numpy as np
client = OpenAI()
EMBEDDING_MODEL = "text-embedding-4"
@dataclass
class Chunk:
doc_id: str
chunk_index: int
text: str
embedding: np.ndarray | None = None
def chunk_document(doc_id: str, text: str, chunk_size: int = 500, overlap: int = 50) -> list[Chunk]:
words = text.split()
chunks = []
start = 0
index = 0
while start < len(words):
end = start + chunk_size
chunk_text = " ".join(words[start:end])
chunks.append(Chunk(doc_id=doc_id, chunk_index=index, text=chunk_text))
start += chunk_size - overlap
index += 1
return chunks
def embed_chunks(chunks: list[Chunk]) -> list[Chunk]:
texts = [c.text for c in chunks]
response = client.embeddings.create(model=EMBEDDING_MODEL, input=texts)
for chunk, item in zip(chunks, response.data):
chunk.embedding = np.array(item.embedding)
return chunks
chunk_document splits on word count with a fixed overlap between consecutive chunks, rather than a hard, non-overlapping cut. The overlap matters because a hard cut can split a sentence or a key phrase exactly at a chunk boundary, causing both resulting chunks to individually lose the context that made the passage meaningful — a 50-word overlap means that content near a boundary appears fully intact in at least one chunk. embed_chunks batches all chunk texts into a single embeddings.create call rather than one call per chunk, since the embeddings endpoint accepts a list of inputs and batching is both faster and cheaper than issuing one request per chunk.
Note: Practical chunk size should be tuned against the embedding model's context limit and the corpus's typical passage length; 500 words is a reasonable starting point for prose documents but not a fixed rule, and
text-embedding-4's exact input limits should be checked against current documentation.
Building a Simple In-Memory Index
class SemanticIndex:
def __init__(self):
self.chunks: list[Chunk] = []
def add_document(self, doc_id: str, text: str) -> None:
new_chunks = chunk_document(doc_id, text)
embed_chunks(new_chunks)
self.chunks.extend(new_chunks)
def search(self, query: str, top_k: int = 5) -> list[tuple[Chunk, float]]:
query_embedding = _embed_single(query)
scored = [
(chunk, _cosine_similarity(query_embedding, chunk.embedding))
for chunk in self.chunks
]
scored.sort(key=lambda pair: pair[1], reverse=True)
return scored[:top_k]
def _embed_single(text: str) -> np.ndarray:
response = client.embeddings.create(model=EMBEDDING_MODEL, input=[text])
return np.array(response.data[0].embedding)
def _cosine_similarity(a: np.ndarray, b: np.ndarray) -> float:
return float(np.dot(a, b) / (np.linalg.norm(a) * np.linalg.norm(b)))
_cosine_similarity measures the angle between two vectors rather than their raw distance, which is the standard choice for text embeddings because embedding magnitude carries little semantic meaning on its own — two vectors pointing in nearly the same direction represent similar meaning regardless of their length, and cosine similarity captures exactly that while ignoring magnitude differences. SemanticIndex.search recomputes similarity against every stored chunk with a linear scan, which is correct and simple but scales linearly with corpus size — acceptable for a few thousand chunks, but the point where a real vector database with approximate nearest-neighbor indexing (HNSW or similar) becomes necessary rather than optional, typically well before a million chunks.
Combining Semantic and Keyword Search
def keyword_score(query: str, text: str) -> float:
query_terms = set(query.lower().split())
text_terms = set(text.lower().split())
if not query_terms:
return 0.0
return len(query_terms & text_terms) / len(query_terms)
def hybrid_search(index: SemanticIndex, query: str, top_k: int = 5, semantic_weight: float = 0.7) -> list[tuple[Chunk, float]]:
query_embedding = _embed_single(query)
scored = []
for chunk in index.chunks:
semantic = _cosine_similarity(query_embedding, chunk.embedding)
keyword = keyword_score(query, chunk.text)
combined = semantic_weight * semantic + (1 - semantic_weight) * keyword
scored.append((chunk, combined))
scored.sort(key=lambda pair: pair[1], reverse=True)
return scored[:top_k]
Pure semantic search occasionally misses exact-match queries — a specific error code, a product SKU, an acronym — because embeddings capture general meaning rather than precise lexical identity, and a short, specific token can end up embedded close to many unrelated passages that share a similar general topic. hybrid_search blends a cheap keyword-overlap score with the semantic score, weighted so semantic similarity still dominates (semantic_weight=0.7 by default) while exact-term matches still get a meaningful boost. This weighting is a genuine tuning knob: a support-ticket search corpus full of specific error codes benefits from a lower semantic_weight, while a corpus of narrative documentation benefits from leaning further toward pure semantic matching.
Returning Results With Document Context
def format_search_results(results: list[tuple[Chunk, float]]) -> list[dict]:
return [
{
"doc_id": chunk.doc_id,
"chunk_index": chunk.chunk_index,
"excerpt": chunk.text[:300] + ("..." if len(chunk.text) > 300 else ""),
"score": round(score, 4),
}
for chunk, score in results
]
Truncating the excerpt to 300 characters keeps a search results list scannable — a full 500-word chunk in a results list defeats the purpose of ranking, since the user still has to read through it to find the relevant sentence. A more advanced version would locate and center the excerpt on the specific sentence most similar to the query rather than always showing the chunk's start, which is a natural next refinement once basic search is working.
Testing Ranking Behavior Without Real Embeddings
def test_cosine_similarity_ranks_closer_vectors_higher():
identical = np.array([1.0, 0.0, 0.0])
similar = np.array([0.9, 0.1, 0.0])
orthogonal = np.array([0.0, 1.0, 0.0])
score_similar = _cosine_similarity(identical, similar)
score_orthogonal = _cosine_similarity(identical, orthogonal)
assert score_similar > score_orthogonal
print("PASS: cosine similarity ranks a near-identical vector above an orthogonal one")
def test_hybrid_search_boosts_exact_keyword_match(monkeypatch_embed):
index = SemanticIndex()
chunk_a = Chunk(doc_id="doc1", chunk_index=0, text="general discussion of cloud pricing models")
chunk_b = Chunk(doc_id="doc2", chunk_index=0, text="error code E-4021 troubleshooting steps")
# Give both chunks identical embeddings so only the keyword term can differentiate them.
shared_vector = np.array([1.0, 0.0, 0.0])
chunk_a.embedding = shared_vector
chunk_b.embedding = shared_vector
index.chunks = [chunk_a, chunk_b]
monkeypatch_embed(lambda text: shared_vector)
results = hybrid_search(index, "E-4021", semantic_weight=0.5)
assert results[0][0].doc_id == "doc2"
print("PASS: exact keyword match outranks an equally-similar embedding when scores tie semantically")
def _make_monkeypatch():
import builtins
module = globals()
def apply(fake_fn):
module["_embed_single"] = fake_fn
return apply
monkeypatch_embed = _make_monkeypatch()
test_cosine_similarity_ranks_closer_vectors_higher()
test_hybrid_search_boosts_exact_keyword_match(monkeypatch_embed)
The first test uses hand-constructed numpy vectors with no API call at all, which is possible because cosine similarity is pure math — testing it needs no dependency injection. The second test does inject a fake _embed_single and deliberately makes both chunks' embeddings identical, isolating the keyword-scoring contribution of hybrid_search from the semantic contribution — this is the key testing technique for a hybrid scoring function: hold one signal constant so the other signal's effect becomes directly observable in the test's assertions.
Extending This Project
Migrate the in-memory index to pgvector or a managed vector database once the corpus exceeds a few thousand chunks, and add a re-ranking pass that sends the top 20 semantic results plus the query to the main language model for a final relevance re-ordering, which typically improves precision at the cost of one additional API call per search.
Common Mistakes
- Embedding entire long documents as a single vector. This dilutes the vector's semantic focus across every topic in the document and produces poor retrieval precision. Chunk first.
- Chunking without any overlap. A hard, non-overlapping split can sever a sentence or concept exactly at a chunk boundary, degrading both resulting chunks. A modest overlap (roughly 10 percent of chunk size) is cheap insurance.
- Relying on pure semantic search for queries containing exact identifiers. Error codes, SKUs, and acronyms are often poorly served by embeddings alone. Blend in a keyword signal for corpora where such terms are common.
Best Practices
- Batch embedding calls across chunks rather than issuing one call per chunk. The embeddings endpoint accepts a list input; use it.
- Return the matched passage, not just the parent document, in search results. This is what actually saves the user time.
- Move to a dedicated vector database once linear-scan search becomes the bottleneck. The ranking math stays the same; only the storage and lookup mechanism needs to change at scale.