Part 2: Enhancing the Motoku LLM Retrieval System with OpenAI Embeddings and Prompt-based Retrieval

In Part 1, we established a foundational Motoko LLM retrieval system using Internet Computer Protocol (ICP), Motoko, and Node.js. Now, we will enhance this system by integrating OpenAI embeddings and adding a feature to retrieve embeddings based on a prompt. This involves utilizing OpenAI's powerful embeddings API to generate embeddings for text data before storing it in the Motoko canister and enabling prompt-based retrieval.

Prerequisites

Before proceeding, ensure you have completed the setup from Part 1 and have the following additional tools and knowledge:

• An OpenAI API key.
• Familiarity with OpenAI's embeddings API.
• Updated Node.js environment with axios installed for making HTTP requests.

Step 4: Integrating OpenAI Embeddings and Prompt-based Retrieval

4.1 Install axios

First, install axios to enable HTTP requests to the OpenAI API:

npm install axios

4.2 Update Environment Configuration

Add your OpenAI API key to the .env file:

OPENAI_API_KEY=<your-openai-api-key>

4.3 Modify the Node.js Server Script

Update server.js to integrate with the OpenAI embeddings API and handle prompt-based retrieval. Replace the existing content with the following:

const express = require('express');
const bodyParser = require('body-parser');
const axios = require('axios');
const { HttpAgent, Actor } = require('@dfinity/agent');
const { idlFactory } = require('./idl/embedding_store.did.js');
require('dotenv').config();

const app = express();
const port = 3000;

app.use(bodyParser.json());

const canisterId = process.env.CANISTER_ID;
const host = process.env.HOST;
const openaiApiKey = process.env.OPENAI_API_KEY;

// Initialize the agent
const agent = new HttpAgent({ host });
agent.fetchRootKey(); // Necessary for local development

// Create an actor instance
const embeddingStore = Actor.createActor(idlFactory, {
    agent,
    canisterId,
});

// Helper function to convert BigInt to a string for JSON serialization
const serializeBigInt = (obj) => {
    if (typeof obj === 'bigint') {
        return obj.toString();
    } else if (Array.isArray(obj)) {
        return obj.map(serializeBigInt);
    } else if (typeof obj === 'object' && obj !== null) {
        return Object.fromEntries(
            Object.entries(obj).map(([k, v]) => [k, serializeBigInt(v)])
        );
    }
    return obj;
};

const getOpenAIEmbedding = async (text) => {
    try {
        const response = await axios.post(
            'https://api.openai.com/v1/embeddings',
            { input: text, model: 'text-embedding-ada-002' },
            {
                headers: {
                    'Authorization': `Bearer ${openaiApiKey}`,
                    'Content-Type': 'application/json'
                }
            }
        );
        return response.data.data[0].embedding;
    } catch (error) {
        console.error('Error fetching embedding from OpenAI:', error);
        throw new Error('Failed to fetch embedding from OpenAI');
    }
};

const cosineSimilarity = (vecA, vecB) => {
    const dotProduct = vecA.reduce((sum, a, idx) => sum + a * vecB[idx], 0);
    const normA = Math.sqrt(vecA.reduce((sum, a) => sum + a * a, 0));
    const normB = Math.sqrt(vecB.reduce((sum, b) => sum + b * b, 0));
    return dotProduct / (normA * normB);
};

app.post('/storeEmbedding', async (req, res) => {
    const { key, text } = req.body;
    try {
        if (key !== process.env.SECRET_KEY) {
            throw new Error('Invalid key');
        }
        // Get embedding from OpenAI
        const embedding = await getOpenAIEmbedding(text);
        await embeddingStore.storeEmbedding(key, text, embedding);
        res.status(200).send('Embedding stored successfully.');
    } catch (error) {
        res.status(500).send(`Error: ${error.message}`);
    }
});

app.post('/retrieveEmbeddings', async (req, res) => {
    const { prompt } = req.body;
    try {
        // Get prompt embedding from OpenAI
        const promptEmbedding = await getOpenAIEmbedding(prompt);

        // Retrieve stored embeddings
        const embeddings = await embeddingStore.getEmbeddings();

        // Calculate similarities
        const results = embeddings.map((embedding) => ({
            text: embedding.text,
            similarity: cosineSimilarity(promptEmbedding, embedding.embedding),
        }));

        // Sort by similarity
        results.sort((a, b) => b.similarity - a.similarity);

        res.status(200).json(serializeBigInt(results));
    } catch (error) {
        res.status(500).send(`Error: ${error.message}`);
    }
});

app.listen(port, () => {
    console.log(`Server is running on http://localhost:${port}`);
});

4.4 Test the Updated System

4.4.1 Storing an Embedding

Store an embedding using the updated endpoint:

curl -X POST http://localhost:3000/storeEmbedding \
     -H "Content-Type: application/json" \
     -d '{"key": "8529741360", "text": "example text"}'

This will generate an embedding for the provided text using OpenAI and store it in the Motoko canister.

4.4.2 Retrieving Embeddings Based on a Prompt

Retrieve embeddings based on a prompt:

curl -X POST http://localhost:3000/retrieveEmbeddings \
     -H "Content-Type: application/json" \
     -d '{"prompt": "example prompt"}'

This will generate an embedding for the prompt using OpenAI, calculate the cosine similarity between the prompt embedding and stored embeddings, and return the results sorted by similarity.

Conclusion

In this enhanced version, we integrated OpenAI embeddings to improve the retrieval system's performance and added prompt-based retrieval functionality. The Node.js server now interacts with the OpenAI API to generate embeddings for the input text before storing them in the Motoko canister, and it can retrieve and rank stored embeddings based on the similarity to a given prompt. This setup leverages the decentralized capabilities of ICP and the advanced natural language processing capabilities of OpenAI, providing a robust solution for LLM retrieval.

As a next step, you could explore implementing more advanced search and filtering capabilities, integrating user authentication, or developing a frontend interface for your system.