diff --git a/agentlightning/adapter/messages.py b/agentlightning/adapter/messages.py index fe6c5912e..4bc731262 100644 --- a/agentlightning/adapter/messages.py +++ b/agentlightning/adapter/messages.py @@ -3,6 +3,7 @@ from __future__ import annotations import json +import logging from collections import defaultdict from typing import TYPE_CHECKING, Any, Dict, Generator, Iterable, List, Optional, Sequence, TypedDict, Union, cast @@ -12,6 +13,8 @@ from .base import TraceAdapter +logger = logging.getLogger(__name__) + if TYPE_CHECKING: from openai.types.chat import ( ChatCompletionFunctionToolParam, @@ -103,6 +106,30 @@ def group_genai_dict(data: Dict[str, Any], prefix: str) -> Union[Dict[str, Any], return result +def _infer_missing_role(msg: Dict[str, Any]) -> Optional[str]: + """Infer an OpenAI chat role for a prompt entry whose ``role`` field is missing. + + Some tracers (notably AgentOps when re-emitting prior turns inside ``gen_ai.prompt.N.*``) + serialize the nested ``tool_calls`` / ``tool_call_id`` subtree without the sibling ``role`` + key. The role is still unambiguous in those cases: + + - A message carrying ``tool_calls`` can only have come from the assistant. + - A message carrying ``tool_call_id`` can only have come from a tool response. + + Args: + msg: A prompt entry parsed from ``gen_ai.prompt.*`` attributes. + + Returns: + ``"assistant"`` or ``"tool"`` when the role can be inferred unambiguously, + otherwise ``None``. + """ + if "tool_calls" in msg: + return "assistant" + if "tool_call_id" in msg: + return "tool" + return None + + def convert_to_openai_messages(prompt_completion_list: List[_RawSpanInfo]) -> Generator[OpenAIMessages, None, None]: """Convert raw trace payloads into OpenAI-style chat messages. @@ -131,7 +158,19 @@ def convert_to_openai_messages(prompt_completion_list: List[_RawSpanInfo]) -> Ge # Extract messages for msg in pc_entry["prompt"]: - role = msg["role"] + role = msg.get("role") + if role is None: + # Some tracers omit ``role`` on re-emitted assistant/tool turns; recover when the + # role is unambiguous, otherwise drop just this message instead of crashing the + # whole rollout's adapter pass. + role = _infer_missing_role(msg) + if role is None: + logger.warning( + "Skipping prompt message with no 'role' and no inferable role hint: %r", + msg, + ) + continue + logger.debug("Inferred missing role %r for prompt message: %r", role, msg) if role == "assistant" and "tool_calls" in msg: # Use the tool_calls directly diff --git a/tests/adapter/test_messages_adapter.py b/tests/adapter/test_messages_adapter.py index c4280884b..028afd5c0 100644 --- a/tests/adapter/test_messages_adapter.py +++ b/tests/adapter/test_messages_adapter.py @@ -1,6 +1,7 @@ # Copyright (c) Microsoft. All rights reserved. import json +import logging from importlib.metadata import version from typing import Any, Dict, Optional @@ -391,3 +392,114 @@ def test_trace_messages_adapter_handles_multiple_tool_calls(): ] assert adapter.adapt(spans) == expected + + +# Regression coverage for https://github.com/microsoft/agent-lightning/issues/425 and +# https://github.com/microsoft/agent-lightning/issues/311. Some tracers (notably AgentOps when +# re-emitting prior turns inside ``gen_ai.prompt.N.*``) drop the ``role`` key on assistant +# tool-call entries and tool-response entries. The adapter must recover instead of crashing +# the entire rollout adapter pass. +@pytest.mark.skipif( + _skip_for_openai_lt_1_100_0, + reason="Requires openai>=1.100.0", +) +def test_trace_messages_adapter_recovers_assistant_role_from_tool_calls() -> None: + tool_call_id = "call_AnIgZ6EdncSTDeJMn3rcKSDM" + tool_name = "get_rmc_percentage_of_sales" + tool_arguments = json.dumps({"category": "Scooter"}) + tool_payload = json.dumps({"value": 0.42}) + + spans = [ + make_span( + "openai.chat.completion", + { + "gen_ai.prompt.0.role": "system", + "gen_ai.prompt.0.content": "You are a financial data analyst.", + "gen_ai.prompt.1.role": "user", + "gen_ai.prompt.1.content": "How does Scooter compare?", + # Re-emitted assistant tool-call entry: ``tool_calls.*`` present, ``role`` missing. + "gen_ai.prompt.2.tool_calls.0.id": tool_call_id, + "gen_ai.prompt.2.tool_calls.0.name": tool_name, + "gen_ai.prompt.2.tool_calls.0.arguments": tool_arguments, + # Re-emitted tool response: ``tool_call_id`` + ``content`` present, ``role`` missing. + "gen_ai.prompt.3.tool_call_id": tool_call_id, + "gen_ai.prompt.3.content": tool_payload, + "gen_ai.completion.0.role": "assistant", + "gen_ai.completion.0.content": "Scooter sales are up 4%.", + "gen_ai.completion.0.finish_reason": "stop", + }, + 0, + ), + ] + + result = TraceToMessages().adapt(spans) + + assert result == [ + { + "messages": [ + {"content": "You are a financial data analyst.", "role": "system"}, + {"content": "How does Scooter compare?", "role": "user"}, + { + "content": None, + "role": "assistant", + "tool_calls": [ + { + "id": tool_call_id, + "type": "function", + "function": {"name": tool_name, "arguments": tool_arguments}, + } + ], + }, + { + "content": tool_payload, + "role": "tool", + "tool_call_id": tool_call_id, + }, + {"content": "Scooter sales are up 4%.", "role": "assistant"}, + ], + "tools": None, + } + ] + + +@pytest.mark.skipif( + _skip_for_openai_lt_1_100_0, + reason="Requires openai>=1.100.0", +) +def test_trace_messages_adapter_skips_unidentifiable_prompt_entry( + caplog: pytest.LogCaptureFixture, +) -> None: + spans = [ + make_span( + "openai.chat.completion", + { + "gen_ai.prompt.0.role": "system", + "gen_ai.prompt.0.content": "You are helpful.", + "gen_ai.prompt.1.role": "user", + "gen_ai.prompt.1.content": "Hi.", + # Garbage entry with neither ``role`` nor any role hint; must not crash. + "gen_ai.prompt.2.unexpected": "noise", + "gen_ai.completion.0.role": "assistant", + "gen_ai.completion.0.content": "Hello.", + "gen_ai.completion.0.finish_reason": "stop", + }, + 0, + ), + ] + + with caplog.at_level(logging.WARNING, logger="agentlightning.adapter.messages"): + result = TraceToMessages().adapt(spans) + + assert result == [ + { + "messages": [ + {"content": "You are helpful.", "role": "system"}, + {"content": "Hi.", "role": "user"}, + {"content": "Hello.", "role": "assistant"}, + ], + "tools": None, + } + ] + assert any( + "no inferable role hint" in record.getMessage() for record in caplog.records + ), "Expected a warning naming the malformed prompt entry"