Access Control

# Document-level access control in RAG
from typing import List, Dict, Optional

class SecureRAGRetriever:
    def __init__(self, vector_store, permission_service):
        self.vector_store = vector_store
        self.permissions = permission_service
    
    def retrieve(
        self, 
        query: str, 
        user_id: str, 
        top_k: int = 5
    ) -> List[Dict]:
        """Retrieve documents respecting user permissions"""
        
        # Get user's accessible document IDs and roles
        user_permissions = self.permissions.get_user_access(user_id)
        allowed_doc_ids = user_permissions["document_ids"]
        user_roles = user_permissions["roles"]
        user_clearance = user_permissions["clearance_level"]
        
        # Build metadata filter for vector search
        metadata_filter = {
            "$or": [
                {"document_id": {"$in": allowed_doc_ids}},
                {"allowed_roles": {"$in": user_roles}},
                {"visibility": "public"}
            ],
            "classification_level": {"$lte": user_clearance}
        }
        
        # Search with access filter
        results = self.vector_store.similarity_search(
            query=query,
            k=top_k,
            filter=metadata_filter
        )
        
        # Log access for audit trail
        self._log_access(user_id, query, [r.id for r in results])
        
        return results

# Role-based tool access for AI agents
from typing import List, Callable
from functools import wraps

# Define tool permissions per role
TOOL_PERMISSIONS = {
    "viewer": ["search_docs", "get_info"],
    "editor": ["search_docs", "get_info", "update_record"],
    "admin": ["search_docs", "get_info", "update_record", "delete_record", "manage_users"]
}

def require_permission(tool_name: str):
    """Decorator to enforce tool access based on user role"""
    def decorator(func: Callable):
        @wraps(func)
        def wrapper(*args, user_context: dict, **kwargs):
            user_role = user_context.get("role", "viewer")
            allowed_tools = TOOL_PERMISSIONS.get(user_role, [])
            
            if tool_name not in allowed_tools:
                raise PermissionError(
                    f"User role '{user_role}' cannot access tool '{tool_name}'"
                )
            
            return func(*args, **kwargs)
        return wrapper
    return decorator

# Example tool with permission check
@require_permission("delete_record")
def delete_record(record_id: str) -> dict:
    # Only admins can execute this
    return {"deleted": record_id}