LightRAG/lightrag/services/tenant_service.py

"""Service for managing tenants and knowledge bases."""

from typing import Optional, List, Dict, Any
import logging
from datetime import datetime

from lightrag.models.tenant import Tenant, KnowledgeBase, TenantConfig, KBConfig
from lightrag.base import BaseKVStorage

logger = logging.getLogger(__name__)


class TenantService:
    """Service for managing tenants and knowledge bases."""
    
    def __init__(self, kv_storage: BaseKVStorage):
        """Initialize tenant service with KV storage backend.
        
        Args:
            kv_storage: Backend storage for tenant/KB metadata
        """
        self.kv_storage = kv_storage
        self.tenant_namespace = "__tenants__"
        self.kb_namespace = "__knowledge_bases__"
    
    async def create_tenant(
        self,
        tenant_name: str,
        description: Optional[str] = None,
        config: Optional[TenantConfig] = None,
        created_by: Optional[str] = None,
        metadata: Optional[Dict[str, Any]] = None,
    ) -> Tenant:
        """Create a new tenant.
        
        Args:
            tenant_name: Display name for the tenant
            description: Optional description
            config: Optional tenant configuration
            created_by: User ID that created the tenant
            metadata: Optional metadata dictionary
            
        Returns:
            Created Tenant object
        """
        import json
        
        tenant = Tenant(
            tenant_name=tenant_name,
            description=description,
            config=config or TenantConfig(),
            created_by=created_by,
            metadata=metadata or {},
        )
        
        # Store tenant in PostgreSQL tenants table for FK integrity
        if hasattr(self.kv_storage, 'db') and self.kv_storage.db is not None:
            try:
                metadata_json = json.dumps(tenant.metadata) if tenant.metadata else '{}'
                # Use query method with RETURNING to insert tenant
                await self.kv_storage.db.query(
                    """
                    INSERT INTO tenants (tenant_id, name, description, metadata, created_at, updated_at)
                    VALUES ($1, $2, $3, $4::jsonb, NOW(), NOW())
                    ON CONFLICT (tenant_id) DO UPDATE SET
                        name = EXCLUDED.name,
                        description = EXCLUDED.description,
                        metadata = EXCLUDED.metadata,
                        updated_at = NOW()
                    RETURNING tenant_id
                    """,
                    [tenant.tenant_id, tenant_name, description or "", metadata_json]
                )
                logger.debug(f"Inserted tenant {tenant.tenant_id} into PostgreSQL tenants table")
            except Exception as e:
                logger.error(f"Failed to insert tenant into PostgreSQL: {e}")
                raise
        
        # Store tenant metadata in KV storage
        tenant_data = tenant.to_dict()
        await self.kv_storage.upsert({
            f"{self.tenant_namespace}:{tenant.tenant_id}": tenant_data
        })
        
        logger.info(f"Created tenant: {tenant.tenant_id} ({tenant_name})")
        return tenant
    
    async def get_tenant(self, tenant_id: str) -> Optional[Tenant]:
        """Retrieve a tenant by ID.
        
        Queries both PostgreSQL tenants table and KV storage to ensure
        tenants created via database initialization scripts are also available.
        
        Args:
            tenant_id: Tenant identifier
            
        Returns:
            Tenant object if found, None otherwise
        """
        # First, try to get tenant from PostgreSQL database if available
        if hasattr(self.kv_storage, 'db') and self.kv_storage.db is not None:
            try:
                logger.debug(f"Attempting to query tenant {tenant_id} from PostgreSQL database")
                # Query the tenants table directly (use $1 for PostgreSQL parameter)
                row = await self.kv_storage.db.query(
                    "SELECT tenant_id, name, description, created_at, updated_at FROM tenants WHERE tenant_id = $1",
                    [tenant_id]
                )
                
                if row:
                    # Create a Tenant object from the database row
                    tenant = Tenant(
                        tenant_id=row['tenant_id'],
                        tenant_name=row['name'],
                        description=row.get('description', ''),
                        created_by=None,  # Not tracked in basic schema
                        metadata={'is_public': True},  # Allow all users to access demo tenants
                    )
                    # Override timestamps from database
                    if 'created_at' in row:
                        tenant.created_at = row['created_at']
                    if 'updated_at' in row:
                        tenant.updated_at = row['updated_at']
                    
                    logger.debug(f"Retrieved tenant {tenant_id} from PostgreSQL database")
                    return tenant
            except Exception as e:
                logger.debug(f"Could not query tenant from PostgreSQL database: {e}")
                # Fall through to KV storage
        
        # Try KV storage as fallback
        data = await self.kv_storage.get_by_id(
            f"{self.tenant_namespace}:{tenant_id}"
        )
        if not data:
            return None
        return self._deserialize_tenant(data)
    
    async def verify_user_access(
        self,
        user_id: str,
        tenant_id: str,
        required_role: str = "viewer"
    ) -> bool:
        """Verify that a user has required role for a specific tenant.
        
        This is a CRITICAL security function that prevents unauthorized
        cross-tenant data access. Checks user-tenant membership table
        with role-based access control.
        
        Args:
            user_id: User identifier from JWT token
            tenant_id: Requested tenant ID
            required_role: Minimum required role (viewer, editor, admin, owner)
            
        Returns:
            True if user has access with required role, False otherwise
        """
        if not user_id or not tenant_id:
            logger.warning("verify_user_access called with empty user_id or tenant_id")
            return False
        
        # SEC-002 FIX: Check for super-admin users from configuration instead of hardcoded "admin"
        # Super-admins are configured via LIGHTRAG_SUPER_ADMIN_USERS environment variable
        super_admins_list = []
        try:
            from lightrag.api.config import SUPER_ADMIN_USERS
            if SUPER_ADMIN_USERS:
                super_admins_list = [u.strip().lower() for u in SUPER_ADMIN_USERS.split(",") if u.strip()]
        except ImportError:
            pass  # Config not available
            
        # Fallback: If no super admins configured, default to "admin" for backward compatibility
        # This ensures the default admin user always has access unless explicitly disabled
        if not super_admins_list:
            import os
            env_super_admins = os.environ.get("LIGHTRAG_SUPER_ADMIN_USERS")
            # If env var is not set, default to "admin". If set to empty string, it means "no super admins"
            if env_super_admins is None:
                super_admins_list = ["admin"]
            elif env_super_admins.strip():
                super_admins_list = [u.strip().lower() for u in env_super_admins.split(",") if u.strip()]
                
        if user_id.lower() in super_admins_list:
            logger.debug(f"Access granted: super-admin user {user_id} has access to all tenants")
            return True
        
        # Check membership table using PostgreSQL function
        if hasattr(self.kv_storage, 'db') and self.kv_storage.db:
            try:
                result = await self.kv_storage.db.query(
                    "SELECT has_tenant_access($1, $2, $3) as has_access",
                    [user_id, tenant_id, required_role]
                )
                # result is an asyncpg Record object (not a list when multirows=False)
                # Access using dict-style: result['has_access']
                if result is not None:
                    # Handle both dict-like Record and raw boolean result
                    if hasattr(result, '__getitem__'):
                        # Try dict-style access first (asyncpg Record)
                        try:
                            has_access = result['has_access']
                        except (KeyError, TypeError):
                            # Fall back to index access if key doesn't work
                            has_access = result[0] if len(result) > 0 else False
                    else:
                        # Direct boolean result
                        has_access = bool(result)
                    
                    if has_access:
                        logger.debug(f"Access granted: user {user_id} has {required_role}+ role for tenant {tenant_id}")
                        return True
                    else:
                        logger.warning(f"Access denied: user {user_id} lacks {required_role} role for tenant {tenant_id}")
                        return False
            except Exception as e:
                # Function might not exist if migration hasn't run - use legacy fallback
                error_msg = str(e)
                if "has_tenant_access" in error_msg and "does not exist" in error_msg:
                    logger.debug(f"has_tenant_access function not found, using legacy access check")
                else:
                    logger.warning(f"Error checking user access: {e}")
                # Fall through to legacy check
        
        # Legacy fallback: Check if tenant is public or user is creator
        tenant = await self.get_tenant(tenant_id)
        if not tenant:
            logger.debug(f"Tenant {tenant_id} not found during access check")
            return False
        
        # Check if tenant is public
        if tenant.metadata.get("is_public", False):
            logger.debug(f"Access granted: tenant {tenant_id} is public")
            return True
        
        # Check if user is the creator
        if tenant.created_by == user_id:
            logger.debug(f"Access granted: user {user_id} is creator of tenant {tenant_id}")
            return True
        
        logger.warning(
            f"Access denied: user {user_id} has no access to tenant {tenant_id}"
        )
        return False
    
    async def add_user_to_tenant(
        self,
        user_id: str,
        tenant_id: str,
        role: str,
        created_by: str
    ) -> dict:
        """Add a user to a tenant with specified role.
        
        Args:
            user_id: User identifier to add
            tenant_id: Tenant identifier
            role: User role (owner, admin, editor, viewer)
            created_by: User who is adding this member
            
        Returns:
            Dictionary with membership information
            
        Raises:
            ValueError: If tenant doesn't exist or role is invalid
        """
        # Validate role
        valid_roles = ['owner', 'admin', 'editor', 'viewer']
        if role not in valid_roles:
            raise ValueError(f"Invalid role: {role}. Must be one of {valid_roles}")
        
        # Verify tenant exists
        tenant = await self.get_tenant(tenant_id)
        if not tenant:
            raise ValueError(f"Tenant {tenant_id} not found")
        
        if not self.kv_storage.db:
            raise RuntimeError("PostgreSQL database required for membership management")
        
        try:
            # Insert membership - use multirows=True to get a list of Records
            results = await self.kv_storage.db.query(
                """
                INSERT INTO user_tenant_memberships (user_id, tenant_id, role, created_by)
                VALUES ($1, $2, $3, $4)
                ON CONFLICT (user_id, tenant_id) 
                DO UPDATE SET role = $3, updated_at = NOW()
                RETURNING id, user_id, tenant_id, role, created_at, created_by, updated_at
                """,
                [user_id, tenant_id, role, created_by],
                multirows=True
            )
            
            if results and len(results) > 0:
                membership = results[0]
                logger.info(f"Added user {user_id} to tenant {tenant_id} with role {role}")
                return {
                    "id": str(membership['id']),
                    "user_id": str(membership['user_id']),
                    "tenant_id": str(membership['tenant_id']),
                    "role": str(membership['role']),
                    "created_at": membership['created_at'].isoformat() if hasattr(membership['created_at'], 'isoformat') else str(membership['created_at']),
                    "created_by": str(membership['created_by']) if membership['created_by'] else None,
                    "updated_at": membership['updated_at'].isoformat() if hasattr(membership['updated_at'], 'isoformat') else str(membership['updated_at'])
                }
            else:
                raise RuntimeError("Failed to add user to tenant")
                
        except Exception as e:
            logger.error(f"Error adding user to tenant: {e}")
            raise
    
    async def remove_user_from_tenant(
        self,
        user_id: str,
        tenant_id: str
    ) -> bool:
        """Remove a user from a tenant.
        
        Args:
            user_id: User identifier to remove
            tenant_id: Tenant identifier
            
        Returns:
            True if removed, False if membership didn't exist
        """
        if not self.kv_storage.db:
            raise RuntimeError("PostgreSQL database required for membership management")
        
        try:
            results = await self.kv_storage.db.query(
                """
                DELETE FROM user_tenant_memberships
                WHERE user_id = $1 AND tenant_id = $2
                RETURNING id
                """,
                [user_id, tenant_id],
                multirows=True
            )
            
            if results and len(results) > 0:
                logger.info(f"Removed user {user_id} from tenant {tenant_id}")
                return True
            else:
                logger.debug(f"No membership found for user {user_id} in tenant {tenant_id}")
                return False
                
        except Exception as e:
            logger.error(f"Error removing user from tenant: {e}")
            raise
    
    async def update_user_role(
        self,
        user_id: str,
        tenant_id: str,
        new_role: str
    ) -> dict:
        """Update a user's role in a tenant.
        
        Args:
            user_id: User identifier
            tenant_id: Tenant identifier
            new_role: New role to assign
            
        Returns:
            Updated membership information
            
        Raises:
            ValueError: If role is invalid or membership doesn't exist
        """
        # Validate role
        valid_roles = ['owner', 'admin', 'editor', 'viewer']
        if new_role not in valid_roles:
            raise ValueError(f"Invalid role: {new_role}. Must be one of {valid_roles}")
        
        if not self.kv_storage.db:
            raise RuntimeError("PostgreSQL database required for membership management")
        
        try:
            results = await self.kv_storage.db.query(
                """
                UPDATE user_tenant_memberships
                SET role = $1, updated_at = NOW()
                WHERE user_id = $2 AND tenant_id = $3
                RETURNING id, user_id, tenant_id, role, created_at, created_by, updated_at
                """,
                [new_role, user_id, tenant_id],
                multirows=True
            )
            
            if results and len(results) > 0:
                membership = results[0]
                logger.info(f"Updated role for user {user_id} in tenant {tenant_id} to {new_role}")
                return {
                    "id": str(membership['id']),
                    "user_id": str(membership['user_id']),
                    "tenant_id": str(membership['tenant_id']),
                    "role": str(membership['role']),
                    "created_at": membership['created_at'].isoformat() if hasattr(membership['created_at'], 'isoformat') else str(membership['created_at']),
                    "created_by": str(membership['created_by']) if membership['created_by'] else None,
                    "updated_at": membership['updated_at'].isoformat() if hasattr(membership['updated_at'], 'isoformat') else str(membership['updated_at'])
                }
            else:
                raise ValueError(f"No membership found for user {user_id} in tenant {tenant_id}")
                
        except Exception as e:
            logger.error(f"Error updating user role: {e}")
            raise
    
    async def get_user_tenants(
        self,
        user_id: str,
        skip: int = 0,
        limit: int = 100
    ) -> dict:
        """Get all tenants a user has access to.
        
        Args:
            user_id: User identifier
            skip: Number of records to skip
            limit: Maximum number of records to return
            
        Returns:
            Dictionary with items and total count
        """
        if not self.kv_storage.db:
            # Fallback: return all public tenants
            all_tenants = await self.list_tenants(skip=skip, limit=limit)
            public_tenants = [
                t for t in all_tenants["items"] 
                if t.metadata.get("is_public", False)
            ]
            return {
                "items": public_tenants,
                "total": len(public_tenants),
                "skip": skip,
                "limit": limit
            }
        
        try:
            # Get tenants with user's membership
            result = await self.kv_storage.db.query(
                """
                SELECT t.*, m.role, m.created_at as member_since
                FROM tenants t
                INNER JOIN user_tenant_memberships m ON t.tenant_id = m.tenant_id
                WHERE m.user_id = $1
                ORDER BY t.created_at DESC
                LIMIT $2 OFFSET $3
                """,
                [user_id, limit, skip],
                multirows=True
            )
            
            # Get total count
            count_result = await self.kv_storage.db.query(
                """
                SELECT COUNT(*) as total
                FROM user_tenant_memberships
                WHERE user_id = $1
                """,
                [user_id]
            )
            
            # count_result is a single Record when multirows=False (default)
            total = count_result['total'] if count_result else 0
            
            tenants = []
            if result:
                for row in result:
                    tenant_dict = dict(row)
                    tenant_dict['user_role'] = tenant_dict.pop('role', None)
                    tenant_dict['member_since'] = tenant_dict.pop('member_since', None)
                    tenants.append(self._deserialize_tenant(tenant_dict))
            
            return {
                "items": tenants,
                "total": total,
                "skip": skip,
                "limit": limit
            }
            
        except Exception as e:
            logger.error(f"Error getting user tenants: {e}")
            raise
    
    async def get_tenant_members(
        self,
        tenant_id: str,
        skip: int = 0,
        limit: int = 100
    ) -> dict:
        """Get all members of a tenant.
        
        Args:
            tenant_id: Tenant identifier
            skip: Number of records to skip
            limit: Maximum number of records to return
            
        Returns:
            Dictionary with items and total count
        """
        if not self.kv_storage.db:
            raise RuntimeError("PostgreSQL database required for membership management")
        
        try:
            result = await self.kv_storage.db.query(
                """
                SELECT user_id, role, created_at, created_by, updated_at
                FROM user_tenant_memberships
                WHERE tenant_id = $1
                ORDER BY created_at DESC
                LIMIT $2 OFFSET $3
                """,
                [tenant_id, limit, skip],
                multirows=True
            )
            
            # Get total count
            count_result = await self.kv_storage.db.query(
                """
                SELECT COUNT(*) as total
                FROM user_tenant_memberships
                WHERE tenant_id = $1
                """,
                [tenant_id]
            )
            
            # count_result is a single Record when multirows=False (default)
            total = count_result['total'] if count_result else 0
            
            members = []
            if result:
                for row in result:
                    members.append({
                        "user_id": str(row['user_id']),
                        "role": str(row['role']),
                        "created_at": row['created_at'].isoformat() if hasattr(row['created_at'], 'isoformat') else str(row['created_at']),
                        "created_by": str(row['created_by']) if row['created_by'] else None,
                        "updated_at": row['updated_at'].isoformat() if row['updated_at'] and hasattr(row['updated_at'], 'isoformat') else (str(row['updated_at']) if row['updated_at'] else None)
                    })
            
            return {
                "items": members,
                "total": total,
                "skip": skip,
                "limit": limit
            }
            
        except Exception as e:
            logger.error(f"Error getting tenant members: {e}")
            raise

    
    async def update_tenant(
        self,
        tenant_id: str,
        **kwargs,
    ) -> Optional[Tenant]:
        """Update a tenant.
        
        Args:
            tenant_id: Tenant identifier
            **kwargs: Fields to update
            
        Returns:
            Updated Tenant object if found, None otherwise
        """
        tenant = await self.get_tenant(tenant_id)
        if not tenant:
            return None
        
        # Update fields
        for key, value in kwargs.items():
            if hasattr(tenant, key):
                setattr(tenant, key, value)
        
        tenant.updated_at = datetime.utcnow()
        
        # Store updated tenant
        tenant_data = tenant.to_dict()
        await self.kv_storage.upsert({
            f"{self.tenant_namespace}:{tenant_id}": tenant_data
        })
        
        logger.info(f"Updated tenant: {tenant_id}")
        return tenant
    
    async def list_tenants(
        self,
        skip: int = 0,
        limit: int = 100,
        search: Optional[str] = None,
        tenant_id_filter: Optional[str] = None,
    ) -> Dict[str, Any]:
        """List all tenants with pagination.
        
        Queries both KV storage and PostgreSQL database to ensure tenants
        created via database initialization scripts are also available.
        
        Args:
            skip: Number of tenants to skip (for pagination)
            limit: Maximum number of tenants to return
            search: Optional search string to filter by name or description
            tenant_id_filter: Optional tenant ID to filter by (for non-admin users)
            
        Returns:
            Dict with 'items' (list of tenants) and 'total' (count) keys
        """
        try:
            all_tenants = []
            
            # First, try to get tenants from PostgreSQL database if available
            if hasattr(self.kv_storage, 'db') and self.kv_storage.db is not None:
                try:
                    logger.debug("Attempting to query tenants from PostgreSQL database")
                    # Query tenants with computed statistics using LEFT JOINs
                    # This ensures we get real kb_count, total_documents, and storage from the database
                    # Note: Documents are stored in lightrag_doc_full with workspace={tenant_id}:{kb_id}
                    stats_query = """
                        SELECT 
                            t.tenant_id, 
                            t.name, 
                            t.description, 
                            t.created_at, 
                            t.updated_at,
                            COALESCE(kb_stats.kb_count, 0) as kb_count,
                            COALESCE(doc_stats.doc_count, 0) as total_documents,
                            COALESCE(doc_stats.total_size_bytes, 0) as total_size_bytes
                        FROM tenants t
                        LEFT JOIN (
                            SELECT tenant_id, COUNT(*) as kb_count 
                            FROM knowledge_bases 
                            GROUP BY tenant_id
                        ) kb_stats ON t.tenant_id = kb_stats.tenant_id
                        LEFT JOIN (
                            SELECT 
                                SPLIT_PART(workspace, ':', 1) as tenant_id,
                                COUNT(*) as doc_count, 
                                COALESCE(SUM(LENGTH(content)), 0) as total_size_bytes
                            FROM lightrag_doc_full 
                            GROUP BY SPLIT_PART(workspace, ':', 1)
                        ) doc_stats ON t.tenant_id = doc_stats.tenant_id
                        ORDER BY t.created_at DESC
                    """
                    rows = await self.kv_storage.db.query(
                        stats_query,
                        multirows=True
                    )
                    
                    if rows:
                        for row in rows:
                            try:
                                # Convert bytes to MB for storage
                                total_size_bytes = row.get('total_size_bytes', 0) or 0
                                total_storage_mb = total_size_bytes / (1024 * 1024)
                                
                                # Create a Tenant object from the database row with computed statistics
                                tenant = Tenant(
                                    tenant_id=row['tenant_id'],
                                    tenant_name=row['name'],
                                    description=row.get('description', ''),
                                    created_by=None,  # Not tracked in basic schema
                                    metadata={},
                                    kb_count=row.get('kb_count', 0) or 0,
                                    total_documents=row.get('total_documents', 0) or 0,
                                    total_storage_mb=total_storage_mb,
                                )
                                # Override timestamps from database
                                if 'created_at' in row:
                                    tenant.created_at = row['created_at']
                                if 'updated_at' in row:
                                    tenant.updated_at = row['updated_at']
                                
                                all_tenants.append(tenant)
                            except Exception as e:
                                logger.error(f"Error processing tenant row: {e}")
                                continue
                    
                    logger.info(f"Retrieved {len(all_tenants)} tenants from PostgreSQL database")
                except Exception as e:
                    logger.debug(f"Could not query tenants from PostgreSQL database: {e}")
                    # Fall through to KV storage
            
            # If no tenants from database, try KV storage
            if not all_tenants:
                logger.debug("Querying tenants from KV storage")
                tenant_keys = []
                if hasattr(self.kv_storage, 'get_by_prefix'):
                    # For storages that support prefix search
                    tenant_keys = await self.kv_storage.get_by_prefix(self.tenant_namespace)
                elif hasattr(self.kv_storage, 'get_all'):
                    # For storages like JsonKVStorage that have get_all
                    all_data = await self.kv_storage.get_all()
                    tenant_keys = [key for key in all_data.keys() if key.startswith(f"{self.tenant_namespace}:")]
                
                # Filter and deserialize tenants from KV storage
                for key in tenant_keys:
                    if not key.startswith(f"{self.tenant_namespace}:"):
                        continue
                    try:
                        data = await self.kv_storage.get_by_id(key)
                        if data:
                            tenant = self._deserialize_tenant(data)
                            
                            # Skip invalid tenants
                            if not tenant.tenant_id:
                                logger.warning(f"Skipping tenant with empty ID from key {key}")
                                continue
                            
                            all_tenants.append(tenant)
                    except Exception as e:
                        logger.error(f"Error deserializing tenant from key {key}: {e}")
                        continue
            
            # Apply filters
            filtered_tenants = []
            for tenant in all_tenants:
                # Apply tenant ID filter
                if tenant_id_filter and tenant.tenant_id != tenant_id_filter:
                    continue
                # Apply search filter
                if search:
                    search_lower = search.lower()
                    if not (search_lower in tenant.tenant_name.lower() or 
                            search_lower in (tenant.description or "").lower()):
                        continue
                filtered_tenants.append(tenant)
            
            # Sort by created_at descending
            filtered_tenants.sort(key=lambda t: t.created_at, reverse=True)
            
            # Apply pagination
            total = len(filtered_tenants)
            paginated_tenants = filtered_tenants[skip:skip + limit]
            
            logger.info(f"Listed {len(paginated_tenants)} tenants out of {total} (skip={skip}, limit={limit})")
            return {
                "items": paginated_tenants,
                "total": total
            }
        except Exception as e:
            logger.error(f"Error listing tenants: {e}")
            return {"items": [], "total": 0}
    
    async def delete_tenant(self, tenant_id: str) -> bool:
        """Delete a tenant.
        
        Args:
            tenant_id: Tenant identifier
            
        Returns:
            True if deleted, False if not found
        """
        tenant = await self.get_tenant(tenant_id)
        if not tenant:
            return False
        
        # Delete all KBs associated with tenant
        kbs_result = await self.list_knowledge_bases(tenant_id)
        kbs_list = kbs_result.get("items", [])
        for kb in kbs_list:
            await self.delete_knowledge_base(tenant_id, kb.kb_id)
        
        # Delete tenant
        await self.kv_storage.delete(
            [f"{self.tenant_namespace}:{tenant_id}"]
        )
        
        logger.info(f"Deleted tenant: {tenant_id}")
        return True
    
    async def create_knowledge_base(
        self,
        tenant_id: str,
        kb_name: str,
        description: Optional[str] = None,
        config: Optional[KBConfig] = None,
        created_by: Optional[str] = None,
    ) -> KnowledgeBase:
        """Create a new knowledge base for a tenant.
        
        Args:
            tenant_id: Parent tenant ID
            kb_name: Display name for KB
            description: Optional description
            config: Optional KB configuration
            created_by: User ID that created the KB
            
        Returns:
            Created KnowledgeBase object
            
        Raises:
            ValueError: If tenant not found
        """
        # Verify tenant exists
        tenant = await self.get_tenant(tenant_id)
        if not tenant:
            raise ValueError(f"Tenant {tenant_id} not found")
        
        kb = KnowledgeBase(
            tenant_id=tenant_id,
            kb_name=kb_name,
            description=description,
            config=config,
            created_by=created_by,
        )
        
        # Store KB metadata
        kb_data = kb.to_dict()
        await self.kv_storage.upsert({
            f"{self.kb_namespace}:{tenant_id}:{kb.kb_id}": kb_data
        })
        
        # Update tenant KB count
        tenant.kb_count += 1
        await self.update_tenant(tenant_id, kb_count=tenant.kb_count)
        
        logger.info(f"Created KB: {kb.kb_id} ({kb_name}) for tenant {tenant_id}")
        return kb
    
    async def get_knowledge_base(
        self,
        tenant_id: str,
        kb_id: str,
    ) -> Optional[KnowledgeBase]:
        """Retrieve a knowledge base.
        
        Args:
            tenant_id: Parent tenant ID
            kb_id: Knowledge base ID
            
        Returns:
            KnowledgeBase object if found, None otherwise
        """
        data = await self.kv_storage.get_by_id(
            f"{self.kb_namespace}:{tenant_id}:{kb_id}"
        )
        if not data:
            return None
        return self._deserialize_kb(data)
    
    async def update_knowledge_base(
        self,
        tenant_id: str,
        kb_id: str,
        **kwargs,
    ) -> Optional[KnowledgeBase]:
        """Update a knowledge base.
        
        Args:
            tenant_id: Parent tenant ID
            kb_id: Knowledge base ID
            **kwargs: Fields to update
            
        Returns:
            Updated KnowledgeBase object if found, None otherwise
        """
        kb = await self.get_knowledge_base(tenant_id, kb_id)
        if not kb:
            return None
        
        # Update fields
        for key, value in kwargs.items():
            if hasattr(kb, key):
                setattr(kb, key, value)
        
        kb.updated_at = datetime.utcnow()
        
        # Store updated KB
        kb_data = kb.to_dict()
        await self.kv_storage.upsert({
            f"{self.kb_namespace}:{tenant_id}:{kb_id}": kb_data
        })
        
        logger.info(f"Updated KB: {kb_id} for tenant {tenant_id}")
        return kb
    
    async def list_knowledge_bases(
        self,
        tenant_id: str,
        skip: int = 0,
        limit: int = 100,
        search: Optional[str] = None,
    ) -> Dict[str, Any]:
        """List all knowledge bases for a tenant with pagination.
        
        Queries both KV storage and PostgreSQL database to ensure KBs
        created via database initialization scripts are also available.
        
        Args:
            tenant_id: Parent tenant ID
            skip: Number of KBs to skip (for pagination)
            limit: Maximum number of KBs to return
            search: Optional search string to filter by name or description
            
        Returns:
            Dict with 'items' (list of KBs) and 'total' (count) keys
        """
        try:
            all_kbs = []
            
            # First, try to get KBs from PostgreSQL database if available
            db_queried = False
            if hasattr(self.kv_storage, 'db'):
                logger.info(f"PGKVStorage.db exists, attempting to query KBs for tenant {tenant_id}")
                if self.kv_storage.db is not None:
                    try:
                        logger.info(f"Querying knowledge bases from PostgreSQL for tenant {tenant_id}")
                        # Query the knowledge_bases table directly
                        rows = await self.kv_storage.db.query(
                            "SELECT kb_id, tenant_id, name, description, created_at, updated_at FROM knowledge_bases WHERE tenant_id = $1 ORDER BY created_at DESC",
                            params=[tenant_id],
                            multirows=True
                        )
                        
                        if rows:
                            for row in rows:
                                try:
                                    # Create a KnowledgeBase object from the database row
                                    kb = KnowledgeBase(
                                        kb_id=row['kb_id'],
                                        tenant_id=row['tenant_id'],
                                        kb_name=row['name'],
                                        description=row.get('description', ''),
                                    )
                                    # Override timestamps from database
                                    if 'created_at' in row:
                                        kb.created_at = row['created_at']
                                    if 'updated_at' in row:
                                        kb.updated_at = row['updated_at']
                                    
                                    all_kbs.append(kb)
                                except Exception as e:
                                    logger.error(f"Error processing KB row: {e}")
                                    continue
                        
                        logger.info(f"Retrieved {len(all_kbs)} knowledge bases from PostgreSQL for tenant {tenant_id}")
                        db_queried = True
                    except Exception as e:
                        logger.warning(f"Could not query knowledge bases from PostgreSQL: {e}")
                        # Fall through to KV storage
                else:
                    logger.debug(f"PGKVStorage.db is None for tenant {tenant_id}")
            else:
                logger.debug(f"Storage doesn't have 'db' attribute (type: {type(self.kv_storage).__name__})")
            
            # If no KBs from database, try KV storage
            if not all_kbs and not db_queried:
                logger.info(f"Querying knowledge bases from KV storage for tenant {tenant_id}")
                kb_keys = []
                if hasattr(self.kv_storage, 'get_by_prefix'):
                    # For storages that support prefix search
                    tenant_prefix = f"{self.kb_namespace}:{tenant_id}:"
                    kb_keys = await self.kv_storage.get_by_prefix(tenant_prefix)
                elif hasattr(self.kv_storage, 'get_all'):
                    # For storages like JsonKVStorage that have get_all
                    all_data = await self.kv_storage.get_all()
                    kb_keys = [key for key in all_data.keys() if key.startswith(f"{self.kb_namespace}:{tenant_id}:")]
                
                # Filter and deserialize KBs from KV storage
                for key in kb_keys:
                    if not key.startswith(f"{self.kb_namespace}:{tenant_id}:"):
                        continue
                    try:
                        data = await self.kv_storage.get_by_id(key)
                        if data:
                            kb = self._deserialize_kb(data)
                            all_kbs.append(kb)
                    except Exception as e:
                        logger.error(f"Error deserializing KB from key {key}: {e}")
                        continue
            
            # Apply search filter
            filtered_kbs = []
            for kb in all_kbs:
                if search:
                    search_lower = search.lower()
                    if not (search_lower in kb.kb_name.lower() or 
                            search_lower in (kb.description or "").lower()):
                        continue
                filtered_kbs.append(kb)
            
            # Sort by created_at descending
            filtered_kbs.sort(key=lambda k: k.created_at, reverse=True)
            
            # Apply pagination
            total = len(filtered_kbs)
            paginated_kbs = filtered_kbs[skip:skip + limit]
            
            logger.info(f"Listed {len(paginated_kbs)} KBs out of {total} for tenant {tenant_id} (skip={skip}, limit={limit})")
            return {
                "items": paginated_kbs,
                "total": total
            }
        except Exception as e:
            logger.error(f"Error listing KBs for tenant {tenant_id}: {e}")
            return {"items": [], "total": 0}
    
    async def delete_knowledge_base(
        self,
        tenant_id: str,
        kb_id: str,
    ) -> bool:
        """Delete a knowledge base.
        
        Args:
            tenant_id: Parent tenant ID
            kb_id: Knowledge base ID
            
        Returns:
            True if deleted, False if not found
        """
        kb = await self.get_knowledge_base(tenant_id, kb_id)
        if not kb:
            return False
        
        # Delete KB
        await self.kv_storage.delete(
            [f"{self.kb_namespace}:{tenant_id}:{kb_id}"]
        )
        
        # Update tenant KB count
        tenant = await self.get_tenant(tenant_id)
        if tenant:
            tenant.kb_count = max(0, tenant.kb_count - 1)
            await self.update_tenant(tenant_id, kb_count=tenant.kb_count)
        
        logger.info(f"Deleted KB: {kb_id} for tenant {tenant_id}")
        return True
    
    def _deserialize_tenant(self, data: Dict[str, Any]) -> Tenant:
        """Convert stored data to Tenant object."""
        import json
        
        # Handle PGKVStorage wrapping
        if "data" in data:
            inner_data = data["data"]
            if isinstance(inner_data, str):
                try:
                    inner_data = json.loads(inner_data)
                except json.JSONDecodeError:
                    logger.warning(f"Failed to decode JSON from tenant data: {inner_data}")
            
            if isinstance(inner_data, dict) and "tenant_id" in inner_data:
                data = inner_data

        if not data.get("tenant_id"):
            logger.warning(f"Deserializing tenant with missing ID. Data keys: {list(data.keys())}")
            
        config_data = data.get("config", {})
        quota_data = data.get("quota", {})
        
        config = TenantConfig(
            llm_model=config_data.get("llm_model", "gpt-4o-mini"),
            embedding_model=config_data.get("embedding_model", "bge-m3:latest"),
            rerank_model=config_data.get("rerank_model"),
            chunk_size=config_data.get("chunk_size", 1200),
            chunk_overlap=config_data.get("chunk_overlap", 100),
            top_k=config_data.get("top_k", 40),
            cosine_threshold=config_data.get("cosine_threshold", 0.2),
            enable_llm_cache=config_data.get("enable_llm_cache", True),
            custom_metadata=config_data.get("custom_metadata", {}),
        )
        
        # Create and return tenant
        tenant = Tenant(
            tenant_id=data.get("tenant_id", ""),
            tenant_name=data.get("tenant_name", ""),
            description=data.get("description"),
            config=config,
            is_active=data.get("is_active", True),
            created_at=datetime.fromisoformat(data.get("created_at")) if data.get("created_at") else datetime.utcnow(),
            updated_at=datetime.fromisoformat(data.get("updated_at")) if data.get("updated_at") else datetime.utcnow(),
            created_by=data.get("created_by"),
            updated_by=data.get("updated_by"),
            metadata=data.get("metadata", {}),
            kb_count=data.get("kb_count", 0),
            total_documents=data.get("total_documents", 0),
            total_storage_mb=data.get("total_storage_mb", 0.0),
        )
        return tenant
    
    def _deserialize_kb(self, data: Dict[str, Any]) -> KnowledgeBase:
        """Convert stored data to KnowledgeBase object."""
        import json

        # Handle PGKVStorage wrapping
        if "data" in data:
            inner_data = data["data"]
            if isinstance(inner_data, str):
                try:
                    inner_data = json.loads(inner_data)
                except json.JSONDecodeError:
                    logger.warning(f"Failed to decode JSON from KB data: {inner_data}")
            
            if isinstance(inner_data, dict) and "kb_id" in inner_data:
                data = inner_data

        config_data = data.get("config")
        config = KBConfig(**config_data) if config_data else None
        
        kb = KnowledgeBase(
            kb_id=data.get("kb_id", ""),
            tenant_id=data.get("tenant_id", ""),
            kb_name=data.get("kb_name", ""),
            description=data.get("description"),
            is_active=data.get("is_active", True),
            status=data.get("status", "ready"),
            document_count=data.get("document_count", 0),
            entity_count=data.get("entity_count", 0),
            relationship_count=data.get("relationship_count", 0),
            chunk_count=data.get("chunk_count", 0),
            storage_used_mb=data.get("storage_used_mb", 0.0),
            last_indexed_at=datetime.fromisoformat(data.get("last_indexed_at")) if data.get("last_indexed_at") else None,
            index_version=data.get("index_version", 1),
            config=config,
            created_at=datetime.fromisoformat(data.get("created_at")) if data.get("created_at") else datetime.utcnow(),
            updated_at=datetime.fromisoformat(data.get("updated_at")) if data.get("updated_at") else datetime.utcnow(),
            created_by=data.get("created_by"),
            updated_by=data.get("updated_by"),
            metadata=data.get("metadata", {}),
        )
        return kb