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

package com.ankurm.batchllm;

import java.math.BigDecimal;
import java.util.List;
import java.util.concurrent.atomic.AtomicLong;

import org.springframework.batch.infrastructure.item.Chunk;
import org.springframework.batch.infrastructure.item.ItemWriter;
import org.springframework.jdbc.core.JdbcTemplate;
import org.springframework.stereotype.Component;
import org.springframework.transaction.support.TransactionSynchronization;
import org.springframework.transaction.support.TransactionSynchronizationManager;

/**
 * Writes the results AND the cost ledger row for the chunk in one transaction, so a ledger row
 * exists if and only if its tickets do. The reader position, the results and the ledger either
 * all commit or all roll back.
 */
@Component
public class ClassificationWriter implements ItemWriter<Classification> {

    private final JdbcTemplate jdbc;
    private final UsageMeter meter;
    private final Pricing pricing;
    private final AtomicLong committedLedgerRows = new AtomicLong();

    public ClassificationWriter(JdbcTemplate jdbc, UsageMeter meter, Pricing pricing) {
        this.jdbc = jdbc;
        this.meter = meter;
        this.pricing = pricing;
    }

    @Override
    public void write(Chunk<? extends Classification> chunk) {
        List<? extends Classification> items = chunk.getItems();
        jdbc.batchUpdate("INSERT INTO TICKET_RESULT (TICKET_ID, CATEGORY, PRIORITY, CONFIDENCE) VALUES (?,?,?,?)",
                items, items.size(), (ps, c) -> {
                    ps.setLong(1, c.ticketId());
                    ps.setString(2, c.category());
                    ps.setInt(3, c.priority());
                    ps.setDouble(4, c.confidence());
                });

        // A ticket skipped in an earlier, uncommitted attempt (the skip listener commits on its own) and
        // classified now must not stay in the skipped table.
        jdbc.batchUpdate("DELETE FROM SKIPPED_TICKET WHERE TICKET_ID = ?", items, items.size(),
                (ps, c) -> ps.setLong(1, c.ticketId()));

        long first = items.stream().mapToLong(Classification::ticketId).min().orElse(0);
        long last = items.stream().mapToLong(Classification::ticketId).max().orElse(0);
        UsageMeter.Snapshot spent = meter.drain();
        if (TransactionSynchronizationManager.isSynchronizationActive()) {
            TransactionSynchronizationManager.registerSynchronization(new TransactionSynchronization() {
                @Override
                public void afterCompletion(int status) {
                    if (status == STATUS_COMMITTED) {
                        long rows = committedLedgerRows.incrementAndGet();
                        if (rows % 50 == 0) {
                            System.out.println("PROGRESS committed-chunks-this-jvm=" + rows + " last-ticket=" + last);
                        }
                    } else {
                        meter.restore(spent); // rolled back: the money was still spent
                    }
                }
            });
        }
        BigDecimal cost = pricing.cost(spent.promptTokens(), spent.completionTokens());
        jdbc.update("""
                INSERT INTO COST_LEDGER (FIRST_TICKET, LAST_TICKET, TICKETS_WRITTEN, TICKETS_SKIPPED, ATTEMPTS,
                  REJECTED_429, UNAVAILABLE_503, RETRIES, BILLED_CALLS, BAD_REPLIES, PROMPT_TOKENS, COMPLETION_TOKENS,
                  COST, PRICE_VERSION) VALUES (?,?,?,?,?,?,?,?,?,?,?,?,?,?)""",
                first, last, items.size(), spent.skipped(), spent.attempts(), spent.rejected429(),
                spent.unavailable503(), spent.retries(), spent.billedCalls(), spent.badReplies(),
                spent.promptTokens(), spent.completionTokens(), cost, pricing.version());
    }

    public long committedLedgerRows() {
        return committedLedgerRows.get();
    }
}