Code: RAG with Embeddings

1# part3/rag_code.py
2from openai import OpenAI
3client = OpenAI()
4documents = [  # facts from a made-up school handbook
5    "The library is open from 8 am to 4 pm on weekdays.",
6    "Students can borrow up to 3 books for 2 weeks.",
7    "The science fair is on 14 March in the school gym.",
8    "Lunch is from 12:00 to 12:45 in the cafeteria.",
9    "Lost-and-found items are kept at the front office.",
10]
11
12def embed(texts):  # each text -> a long list of numbers
13    result = client.embeddings.create(
14        model="text-embedding-3-small", input=texts)
15    return [item.embedding for item in result.data]
16
17def similarity(a, b):  # 1 = same direction (very alike)
18    dot = sum(x * y for x, y in zip(a, b))
19    length_a = sum(x * x for x in a) ** 0.5
20    length_b = sum(x * x for x in b) ** 0.5
21    return dot / (length_a * length_b)
22
23# 1. Retrieve: find the 2 documents closest in meaning
24doc_vectors = embed(documents)  # embed the handbook once
25question = "How many books can I borrow?"
26question_vector = embed([question])[0]
27scores = [similarity(question_vector, doc_vector)
28          for doc_vector in doc_vectors]
29ranked = sorted(zip(scores, documents), reverse=True)
30for score, doc in ranked:
31    print(round(score, 2), doc)
32top_docs = [doc for score, doc in ranked[:2]]
33
34# 2. Augment: put the top 2 into the prompt
35context = "\n".join(top_docs)
36system_prompt = """Answer using ONLY the context.
37If the answer isn't there, say you don't know."""
38user_message = f"Context:\n{context}\n\nQuestion: {question}"
39
40# 3. Generate: the model answers from the context
41messages = [{"role": "system", "content": system_prompt},
42            {"role": "user", "content": user_message}]
43response = client.chat.completions.create(
44    model="gpt-4o-mini", messages=messages)
45print("Answer:", response.choices[0].message.content)