CRUD в векторных БД: upsert, query, delete, metadata filtering

import chromadb
from chromadb.utils import embedding_functions

client = chromadb.PersistentClient(path="./chroma_db")
ef = embedding_functions.OpenAIEmbeddingFunction(
    api_key="sk-...",
    model_name="text-embedding-3-small",
)
collection = client.get_or_create_collection("docs", embedding_function=ef)

# ── Upsert одного документа ───────────────────────────────────────────
collection.upsert(
    ids=["doc-001"],
    documents=["Введение в RAG: основные принципы..."],
    metadatas=[{"source": "handbook.pdf", "chapter": 1, "lang": "ru"}],
)

# ── Bulk upsert (батчами по 100 — Chroma не любит большие батчи) ──────
BATCH = 100
for i in range(0, len(docs), BATCH):
    batch = docs[i:i+BATCH]
    collection.upsert(
        ids=[d["id"] for d in batch],
        documents=[d["text"] for d in batch],
        metadatas=[d["meta"] for d in batch],
    )

# ── Обновить только метаданные (без пересчёта embedding) ──────────────
collection.update(
    ids=["doc-001"],
    metadatas=[{"source": "handbook.pdf", "chapter": 1, "lang": "ru", "status": "reviewed"}],
)

# ── Получить объект по ID ─────────────────────────────────────────────
result = collection.get(ids=["doc-001"], include=["documents", "metadatas"])
print(result["metadatas"][0])

from qdrant_client import QdrantClient
from qdrant_client.models import PointStruct, PointIdsList, SetPayload
import uuid

client = QdrantClient(url="http://localhost:6333")

# ── Upsert (upsert = insert or replace) ──────────────────────────────
client.upsert(
    collection_name="docs",
    points=[
        PointStruct(
            id=str(uuid.uuid4()),          # UUID или int
            vector=[0.1, 0.2, ...],        # embedding
            payload={
                "text": "Введение в RAG...",
                "source": "handbook.pdf",
                "chapter": 1,
                "lang": "ru",
            },
        )
    ],
)

# ── Batch upsert ───────────────────────────────────────────────────────
points = [
    PointStruct(id=d["id"], vector=d["embedding"], payload=d["meta"])
    for d in docs
]
# upload_points автоматически бьёт на батчи по 64
client.upload_points(
    collection_name="docs",
    points=points,
    batch_size=64,
    parallel=4,     # параллельных воркеров
)

# ── Обновить только payload (без пересчёта вектора) ───────────────────
client.set_payload(
    collection_name="docs",
    payload={"status": "reviewed", "reviewer": "ivan"},
    points=["doc-uuid-001"],  # или Filter(...)
)

# ── Удалить конкретные ключи из payload ───────────────────────────────
client.delete_payload(
    collection_name="docs",
    keys=["reviewer"],
    points=["doc-uuid-001"],
)

# ── Получить объекты по ID ────────────────────────────────────────────
results = client.retrieve(
    collection_name="docs",
    ids=["doc-uuid-001"],
    with_payload=True,
    with_vectors=False,  # вектор обычно не нужен
)

import weaviate
import weaviate.classes as wvc
from weaviate.util import generate_uuid4

with weaviate.connect_to_local() as client:
    articles = client.collections.use("Articles")

    # ── Insert (= upsert с явным UUID) ────────────────────────────────
    # Если UUID уже существует — Weaviate выдаст ошибку, используйте replace
    uuid = generate_uuid4()
    articles.data.insert(
        properties={"title": "RAG Guide", "content": "...", "lang": "ru"},
        uuid=uuid,
    )

    # ── Replace (полная замена объекта по UUID) ────────────────────────
    articles.data.replace(
        uuid=uuid,
        properties={"title": "RAG Guide v2", "content": "Новый контент", "lang": "ru"},
    )

    # ── Partial update (обновить только указанные поля) ───────────────
    articles.data.update(
        uuid=uuid,
        properties={"lang": "en"},   # остальные поля не меняются
    )

    # ── Bulk upsert через insert_many ─────────────────────────────────
    objects = [
        wvc.data.DataObject(
            properties={"title": d["title"], "content": d["text"]},
            uuid=d.get("uuid") or generate_uuid4(),
        )
        for d in docs
    ]
    result = articles.data.insert_many(objects)
    if result.has_errors:
        for err in result.errors.values():
            print(f"Error: {err.message}")

    # ── Паттерн настоящего upsert (проверить → вставить/заменить) ─────
    def upsert_article(client, collection_name: str, doc: dict, doc_uuid: str):
        col = client.collections.use(collection_name)
        try:
            col.data.replace(uuid=doc_uuid, properties=doc)
        except weaviate.exceptions.UnexpectedStatusCodeError:
            col.data.insert(properties=doc, uuid=doc_uuid)

-- ── Схема таблицы ─────────────────────────────────────────────────────
CREATE TABLE chunks (
    id         TEXT PRIMARY KEY,
    content    TEXT NOT NULL,
    embedding  vector(1536),
    source     TEXT,
    chapter    INT,
    lang       TEXT DEFAULT 'ru',
    updated_at TIMESTAMPTZ DEFAULT now()
);

-- ── Upsert: INSERT ... ON CONFLICT DO UPDATE ───────────────────────────
INSERT INTO chunks (id, content, embedding, source, chapter, lang)
VALUES ($1, $2, $3, $4, $5, $6)
ON CONFLICT (id) DO UPDATE SET
    content    = EXCLUDED.content,
    embedding  = EXCLUDED.embedding,
    source     = EXCLUDED.source,
    chapter    = EXCLUDED.chapter,
    updated_at = now();

-- ── Обновить только метаданные (без пересчёта embedding) ──────────────
UPDATE chunks
SET lang = 'en', updated_at = now()
WHERE id = $1;

-- ── Batch upsert через COPY + upsert ──────────────────────────────────
-- Эффективнее для больших объёмов:
CREATE TEMP TABLE chunks_staging (LIKE chunks INCLUDING ALL);
COPY chunks_staging FROM STDIN WITH (FORMAT CSV);
INSERT INTO chunks
    SELECT * FROM chunks_staging
ON CONFLICT (id) DO UPDATE SET
    content = EXCLUDED.content,
    embedding = EXCLUDED.embedding,
    updated_at = now();
DROP TABLE chunks_staging;

import asyncpg
import numpy as np

async def upsert_chunks(pool: asyncpg.Pool, chunks: list[dict]):
    """Bulk upsert чанков с embeddingами."""
    async with pool.acquire() as conn:
        await conn.executemany(
            """
            INSERT INTO chunks (id, content, embedding, source, lang)
            VALUES ($1, $2, $3::vector, $4, $5)
            ON CONFLICT (id) DO UPDATE SET
                content   = EXCLUDED.content,
                embedding = EXCLUDED.embedding,
                source    = EXCLUDED.source,
                updated_at = now()
            """,
            [
                (c["id"], c["content"], str(c["embedding"].tolist()), c["source"], c.get("lang", "ru"))
                for c in chunks
            ],
        )

# ── Chroma ────────────────────────────────────────────────────────────
results = collection.query(
    query_texts=["как работает RAG"],
    n_results=10,
    where={"lang": {"$eq": "ru"}},                    # metadata filter
    where_document={"$contains": "retrieval"},         # filter по тексту документа
    include=["documents", "metadatas", "distances"],
)
# Порог: results["distances"][0][i] < 0.5

# ── Qdrant ────────────────────────────────────────────────────────────
from qdrant_client.models import Filter, FieldCondition, MatchValue, Range

results = client.query_points(
    collection_name="docs",
    query=query_vector,               # готовый embedding
    query_filter=Filter(
        must=[
            FieldCondition(key="lang", match=MatchValue(value="ru")),
            FieldCondition(key="chapter", range=Range(gte=1, lte=5)),
        ]
    ),
    score_threshold=0.7,              # отсечь нерелевантное
    limit=10,
    with_payload=True,
)

# ── Weaviate ──────────────────────────────────────────────────────────
from weaviate.classes.query import Filter, MetadataQuery

with weaviate.connect_to_local() as client:
    results = client.collections.use("Articles").query.hybrid(
        query="как работает RAG",
        alpha=0.75,
        where=Filter.by_property("lang").equal("ru") &
              Filter.by_property("chapter").greater_or_equal(1),
        limit=10,
        return_metadata=MetadataQuery(score=True, distance=True),
    )

# ── pgvector ──────────────────────────────────────────────────────────
rows = await conn.fetch(
    """
    SELECT id, content, source,
           embedding <=> $1 AS distance
    FROM chunks
    WHERE lang = $2
      AND chapter BETWEEN $3 AND $4
      AND embedding <=> $1 < $5        -- порог близости
    ORDER BY distance
    LIMIT $6
    """,
    str(query_vector.tolist()), "ru", 1, 5, 0.5, 10,
)

def search_with_filter(
    query_vec: list[float],
    metadata_filter: dict,
    limit: int = 10,
    oversampling: int = 5,         # запрашиваем в 5 раз больше
) -> list[dict]:
    """
    Реализует post-filter с oversampling для pgvector.
    Если фильтры очень строгие — увеличьте oversampling до 10-20.
    """
    candidates = limit * oversampling  # запрашиваем с запасом

    rows = conn.execute(
        """
        SELECT id, content, source,
               embedding <=> %s AS distance
        FROM chunks
        ORDER BY distance           -- сначала ищем без фильтра
        LIMIT %s                    -- берём с запасом
        """,
        (str(query_vec), candidates),
    ).fetchall()

    # Post-filter в Python
    filtered = [
        r for r in rows
        if all(r[k] == v for k, v in metadata_filter.items())
    ]
    return filtered[:limit]


# ── Qdrant: настройка oversampling ────────────────────────────────────
from qdrant_client.models import SearchParams

results = client.query_points(
    collection_name="docs",
    query=query_vector,
    query_filter=Filter(must=[FieldCondition(key="lang", match=MatchValue(value="ru"))]),
    search_params=SearchParams(
        hnsw_ef=128,              # точность HNSW traversal
        exact=False,              # False = HNSW, True = brute-force (100% recall)
    ),
    limit=10,
    # Qdrant сам обрабатывает фильтр in-search — oversampling не нужен
)

# ── Qdrant: создать индексы на payload ───────────────────────────────
from qdrant_client.models import PayloadSchemaType

client.create_payload_index(
    collection_name="docs",
    field_name="lang",
    field_schema=PayloadSchemaType.KEYWORD,   # hash-индекс для exact match
)
client.create_payload_index(
    collection_name="docs",
    field_name="chapter",
    field_schema=PayloadSchemaType.INTEGER,   # range-индекс для числовых сравнений
)
client.create_payload_index(
    collection_name="docs",
    field_name="published_at",
    field_schema=PayloadSchemaType.DATETIME,  # для сравнений дат
)

# ── pgvector: стандартные PostgreSQL индексы ──────────────────────────
# Запустить один раз после создания таблицы
await conn.execute("CREATE INDEX idx_chunks_lang    ON chunks(lang)")
await conn.execute("CREATE INDEX idx_chunks_chapter ON chunks(chapter)")
await conn.execute("CREATE INDEX idx_chunks_date    ON chunks(published_at)")
# Для полнотекстового поиска
await conn.execute("""
    CREATE INDEX idx_chunks_fts ON chunks
    USING GIN (to_tsvector('russian', content))
""")

# ── Weaviate: настройка при создании схемы ────────────────────────────
from weaviate.classes.config import Property, DataType
Property(
    name="lang",
    data_type=DataType.TEXT,
    index_filterable=True,    # hash-индекс для фильтров
    index_searchable=False,   # не нужен BM25 для lang
),
Property(
    name="chapter",
    data_type=DataType.INT,
    index_filterable=True,
    index_range_filters=True, # дополнительный B-tree для range запросов
),

# ── Chroma ────────────────────────────────────────────────────────────
# По ID
collection.delete(ids=["doc-001", "doc-002"])

# По фильтру (все документы из источника)
collection.delete(where={"source": {"$eq": "old_handbook.pdf"}})

# По тексту документа
collection.delete(where_document={"$contains": "устаревшая информация"})


# ── Qdrant ────────────────────────────────────────────────────────────
from qdrant_client.models import PointIdsList, FilterSelector, Filter, FieldCondition, MatchValue

# По списку ID
client.delete(
    collection_name="docs",
    points_selector=PointIdsList(points=["uuid-001", "uuid-002"]),
)

# По фильтру — удалить все из старого источника
client.delete(
    collection_name="docs",
    points_selector=FilterSelector(
        filter=Filter(
            must=[FieldCondition(key="source", match=MatchValue(value="old_handbook.pdf"))]
        )
    ),
)

# Принудительная оптимизация после массового удаления
client.update_collection(
    collection_name="docs",
    optimizer_config={"deleted_threshold": 0.2},   # оптимизировать при 20% удалённых
)


# ── Weaviate ──────────────────────────────────────────────────────────
from weaviate.classes.query import Filter

with weaviate.connect_to_local() as client:
    articles = client.collections.use("Articles")

    # По ID
    articles.data.delete_by_id("object-uuid-here")

    # По фильтру
    result = articles.data.delete_many(
        where=Filter.by_property("source").equal("old_handbook.pdf"),
        dry_run=True,    # сначала посмотреть сколько будет удалено
    )
    print(f"Будет удалено: {result.matches} объектов")

    result = articles.data.delete_many(
        where=Filter.by_property("source").equal("old_handbook.pdf"),
        # dry_run=False по умолчанию — реальное удаление
    )


# ── pgvector ──────────────────────────────────────────────────────────
# По ID
await conn.execute("DELETE FROM chunks WHERE id = $1", "doc-001")

# По фильтру (с RETURNING чтобы узнать сколько удалено)
deleted = await conn.fetch(
    "DELETE FROM chunks WHERE source = $1 RETURNING id",
    "old_handbook.pdf"
)
print(f"Удалено: {len(deleted)} чанков")

# Очистить Dead tuples после массового удаления
await conn.execute("VACUUM ANALYZE chunks")

from datetime import datetime, UTC
from functools import wraps

# ── Паттерн: soft delete через метаданные ─────────────────────────────
# При "удалении" — обновляем payload, не трогаем вектор
def soft_delete_qdrant(client, collection_name: str, doc_ids: list[str]):
    """Помечает объекты как удалённые. Физически не трогает вектор."""
    client.set_payload(
        collection_name=collection_name,
        payload={
            "is_deleted": True,
            "deleted_at": datetime.now(UTC).isoformat(),
        },
        points=doc_ids,
    )

# ── Запросы всегда исключают deleted ─────────────────────────────────
def active_filter(base_filter=None):
    """Добавляет фильтр is_deleted=false к любому запросу."""
    from qdrant_client.models import Filter, FieldCondition, MatchValue
    not_deleted = FieldCondition(
        key="is_deleted",
        match=MatchValue(value=False),
    )
    if base_filter is None:
        return Filter(must=[not_deleted])
    return Filter(
        must=[not_deleted, *base_filter.must],
        should=base_filter.should,
        must_not=base_filter.must_not,
    )

# Использование:
results = client.query_points(
    collection_name="docs",
    query=query_vector,
    query_filter=active_filter(
        Filter(must=[FieldCondition(key="lang", match=MatchValue(value="ru"))])
    ),
    limit=10,
)

# ── Периодическая физическая очистка (раз в сутки/неделю) ─────────────
def purge_deleted(client, collection_name: str):
    """Физически удаляет soft-deleted объекты."""
    from qdrant_client.models import FilterSelector, Filter, FieldCondition, MatchValue
    client.delete(
        collection_name=collection_name,
        points_selector=FilterSelector(
            filter=Filter(must=[
                FieldCondition(key="is_deleted", match=MatchValue(value=True))
            ])
        ),
    )

from abc import ABC, abstractmethod
from dataclasses import dataclass

@dataclass
class SearchResult:
    id: str
    text: str
    metadata: dict
    score: float

class VectorStore(ABC):
    """Абстракция над векторной базой данных."""

    @abstractmethod
    def upsert(self, documents: list[dict]) -> None:
        """
        documents: [{"id": str, "text": str, "embedding": list[float], "metadata": dict}]
        Каждый документ upsert'ится по его id.
        """
        ...

    @abstractmethod
    def search(
        self,
        query_embedding: list[float],
        filters: dict | None = None,
        limit: int = 10,
        score_threshold: float | None = None,
    ) -> list[SearchResult]:
        """Векторный поиск с опциональным metadata-фильтром."""
        ...

    @abstractmethod
    def delete(self, ids: list[str]) -> None:
        """Удалить объекты по ID."""
        ...

    def delete_by_filter(self, filters: dict) -> int:
        """Удалить по фильтру. Возвращает количество удалённых. Опционально."""
        raise NotImplementedError


# ── Реализация для Qdrant ─────────────────────────────────────────────
class QdrantStore(VectorStore):
    def __init__(self, url: str, collection: str):
        from qdrant_client import QdrantClient
        self.client = QdrantClient(url=url)
        self.collection = collection

    def upsert(self, documents: list[dict]) -> None:
        from qdrant_client.models import PointStruct
        self.client.upsert(
            collection_name=self.collection,
            points=[
                PointStruct(
                    id=d["id"],
                    vector=d["embedding"],
                    payload={"text": d["text"], **d.get("metadata", {})},
                )
                for d in documents
            ],
        )

    def search(self, query_embedding, filters=None, limit=10, score_threshold=None):
        from qdrant_client.models import Filter, FieldCondition, MatchValue
        q_filter = None
        if filters:
            q_filter = Filter(must=[
                FieldCondition(key=k, match=MatchValue(value=v))
                for k, v in filters.items()
            ])
        results = self.client.query_points(
            collection_name=self.collection,
            query=query_embedding,
            query_filter=q_filter,
            score_threshold=score_threshold,
            limit=limit,
            with_payload=True,
        )
        return [
            SearchResult(
                id=str(r.id),
                text=r.payload.get("text", ""),
                metadata={k: v for k, v in r.payload.items() if k != "text"},
                score=r.score,
            )
            for r in results.points
        ]

    def delete(self, ids: list[str]) -> None:
        from qdrant_client.models import PointIdsList
        self.client.delete(
            collection_name=self.collection,
            points_selector=PointIdsList(points=ids),
        )


# ── Использование: можно менять бэкенд без изменения кода ─────────────
def build_index(store: VectorStore, chunks: list[dict]):
    """Работает с любой реализацией VectorStore."""
    store.upsert(chunks)

def retrieve(store: VectorStore, query_emb: list[float], lang: str) -> list[SearchResult]:
    return store.search(query_emb, filters={"lang": lang}, limit=10, score_threshold=0.7)

# Dev:
store = QdrantStore("http://localhost:6333", "docs")
# Prod — просто меняем:
# store = PGVectorStore(dsn="postgresql://...", table="chunks")

# ── Объект vectordb = ID + Vector + Payload ───────────────────────────

# ── Детерминированный ID (для idempotent upsert) ─────────────────────
import hashlib
doc_id = hashlib.md5(f"{source_url}::{chunk_index}".encode()).hexdigest()

# ── Upsert ────────────────────────────────────────────────────────────
# Chroma:   collection.upsert(ids, documents, metadatas)
# Qdrant:   client.upsert(collection_name, points=[PointStruct(...)])
#           client.upload_points(..., batch_size=64, parallel=4)
# Weaviate: collection.data.insert_many([DataObject(...)])
# pgvector: INSERT ... ON CONFLICT (id) DO UPDATE SET ...

# ── Обновить только метаданные (без пересчёта вектора) ───────────────
# Qdrant:   client.set_payload(payload={...}, points=[id])
# Weaviate: collection.data.update(uuid, properties={...})
# pgvector: UPDATE chunks SET status=$1 WHERE id=$2

# ── Query ─────────────────────────────────────────────────────────────
# Chroma:   collection.query(query_texts, where={...}, n_results)
# Qdrant:   client.query_points(query=vec, query_filter=Filter(...), score_threshold)
# Weaviate: collection.query.hybrid(query, alpha, where, limit)
# pgvector: SELECT ... WHERE ... ORDER BY emb<=>$1 LIMIT $2

# ── Metadata filtering стратегии ──────────────────────────────────────
# Pre-filter (Chroma):  WHERE применяется ДО HNSW поиска
#   + быстро, - ухудшает recall на разреженных фильтрах
# Post-filter (pgvector): HNSW ищет topK×oversampling, потом WHERE
#   + хороший recall, - нужен oversampling (×3-10)
# In-filter (Qdrant):  фильтр интегрирован в обход графа
#   + оптимально, автоматический выбор стратегии

# ── Indексы на payload (создать ДО заполнения!) ───────────────────────
# Qdrant:   client.create_payload_index(field_name, field_schema=KEYWORD/INTEGER/...)
# pgvector: CREATE INDEX ON chunks(lang); CREATE INDEX ON chunks(chapter)
# Weaviate: Property(index_filterable=True, index_range_filters=True)

# ── Delete ────────────────────────────────────────────────────────────
# Chroma:   collection.delete(ids=[...]) / collection.delete(where={...})
# Qdrant:   client.delete(points_selector=PointIdsList/FilterSelector)
# Weaviate: collection.data.delete_by_id(uuid) / .delete_many(where=...)
# pgvector: DELETE FROM chunks WHERE id=$1 / WHERE условие

# ── После массового удаления ──────────────────────────────────────────
# pgvector: VACUUM ANALYZE chunks
# Qdrant:   optimizer_config={"deleted_threshold": 0.2}
# Weaviate: фоновая компакция автоматическая

# ── Soft delete паттерн ───────────────────────────────────────────────
# set_payload({"is_deleted": True})
# Все запросы: must=[FieldCondition(key="is_deleted", match=False)]
# Периодически: purge физически удалённые объекты