面向RAG聊天机器人的Chroma向量AWS S3最优存取方案咨询

问题概述 #

我正在开发一款基于ChromaDB的RAG聊天机器人，包含两大核心模块：

• 管理员页面：上传PDF文档，通过ChromaDB转换为向量后存储到AWS S3
• 聊天页面：供用户与机器人交互，实现高效检索

目前的目标是优化PDF转Chroma向量的流程，并将向量无缝集成到S3中，但在序列化/反序列化Chroma向量到SQLite3时遇到了问题。想请教在AWS S3环境下处理Chroma向量的最佳实践：是继续探索SQLite3存储序列化向量，还是考虑FAISS等替代方案以获得更好的性能与可扩展性？

当前实现代码 #

序列化SQLite存储函数 #

def create_sqlite_database(db_conn, documents):
    # 创建表并插入数据到SQLite
    cursor = db_conn.cursor()
    cursor.execute('''CREATE TABLE IF NOT EXISTS embeddings (
                        id INTEGER PRIMARY KEY AUTOINCREMENT,
                        document_id TEXT,
                        embedding BLOB
                    )''')

    # 插入数据
    for doc_id, embedding in enumerate(documents):
        cursor.execute("INSERT INTO embeddings (document_id, embedding) VALUES (?, ?)",
                       (str(doc_id), pickle.dumps(embedding)))

    db_conn.commit()

向量提取与存储主函数 #

def extract_and_store_embeddings(bucket_name, pdf_key, region='ap-south-1'):
    try:
        # 调试信息
        st.write(f"Bucket Name: {bucket_name}")
        st.write(f"PDF Key: {pdf_key}")
        st.write(f"Region: {region}")

        # 从S3加载PDF（使用Amazon Textract）
        loader = AmazonTextractPDFLoader(f's3://{bucket_name}/{pdf_key}', region_name=region)
        st.write("Loader初始化完成")
        text_documents = loader.load()
        st.write("文档加载完成")

        # 文档分块
        txt_spt = RecursiveCharacterTextSplitter(chunk_size=100000, chunk_overlap=10000)
        documents = txt_spt.split_documents(text_documents)
        st.write("文档分块完成")

        # 初始化Chroma DB与嵌入模型
        embeddings = OpenAIEmbeddings()
        st.write("嵌入模型初始化完成")
        db = Chroma.from_documents(documents, embeddings, persist_directory="./chroma_db")
        st.write("Chroma DB创建完成")

        # 创建临时SQLite数据库
        temp_db_file = tempfile.NamedTemporaryFile(delete=False, suffix='.sqlite')
        db_conn = sqlite3.connect(temp_db_file.name)
        st.write("SQLite数据库连接成功")

        # 插入数据到SQLite
        create_sqlite_database(db_conn, documents)
        st.write("数据插入SQLite完成")

        # 上传SQLite备份到S3
        s3key = f"{os.path.splitext(pdf_key)[0]}.backup.sqlite"
        s3_client = boto3.client('s3')
        s3_client.upload_file(temp_db_file.name, bucket_name, s3key)
        st.success(f"SQLite备份已上传至S3：s3://{bucket_name}/{s3key}")

        return db
    except boto3.exceptions.S3UploadFailedError as e:
        st.error(f"S3上传失败: {e}")
    except NoCredentialsError as e:
        st.error(f"未找到AWS凭证: {e}")
    except PartialCredentialsError as e:
        st.error(f"AWS凭证不完整: {e}")
    except Exception as e:
        st.error(f"发生错误: {e}")

最佳实践建议 #

为什么不推荐继续使用SQLite方案 #

• S3与SQLite的天然不兼容：SQLite是单文件型数据库，依赖本地文件系统的随机读写能力，而S3是对象存储，仅支持完整对象的上传/下载，无法直接修改文件内容。每次访问S3上的SQLite文件都需要下载整个文件，不仅性能低下，并发场景下还容易出现数据损坏。
• Pickle序列化的隐患：Pickle序列化存在安全风险（可能被注入恶意代码），且版本兼容性差，不同Python环境或依赖版本下可能无法正常反序列化向量数据。
• 冗余的工作流：ChromaDB本身已经内置了向量持久化、元数据管理功能，额外将向量转存到SQLite属于重复工作，增加了系统复杂度和出错概率。

ChromaDB与S3集成的正确方式 #

直接利用ChromaDB的原生持久化功能，将生成的向量数据直接上传到S3，无需中间转换：

• 生成并持久化向量：创建Chroma实例后调用db.persist()，将向量和元数据写入本地persist_directory。
• 上传到S3：将persist_directory下的所有文件按目录结构上传到S3指定路径。
• 聊天端加载：从S3下载持久化文件到本地临时目录，通过Chroma(persist_directory="./temp_chroma", embedding_function=embeddings)加载向量库。

优化后的核心代码片段：

def extract_and_store_embeddings(bucket_name, pdf_key, region='ap-south-1'):
    try:
        # 省略文档加载、分块、嵌入初始化步骤...

        # 初始化Chroma并持久化
        db = Chroma.from_documents(documents, embeddings, persist_directory="./chroma_db")
        db.persist()
        st.write("Chroma向量已持久化到本地")

        # 将Chroma持久化目录上传到S3
        s3_prefix = f"{os.path.splitext(pdf_key)[0]}/chroma_db/"
        s3_client = boto3.client('s3')
        # 遍历目录上传所有文件
        for root, dirs, files in os.walk("./chroma_db"):
            for file in files:
                local_path = os.path.join(root, file)
                s3_path = os.path.join(s3_prefix, os.path.relpath(local_path, "./chroma_db"))
                s3_client.upload_file(local_path, bucket_name, s3_path)
        
        st.success(f"Chroma向量库已上传至S3：s3://{bucket_name}/{s3_prefix}")
        return db
    # 异常处理部分...

替代方案评估 #

• FAISS：FAISS是高性能向量检索库，但仅专注于向量相似度计算，缺乏元数据管理、持久化、索引自动维护等数据库级功能，需要自行实现文档关联、存储逻辑，适合对检索性能要求极高但无需复杂元数据操作的场景。对于新手而言，ChromaDB的封装性更强，更易上手。
• AWS托管向量数据库：如果未来需要更高的可扩展性和性能，可考虑迁移到AWS原生服务：
  • Amazon OpenSearch Serverless：支持向量检索，自带S3快照备份，适合大规模数据场景。
  • Amazon Bedrock向量存储：与Bedrock大模型深度集成，简化RAG流程。
• ChromaDB云托管版：无需自行管理基础设施，直接使用Chroma提供的云服务，支持多环境同步。

额外优化建议 #

• 调整文档分块参数：当前chunk_size=100000过大，会导致检索精度下降，建议调整为chunk_size=1000-2000，chunk_overlap=100-200，平衡检索精度和效率。
• 批量操作优化：避免循环插入数据，尽量使用Chroma的批量API提升处理速度。
• 临时文件清理：上传完成后及时清理本地临时文件，避免占用存储空间。

内容的提问来源于stack exchange，提问作者Rishil Boddula