Integrate OceanBase Cloud vector retrieval with LlamaIndex

OceanBase Database supports vector data storage, vector indexes, and embedding vector retrieval in V4.3.3 and later. You can store vectorized data in OceanBase Database for further retrieval.

LlamaIndex is a framework for building context-augmented generative AI applications by using large language models (LLMs), including proxies and workflows. It provides a wealth of capabilities such as data connectors, data indexes, proxies, observability/assessment integration, and workflows.

This topic describes how to integrate the vector retrieval feature of OceanBase Cloud, Tongyi Qianwen (Qwen), and LlamaIndex for Document Question Answering (DQA).

Prerequisites

• A transactional cluster instance of the MySQL compatible mode is available in your environment.

• To use a cluster instance, you first need to create a tenant by referring to Create a tenant.

• You have created a MySQL-compatible tenant, a database, and an account, and granted the read and write permissions to the database account. For more information, see Create an account and Create a database (MySQL compatible mode).

• You have been granted the project admin or instance admin role to perform read and write operations on instances in the project. If you do not have the required permissions, contact the organization admin.

• You have installed Python 3.9 or later and pip. If Python installed on your server is of an early version, you can use Miniconda to build a new environment of Python 3.9 or later. For more information, see Installing Miniconda.

• You have installed LlamaIndex. The installation commands are as follows:

  python3 -m pip install llama-index-vector-stores-oceanbase llama-index
  python3 -m pip install llama-index-embeddings-dashscope
  python3 -m pip install llama-index-llms-dashscope
  

Step 1: Obtain the database connection information

1. Log in to the OceanBase Cloud console.

2. In the instance list page, expand the the information of the target instance.

3. Select Connect > Get Connection String under the target tenant.

4. In the pop-up window, select Public Network as the connection method.

5. Follow the prompts in the pop-up window to obtain the public endpoint and the connection string.

Step 2: Register an LLM account

Register an Alibaba Cloud Model Studio account, activate the model service, and obtain an API key.

Specify the API key for the relevant environment variable.

export DASHSCOPE_API_KEY="YOUR_DASHSCOPE_API_KEY"

Step 3: Build an AI assistant

Download sample data

mkdir -p '/root/llamaindex/paul_graham/'
wget 'https://raw.githubusercontent.com/run-llama/llama_index/main/docs/docs/examples/data/paul_graham/paul_graham_essay.txt' -O '/root/llamaindex/paul_graham/paul_graham_essay.txt'

Load the data

import os
from pyobvector import ObVecClient
from llama_index.core import Settings
from llama_index.embeddings.dashscope import DashScopeEmbedding
from llama_index.core import (
SimpleDirectoryReader,
load_index_from_storage,
VectorStoreIndex,
StorageContext,
)
from llama_index.vector_stores.oceanbase import OceanBaseVectorStore
from llama_index.llms.dashscope import DashScope, DashScopeGenerationModels
#set ob client
client = ObVecClient(uri="127.0.0.1:2881", user="root@test",password="",db_name="test")
# Global Settings
Settings.embed_model = DashScopeEmbedding()
# config llm model
dashscope_llm = DashScope(
    model_name=DashScopeGenerationModels.QWEN_MAX,
    api_key=os.environ.get("DASHSCOPE_API_KEY", ""),
)
# load documents
documents = SimpleDirectoryReader("/root/llamaindex/paul_graham/").load_data()
oceanbase = OceanBaseVectorStore(
    client=client,
    dim=1536,
    drop_old=True,
    normalize=True,
)

storage_context = StorageContext.from_defaults(vector_store=oceanbase)
index = VectorStoreIndex.from_documents(
    documents, storage_context=storage_context
)

Perform vector search

Here is an example of searching for "What did the author do growing up?" from paul_graham_essay.txt:

# set Logging to DEBUG for more detailed outputs
query_engine = index.as_query_engine(llm=dashscope_llm)
res = query_engine.query("What did the author do growing up?")
res.response

The expected output is as follows:

'Growing up, the author worked on writing and programming outside of school. In terms of writing, he wrote short stories, which he now considers to be awful, as they had very little plot and focused mainly on characters with strong feelings. For programming, he started in 9th grade by trying to write programs on an IBM 1401 at his school, using an early version of Fortran. Later, after getting a TRS-80 microcomputer, he began to write more practical programs, including simple games, a program to predict the flight height of model rockets, and a word processor that his father used for writing.'