Metadata Filtering: фильтрация до векторного поиска

from dataclasses import dataclass, field
from datetime import datetime
from typing import Any
import numpy as np
from sentence_transformers import SentenceTransformer


@dataclass
class Document:
    id: int
    text: str
    vector: np.ndarray
    # Metadata (payload)
    tenant_id: str
    language: str
    doc_type: str
    created_at: datetime
    access_level: int
    tags: list[str]
    is_verified: bool


class SimpleVectorStore:
    """Минималистичное хранилище с поддержкой metadata filtering."""

    def __init__(self):
        self.docs: dict[int, Document] = {}
        # Payload индексы для быстрой фильтрации
        self._idx_tenant:   dict[str, set[int]] = {}    # hashmap
        self._idx_language: dict[str, set[int]] = {}    # hashmap
        self._idx_doc_type: dict[str, set[int]] = {}    # hashmap
        # Числовые/дата поля — просто храним для range-фильтров (в prod: B-tree)
        # Сортированный список для range-запросов
        self._idx_created_at: list[tuple[datetime, int]] = []  # sorted by datetime

    def add(self, doc: Document) -> None:
        self.docs[doc.id] = doc
        # Обновляем инвертированные индексы
        self._idx_tenant.setdefault(doc.tenant_id, set()).add(doc.id)
        self._idx_language.setdefault(doc.language, set()).add(doc.id)
        self._idx_doc_type.setdefault(doc.doc_type, set()).add(doc.id)
        self._idx_created_at.append((doc.created_at, doc.id))
        self._idx_created_at.sort(key=lambda x: x[0])

    def _filter_to_ids(
        self,
        tenant_id: str | None = None,
        language: str | None = None,
        doc_types: list[str] | None = None,
        created_after: datetime | None = None,
        max_access_level: int | None = None,
        require_verified: bool | None = None,
        tags: list[str] | None = None,
    ) -> set[int] | None:
        """
        Возвращает set[id] документов, удовлетворяющих всем условиям.
        None означает «без фильтра» (все документы).
        """
        candidate_sets: list[set[int]] = []

        if tenant_id is not None:
            candidate_sets.append(self._idx_tenant.get(tenant_id, set()))

        if language is not None:
            candidate_sets.append(self._idx_language.get(language, set()))

        if doc_types is not None:
            # OR по списку типов, затем пересечение с остальными
            union: set[int] = set()
            for dt in doc_types:
                union |= self._idx_doc_type.get(dt, set())
            candidate_sets.append(union)

        # Пересечение всех conditions (AND)
        if not candidate_sets:
            base_ids = set(self.docs.keys())
        else:
            base_ids = candidate_sets[0]
            for s in candidate_sets[1:]:
                base_ids = base_ids & s

        # Диапазонные фильтры — применяем поверх base_ids
        if created_after is not None:
            base_ids = {
                doc_id for doc_id in base_ids
                if self.docs[doc_id].created_at >= created_after
            }

        if max_access_level is not None:
            base_ids = {
                doc_id for doc_id in base_ids
                if self.docs[doc_id].access_level <= max_access_level
            }

        if require_verified is not None:
            base_ids = {
                doc_id for doc_id in base_ids
                if self.docs[doc_id].is_verified == require_verified
            }

        if tags is not None:
            base_ids = {
                doc_id for doc_id in base_ids
                if any(t in self.docs[doc_id].tags for t in tags)
            }

        return base_ids

    def search(
        self,
        query_vector: np.ndarray,
        top_k: int = 5,
        **filter_kwargs,
    ) -> list[tuple[Document, float]]:
        """Pre-filtering: сначала фильтр, затем ANN по кандидатам."""
        candidate_ids = self._filter_to_ids(**filter_kwargs)

        # Векторный поиск только по кандидатам
        results = []
        for doc_id in candidate_ids:
            doc = self.docs[doc_id]
            score = float(query_vector @ doc.vector)  # cosine (при норм. векторах)
            results.append((doc, score))

        results.sort(key=lambda x: x[1], reverse=True)
        return results[:top_k]


# ── Пример использования ────────────────────────────────────────────────
model = SentenceTransformer("all-MiniLM-L6-v2")
store = SimpleVectorStore()

documents_data = [
    ("Политика отпусков компании ACME 2024", "acme", "ru", "policy", 1, ["hr","benefits"], True),
    ("ACME Vacation Policy 2024",             "acme", "en", "policy", 1, ["hr","benefits"], True),
    ("Политика отпусков компании Beta 2024",  "beta", "ru", "policy", 1, ["hr"],           True),
    ("ACME техническая документация API",     "acme", "ru", "technical", 0, ["api","dev"], True),
    ("ACME черновик новой политики",          "acme", "ru", "draft",  1, ["hr"],           False),
    ("ACME компенсации и льготы C-level",     "acme", "ru", "policy", 3, ["hr","comp"],    True),
]

for i, (text, tenant, lang, doc_type, access, tags, verified) in enumerate(documents_data):
    vec = model.encode(text, normalize_embeddings=True)
    store.add(Document(
        id=i, text=text, vector=vec,
        tenant_id=tenant, language=lang, doc_type=doc_type,
        created_at=datetime(2024, 1, 1),
        access_level=access, tags=tags, is_verified=verified,
    ))

# Поиск для обычного сотрудника ACME (access_level=1)
query_vec = model.encode("политика отпусков", normalize_embeddings=True)
results = store.search(
    query_vec,
    top_k=3,
    tenant_id="acme",           # только документы ACME
    language="ru",              # только русский язык
    doc_types=["policy","faq"], # только политики и FAQ
    max_access_level=1,         # уровень доступа ≤ 1
    require_verified=True,      # только проверенные
)

for doc, score in results:
    print(f"[{score:.3f}] {doc.text}")

[0.891] Политика отпусков компании ACME 2024
# ACME Vacation Policy (english) — отфильтрован по language="ru"
# Beta policy — отфильтрован по tenant_id="acme"
# draft — отфильтрован по require_verified=True
# C-level — отфильтрован по max_access_level=1

from qdrant_client import QdrantClient
from qdrant_client.models import (
    Distance, VectorParams, PointStruct,
    Filter, FieldCondition, MatchValue, MatchAny, Range,
    PayloadSchemaType,
)
from datetime import datetime

client = QdrantClient(":memory:")

COLLECTION = "kb"
client.create_collection(
    collection_name=COLLECTION,
    vectors_config=VectorParams(size=384, distance=Distance.COSINE),
)

# ── Создаём payload index на нужных полях ──────────────────────────────
# БЕЗ этого фильтрация = полный скан (O(N)), медленно при большом N

client.create_payload_index(COLLECTION, "tenant_id",    PayloadSchemaType.KEYWORD)
client.create_payload_index(COLLECTION, "language",     PayloadSchemaType.KEYWORD)
client.create_payload_index(COLLECTION, "doc_type",     PayloadSchemaType.KEYWORD)
client.create_payload_index(COLLECTION, "access_level", PayloadSchemaType.INTEGER)
client.create_payload_index(COLLECTION, "created_at",   PayloadSchemaType.FLOAT)
client.create_payload_index(COLLECTION, "is_verified",  PayloadSchemaType.BOOL)
client.create_payload_index(COLLECTION, "tags",         PayloadSchemaType.KEYWORD)

# ── Загрузка точек ─────────────────────────────────────────────────────
from sentence_transformers import SentenceTransformer
model = SentenceTransformer("all-MiniLM-L6-v2")

docs = [
    {"id": 1, "text": "Политика отпусков ACME 2024",       "tenant": "acme", "lang": "ru", "type": "policy",    "access": 1, "tags": ["hr","benefits"], "verified": True,  "ts": 1704067200},
    {"id": 2, "text": "ACME Vacation Policy 2024",          "tenant": "acme", "lang": "en", "type": "policy",    "access": 1, "tags": ["hr","benefits"], "verified": True,  "ts": 1704067200},
    {"id": 3, "text": "Политика Beta Corporation",          "tenant": "beta", "lang": "ru", "type": "policy",    "access": 1, "tags": ["hr"],            "verified": True,  "ts": 1704067200},
    {"id": 4, "text": "ACME API техническая документация",  "tenant": "acme", "lang": "ru", "type": "technical", "access": 0, "tags": ["api","dev"],     "verified": True,  "ts": 1704067200},
    {"id": 5, "text": "ACME черновик политики (draft)",     "tenant": "acme", "lang": "ru", "type": "draft",     "access": 1, "tags": ["hr"],            "verified": False, "ts": 1704067200},
    {"id": 6, "text": "ACME компенсации топ-менеджмента",  "tenant": "acme", "lang": "ru", "type": "policy",    "access": 3, "tags": ["hr","comp"],     "verified": True,  "ts": 1704067200},
]

points = [
    PointStruct(
        id=d["id"],
        vector=model.encode(d["text"], normalize_embeddings=True).tolist(),
        payload={
            "text":         d["text"],
            "tenant_id":    d["tenant"],
            "language":     d["lang"],
            "doc_type":     d["type"],
            "access_level": d["access"],
            "tags":         d["tags"],
            "is_verified":  d["verified"],
            "created_at":   d["ts"],
        },
    )
    for d in docs
]
client.upsert(collection_name=COLLECTION, points=points)

# ── Поиск с фильтром ───────────────────────────────────────────────────
query_vector = model.encode("политика отпусков", normalize_embeddings=True).tolist()

results = client.search(
    collection_name=COLLECTION,
    query_vector=query_vector,
    limit=5,
    query_filter=Filter(
        must=[
            FieldCondition(key="tenant_id",   match=MatchValue(value="acme")),
            FieldCondition(key="language",    match=MatchValue(value="ru")),
            FieldCondition(key="is_verified", match=MatchValue(value=True)),
            FieldCondition(key="access_level",range=Range(lte=1)),
        ],
        must_not=[
            FieldCondition(key="doc_type", match=MatchValue(value="draft")),
        ],
        should=[
            # Повышаем приоритет документов с тегами hr или benefits
            # (should ≠ обязательное условие; ни одно из should может не выполняться)
            FieldCondition(key="tags", match=MatchAny(any=["hr", "benefits"])),
        ],
    ),
)

for r in results:
    print(f"id={r.id}  score={r.score:.3f}  {r.payload['text']}")

from qdrant_client.models import DatetimeRange
from datetime import datetime, timezone

# ── Range по дате: документы за последний год ──────────────────────────
cutoff = datetime(2024, 1, 1, tzinfo=timezone.utc).timestamp()

results = client.search(
    collection_name=COLLECTION,
    query_vector=query_vector,
    limit=5,
    query_filter=Filter(
        must=[
            FieldCondition(key="tenant_id", match=MatchValue(value="acme")),
            FieldCondition(
                key="created_at",
                range=Range(gte=cutoff),   # created_at >= 2024-01-01
            ),
        ]
    ),
)

# ── Nested AND/OR: (type==policy OR type==faq) AND (lang==ru OR lang==en) ──
from qdrant_client.models import Filter as QFilter

results = client.search(
    collection_name=COLLECTION,
    query_vector=query_vector,
    limit=5,
    query_filter=QFilter(
        must=[
            FieldCondition(key="tenant_id", match=MatchValue(value="acme")),
            # Вложенный should работает как OR внутри must (AND):
            QFilter(
                should=[
                    FieldCondition(key="doc_type", match=MatchValue(value="policy")),
                    FieldCondition(key="doc_type", match=MatchValue(value="faq")),
                ]
            ),
        ]
    ),
)

# ── Scroll (получить все документы клиента без векторного поиска) ──────
scroll_result, next_cursor = client.scroll(
    collection_name=COLLECTION,
    scroll_filter=Filter(
        must=[FieldCondition(key="tenant_id", match=MatchValue(value="acme"))]
    ),
    limit=100,
    with_payload=True,
    with_vectors=False,
)
print(f"Всего документов ACME: {len(scroll_result)}")

import weaviate
import weaviate.classes as wvc
from weaviate.classes.query import Filter, MetadataQuery

client = weaviate.connect_to_local()

kb = client.collections.get("KnowledgeBase")

# ── Простой фильтр ─────────────────────────────────────────────────────
results = kb.query.near_text(
    query="политика отпусков",
    limit=5,
    filters=Filter.by_property("tenant_id").equal("acme"),
)

# ── Составной AND-фильтр ───────────────────────────────────────────────
results = kb.query.near_text(
    query="политика отпусков",
    limit=5,
    filters=(
        Filter.by_property("tenant_id").equal("acme")
        & Filter.by_property("language").equal("ru")
        & Filter.by_property("is_verified").equal(True)
        & Filter.by_property("access_level").less_or_equal(1)
    ),
    return_metadata=MetadataQuery(distance=True),
)

for obj in results.objects:
    print(f"distance={obj.metadata.distance:.3f}  {obj.properties['text']}")

# ── OR-фильтр ──────────────────────────────────────────────────────────
results = kb.query.near_text(
    query="техническая документация",
    limit=5,
    filters=(
        Filter.by_property("tenant_id").equal("acme")
        & (
            Filter.by_property("doc_type").equal("technical")
            | Filter.by_property("doc_type").equal("faq")
        )
    ),
)

# ── Range по числовому полю ────────────────────────────────────────────
results = kb.query.near_text(
    query="новости компании",
    limit=5,
    filters=(
        Filter.by_property("tenant_id").equal("acme")
        & Filter.by_property("created_at").greater_or_equal(1704067200)  # 2024-01-01
    ),
)

# ── Фильтр по массиву tags (contains any) ─────────────────────────────
results = kb.query.near_text(
    query="льготы сотрудников",
    limit=5,
    filters=(
        Filter.by_property("tenant_id").equal("acme")
        & Filter.by_property("tags").contains_any(["hr", "benefits"])
    ),
)

client.close()

import chromadb
from sentence_transformers import SentenceTransformer

client = chromadb.Client()
model = SentenceTransformer("all-MiniLM-L6-v2")

collection = client.create_collection("kb")

# Добавляем документы с метаданными
texts = [
    "Политика отпусков ACME 2024",
    "ACME API документация",
    "Политика Beta Corporation",
    "ACME черновик политики",
]
metas = [
    {"tenant_id": "acme", "language": "ru", "doc_type": "policy",    "access_level": 1, "is_verified": True},
    {"tenant_id": "acme", "language": "ru", "doc_type": "technical", "access_level": 0, "is_verified": True},
    {"tenant_id": "beta", "language": "ru", "doc_type": "policy",    "access_level": 1, "is_verified": True},
    {"tenant_id": "acme", "language": "ru", "doc_type": "draft",     "access_level": 1, "is_verified": False},
]

collection.add(
    documents=texts,
    embeddings=model.encode(texts, normalize_embeddings=True).tolist(),
    metadatas=metas,
    ids=[str(i) for i in range(len(texts))],
)

query_vec = model.encode(["политика отпусков"], normalize_embeddings=True).tolist()

# ── Простой равенство ──────────────────────────────────────────────────
results = collection.query(
    query_embeddings=query_vec,
    n_results=3,
    where={"tenant_id": "acme"},   # сокращение для {"tenant_id": {"$eq": "acme"}}
)

# ── AND через $and ─────────────────────────────────────────────────────
results = collection.query(
    query_embeddings=query_vec,
    n_results=3,
    where={
        "$and": [
            {"tenant_id":   {"$eq": "acme"}},
            {"language":    {"$eq": "ru"}},
            {"is_verified": {"$eq": True}},
            {"access_level":{"$lte": 1}},
            {"doc_type":    {"$ne": "draft"}},
        ]
    },
)

# ── OR через $or ───────────────────────────────────────────────────────
results = collection.query(
    query_embeddings=query_vec,
    n_results=3,
    where={
        "$and": [
            {"tenant_id": {"$eq": "acme"}},
            {"$or": [
                {"doc_type": {"$eq": "policy"}},
                {"doc_type": {"$eq": "faq"}},
            ]},
        ]
    },
)

# Операторы: $eq, $ne, $gt, $gte, $lt, $lte, $in, $nin
results = collection.query(
    query_embeddings=query_vec,
    n_results=3,
    where={"doc_type": {"$in": ["policy", "faq", "technical"]}},
)

for i, doc in enumerate(results["documents"][0]):
    meta = results["metadatas"][0][i]
    dist = results["distances"][0][i]
    print(f"[{dist:.3f}] {doc[:60]}  ({meta['doc_type']})")

from dataclasses import dataclass

@dataclass
class UserContext:
    user_id: str
    tenant_id: str
    role_level: int          # 0=guest, 1=employee, 2=manager, 3=admin
    allowed_languages: list[str]
    allowed_doc_types: list[str] | None = None  # None = все типы


def build_qdrant_filter(ctx: UserContext) -> Filter:
    """Строим фильтр на основе контекста пользователя."""
    must_conditions = [
        # Мультитенантная изоляция — первым делом
        FieldCondition(key="tenant_id",    match=MatchValue(value=ctx.tenant_id)),
        # Только проверенные документы
        FieldCondition(key="is_verified",  match=MatchValue(value=True)),
        # RBAC: access_level <= роль пользователя
        FieldCondition(key="access_level", range=Range(lte=ctx.role_level)),
    ]

    # Языковой фильтр
    if len(ctx.allowed_languages) == 1:
        must_conditions.append(
            FieldCondition(key="language", match=MatchValue(value=ctx.allowed_languages[0]))
        )
    else:
        must_conditions.append(
            FieldCondition(key="language", match=MatchAny(any=ctx.allowed_languages))
        )

    # Фильтр по типу документа (если задан)
    if ctx.allowed_doc_types:
        must_conditions.append(
            FieldCondition(key="doc_type", match=MatchAny(any=ctx.allowed_doc_types))
        )

    return Filter(
        must=must_conditions,
        must_not=[
            # Никогда не показываем черновики
            FieldCondition(key="doc_type", match=MatchValue(value="draft")),
        ],
    )


def rag_search(query: str, user_ctx: UserContext, top_k: int = 5):
    """RAG-поиск с автоматическим применением ACL-фильтра."""
    query_vector = model.encode(query, normalize_embeddings=True).tolist()
    filter_ = build_qdrant_filter(user_ctx)

    results = client.search(
        collection_name=COLLECTION,
        query_vector=query_vector,
        limit=top_k,
        query_filter=filter_,
    )
    return results


# ── Использование ──────────────────────────────────────────────────────
employee_ctx = UserContext(
    user_id="ivan@acme.com",
    tenant_id="acme",
    role_level=1,
    allowed_languages=["ru", "en"],
)

manager_ctx = UserContext(
    user_id="anna@acme.com",
    tenant_id="acme",
    role_level=2,
    allowed_languages=["ru"],
    allowed_doc_types=["policy", "faq", "technical", "legal"],
)

# Сотрудник не увидит C-level документы
employee_results = rag_search("компенсации топ-менеджмента", employee_ctx)
# Менеджер увидит
manager_results  = rag_search("компенсации топ-менеджмента", manager_ctx)