Repository: UKGovernmentBEIS/inspect_ai
Issue: Scoring crashes in transcript JSON change tracking on tasks involving model grading
Root cause: Validation Gap

Discussion on Slack [here](https://inspectcommunity.slack.com/archives/C0805HJTK8U/p1734010576953179)

@art-dsit, @Jay-Bailey-DSIT, and I ran into an issue where certain tasks will crash with variations of the error below. @Jay-Bailey-DSIT and my tasks both used model grading. Note the strange tilde character in the KeyError. 

What's strange is the scoring appears to almost succeed, generating scores which I can see, but showing as an error on the transcript. 

I can provide the tasks upon request to help debugging this. 

```
┌───────────────────────────────────────── Traceback (most recent call last) ──────────────────────────────────────────┐
│ /home/ubuntu/AS-Evals/.venv/lib/python3.12/site-packages/inspect_ai/_eval/task/run.py:622 in task_run_sample         │
│                                                                                                                      │
│ /usr/lib/python3.12/contextlib.py:144 in __exit__                                                                    │
│                                                                                                                      │
│   141 │   def __exit__(self, typ, value, traceback):                                                                 │
│   142 │   │   if typ is None:                                                                                        │
│   143 │   │   │   try:                                                                                               │
│ > 144 │   │   │   │   next(self.gen)                                                                                 │
│   145 │   │   │   except StopIteration:                                                                              │
│   146 │   │   │   │   return False                                                                                   │
│   147 │   │   │   else:                                                                                              │
│                                                                                                                      │
│ /home/ubuntu/AS-Evals/.venv/lib/python3.12/site-packages/inspect_ai/log/_transcript.py:348 in step                   │
│                                                                                                                      │
│ /usr/lib/python3.12/contextlib.py:144 in __exit__                                                                    │
│                                                                                                                      │
│   141 │   def __exit__(self, typ, value, traceback):                                                                 │
│   142 │   │   if typ is None:                                                                                        │
│   143 │   │   │   try:                                                                                               │
│ > 144 │   │   │   │   next(self.gen)                                                                                 │
│   145 │   │   │   except StopIteration:                                                                              │
│   146 │   │   │   │   return False                                                                                   │
│   147 │   │   │   else:                                                                                              │
│                                                                                                                      │
│ /home/ubuntu/AS-Evals/.venv/lib/python3.12/site-packages/inspect_ai/log/_transcript.py:388 in track_state_changes    │
│                                                                                                                      │
│ /home/ubuntu/AS-Evals/.venv/lib/python3.12/site-packages/inspect_ai/_util/json.py:107 in json_changes                │
└──────────────────────────────────────────────────────────────────────────────────────────────────────────────────────┘
KeyError: 'anthropic~1claude-3-5-sonnet-20241022'
```


Could you post a minimal example that fails? (will help us to not have to speculatively reconstruct the failure case)

Sounds good, I've started to try and look for a minimal repro.

Just to echo here what I posted to slack earlier, Jay and I found that [this line](https://github.com/UKGovernmentBEIS/inspect_ai/blob/main/src/inspect_ai/_util/json.py#L97) seems to be where the key modification is happening. It seems that the make_patch function uses the jsonpointer library which converts '/' to '~1' when e.g. we have a model name like 'anthropic/claude...' as a key in the before and after dicts. This then causes problems when applying the patch later on on line 107.

This works as a (temporary) fix (tested for remove_guardrails):
```
def json_changes(
    before: dict[str, Any], after: dict[str, Any]
) -> list[JsonChange] | None:
    patch = jsonpatch.make_patch(before, after)
    if patch:
        changes: list[JsonChange] = []
        for change in cast(list[Any], patch):
            json_change = JsonChange(**change)
            if json_change.op == "replace":
                paths = json_change.path.split("/")[1:]
                replaced = before
                for path in paths:
                    index: Any = int(path) if path.isnumeric() else path
                    try:
                        replaced = replaced[index]
                    except KeyError:
                        replaced = replaced[index.replace("~1", "/")]
                json_change.replaced = replaced
            changes.append(json_change)
        return changes
    else:
        return None
```


something else we have noticed is that the problem seems to come from the json version of the TaskState with this:

```
'completed': True, 'metadata': {'model_usage': {'anthropic/claude-3-5-sonnet-20241022': {'input_tokens': 4, 'output_tokens': 151, 'total_tokens': 2094, 'input_tokens_cache_write': 1939, 'input_tokens_cache_read': 0
            }
        }
    }
```

and this seems to be related to use of a grader model in the scorer. The `before` dict contains the above, then `after` contains the same but with updated token counts. This only seems to happen on k8s

Here is what I think is a minimal repro example:

```
from inspect_ai import Task, task
from inspect_ai.dataset import Sample
from inspect_ai.model import get_model
from inspect_ai.model._model import sample_model_usage
from inspect_ai.scorer import Score, Scorer, Target, mean, scorer, stderr
from inspect_ai.solver import Generate, Solver, TaskState, solver


@solver
def example_solver() -> Solver:
    # `agents.autonomous_systems_agent.autonomous_systems_agent` has this same problem in the presence of a grader model set up like this because it also uses sample_model_usage in the same way
    async def _solver(state: TaskState, generate: Generate) -> TaskState:
        model = get_model()
        await model.generate(input=state.input)
        state.metadata["model_usage"] = sample_model_usage()

        """
        The above line results in a dict with key e.g. "openai/gpt-4o-2024-08-06" and value ModelUsage object being assigned to the "model_usage" key in the state.metadata dict.

        I think the later call of the grader model in the scorer then causes an update to this ModelUsage object, which then means that when the inspect_ai._util.json.json_changes function is called during the state update it runs into trouble due to the '/' in the key.
        """

        return state
    return _solver


@scorer(metrics={"accuracy": [mean(), stderr()]})
def example_scorer() -> Scorer:
    # Using a different model to the solver here prevents the bug from appearing
    grader_model = get_model()

    async def _score(state: TaskState, target: Target) -> Score:
        await grader_model.generate(input="Say no")

        return Score(value={"accuracy": 1.})
    return _score


@task
def example_task() -> Task:
    return Task(
        dataset=[Sample(input="Say yes")],
        solver=example_solver(),
        scorer=example_scorer(),
    )
```

Important points are: 

- we call the same model for both the solver and the scorer (as a grader model)
- tracking of model usage in the task state per this line `state.metadata["model_usage"] = sample_model_usage()`. 

My understanding is that tracking the model usage like this creates a dict within the state with the model name as a key. The model name contains '/'. When the same model is called the second time in the scorer, this triggers a call of `inspect_ai._util.json.json_changes` due to the change to the model usage in the state. But json_changes doesn't properly handle update where the key of the updated value contains '/'. This is because the jsonpointer library is used, which converts '/' to '~1' and breaks the reference.

Thank you for sharing the reproducible case - I can reproduce it locally and am taking a look now!

Resolved here: https://github.com/UKGovernmentBEIS/inspect_ai/commit/a4521ba28a55226f7af45d60194c8daf028e6015