大模型Rag - 向量数据库索引

一、安装与部署

向量数据库 Chroma是一个简单易用的 AI 原生向量数据库，安装十分简单，使用 pip install chromadb 即可完成安装。
安装完成后，可以使用如下命令启动服务端：

chromadb run --path ./data --host 127.0.0.1 --port 8000

其中：
--path 是数据库文件的本地存储路径。
--host 是对外提供服务的 IP。
--port 是服务端口号。

服务启动后，可以在指定路径查看本地持久化存储的数据。

二、连接数据库

客户端通过 chromadb.HttpClient 与 Chroma 服务端连接：

import chromadb
from chromadb.utils import embedding_functions
# 链接数据库
client = chromadb.HttpClient(host='localhost', port=8000)
# 指定对应的 embedding 模型
model_path = './gte-large-zh/'
embedding_fn = SentenceTransformerEmbeddingFunction(model_name=name_path)
collection = client.create_collection(
    name='rag_db',
    embedding_function=embedding_fn,
    metadata={"hnsw:space": "cosine"}  # 设置为余弦相似度和索引
)
# 创建向量集合

docs = ["第一句话。", "第二句话。", "第三句话。"]
ids = ["id1", "id2", "id3"]
metadatas = [
    {"chapter": 1, "verse": 5},
    {"chapter": 2, "verse": 5},
    {"chapter": 3, "verse": 5},
]

collection.add(documents=docs, ids=['1','2','3'], metadatas=metadatas)

# 检索
get_collection = chroma_client.get_collection(name='rag_db',embedding_function=embedding_fn)

id_result = get_collection.get(ids=['1'],include=["documents","embeddings","metadatas"])

# 向量检索
query = '问题'
get_collection.query(
	query_texts=query,# 向量检索
    n_results=2,
    include=["documents","embeddings","metadatas"],
    where={"verse":5})# 元数据检索