Chroma向量数据库

import os
import chromadb
from chromadb.utils.embedding_functions import OpenAIEmbeddingFunction
# 自定义嵌入函数
from chromadb import Documents, EmbeddingFunction, Embeddings
import httpx

# 创建持久化存储客户端
client = chromadb.PersistentClient(path="chroma_data")

# 自定义一个ollama的嵌入函数
class OllamaEmbeddingFunction(EmbeddingFunction):
    def __init__(self, 
         model: str = "dengcao/Qwen3-Embedding-4B:Q5_K_M",
         base_url: str = "http://127.0.0.1:11434",
         timeout: float = 60.0):
        """
        使用 httpx 同步客户端调用 Ollama 的 /api/embeddings
        """
        self.model = model
        self.base_url = base_url.rstrip("/")
        self.timeout = timeout
        # 复用连接，提升性能
        self.client = httpx.Client(timeout=self.timeout)

    def _embed_one(self, text: str) -> list[float]:
        url = f"{self.base_url}/api/embeddings"
        # Ollama 的字段是 "prompt"
        resp = self.client.post(url, json={"model": self.model, "prompt": text})
        resp.raise_for_status()
        data = resp.json()
        print(data)
        # 常见返回是 {"embedding": [...]}
        if "embedding" in data:
            return data["embedding"]
        raise RuntimeError(f"向量结果错误: {data}")

    def __call__(self, input: Documents) -> Embeddings:
        # Chroma 会把 documents 列表传进来
        if not isinstance(input, list):
            input = [input]
        return [self._embed_one(x) for x in input]



# 创建集合，并指定嵌入模型
collection = client.get_or_create_collection(
    name="my_collection",
    # 指定嵌入模型
    embedding_function=OpenAIEmbeddingFunction(
        api_key=os.getenv("OPENAI_API_KEY"),
        model_name="text-embedding-3-small"
    ),

    # # 指定自定义ollama的嵌入模型
    # embedding_function=OllamaEmbeddingFunction(),
    # embedding_function=None, 可以不指定嵌入模型，而是在写入数据的时候直接写入向量
    # 指定集合的元数据
    metadata={
        "description": "my first Chroma collection",
        "created": "2025-08-07"
    }
)

# # 删除集合
# client.delete_collection(name="my-collection")


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

# 直接写入向量
collection.add(
    ids=["id1", "id2", "id3"],
    embeddings=[[1.1, 2.3, 3.2], [4.5, 6.9, 4.4], [1.1, 2.3, 3.2]],
    documents=["doc1", "doc2", "doc3"],
    metadatas=[{"chapter": 3, "verse": 16}, {"chapter": 3, "verse": 5}, {"chapter": 29, "verse": 11}],
)

# 更新数据
collection.upsert(
    ids=["id1"],
    embeddings=[[1.1, 2.3, 3.2]],
    metadatas=[{"chapter": 3, "verse": 16}],
    documents=["doc1"],
)

# 删除数据
collection.delete(
    ids=["id3"],
)

# 查询
collection.query(
    query_embeddings=[[11.1, 12.1, 13.1],[1.1, 2.3, 3.2]], # 查询向量
    n_results=5, # 返回5个结果，不写默认返回10个
    ids=["id1", "id2"], # 指定返回的id
    limit=2, # 指定检索的文档数量，默认100
    offset=1, # 指定返回的文档偏移量，默认0
    where={"page": 10}, # 查询元数据字段 'page' 等于 10 的记录
    where_document={"$contains": "search string"} # 查询记录文档中包含搜索字符串的记录
)

# 查询文本
results = collection.query(
    query_texts=["first query", "second query"],
    include=["documents", "metadatas"], # 指定返回的文档和元数据
    )

# 列表查询，也有where参数
results = collection.get(
    # where={"type": "news"},
    where={"page": { "$gt": 10 }}, # 使用运算符过滤，大于10
)

# uv pip install langchain-chroma langchain-openai
import os
# uv pip install langchain-ollama
from langchain_ollama import OllamaEmbeddings
from langchain_chroma import Chroma
from langchain_core.documents import Document
from uuid import uuid4

# --------- 使用openai嵌入模型 ----------
# os.environ["OPENAI_API_KEY"] = "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" 
# # 初始化模型
# from langchain_openai import OpenAIEmbeddings
# embeddings = OpenAIEmbeddings(model="text-embedding-3-large")

# --------- 使用自定义的ollama嵌入模型 ----------
embeddings = OllamaEmbeddings(
    model="dengcao/Qwen3-Embedding-4B:Q5_K_M",         # 模型名称
    base_url="http://127.0.0.1:11434"  # 设置本地 Ollama 嵌入 API 的地址
)

# 初始化实列
vector_store = Chroma(
    collection_name="example_collection", # 指定集合名称
    embedding_function=embeddings, # 指定模型
    persist_directory="./chroma_langchain_db", # 指定本地路径
    # host="localhost", # 或者连接到远程的ChromaDB
)


# ===========写入数据=============
def add_data():
    document_1 = Document(
        page_content="I had chocolate chip pancakes and scrambled eggs for breakfast this morning.",
        metadata={"type": "tweet"},
        id=1,
    )

    document_2 = Document(
        page_content="The weather forecast for tomorrow is cloudy and overcast, with a high of 62 degrees.",
        metadata={"type": "news"},
        id=2,
    )

    documents = [
        document_1,
        document_2,
    ]

    uuids = [str(uuid4()) for _ in range(len(documents))]
    vector_store.add_documents(documents=documents, ids=uuids)

# ===========更新数据=============
def update_data(uuids):
    updated_document_1 = Document(
        page_content="I had chocolate chip pancakes and fried eggs for breakfast this morning.",
        metadata={"type": "tweet"},
        id=1,
    )

    updated_document_2 = Document(
        page_content="The weather forecast for tomorrow is sunny and warm, with a high of 82 degrees.",
        metadata={"type": "news"},
        id=2,
    )
    # 更新一条数据
    vector_store.update_document(document_id=uuids[0], document=updated_document_1)

    # 更新多条
    documents = [updated_document_1, updated_document_2]
    vector_store.update_documents(
        ids=uuids[:2], documents=documents
    )

# ===========删除数据=============
def delete_data(uuids):
    vector_store.delete(ids=uuids[-1])

# ===========查询数据=============
# 简单相似搜索
def similarity_search(query, k, filter):
    results = vector_store.similarity_search(
        query,
        k=k,
        filter=filter,
    )
    print(results)
    for res in results:
        print(f"* {res.page_content} [{res.metadata}]")

# 按评分相似搜索
def similarity_search_with_score(query, k, filter):
    results = vector_store.similarity_search_with_score(
        query,
        k=k,
        filter=filter,
    )
    for res, score in results:
        print(f"* [SIM={score:3f}] {res.page_content} [{res.metadata}]")

# 传递向量搜索
def similarity_search_by_vector(query, k):
    results = vector_store.similarity_search_by_vector(
        embedding=embeddings.embed_query(query),
        k=k,
    )
    for doc in results:
        print(f"* {doc.page_content} [{doc.metadata}]")

# ============列出数据=============

def list_data():
    raw = vector_store.get(
        # include=["documents", "metadatas"]   # 选填需要哪些字段
    )
    # all_docs = list(zip(raw["ids"], raw["documents"], raw["metadatas"]))
    print(raw)
    return raw

def list_by_page(
    page:int = 1,
    page_size:int = 20,
    order:str = "desc",        # "asc" / "desc"
    where:dict|None = None,    # 过滤条件
):
    offset = (page - 1) * page_size

    # ① 直接分页取回来，只要 documents / metadatas；ids 会自动返回
    batch = vector_store.get(
        where=where,
        limit=page_size,
        offset=offset,
        include=["documents", "metadatas"],   # ids 默认带回
    )

    if not batch["ids"]:                     # 本页没有数据
        return batch

    # ② 用索引列表做排序，完全不触碰 zip
    idx_sorted = sorted(
        range(len(batch["metadatas"])),
        key=lambda i: batch["metadatas"][i].get("created_at", 0),
        reverse=(order == "desc"),
    )

    # ③ 根据排序后的索引重新排列三个数组
    for key in ("ids", "documents", "metadatas"):
        batch[key] = [batch[key][i] for i in idx_sorted]
    print(batch)
    return batch    # 结构仍保持 Chroma 的 get 返回格式

# 按字段过滤(例如某个字段为类型)
def list_by_type(doc_type:str, page:int=1, page_size:int=20):
    where = {"type": doc_type}        # 比如 "tweet" / "news"
    offset = (page - 1) * page_size
    batch = vector_store.get(
        where=where,                  # where 过滤
        limit=page_size,
        offset=offset,
        include=["documents", "metadatas"]  # 移除 "ids"，因为 ids 是默认返回的
    )
    print(batch)
    return batch

# 测试
if __name__ == "__main__":
    # add_data()
    # list_by_page()
    # list_by_type("news")
    print(embeddings.embed_query("你好，世界"))