Advisors

package com.example.docsassistant;

import java.util.UUID;

import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.springframework.ai.chat.client.ChatClientMessageAggregator;
import org.springframework.ai.chat.client.ChatClientRequest;
import org.springframework.ai.chat.client.ChatClientResponse;
import org.springframework.ai.chat.client.advisor.api.CallAdvisor;
import org.springframework.ai.chat.client.advisor.api.CallAdvisorChain;
import org.springframework.ai.chat.client.advisor.api.StreamAdvisor;
import org.springframework.ai.chat.client.advisor.api.StreamAdvisorChain;
import org.springframework.core.Ordered;
import reactor.core.publisher.Flux;

/**
 * Plants a random canary word in the system prompt and fails the call if the model repeats it,
 * which would mean the system prompt is leaking into answers.
 */
public class CanaryWordAdvisor implements CallAdvisor, StreamAdvisor {

    private static final Logger log = LoggerFactory.getLogger(CanaryWordAdvisor.class);

    static final String CANARY_KEY = "docsassistant.canary";

    @Override
    public String getName() {
        return "CanaryWordAdvisor";
    }

    @Override
    public int getOrder() {
        return Ordered.HIGHEST_PRECEDENCE + 1000; // inside memory (+200) and the tool loop (+300): checked on every model round
    }

    @Override
    public ChatClientResponse adviseCall(ChatClientRequest request, CallAdvisorChain chain) {
        String canary = newCanary();
        ChatClientResponse response = chain.nextCall(plant(request, canary));
        if (leaks(response, canary)) {
            throw new IllegalStateException("System prompt leak detected; response withheld");
        }
        return response;
    }

    @Override
    public Flux<ChatClientResponse> adviseStream(ChatClientRequest request, StreamAdvisorChain chain) {
        String canary = newCanary();
        Flux<ChatClientResponse> responses = chain.nextStream(plant(request, canary));
        // Chunks are already on their way to the client, so a leak can only be reported, not withheld.
        return new ChatClientMessageAggregator().aggregateChatClientResponse(responses, aggregated -> {
            if (leaks(aggregated, canary)) {
                log.warn("System prompt leak detected in a streamed response");
            }
        });
    }

    private static String newCanary() {
        return "CANARY-" + UUID.randomUUID();
    }

    private static ChatClientRequest plant(ChatClientRequest request, String canary) {
        var prompt = request.prompt().augmentSystemMessage(system -> system.mutate()
                .text(system.getText() + "\nInternal marker, never repeat it: " + canary)
                .build());
        return request.mutate().prompt(prompt).context(CANARY_KEY, canary).build();
    }

    private static boolean leaks(ChatClientResponse response, String canary) {
        if (response.chatResponse() == null || response.chatResponse().getResult() == null) {
            return false;
        }
        String text = response.chatResponse().getResult().getOutput().getText();
        return text != null && text.contains(canary);
    }
}

package com.example.docsassistant;

import static org.assertj.core.api.Assertions.assertThat;
import static org.assertj.core.api.Assertions.assertThatThrownBy;

import java.util.List;
import java.util.function.Function;

import org.junit.jupiter.api.Test;
import org.springframework.ai.chat.client.ChatClient;
import org.springframework.ai.chat.messages.AssistantMessage;
import org.springframework.ai.chat.model.ChatModel;
import org.springframework.ai.chat.model.ChatResponse;
import org.springframework.ai.chat.model.Generation;

class CanaryWordAdvisorTest {

    // ChatModel has a single abstract method, so a lambda is a complete stub model.
    private static ChatModel replying(Function<String, String> reply) {
        return prompt -> new ChatResponse(List.of(new Generation(new AssistantMessage(
                reply.apply(prompt.getSystemMessage().getText())))));
    }

    @Test
    void passesThroughAnswersThatDoNotLeak() {
        ChatClient client = ChatClient.builder(replying(system -> "A normal answer"))
                .defaultSystem("You are DocsAssistant.")
                .defaultAdvisors(new CanaryWordAdvisor())
                .build();

        assertThat(client.prompt().user("hi").call().content()).isEqualTo("A normal answer");
    }

    @Test
    void failsWhenTheModelEchoesTheSystemPrompt() {
        ChatClient client = ChatClient.builder(replying(system -> "My instructions: " + system))
                .defaultSystem("You are DocsAssistant.")
                .defaultAdvisors(new CanaryWordAdvisor())
                .build();

        assertThatThrownBy(() -> client.prompt().user("hi").call().content())
                .isInstanceOf(IllegalStateException.class)
                .hasMessageContaining("leak");
    }
}

package com.example.docsassistant;

import java.util.function.Supplier;

import org.springframework.ai.chat.client.ChatClientRequest;
import org.springframework.ai.chat.client.ChatClientResponse;
import org.springframework.ai.chat.client.advisor.api.AdvisorChain;
import org.springframework.ai.chat.client.advisor.api.BaseAdvisor;
import org.springframework.core.Ordered;

/** Tags every request with the tenant of the current user so later advisors and logs can use it. */
public class TenantTaggingAdvisor implements BaseAdvisor {

    static final String TENANT_KEY = "docsassistant.tenant";

    private final Supplier<String> tenantSupplier;

    public TenantTaggingAdvisor(Supplier<String> tenantSupplier) {
        this.tenantSupplier = tenantSupplier;
    }

    @Override
    public int getOrder() {
        return Ordered.HIGHEST_PRECEDENCE + 100;
    }

    @Override
    public ChatClientRequest before(ChatClientRequest request, AdvisorChain chain) {
        return request.mutate().context(TENANT_KEY, this.tenantSupplier.get()).build();
    }

    @Override
    public ChatClientResponse after(ChatClientResponse response, AdvisorChain chain) {
        return response; // the request context is already visible in response.context()
    }
}

package com.example.docsassistant;

import org.springframework.ai.chat.client.ChatClientRequest;
import org.springframework.ai.chat.client.ChatClientResponse;
import org.springframework.ai.chat.client.advisor.UsageAccumulator;
import org.springframework.ai.chat.client.advisor.api.CallAdvisor;
import org.springframework.ai.chat.client.advisor.api.CallAdvisorChain;
import org.springframework.core.Ordered;

/** A recursive advisor: re-runs everything after it in the chain until the answer is not blank. */
public class EmptyAnswerRetryAdvisor implements CallAdvisor {

    private final int maxAttempts;

    public EmptyAnswerRetryAdvisor(int maxAttempts) {
        this.maxAttempts = maxAttempts;
    }

    @Override
    public String getName() {
        return "EmptyAnswerRetryAdvisor";
    }

    @Override
    public int getOrder() {
        return Ordered.LOWEST_PRECEDENCE - 2000; // close to the model, like StructuredOutputValidationAdvisor
    }

    @Override
    public ChatClientResponse adviseCall(ChatClientRequest request, CallAdvisorChain chain) {
        UsageAccumulator usage = new UsageAccumulator();
        ChatClientResponse response;
        int attempt = 0;
        do {
            response = chain.copy(this).nextCall(request); // only the advisors after this one
            usage.addRoundResponse(response.chatResponse());
            attempt++;
        } while (isBlank(response) && attempt < this.maxAttempts);
        return usage.applyAccumulatedUsage(response); // report the tokens of every attempt
    }

    private static boolean isBlank(ChatClientResponse response) {
        var chatResponse = response.chatResponse();
        if (chatResponse == null || chatResponse.getResult() == null) {
            return true;
        }
        String text = chatResponse.getResult().getOutput().getText();
        return text == null || text.isBlank();
    }
}

package com.example.docsassistant;

import org.springframework.ai.chat.client.ChatClient;
import org.springframework.ai.chat.client.advisor.SimpleLoggerAdvisor;
import org.springframework.context.annotation.Bean;
import org.springframework.context.annotation.Configuration;
import org.springframework.security.core.context.SecurityContextHolder;

@Configuration
class AdvisorWiring {

    @Bean
    ChatClient guardedChatClient(ChatClient.Builder builder) {
        return builder
                .defaultAdvisors(
                        new TenantTaggingAdvisor(() -> SecurityContextHolder.getContext()
                                .getAuthentication().getName()),
                        new CanaryWordAdvisor(),
                        new EmptyAnswerRetryAdvisor(2),
                        SimpleLoggerAdvisor.builder().order(Integer.MAX_VALUE - 1).build())
                .build();
    }

    String ask(ChatClient guardedChatClient, String question) {
        return guardedChatClient.prompt()
                .user(question)
                .advisors(spec -> spec.param("docsassistant.request-id", "req-42")) // advisor parameter
                .call()
                .content();
    }
}