Integrate OceanBase vector with LlamaIndex

OceanBase Database provides vector storage, vector indexing, and embedding vector search capabilities. You can store vectorized data in OceanBase Database for subsequent retrieval and search.

LlamaIndex is a framework for building context-enhanced AI applications using LLMs (including agents and workflows). It provides data connectors, data indexing, agents, observability and evaluation integration, and workflows.

This topic describes how to integrate OceanBase Database's vector search, Qwen, and LlamaIndex to implement document-based question answering.

Prerequisites

• You have deployed OceanBase Database V4.4.0 or later and created a MySQL-compatible tenant. After you create a tenant, follow the steps below. For more information, see Create a tenant.

  • You have a MySQL-compatible tenant, database, and account, and granted the read and write privileges to the database account.
  • You have set the ob_vector_memory_limit_percentage parameter in the tenant to enable vector search. We recommend that you set the value to 30 for OceanBase Database versions earlier than V4.3.5 BP3, and to 0 for V4.3.5 BP3 and later. For more information about this parameter, see ob_vector_memory_limit_percentage.

• You have installed Python 3.9 or later.

• You have installed the required dependencies:

  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
  

• You have obtained the Qwen API key.

Step 1: Obtain the database connection information

Obtain the database connection string from the deployment personnel or administrator of OceanBase Database, for example:

obclient -h$host -P$port -u$user_name -p$password -D$database_name

Parameter description:

• $host: the IP address for connecting to OceanBase Database. For connection through OceanBase Database Proxy (ODP), use the IP address of an ODP. For direct connection, use the IP address of an OBServer node.

• $port: the port for connecting to OceanBase Database. For connection through ODP, the default value is 2883, which can be customized when ODP is deployed. For direct connection, the default value is 2881, which can be customized when OceanBase Database is deployed.

• $database_name: the name of the database to be accessed.

  Notice

  The user for connecting to the tenant must have the CREATE, INSERT, DROP, and SELECT privileges on the database. For more information about user privileges, see Privilege types in MySQL-compatible mode.

• $user_name: the tenant account. For connection through ODP, the format is username@tenant name#cluster name or cluster name:tenant name:username. For direct connection, the format is username@tenant name.

• $password: the password of the account.

For more information about the connection string, see Connect to an OceanBase tenant by using OBClient.

Step 2: Build your AI assistant

Set the Qwen API key environment variable

Obtain the Qwen API key and configure the API key to an environment variable.

export DASHSCOPE_API_KEY="YOUR_DASHSCOPE_API_KEY"

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 data text

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
)

Vector search

This step shows how to query "What did the author do growing up?" from the document 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

Expected result:

'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.'