LlamaIndex 使用Chroma作为向量数据库

在 LlamaIndex 构建本地文件向量数据库 当中，我们将由文本转换的向量数据保存在了本地的文件目录中，这种方式适合于小量的数据。如果数据量较大，那么可以使用更专业的向量数据库，例如开源的 chroma，这篇文章修改了前面的例子，演示如何使用chroma来保存向量数据，并执行查询。

构建chroma向量数据库

在工作目录下，创建文件 txt_chroma.py，用于构建chroma向量库：

# 构建向量库，使用chroma from dotenv import load_dotenv load_dotenv() # 选择embedding模型 from llama_index.core import Settings from llama_index.embeddings.openai import OpenAIEmbedding Settings.embed_model = OpenAIEmbedding(model="text-embedding-3-large") # 1. 加载数据 from llama_index.core import SimpleDirectoryReader documents = SimpleDirectoryReader(input_files=["./player.txt"]).load_data() # 2. 初始化chroma客户端，设置存储路径 import chromadb chroma_client = chromadb.PersistentClient(path="./chroma_db") # 3. 创建chroma集合 from llama_index.vector_stores.chroma import ChromaVectorStore chroma_collection = chroma_client.get_or_create_collection("player") # 4. 创建 Chroma 向量存储实例 from llama_index.core import StorageContext from llama_index.vector_stores.chroma import ChromaVectorStore # 绑定 Chroma 集合到向量存储 vector_store = ChromaVectorStore(chroma_collection=chroma_collection) # 初始化存储上下文 storage_context = StorageContext.from_defaults(vector_store=vector_store) # 5. 构建向量索引 from llama_index.core import VectorStoreIndex vector_index = VectorStoreIndex.from_documents( documents=documents, storage_context=storage_context ) # 6. 存储 vector_index.storage_context.persist(persist_dir="./chroma_db")

上面代码的核心部分，在于构建向量索引 vector_index 时，传入自行构建的 storage_context 对象。而这个 storage_context对象，又关联了chroma向量存储。

进行查询

在工作目录下，再创建一个新的文件：txt_chroma_query.py，对上一步构建的chroma向量库进行查询：

# 对chroma向量数据库进行查询 # 配置模型 from llama_index.core import Settings, StorageContext from llama_index.llms.openai import OpenAI from llama_index.embeddings.openai import OpenAIEmbedding Settings.embed_model = OpenAIEmbedding(model="text-embedding-3-large") Settings.llm = OpenAI(model="gpt-4o") # 1. 加载本地文件向量数据库 chroma import chromadb db = chromadb.PersistentClient(path="./chroma_db") collection = db.get_collection("player") # 2. 构建向量索引 from llama_index.vector_stores.chroma import ChromaVectorStore from llama_index.core import VectorStoreIndex vector_store = ChromaVectorStore(chroma_collection=collection) storage_context = StorageContext.from_defaults(vector_store=vector_store) vector_index = VectorStoreIndex.from_vector_store( storage_context=storage_context, vector_store=vector_store, ) # 3. 构建查询引擎 index_count = collection.count() query_engine = vector_index.as_query_engine( similarity_top_k= index_count if index_count >=1 else 1, ) response = query_engine.query("ShadowBlade_007 注册了多久了？") # 4. 打印结果 print(response)

上面代码中，创建query_engine时，我们设置了similarity_top_k参数。这是因为在 query_engine = vector_index.as_query_engine() 中，LlamaIndex 默认的相似性检索参数为 similarity_top_k=2（即返回最相关的 2 个结果）。

但是当前 Chroma 集合 player 中仅包含 ​​1 个向量化文档​​（例如 player.txt 文件内容过短或未分块），导致索引容量不足。因此，如果采用默认配置，就会输出一段异常信息：Number of requested results 2 is greater than number of elements in index 1, updating n_results = 1

这主要还是 player.txt 内容过短（因为我们只有10个玩家的示例信息），LlamaIndex 默认分块器 SentenceSplitter 未生成足够节点（Node），导致索引中仅有 1 个向量

最终的输出结果是类似的，这里就不再演示了。

感谢阅读，希望这篇文章能给你带来帮助！