Retrieval

With both the knowledge base and the query turned into vectors, this lesson finds which stored chunks are actually relevant to the question, by measuring how similar their vectors are.

Cosine similarity

The standard way to compare two vectors for similarity is cosine similarity: the cosine of the angle between them, ranging from -1 (opposite) to 1 (identical direction), regardless of each vector's magnitude. For sparse Counter-based vectors, it only needs to look at the words that appear in both:

Pythonshares state with other Python blocks

import re
import math
from collections import Counter

def toy_embed(text):
    words = re.findall(r"[a-z']+", text.lower())
    return Counter(words)

def cosine_similarity(vec1, vec2):
    common_words = set(vec1) & set(vec2)
    dot_product = sum(vec1[word] * vec2[word] for word in common_words)
    magnitude1 = math.sqrt(sum(count * count for count in vec1.values()))
    magnitude2 = math.sqrt(sum(count * count for count in vec2.values()))
    if magnitude1 == 0 or magnitude2 == 0:
        return 0.0
    return dot_product / (magnitude1 * magnitude2)

a = toy_embed('Password reset emails expire after one hour.')
b = toy_embed('How do I reset my password?')
c = toy_embed('Kodstigen is free for all students.')

print('a vs b (related):', cosine_similarity(a, b))
print('a vs c (unrelated):', cosine_similarity(a, c))

import re
import math
from collections import Counter

def toy_embed(text):
    words = re.findall(r"[a-z']+", text.lower())
    return Counter(words)

def cosine_similarity(vec1, vec2):
    common_words = set(vec1) & set(vec2)
    dot_product = sum(vec1[word] * vec2[word] for word in common_words)
    magnitude1 = math.sqrt(sum(count * count for count in vec1.values()))
    magnitude2 = math.sqrt(sum(count * count for count in vec2.values()))
    if magnitude1 == 0 or magnitude2 == 0:
        return 0.0
    return dot_product / (magnitude1 * magnitude2)

a = toy_embed('Password reset emails expire after one hour.')
b = toy_embed('How do I reset my password?')
c = toy_embed('Kodstigen is free for all students.')

print('a vs b (related):', cosine_similarity(a, b))
print('a vs c (unrelated):', cosine_similarity(a, c))

Run it, the password-related pair should score noticeably higher than the unrelated pair, exactly what makes retrieval work.

Retrieving the top-K matches

Scoring the query against every stored chunk and keeping only the best few (top-K) is the retrieval step itself:

Pythonshares state with other Python blocks

class VectorStore:
    def __init__(self):
        self.entries = []

    def add(self, text, vector, metadata=None):
        self.entries.append({'text': text, 'vector': vector, 'metadata': metadata or {}})

store = VectorStore()
store.add('Password reset emails expire after one hour.', toy_embed('Password reset emails expire after one hour.'), {'source': 'faq.md'})
store.add('Kodstigen is currently free for all students.', toy_embed('Kodstigen is currently free for all students.'), {'source': 'billing.md'})
store.add('New students should start with the Python path.', toy_embed('New students should start with the Python path.'), {'source': 'onboarding.md'})

def retrieve(store, query_vector, top_k=2):
    scored = [
        (cosine_similarity(query_vector, entry['vector']), entry)
        for entry in store.entries
    ]
    scored.sort(key=lambda pair: pair[0], reverse=True)
    return [entry for score, entry in scored[:top_k]]

query_vector = toy_embed('How do I reset my password?')
results = retrieve(store, query_vector, top_k=2)
for entry in results:
    print(entry['metadata']['source'], '->', entry['text'][:60])

class VectorStore:
    def __init__(self):
        self.entries = []

    def add(self, text, vector, metadata=None):
        self.entries.append({'text': text, 'vector': vector, 'metadata': metadata or {}})

store = VectorStore()
store.add('Password reset emails expire after one hour.', toy_embed('Password reset emails expire after one hour.'), {'source': 'faq.md'})
store.add('Kodstigen is currently free for all students.', toy_embed('Kodstigen is currently free for all students.'), {'source': 'billing.md'})
store.add('New students should start with the Python path.', toy_embed('New students should start with the Python path.'), {'source': 'onboarding.md'})

def retrieve(store, query_vector, top_k=2):
    scored = [
        (cosine_similarity(query_vector, entry['vector']), entry)
        for entry in store.entries
    ]
    scored.sort(key=lambda pair: pair[0], reverse=True)
    return [entry for score, entry in scored[:top_k]]

query_vector = toy_embed('How do I reset my password?')
results = retrieve(store, query_vector, top_k=2)
for entry in results:
    print(entry['metadata']['source'], '->', entry['text'][:60])

sort(..., reverse=True) puts the highest similarity scores first, then [:top_k] keeps only that many, everything else is discarded, only the most relevant handful of chunks moves on to the next step.

TIP

top_k is another tunable knob: too small and you might miss a relevant chunk, too large and you dilute the final prompt with irrelevant text (and pay for more tokens). 2-5 is a common starting point for small knowledge bases.

Retrieval

Cosine similarity

Retrieving the top-K matches

📝 Retrieval Quiz