Building a Simple Semantic Search Engine in Python
Building a Simple Semantic Search Engine in Python
This lesson assembles the pieces from the previous four lessons — text preparation, embedding generation, and cosine similarity — into a small, complete, reusable search engine class. Unlike Unit 10's notes project, which fed retrieved text into a language model to generate an answer (retrieval-augmented generation), this project stops at retrieval itself: given a query, return the most relevant documents, ranked, with their scores. That is the shape of an actual search feature — a "search this knowledge base" box, a "find similar listings" feature, a "related articles" widget — where the goal is to surface the right documents to a human, not to generate new text from them. The corpus used here is a small set of recipe descriptions, chosen because it is a different kind of content from the notes-style text in Unit 10 and demonstrates that the same engine works on any domain of short, descriptive text.
Designing the Engine's Interface First
Before writing the implementation, it helps to decide what the class needs to do, because that shapes every internal choice:
- Index documents: accept a list of texts (and optional IDs/metadata), clean them, embed them, and store the vectors alongside the original text.
- Search: accept a query string, embed it the same way documents were embedded, compare it against every stored vector, and return the top matches ranked by similarity.
- Be testable without real API calls: like the
EmbeddingClientin Lesson 2, the engine should accept its embedding function as a dependency rather than hard-coding the OpenAI client inside it.
The Core Implementation
import re
from dataclasses import dataclass, field
def clean_text(text: str) -> str:
text = re.sub(r"<[^>]+>", " ", text)
return re.sub(r"\s+", " ", text).strip()
def cosine_similarity(a: list[float], b: list[float]) -> float:
dot = sum(x * y for x, y in zip(a, b))
mag_a = sum(x * x for x in a) ** 0.5
mag_b = sum(x * x for x in b) ** 0.5
if mag_a == 0 or mag_b == 0:
return 0.0
return dot / (mag_a * mag_b)
@dataclass
class Document:
doc_id: str
text: str
embedding: list[float] = field(default_factory=list)
class SemanticSearchEngine:
"""A minimal in-memory semantic search engine.
`embed_fn` takes a list of strings and returns a list of embedding
vectors in the same order — this is the same contract implemented
by `EmbeddingClient.embed` from Lesson 2, so a real OpenAI-backed
client can be swapped in without changing this class.
"""
def __init__(self, embed_fn):
self._embed_fn = embed_fn
self._documents: list[Document] = []
def index(self, texts: list[str], ids: list[str] | None = None) -> None:
ids = ids or [str(i) for i in range(len(texts))]
cleaned = [clean_text(t) for t in texts]
vectors = self._embed_fn(cleaned)
for doc_id, original_text, vector in zip(ids, texts, vectors):
self._documents.append(Document(doc_id=doc_id, text=original_text, embedding=vector))
def search(self, query: str, top_k: int = 3) -> list[tuple[Document, float]]:
if not self._documents:
return []
query_vector = self._embed_fn([clean_text(query)])[0]
scored = [
(doc, cosine_similarity(query_vector, doc.embedding))
for doc in self._documents
]
scored.sort(key=lambda pair: pair[1], reverse=True)
return scored[:top_k]
A few design decisions are worth explaining. Document is a small dataclass that keeps the original (uncleaned) text alongside its embedding, because the cleaned version is only needed at embedding time — search results should show the reader the real text, not the stripped-down version fed to the model. index() cleans every text with the same clean_text function used in Lesson 3 before embedding, and search() applies the identical cleaning to the query, which matters because Lesson 3 established that indexing and querying must be preprocessed consistently. embed_fn is injected through the constructor rather than imported directly, which is what makes this class testable without hitting the network, shown next.
Testing the Engine With a Fake Embedder
def fake_embed_fn(texts: list[str]) -> list[list[float]]:
"""A deterministic fake embedder for tests.
Produces a vector that encodes, crudely, whether each of a few
known keywords appears in the text — enough structure to make
similarity rankings meaningful in a test, without calling any API.
"""
keywords = ["chocolate", "chicken", "spicy", "dessert"]
vectors = []
for text in texts:
lowered = text.lower()
vectors.append([1.0 if kw in lowered else 0.0 for kw in keywords])
return vectors
def test_semantic_search_engine_ranks_relevant_first():
engine = SemanticSearchEngine(embed_fn=fake_embed_fn)
engine.index(
texts=[
"Rich chocolate lava cake, a warm dessert.",
"Grilled chicken with lemon and herbs.",
"Spicy chicken curry with coconut milk.",
],
ids=["cake", "grilled_chicken", "curry"],
)
results = engine.search("a spicy chicken dish", top_k=2)
top_id = results[0][0].doc_id
assert top_id == "curry"
assert results[0][1] >= results[1][1]
print("PASS: SemanticSearchEngine ranks the most relevant document first")
test_semantic_search_engine_ranks_relevant_first()
fake_embed_fn is not a real embedding model — it is a small, deterministic stand-in that maps keyword presence to fixed vector positions, purely so the test has predictable, explainable similarity relationships to assert against. This is the same dependency-injection testing pattern used throughout this course: the test verifies that SemanticSearchEngine's indexing and ranking logic is correct, independent of whatever a real embedding model would actually produce. A separate, smaller set of manual or integration checks (not run automatically, and not shown as unit tests) would be the place to confirm that the real text-embedding-4 model produces sensible rankings on real recipe text — mixing that concern into an automated test suite would make the suite slow, costly, and dependent on network access.
Wiring In the Real Embedding Client
Swapping the fake embedder for the real one from Lesson 2 requires no change to SemanticSearchEngine itself — only the function passed into its constructor changes.
from openai import OpenAI
client = OpenAI()
def real_embed_fn(texts: list[str]) -> list[list[float]]:
response = client.embeddings.create(model="text-embedding-4", input=texts)
return [item.embedding for item in response.data]
engine = SemanticSearchEngine(embed_fn=real_embed_fn)
engine.index(
texts=[
"Rich chocolate lava cake with a molten center, served warm.",
"Grilled chicken breast marinated in lemon, garlic, and herbs.",
"Spicy Thai-style chicken curry with coconut milk and chili.",
"No-bake cheesecake with a graham cracker crust.",
],
ids=["lava_cake", "grilled_chicken", "thai_curry", "cheesecake"],
)
for doc, score in engine.search("something sweet and creamy", top_k=2):
print(f"{score:.3f} {doc.doc_id} — {doc.text}")
This is the same class, unchanged, now backed by a real embedding model. real_embed_fn matches the exact contract SemanticSearchEngine expects (a list of strings in, a list of vectors out), which is precisely why the class did not need to know or care whether it was talking to a fake function or a real API during development and testing.
Extending the Engine: Returning Structured Results
A raw list of (Document, float) tuples is fine internally, but an application layer (an API endpoint, a UI) usually wants plain, serializable data.
def search_as_dicts(engine: SemanticSearchEngine, query: str, top_k: int = 3) -> list[dict]:
results = engine.search(query, top_k=top_k)
return [
{"id": doc.doc_id, "text": doc.text, "score": round(score, 4)}
for doc, score in results
]
Separating this formatting step from SemanticSearchEngine.search() itself keeps the core class focused purely on retrieval logic, while formatting concerns (rounding scores, converting to dictionaries, adding pagination) live at the boundary where the engine meets the rest of an application — a separation that pays off once, as in Lesson 7, the engine needs to also support filtering results by metadata before they are returned.
Common Mistakes
- Re-embedding the entire corpus on every search.
SemanticSearchEngine.index()embeds documents once, up front;search()only embeds the (single, short) query. Recomputing document embeddings on every query call wastes cost and time for no benefit, since the documents have not changed. - Forgetting to keep the mapping between a document's text/ID and its vector. Without storing them together (as
Documentdoes here), a search result is just a bare score with nothing to show the user. - Hard-coding the
OpenAIclient inside the search class. This makes the class impossible to unit test without real API calls and real cost, and harder to reuse if the embedding backend ever changes.
Best Practices
- Keep the search engine's core logic free of any specific embedding provider. Accepting an
embed_fn(or anEmbeddingClient-like object) as a dependency, as done here, keeps the retrieval logic portable and testable. - Preserve the original, uncleaned text for display, even when a cleaned version is used for embedding. Users should see readable, natural results, not the stripped-down text sent to the model.
- Separate ranking logic from presentation formatting. Returning rich internal objects from
search()and converting them to plain dictionaries or JSON at the boundary (assearch_as_dictsdoes) keeps each function focused and easy to test independently.