fix(multimodal): preserve distinct image prompts

agentlightning/_prompt_identity.py

@@ -0,0 +1,44 @@
+# Copyright (c) Microsoft. All rights reserved.
+
+"""Prompt identity shared by event deduplication and image alignment."""
+
+from __future__ import annotations
+
+import json
+from typing import Any
+
+PromptKey = tuple[tuple[int, ...], tuple[str, ...]]
+
+
+def extract_image_urls_from_messages(messages: Any) -> list[str]:
+    if not isinstance(messages, list):
+        return []
+    image_urls: list[str] = []
+    for message in messages:
+        if not isinstance(message, dict):
+            continue
+        content = message.get("content")
+        if isinstance(content, str):
+            try:
+                content = json.loads(content)  # Some clients serialize content parts as JSON text.
+            except json.JSONDecodeError:
+                continue
+        if isinstance(content, list):
+            for content_part in content:
+                if not isinstance(content_part, dict) or content_part.get("type") != "image_url":
+                    continue
+                image_url = content_part.get("image_url")
+                if isinstance(image_url, dict) and isinstance(image_url.get("url"), str):
+                    image_urls.append(image_url["url"])
+    return image_urls
+
+
+def prompt_key(prompt_token_ids: Any, image_urls: list[str]) -> PromptKey | None:
+    """Include ordered image references: placeholder token IDs omit image content."""
+    if (
+        not isinstance(prompt_token_ids, list)
+        or not prompt_token_ids
+        or any(type(token_id) is not int for token_id in prompt_token_ids)
+    ):
+        return None
+    return tuple(prompt_token_ids), tuple(image_urls)

agentlightning/server/routes/events.py

@@ -11,6 +11,7 @@
 from fastapi import APIRouter, Query
 from fastapi.exceptions import HTTPException
 
+from agentlightning._prompt_identity import PromptKey, extract_image_urls_from_messages, prompt_key
 from agentlightning.schemas import DEFAULT_ATTEMPT_ID, Event, EventCreate
 from agentlightning.server.store import _events, _rollouts
 
@@ -28,10 +29,18 @@ def _drop_superseded_routed_experts(events: list[Event], event_type: str, data:
 
     current = _trim_model_request(data)
     prompt_ids = current["prompt_token_ids"]
+    request = data.get("request")
+    messages = request.get("messages") if isinstance(request, dict) else None
+    image_urls = extract_image_urls_from_messages(messages)
     for index in range(len(events) - 1, -1, -1):
         previous = events[index]
         if previous.event_type != "model_request":
             continue
+        previous_request = previous.data.get("request")
+        previous_messages = previous_request.get("messages") if isinstance(previous_request, dict) else None
+        if image_urls != extract_image_urls_from_messages(previous_messages):
+            # Equal placeholder IDs do not imply equal visual inputs or routes.
+            return
         previous_data = _trim_model_request(previous.data)
         previous_context = previous_data["prompt_token_ids"] + previous_data["response_token_ids"]
         same_prompt = prompt_ids == previous_data["prompt_token_ids"]
@@ -209,23 +218,22 @@ def _without_routed_experts(event: Event) -> Event:
     return event.model_copy(update={"data": data})
 
 
-def _dedupe_model_requests_by_prompt_token_ids(events: list[Event]) -> list[Event]:
-    """Keep the last request for each valid, non-empty prompt-token key."""
-    last_index_by_prompt: dict[tuple[int, ...], int] = {}
+def _dedupe_model_requests_by_prompt(events: list[Event]) -> list[Event]:
+    """Keep the last raw request for each valid prompt-token and image key."""
+    last_index_by_prompt: dict[PromptKey, int] = {}
     kept_indexes: set[int] = set()
     for index, event in enumerate(events):
         if event.event_type != "model_request":
             continue
-        prompt_token_ids = event.data.get("prompt_token_ids", [])
-        if (
-            not isinstance(prompt_token_ids, list)
-            or not prompt_token_ids
-            or any(type(token_id) is not int for token_id in prompt_token_ids)
-        ):
+        prompt_token_ids = _trim_model_request(event.data)["prompt_token_ids"]
+        request = event.data.get("request")
+        messages = request.get("messages") if isinstance(request, dict) else None
+        key = prompt_key(prompt_token_ids, extract_image_urls_from_messages(messages))
+        if key is None:
             # Skip deduplication, not the request, when no valid key exists.
             kept_indexes.add(index)
             continue
-        last_index_by_prompt[tuple(prompt_token_ids)] = index
+        last_index_by_prompt[key] = index
 
     kept_indexes.update(last_index_by_prompt.values())
     return [event for index, event in enumerate(events) if event.event_type != "model_request" or index in kept_indexes]
@@ -250,8 +258,8 @@ async def query_events(
         event_type=event_type,
     )
     if format == "triplet":
+        events = _dedupe_model_requests_by_prompt(events)
         events = [_to_triplet_format(e) for e in events]
-        events = _dedupe_model_requests_by_prompt_token_ids(events)
     if not include_routed_experts:
         events = [_without_routed_experts(event) for event in events]
     return events

agentlightning/verl/agl_rollout_manager.py

@@ -4,7 +4,6 @@
 
 from __future__ import annotations
 
-import json  # [multimodal-patch]
 import time
 import traceback
 import uuid
@@ -18,6 +17,8 @@
 from httpx_retries import Retry, RetryTransport
 from pydantic import BaseModel, Field
 
+from agentlightning._prompt_identity import PromptKey, prompt_key
+from agentlightning._prompt_identity import extract_image_urls_from_messages as _extract_image_urls_from_messages
 from agentlightning.client import AgentLightningSyncClient
 from agentlightning.schemas import (
     TERMINAL_STATES,
@@ -141,33 +142,8 @@ def _without_routed_experts(event: Event) -> Event:
     return event.model_copy(update={"data": data})
 
 
-# [multimodal-patch] Extract image URLs from OpenAI-style chat messages, in order of
-# appearance (ported from agent-lightning v0.3.0 TripletAdapter.extract_prompt_image_urls).
-def _extract_image_urls_from_messages(messages: Any) -> list[str]:
-    if not isinstance(messages, list):
-        return []
-    image_urls: list[str] = []
-    for message in messages:
-        if not isinstance(message, dict):
-            continue
-        content = message.get("content")
-        if isinstance(content, str):
-            try:
-                content = json.loads(content)  # Some clients serialize content parts as JSON text.
-            except json.JSONDecodeError:
-                continue
-        if isinstance(content, list):
-            for content_part in content:
-                if not isinstance(content_part, dict) or content_part.get("type") != "image_url":
-                    continue
-                image_url = content_part.get("image_url")
-                if isinstance(image_url, dict) and isinstance(image_url.get("url"), str):
-                    image_urls.append(image_url["url"])
-    return image_urls
-
-
 # [multimodal-patch] Recover per-triplet image URLs from raw model_request events.
-# Replicates the server triplet view (dedupe by prompt_token_ids, keep last) and the
+# Replicates the server triplet view (dedupe by prompt token ids and images, keep last) and the
 # manager-side filtering below, so the result aligns one-to-one with the kept triplets.
 # Returns None when the rollout is text-only or when alignment cannot be guaranteed
 # (caller then leaves image_urls unset).
@@ -221,13 +197,14 @@ def _aligned_image_urls(raw_events: list[Event], n_triplets: int) -> list[list[s
 
     # Mirror the server-side dedupe: keep the last request for each valid prompt key.
     # Missing or malformed ids cannot establish that two requests are duplicates.
-    last_index_by_prompt: dict[tuple[int, ...], int] = {}
+    last_index_by_prompt: dict[PromptKey, int] = {}
     kept_indexes: set[int] = set()
-    for index, (_, prompt_token_ids, _, _) in enumerate(requests):
-        if not prompt_token_ids or any(type(token_id) is not int for token_id in prompt_token_ids):
+    for index, (_, prompt_token_ids, _, image_urls) in enumerate(requests):
+        key = prompt_key(prompt_token_ids, image_urls)
+        if key is None:
             kept_indexes.add(index)
             continue
-        last_index_by_prompt[tuple(prompt_token_ids)] = index
+        last_index_by_prompt[key] = index
     kept_indexes.update(last_index_by_prompt.values())
 
     aligned: list[list[str] | None] = []

docs/80-example-multimodal-qa.md

@@ -8,6 +8,8 @@ Multimodal QA is a minimal end-to-end example of **multimodal (image) reinforcem
 
 The example exercises the multimodal training data path: image URLs are recovered from the raw `model_request` events, aligned with the training triplets, and turned into per-row `multi_modal_inputs` (`pixel_values` + `image_grid_thw`) with mrope `(batch, 4, seq_len)` position ids, so the training forward pass actually consumes the image features.
 
+Within a rollout, requests are deduplicated only when both their prompt token IDs and ordered image URLs match. Different images can produce identical image placeholder tokens, so their transitions are kept separately; repeated requests with the same image URLs retain the last response.
+
 ## Data Preparation
 
 The dataset is generated with PIL at startup — there is nothing to download:

tests/server/test_endpoints.py

@@ -314,6 +314,47 @@ def post_model_request(prompt_token_ids: list[int], response_token_ids: list[int
     assert [event["data"]["response_token_ids"] for event in triplet_events[1:]] == [[20], [30]]
 
 
+@pytest.mark.parametrize("serialized_content", [False, True])
+def test_triplet_events_distinguish_image_prompts(
+    client: TestClient, auth_headers: dict[str, str], serialized_content: bool
+):
+    import json
+
+    rollout_id = _rollout(client, auth_headers)["rollout_id"]
+    image_a = "data:image/png;base64,YQ=="
+    image_b = "data:image/png;base64,Yg=="
+    images_by_request = [[image_a], [image_b], [image_a], [image_a, image_b], [image_b, image_a]]
+    for index, images in enumerate(images_by_request):
+        content = [{"type": "image_url", "image_url": {"url": url}} for url in images]
+        response = client.post(
+            f"/api/rollouts/{rollout_id}/attempt/0/events",
+            json={
+                "event_type": "model_request",
+                "data": {
+                    "request": {
+                        "messages": [
+                            {"role": "user", "content": json.dumps(content) if serialized_content else content}
+                        ]
+                    },
+                    "response": {"prompt_token_ids": [1, 2], "choices": [{"token_ids": [index]}]},
+                },
+            },
+            headers=auth_headers,
+        )
+        assert response.status_code == 200
+
+    response = client.get(f"/api/rollouts/{rollout_id}/events?format=triplet", headers=auth_headers)
+    assert response.status_code == 200
+    events = response.json()
+    assert [event["data"]["response_token_ids"] for event in events] == [[1], [2], [3], [4]]
+    # Image identity is needed only while deduplicating the raw events.
+    assert all(
+        set(event["data"]) == {"prompt_token_ids", "response_token_ids", "response_log_probs", "server"}
+        for event in events
+    )
+    assert "data:image" not in response.text
+
+
 @pytest.mark.parametrize("prompt_token_ids", [None, [], [[1]], ["1"], [True], [1.0], "bad-ids", {"token": 1}])
 def test_triplet_events_do_not_dedupe_without_valid_prompt_tokens(
     client: TestClient,

tests/verl/test_agl_rollout_manager.py

@@ -198,6 +198,7 @@ def test_build_completed_rollout_aligns_image_urls_with_triplets() -> None:
         _event("reward", {"value": 1.0}),
     ]
     triplet_events = [
+        _triplet_model_request_event([1, 2], [3, 4]),
         _triplet_model_request_event([1, 2], [30, 40]),
         _triplet_model_request_event([7, 8], [9]),
         _triplet_model_request_event([10], [11]),
@@ -217,8 +218,8 @@ def test_build_completed_rollout_aligns_image_urls_with_triplets() -> None:
     )
 
     assert completed.triplets is not None
-    assert len(completed.triplets) == 3
-    assert [triplet.image_urls for triplet in completed.triplets] == [[_IMG2], [_IMG, _IMG2], None]
+    assert len(completed.triplets) == 4
+    assert [triplet.image_urls for triplet in completed.triplets] == [[_IMG], [_IMG2], [_IMG, _IMG2], None]
 
 
 @pytest.mark.parametrize("prompt_token_ids", [None, [], [[1]], ["1"], [True], [1.0], "bad-ids", {"token": 1}])
@@ -267,3 +268,97 @@ def test_aligned_image_urls_text_only_rollout_returns_none_without_warning(
     # Text-only rollouts keep the exact original behavior: no alignment, no warning.
     assert _aligned_image_urls(raw_events, 2) is None
     assert capsys.readouterr().out == ""
+
+
+@pytest.mark.parametrize("routed_experts", [False, True])
+def test_distinct_images_survive_server_to_training_batch(routed_experts: bool) -> None:
+    import base64
+    import io
+    from typing import cast
+
+    import numpy as np
+    import torch
+    from fastapi.testclient import TestClient
+    from PIL import Image
+
+    from agentlightning.client import AgentLightningSyncClient
+    from agentlightning.server.app import create_app
+    from agentlightning.verl.rollout_adapter import RolloutAdapter
+
+    class ImageProcessor:
+        def __call__(self, text: object, images: list, return_tensors: str) -> dict:
+            return {"pixel_values": torch.tensor([image.getpixel((0, 0)) for image in images])}
+
+    urls = []
+    for color in ("red", "blue"):
+        buffer = io.BytesIO()
+        Image.new("RGB", (2, 2), color).save(buffer, format="PNG")
+        urls.append("data:image/png;base64," + base64.b64encode(buffer.getvalue()).decode())
+
+    app = create_app(
+        {
+            "key": "test",
+            "default_proxy": {"model_name": "test-model", "train": {"temperature": 1}, "val": {"temperature": 0}},
+        }
+    )
+    with TestClient(app, headers={"Authorization": "Bearer test"}) as client:
+        rollout_id = client.post("/api/rollouts", json=[{"input": {}}]).json()[0]["rollout_id"]
+        try:
+            client.patch(
+                f"/api/rollouts/{rollout_id}", json={"status": {"state": "running", "last_attempt_id": "0"}}
+            ).raise_for_status()
+            for index, url in enumerate([urls[0], urls[0], urls[1]]):
+                event = _raw_model_request_event([url], [1, 2], [index + 3])
+                event.data["response"]["choices"][0]["logprobs"] = {"content": [{"logprob": -0.1}]}
+                if routed_experts:
+                    buffer = io.BytesIO()
+                    np.save(buffer, np.array([index * 10, index * 10 + 1], dtype=np.uint8).reshape(2, 1, 1))
+                    event.data["routed_experts"] = base64.b64encode(buffer.getvalue()).decode()
+                response = client.post(
+                    f"/api/rollouts/{rollout_id}/attempt/0/events",
+                    json={"event_type": "model_request", "data": event.data},
+                )
+                assert response.status_code == 200
+            client.post(
+                f"/api/rollouts/{rollout_id}/attempt/0/events",
+                json={"event_type": "reward", "data": {"value": 0.75}},
+            ).raise_for_status()
+            client.patch(
+                f"/api/rollouts/{rollout_id}", json={"status": {"state": "succeeded", "last_attempt_id": "0"}}
+            ).raise_for_status()
+            manager = object.__new__(AglRolloutManagerBase)
+            manager.client = cast(AgentLightningSyncClient, client)
+            completed = manager._build_completed_rollout(
+                EnqueuedRollout(data_id="sample", rollout_id=rollout_id, step=0, sample_idx_in_step=0, enqueue_time=0),
+                manager._get_rollout(rollout_id),
+            )
+            assert completed.triplets is not None
+            assert [triplet.image_urls for triplet in completed.triplets] == [[urls[0]], [urls[1]]]
+
+            adapter = RolloutAdapter(
+                max_prompt_length=4,
+                max_response_length=2,
+                device=torch.device("cpu"),
+                pad_token_id=0,
+                processor=ImageProcessor(),
+                require_routed_experts=routed_experts,
+            )
+            batch, _ = adapter.get_train_data_batch([completed])
+            assert batch.batch is not None
+            assert batch.batch["responses"].tolist() == [[4, 0], [5, 0]]
+            assert batch.batch["is_drop_mask"].tolist() == [False, False]
+            assert batch.batch["token_level_scores"].sum(-1).tolist() == [0.75, 0.75]
+            torch.testing.assert_close(batch.batch["rollout_log_probs"], torch.tensor([[-0.1, 0.0], [-0.1, 0.0]]))
+            assert [row["pixel_values"].tolist() for row in batch.non_tensor_batch["multi_modal_inputs"]] == [
+                [[255, 0, 0]],
+                [[0, 0, 255]],
+            ]
+            if routed_experts:
+                assert batch.batch["routed_experts"][:, :, 0, 0].tolist() == [
+                    [0, 0, 10, 11, 0, 0],
+                    [0, 0, 20, 21, 0, 0],
+                ]
+                raw_events = client.get(f"/api/rollouts/{rollout_id}/events").json()
+                assert ["routed_experts" in event["data"] for event in raw_events[:3]] == [False, True, True]
+        finally:
+            client.delete(f"/api/rollouts/{rollout_id}").raise_for_status()