应用实践 | 探索如何将 OceanBase 作为向量数据库训练问答机器人

npm install dotenv express lodash mysql2 openai sequelize --save

CREATETABLE article_vectors ( id INT AUTO_INCREMENT, content_vector JSON, content TEXT CHARACTERSET utf8mb4 COLLATE utf8mb4_0900_ai_ci,PRIMARY KEY (id));

const { DataTypes, Model } = require('sequelize');class ArticleVector extends Model {}module.exports = (sequelize) => {    ArticleVector.init(        {            id: {                type: DataTypes.INTEGER,                autoIncrement: true,                primaryKey: true,            },            content_vector: {                type: DataTypes.JSON,            },            content: {                type: DataTypes.TEXT,            },        },        {            sequelize,            modelName: 'ArticleVector',            tableName: 'article_vectors',            timestamps: false,        },        {            sequelize,            modelName: 'ArticleVector',            tableName: 'content',            timestamps: false,        }    );    return ArticleVector;};

async function fetchRepoContents(    repo,    path = '',    branch = 'main',    limit = 100) {    const url = `https://api.github.com/repos/${repo}/contents/${path}?ref=${branch}`;    const accessToken = process.env.GITHUB_TOKEN;    const response = await fetch(url, {        headers: {            Authorization: `Bearer ${accessToken}`,            Accept: 'application/vnd.github+json',            'X-GitHub-Api-Version': '2022-11-28',        },    });    const data = await response.json();// If the path is a directory, recursively fetch the contents    if (Array.isArray(data)) {        let count = 0;        const files = [];        for (const item of data) {            if (count >= limit) {                break;            }            if (item.type === 'dir') {                const dirFiles = await fetchRepoContents(                    repo,                    item.path,                    branch,                    limit - count                );                files.push(...dirFiles);                count += dirFiles.length;            } else if (item.type === 'file' && item.name.endsWith('.md')) {                files.push(item);                count++;            }        }        return files;    }    return [];}/** * @param {RequestInfo | URL} url */async function fetchFileContent(url) {    const response = await fetch(url);    const content = await response.text();    return content;}

async function articleEmbedding(repo, path = '', branch, limit) {    const contents = await fetchRepoContents(repo, path, branch, limit);    console.log(contents.length);    contents.forEach(async (item) => {        const content = await fetchFileContent(item.download_url);        await storeEmbedding(content);    });}

const { Sequelize } = require('sequelize');const dotenv = require('dotenv');dotenv.config();const env = process.env.NODE_ENV || 'development';const config = require(__dirname + '/../config/config.json')[env];const { Configuration, OpenAIApi } = require('openai');const configuration = new Configuration({    apiKey: process.env.OPENAI_KEY,});const openai = new OpenAIApi(configuration);

async function embedArticle(article) {    // Create article embeddings using OpenAI    try {        const result = await openai.createEmbedding({            model: 'text-embedding-ada-002',            input: article,        });        // get the embedding        const embedding = result.data.data[0].embedding;        return embedding;    } catch (error) {        if (error.response) {            console.log(error.response.status);            console.log(error.response.data);        } else {            console.log(error.message);        }    }}

let sequelize;if (config.use_env_variable) {    sequelize = new Sequelize(process.env[config.use_env_variable], config);} else {    sequelize = new Sequelize(        config.database,        config.username,        config.password,        config    );}const ArticleVectorModel = require('../models/ArticleVector');const ArticleVector = ArticleVectorModel(sequelize);async function storeEmbedding(article) {    try {        await sequelize.authenticate();        console.log('Connection has been established successfully.');        const articleVector = await ArticleVector.create({            content_vector: await embedArticle(article),            content: article,        });        console.log('New article vector created:', articleVector.id);    } catch (err) {        console.error('Error:', err);    }}

GITHUB_TOKEN="Get your Github token from https://docs.github.com/en/rest/guides/getting-started-with-the-rest-api?apiVersion=2022-11-28"OPENAI_KEY ="YOUR_OPENAI_KEY"DOC_OWNER="OceanBase, you can change this to another product's name so that the chatbot knows which product it is dealing with"MODEL="gpt-4 and gpt-3.5-turbo are supported"

async function getSimilarDoc(question) {    let sequelize;    if (config.use_env_variable) {        sequelize = new Sequelize(process.env[config.use_env_variable], config);    } else {        sequelize = new Sequelize(            config.database,            config.username,            config.password,            config        );    }    const ArticleVector = ArticleVectorModel(sequelize);    // get all rows in article_vector table    const vectors = await ArticleVector.findAll();    sequelize.close();    // console.log(vectors[0]);// Calculate cosine similarity for each vector    const embeddedQuestion = await embedArticle(question);    // Calculate cosine similarity for each vector    // Calculate cosine similarity for each vector    const similarities = vectors.map((vector) => {        const similarity = cosineSimilarity(            embeddedQuestion,            vector.content_vector        );        return {            id: vector.id,            content: vector.content,            similarity: similarity,        };    });    const mostSimilarDoc = _.orderBy(similarities, ['similarity'], ['desc'])[0];    // console.log(mostSimilarDoc);    return mostSimilarDoc;}

function cosineSimilarity(vecA, vecB) {    const dotProduct = _.sum(_.zipWith(vecA, vecB, (a, b) => a * b));    const magnitudeA = Math.sqrt(_.sum(_.map(vecA, (a) => Math.pow(a, 2))));    const magnitudeB = Math.sqrt(_.sum(_.map(vecB, (b) => Math.pow(b, 2))));if (magnitudeA === 0 || magnitudeB === 0) {        return 0;    }    return dotProduct / (magnitudeA * magnitudeB);}

async function askAI(question) {    const mostSimilarDoc = await getSimilarDoc(question);    const context = mostSimilarDoc.content;    try {        const result = await openai.createChatCompletion({            model: process.env.MODEL,            messages: [                {                    role: 'system',                    content: `You are ${                        process.env.DOC_OWNER || 'a'                    } documentation assistant. You will be provided a reference docs article and a question, please answer the question based on the article provided. Please don't make up anything.`,                },                {                    role: 'user',                    content: `Here is an article related to the question: \n ${context} \n Answer this question: \n ${question}`,                },            ],        });        // get the answer        const answer = result.data.choices[0].message.content;        console.log(answer);        return {            answer: answer,            docId: mostSimilarDoc.id,        };    } catch (error) {        if (error.response) {            console.log(error.response.status);            console.log(error.response.data);        } else {            console.log(error.message);        }    }}

const express = require('express');const app = express();const db = require('./models');app.use(express.json());

db.sequelize.sync().then((req) => {    app.listen(3000, () => {        console.log('Server running at port 3000...');    });});

app.post('/train', async (req, res, next) => {    try {        await articleEmbedding(            req.body.repo,            req.body.path,            req.body.branch,            req.body.limit        );        res.json({            status: 'Success',            message: 'Docs trained successfully.',        });    } catch (e) {        next(e); // Pass the error to the error handling middleware    }});

{  "repo":"oceanbase/oceanbase-doc",  "path":"en-US",  "branch": "V4.1.0",  "limit": 1000}

app.post('/ask', async (req, res, next) => {    try {        const answer = await askAI(req.body.question);        res.json({            status: 'Success',            answer: answer.answer,            docId: answer.docId,        });    } catch (e) {        next(e); // Pass the error to the error handling middleware    }});