spring-ai/batch-llm/src/main/java/com/ankurm/batchllm/SimulatedProvider.java

package com.ankurm.batchllm;

import java.util.ArrayDeque;
import java.util.List;
import java.util.concurrent.ConcurrentHashMap;
import java.util.concurrent.atomic.AtomicInteger;
import java.util.regex.Matcher;
import java.util.regex.Pattern;

import org.springframework.ai.chat.messages.AssistantMessage;
import org.springframework.ai.chat.metadata.ChatResponseMetadata;
import org.springframework.ai.chat.metadata.DefaultUsage;
import org.springframework.ai.chat.model.ChatModel;
import org.springframework.ai.chat.model.ChatResponse;
import org.springframework.ai.chat.model.Generation;
import org.springframework.ai.chat.prompt.Prompt;

import com.ankurm.batchllm.ProviderMeter.Slot;

/**
 * A hand-written, deterministic {@link ChatModel} -- the same idea as the ScriptedChatModel in the
 * other modules of this repository, but driven by a fault plan instead of a list of replies.
 * It is NOT a language model: it reads the ticket number out of the prompt and decides, from a
 * hash of (seed, ticket, call number), whether this call returns valid JSON, malformed JSON, a
 * 429, a 503 or a slow answer. No network, no key, no vendor.
 *
 * <p>It also enforces its own quota: more than {@code limitPerSecond} calls in any sliding second
 * get a 429, which is what a client-side rate limiter exists to avoid.
 *
 * <p>Token counts are ESTIMATED as characters / 4 and returned through the same
 * {@code ChatResponseMetadata.usage} field a real provider fills in.
 */
public class SimulatedProvider implements ChatModel {

    private static final Pattern TICKET_ID = Pattern.compile("Ticket #(\\d+)");
    private static final List<String> CATEGORIES = List.of("BILLING", "TECHNICAL", "ACCOUNT", "SHIPPING", "FEATURE_REQUEST", "OTHER");

    private final ProviderFaults faults;
    private final ProviderMeter meter;
    private final ConcurrentHashMap<Long, AtomicInteger> callsPerTicket = new ConcurrentHashMap<>();
    private final ConcurrentHashMap<Long, AtomicInteger> billedPerTicket = new ConcurrentHashMap<>();
    private final ArrayDeque<Long> window = new ArrayDeque<>();
    private volatile int latencyDivisor = 1;

    public SimulatedProvider(ProviderFaults faults, ProviderMeter meter) {
        this.faults = faults;
        this.meter = meter;
    }

    /** Test hook: make every sleep this many times shorter. */
    public void setLatencyDivisor(int divisor) {
        this.latencyDivisor = Math.max(1, divisor);
    }

    public ProviderMeter meter() {
        return meter;
    }

    /** Total calls the model has seen for one ticket, rejected ones included. */
    public int callsFor(long ticketId) {
        AtomicInteger n = callsPerTicket.get(ticketId);
        return n == null ? 0 : n.get();
    }

    private double unit(long ticket, long salt) {
        long h = TicketGenerator.mix(TicketGenerator.mix(faults.seed() * 31 + salt) ^ ticket);
        return (h >>> 11) * 0x1.0p-53;
    }

    @Override
    public ChatResponse call(Prompt prompt) {
        String text = prompt.getContents();
        Matcher m = TICKET_ID.matcher(text);
        if (!m.find()) {
            throw new IllegalArgumentException("simulated provider expects 'Ticket #<id>' in the prompt");
        }
        long id = Long.parseLong(m.group(1));
        int callNo = callsPerTicket.computeIfAbsent(id, k -> new AtomicInteger()).incrementAndGet();
        meter.add(Slot.RECEIVED, 1);

        if (overQuota()) {
            meter.add(Slot.REJECTED_429, 1);
            meter.add(Slot.WINDOW_429, 1);
            throw new ProviderRateLimitException("429 Too Many Requests (quota: " + faults.limitPerSecond() + " calls per second)");
        }
        if (unit(id, 1000L + callNo) < faults.random429()) {
            meter.add(Slot.REJECTED_429, 1);
            throw new ProviderRateLimitException("429 Too Many Requests (injected)");
        }
        if (unit(id, 2000L + callNo) < faults.random503()) {
            meter.add(Slot.UNAVAILABLE_503, 1);
            throw new ProviderUnavailableException("503 Service Unavailable (injected)");
        }

        boolean slow = unit(id, 3000L + callNo) < faults.slowShare();
        long millis = slow ? faults.slowMillis() : 1 + (long) (unit(id, 4000L + callNo) * 3);
        try {
            Thread.sleep(Math.max(0, millis / latencyDivisor));
        } catch (InterruptedException e) {
            Thread.currentThread().interrupt();
            throw new IllegalStateException(e);
        }

        int billedNo = billedPerTicket.computeIfAbsent(id, k -> new AtomicInteger()).incrementAndGet();
        String reply = reply(id, billedNo);
        int promptTokens = (text.length() + 3) / 4;
        int completionTokens = (reply.length() + 3) / 4;
        meter.add(Slot.BILLED, 1);
        meter.add(Slot.PROMPT_TOKENS, promptTokens);
        meter.add(Slot.COMPLETION_TOKENS, completionTokens);

        ChatResponseMetadata metadata = ChatResponseMetadata.builder()
                .model("simulated-classifier")
                .usage(new DefaultUsage(promptTokens, completionTokens))
                .build();
        return new ChatResponse(List.of(new Generation(new AssistantMessage(reply))), metadata);
    }

    private synchronized boolean overQuota() {
        long now = System.nanoTime();
        while (!window.isEmpty() && now - window.peekFirst() > 1_000_000_000L) {
            window.pollFirst();
        }
        if (window.size() >= faults.limitPerSecond()) {
            return true;
        }
        window.addLast(now);
        return false;
    }

    /** The reply for ticket {@code id} on its {@code billedNo}-th billed answer. */
    String reply(long id, int billedNo) {
        String category = CATEGORIES.get((int) Math.floorMod(TicketGenerator.mix(id * 1_000_003L + 17), 6L));
        int priority = 1 + (int) (unit(id, 5000L) * 4);
        double confidence = Math.round((0.60 + unit(id, 6000L) * 0.39) * 100) / 100.0;
        String good = "{\"category\":\"" + category + "\",\"priority\":" + priority + ",\"confidence\":" + confidence + "}";
        boolean permanent = unit(id, 7000L) < faults.permanentBad();
        boolean flaky = !permanent && billedNo == 1 && unit(id, 8000L) < faults.flakyBad();
        if (!permanent && !flaky) {
            // one reply in ten comes wrapped in a markdown fence: valid after the fence is stripped
            return unit(id, 9000L) < 0.10 ? "```json\n" + good + "\n```" : good;
        }
        return switch ((int) (unit(id, 10_000L) * 4)) {
            case 0 -> good.substring(0, good.length() - 14);                       // truncated
            case 1 -> "Sure! Here is the classification you asked for: " + good;   // prose in front
            case 2 -> good.replace(category, category.charAt(0) + category.substring(1).toLowerCase() + " & Payments"); // enum value not in the schema
            default -> "{\"category\":\"" + category + "\",\"priority\":" + priority + "}"; // field missing
        };
    }
}