Fix BatchInferenceAPI caching wrong responses for chunked batches

[v0ropaev]

Problem

When BatchInferenceAPI.__call__ splits the uncached prompts into more than one batch, the responses returned by that call are correct, but the cache gets corrupted: the first len(last_chunk) uncached prompts are saved with the last chunk's completions. Any later run that hits the cache silently returns another prompt's answer.

It happens whenever the uncached prompts end up in more than one batch, e.g. chunk smaller than the number of uncached prompts, more than 100k prompts / 250MB for Anthropic, or more than 50k prompts for OpenAI.

Cause

The multi-chunk branch already saves each chunk inside its loop. The "Save responses to cache" block after the if/else then ran again for both paths with zip(uncached_prompts, responses), and in the multi-chunk path responses is the loop variable left over from the last chunk. This came in with #38 (Redis caching), which merged the tails of the two branches; before that, the single-batch branch had its own save and the multi-chunk branch returned early.

Repro (no API calls)

import asyncio
import tempfile
from pathlib import Path

from safetytooling.apis.batch_api import BatchInferenceAPI
from safetytooling.data_models import ChatMessage, LLMResponse, MessageRole, Prompt


async def fake_batch(model_id, prompts, max_tokens, **kwargs):
    # Echo each prompt back as its completion
    return [
        LLMResponse(model_id=model_id, completion=p.messages[0].content, stop_reason="stop_sequence", cost=0)
        for p in prompts
    ], "fake_batch"


async def main():
    tmp = Path(tempfile.mkdtemp())
    kwargs = dict(log_dir=tmp, prompt_history_dir=None, cache_dir=tmp / "cache", anthropic_api_key="x", openai_api_key="x")
    prompts = [Prompt(messages=[ChatMessage(role=MessageRole.user, content=f"Say {i}")]) for i in range(4)]

    api = BatchInferenceAPI(**kwargs)
    api._anthropic_batch = fake_batch
    responses, _ = await api(model_id="claude-3-5-haiku-20241022", prompts=prompts, max_tokens=10, chunk=2)
    print("first run: ", [r.completion for r in responses])

    responses, batch_id = await BatchInferenceAPI(**kwargs)(
        model_id="claude-3-5-haiku-20241022", prompts=prompts, max_tokens=10, chunk=2
    )
    print("cached run:", batch_id, [r.completion for r in responses])


asyncio.run(main())

On main (library prints omitted):

first run:  ['Say 0', 'Say 1', 'Say 2', 'Say 3']
cached run: cached ['Say 2', 'Say 3', 'Say 2', 'Say 3']

With this PR the cached run returns ['Say 0', 'Say 1', 'Say 2', 'Say 3'].

Fix

Move the post-loop save back into the single-batch else branch, so each path saves its responses exactly once. The multi-chunk path keeps its per-chunk save, which pairs each chunk's prompts with that chunk's responses. As a side effect the multi-chunk path no longer rewrites part of the cache a second time.

Tests

New offline tests/test_batch_api_cache.py (kept out of the slow, real-API tests/test_batch_api.py): a fake _anthropic_batch, then a fresh BatchInferenceAPI has to read every prompt's own completion back from the cache, for single_batch and two_chunks. The two_chunks case fails on main as above; the single_batch case guards the moved block.

• pytest tests/test_batch_api_cache.py: 2 passed
• pytest -n 6 without API keys: the only failures are the same 8 key-requiring tests as on main
• ruff check . (0.9.0) and black --check . (24.10.0): clean

Caches written by earlier multi-chunk runs may already contain wrong entries; those need to be cleared (or a new seed used), this PR doesn't repair them.

Separate from this, and not changed here: for non-Anthropic models __call__ never copies max_tokens into kwargs, so the OpenAI batch body gets max_tokens: null and the cache key's max_tokens is always None. Fixing that changes cache keys, so I'd rather do it in its own PR if you want it.