当前位置：首页 > news >正文

chromadb向量数据库使用（2）

news 来源：原创 2025/5/1 6:18:28

- 代码
- 代码解释
- - **1. 导入 chatGLM 嵌入函数**
  - **2. 创建 ChromaDB 客户端和集合**
  - **3. 查询集合中的数据数量**
  - **4. 添加数据到集合**
  - **5. 获取已存储的文档**
  - **6. 更新文档**
  - **7. 再次获取数据，验证更新结果**
  - **8. 删除某个文档**
  - **9. 获取已删除的文档**
  - **总结**

代码

import chromadb.utils.embedding_functions as embedding_functions

openai_ef = embedding_functions.OpenAIEmbeddingFunction(
                api_key="your api key",
                api_base="https://open.bigmodel.cn/api/paas/v4",
                model_name="embedding-3"
            )
import chromadb
client = chromadb.Client()
collection = client.get_or_create_collection(name="my_collection", embedding_function=openai_ef)

collection.count()

collection.add(
    documents=["lorem ipsum...", "doc2", "doc3"],
    metadatas=[{"chapter": "3", "verse": "16"}, {"chapter": "3", "verse": "5"}, {"chapter": "29", "verse": "11"}],
    ids=["id1", "id2", "id3"]
)

collection.get(
	ids=["id1", "id2", "id3"]
)

{'ids': ['id1', 'id2', 'id3'],
 'embeddings': None,
 'documents': ['lorem ipsum...', 'doc2', 'doc3'],
 'uris': None,
 'data': None,
 'metadatas': [{'chapter': '3', 'verse': '16'},
  {'chapter': '3', 'verse': '5'},
  {'chapter': '29', 'verse': '11'}],
 'included': [<IncludeEnum.documents: 'documents'>,
  <IncludeEnum.metadatas: 'metadatas'>]}

collection.update(
    ids=["id1", "id2", "id3"],
    metadatas=[{"chapter": "3", "verse": "16"}, {"chapter": "3", "verse": "5"}, {"chapter": "29", "verse": "11"}],
    documents=["doc1", "doc2", "doc3"]
)

collection.get(
	ids=["id1", "id2", "id3"]
)

{'ids': ['id1', 'id2', 'id3'],
 'embeddings': None,
 'documents': ['doc1', 'doc2', 'doc3'],
 'uris': None,
 'data': None,
 'metadatas': [{'chapter': '3', 'verse': '16'},
  {'chapter': '3', 'verse': '5'},
  {'chapter': '29', 'verse': '11'}],
 'included': [<IncludeEnum.documents: 'documents'>,
  <IncludeEnum.metadatas: 'metadatas'>]}

collection.delete(
    ids=["id1"]
)

collection.get(
	ids=["id1"]
)

{'ids': [],
 'embeddings': None,
 'documents': [],
 'uris': None,
 'data': None,
 'metadatas': [],
 'included': [<IncludeEnum.documents: 'documents'>,
  <IncludeEnum.metadatas: 'metadatas'>]}

代码解释

1. 导入 chatGLM 嵌入函数

import chromadb.utils.embedding_functions as embedding_functions

openai_ef = embedding_functions.OpenAIEmbeddingFunction(
                api_key="your api key",
                api_base="https://open.bigmodel.cn/api/paas/v4",
                model_name="embedding-3"
            )

解释：

embedding_functions.OpenAIEmbeddingFunction 用于创建 chatGLM 的文本嵌入函数，API 请求会发送到 https://open.bigmodel.cn/api/paas/v4，使用 embedding-3 作为嵌入模型。
这里的 api_key 需要替换为你的实际 API 密钥。

2. 创建 ChromaDB 客户端和集合

import chromadb
client = chromadb.Client()
collection = client.get_or_create_collection(name="my_collection", embedding_function=openai_ef)

解释：

chromadb.Client() 创建一个 ChromaDB 客户端。
client.get_or_create_collection() 用于获取或创建一个名为 "my_collection" 的集合，并使用 openai_ef 作为嵌入函数。

3. 查询集合中的数据数量

collection.count()

输出：

解释：

collection.count() 返回集合中的文档数量。
由于此时还没有向集合添加任何数据，所以返回 0。

4. 添加数据到集合

collection.add(
    documents=["lorem ipsum...", "doc2", "doc3"],
    metadatas=[{"chapter": "3", "verse": "16"}, {"chapter": "3", "verse": "5"}, {"chapter": "29", "verse": "11"}],
    ids=["id1", "id2", "id3"]
)

解释：

documents：存储要添加的文档内容，如 "lorem ipsum..."。
metadatas：为每个文档添加元数据，例如 "chapter" 和 "verse"。
ids：为每个文档指定唯一的 id（id1，id2，id3）。

5. 获取已存储的文档

collection.get(
	ids=["id1", "id2", "id3"]
)

输出：

{
    'ids': ['id1', 'id2', 'id3'],
    'embeddings': None,
    'documents': ['lorem ipsum...', 'doc2', 'doc3'],
    'uris': None,
    'data': None,
    'metadatas': [{'chapter': '3', 'verse': '16'},
                  {'chapter': '3', 'verse': '5'},
                  {'chapter': '29', 'verse': '11'}],
    'included': ['documents', 'metadatas']
}

解释：

documents：返回存储的文档内容。
metadatas：返回对应的元数据。
ids：返回请求的 id。
embeddings 为空，因为当前存储时没有计算嵌入（可能是 OpenAI API 限制或者设置问题）。

6. 更新文档

collection.update(
    ids=["id1", "id2", "id3"],
    metadatas=[{"chapter": "3", "verse": "16"}, {"chapter": "3", "verse": "5"}, {"chapter": "29", "verse": "11"}],
    documents=["doc1", "doc2", "doc3"]
)

解释：

更新 id1 对应的 document 从 "lorem ipsum..." 变为 "doc1"。
其他 id 保持不变。

7. 再次获取数据，验证更新结果

collection.get(
	ids=["id1", "id2", "id3"]
)

输出：

{
    'ids': ['id1', 'id2', 'id3'],
    'embeddings': None,
    'documents': ['doc1', 'doc2', 'doc3'],
    'uris': None,
    'data': None,
    'metadatas': [{'chapter': '3', 'verse': '16'},
                  {'chapter': '3', 'verse': '5'},
                  {'chapter': '29', 'verse': '11'}],
    'included': ['documents', 'metadatas']
}

解释：

id1 的 document 成功更新为 "doc1"，其他数据未改变。

8. 删除某个文档

collection.delete(
    ids=["id1"]
)

解释：

删除 id1 对应的文档。

9. 获取已删除的文档

collection.get(
	ids=["id1"]
)

输出：

{
    'ids': [],
    'embeddings': None,
    'documents': [],
    'uris': None,
    'data': None,
    'metadatas': [],
    'included': ['documents', 'metadatas']
}