Simple retrieval augmented generation#

In this notebook we see how retrieval augmented generation (RAG) works using OpenAI and numpy. This implementation avoids using complex libraries intentionally. To keep the code simple, we are using Euclidean distances to determine related entries in the knowledge base. Maximum inner product search is more common in the field though.

import numpy as np
import openai
from IPython.display import Markdown, display
from llama_index.embeddings.huggingface import HuggingFaceEmbedding
from dotenv import load_dotenv
load_dotenv()

def show(text):
    display(Markdown(text))

We aim to answer this question:

question = "What product can I buy for brushing teeth?"

… using these product descriptions (and more):

with open('household-products.md', 'r') as file:
    all_product_descriptions = file.read()

splits = all_product_descriptions.split("\n")
[show(s) for s in splits[:5]];

  • Frostmancer - A freezer that keeps food frozen with steady cold storage. Good for bulk meals, ice cream, and leftovers.

  • Chillwyrm - A refrigerator that keeps groceries fresh and drinks cold. It has room for everyday staples.

  • Dustbane - A vacuum cleaner for carpets, rugs, and hard floors. It pulls up crumbs, dust, and pet hair.

  • Brushwick - An electric toothbrush that helps clean teeth with powered bristles. It makes daily brushing easier.

  • Zapoven - A microwave oven for reheating meals and drinks quickly. Useful for leftovers and quick snacks.

Vector embeddings#

To make our code snippets searchable, we need to created vector embedding form them, we need to turn them into vectors.

e = HuggingFaceEmbedding(model_name="intfloat/multilingual-e5-large-instruct")
vector = e.get_text_embedding("Hello world")

len(vector)

1024

vector[:3]

[0.00521087646484375, 0.0252685546875, 0.007358551025390625]

Vector store#

We also need a vector store, which is basically just a dictionary that allows us to quickly find a text given a corresponding vector, or a vector that has a short distance to it.

class VectorStore:
    def __init__(self, texts=None):
        self._store = {}
        if texts is not None:
            for text in texts:
                self._store[tuple(e.get_text_embedding(text))] = text
    
    def search(self, text, n_best_results=3):
        single_vector = e.get_text_embedding(text)
        
        # Step 1: Compute Euclidean distances
        distances = [(np.linalg.norm(np.asarray(single_vector) - np.asarray(vector)), vector) 
                     for vector in self._store.keys()]

        # Step 2: Sort distances and get the n vectors with the shortest distances
        distances.sort()  # Sort based on the first element in the tuple (distance)
        closest_vectors = [vec for _, vec in distances[:n_best_results]]
        
        return [self._store[tuple(v)] for v in closest_vectors]
    
    def get_text(self, vector):
        return self._store[vector]

vectore_store = VectorStore(splits)

Searching the vector store#

We can then search in the store for vectors and corresponding texts that are close by a given question.

question

'What product can I buy for brushing teeth?'

vector = e.get_text_embedding(question)
vector[:3]

[-0.00042939186096191406, 0.03106689453125, -0.056121826171875]

related_products = vectore_store.search(question)
show("\n\n".join(related_products))

  • Brushwick - An electric toothbrush that helps clean teeth with powered bristles. It makes daily brushing easier.

  • SugarBrush - A toothbrush made of pure sugar, for the extra sweet smile.

  • Purepixie - An air purifier that filters dust and particles from indoor air. It helps freshen living spaces.

Prompting an LLM#

We will also need access to a large language model (LLM) to combine the code snippets and the question to retrieve an answer to our question that involves the code snippets.

def prompt_llm_server(message:str, model="meta-llama/Llama-3.3-70B-Instruct"):
    """A prompt helper function that sends a message to an LLM-server
    and returns only the text response.
    """
    import os
    import openai
    
    # convert message in the right format if necessary
    if isinstance(message, str):
        message = [{"role": "user", "content": message}]
        
    # setup connection to the LLM-server
    client = openai.OpenAI(base_url="https://llm.scads.ai/v1", 
                           api_key=os.environ.get('SCADSAI_API_KEY'))
    
    # submit prompt
    response = client.chat.completions.create(
        model=model,
        messages=message
    )
    
    # extract answer
    return response.choices[0].message.content

prompt_llm_server("Hello world")

"Hello! It's nice to meet you. Is there something I can help you with or would you like to chat?"

We can then assemble selected product descriptions and question to a prompt.

context = "\n\n".join(related_products)

prompt = f"""
Answer the question by the very end and consider given product descriptions. 
Choose at least one of the listed products to answer the question.

## Code snippets
{context}

## Question
{question}
"""

print(prompt)

Answer the question by the very end and consider given product descriptions. 
Choose at least one of the listed products to answer the question.

## Code snippets
* Brushwick - An electric toothbrush that helps clean teeth with powered bristles. It makes daily brushing easier.

* SugarBrush - A toothbrush made of pure sugar, for the extra sweet smile.

* Purepixie - An air purifier that filters dust and particles from indoor air. It helps freshen living spaces.

## Question
What product can I buy for brushing teeth?

Answering our question#

Eventually we can answer our question

answer = prompt_llm_server(prompt)

show(answer)

To brush your teeth, you can consider the following products from the list:

  1. Brushwick - An electric toothbrush that helps clean teeth with powered bristles, making daily brushing easier.

  2. SugarBrush - A toothbrush made of pure sugar, for the extra sweet smile.

Between these two options, Brushwick seems like a more practical and hygienic choice for brushing teeth, as it uses powered bristles for cleaning and is designed for daily brushing. SugarBrush, being made of pure sugar, might not be the best option for oral hygiene due to the potential for sugar to contribute to tooth decay.

Therefore, the product you can buy for brushing teeth is Brushwick.

Prompting without RAG#

In comparison, we send the same question together with minimal instructions to ChatGPT without out additional code-snippets.

answer = prompt_llm_server(f"""
Answer this question:
{question}
""")

show(answer)

There are several products you can buy for brushing teeth, including:

  1. Toothbrush: You can choose from manual or electric toothbrushes. Manual toothbrushes are inexpensive and easy to use, while electric toothbrushes can be more effective at removing plaque and improving gum health.

  2. Toothpaste: There are many types of toothpaste available, including:

    • Fluoride toothpaste: helps prevent tooth decay and strengthen tooth enamel.

    • Whitening toothpaste: contains mild abrasives to help remove surface stains and whiten teeth.

    • Sensitive toothpaste: designed for people with sensitive teeth and gums.

    • Natural toothpaste: made with natural ingredients and often free from harsh chemicals.

  3. Toothbrush heads (for electric toothbrushes): You can buy replacement heads for your electric toothbrush, which come in different sizes and types (e.g., standard, compact, or specialized for specific oral care needs).

  4. Interdental brushes: These small brushes are designed to clean between teeth and around dental work, such as bridges and implants.

  5. Dental floss: Used to remove food particles and plaque from between teeth and under the gumline.

  6. Mouthwash: An optional product that can help kill bacteria, freshen breath, and provide additional protection against tooth decay and gum disease.

When choosing a product, consider your oral care needs, preferences, and any recommendations from your dentist or hygienist.

Exercise#

Modify the question and ask for other household products. Compare the reponse by both approaches, with and without RAG.