Retrieving Relevant Chunks and Building LLM Prompts in JavaScript

Introduction

Welcome to the third lesson in our journey through Scaling Up RAG with Vector Databases! Well done, you're halfway through this course. In the previous lesson, you learned how to split or chunk your text data and store those chunks in a vector database collection. Now, we'll delve into retrieving the most relevant chunks for any given query and building an LLM prompt to produce more accurate, context-driven answers.

Retrieving the Most Relevant Chunks

Join the 1M+ learners on CodeSignal

Be a part of our community of 1M+ users who develop and demonstrate their skills on CodeSignal

JavaScript

async function retrieveTopChunks(query, collection, topK = 2) {
    /**
     * Retrieves the topK chunks relevant to the given query from 'collection'.
     * Returns an array of retrieved chunks, each containing 'chunk' text,
     * 'docId', and 'distance'.
     */
    // Search for topK results matching the user's query
    const results = await collection.query({
        queryTexts: [query],
        nResults: topK
    });

    const retrievedChunks = [];

    // Safeguard in case no results are found
    if (!results.documents || !results.documents[0]) {
        return retrievedChunks;
    }

    // Gather each retrieved chunk, along with its distance score
    for (let i = 0; i < results.documents[0].length; i++) {
        retrievedChunks.push({
            chunk: results.documents[0][i],
            docId: results.ids[0][i],
            distance: results.distances[0][i]
        });
    }
    return retrievedChunks;
}

JavaScript

function buildPrompt(query, retrievedChunks) {
    /**
     * Constructs a prompt by restating the 'query' and adding the retrieved chunks
     * as an inline context for the LLM.
     */
    let prompt = `Question: ${query}\nAnswer using only the following context:\n`;
    retrievedChunks.forEach(rc => {
        prompt += `- ${rc.chunk}\n`;
    });
    prompt += "Answer:";
    return prompt;
}

function buildPrompt(query, retrievedChunks) {
    /**
     * Constructs a prompt by restating the 'query' and adding the retrieved chunks
     * as an inline context for the LLM.
     */
    let prompt = `Question: ${query}\nAnswer using only the following context:\n`;
    retrievedChunks.forEach(rc => {
        prompt += `- ${rc.chunk}\n`;
    });
    prompt += "Answer:";
    return prompt;
}

JavaScript

import fs from 'fs';
import path from 'path';
import { ChromaClient, OpenAIEmbeddingFunction } from 'chromadb';

// Load corpus data from JSON file
const corpusData = JSON.parse(fs.readFileSync('data/corpus.json', 'utf8'));

const embedder = new OpenAIEmbeddingFunction({
    model_name: "text-embedding-ada-002"
});

const chroma = new ChromaClient({ path: "http://localhost:8000" });
const collections = await chroma.listCollections();
let collection;
const collectionName = "rag_collection";

if (collections.find((coll) => coll === collectionName)) {
    console.log(`Collection "${collectionName}" already exists. Retrieving it...`);
    collection = await chroma.getCollection({ name: collectionName });
    if (!collection.embeddingFunction) {
        collection.embeddingFunction = embedder;
    }
} else {
    console.log(`Collection "${collectionName}" does not exist. Creating collection...`);
    collection = await chroma.createCollection({
        name: collectionName,
        embeddingFunction: embedder,
    });
}

// Batch add documents from the corpus data
const documents = corpusData.map(doc => doc.content);
const ids = corpusData.map(doc => `chunk_${doc.id}_0`);
collection.add({ documents, ids });

import fs from 'fs';
import path from 'path';
import { ChromaClient, OpenAIEmbeddingFunction } from 'chromadb';

// Load corpus data from JSON file
const corpusData = JSON.parse(fs.readFileSync('data/corpus.json', 'utf8'));

const embedder = new OpenAIEmbeddingFunction({
    model_name: "text-embedding-ada-002"
});

const chroma = new ChromaClient({ path: "http://localhost:8000" });
const collections = await chroma.listCollections();
let collection;
const collectionName = "rag_collection";

if (collections.find((coll) => coll === collectionName)) {
    console.log(`Collection "${collectionName}" already exists. Retrieving it...`);
    collection = await chroma.getCollection({ name: collectionName });
    if (!collection.embeddingFunction) {
        collection.embeddingFunction = embedder;
    }
} else {
    console.log(`Collection "${collectionName}" does not exist. Creating collection...`);
    collection = await chroma.createCollection({
        name: collectionName,
        embeddingFunction: embedder,
    });
}

// Batch add documents from the corpus data
const documents = corpusData.map(doc => doc.content);
const ids = corpusData.map(doc => `chunk_${doc.id}_0`);
collection.add({ documents, ids });

JavaScript

async function main() {
    const query = "What are some recent technological breakthroughs?";
    const retrievedDocs = await retrieveTopChunks(query, collection, 5);
    const finalPrompt = buildPrompt(query, retrievedDocs);
    const answer = await getLLMResponse(finalPrompt);

    console.log("Prompt:\n");
    console.log(finalPrompt);
    console.log("\nLLM Answer:", answer);
}

main();

async function main() {
    const query = "What are some recent technological breakthroughs?";
    const retrievedDocs = await retrieveTopChunks(query, collection, 5);
    const finalPrompt = buildPrompt(query, retrievedDocs);
    const answer = await getLLMResponse(finalPrompt);

    console.log("Prompt:\n");
    console.log(finalPrompt);
    console.log("\nLLM Answer:", answer);
}

main();

Prompt:

Question: What are some recent technological breakthroughs?
Answer using only the following context:
- The Industrial Revolution brought significant technological and social changes. It reshaped economies and altered the fabric of society. Scholars examine its impact on labor, innovation, and modern industrial practices.
- Breakthroughs in renewable energy technologies are reducing global dependence on fossil fuels. Solar and wind systems are becoming more efficient and affordable. These innovations are crucial to combating climate change and ensuring a sustainable future.
- The digital revolution is transforming how we approach health and wellness. Technological innovations, from fitness trackers to health apps, are empowering individuals to manage their well-being. This integration of technology and lifestyle is reshaping daily habits for a healthier future.
- Advances in medical technology are revolutionizing patient care through new diagnostic and treatment methods. Breakthroughs in imaging and robotics are enhancing the precision of medical procedures. Healthcare professionals are optimistic about the potential for improved outcomes.
- Scientists are developing renewable materials that could replace traditional plastics. Innovations in biopolymers are leading to sustainable manufacturing practices. These breakthroughs promise to reduce environmental waste and support a circular economy.
Answer:

LLM Answer: Recent technological breakthroughs include advancements in renewable energy technologies, which are making solar and wind systems more efficient and affordable, thereby reducing global dependence on fossil fuels and aiding in the fight against climate change. Additionally, the digital revolution is enhancing health and wellness through innovations like fitness trackers and health apps, empowering individuals to better manage their well-being. In the medical field, new diagnostic and treatment methods, along with improvements in imaging and robotics, are revolutionizing patient care and enhancing the precision of medical procedures. Furthermore, scientists are developing renewable materials, such as biopolymers, to replace traditional plastics, promoting sustainable manufacturing practices and supporting a circular economy.

Prompt:

Question: What are some recent technological breakthroughs?
Answer using only the following context:
- The Industrial Revolution brought significant technological and social changes. It reshaped economies and altered the fabric of society. Scholars examine its impact on labor, innovation, and modern industrial practices.
- Breakthroughs in renewable energy technologies are reducing global dependence on fossil fuels. Solar and wind systems are becoming more efficient and affordable. These innovations are crucial to combating climate change and ensuring a sustainable future.
- The digital revolution is transforming how we approach health and wellness. Technological innovations, from fitness trackers to health apps, are empowering individuals to manage their well-being. This integration of technology and lifestyle is reshaping daily habits for a healthier future.
- Advances in medical technology are revolutionizing patient care through new diagnostic and treatment methods. Breakthroughs in imaging and robotics are enhancing the precision of medical procedures. Healthcare professionals are optimistic about the potential for improved outcomes.
- Scientists are developing renewable materials that could replace traditional plastics. Innovations in biopolymers are leading to sustainable manufacturing practices. These breakthroughs promise to reduce environmental waste and support a circular economy.
Answer:

LLM Answer: Recent technological breakthroughs include advancements in renewable energy technologies, which are making solar and wind systems more efficient and affordable, thereby reducing global dependence on fossil fuels and aiding in the fight against climate change. Additionally, the digital revolution is enhancing health and wellness through innovations like fitness trackers and health apps, empowering individuals to better manage their well-being. In the medical field, new diagnostic and treatment methods, along with improvements in imaging and robotics, are revolutionizing patient care and enhancing the precision of medical procedures. Furthermore, scientists are developing renewable materials, such as biopolymers, to replace traditional plastics, promoting sustainable manufacturing practices and supporting a circular economy.

JavaScript

function buildPrompt(query, retrievedChunks) {
    /**
     * Constructs a prompt by restating the 'query' and adding the retrieved chunks
     * as an inline context for the LLM.
     */
    let prompt = `Question: ${query}\nAnswer using only the following context:\n`;
    retrievedChunks.forEach(rc => {
        prompt += `- ${rc.chunk}\n`;
    });
    prompt += "Answer:";
    return prompt;
}

JavaScript

import fs from 'fs';
import path from 'path';
import { ChromaClient, OpenAIEmbeddingFunction } from 'chromadb';

// Load corpus data from JSON file
const corpusData = JSON.parse(fs.readFileSync('data/corpus.json', 'utf8'));

const embedder = new OpenAIEmbeddingFunction({
    model_name: "text-embedding-ada-002"
});

const chroma = new ChromaClient({ path: "http://localhost:8000" });
const collections = await chroma.listCollections();
let collection;
const collectionName = "rag_collection";

if (collections.find((coll) => coll === collectionName)) {
    console.log(`Collection "${collectionName}" already exists. Retrieving it...`);
    collection = await chroma.getCollection({ name: collectionName });
    if (!collection.embeddingFunction) {
        collection.embeddingFunction = embedder;
    }
} else {
    console.log(`Collection "${collectionName}" does not exist. Creating collection...`);
    collection = await chroma.createCollection({
        name: collectionName,
        embeddingFunction: embedder,
    });
}

// Batch add documents from the corpus data
const documents = corpusData.map(doc => doc.content);
const ids = corpusData.map(doc => `chunk_${doc.id}_0`);
collection.add({ documents, ids });

JavaScript

async function main() {
    const query = "What are some recent technological breakthroughs?";
    const retrievedDocs = await retrieveTopChunks(query, collection, 5);
    const finalPrompt = buildPrompt(query, retrievedDocs);
    const answer = await getLLMResponse(finalPrompt);

    console.log("Prompt:\n");
    console.log(finalPrompt);
    console.log("\nLLM Answer:", answer);
}

main();

Prompt:

Question: What are some recent technological breakthroughs?
Answer using only the following context:
- The Industrial Revolution brought significant technological and social changes. It reshaped economies and altered the fabric of society. Scholars examine its impact on labor, innovation, and modern industrial practices.
- Breakthroughs in renewable energy technologies are reducing global dependence on fossil fuels. Solar and wind systems are becoming more efficient and affordable. These innovations are crucial to combating climate change and ensuring a sustainable future.
- The digital revolution is transforming how we approach health and wellness. Technological innovations, from fitness trackers to health apps, are empowering individuals to manage their well-being. This integration of technology and lifestyle is reshaping daily habits for a healthier future.
- Advances in medical technology are revolutionizing patient care through new diagnostic and treatment methods. Breakthroughs in imaging and robotics are enhancing the precision of medical procedures. Healthcare professionals are optimistic about the potential for improved outcomes.
- Scientists are developing renewable materials that could replace traditional plastics. Innovations in biopolymers are leading to sustainable manufacturing practices. These breakthroughs promise to reduce environmental waste and support a circular economy.
Answer:

LLM Answer: Recent technological breakthroughs include advancements in renewable energy technologies, which are making solar and wind systems more efficient and affordable, thereby reducing global dependence on fossil fuels and aiding in the fight against climate change. Additionally, the digital revolution is enhancing health and wellness through innovations like fitness trackers and health apps, empowering individuals to better manage their well-being. In the medical field, new diagnostic and treatment methods, along with improvements in imaging and robotics, are revolutionizing patient care and enhancing the precision of medical procedures. Furthermore, scientists are developing renewable materials, such as biopolymers, to replace traditional plastics, promoting sustainable manufacturing practices and supporting a circular economy.

Introduction

Retrieving the Most Relevant Chunks

Join the 1M+ learners on CodeSignal

Building a Prompt for the LLM

Integrating the Corpus and Creating the Collection

Querying the Database and Generating Answers

Examining the Output

Conclusion and Next Steps