可以用chromadb自带的api,也可以用langchain支持的,封装程度更高,更加结构化。

初始化Chroma

首先要知道初始化chroma是做什么用的。
相当于提供一个库,然后后续调用方法、根据传入的参数进行查询。
所以初始化,相当于提供源数据(非提问数据)。

embedding是向量化的工具,可以是本地大模型,也可以是在线向量化接口。

初始化Chroma-1. 本地持久化(最常用)
from langchain_chroma import Chroma
from langchain_openai import OpenAIEmbeddings

embedding = OpenAIEmbeddings()

# 初始化:数据存在 ./chroma_db 目录
db = Chroma(
    embedding_function=embedding,
    persist_directory="./chroma_db",  # 本地持久化
)

1和4的区别就是如果不指定集合,那么就用默认集合。
langchain Chroma 会默认collection,叫做langchain
原生chroma 无默认值,必须指定。

建议最好显示的指定默认值。

初始化Chroma-2. 从Documents直接创建 + 持久化

代码:

db = Chroma.from_documents(
    documents=docs,
    embedding=embedding,
    persist_directory="./chroma_db",
)

当然也可以加集合:
db = Chroma.from_documents(
    documents=docs,
    embedding=embedding,
    persist_directory="./chroma_db",
    collection_name="my_collection"  # ✅ 直接加在这里
)

注:这里容易有疑问,documents已经提供了源数据,为什么还需要数据库呢?

是这样的,documents文本并不支持向量查询,会先存到库里,然后再向量化查询。

初始化Chroma-3. 从 texts + metadatas 创建

代码:

texts = ["内容1", "内容2"]
metadatas = [{"source": "A"}, {"source": "B"}]

db = Chroma.from_texts(
    texts=texts,
    metadatas=metadatas,
    embedding=embedding,
    persist_directory="./chroma_db"
)

注:documents就相当于texts和metadatas,只不过换了一种方式而已。
对照关系:

Document(
    page_content="这是文本内容 → 对应 texts",  # ✅ 就是 texts
    metadata={"source": "xxx", "author": "yyy"}   # ✅ 就是 meta
)
初始化Chroma-4. 连接已存在的 Chroma 集合(指定 collection_name)
db = Chroma(
    embedding_function=embedding,
    persist_directory="./chroma_db",
    collection_name="user_123_docs",  # 隔离不同数据
)
初始化Chroma-5. 内存模式(临时,不保存)
db = Chroma(
    embedding_function=embedding,
    # 不写 persist_directory 就是内存模式
)
persist_directory参数可以不填吗?

可以,如果不填,并不会有默认的文件路径,而是表示内存模式,不会保存

初始化Chroma-6. 连接远程 Chroma 服务器(生产环境)
import chromadb

client = chromadb.HttpClient(
    host="192.168.1.100",
    port=8000
)

db = Chroma(
    client=client,  # 使用远程客户端
    collection_name="my_collection",
    embedding_function=embedding
)
查询(推荐版)

代码:

from fastapi import APIRouter, HTTPException
from pydantic import BaseModel
from typing import Optional, List, Dict, Any
from core.chroma_pool import ChromaPool
from schemas.knowledge_base import KnowledgeAddRequest, KnowledgeQueryRequest
# 引入 LangChain Chroma 和嵌入函数(根据实际场景替换)
from langchain_chroma import Chroma
from langchain.embeddings.base import Embeddings
from langchain_core.documents import Document

# 初始化 Chroma 连接池
chromaPool = ChromaPool()
router = APIRouter()


# --- 核心工具函数:统一格式化输出 ---
def format_results(docs: List[Document], scores: List[float] = None) -> Dict:
    """
    将 LangChain Document 列表转换为原代码所需的格式
    """
    return {
        "ids": [doc.id for doc in docs],
        "documents": [doc.page_content for doc in docs],
        "metadatas": [doc.metadata for doc in docs],
        "distances": scores if scores is not None else [] # 如果没有分数传空列表
    }

@router.post("/knowledge_base/query")
def query(req: KnowledgeQueryRequest):
    try:
        chroma = chromaPool.get_db("knowledge_base")

        # 1. 预处理查询词
        clean_query = req.query.strip() if req.query else ""
        is_valid_query = bool(clean_query)

        final_docs = []
        final_scores = None  # 用于存储 distances

        if not is_valid_query:
            # ==========================
            # 🟢 场景 A: 无查询词 -> 走原生 get (查全部/纯筛选)
            # 注意:原生 get 没有距离分数的概念,所以 distances 为空
            # ==========================
            raw_results = chroma._collection.get(
                where=req.filter,
                limit=req.top_k,
                include=['documents', 'metadatas']
            )
            # 手动转为 Document 对象以便统一处理
            final_docs = [
                Document(page_content=doc, metadata=meta)
                for doc, meta in zip(raw_results['documents'], raw_results['metadatas'])
            ]

        else:
            # ==========================
            # 🔵 场景 B: 有查询词 -> 走带分数的相似度搜索
            # 这里不使用 retriever.invoke() 因为它会丢弃分数
            # 直接使用 vectorstore 的方法并传入 filter,效果和 as_retriever 配置 filter 是一样的
            # ==========================
            results_with_score = chroma.similarity_search_with_score(
                query=clean_query,
                k=req.top_k,
                filter=req.filter  # 直接支持 filter 字典
            )

            # 解包结果 (LangChain 返回的是 [(doc, score), ...])
            final_docs = [res[0] for res in results_with_score]
            final_scores = [float(res[1]) for res in results_with_score]

        # 2. 统一返回格式
        formatted_result = format_results(final_docs, final_scores)

        return {"code": 200, "msg": "成功", "data": formatted_result}

    except Exception as e:
        raise HTTPException(status_code=500, detail=f"查询失败:{str(e)}")

@router.post("/knowledge_base/query_all")
def query_all():
    try:
        chroma = chromaPool.get_db("knowledge_base")

        # ✅ 正确写法:include 里 不要写 ids!
        all_data = chroma.get(include=["documents", "metadatas"])

        # 格式适配(ids chroma 会自动返回)
        formatted_result = {
            "ids": all_data["ids"],  # 自带,不用写进 include
            "documents": all_data["documents"],
            "metadatas": all_data["metadatas"],
            "distances": []  # 无分数,返回空数组
        }
        return {"code": 200, "msg": "成功", "data": formatted_result}
    except Exception as e:
        raise HTTPException(status_code=500, detail=f"查询全部失败:{str(e)}")

similarity_search()相似度查询

代码:

docs = vector_store.similarity_search(query_text, k=2)

query_text # 表示输入的文本,例如 今天天气怎么样?
k # 返回的条数,默认是4条

similarity_search()的实现原理

伪代码:

def similarity_search(query):
    # 👇 这里就是实时转换!
    query_vector = embedding.embed_query(query) # 先转换为向量

    # 👇 去库里面找最相似的
    results = chroma.search(query_vector) # 再去库里根据向量查询
    return results
similarity_search() 根据metadata元数据查询

当然支持,filter字段就是配置元数据查询的。
如果配置了filter,filter是在向量前查询的。

json中只有一个字段:
docs = db.similarity_search(
    query="你的查询内容",
    filter={"user_id": "1001"}
)

json中有多个字段(and关系):
docs = db.similarity_search(
    query="查询内容",
    filter={
        "user_id": "1001",
        "source": "文档A",
        "status": "active"
    }
)

运算符查询:
docs = db.similarity_search(
    query="查询内容",
    filter={
        "user_id": "1001",
        "create_time": {"$gt": "2025-01-01"}
    }
)
支持的运算符
filter={
    "age": {"$gt": 18},      # 大于
    "age": {"$lt": 60},      # 小于
    "age": {"$gte": 18},     # 大于等于
    "age": {"$lte": 60},     # 小于等于
    "name": {"$ne": "张三"}, # 不等于
    "tag": {"$in": ["重要", "紧急"]},  # 在列表里
    "tag": {"$nin": ["过期"]}        # 不在列表里
}
根据向量直接查询?

支持,用similarity_search_by_vector()方法即可,这个专门用来查向量的。

生产中这种场景极少,一般都是传文本,后台自动向量化。只有一种情况,即前端传递少量向量化数值到后台。

入参类:

class KnowledgeQueryByVectorRequest(BaseModel):
    vectorList: List[float] = Field(..., description="查询向量数组")
    top_k: int = Field(3, ge=1, le=50)
    filter: Optional[Dict[str, Any]] = None

接口方法(路由方法):

@app.post("/query_by_vector")
def query_by_vector(req: KnowledgeQueryByVectorRequest):
    docs = vectorstore.similarity_search_by_vector(
        embedding=req.vectorList,
        k=req.top_k,
        filter=req.filter
    )
    return docs
支持只根据元数据查询,不根据向量查询吗?

similarity_search()方法是不支持的,但是原生的chroma(或者通过Chroma获取到的原生chroma是支持的),代码:

# 原生 Chroma 接口
results = db._collection.get(
    where={"user_id": "1001"},  # 元数据过滤
    include=["documents", "metadatas"]
)
print(results["documents"])
print(results["metadatas"])
根据id查询get_by_ids()
docs = db.get_by_ids(["id_123", "id_456"])  

查单条:
doc = db.get_by_ids(["my_id"])[0]

返回的是Document对象。

注:get_by_ids()是精确查询,并不走向量,所以速度极快,相当于索引。

如果还想用chroma的客户端怎么办?

langchain是支持原生客户端的。

collection = db._collection # 获取原生客户端。
# 根据 ID 查询(支持返回向量)
result = collection.get(
    ids=["id_123"],
    include=["documents", "metadatas", "embeddings"]
)
完整例子-用于参考
import os
from langchain_chroma import Chroma
from langchain_community.embeddings import DashScopeEmbeddings
from langchain_core.documents import Document

# 1. 配置 API Key (请替换为你的真实 Key,或者在环境变量中设置)
# 建议 export DASHSCOPE_API_KEY="sk-xxxxx"


API_KEY = os.getenv("DASHSCOPE_API_KEY", "YOUR API KEY")
def main():
    print("🚀 正在初始化阿里云百炼 Embedding 模型...")

    # 2. 初始化 Embedding 模型
    # model="text-embedding-v4" 是阿里最新的模型,效果好,默认维度 1024
    embeddings = DashScopeEmbeddings(
        model="text-embedding-v4",
        dashscope_api_key=API_KEY
    )

    # 3. 准备数据
    documents = [
        Document(page_content="阿里云百炼提供了强大的文本向量化能力。", metadata={"type": "tech"}),
        Document(page_content="ChromaDB 是一个非常适合本地开发的向量数据库。", metadata={"type": "tech"}),
        Document(page_content="大熊猫是中国的国宝,主要吃竹子。", metadata={"type": "animal"}),
        Document(page_content="Python 是一门非常流行的高级编程语言。", metadata={"type": "code"}),
    ]

    print("\n📚 正在向量化并存入 Chroma 数据库...")
    # 4. 构建数据库
    # 注意:这里直接使用 from_documents,LangChain 会自动调用百炼接口把文本转成向量
    vector_store = Chroma.from_documents(
        documents=documents,
        embedding=embeddings,
        persist_directory="./chroma_db_bailian"  # 指定本地存储路径
    )

    print("\n🔍 开始查询测试...")

    # --- 查询 1: 语义搜索 ---
    query_text = "怎么把数据存进数据库?"
    print(f"用户提问: {query_text}")

    # similarity_search 会自动调用百炼接口把问题变成向量,然后去库里找
    docs = vector_store.similarity_search(query_text, k=2)

    for i, doc in enumerate(docs):
        print(f"[{i + 1}] 内容: {doc.page_content} (来源: {doc.metadata['type']})")

    # --- 查询 2: 带分数的搜索 ---
    print("\n📊 相似度分数详情:")
    docs_with_score = vector_store.similarity_search_with_score(query_text, k=2)

    for doc, score in docs_with_score:
        # Chroma 默认使用 L2 距离,分数越小越相似
        print(f"内容: {doc.page_content} -> 距离分数: {score:.4f}")


if __name__ == "__main__":
    try:
        main()
    except Exception as e:
        print(f"❌ 发生错误: {e}")
        print("💡 提示: 请检查你的 DASHSCOPE_API_KEY 是否正确,以及网络连接是否正常。")
Logo

AtomGit 是由开放原子开源基金会联合 CSDN 等生态伙伴共同推出的新一代开源与人工智能协作平台。平台坚持“开放、中立、公益”的理念,把代码托管、模型共享、数据集托管、智能体开发体验和算力服务整合在一起,为开发者提供从开发、训练到部署的一站式体验。

更多推荐