ChromaDB: локальная векторная БД для разработки

pip install chromadb

# Для работы с OpenAI-эмбеддерами
pip install chromadb openai

# Для sentence-transformers (локальные модели)
pip install chromadb sentence-transformers

import chromadb

client = chromadb.PersistentClient(path="./chroma_db")

# ── Три способа открыть коллекцию ────────────────────────────────

# 1. Создать новую (ошибка если уже существует)
col = client.create_collection("docs")

# 2. Открыть существующую (ошибка если не существует)
col = client.get_collection("docs")

# 3. Открыть или создать — самый частый вариант
col = client.get_or_create_collection("docs")

# ── Параметры при создании ────────────────────────────────────────

col = client.get_or_create_collection(
    name="docs",
    metadata={
        # Метрика расстояния (нельзя изменить после создания!)
        "hnsw:space": "cosine",         # "cosine" | "l2" | "ip"

        # HNSW-параметры (тоже нельзя изменить потом)
        "hnsw:M": 16,                   # связность графа (по умолчанию 16)
        "hnsw:construction_ef": 100,    # точность построения (по умолчанию 100)
        "hnsw:search_ef": 100,          # точность поиска (по умолчанию 10!)
        "hnsw:num_threads": 4,          # потоки индексирования
    },
)

# ── Управление коллекциями ────────────────────────────────────────

# Список всех коллекций
collections = client.list_collections()
print([c.name for c in collections])

# Информация о коллекции
print(col.name)          # "docs"
print(col.metadata)      # {'hnsw:space': 'cosine', ...}
print(col.count())       # количество документов

# Удалить коллекцию (вместе с данными!)
client.delete_collection("old_docs")

import chromadb
import numpy as np

# ── Создание коллекции с правильными параметрами ─────────────────

client = chromadb.PersistentClient(path="./chroma_db")

col = client.get_or_create_collection(
    name="production_docs",
    metadata={
        "hnsw:space": "cosine",
        "hnsw:M": 16,
        "hnsw:construction_ef": 200,   # точный индекс при добавлении
        "hnsw:search_ef": 200,         # точный поиск при запросах
    },
)

# ── Измерение реального recall ────────────────────────────────────

def measure_recall_at_k(
    col,
    query_vecs: np.ndarray,
    ground_truth_ids: list[list[str]],
    k: int = 10,
    ef_search_values: list[int] = [10, 50, 100, 200, 500],
) -> dict:
    """
    Измерить recall@k для разных значений ef_search.
    ground_truth_ids[i] = список ID релевантных документов для запроса i.
    """
    results = {}
    for ef in ef_search_values:
        hits = 0
        total = 0
        for qvec, gt_ids in zip(query_vecs, ground_truth_ids):
            res = col.query(
                query_embeddings=[qvec.tolist()],
                n_results=k,
                include=["distances"],
                # Можно переопределить ef прямо в query:
                # (работает в Chroma >= 0.4.10)
                # query_params={"ef": ef},  # если ваша версия поддерживает
            )
            returned_ids = set(res["ids"][0])
            hits += len(returned_ids & set(gt_ids))
            total += len(gt_ids)
        results[ef] = hits / total if total > 0 else 0.0
        print(f"ef={ef:4d}: recall@{k} = {results[ef]:.3f}")
    return results

import chromadb
from datetime import datetime

client = chromadb.PersistentClient(path="./chroma_db")
col = client.get_or_create_collection("docs", metadata={"hnsw:space": "cosine"})

# ── ADD: добавить документы ───────────────────────────────────────
col.add(
    ids=["doc_001", "doc_002", "doc_003"],
    documents=[
        "FastAPI — async веб-фреймворк для Python",
        "Django — полнофункциональный фреймворк с ORM",
        "Flask — микрофреймворк для небольших приложений",
    ],
    metadatas=[
        {"source": "docs", "category": "web", "year": 2024, "tokens": 42},
        {"source": "docs", "category": "web", "year": 2023, "tokens": 38},
        {"source": "blog", "category": "web", "year": 2022, "tokens": 35},
    ],
    # embeddings=[...],  # если не передать — Chroma эмбеддирует сам
)

# ── UPSERT: добавить или обновить (рекомендуется) ─────────────────
col.upsert(
    ids=["doc_001"],
    documents=["FastAPI: быстрый async REST-фреймворк (обновлено)"],
    metadatas=[{"source": "docs", "category": "web", "year": 2024, "tokens": 48}],
)

# ── UPDATE: частичное обновление метаданных ───────────────────────
col.update(
    ids=["doc_002"],
    metadatas=[{"source": "docs", "category": "web", "year": 2024, "tokens": 38}],
    # documents и embeddings не передаём — они не изменятся
)

# ── GET: получить без поиска ──────────────────────────────────────
# По конкретным IDs
result = col.get(ids=["doc_001", "doc_003"])
print(result["documents"])   # ['FastAPI...', 'Flask...']
print(result["metadatas"])   # [{'source': 'docs', ...}, ...]

# По фильтру
result = col.get(
    where={"category": "web", "year": {"$gte": 2023}},
    include=["documents", "metadatas"],
    limit=10,
    offset=0,   # для пагинации
)

# ── DELETE: удалить ───────────────────────────────────────────────
col.delete(ids=["doc_003"])

# Удалить всё из источника (осторожно!)
col.delete(where={"source": "blog"})

print(f"Документов в коллекции: {col.count()}")

# ── Батчевое добавление (важно для больших корпусов) ──────────────
def batch_upsert(col, documents: list[dict], batch_size: int = 500):
    """
    documents: список словарей {'id', 'text', 'metadata', 'embedding'}
    Chroma падает или замедляется при батчах > 5000.
    500–1000 — оптимальный размер.
    """
    for i in range(0, len(documents), batch_size):
        batch = documents[i : i + batch_size]
        col.upsert(
            ids=[d["id"] for d in batch],
            documents=[d["text"] for d in batch],
            metadatas=[d["metadata"] for d in batch],
            embeddings=[d["embedding"] for d in batch],
        )
        print(f"\rИндексировано: {min(i + batch_size, len(documents))}/{len(documents)}", end="")
    print()

import chromadb
import numpy as np

client = chromadb.PersistentClient(path="./chroma_db")
col = client.get_collection("docs")

# ── Запрос по тексту (Chroma эмбеддирует сам) ─────────────────────
results = col.query(
    query_texts=["как установить FastAPI"],
    n_results=5,
)

# ── Запрос по готовым эмбеддингам (рекомендуется в production) ────
query_embedding = embed_query("как установить FastAPI")  # ваша функция

results = col.query(
    query_embeddings=[query_embedding.tolist()],   # список векторов!
    n_results=5,
)

# ── Батчевый запрос: несколько запросов за раз ────────────────────
results = col.query(
    query_embeddings=[
        embed_query("установка FastAPI").tolist(),
        embed_query("Django ORM запросы").tolist(),
    ],
    n_results=3,  # топ-3 для каждого запроса
)
# results["ids"][0] → топ-3 для первого запроса
# results["ids"][1] → топ-3 для второго запроса

# ── Структура результата ──────────────────────────────────────────
print(results.keys())
# dict_keys(['ids', 'distances', 'metadatas', 'embeddings', 'documents', 'uris', 'data'])

# Для одного запроса:
print(results["ids"][0])         # ['doc_001', 'doc_002', 'doc_003']
print(results["distances"][0])   # [0.12, 0.34, 0.45] — меньше = ближе (для cosine: 1-cosine)
print(results["documents"][0])   # ['FastAPI...', 'Django...', ...]
print(results["metadatas"][0])   # [{'source': 'docs', ...}, ...]

# Только то, что нужно — меньше трафика, быстрее
results = col.query(
    query_embeddings=[query_vec.tolist()],
    n_results=10,
    include=["documents", "metadatas", "distances"],  # без embeddings — они тяжёлые
)

# Получить векторы (например, для дополнительной обработки)
results = col.query(
    query_embeddings=[query_vec.tolist()],
    n_results=5,
    include=["documents", "metadatas", "distances", "embeddings"],
)
# results["embeddings"][0] → list of 5 vectors, каждый 1536-dim

# Только IDs и distances (минимальная нагрузка — для reranking-пайплайна)
results = col.query(
    query_embeddings=[query_vec.tolist()],
    n_results=50,                           # широкий поиск
    include=["ids", "distances"],           # только для reranker
)
ids_for_reranking = results["ids"][0]

import chromadb

client = chromadb.PersistentClient(path="./chroma_db")
col = client.get_collection("docs")

# ── Простые фильтры ───────────────────────────────────────────────

# Только из источника "wiki"
results = col.query(
    query_embeddings=[query_vec.tolist()],
    n_results=5,
    where={"source": "wiki"},  # краткая форма $eq
)

# Только свежие документы
results = col.query(
    query_embeddings=[query_vec.tolist()],
    n_results=5,
    where={"year": {"$gte": 2023}},
)

# ── Диапазон ──────────────────────────────────────────────────────

# Документы 2022–2024 года с небольшим количеством токенов
results = col.query(
    query_embeddings=[query_vec.tolist()],
    n_results=5,
    where={
        "$and": [
            {"year": {"$gte": 2022}},
            {"year": {"$lte": 2024}},
            {"tokens": {"$lte": 400}},
        ]
    },
)

# ── Несколько источников ($in) ─────────────────────────────────────

results = col.query(
    query_embeddings=[query_vec.tolist()],
    n_results=10,
    where={"source": {"$in": ["wiki", "docs", "github"]}},
)

# ── Исключить категорию ($nin) ─────────────────────────────────────

results = col.query(
    query_embeddings=[query_vec.tolist()],
    n_results=5,
    where={"category": {"$nin": ["deprecated", "draft"]}},
)

# ── Комбинация $and + $or ─────────────────────────────────────────

# (source IN [wiki, docs]) AND (year >= 2023) AND (tokens <= 500)
results = col.query(
    query_embeddings=[query_vec.tolist()],
    n_results=10,
    where={
        "$and": [
            {"source": {"$in": ["wiki", "docs"]}},
            {"year": {"$gte": 2023}},
            {"tokens": {"$lte": 500}},
        ]
    },
)

# ── $or: из разных источников разных лет ─────────────────────────

results = col.query(
    query_embeddings=[query_vec.tolist()],
    n_results=10,
    where={
        "$or": [
            {"$and": [{"source": "wiki"}, {"year": {"$gte": 2024}}]},
            {"$and": [{"source": "docs"}, {"verified": True}]},
        ]
    },
)

# ── where_document: текстовый фильтр по документу ────────────────

# Только документы, содержащие точную строку
results = col.query(
    query_embeddings=[query_vec.tolist()],
    n_results=5,
    where_document={"$contains": "FastAPI"},
)

# НЕ содержит — исключить документы с ошибками/deprecated
results = col.query(
    query_embeddings=[query_vec.tolist()],
    n_results=5,
    where_document={"$not_contains": "deprecated"},
)

# ── Комбинация where + where_document ────────────────────────────

# Из доверенных источников И содержащие нужный термин
results = col.query(
    query_embeddings=[query_vec.tolist()],
    n_results=10,
    where={"source": {"$in": ["wiki", "docs"]}},
    where_document={"$contains": "async"},
)

# ── Практика: точный поиск по артикулу или UUID ───────────────────
# Embedding-модели плохо работают с кодами, UUID, артикулами —
# они «размываются» в семантическом пространстве.
# Решение: сохранять артикулы в метаданных + фильтровать через where.

results = col.query(
    query_embeddings=[query_vec.tolist()],
    n_results=5,
    where={"product_id": "SKU-20481"},    # точное совпадение через SQLite
)

from dataclasses import dataclass
from typing import Optional
import chromadb


@dataclass
class RetrievedDoc:
    id: str
    document: str
    metadata: dict
    distance: float

    @property
    def cosine_similarity(self) -> float:
        """Chroma возвращает distance = 1 - cosine_sim для метрики cosine."""
        return 1.0 - self.distance


def retrieve(
    col,
    query_embedding: list[float],
    n_results: int = 10,
    max_distance: float = 0.4,   # порог: distance <= 0.4 → cosine_sim >= 0.6
    where: Optional[dict] = None,
    where_document: Optional[dict] = None,
) -> list[RetrievedDoc]:
    """
    Поиск с отсечением нерелевантных результатов.

    max_distance=0.4 при metric=cosine означает:
      cosine_similarity = 1 - 0.4 = 0.6 — минимальная релевантность.
    """
    # Запрашиваем больше, чем нужно — компенсация за фильтрацию
    raw_n = min(n_results * 3, col.count()) if col.count() > 0 else n_results

    kwargs: dict = {
        "query_embeddings": [query_embedding],
        "n_results": raw_n,
        "include": ["documents", "metadatas", "distances"],
    }
    if where:
        kwargs["where"] = where
    if where_document:
        kwargs["where_document"] = where_document

    res = col.query(**kwargs)

    docs = []
    for doc_id, doc_text, meta, dist in zip(
        res["ids"][0],
        res["documents"][0],
        res["metadatas"][0],
        res["distances"][0],
    ):
        if dist > max_distance:
            break  # результаты отсортированы по distance → дальше только хуже
        docs.append(RetrievedDoc(
            id=doc_id,
            document=doc_text,
            metadata=meta,
            distance=dist,
        ))
        if len(docs) >= n_results:
            break

    return docs


# ── Пример использования ──────────────────────────────────────────
client = chromadb.PersistentClient(path="./chroma_db")
col = client.get_collection("docs")

query_vec = embed_query("установка async фреймворка")  # ваша функция

docs = retrieve(
    col,
    query_embedding=query_vec.tolist(),
    n_results=5,
    max_distance=0.35,             # строгий порог: cosine_sim >= 0.65
    where={"source": {"$in": ["wiki", "docs"]}},
)

if not docs:
    print("Релевантных документов не найдено — отвечаем 'не знаю'")
else:
    context = "\n\n".join(
        f"[{d.metadata.get('source', '?')}] {d.document}" for d in docs
    )
    print(f"Найдено {len(docs)} документов, топ: cosine={docs[0].cosine_similarity:.3f}")

import chromadb
from chromadb.utils.embedding_functions import (
    OpenAIEmbeddingFunction,
    SentenceTransformerEmbeddingFunction,
    DefaultEmbeddingFunction,   # all-MiniLM-L6-v2, бесплатно
)

client = chromadb.PersistentClient(path="./chroma_db")

# ── Встроенные embedding functions ───────────────────────────────

# OpenAI (нужен API key)
openai_ef = OpenAIEmbeddingFunction(
    api_key="sk-...",
    model_name="text-embedding-3-small",
)
col_openai = client.get_or_create_collection(
    "docs_openai",
    embedding_function=openai_ef,
    metadata={"hnsw:space": "cosine"},
)

# Sentence-transformers (локально)
st_ef = SentenceTransformerEmbeddingFunction(
    model_name="intfloat/multilingual-e5-large",
    device="cpu",
)
col_local = client.get_or_create_collection(
    "docs_local",
    embedding_function=st_ef,
    metadata={"hnsw:space": "cosine"},
)

# Default (all-MiniLM-L6-v2) — только для прототипов
col_default = client.get_or_create_collection("docs_dev")
# Без embedding_function — используется DefaultEmbeddingFunction

# ── Кастомная embedding function ──────────────────────────────────
from chromadb import EmbeddingFunction, Embeddings
import numpy as np
from openai import AsyncOpenAI


class BGEEmbeddingFunction(EmbeddingFunction):
    """
    Кастомный эмбеддер для BGE-M3.
    Наследуемся от EmbeddingFunction и реализуем __call__.
    """

    def __init__(self, device: str = "cpu"):
        from FlagEmbedding import BGEM3FlagModel
        self.model = BGEM3FlagModel("BAAI/bge-m3", use_fp16=True)

    def __call__(self, input: list[str]) -> Embeddings:
        """input — список текстов. Возвращаем список list[float]."""
        output = self.model.encode(
            input,
            return_dense=True,
            return_sparse=False,
        )
        vecs = output["dense_vecs"]  # np.ndarray (N, 1024)
        return vecs.tolist()


col_bge = client.get_or_create_collection(
    "docs_bge",
    embedding_function=BGEEmbeddingFunction(),
    metadata={
        "hnsw:space": "cosine",
        "hnsw:M": 16,
        "hnsw:search_ef": 200,
    },
)

# Теперь add/query работают без явных embeddings:
col_bge.add(
    ids=["1", "2"],
    documents=["FastAPI setup", "Django intro"],
    metadatas=[{"source": "docs"}, {"source": "docs"}],
)
# Chroma вызовет BGEEmbeddingFunction автоматически

"""
pip install langchain-chroma langchain-openai
"""

from langchain_chroma import Chroma
from langchain_openai import OpenAIEmbeddings
from langchain.schema import Document

embeddings = OpenAIEmbeddings(model="text-embedding-3-small")

# ── Создание/подключение ──────────────────────────────────────────

# Создать и заполнить из списка документов
vectorstore = Chroma.from_documents(
    documents=[
        Document(page_content="FastAPI async веб-фреймворк", metadata={"source": "docs"}),
        Document(page_content="Django полный стек", metadata={"source": "docs"}),
    ],
    embedding=embeddings,
    persist_directory="./chroma_db",
    collection_name="docs",
    collection_metadata={"hnsw:space": "cosine", "hnsw:search_ef": 200},
)

# Подключиться к существующей
vectorstore = Chroma(
    persist_directory="./chroma_db",
    collection_name="docs",
    embedding_function=embeddings,
    collection_metadata={"hnsw:space": "cosine"},
)

# ── Поиск ─────────────────────────────────────────────────────────

# Similarity search (топ-k)
docs = vectorstore.similarity_search("как установить FastAPI", k=5)

# С оценками (distance)
docs_with_scores = vectorstore.similarity_search_with_score("async фреймворк", k=5)
for doc, score in docs_with_scores:
    print(f"distance={score:.3f}: {doc.page_content[:60]}")

# С фильтрами метаданных
docs = vectorstore.similarity_search(
    "ORM запросы",
    k=5,
    filter={"source": "docs", "year": {"$gte": 2023}},
)

# ── as_retriever() — для Chain/Agent ─────────────────────────────

# Базовый retriever (top-k)
retriever = vectorstore.as_retriever(
    search_type="similarity",
    search_kwargs={"k": 5},
)

# С порогом сходства
retriever = vectorstore.as_retriever(
    search_type="similarity_score_threshold",
    search_kwargs={
        "score_threshold": 0.6,   # cosine_similarity >= 0.6 (не distance!)
        "k": 5,
    },
)

# MMR — максимальная маргинальная релевантность
# Balances relevance vs diversity: убирает дубликаты из результатов
retriever = vectorstore.as_retriever(
    search_type="mmr",
    search_kwargs={
        "k": 5,
        "fetch_k": 20,     # изначально берём 20, потом MMR отбирает 5
        "lambda_mult": 0.5, # 0=максимальное разнообразие, 1=максимальная релевантность
    },
)

# ── RetrievalQA chain ─────────────────────────────────────────────
from langchain.chains import RetrievalQA
from langchain_openai import ChatOpenAI

llm = ChatOpenAI(model="gpt-4o-mini")
qa_chain = RetrievalQA.from_chain_type(
    llm=llm,
    retriever=retriever,
    return_source_documents=True,
)

answer = qa_chain.invoke({"query": "Что такое FastAPI?"})
print(answer["result"])
print(answer["source_documents"])

# ── Запуск через Docker ───────────────────────────────────────────
docker run -d \
  -p 8000:8000 \
  -v $(pwd)/chroma_data:/chroma/chroma \
  chromadb/chroma:latest

# ── Или через chroma CLI ──────────────────────────────────────────
pip install chromadb
chroma run --path ./chroma_data --port 8000

import chromadb
from chromadb.config import Settings

# ── Подключение к серверу ─────────────────────────────────────────
client = chromadb.HttpClient(
    host="localhost",
    port=8000,
    settings=Settings(
        chroma_client_auth_provider="chromadb.auth.token.TokenAuthClientProvider",
        chroma_client_auth_credentials="your-token",  # если настроена авторизация
    ),
)

# ── Тот же API, другой бэкенд ─────────────────────────────────────
col = client.get_or_create_collection("docs")
col.add(ids=["1"], documents=["..."], metadatas=[{"source": "test"}])
results = col.query(query_texts=["запрос"], n_results=3)

# ── Асинхронный клиент (Chroma >= 0.5) ───────────────────────────
import asyncio
import chromadb

async def async_example():
    client = await chromadb.AsyncHttpClient(host="localhost", port=8000)
    col = await client.get_or_create_collection("docs")

    await col.add(
        ids=["1"],
        documents=["FastAPI tutorial"],
        metadatas=[{"source": "docs"}],
    )

    results = await col.query(
        query_texts=["установка FastAPI"],
        n_results=5,
    )
    return results

asyncio.run(async_example())

import chromadb
from chromadb.config import Settings
import os

# ── Отключить телеметрию ──────────────────────────────────────────
# По умолчанию Chroma отправляет анонимную телеметрию в Posthog.
# Отключается через переменную окружения или Settings:

os.environ["ANONYMIZED_TELEMETRY"] = "False"

client = chromadb.PersistentClient(
    path="./chroma_db",
    settings=Settings(anonymized_telemetry=False),
)

# ── Параметры коллекции для production ───────────────────────────
COLLECTION_CONFIG = {
    "hnsw:space": "cosine",
    "hnsw:M": 16,
    "hnsw:construction_ef": 200,  # плотный индекс при добавлении
    "hnsw:search_ef": 200,        # точный поиск (не дефолтные 10!)
    "hnsw:num_threads": 4,        # потоки для добавления
}

col = client.get_or_create_collection("docs", metadata=COLLECTION_CONFIG)

# ── Паттерн: изолированные коллекции на тенант ───────────────────

def get_tenant_collection(client, tenant_id: str):
    """Отдельная коллекция на каждого пользователя/клиента."""
    return client.get_or_create_collection(
        name=f"docs_{tenant_id}",
        metadata={**COLLECTION_CONFIG},
    )

# ── Паттерн: версионирование коллекций ───────────────────────────

def reindex_collection(client, old_name: str, new_name: str, new_embed_fn):
    """
    Переиндексация при смене модели.
    1. Создаём новую коллекцию
    2. Переиндексируем все документы
    3. Переименовывание: не поддерживается в Chroma —
       меняем имя в конфиге приложения, старую удаляем
    """
    old_col = client.get_collection(old_name)
    new_col = client.get_or_create_collection(new_name, metadata=COLLECTION_CONFIG)

    # Получаем все данные батчами
    batch_size = 500
    offset = 0
    while True:
        data = old_col.get(limit=batch_size, offset=offset,
                           include=["documents", "metadatas"])
        if not data["ids"]:
            break

        new_embeddings = new_embed_fn(data["documents"])
        new_col.upsert(
            ids=data["ids"],
            documents=data["documents"],
            metadatas=data["metadatas"],
            embeddings=new_embeddings,
        )
        offset += batch_size
        print(f"\rПереиндексировано: {offset}", end="")

    print(f"\nГотово. Можно удалить коллекцию '{old_name}'")