ChromaDB是一个专为AI时代设计的开源向量数据库，核心功能是高效存储和检索高维向量数据（embeddings）。与传统基于关键词匹配的数据库不同，它通过计算向量间的余弦相似度或欧氏距离来实现语义级搜索。

构建数据库客户端

词向量嵌入使用本地推理模型Jerry0/text2vec-large-chinese
需要安装sentence_transformers

>pip install sentence_transformers

代码 

"""
ChromaDB客户端实现
"""
import os
from chromadb.utils import embedding_functions
from config import settings
from typing import Optional, Dict, List
from chromadb import Client

class ChromaClient:
    def __init__(self,collection_name: str = "cohere_collection"):
        """初始化ChromaDB客户端"""
        model_name =  os.path.dirname(os.path.dirname(os.path.abspath(__file__))) + "/ai-modelscope"
        print("模型路径：", model_name)
        self.custom_ef = embedding_functions.SentenceTransformerEmbeddingFunction(
            model_name=model_name,
            device="cpu"  # 使用 GPU 加速
        )
        self.collection_name= collection_name
        self.client= Client()
        self.collection= self.client.get_or_create_collection(
            name=collection_name,
            embedding_function=self.custom_ef
        )
   
    def turn_collection(self, collection_name: str):
        """切换集合"""
        self.collection_name= collection_name
        self.collection= self.client.get_or_create_collection(
            name=collection_name,
            embedding_function=self.custom_ef
        )
    def insert_documents(self,documents: List[str], metadatas: Optional[List[Dict]] = None, ids: Optional[List[str]] = None) -> None:
        """插入文档到集合"""
        self.collection.add(
            documents=documents,
            metadatas=metadatas,
            ids=ids
        )
    
    def query(self,  query_texts: List[str], n_results: int = 5) -> Dict:
        """查询相似文档"""
        return self.collection.query(
            query_texts=query_texts,
            n_results=n_results
        )

    def __enter__(self):
        """上下文管理器入口"""
        return self

    def __exit__(self, exc_type, exc_val, exc_tb):
        """上下文管理器出口"""
        print("ChromaDB客户端已关闭")


#定义实例获取与释放
def get_chroma_client():
    chroma = ChromaClient()
    try:
        chroma.client.heartbeat()  # 测试连接
        yield chroma
    finally:
        chroma.client.persist()  # 确保数据持久化
        chroma.client.stop()  # 停止客户端

测试客户端使用代码

from core.chroma_client import ChromaClient
import random

def generate_test_documents(num_docs=10, chunks_per_doc=3):
    """生成学术文档测试chunks
    
    生成包含完整学术文档结构的测试数据，包括标题、摘要、章节等内容
    
    Args:
        num_docs: 要生成的文档数量 (默认: 10)
        chunks_per_doc: 每个文档分割的chunk数量 (默认: 3)
        
    Returns:
        tuple: (documents, metadatas, ids)
    """
    documents = []
    metadatas = []
    ids = []
    
    academic_templates = [
        "论文标题: {}\n摘要: 本研究探讨了{}领域的关键问题。通过{}方法，我们发现{}",
        "研究背景: {}\n研究方法: {}\n初步结果: {}",
        "文献综述: {}\n理论框架: {}\n研究假设: {}",
        "数据收集: {}\n分析方法: {}\n研究限制: {}",
        "结论: {}\n研究贡献: {}\n未来方向: {}"
    ]
    
    research_topics = [
        "人工智能伦理", "量子计算", "生物信息学", 
        "气候模型", "神经网络优化", "区块链技术"
    ]
    
    for doc_idx in range(1, num_docs + 1):
        topic = random.choice(research_topics)
        for chunk_idx in range(1, chunks_per_doc + 1):
            # 生成学术文档内容
            template = academic_templates[(chunk_idx-1) % len(academic_templates)]
            content = template.format(
                f"{topic}研究{doc_idx}",
                f"{topic}领域",
                f"{random.choice(['定量','定性','混合'])}研究方法",
                f"显著{random.choice(['正','负'])}相关关系"
            )
            
            documents.append(content)
            
            # 元数据
            metadatas.append({
                "source": f"academic_paper_{doc_idx}",
                "chunk_num": chunk_idx,
                "doc_id": f"paper_{doc_idx}",
                "doc_type": "academic",
                "topic": topic,
                "section": ["title","abstract","method","result","conclusion"][(chunk_idx-1)%5]
            })
            
            ids.append(f"paper_{doc_idx}_section_{chunk_idx}")
    
    return documents, metadatas, ids

def test_chroma_operations():
    """测试ChromaDB的文档chunk和召回功能"""
    # 1. 生成测试数据
    documents, metadatas, ids = generate_test_documents()
    
    # 2. 使用上下文管理器操作数据库
    with ChromaClient("document_chunks_test") as client:
        # 插入文档chunks
        client.insert_documents(
            documents=documents,
            metadatas=metadatas,
            ids=ids
        )
        
        # 测试召回功能 - 查询特定文档的chunks
        test_query = "研究背景: 区块链技术研究"
        print(f"\n查询测试: '{test_query}'")
        results = client.query(
            query_texts=[test_query],
            n_results=5
        )
        
        # 打印查询结果
        print("\n召回结果:")
        for i, doc in enumerate(results['documents'][0]):
            print(f"{i+1}. {doc} (相似度: {results['distances'][0][i]:.4f})\n")
            print(f"元数据: {results['metadatas'][0][i]}\n")

if __name__ == "__main__":
    test_chroma_operations()

输出结果

查询测试: '研究背景: 区块链技术研究'

召回结果:
1. 研究背景: 区块链技术研究1
研究方法: 区块链技术领域
初步结果: 定量研究方法 (相似度: 0.3923)

(----: [[0.39233890175819397, 0.4587123990058899, 0.6135913133621216, 0.8047057390213013, 0.8106489777565002]])

元数据: {'chunk_num': 2, 'doc_id': 'paper_1', 'doc_type': 'academic', 'section': 'abstract', 'source': 'academic_paper_1', 'topic': '区块链技术'}

2. 文献综述: 区块链技术研究1
理论框架: 区块链技术领域
研究假设: 定性研究方法 (相似度: 0.4587)

(----: [[0.39233890175819397, 0.4587123990058899, 0.6135913133621216, 0.8047057390213013, 0.8106489777565002]])

元数据: {'chunk_num': 3, 'doc_id': 'paper_1', 'doc_type': 'academic', 'section': 'method', 'source': 'academic_paper_1', 'topic': '区块链技术'}

3. 论文标题: 区块链技术研究1
摘要: 本研究探讨了区块链技术领域领域的关键问题。通过混合研究方法方法，我们发现显著正相关关系 (相似度: 0.6136)

(----: [[0.39233890175819397, 0.4587123990058899, 0.6135913133621216, 0.8047057390213013, 0.8106489777565002]])

元数据: {'chunk_num': 1, 'doc_id': 'paper_1', 'doc_type': 'academic', 'section': 'title', 'source': 'academic_paper_1', 'topic': '区块链技术'}

4. 研究背景: 量子计算研究2
研究方法: 量子计算领域
初步结果: 定性研究方法 (相似度: 0.8047)

(----: [[0.39233890175819397, 0.4587123990058899, 0.6135913133621216, 0.8047057390213013, 0.8106489777565002]])

元数据: {'chunk_num': 2, 'doc_id': 'paper_2', 'doc_type': 'academic', 'section': 'abstract', 'source': 'academic_paper_2', 'topic': '量子计算'}

5. 研究背景: 量子计算研究6
研究方法: 量子计算领域
初步结果: 定性研究方法 (相似度: 0.8106)

(----: [[0.39233890175819397, 0.4587123990058899, 0.6135913133621216, 0.8047057390213013, 0.8106489777565002]])

元数据: {'chunk_num': 2, 'doc_id': 'paper_6', 'doc_type': 'academic', 'section': 'abstract', 'source': 'academic_paper_6', 'topic': '量子计算'}

ChromaDB客户端已关闭

原理说明

在ChromaDB的查询结果中， distances 数组表示的是查询字符串与每个召回结果之间的余弦距离(Cosine Distance)。具体来说：

1. 余弦距离范围是0到2：
   - 0表示完全相似(余弦相似度为1)
   - 2表示完全不相似(余弦相似度为-1)
   - 1表示正交(余弦相似度为0)
2. 在您的结果中，所有距离值都接近1(0.992到1.105之间)，这表明：
   - 查询字符串与召回结果之间的相似度中等
   - 结果之间区分度不高，可能是因为测试数据模板相似性较高
3. 结果分析：

        print(f"{i+1}. {doc} (距离: {results['distances'][0][i]:.4f})")