diff --git a/plugins/observability/nemo_relay/__init__.py b/plugins/observability/nemo_relay/__init__.py index 774d2712c5f6..afdec5d8ac94 100644 --- a/plugins/observability/nemo_relay/__init__.py +++ b/plugins/observability/nemo_relay/__init__.py @@ -20,6 +20,11 @@ logger = logging.getLogger(__name__) _INIT_FAILED = object() _LOCK = threading.RLock() _RUNTIME: "_Runtime | object | None" = None +_RELAY_LLM_SURFACE_BY_API_MODE = { + "anthropic_messages": "anthropic.messages", + "chat_completions": "openai.chat_completions", + "codex_responses": "openai.responses", +} @dataclass @@ -358,6 +363,8 @@ class _Runtime: normalize_payload: Callable[[Any], Any], shape_response: Callable[[Any], Any], make_managed_execute: Callable[[Callable[[Any], Any]], Any], + *, + preserve_raw_response: bool, ) -> Any: # NeMo Relay's native managed execution may wrap a failing callback as an # internal runtime error, hiding the real downstream provider/tool @@ -387,7 +394,9 @@ class _Runtime: if downstream_error is not None and _is_relay_wrapped_callback_error(exc, callback_error): raise downstream_error raise - return raw_response["value"] if raw_response["set"] else managed_result + if preserve_raw_response and raw_response["set"]: + return raw_response["value"] + return managed_result def execute_llm(self, kwargs: dict[str, Any]) -> Any: state = self.ensure_session(kwargs) @@ -404,7 +413,7 @@ class _Runtime: def _make_managed(impl: Callable[[Any], Any]) -> Any: async def _managed_execute() -> Any: result = self.nemo_relay.llm.execute( - str(kwargs.get("provider") or "llm"), + _relay_llm_surface(kwargs), request, impl, handle=state.handle, @@ -426,7 +435,7 @@ class _Runtime: return _managed_execute() return self._run_managed_with_downstream_preservation( - next_call, _normalize, _llm_response_payload, _make_managed + next_call, _normalize, _llm_response_payload, _make_managed, preserve_raw_response=True ) def execute_tool(self, kwargs: dict[str, Any]) -> Any: @@ -464,7 +473,7 @@ class _Runtime: return _managed_execute() return self._run_managed_with_downstream_preservation( - next_call, _normalize, _jsonable, _make_managed + next_call, _normalize, _jsonable, _make_managed, preserve_raw_response=False ) @@ -923,6 +932,14 @@ def _tool_key(kwargs: dict[str, Any]) -> str: ) +def _relay_llm_surface(kwargs: dict[str, Any]) -> str: + api_mode = str(kwargs.get("api_mode") or "").strip().lower() + return _RELAY_LLM_SURFACE_BY_API_MODE.get( + api_mode, + str(kwargs.get("provider") or "llm"), + ) + + def _metadata(kwargs: dict[str, Any]) -> dict[str, Any]: keys = ( "telemetry_schema_version", diff --git a/tests/plugins/test_nemo_relay_plugin.py b/tests/plugins/test_nemo_relay_plugin.py index 263e98c74b1c..87e0a34135e3 100644 --- a/tests/plugins/test_nemo_relay_plugin.py +++ b/tests/plugins/test_nemo_relay_plugin.py @@ -603,6 +603,92 @@ def test_nemo_relay_plugin_activates_and_owns_dynamic_plugins(tmp_path, monkeypa assert event_names.index("atif.deregister") < event_names.index("plugin.activation.close") +@pytest.mark.parametrize( + ("provider", "api_mode", "expected_surface", "should_rewrite"), + [ + ("custom", "chat_completions", "openai.chat_completions", True), + ("openai-codex", "codex_responses", "openai.responses", True), + ("anthropic", "anthropic_messages", "anthropic.messages", True), + ("custom", "anthropic_messages", "anthropic.messages", True), + ("bedrock", "bedrock_converse", "bedrock", False), + ], +) +def test_nemo_relay_managed_llm_uses_wire_protocol_for_interceptor_dispatch( + tmp_path, + monkeypatch, + provider, + api_mode, + expected_surface, + should_rewrite, +): + fake = _FakeNemoRelay() + supported_surfaces = { + "anthropic.messages", + "openai.chat_completions", + "openai.responses", + } + + def execute(name, request, func, **kwargs): + fake.events.append(("llm.execute.start", name, request.content, kwargs)) + content = dict(request.content) + if name in supported_surfaces: + content["rewritten_for"] = name + result = func(_FakeLLMRequest(request.headers, content)) + fake.events.append(("llm.execute.end", name, result, kwargs)) + return result + + fake.llm.execute = execute + plugin = _fresh_plugin(monkeypatch, fake) + _enable_dynamic_plugin(tmp_path, monkeypatch) + + result = plugin.on_llm_execution_middleware( + session_id="s1", + provider=provider, + api_mode=api_mode, + model="fixture", + request={"messages": [{"role": "user", "content": "hi"}]}, + next_call=lambda request: request, + ) + + execute_start = next( + event for event in fake.events if event[0] == "llm.execute.start" + ) + assert execute_start[1] == expected_surface + assert execute_start[3]["metadata"]["provider"] == provider + assert execute_start[3]["metadata"]["api_mode"] == api_mode + if should_rewrite: + assert result["rewritten_for"] == expected_surface + else: + assert "rewritten_for" not in result + + +def test_nemo_relay_managed_tool_returns_post_interceptor_result(tmp_path, monkeypatch): + fake = _FakeNemoRelay() + + def execute(name, args, func, **kwargs): + fake.events.append(("tool.execute.start", name, args, kwargs)) + raw = func({"intercepted": True, **args}) + result = {"compressed": True, "raw": raw} + fake.events.append(("tool.execute.end", name, result, kwargs)) + return result + + fake.tools.execute = execute + plugin = _fresh_plugin(monkeypatch, fake) + _enable_dynamic_plugin(tmp_path, monkeypatch) + + result = plugin.on_tool_execution_middleware( + session_id="s1", + tool_name="fixture-tool", + args={"value": 1}, + next_call=lambda args: {"tool_output": args}, + ) + + assert result == { + "compressed": True, + "raw": {"tool_output": {"intercepted": True, "value": 1}}, + } + + def test_nemo_relay_plugin_activates_before_registering_managed_middleware(tmp_path, monkeypatch): fake = _FakeNemoRelay() plugin = _fresh_plugin(monkeypatch, fake)