Skip to content

向量数据库

向量数据库用于存储和检索高维向量数据,是 RAG 系统的核心组件。

概述

向量数据库是专门设计用于存储、索引和检索高维向量的数据库系统。

常用向量数据库

Chroma

Chroma 是一个轻量级开源向量数据库。

python
import chromadb
from chromadb.config import Settings

# 初始化客户端
client = chromadb.Client(Settings(
    chroma_db_impl="duckdb+parquet",
    persist_directory="./chroma_db"
))

# 创建集合
collection = client.create_collection(name="knowledge_base")

# 添加文档
collection.add(
    documents=["RAG 是检索增强生成技术", "向量数据库用于存储嵌入向量"],
    metadatas=[{"source": "doc1"}, {"source": "doc2"}],
    ids=["1", "2"]
)

# 查询
results = collection.query(
    query_texts=["什么是 RAG?"],
    n_results=2
)

print(results)

Pinecone

Pinecone 是托管式向量数据库服务。

python
import pinecone

# 初始化客户端
pinecone.init(api_key="your-api-key", environment="us-west1-gcp")

# 创建索引
pinecone.create_index("knowledge-base", dimension=384)

# 连接索引
index = pinecone.Index("knowledge-base")

# 插入向量
index.upsert([
    ("vec1", [0.1, 0.2, 0.3], {"source": "doc1"}),
    ("vec2", [0.4, 0.5, 0.6], {"source": "doc2"})
])

# 查询
results = index.query(
    vector=[0.1, 0.2, 0.3],
    top_k=5,
    include_metadata=True
)

print(results)

Weaviate

Weaviate 是开源向量数据库,支持 GraphQL 查询。

python
import weaviate

# 初始化客户端
client = weaviate.Client("http://localhost:8080")

# 创建类
client.schema.create_class({
    "class": "Document",
    "properties": [
        {"name": "content", "dataType": ["string"]},
        {"name": "source", "dataType": ["string"]}
    ]
})

# 添加对象
client.data_object.create({
    "content": "RAG 是检索增强生成技术",
    "source": "doc1"
}, "Document")

# 查询
result = client.query.get("Document", ["content", "source"]).do()
print(result)

索引类型

HNSW

Hierarchical Navigable Small Worlds 是一种高效的近似最近邻算法。

python
# 使用 HNSW 索引
collection = client.create_collection(
    name="knowledge_base",
    metadata={"hnsw:space": "cosine"}
)

IVF

Inverted File Index 适合大规模数据集。

python
# 使用 IVF 索引
collection = client.create_collection(
    name="knowledge_base",
    metadata={"ivf:nlist": 100}
)

FLAT

精确最近邻搜索,适合小规模数据集。

python
# 使用 FLAT 索引
collection = client.create_collection(
    name="knowledge_base",
    metadata={"hnsw:enabled": False}
)

相似度度量

余弦相似度

python
import numpy as np

def cosine_similarity(vec1, vec2):
    return np.dot(vec1, vec2) / (np.linalg.norm(vec1) * np.linalg.norm(vec2))

欧氏距离

python
def euclidean_distance(vec1, vec2):
    return np.linalg.norm(np.array(vec1) - np.array(vec2))

点积

python
def dot_product(vec1, vec2):
    return np.dot(vec1, vec2)

最佳实践

向量维度选择

python
# 根据任务选择合适的维度
# 小规模任务:384维(all-MiniLM-L6-v2)
# 大规模任务:768维(all-mpnet-base-v2)
# 高质量任务:1024维+(text-embedding-3-large)

批量插入

python
# 批量插入提高效率
batch_size = 1000

for i in range(0, len(documents), batch_size):
    batch = documents[i:i+batch_size]
    collection.add(
        documents=batch,
        ids=[str(j) for j in range(i, i+len(batch))]
    )

定期更新

python
# 定期更新向量数据库
def update_knowledge_base(new_documents):
    collection.add(
        documents=new_documents,
        ids=[str(time.time()) + str(i) for i in range(len(new_documents))]
    )