类功能

功能描述

提供向量数据库。

函数原型

from mx_rag.storage.vectorstore import OpenGaussDB
OpenGaussDB(engine, collection_name, search_mode, index_type, metric_type)

参数说明

参数名	数据类型	可选/必选	说明
engine	Engine	必选	Engine实例。说明： Engine由用户控制传入，请使用安全的连接方式。
collection_name	str	可选	集合名称，不能为空，最大长度为1024，默认为“vectorstore”。
search_mode	SearchMode	可选	检索模式，当前支持三种模式，包括稠密检索（DENSE），稀疏检索(SPARSE)和混合检索（HYBRID）。类型介绍可参见SearchMode。
index_type	str	可选	向量检索类型，当前支持IVFFLAT，HNSW，默认为HNSW，该字段稠密检索和混合检索模式时，针对稠密向量有效。稀疏向量检索类型为HNSW，不支持配置。
metric_type	str	可选	向量距离计算方式，支持IP，L2，COSINE，默认为IP。

返回类型

数据类型	说明
OpenGaussDB	OpenGaussDB对象。

调用示例

import numpy as np
from mx_rag.storage.vectorstore import OpenGaussDB, SearchMode
from sqlalchemy import URL, create_engine

# OpenGauss
username = "demo"
password = "<password here>"
host = "<host here>"
port = "<port here>"
database = "testdb"

# vector config
dim = 128 
n_emb = 1000

url = URL.create(
   "opengauss+psycopg2",
   username=username,
   password=password,
   host=host,
   port=port,
   database=database
)
print(url)
# create an engine
engine = create_engine(url, pool_size=20, max_overflow=10, pool_pre_ping=True) 
# search mode defaults to DENSE
# similarity strategy defaults to FLAT_IP
dense_store = OpenGaussDB.create(
    engine=engine,
    dense_dim=dim
)

# add vectors
dense_embeddings = np.random.randn(n_emb, dim)
ids = list(range(n_emb))
dense_store.add(ids, dense_embeddings)

# search vectors
res = dense_store.search(dense_embeddings[:3].tolist(), k=3)
print(res)

# delete vectors
count = dense_store.delete(ids)
print(count)

# update vector
dense_store.update([1], dense_embeddings[:1])

# drop table
dense_store.drop_collection()

父主题： OpenGaussDB类