Access Control & Privacy

import os
from dataclasses import dataclass, field

import jwt

JWKS = jwt.PyJWKClient(os.environ["OIDC_JWKS_URL"])
ISSUER, AUDIENCE = os.environ["OIDC_ISSUER"], "docsassistant-api"

# group -> ACL tags readable by members; versioned with the service
POLICY: dict[str, set[str]] = {
    "staff": {"public", "internal"},
    "backend-team": {"backend"},
    "security-team": {"security"},
}


@dataclass(frozen=True)
class Caller:
    user_id: str
    readable_acls: list[str] = field(default_factory=lambda: ["public"])


def readable_acls(groups: list[str]) -> list[str]:
    tags = {"public"}
    for group in groups:
        tags |= POLICY.get(group, set())
    return sorted(tags)


def caller_from_token(token: str | None) -> Caller:
    if not token:                                    # no shared "anonymous" identity: the API answers 401
        raise jwt.InvalidTokenError("missing bearer token")
    key = JWKS.get_signing_key_from_jwt(token)
    claims = jwt.decode(token, key.key, algorithms=["RS256", "ES256"], audience=AUDIENCE, issuer=ISSUER)
    return Caller(user_id=claims["sub"], readable_acls=readable_acls(claims.get("groups", [])))


def acl_filter(caller: Caller, extra: dict | None = None) -> dict:
    """PGVector filter; the ACL clause comes only from the verified caller."""
    clauses = [{"acl": {"$in": caller.readable_acls}}]
    if extra:
        if "acl" in extra:
            raise ValueError("Client filters must not name the acl field")
        clauses.append(extra)
    return {"$and": clauses} if len(clauses) > 1 else clauses[0]

package com.example.docsassistant.rag.acl;

import java.util.List;
import java.util.Map;
import java.util.Set;
import java.util.TreeSet;

import org.springframework.ai.chat.client.ChatClient;
import org.springframework.ai.rag.advisor.RetrievalAugmentationAdvisor;
import org.springframework.ai.rag.retrieval.search.VectorStoreDocumentRetriever;
import org.springframework.ai.vectorstore.VectorStore;
import org.springframework.ai.vectorstore.filter.Filter;
import org.springframework.ai.vectorstore.filter.FilterExpressionBuilder;
import org.springframework.security.oauth2.jwt.Jwt;
import org.springframework.stereotype.Service;

@Service
public class AclAwareAskService {

    static final Map<String, Set<String>> POLICY = Map.of(
            "staff", Set.of("public", "internal"),
            "backend-team", Set.of("backend"),
            "security-team", Set.of("security"));

    private final ChatClient chatClient;

    public AclAwareAskService(ChatClient.Builder builder, VectorStore vectorStore) {
        this.chatClient = builder.defaultAdvisors(RetrievalAugmentationAdvisor.builder()
                        .documentRetriever(VectorStoreDocumentRetriever.builder()
                                .vectorStore(vectorStore)
                                .filterExpression(new FilterExpressionBuilder().eq("acl", "public").build()) // default
                                .topK(6)
                                .build())
                        .build())
                .build();
    }

    public static List<String> readableAcls(Jwt jwt) {
        Set<String> tags = new TreeSet<>(Set.of("public"));
        List<String> groups = jwt.getClaimAsStringList("groups");
        if (groups != null) {
            groups.forEach(g -> tags.addAll(POLICY.getOrDefault(g, Set.of())));
        }
        return List.copyOf(tags);
    }

    static Filter.Expression aclFilter(Jwt jwt, String component) {
        var b = new FilterExpressionBuilder();
        var acl = b.in("acl", readableAcls(jwt).toArray());
        return component == null ? acl.build() : b.and(acl, b.eq("component", component)).build();
    }

    public String ask(Jwt caller, String question, String component) {
        return chatClient.prompt()
                .user(question)
                .advisors(a -> a.param(VectorStoreDocumentRetriever.FILTER_EXPRESSION, aclFilter(caller, component)))
                .call()
                .content();
    }
}

package com.example.docsassistant.l4j.acl;

import static dev.langchain4j.store.embedding.filter.MetadataFilterBuilder.metadataKey;

import java.util.List;

import dev.langchain4j.data.segment.TextSegment;
import dev.langchain4j.invocation.InvocationParameters;
import dev.langchain4j.model.chat.ChatModel;
import dev.langchain4j.model.embedding.EmbeddingModel;
import dev.langchain4j.rag.content.retriever.EmbeddingStoreContentRetriever;
import dev.langchain4j.service.AiServices;
import dev.langchain4j.service.UserMessage;
import dev.langchain4j.store.embedding.EmbeddingStore;
import dev.langchain4j.store.embedding.filter.Filter;

public final class AclAwareAssistant {

    public interface DocsAssistant {
        String ask(@UserMessage String question, InvocationParameters caller);
    }

    private AclAwareAssistant() {
    }

    static Filter aclFilter(InvocationParameters params) {
        List<String> readable = params == null ? List.of("public") : params.getOrDefault("readableAcls", List.of("public"));
        return metadataKey("acl").isIn(readable);
    }

    public static DocsAssistant create(ChatModel model, EmbeddingModel embeddings, EmbeddingStore<TextSegment> store) {
        var retriever = EmbeddingStoreContentRetriever.builder()
                .embeddingStore(store)
                .embeddingModel(embeddings)
                .maxResults(6)
                .dynamicFilter(query -> aclFilter(query.metadata().invocationParameters()))
                .build();
        return AiServices.builder(DocsAssistant.class)
                .chatModel(model)
                .contentRetriever(retriever)
                .build();
    }

    /** Called by the web layer after verifying the token. */
    public static InvocationParameters forCaller(List<String> readableAcls) {
        return InvocationParameters.from("readableAcls", readableAcls);
    }
}