LightRAG/lightrag/kg/vector_tenant_support.py

# Vector Database Multi-Tenant Support Module
# Supports: Qdrant, Milvus, FAISS, Nano Vector DB

from typing import Any, Dict, List


class VectorTenantHelper:
    """Helper class for vector DB multi-tenant operations"""

    @staticmethod
    def add_tenant_metadata(
        payload: Dict[str, Any], tenant_id: str, kb_id: str
    ) -> Dict[str, Any]:
        """Add tenant_id and kb_id to vector payload/metadata"""
        payload["tenant_id"] = tenant_id
        payload["kb_id"] = kb_id
        return payload

    @staticmethod
    def get_tenant_filter(tenant_id: str, kb_id: str) -> Dict[str, Any]:
        """Create a filter for tenant isolation in vector DB queries"""
        return {
            "must": [
                {"key": "tenant_id", "match": {"value": tenant_id}},
                {"key": "kb_id", "match": {"value": kb_id}},
            ]
        }

    @staticmethod
    def get_tenant_filter_milvus(tenant_id: str, kb_id: str) -> str:
        """Create a Milvus WHERE clause for tenant isolation"""
        return f'tenant_id == "{tenant_id}" && kb_id == "{kb_id}"'

    @staticmethod
    def make_tenant_id(tenant_id: str, kb_id: str, original_id: str) -> str:
        """Create a tenant-scoped vector ID"""
        return f"{tenant_id}:{kb_id}:{original_id}"

    @staticmethod
    def parse_tenant_id(tenant_id_str: str) -> Dict[str, str]:
        """Parse a tenant-scoped vector ID"""
        parts = tenant_id_str.split(":", 2)
        if len(parts) == 3:
            return {"tenant_id": parts[0], "kb_id": parts[1], "original_id": parts[2]}
        return {"original_id": tenant_id_str}

    @staticmethod
    def create_tenant_collection_name(
        base_name: str, tenant_id: str, kb_id: str
    ) -> str:
        """Create a tenant-scoped collection name"""
        return f"{base_name}_{tenant_id}_{kb_id}".replace("-", "_")

    @staticmethod
    def get_tenant_collection_pattern(
        base_name: str, tenant_id: str, kb_id: str
    ) -> str:
        """Get a pattern for finding tenant-specific collections"""
        return f"{base_name}_{tenant_id}_{kb_id}*"


class QdrantTenantHelper(VectorTenantHelper):
    """Qdrant-specific tenant helper"""

    @staticmethod
    def build_qdrant_filter(
        tenant_id: str, kb_id: str, additional_filter: Dict = None
    ) -> Dict[str, Any]:
        """Build a Qdrant filter for tenant isolation"""
        must_conditions = [
            {"key": "tenant_id", "match": {"value": tenant_id}},
            {"key": "kb_id", "match": {"value": kb_id}},
        ]

        if additional_filter:
            if "must" in additional_filter:
                must_conditions.extend(additional_filter["must"])

            result = {"must": must_conditions}

            # Copy over any other filter conditions
            for key in ["should", "must_not"]:
                if key in additional_filter:
                    result[key] = additional_filter[key]

            return result

        return {"must": must_conditions}

    @staticmethod
    def update_qdrant_payload(
        payload: Dict[str, Any], tenant_id: str, kb_id: str
    ) -> Dict[str, Any]:
        """Ensure Qdrant payload includes tenant metadata"""
        return VectorTenantHelper.add_tenant_metadata(payload, tenant_id, kb_id)


class MilvusTenantHelper(VectorTenantHelper):
    """Milvus-specific tenant helper"""

    @staticmethod
    def build_milvus_expr(
        tenant_id: str, kb_id: str, additional_expr: str = None
    ) -> str:
        """Build a Milvus WHERE expression for tenant isolation"""
        expr = f'tenant_id == "{tenant_id}" && kb_id == "{kb_id}"'

        if additional_expr:
            expr += f" && ({additional_expr})"

        return expr

    @staticmethod
    def insert_with_tenant(
        collection, data: List[Dict[str, Any]], tenant_id: str, kb_id: str
    ):
        """Insert data with tenant metadata into Milvus"""
        for item in data:
            item["tenant_id"] = tenant_id
            item["kb_id"] = kb_id
        return collection.insert(data)


class FAISSTenantHelper(VectorTenantHelper):
    """FAISS-specific tenant helper"""

    @staticmethod
    def create_tenant_index_name(base_name: str, tenant_id: str, kb_id: str) -> str:
        """Create a tenant-scoped FAISS index name"""
        sanitized_tenant = tenant_id.replace("-", "_").replace(":", "_")
        sanitized_kb = kb_id.replace("-", "_").replace(":", "_")
        return f"{base_name}_{sanitized_tenant}_{sanitized_kb}"

    @staticmethod
    def create_tenant_metadata_list(
        num_vectors: int, tenant_id: str, kb_id: str, base_metadata: List[Dict] = None
    ) -> List[Dict[str, Any]]:
        """Create metadata list for FAISS vectors with tenant info"""
        metadata_list = []

        for i in range(num_vectors):
            metadata = {"tenant_id": tenant_id, "kb_id": kb_id, "index": i}

            if base_metadata and i < len(base_metadata):
                metadata.update(base_metadata[i])

            metadata_list.append(metadata)

        return metadata_list

    @staticmethod
    def filter_metadata_by_tenant(
        metadata_list: List[Dict[str, Any]], tenant_id: str, kb_id: str
    ) -> List[int]:
        """Filter metadata list and return matching indices"""
        matching_indices = []

        for i, metadata in enumerate(metadata_list):
            if (
                metadata.get("tenant_id") == tenant_id
                and metadata.get("kb_id") == kb_id
            ):
                matching_indices.append(i)

        return matching_indices


class NanoVectorTenantHelper(VectorTenantHelper):
    """Nano Vector DB-specific tenant helper"""

    @staticmethod
    def build_nano_filter(tenant_id: str, kb_id: str) -> Dict[str, Any]:
        """Build a Nano Vector DB filter for tenant isolation"""
        return {"tenant_id": tenant_id, "kb_id": kb_id}

    @staticmethod
    def update_nano_document(
        doc: Dict[str, Any], tenant_id: str, kb_id: str
    ) -> Dict[str, Any]:
        """Update document to include tenant metadata"""
        if "metadata" not in doc:
            doc["metadata"] = {}

        doc["metadata"]["tenant_id"] = tenant_id
        doc["metadata"]["kb_id"] = kb_id

        return doc