Python与向量数据库

什么是向量数据库

向量数据库是专门用来存储和检索向量（一组浮点数）的数据库。在大语言模型和 AI 应用中，文本、图片等数据会被转换成高维向量（比如 1536 维的浮点数组），向量数据库能快速找到与某个向量最相似的其他向量。

这种”相似性搜索”是 RAG（检索增强生成）技术的核心。当你向 AI 提问时，系统先把你的问题转成向量，在向量数据库中找到最相关的文档片段，再把这些片段提供给大语言模型来生成回答。

基础概念

向量嵌入（Embedding）

向量嵌入是把非结构化数据（文本、图片、音频）转换成数值向量的过程。例如，一句话”今天天气很好”经过嵌入模型处理后，可能变成一个 1536 维的浮点数组。语义相近的文本，其向量在空间中的距离也相近。

相似度计算

向量数据库通过计算向量之间的距离来判断相似度。常用的距离度量有：

• 余弦相似度：计算两个向量夹角的余弦值，范围 -1 到 1，值越大越相似
• 欧氏距离：两个向量在空间中的直线距离，值越小越相似
• 内积（点积）：两个向量对应位置相乘再求和

RAG（检索增强生成）

RAG 的工作流程：用户提问 -> 问题转成向量 -> 在向量数据库中检索相关文档 -> 将文档和问题一起交给大语言模型 -> 生成回答。这样大语言模型就能基于你自己的数据来回答问题，而不仅限于训练数据。

常见的向量数据库

• ChromaDB：轻量级，适合本地开发和原型验证
• Pinecone：全托管云服务，适合生产环境
• Milvus：高性能分布式向量数据库，适合大规模数据
• Qdrant：用 Rust 编写的高性能向量搜索引擎
• Weaviate：支持混合搜索（向量+关键词）

快速上手

安装 ChromaDB

# 安装 ChromaDB
pip install chromadb

# 安装 OpenAI SDK（用于生成嵌入向量）
pip install openai

最简单的向量存储和检索

import chromadb

# 创建内存模式的客户端（数据不持久化）
client = chromadb.Client()

# 创建或获取一个集合（类似数据库中的表）
collection = client.get_or_create_collection("my_docs")

# 添加文档（ChromaDB 会自动使用默认嵌入模型）
collection.add(
    documents=["Python 是一种编程语言", "Java 也是一种编程语言", "今天天气很好"],
    ids=["doc1", "doc2", "doc3"]
)

# 查询与"编程"最相关的文档
results = collection.query(
    query_texts=["编程语言有哪些"],
    n_results=2  # 返回最相似的 2 个结果
)

print(results["documents"])
# 输出: [['Python 是一种编程语言', 'Java 也是一种编程语言']]

使用持久化存储

上面的例子数据只存在于内存中，程序关闭后就丢失了。使用持久化客户端可以将数据保存到磁盘：

import chromadb

# 创建持久化客户端（数据保存到指定目录）
client = chromadb.PersistentClient(path="./chroma_data")

# 创建集合
collection = client.get_or_create_collection("my_docs")

# 添加文档
collection.add(
    documents=["第一篇文档的内容", "第二篇文档的内容"],
    ids=["1", "2"]
)

# 下次启动程序时，数据依然存在

详细用法

手动指定嵌入向量

如果你有自己的嵌入模型，可以手动计算向量并存入 ChromaDB：

import chromadb

client = chromadb.Client()
collection = client.get_or_create_collection("custom_embeddings")

# 手动提供嵌入向量（这里用模拟数据，实际应使用嵌入模型生成）
# 假设每个文档对应一个 3 维向量
collection.add(
    documents=["苹果是一种水果", "香蕉也是一种水果", "汽车是交通工具"],
    embeddings=[
        [0.9, 0.1, 0.0],  # 苹果的向量
        [0.85, 0.15, 0.0],  # 香蕉的向量（和苹果相近）
        [0.0, 0.1, 0.9],  # 汽车的向量（和水果相远）
    ],
    ids=["1", "2", "3"]
)

# 用向量进行查询
results = collection.query(
    query_embeddings=[[0.9, 0.1, 0.0]],  # 查询和苹果相近的文档
    n_results=2
)

print(results["documents"])
# 输出: [['苹果是一种水果', '香蕉也是一种水果']]

添加元数据

元数据是附加在文档上的额外信息，可以用来过滤搜索结果：

import chromadb

client = chromadb.Client()
collection = client.get_or_create_collection("articles")

# 添加带元数据的文档
collection.add(
    documents=["Python 入门教程", "Java 高级编程", "Python 数据分析"],
    metadatas=[
        {"category": "tutorial", "language": "python", "level": "beginner"},
        {"category": "book", "language": "java", "level": "advanced"},
        {"category": "tutorial", "language": "python", "level": "intermediate"},
    ],
    ids=["1", "2", "3"]
)

# 查询时通过元数据过滤：只搜索 Python 相关的文档
results = collection.query(
    query_texts=["学习编程"],
    n_results=10,
    where={"language": "python"}  # 过滤条件
)

print(results["documents"])
# 输出: [['Python 入门教程', 'Python 数据分析']]

更新和删除文档

import chromadb

client = chromadb.Client()
collection = client.get_or_create_collection("my_docs")

# 先添加文档
collection.add(
    documents=["原始内容"],
    ids=["1"]
)

# 更新文档（id 不变，内容替换）
collection.update(
    documents=["更新后的内容"],
    ids=["1"]
)

# 也可以用 upsert（存在则更新，不存在则添加）
collection.upsert(
    documents=["再次更新的内容"],
    ids=["1"]
)

# 删除文档
collection.delete(ids=["1"])

# 删除整个集合
client.delete_collection("my_docs")

使用 OpenAI 嵌入模型

在实际项目中，通常使用 OpenAI 的嵌入模型来生成向量：

import chromadb
from openai import OpenAI

# 初始化 OpenAI 客户端
openai_client = OpenAI()

def get_embedding(text: str) -> list[float]:
    """使用 OpenAI 的嵌入模型将文本转换为向量"""
    response = openai_client.embeddings.create(
        model="text-embedding-3-small",
        input=text
    )
    return response.data[0].embedding

# 创建 ChromaDB 集合
chroma_client = chromadb.PersistentClient(path="./chroma_data")
collection = chroma_client.get_or_create_collection("knowledge_base")

# 添加文档
docs = [
    "Python 是一种解释型编程语言",
    "FastAPI 是一个高性能的 Python Web 框架",
    "Docker 可以将应用容器化部署",
]

for i, doc in enumerate(docs):
    embedding = get_embedding(doc)
    collection.upsert(
        documents=[doc],
        embeddings=[embedding],
        ids=[str(i)]
    )

# 查询
query = "Python Web 开发"
query_embedding = get_embedding(query)

results = collection.query(
    query_embeddings=[query_embedding],
    n_results=2
)

print(results["documents"])

构建简单的 RAG 系统

下面是一个完整的 RAG 系统示例，包含文档入库、检索和生成回答：

import chromadb
from openai import OpenAI

openai_client = OpenAI()
chroma_client = chromadb.PersistentClient(path="./rag_data")
collection = chroma_client.get_or_create_collection("rag_docs")

def get_embedding(text: str) -> list[float]:
    """生成文本的嵌入向量"""
    response = openai_client.embeddings.create(
        model="text-embedding-3-small",
        input=text
    )
    return response.data[0].embedding

def add_documents(texts: list[str]):
    """将文档添加到向量数据库"""
    for i, text in enumerate(texts):
        embedding = get_embedding(text)
        collection.upsert(
            documents=[text],
            embeddings=[embedding],
            ids=[f"doc_{i}"],
            metadatas=[{"source": "user_input"}]
        )

def search_documents(query: str, n_results: int = 3) -> list[str]:
    """检索与查询最相关的文档"""
    query_embedding = get_embedding(query)
    results = collection.query(
        query_embeddings=[query_embedding],
        n_results=n_results
    )
    return results["documents"][0]

def ask_question(question: str) -> str:
    """RAG 问答：检索相关文档，然后让大模型生成回答"""
    # 第一步：检索相关文档
    relevant_docs = search_documents(question)

    # 第二步：构建提示词
    context = "\n".join(relevant_docs)
    prompt = f"""请根据以下参考资料回答问题。如果参考资料中没有相关信息，请说明。

参考资料：
{context}

问题：{question}

回答："""

    # 第三步：调用大语言模型生成回答
    response = openai_client.chat.completions.create(
        model="gpt-4o-mini",
        messages=[{"role": "user", "content": prompt}]
    )

    return response.choices[0].message.content

# 使用示例
# 先添加知识库文档
add_documents([
    "Python 的 GIL（全局解释器锁）使得同一时刻只有一个线程执行 Python 字节码。",
    "FastAPI 支持 async/await 异步编程，可以高效处理并发请求。",
    "Celery 是 Python 的分布式任务队列，常用于处理耗时操作如发送邮件、生成报表。",
])

# 然后提问
answer = ask_question("Python 如何处理异步任务？")
print(answer)

常见场景

文档问答系统

将公司内部文档（PDF、Word、Wiki）分块后存入向量数据库，员工可以通过自然语言提问，系统自动检索相关内容并生成回答。

# 将长文档切分成小块
def split_text(text: str, chunk_size: int = 500, overlap: int = 50) -> list[str]:
    """将长文本切分成重叠的小块"""
    chunks = []
    start = 0
    while start < len(text):
        end = start + chunk_size
        chunks.append(text[start:end])
        start += chunk_size - overlap  # 重叠部分确保不丢失上下文
    return chunks

# 处理长文档
long_text = "这是一篇很长的文档..."  # 你的文档内容
chunks = split_text(long_text)

# 将每个块存入向量数据库
for i, chunk in enumerate(chunks):
    embedding = get_embedding(chunk)
    collection.upsert(
        documents=[chunk],
        embeddings=[embedding],
        ids=[f"chunk_{i}"],
        metadatas=[{"chunk_index": i}]
    )

语义搜索

传统的关键词搜索只能匹配包含相同词的文档，语义搜索能理解用户意图，找到意思相近但不包含相同关键词的文档：

# 添加文档
collection.add(
    documents=[
        "如何提高代码质量",
        "代码审查的最佳实践",
        "数据库性能优化方法",
    ],
    ids=["1", "2", "3"]
)

# 语义搜索：即使用词不同，也能找到相关文档
results = collection.query(
    query_texts=["怎样写出更好的程序"],  # 没有用"代码质量"这个词
    n_results=2
)
# 仍然能找到"如何提高代码质量"和"代码审查的最佳实践"

推荐系统

通过向量相似度为用户推荐内容：

# 假设用户喜欢某篇文章，找到相似的文章推荐
liked_article = "Python 异步编程入门"
results = collection.query(
    query_texts=[liked_article],
    n_results=5
)

# 返回与用户喜欢的文章最相似的内容
for doc in results["documents"][0]:
    print(f"推荐: {doc}")

注意事项与常见错误

嵌入模型的选择

嵌入模型的质量直接决定搜索效果。不同的嵌入模型生成的向量维度和语义空间不同，不能混用。一旦选定了嵌入模型，整个项目都应该使用同一个模型。

文档分块策略

文档太大时需要切分成小块再存入向量数据库。分块太小会丢失上下文，太大会导致检索不精确。一般建议每块 300-1000 个字符，相邻块之间有 50-100 个字符的重叠。

向量维度必须一致

同一个集合中所有向量的维度必须相同。如果你使用的嵌入模型输出 1536 维向量，那么所有文档和查询的向量都必须是 1536 维。

ChromaDB 不适合大规模生产环境

ChromaDB 适合开发和小规模应用。如果数据量超过百万级别或需要高可用性，应该考虑 Milvus、Pinecone 等专业向量数据库。

搜索结果需要验证

向量搜索返回的是”语义相似”的结果，但不一定是”正确”的结果。在 RAG 系统中，应该让大语言模型判断检索到的内容是否真的与问题相关。

进阶用法

使用 Qdrant 向量数据库

Qdrant 是一个高性能的开源向量数据库，适合生产环境：

# 安装 Qdrant 客户端
pip install qdrant-client

from qdrant_client import QdrantClient
from qdrant_client.models import Distance, VectorParams, PointStruct

# 连接 Qdrant（本地模式）
client = QdrantClient(":memory:")

# 也可以连接远程服务器
# client = QdrantClient(url="http://localhost:6333")

# 创建集合
client.create_collection(
    collection_name="my_docs",
    vectors_config=VectorParams(size=3, distance=Distance.COSINE),
)

# 插入向量
client.upsert(
    collection_name="my_docs",
    points=[
        PointStruct(id=1, vector=[0.9, 0.1, 0.0], payload={"text": "苹果是一种水果"}),
        PointStruct(id=2, vector=[0.85, 0.15, 0.0], payload={"text": "香蕉也是一种水果"}),
        PointStruct(id=3, vector=[0.0, 0.1, 0.9], payload={"text": "汽车是交通工具"}),
    ]
)

# 搜索
results = client.search(
    collection_name="my_docs",
    query_vector=[0.9, 0.1, 0.0],
    limit=2
)

for result in results:
    print(f"相似度: {result.score:.4f}, 内容: {result.payload['text']}")

混合搜索

混合搜索结合了向量搜索和关键词搜索的优点，既能理解语义，又能精确匹配关键词：

import chromadb

client = chromadb.Client()
collection = client.get_or_create_collection("hybrid_search")

# 添加文档
collection.add(
    documents=[
        "Python 3.12 发布于 2023 年 10 月",
        "Python 是最流行的编程语言之一",
        "Java 21 也发布了新特性",
    ],
    ids=["1", "2", "3"]
)

# 向量搜索（语义匹配）
vector_results = collection.query(
    query_texts=["Python 最新版本"],
    n_results=2
)

# 结合关键词过滤
filtered_results = collection.query(
    query_texts=["编程语言"],
    n_results=10,
    where_document={"$contains": "Python"}  # 文档必须包含"Python"
)

使用 LangChain 集成向量数据库

LangChain 是一个流行的 LLM 应用开发框架，提供了向量数据库的统一接口：

pip install langchain langchain-chroma langchain-openai

from langchain_chroma import Chroma
from langchain_openai import OpenAIEmbeddings
from langchain_openai import ChatOpenAI
from langchain.chains import RetrievalQA

# 创建嵌入模型
embeddings = OpenAIEmbeddings(model="text-embedding-3-small")

# 创建向量数据库
vectorstore = Chroma.from_texts(
    texts=[
        "Python 的列表推导式是一种简洁的创建列表的方式",
        "字典是 Python 中最常用的数据结构之一",
        "装饰器可以在不修改函数代码的情况下扩展函数功能",
    ],
    embedding=embeddings,
    collection_name="python_knowledge"
)

# 创建检索器
retriever = vectorstore.as_retriever(search_kwargs={"k": 2})

# 创建问答链
llm = ChatOpenAI(model="gpt-4o-mini")
qa_chain = RetrievalQA.from_chain_type(
    llm=llm,
    retriever=retriever,
    return_source_documents=True
)

# 提问
result = qa_chain.invoke({"query": "Python 中如何扩展函数功能？"})
print(result["result"])