Langchain中集成Amazon DocumentDB作为向量存储的指南

pip install langchain-community pymongo boto3

import os
from pymongo import MongoClient
from langchain_community.vectorstores.docdb import DocumentDBVectorSearch
from langchain_community.embeddings import BedrockEmbeddings
from langchain.text_splitter import CharacterTextSplitter
from langchain_core.documents import Document

# DocumentDB 连接信息
# 推荐使用环境变量或安全配置管理这些敏感信息
DOCDB_CONNECTION_STRING = os.environ.get("DOCDB_CONNECTION_STRING")
DOCDB_DB_NAME = "your_database_name"
DOCDB_COLLECTION_NAME = "your_collection_name"
DOCDB_INDEX_NAME = "vector_index" # 确保在DocumentDB中已创建此向量索引

if not DOCDB_CONNECTION_STRING:
    raise ValueError("DOCDB_CONNECTION_STRING 环境变量未设置。")

# 初始化MongoDB客户端
client = MongoClient(DOCDB_CONNECTION_STRING)
db = client[DOCDB_DB_NAME]
collection = db[DOCDB_COLLECTION_NAME]

# 初始化嵌入模型
# 这里以Amazon Bedrock Embeddings为例
bedrock_embeddings = BedrockEmbeddings(
    model_id="amazon.titan-embed-text-v1",
    region_name="us-east-1" # 根据您的AWS区域调整
)

# 示例文档
raw_documents = [
    "Langchain是一个用于开发LLM应用的框架。",
    "Amazon DocumentDB提供了MongoDB兼容的数据库服务，现在支持向量搜索。",
    "向量搜索是检索非结构化数据的重要技术。",
    "Bedrock Embeddings是Amazon提供的文本嵌入服务。"
]

# 文档分割器（如果需要处理长文档）
text_splitter = CharacterTextSplitter(chunk_size=1000, chunk_overlap=0)
docs = []
for i, text in enumerate(raw_documents):
    # 将原始文本包装成Langchain的Document对象
    # 也可以使用text_splitter.create_documents([text])来处理更长的文本
    docs.append(Document(page_content=text, metadata={"source": f"doc_{i}"}))

# 如果文档是字符串列表，也可以直接传入
# docs = text_splitter.create_documents(raw_documents)

print("开始上传文档到DocumentDB...")
vector_store = DocumentDBVectorSearch.from_documents(
    documents=docs,
    embedding=bedrock_embeddings,
    collection=collection, # 传入pymongo的collection对象
    index_name=DOCDB_INDEX_NAME,
)
print("文档上传成功！")

# 执行相似度搜索
query = "什么是Langchain和向量搜索？"
print(f"\n执行相似度搜索，查询: '{query}'")
results = vector_store.similarity_search(query, k=2)

print("\n搜索结果:")
for i, doc in enumerate(results):
    print(f"--- 结果 {i+1} ---")
    print(f"内容: {doc.page_content}")
    print(f"元数据: {doc.metadata}")

# 也可以将其用作检索器
retriever = vector_store.as_retriever(search_kwargs={"k": 1})
retrieved_docs = retriever.invoke(query)
print(f"\n使用检索器获取结果 (k=1):")
for i, doc in enumerate(retrieved_docs):
    print(f"--- 检索结果 {i+1} ---")
    print(f"内容: {doc.page_content}")
    print(f"元数据: {doc.metadata}")