Integrate OceanBase vector search with Google Gemini

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

export OCEANBASE_DATABASE_URL=YOUR_OCEANBASE_DATABASE_URL
export OCEANBASE_DATABASE_USER=YOUR_OCEANBASE_DATABASE_USER
export OCEANBASE_DATABASE_DB_NAME=YOUR_OCEANBASE_DATABASE_DB_NAME
export OCEANBASE_DATABASE_PASSWORD=YOUR_OCEANBASE_DATABASE_PASSWORD
export GEMINI_API_KEY=YOUR_GEMINI_API_KEY

from google import genai
from glob import glob
from tqdm import tqdm
import os
from google.genai import types
from sqlalchemy import Column, Integer, String
from pyobvector import ObVecClient, VECTOR, IndexParam, l2_distance

documents = [
"Artificial intelligence was founded as an academic discipline in 1956.",
"Alan Turing was the first person to conduct substantial research in AI.",
"Born in Maida Vale, London, Turing was raised in southern England.",
]

genai_client=genai.Client(api_key=os.environ["GEMINI_API_KEY"])
def embed_documents_in_batches(documents, batch_size=100):
    all_embeddings = []
    for i in range(0, len(documents), batch_size):
        batch = documents[i:i + batch_size]
        response = genai_client.models.embed_content(
            model="text-embedding-004",
            contents=batch,
            config=types.EmbedContentConfig(output_dimensionality=768)
        )
        all_embeddings.extend(response.embeddings)
            
    return all_embeddings


embeddings = embed_documents_in_batches(documents)

data = []
for i, line in enumerate(tqdm(documents, desc="Creating embeddings")):
    data.append({"text": line ,"embedding": embeddings[i].values})

OCEANBASE_DATABASE_URL = os.getenv('OCEANBASE_DATABASE_URL')
OCEANBASE_DATABASE_USER = os.getenv('OCEANBASE_DATABASE_USER')
OCEANBASE_DATABASE_DB_NAME = os.getenv('OCEANBASE_DATABASE_DB_NAME')
OCEANBASE_DATABASE_PASSWORD = os.getenv('OCEANBASE_DATABASE_PASSWORD')

client = ObVecClient(uri=OCEANBASE_DATABASE_URL, user=OCEANBASE_DATABASE_USER,password=OCEANBASE_DATABASE_PASSWORD,db_name=OCEANBASE_DATABASE_DB_NAME)

table_name = "gemini_oceanbase_demo_documents"
client.drop_table_if_exist(table_name)

cols = [
    Column("id", Integer, primary_key=True, autoincrement=True),
    Column("text", String(10000), nullable=False),
    Column("embedding", VECTOR(768))
]

# Create vector index
vector_index_params = IndexParam(
    index_name="idx_question_embedding",
    field_name="embedding",  
    index_type="HNSW",
    distance_metric="l2"
)

client.create_table_with_index_params(
    table_name=table_name,
    columns=cols, 
    vidxs=[vector_index_params]
)
print('- Inserting Data to OceanBase...')
client.insert(table_name, data=data)

question = "When was artificial intelligence founded?"

quest_embed = genai_client.models.embed_content(model="text-embedding-004", contents=question)
search_res = client.ann_search(
    table_name,
    vec_data=quest_embed.embeddings[0].values,
    vec_column_name="embedding",
    distance_func=l2_distance,
    with_dist=True,
    topk=1, 
    output_column_names=["id","text"], 
)

print('- The Most Relevant Document and Its Distance to the Query:')
for row in search_res.fetchall():
    print(f'  - ID: {row[0]}\n'
          f'    content: {row[1]}\n'
          f'    distance: {row[2]}')

  - ID: 1
    content: Artificial intelligence was founded as an academic discipline in 1956.
    distance: 0.6019276093082409