Chroma 向量数据库使用示例

Chroma 向量数据库使用示例

实现原理

利用向量数据库进行高效的相似性搜索。具体步骤如下：

1. 文本向量化 ：自然语言文本本身无法直接进行相似度计算。通过文本嵌入模型（如 sentence-transformers 提供的模型），可以将文本映射到高维向量空间。在这个空间中，语义相似的文本其对应的向量在空间位置上也更接近。
2. 向量存储：Chroma 数据库专门用于存储和管理这些高维向量。它会为每个向量建立索引，以便快速检索。
3. 相似度计算：当进行查询时，查询文本同样被转换为向量。向量数据库通过计算查询向量与库中存储向量之间的距离（例如余弦相似度、欧氏距离等），来找出最相似的向量，进而找到对应的原始文本。
4. 持久化存储：Chroma 支持将数据库持久化到本地磁盘，这意味着即使脚本停止运行，数据也不会丢失，下次运行时可以重新加载。

实现代码

import chromadb
from chromadb.config import Settings
import os
from sentence_transformers import SentenceTransformer # 引入 sentence-transformers

# 获取当前脚本文件的绝对路径
script_path = os.path.abspath(__file__)
# 获取脚本文件所在的目录
script_dir = os.path.dirname(script_path)

# 1. 初始化Chroma数据库
persist_dir = os.path.join(script_dir, "chroma_db_demo")
print("Chroma 持久化目录绝对路径：", os.path.abspath(persist_dir))

client = chromadb.PersistentClient(path=persist_dir)
collection = client.get_or_create_collection("demo_collection")

# 2. 示例数据
texts = [
    "人工智能正在改变世界。",
    "向量数据库可以高效存储和检索向量。",
    "大语言模型能够理解和生成自然语言。",
    "数据可视化有助于洞察数据规律。"
]
metadatas = [{"id": i} for i in range(len(texts))]
ids = [f"doc_{i}" for i in range(len(texts))]

# 3. 加载预训练的文本嵌入模型 (确保已安装 sentence-transformers: pip install sentence-transformers)
# 这里我们选用一个常用的中文模型：shibing624/text2vec-base-chinese
# 模型会自动下载到本地缓存，通常在 ~/.cache/torch/sentence_transformers/
model_name = 'shibing624/text2vec-base-chinese'
print(f"正在加载模型: {model_name}，请稍候...")
model = SentenceTransformer(model_name)
print("模型加载完成。")

# 使用模型生成文本的嵌入向量
print("正在为示例文本生成嵌入向量...")
embeddings = model.encode(texts).tolist()
print("嵌入向量生成完成。")

# 4. 插入数据库
collection.add(
    documents=texts,
    embeddings=embeddings,
    metadatas=metadatas,
    ids=ids
)
print("数据已插入Chroma向量数据库（使用随机向量）。")

# 5. 检索相似文本（使用真实查询文本和其嵌入向量）
query_text = "什么是大型语言模型？" # 定义一个实际的查询文本
print(f"\n查询文本: {query_text}")
print("正在为查询文本生成嵌入向量...")
query_embedding = model.encode([query_text]).tolist() # 将查询文本转换为嵌入向量
print("查询嵌入向量生成完成。")
results = collection.query(
    query_embeddings=query_embedding,
    n_results=2
)
print("检索结果：")
for doc, score in zip(results['documents'][0], results['distances'][0]):
    print(f"文本: {doc}，相似度: {score:.4f}")

# # 6. 删除示例
# collection.delete(ids=["doc_0"])
# print("已删除doc_0。")

# 7. 持久化（可选）
# client.persist()  # 删除此行，Chroma会自动持久化
print("数据库已持久化（自动）。")
def print_dir_tree(root_dir):
    print(f"\nChroma 持久化目录结构（{root_dir}）：")
    for root, dirs, files in os.walk(root_dir):
        level = root.replace(root_dir, '').count(os.sep)
        indent = ' ' * 4 * level
        print(f"{indent}{os.path.basename(root)}/")
        subindent = ' ' * 4 * (level + 1)
        for f in files:
            print(f"{subindent}{f}")
print_dir_tree(persist_dir)


# 8. 验证持久化
results = collection.get(ids=["doc_0"])
if results['documents'] and results['documents'][0]:
    print(f"doc_0 内容: {results['documents'][0]}")
else:
    print("doc_0 不存在，已被删除。")

功能概述

1. 初始化 Chroma 向量数据库：在本地文件系统创建一个持久化的向量数据库。
2. 加载文本嵌入模型 ：使用 sentence-transformers 库加载预训练的中文文本嵌入模型 (shibing624/text2vec-base-chinese)，用于将文本转换为向量。
3. 数据准备与向量化：定义示例文本数据，并使用加载的模型将这些文本转换为嵌入向量。
4. 数据存储：将文本、对应的嵌入向量、元数据和唯一ID存入 Chroma 数据库的集合中。
5. 相似性检索：定义一个查询文本，将其转换为嵌入向量，然后在数据库中检索最相似的文本。
6. 持久化验证：展示 Chroma 数据库的自动持久化特性，并验证数据是否成功存储。

数据流转过程

graph TD A[开始] --> B(定义示例文本和元数据); B --> C{加载/初始化 Chroma 客户端}; C --> D{获取或创建集合 "demo_collection"}; D --> E{加载 SentenceTransformer 模型 (shibing624/text2vec-base-chinese)}; E --> F(将示例文本转换为嵌入向量); F --> G{将文本、向量、元数据、ID 添加到集合}; G --> H(定义查询文本); H --> I(将查询文本转换为嵌入向量); I --> J{在集合中查询相似向量 (n_results=2)}; J --> K(输出检索结果: 文本及相似度得分); K --> L(打印持久化目录结构); L --> M{验证持久化: 获取 doc_0}; M --> N(输出 doc_0 内容或不存在信息); N --> O[结束];

• 初始化阶段 ：
• 脚本首先确定 Chroma 数据库的持久化存储路径 ( chroma_db_demo )。
• 然后，初始化一个 PersistentClient 对象，指向该路径。
• 通过客户端获取或创建一个名为 demo_collection 的集合，用于存储数据。
• 数据准备与模型加载 ：
• 定义一组示例文本 ( texts )、对应的元数据 ( metadatas ) 和唯一ID ( ids )。
• 加载 sentence-transformers 提供的 shibing624/text2vec-base-chinese 模型。这个模型擅长将中文文本转换为固定维度的向量。
• 数据向量化与存储 ：
• 使用加载的模型，将 texts 列表中的每个文本句子转换为一个嵌入向量，形成 embeddings 列表。
• 调用 collection.add() 方法，将原始文本、生成的嵌入向量、元数据和ID批量添加到 demo_collection 集合中。
• 查询与检索 ：
• 定义一个查询文本 query_text 。
• 同样使用 model.encode() 将查询文本转换为查询嵌入向量 query_embedding 。
• 调用 collection.query() 方法，传入查询嵌入向量和期望返回的结果数量 ( n_results=2 )。
• Chroma 数据库会计算查询向量与集合中所有存储向量的相似度，并返回最相似的两个结果。
• 结果展示与持久化验证 ：
• 打印检索到的文本及其相似度得分。
• 打印 Chroma 持久化目录的树状结构，展示其内部文件组织。
• 通过 collection.get(ids=["doc_0"]) 尝试获取之前插入的一条数据，以验证持久化是否成功以及数据是否可访问。

---

核心代码功能解释

1. 环境准备与导入 ：

import chromadb
from chromadb.config import Settings
import os
from sentence_transformers import SentenceTransformer # 引入 sentence-transformers

• chromadb : Chroma 向量数据库的核心库。
• os : 用于处理文件路径。
• sentence_transformers.SentenceTransformer : 用于加载和使用文本嵌入模型。

2. Chroma 客户端和集合初始化 ：

  # 获取当前脚本文件的绝对路径
  script_path = os.path.abspath(__file__)
  # 获取脚本文件所在的目录
  script_dir = os.path.dirname(script_path)
  
  # 1. 初始化Chroma数据库
  persist_dir = os.path.join(script_dir, 
  "chroma_db_demo")
  print("Chroma 持久化目录绝对路径：", os.path.abspath
  (persist_dir))
  
  client = chromadb.PersistentClient
  (path=persist_dir)
  collection = client.get_or_create_collection
  ("demo_collection")

• os.path.abspath(file) 和 os.path.dirname() : 获取当前脚本所在的目录，确保持久化路径是相对脚本位置的。
• chromadb.PersistentClient(path=persist_dir) : 创建一个持久化客户端实例，数据将保存在 persist_dir 指定的目录下。
• client.get_or_create_collection("demo_collection") : 获取名为 demo_collection 的集合。如果该集合不存在，则会自动创建它。

3. 加载文本嵌入模型 ：

   model_name = 'shibing624/text2vec-base-chinese'
   print(f"正在加载模型: {model_name}，请稍候...")
   model = SentenceTransformer(model_name)
   print("模型加载完成。")

   • SentenceTransformer(model_name) : 从 Hugging Face Hub (默认) 或本地缓存加载指定的预训练模型。第一次运行时，如果本地没有，会自动下载。

4. 文本向量化 ：

   # 使用模型生成文本的嵌入向量
   print("正在为示例文本生成嵌入向量...")
   embeddings = model.encode(texts).tolist()
   print("嵌入向量生成完成。")

   • model.encode(texts) : 接收一个文本列表，返回一个包含这些文本对应嵌入向量的 NumPy 数组。 .tolist() 将其转换为 Python 列表。

5. 数据插入 ：

   collection.add(
       documents=texts,
       embeddings=embeddings,
       metadatas=metadatas,
       ids=ids
   )

   • collection.add() : 将数据批量添加到集合中。需要提供：
     • documents : 原始文本文档列表。
     • embeddings : 与 documents 对应的嵌入向量列表。
     • metadatas : 与每个文档关联的元数据列表（字典形式）。
     • ids : 每个文档的唯一标识符列表。

6. 相似性查询 ：

   query_text = "什么是大型语言模型？"
   query_embedding = model.encode([query_text]).
   tolist()
   results = collection.query(
       query_embeddings=query_embedding,
       n_results=2
   )

   • 首先将查询文本 query_text 也转换为嵌入向量 query_embedding 。
   • collection.query() : 执行查询。
     • query_embeddings : 查询向量列表（即使只有一个查询，也需要是列表形式）。
     • n_results : 指定返回最相似结果的数量。
   • 返回的 results 是一个字典，其中 results['documents'][0] 包含检索到的文档， results['distances'][0] 包含对应的相似度得分（Chroma 返回的是距离，值越小越相似）。

7. 持久化目录打印与验证 ：

   def print_dir_tree(root_dir):
       # ... (省略打印目录树的实现)
   print_dir_tree(persist_dir)
   
   results = collection.get(ids=["doc_0"])
   if results['documents'] and results['documents']
   [0]:
       print(f"doc_0 内容: {results['documents'][0]}
       ")
   else:
       print("doc_0 不存在，已被删除。")

   • print_dir_tree : 一个辅助函数，用于可视化 Chroma 数据库在磁盘上的存储结构。
   • collection.get(ids=["doc_0"]) : 通过 ID 获取特定的文档，用于验证数据是否正确持久化并可以被检索。

进一步优化

"""
离线环境下 Chroma 向量数据库使用示例 (面向对象重构版)
依赖: chromadb, numpy, sentence-transformers, python-dotenv
安装: pip install chromadb numpy sentence-transformers python-dotenv
"""

import chromadb
from chromadb.config import Settings as ChromaSettings # 避免与自定义Settings冲突
import os
import logging
from logging.handlers import RotatingFileHandler
import argparse
from sentence_transformers import SentenceTransformer
import sys
from dotenv import load_dotenv
from typing import List, Dict, Any, Optional

# --- 全局日志记录器 --- (在主程序块中配置)
logger = logging.getLogger(__name__)

class ConfigManager:
    """配置管理器类，负责加载和解析配置。"""
    def __init__(self):
        load_dotenv() # 加载 .env 文件
        self.args = self._parse_arguments()

    def _parse_arguments(self) -> argparse.Namespace:
        """解析命令行参数，并从环境变量获取默认值。"""
        default_model_name = os.getenv('MODEL_NAME', 'shibing624/text2vec-base-chinese')
        default_persist_dir = os.getenv('PERSIST_DIR', 'chroma_db_demo')
        default_log_level = os.getenv('LOG_LEVEL', 'INFO').upper()
        default_log_file = os.getenv('LOG_FILE', 'chroma_demo.log')
        default_collection_name = os.getenv('COLLECTION_NAME', 'demo_collection')

        parser = argparse.ArgumentParser(description='ChromaDB 向量数据库示例脚本 (面向对象版).')
        parser.add_argument('--model_name', type=str, default=default_model_name,
                            help=f'Sentence Transformer 模型的名称 (默认: {default_model_name})')
        parser.add_argument('--persist_dir', type=str, default=default_persist_dir,
                            help=f'ChromaDB 持久化数据的目录名称 (默认: {default_persist_dir})')
        parser.add_argument('--log_level', type=str, default=default_log_level,
                            choices=['DEBUG', 'INFO', 'WARNING', 'ERROR', 'CRITICAL'],
                            help=f'日志记录级别 (默认: {default_log_level})')
        parser.add_argument('--log_file', type=str, default=default_log_file,
                            help=f'日志输出文件的名称 (默认: {default_log_file})')
        parser.add_argument('--collection_name', type=str, default=default_collection_name,
                            help=f'ChromaDB 集合的名称 (默认: {default_collection_name})')

        # 解决Jupyter等环境下argparse的问题
        if len(sys.argv) <= 1 and not any(arg.startswith('--') for arg in sys.argv):
            return parser.parse_args([])
        return parser.parse_args()

    @property
    def model_name(self) -> str:
        return self.args.model_name

    @property
    def persist_dir(self) -> str:
        return self.args.persist_dir

    @property
    def log_level(self) -> str:
        return self.args.log_level

    @property
    def log_file(self) -> str:
        return self.args.log_file
    
    @property
    def collection_name(self) -> str:
        return self.args.collection_name

class VectorDBManager:
    """向量数据库管理器类，封装了ChromaDB和SentenceTransformer的相关操作。"""
    def __init__(self, config: ConfigManager):
        self.config = config
        self.script_dir = os.path.dirname(os.path.abspath(__file__))
        self.persist_path = os.path.join(self.script_dir, self.config.persist_dir)
        self.model: Optional[SentenceTransformer] = None
        self.client: Optional[chromadb.PersistentClient] = None
        self.collection: Optional[chromadb.API.Collection] = None # API.Collection 是正确的类型提示
        self._initialize_model()
        self._initialize_db()

    def _initialize_model(self):
        """初始化SentenceTransformer模型。"""
        logger.info(f"正在加载模型: {self.config.model_name}，请稍候...")
        try:
            self.model = SentenceTransformer(self.config.model_name)
            logger.info(f"模型 '{self.config.model_name}' 加载完成。")
        except OSError as e_os:
            logger.error(f"加载模型 '{self.config.model_name}' 失败: {e_os}")
            raise # 重新抛出异常，让调用者处理
        except Exception as e:
            logger.error(f"加载模型 '{self.config.model_name}' 时发生未知错误: {e}")
            raise

    def _initialize_db(self):
        """初始化ChromaDB客户端和集合。"""
        logger.info(f"Chroma 持久化目录绝对路径：{os.path.abspath(self.persist_path)}")
        try:
            self.client = chromadb.PersistentClient(path=self.persist_path)
            # 尝试删除旧集合，确保维度匹配
            try:
                self.client.delete_collection(self.config.collection_name)
                logger.info(f"已尝试删除旧的 '{self.config.collection_name}' 集合以避免维度冲突。")
            except Exception as e_del: # 捕获更通用的异常，因为集合不存在时也可能抛出特定错误
                logger.info(f"尝试删除集合 '{self.config.collection_name}' 时发生错误 (可能是集合不存在): {e_del}")
            
            self.collection = self.client.get_or_create_collection(self.config.collection_name)
            logger.info(f"Chroma数据库客户端和集合 '{self.config.collection_name}' 初始化/创建成功。")
        except chromadb.errors.InvalidDimensionException as e_dim:
            logger.error(f"ChromaDB维度错误: {e_dim}. 请删除旧的持久化目录 '{self.config.persist_dir}' 后重试。")
            raise
        except chromadb.errors.ChromaAPIError as e_api:
            logger.error(f"ChromaDB API 错误: {e_api}")
            raise
        except PermissionError as e_perm:
            logger.error(f"权限错误：无法访问或写入持久化目录 '{self.persist_path}'. 错误: {e_perm}")
            raise
        except FileNotFoundError as e_fnf:
            logger.error(f"文件未找到错误：持久化路径 '{self.persist_path}' 可能不存在。错误: {e_fnf}")
            raise
        except Exception as e:
            logger.error(f"初始化Chroma数据库失败: {e}")
            raise

    def encode_texts(self, texts: List[str]) -> Optional[List[List[float]]]:
        """使用加载的模型将文本列表编码为嵌入向量列表。"""
        if not self.model:
            logger.error("模型未初始化，无法生成嵌入向量。")
            return None
        logger.info(f"正在为 {len(texts)} 条文本生成嵌入向量...")
        try:
            embeddings = self.model.encode(texts).tolist()
            logger.info("嵌入向量生成完成。")
            return embeddings
        except AttributeError as e_attr:
            logger.error(f"生成嵌入向量失败：模型对象可能未正确初始化。错误: {e_attr}")
        except Exception as e:
            logger.error(f"生成嵌入向量时发生未知错误: {e}")
        return None

    def add_documents(self, documents: List[str], metadatas: List[Dict[str, Any]], ids: List[str]) -> bool:
        """向集合中添加文档及其嵌入。添加前会检查ID是否已存在。"""
        if not self.collection:
            logger.error("数据库集合未初始化，无法添加文档。")
            return False

        if not documents or not ids or len(documents) != len(ids):
            logger.warning("输入文档、ID列表为空或长度不匹配，跳过添加。")
            return False

        # --- 检查已存在的ID并过滤 ---
        try:
            # 从集合中获取输入的ids，ChromaDB只会返回存在的ids
            existing_docs = self.collection.get(ids=ids)
            # 从返回结果中提取已存在的ids
            existing_ids = set(existing_docs.get('ids', []))
            logger.info(f"在集合 '{self.config.collection_name}' 中找到 {len(existing_ids)} 个已存在的ID。")

            # 过滤掉已存在的文档
            new_documents = []
            new_metadatas = []
            new_ids = []
            for i in range(len(ids)):
                if ids[i] not in existing_ids: # 只有ID不在已存在集合中才认为是新的
                    new_documents.append(documents[i])
                    new_metadatas.append(metadatas[i])
                    new_ids.append(ids[i])

            if not new_ids:
                logger.info("所有待添加的文档ID均已存在，无需重复添加。")
                return True # 即使没有新文档添加，也认为是成功的，因为没有错误发生

            logger.info(f"正在添加 {len(new_ids)} 个新文档...")

        except chromadb.errors.ChromaError as e_get:
            logger.error(f"检查集合中已存在的ID时发生ChromaDB错误: {e_get}. 跳过添加。")
            return False
        except Exception as e_get_other:
            logger.error(f"检查集合中已存在的ID时发生未知错误: {e_get_other}. 跳过添加。")
            return False
        # --- 结束检查过滤逻辑 ---

        embeddings = self.encode_texts(new_documents)
        if embeddings is None:
            logger.error("无法为新文档生成嵌入向量，跳过添加。")
            return False

        # --- 使用过滤后的新文档和ID进行添加 ---
        try:
            self.collection.add(
                documents=new_documents,
                embeddings=embeddings,
                metadatas=new_metadatas,
                ids=new_ids
            )
            logger.info(f"成功向集合 '{self.config.collection_name}' 添加 {len(new_documents)} 个新文档。")
            return True
        except chromadb.errors.DuplicateIDError as e_dup:
            # 理论上过滤后不应该出现此错误，如果出现可能是过滤逻辑问题或并发问题
            logger.warning(f"添加新文档时发生 DuplicateIDError: {e_dup}. 这可能是过滤逻辑的异常情况。")
            return False
        except chromadb.errors.ChromaError as e_api:
            logger.error(f"向Chroma向量数据库插入数据时发生ChromaDB错误: {e_api}")
            return False
        except Exception as e:
            logger.error(f"向Chroma向量数据库插入数据失败: {e}")
            return False

    def query_documents(self, query_text: str, n_results: int = 2) -> Optional[Dict[str, list]]:
        """根据查询文本检索相似的文档。"""
        if not self.collection or not self.model:
            logger.error("数据库集合或模型未初始化，无法查询。")
            return None

        logger.info(f"查询文本: {query_text}")
        query_embedding = self.encode_texts([query_text])
        if query_embedding is None:
            return None

        try:
            results = self.collection.query(
                query_embeddings=query_embedding,
                n_results=n_results
            )
            logger.info("检索完成。")
            return results
        except chromadb.errors.ChromaAPIError as e_api:
            logger.error(f"检索相似文本时发生ChromaDB API错误: {e_api}")
        except Exception as e:
            logger.error(f"检索相似文本时发生未知错误: {e}")
        return None

    def get_document_by_id(self, doc_id: str) -> Optional[Dict[str, list]]:
        """根据ID获取文档。"""
        if not self.collection:
            logger.error("数据库集合未初始化，无法获取文档。")
            return None
        try:
            return self.collection.get(ids=[doc_id])
        except chromadb.errors.ChromaAPIError as e_api:
            logger.error(f"通过ID '{doc_id}' 获取文档时发生ChromaDB API错误: {e_api}")
        except Exception as e:
            logger.error(f"通过ID '{doc_id}' 获取文档时发生未知错误: {e}")
        return None

    def print_persist_dir_tree(self):
        """打印持久化目录的结构树。"""
        logger.info(f"\nChroma 持久化目录结构（{self.persist_path}）：")
        for root, _, files in os.walk(self.persist_path):
            level = root.replace(self.persist_path, '').count(os.sep)
            indent = ' ' * 4 * level
            logger.info(f"{indent}{os.path.basename(root)}/")
            subindent = ' ' * 4 * (level + 1)
            for f in files:
                logger.info(f"{subindent}{f}")

def setup_logging(log_level_str: str, log_file_name: str):
    """配置全局日志记录器。"""
    log_level_numeric = getattr(logging, log_level_str.upper(), logging.INFO)
    
    script_dir = os.path.dirname(os.path.abspath(__file__))
    log_file_path = os.path.join(script_dir, log_file_name)

    logger.setLevel(log_level_numeric)

    console_handler = logging.StreamHandler()
    console_handler.setLevel(log_level_numeric)

    file_handler = RotatingFileHandler(
        log_file_path, 
        maxBytes=1024*1024, # 1 MB
        backupCount=3,
        encoding='utf-8'
    )
    file_handler.setLevel(log_level_numeric)

    formatter = logging.Formatter('%(asctime)s - %(name)s - %(levelname)s - %(message)s')
    console_handler.setFormatter(formatter)
    file_handler.setFormatter(formatter)

    if not logger.handlers: # 防止重复添加handler
        logger.addHandler(console_handler)
        logger.addHandler(file_handler)
    
    logger.info(f"日志将输出到控制台和文件: {log_file_path}")


def main():
    """主执行函数。"""
    try:
        config_manager = ConfigManager()
        setup_logging(config_manager.log_level, config_manager.log_file)
        logger.info(f"脚本将使用以下配置: 模型名称='{config_manager.model_name}', "
                    f"持久化目录='{config_manager.persist_dir}', "
                    f"日志级别='{config_manager.log_level}', "
                    f"集合名称='{config_manager.collection_name}'")

        db_manager = VectorDBManager(config_manager)

    except Exception as e:
        logger.critical(f"脚本初始化失败: {e}", exc_info=True)
        sys.exit(1)

    # 示例数据
    texts = [
        "人工智能正在改变世界。",
        "向量数据库可以高效存储和检索向量。",
        "大语言模型能够理解和生成自然语言。",
        "数据可视化有助于洞察数据规律。"
    ]
    metadatas = [{"id": i} for i in range(len(texts))]
    ids = [f"doc_{i}" for i in range(len(texts))]

    # 添加文档
    if db_manager.add_documents(texts, metadatas, ids):
        logger.info("示例数据添加成功。")
    else:
        logger.error("示例数据添加失败。")

    # 查询文档
    query_text = "什么是大型语言模型？"
    results = db_manager.query_documents(query_text, n_results=2)
    if results and results.get('documents') and results['documents'][0]:
        logger.info("检索结果：")
        for doc, score in zip(results['documents'][0], results['distances'][0]):
            logger.info(f"  文本: {doc}，相似度: {score:.4f}")
    else:
        logger.warning("未找到相似的检索结果或结果格式不正确。")

    # 验证持久化
    retrieved_doc = db_manager.get_document_by_id("doc_0")
    # 检查 retrieved_doc 是否非空且包含 documents 列表，且该列表非空
    if retrieved_doc and retrieved_doc.get('documents') and retrieved_doc['documents']:
        logger.info(f"通过ID获取 'doc_0' 内容: {retrieved_doc['documents'][0]}")
    else:
        logger.info("'doc_0' 不存在或已被删除，或get返回结果格式不正确。")

    db_manager.print_persist_dir_tree()
    logger.info("脚本执行完毕。")

if __name__ == "__main__":
    main()

主要提升点包括：

1、面向对象结构 : 将代码组织在 ConfigManager (配置管理) 和 VectorDBManager (向量数据库操作) 类中。这提高了代码的模块化、可读性、可重用性和可维护性。不同的功能职责更清晰。

2、灵活的配置管理 : 引入了 ConfigManager 类，支持从 .env 文件加载配置，并且可以使用命令行参数覆盖环境变量。这使得脚本的配置更加灵活，无需修改代码即可适应不同的环境和需求。

3、标准化的日志记录 : 使用 Python 标准库的 logging 模块进行日志记录。通过 setup_logging 函数进行统一配置，支持不同的日志级别 (INFO, ERROR 等)，可以将日志同时输出到控制台和文件，并且支持日志文件轮转。这极大地提高了脚本的可观测性和问题排查效率。

3、更详细和健壮的错误处理 : 在关键操作（如模型加载、DB 初始化、数据添加、查询）中使用了更细粒度的 try...except 块，捕获了特定类型的异常（如 OSError, chromadb.errors.ChromaError, PermissionError, FileNotFoundError, DuplicateIDError 等），并记录了更详细的错误信息，有时还会重新抛出异常以便调用者处理。这使得脚本在遇到问题时能够提供更准确的反馈。

4、优化的数据添加逻辑 (基于 ID 去重): 在 add_documents 方法内部封装了基于 ID 的去重逻辑，通过查询现有 ID 并过滤的方式，确保只有尚未存在于集合中的文档才会被实际添加到数据库，并输出相应的日志提示。这避免了重复添加相同 ID 的数据。

5、结构化的数据库操作方法 : 在 VectorDBManager 类中提供了封装好的方法 (encode_texts, add_documents, query_documents, get_document_by_id, print_persist_dir_tree)，使得主函数逻辑更清晰，对 ChromaDB 库的直接依赖性降低，更容易理解和修改各个操作步骤。

通过引入面向对象设计、灵活的配置、标准化的日志和更详细的错误处理，将一个简单的过程式脚本重构为一个更具工程化、更健壮且易于维护的向量数据库示例应用。

对进一步学习或实践的建议

1. 探索不同的嵌入模型 ：尝试
   sentence-transformers 库中其他的中文或多语言模型，观察
   它们在不同任务上的表现和生成的向量维度。
2. 理解相似度/距离度量 ：ChromaDB 默认使用 L2 距离。
   您可以研究一下其他度量方式，如余弦相似度 (cosine
   similarity) 或内积 (inner product)，以及如何在
   ChromaDB 中配置它们（如果支持）。
3. 元数据过滤 ：学习如何在 collection.query() 中使
   用 where 参数，根据元数据进行过滤查询，例如只检索特定来源
   或时间的文档。
4. 更新和删除操作 ：实践 collection.update() 和
   collection.delete() 方法，了解如何管理数据库中的数据。
5. 批量操作与性能 ：当处理大量数据时，了解批量添加
   (add) 和批量查询 (query) 的性能优势。
6. ChromaDB 的更多特性 ：查阅 ChromaDB 的官方文档，
   了解更多高级功能，如多模态数据支持、更复杂的集合配置等。
7. 实际应用场景 ：思考如何将这个基础示例扩展到实际应用
   中，例如：
   • 构建一个简单的文档搜索引擎。
   • 实现一个基于语义的问答机器人原型。
   • 为一组产品描述实现相似产品推荐功能。