跳到主要内容

seekdb 与 Google Gemini 集成

seekdb 提供了向量类型存储、向量索引、embedding 向量搜索的能力。可以将向量化后的数据存储在 seekdb ，供下一步的搜索使用。

Google Gemini AI 是一系列由 Google 开发的多模态大型语言模型（LLMs）。它被设计为能够理解和处理多种类型的数据，包括文本、代码、图像、音频和视频等。

前提条件​

  • 您已经部署了 seekdb。

  • 您的环境中已存在可以使用的数据库和账号，并已对数据库账号授读写权限。

  • 安装 Python 3.11 及以上版本。

  • 安装依赖。

    python3 -m pip install cffi pyseekdb requests google-genai

步骤一：获取数据库连接信息​

联系 seekdb 部署人员或者管理员获取相应的数据库连接串，例如：

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

参数说明：

  • $host：提供 seekdb 连接 IP 地址。

  • $port：提供 seekdb 连接端口，默认是 2881。

  • $database_name：需要访问的数据库名称。

    提示

    连接的用户需要拥有该数据库的 CREATE、INSERT、DROP 和 SELECT 权限。

  • $user_name：提供数据库连接账户。

  • $password：提供账户密码。

步骤二：构建您的 AI 助手​

设置环境变量​

获取 Google Gemini 的 API 和 seekdb 连接信息并配置环境变量中。

export SEEKDB_DATABASE_URL=YOUR_SEEKDB_DATABASE_URL
export SEEKDB_DATABASE_USER=YOUR_SEEKDB_DATABASE_USER
export SEEKDB_DATABASE_DB_NAME=YOUR_SEEKDB_DATABASE_DB_NAME
export SEEKDB_DATABASE_PASSWORD=YOUR_SEEKDB_DATABASE_PASSWORD
export GEMINI_API_KEY=YOUR_GEMINI_API_KEY

示例代码片段​

加载数据​

这里以 text-embedding-004 为例，使用 Google Gemini AI API 生成向量数据：

from google import genai
from glob import glob
from tqdm import tqdm
import os
from google.genai import types
import requests,os,json,ollama,pyseekdb
from tqdm import tqdm
from pyseekdb import HNSWConfiguration

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 generate_embeddings(text):
    """Generate embeddings for a single text using Gemini API"""
    response = genai_client.models.embed_content(
        model="text-embedding-004",
        contents=[text],
        config=types.EmbedContentConfig(output_dimensionality=768)
    )
    return response.embeddings[0].values


ids = []
embeddings = []
documents_list = []

for i, text in enumerate(tqdm(documents, desc="Creating embeddings")):
    # Generate the embedding for the text via Gemini API
    embedding = generate_embeddings(text)
    ids.append(f"{i+1}")
    embeddings.append(embedding)
    documents_list.append(text)

print(f"Successfully processed {len(documents_list)} texts")

定义表并将数据存入 seekdb​

创建一个名为 gemini_seekdb_demo_documents 的表，并数据存入 seekdb：

SEEKDB_DATABASE_HOST = os.getenv('SEEKDB_DATABASE_HOST')
SEEKDB_DATABASE_PORT = int(os.getenv('SEEKDB_DATABASE_PORT', 2881)) 
SEEKDB_DATABASE_USER = os.getenv('SEEKDB_DATABASE_USER')
SEEKDB_DATABASE_DB_NAME = os.getenv('SEEKDB_DATABASE_DB_NAME')
SEEKDB_DATABASE_PASSWORD = os.getenv('SEEKDB_DATABASE_PASSWORD')

client = pyseekdb.Client(host=SEEKDB_DATABASE_HOST, port=SEEKDB_DATABASE_PORT, database=SEEKDB_DATABASE_DB_NAME, user=SEEKDB_DATABASE_USER, password=SEEKDB_DATABASE_PASSWORD)

table_name = "gemini_seekdb_demo_documents"
config = HNSWConfiguration(dimension=768, distance='l2')  
collection = client.create_collection(
    name=table_name,
    configuration=config,
    embedding_function=None
)

print('- Inserting Data to seekdb...')
collection.add(
    ids=ids,
    embeddings=embeddings,
    documents=documents
)

语义搜索​

通过 text-embedding-004 模型生成查询文本的向量，然后根据查询文本的向量与向量表中的每个向量的 l2 距离，搜索最相关的文档：

query = 'When was artificial intelligence founded?'
quest_embed = genai_client.models.embed_content(model="text-embedding-004", contents=query)
res = collection.query(
    query_embeddings=quest_embed.embeddings[0].values,
    n_results=1
)

print('- The Most Relevant Document and Its Distance to the Query:')
for i, (doc_id, document, distance) in enumerate(zip(
    res['ids'][0], 
    res['documents'][0], 
    res['distances'][0]
)):
    print(f'  - ID: {doc_id}')
    print(f'    content: {document}')
    print(f'    distance: {distance:.6f}')

预期结果​

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

Contents

编辑此页