diff --git a/TELEMETRY-CONTRACT.md b/TELEMETRY-CONTRACT.md index ef183049..4defc9be 100644 --- a/TELEMETRY-CONTRACT.md +++ b/TELEMETRY-CONTRACT.md @@ -649,6 +649,16 @@ them so the graph span and the `invoke_agent` root share one name for the graph everything LaunchDarkly owns on a span sits under `launchdarkly.`. `ld.ai.` stays reserved for metric and event keys. +Every `launchdarkly.graph` span, from `graph()` and from each `native_graph` adapter, goes through +`set_ld_span_attributes`, so it carries the identity in section 2 and the `feature_flag` event. +Its config key is the graph key. The graph-level events (`$ld:ai:graph:invocation_success`, +`invocation_failure`, `duration:total` and `total_tokens`) use the same track data, so their +`configKey` is the graph key too, whichever runner produced them. Node-level events keep the node's +own config key. + +A native adapter has no graph variation metadata, so its graph span and graph events report an +empty `variationKey` and version 1. `graph()` reports the real values. + Nothing else in `graph.py` or `native_graph.py` needs work. --- diff --git a/packages/claude-agents/src/launchdarkly_ai_claude_agents/native_graph.py b/packages/claude-agents/src/launchdarkly_ai_claude_agents/native_graph.py index b0ffadb9..98337a97 100644 --- a/packages/claude-agents/src/launchdarkly_ai_claude_agents/native_graph.py +++ b/packages/claude-agents/src/launchdarkly_ai_claude_agents/native_graph.py @@ -18,8 +18,10 @@ NativeTool, get_client, make_track_data, + set_ld_span_attributes, to_ld_context, ) +from launchdarkly_ai_server.utils import make_graph_track_data try: from opentelemetry import trace @@ -239,12 +241,19 @@ async def invoke( tracer_name = "@launchdarkly/ai-claude-agents" if _HAS_OTEL: span = trace.get_tracer(tracer_name).start_span("launchdarkly.graph") - span.set_attribute("launchdarkly.graph.key", def_obj.key) else: span = None start_time = time.monotonic() run_id = str(uuid.uuid4()) + if span: + set_ld_span_attributes( + span, + { + "__ld": make_graph_track_data(def_obj.key, run_id), + "ldContext": raw_ld_context, + }, + ) path: list[str] = [] total_usage = {"input": 0, "output": 0, "total": 0} subagent_tool_ctx: dict[str, Any] = {} @@ -369,7 +378,7 @@ async def _subagent_execute( span.set_status(SpanStatusCode.ERROR, str(exc)) span.end() if ld_context: - td = make_track_data(root, def_obj.key, run_id) + td = make_graph_track_data(def_obj.key, run_id) get_client().track( "$ld:ai:graph:invocation_failure", ld_context, td, 1 ) @@ -411,18 +420,18 @@ async def _subagent_execute( span.end() if ld_context: - root_td = make_track_data(root, def_obj.key, run_id) + graph_td = make_graph_track_data(def_obj.key, run_id) client = get_client() client.track( - "$ld:ai:graph:duration:total", ld_context, root_td, graph_dur + "$ld:ai:graph:duration:total", ld_context, graph_td, graph_dur ) client.track( "$ld:ai:graph:total_tokens", ld_context, - root_td, + graph_td, total_usage["total"], ) - client.track("$ld:ai:graph:invocation_success", ld_context, root_td, 1) + client.track("$ld:ai:graph:invocation_success", ld_context, graph_td, 1) return {"response": final_output, "usage": total_usage} diff --git a/packages/claude-agents/tests/test_native_graph.py b/packages/claude-agents/tests/test_native_graph.py index e01900f4..0dfc7c85 100644 --- a/packages/claude-agents/tests/test_native_graph.py +++ b/packages/claude-agents/tests/test_native_graph.py @@ -587,6 +587,124 @@ async def test_emits_invocation_success_on_completion(self) -> None: assert root_nodes[0]["index"] == 0 assert all(evt != "$ld:ai:graph:path" for evt, _ in track_calls) + @pytest.mark.asyncio + async def test_graph_span_carries_config_identity( + self, monkeypatch: pytest.MonkeyPatch + ) -> None: + """The graph span must identify the config so Monitoring can link the trace.""" + import launchdarkly_ai_claude_agents.native_graph as ng_mod + + monkeypatch.setenv("LD_ENVIRONMENT_ID", "env-123") + mock_span = MagicMock() + mock_trace = MagicMock() + mock_trace.get_tracer.return_value.start_span.return_value = mock_span + + mock_sdk = _make_sdk_mock("done") + graph_def = _make_graph_def() + ctx = {"kind": "user", "key": "u1"} + + with patch( + "importlib.import_module", + side_effect=lambda n: ( + mock_sdk if n == "claude_agent_sdk" else __import__(n) + ), + ): + with patch.object(ng_mod, "get_client", return_value=MagicMock()): + with patch.object(ng_mod, "trace", mock_trace): + with patch.object(ng_mod, "_HAS_OTEL", True): + await to_claude_agents( + _make_def_promise(graph_def), + opts={"context": ctx}, + ).invoke("hi") + + attrs = {c[0][0]: c[0][1] for c in mock_span.set_attribute.call_args_list} + assert attrs["launchdarkly.operation.type"] == "gen_ai" + assert attrs["launchdarkly.config.key"] == "test-graph" + assert attrs["launchdarkly.graph.key"] == "test-graph" + assert attrs["launchdarkly.run.id"] + assert "launchdarkly.variation.key" in attrs + assert attrs["context.contextKeys.user"] == "u1" + + events = {c[0][0]: c[0][1] for c in mock_span.add_event.call_args_list} + assert events["feature_flag"] == { + "feature_flag.key": "test-graph", + "feature_flag.provider.name": "LaunchDarkly", + "feature_flag.set.id": "env-123", + "feature_flag.context.id": "u1", + "feature_flag.contextKeys": '{"user":"u1"}', + } + + @pytest.mark.asyncio + async def test_tracking_events_carry_the_environment_id( + self, monkeypatch: pytest.MonkeyPatch + ) -> None: + monkeypatch.setenv("LD_ENVIRONMENT_ID", "env-123") + payloads: list[Any] = [] + mock_ld_client = MagicMock() + mock_ld_client.track = MagicMock( + side_effect=lambda evt, ctx, data, val: payloads.append(data) + ) + + mock_sdk = _make_sdk_mock("done") + graph_def = _make_graph_def() + ctx = {"kind": "user", "key": "u1"} + + with patch( + "importlib.import_module", + side_effect=lambda n: ( + mock_sdk if n == "claude_agent_sdk" else __import__(n) + ), + ): + with patch.object(_claude_ng, "get_client", return_value=mock_ld_client): + await to_claude_agents( + _make_def_promise(graph_def), + opts={"context": ctx}, + ).invoke("hi") + + assert payloads + for data in payloads: + assert data["environmentId"] == "env-123" + + @pytest.mark.asyncio + async def test_graph_events_are_keyed_to_the_graph(self) -> None: + """Graph-level events carry the graph key, like the graph span and graph().""" + calls: list[tuple[str, Any]] = [] + mock_ld_client = MagicMock() + mock_ld_client.track = MagicMock( + side_effect=lambda evt, ctx, data, val: calls.append((evt, data)) + ) + + mock_sdk = _make_sdk_mock("done") + graph_def = _make_graph_def() + ctx = {"kind": "user", "key": "u1"} + + with patch( + "importlib.import_module", + side_effect=lambda n: ( + mock_sdk if n == "claude_agent_sdk" else __import__(n) + ), + ): + with patch.object(_claude_ng, "get_client", return_value=mock_ld_client): + await to_claude_agents( + _make_def_promise(graph_def), + opts={"context": ctx}, + ).invoke("hi") + + graph_events = [ + (e, d) + for e, d in calls + if e + in ( + "$ld:ai:graph:invocation_success", + "$ld:ai:graph:duration:total", + "$ld:ai:graph:total_tokens", + ) + ] + assert len(graph_events) == 3 + for _, data in graph_events: + assert data["configKey"] == "test-graph" + assert data["graphKey"] == "test-graph" + @pytest.mark.asyncio async def test_emits_invocation_failure_on_error(self) -> None: track_calls: list[str] = [] diff --git a/packages/client/src/launchdarkly_ai_server/graph.py b/packages/client/src/launchdarkly_ai_server/graph.py index 02e81f55..042d6776 100644 --- a/packages/client/src/launchdarkly_ai_server/graph.py +++ b/packages/client/src/launchdarkly_ai_server/graph.py @@ -27,7 +27,13 @@ UsageDict, VariationMeta, ) -from .utils import end_span_once, model_stamps_from_meta, select_handler, to_ld_context +from .utils import ( + end_span_once, + model_stamps_from_meta, + select_handler, + set_ld_span_attributes, + to_ld_context, +) logger = logging.getLogger(__name__) @@ -109,7 +115,7 @@ async def _build_graph( ) -> tuple[GraphDefinition, TrackData, Callable[..., Any]]: from .judges import run_judges from .lifecycle import extract_variation, get_client - from .tracking import execute_and_track + from .tracking import _try_get_environment_id, execute_and_track result = await _fetch_graph_variation(key, context) # Convert once so all inner track() calls use an ldclient.Context object. @@ -128,6 +134,9 @@ async def _build_graph( **model_stamps_from_meta(meta), "graphKey": key, } + _environment_id = _try_get_environment_id() + if _environment_id: + graph_track_data["environmentId"] = _environment_id if not enabled or not topology: return ( @@ -914,7 +923,9 @@ async def invoke( tracer = trace.get_tracer("@launchdarkly/ai-server") with tracer.start_as_current_span("launchdarkly.graph") as span: - span.set_attribute("launchdarkly.graph.key", self._key) + set_ld_span_attributes( + span, {"__ld": graph_track_data, "ldContext": context} + ) start_time = time.monotonic() total_usage = {"input": 0, "output": 0, "total": 0} @@ -1130,7 +1141,7 @@ async def _stream_events( tracer = trace.get_tracer("@launchdarkly/ai-server") span = tracer.start_span("launchdarkly.graph", context=caller_context) - span.set_attribute("launchdarkly.graph.key", self._key) + set_ld_span_attributes(span, {"__ld": graph_track_data, "ldContext": context}) span_context = set_span_in_context(span, caller_context) ended: set[int] = set() start_time = time.monotonic() diff --git a/packages/client/src/launchdarkly_ai_server/utils.py b/packages/client/src/launchdarkly_ai_server/utils.py index 2958de6e..8a0c3b0a 100644 --- a/packages/client/src/launchdarkly_ai_server/utils.py +++ b/packages/client/src/launchdarkly_ai_server/utils.py @@ -597,9 +597,11 @@ def make_track_data(node: GraphNode, graph_key: str, run_id: str) -> dict[str, A Builds the standard tracking payload for a graph node event. Shared by all native graph adapters (openai-agents, claude-agents, langchain-agents). """ + from .tracking import _try_get_environment_id # late import: tracking imports utils + meta = node.meta if isinstance(node.meta, dict) else {} config = node.config if isinstance(node.config, dict) else {} - return { + track_data: dict[str, Any] = { "runId": run_id, "configKey": node.key, "variationKey": meta.get("variationKey", ""), @@ -609,6 +611,40 @@ def make_track_data(node: GraphNode, graph_key: str, run_id: str) -> dict[str, A **model_stamps_from_meta(meta), "graphKey": graph_key, } + environment_id = _try_get_environment_id() + if environment_id: + track_data["environmentId"] = environment_id + return track_data + + +def make_graph_track_data(graph_key: str, run_id: str) -> dict[str, Any]: + """ + Builds the tracking payload for a native graph run as a whole. + + The graph flag is itself the AI Config the trace belongs to, so its key is + both the config key and the graph key -- the same choice ``graph.py`` makes + for the SDK's own graph runner. There is no model or variation to report: a + graph flag carries a topology, and ``GraphDefinition`` does not expose its + ``_ldMeta``. + + Native graph adapters pass this to ``set_ld_span_attributes`` so the + ``launchdarkly.graph`` span can be found by an AI Config Monitoring query. + """ + from .tracking import _try_get_environment_id # late import: tracking imports utils + + track_data: dict[str, Any] = { + "runId": run_id, + "configKey": graph_key, + "variationKey": "", + "version": 1, + "modelName": "", + "providerName": "", + "graphKey": graph_key, + } + environment_id = _try_get_environment_id() + if environment_id: + track_data["environmentId"] = environment_id + return track_data def _usable_context_key(value: Any) -> str | None: diff --git a/packages/client/tests/test_graph.py b/packages/client/tests/test_graph.py index a172e021..c3379682 100644 --- a/packages/client/tests/test_graph.py +++ b/packages/client/tests/test_graph.py @@ -259,6 +259,21 @@ async def test_graph_invocation_success_tracked( events = [c[0][0] for c in mock_ld_client.track.call_args_list] assert "$ld:ai:graph:invocation_success" in events + async def test_graph_events_carry_the_environment_id( + self, mock_ld_client: MagicMock, monkeypatch: pytest.MonkeyPatch + ) -> None: + monkeypatch.setenv("LD_ENVIRONMENT_ID", "env-abc") + g = graph("graph-key", handlers=[_make_handler()]) + await g.invoke("hi", CONTEXT) + graph_calls = [ + c + for c in mock_ld_client.track.call_args_list + if str(c[0][0]).startswith("$ld:ai:graph:") + ] + assert graph_calls + for c in graph_calls: + assert c[0][2]["environmentId"] == "env-abc" + async def test_graph_invocation_failure_tracked_on_error( self, mock_ld_client: MagicMock ) -> None: diff --git a/packages/client/tests/test_graph_stream.py b/packages/client/tests/test_graph_stream.py index 2b1cd47c..764dbef3 100644 --- a/packages/client/tests/test_graph_stream.py +++ b/packages/client/tests/test_graph_stream.py @@ -818,6 +818,37 @@ async def fn( for hs in handler_spans: assert hs.get_span_context().trace_id == gctx.trace_id + @pytest.mark.parametrize("mode", ["invoke", "stream"]) + async def test_graph_span_carries_config_identity( + self, mock_ld_client: MagicMock, mode: str + ) -> None: + """The graph() span is found by config key, like a native adapter's.""" + g = graph("graph-key", handlers=[_make_streaming_handler(["ok"])]) + if mode == "invoke": + await g.invoke("hi", CONTEXT) + else: + await _collect(g.stream("hi", CONTEXT)) + + graph_spans = [s for s in _finished() if s.name == "launchdarkly.graph"] + assert len(graph_spans) == 1 + attrs = graph_spans[0].attributes or {} + assert attrs["launchdarkly.operation.type"] == "gen_ai" + assert attrs["launchdarkly.config.key"] == "graph-key" + assert attrs["launchdarkly.graph.key"] == "graph-key" + + success = [ + c[0][2] + for c in mock_ld_client.track.call_args_list + if c[0][0] == "$ld:ai:graph:invocation_success" + ] + assert len(success) == 1 + assert attrs["launchdarkly.run.id"] == success[0]["runId"] + assert attrs["launchdarkly.variation.key"] == success[0]["variationKey"] + + flag_events = [e for e in graph_spans[0].events if e.name == "feature_flag"] + assert len(flag_events) == 1 + assert (flag_events[0].attributes or {})["feature_flag.key"] == "graph-key" + async def test_graph_parent_captured_at_stream_call_time( self, mock_ld_client: MagicMock ) -> None: diff --git a/packages/client/tests/test_utils.py b/packages/client/tests/test_utils.py index 88f05298..ace49b27 100644 --- a/packages/client/tests/test_utils.py +++ b/packages/client/tests/test_utils.py @@ -18,6 +18,7 @@ parse_template, parse_usage, ) +from launchdarkly_ai_server.utils import make_graph_track_data # --------------------------------------------------------------------------- # ?3.1 parse_template @@ -305,6 +306,47 @@ def test_omits_model_key_and_version_when_absent(self) -> None: assert "modelKey" not in td assert "modelVersion" not in td + def test_carries_the_environment_id(self, monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.setenv("LD_ENVIRONMENT_ID", "env-abc") + td = make_track_data( + self._node({"variationKey": "v1", "version": 1}), "graph-key", "run-1" + ) + assert td["environmentId"] == "env-abc" + + def test_omits_environment_id_when_unset( + self, monkeypatch: pytest.MonkeyPatch + ) -> None: + monkeypatch.delenv("LD_ENVIRONMENT_ID", raising=False) + td = make_track_data( + self._node({"variationKey": "v1", "version": 1}), "graph-key", "run-1" + ) + assert "environmentId" not in td + + +class TestMakeGraphTrackData: + """The payload a native graph adapter puts on its ``launchdarkly.graph`` span.""" + + def test_uses_the_graph_key_as_the_config_key( + self, monkeypatch: pytest.MonkeyPatch + ) -> None: + monkeypatch.setenv("LD_ENVIRONMENT_ID", "env-abc") + assert make_graph_track_data("graph-key", "run-1") == { + "runId": "run-1", + "configKey": "graph-key", + "variationKey": "", + "version": 1, + "modelName": "", + "providerName": "", + "graphKey": "graph-key", + "environmentId": "env-abc", + } + + def test_omits_environment_id_when_unset( + self, monkeypatch: pytest.MonkeyPatch + ) -> None: + monkeypatch.delenv("LD_ENVIRONMENT_ID", raising=False) + assert "environmentId" not in make_graph_track_data("graph-key", "run-1") + class TestModelStampsFromMeta: """§3.10 model stamps — malformed ``modelVersion`` is omitted, never raises.""" diff --git a/packages/langchain-agents/src/launchdarkly_ai_langchain_agents/native_graph.py b/packages/langchain-agents/src/launchdarkly_ai_langchain_agents/native_graph.py index b0635866..8ab8d09a 100644 --- a/packages/langchain-agents/src/launchdarkly_ai_langchain_agents/native_graph.py +++ b/packages/langchain-agents/src/launchdarkly_ai_langchain_agents/native_graph.py @@ -19,8 +19,10 @@ get_client, make_track_data, parse_template, + set_ld_span_attributes, to_ld_context, ) +from launchdarkly_ai_server.utils import make_graph_track_data from .messages import to_lang_chain_messages @@ -167,13 +169,6 @@ async def invoke( else None ) - tracer_name = "@launchdarkly/ai-langchain-agents" - if _HAS_OTEL: - span = trace.get_tracer(tracer_name).start_span("launchdarkly.graph") - span.set_attribute("launchdarkly.graph.key", def_obj.key) - else: - span = None - start_time = time.monotonic() run_id = str(uuid.uuid4()) path: list[str] = [] @@ -368,6 +363,20 @@ async def _pre_visit(node_key: str) -> None: else [HumanMessage(input_text)] ) + tracer_name = "@launchdarkly/ai-langchain-agents" + if _HAS_OTEL: + span = trace.get_tracer(tracer_name).start_span("launchdarkly.graph") + else: + span = None + if span: + set_ld_span_attributes( + span, + { + "__ld": make_graph_track_data(def_obj.key, run_id), + "ldContext": raw_ld_context, + }, + ) + try: result = await compiled.ainvoke({"messages": initial_messages}) if span: @@ -378,7 +387,7 @@ async def _pre_visit(node_key: str) -> None: span.set_status(SpanStatusCode.ERROR, str(exc)) span.end() if ld_context: - td = make_track_data(root, def_obj.key, run_id) + td = make_graph_track_data(def_obj.key, run_id) get_client().track("$ld:ai:graph:invocation_failure", ld_context, td, 1) raise @@ -412,13 +421,13 @@ def _content_str(msg: Any) -> str: span.end() if ld_context: - root_td = make_track_data(root, def_obj.key, run_id) + graph_td = make_graph_track_data(def_obj.key, run_id) client = get_client() - client.track("$ld:ai:graph:duration:total", ld_context, root_td, duration) + client.track("$ld:ai:graph:duration:total", ld_context, graph_td, duration) client.track( - "$ld:ai:graph:total_tokens", ld_context, root_td, total_usage["total"] + "$ld:ai:graph:total_tokens", ld_context, graph_td, total_usage["total"] ) - client.track("$ld:ai:graph:invocation_success", ld_context, root_td, 1) + client.track("$ld:ai:graph:invocation_success", ld_context, graph_td, 1) return {"response": final_output, "usage": total_usage} diff --git a/packages/langchain-agents/tests/test_native_graph.py b/packages/langchain-agents/tests/test_native_graph.py index 53457ddc..2a79e856 100644 --- a/packages/langchain-agents/tests/test_native_graph.py +++ b/packages/langchain-agents/tests/test_native_graph.py @@ -436,6 +436,158 @@ async def test_invocation_success_and_duration_tracked(self) -> None: assert "$ld:ai:graph:invocation_success" in track_calls assert "$ld:ai:graph:path" not in track_calls + @pytest.mark.asyncio + async def test_graph_span_carries_config_identity( + self, monkeypatch: pytest.MonkeyPatch + ) -> None: + """The graph span must identify the config so Monitoring can link the trace.""" + import launchdarkly_ai_langchain_agents.native_graph as ng_mod + + monkeypatch.setenv("LD_ENVIRONMENT_ID", "env-123") + mock_span = MagicMock() + mock_trace = MagicMock() + mock_trace.get_tracer.return_value.start_span.return_value = mock_span + + ai_msg = _make_ai_msg("done") + mocks = _make_langgraph_mocks(ai_msg) + graph_def = _make_graph_def() + ctx = {"kind": "user", "key": "u1"} + + with _patch_imports(mocks): + with patch.object(ng_mod, "get_client", return_value=MagicMock()): + with patch.object(ng_mod, "trace", mock_trace): + with patch.object(ng_mod, "_HAS_OTEL", True): + await to_lang_graph( + _make_def_promise(graph_def), + opts={"context": ctx}, + ).invoke("hi") + + attrs = {c[0][0]: c[0][1] for c in mock_span.set_attribute.call_args_list} + assert attrs["launchdarkly.operation.type"] == "gen_ai" + assert attrs["launchdarkly.config.key"] == "test-graph" + assert attrs["launchdarkly.graph.key"] == "test-graph" + assert attrs["launchdarkly.run.id"] + assert "launchdarkly.variation.key" in attrs + assert attrs["context.contextKeys.user"] == "u1" + + events = {c[0][0]: c[0][1] for c in mock_span.add_event.call_args_list} + assert events["feature_flag"] == { + "feature_flag.key": "test-graph", + "feature_flag.provider.name": "LaunchDarkly", + "feature_flag.set.id": "env-123", + "feature_flag.context.id": "u1", + "feature_flag.contextKeys": '{"user":"u1"}', + } + + @pytest.mark.asyncio + async def test_tracking_events_carry_the_environment_id( + self, monkeypatch: pytest.MonkeyPatch + ) -> None: + monkeypatch.setenv("LD_ENVIRONMENT_ID", "env-123") + payloads: list[Any] = [] + mock_ld_client = MagicMock() + mock_ld_client.track = MagicMock( + side_effect=lambda evt, ctx, data, val: payloads.append(data) + ) + + ai_msg = _make_ai_msg("done") + mocks = _make_langgraph_mocks(ai_msg) + graph_def = _make_graph_def() + ctx = {"kind": "user", "key": "u1"} + + with _patch_imports(mocks): + with patch( + "launchdarkly_ai_langchain_agents.native_graph.get_client", + return_value=mock_ld_client, + ): + await to_lang_graph( + _make_def_promise(graph_def), + opts={"context": ctx}, + ).invoke("hi") + node_fn = mocks["_node_fns"]["root"] + await node_fn({"messages": []}) + + assert payloads + for data in payloads: + assert data["environmentId"] == "env-123" + + @pytest.mark.asyncio + async def test_graph_events_are_keyed_to_the_graph(self) -> None: + """Graph-level events carry the graph key, like the graph span and graph().""" + calls: list[tuple[str, Any]] = [] + mock_ld_client = MagicMock() + mock_ld_client.track = MagicMock( + side_effect=lambda evt, ctx, data, val: calls.append((evt, data)) + ) + + mocks = _make_langgraph_mocks(_make_ai_msg("done")) + graph_def = _make_graph_def() + ctx = {"kind": "user", "key": "u1"} + + with _patch_imports(mocks): + with patch( + "launchdarkly_ai_langchain_agents.native_graph.get_client", + return_value=mock_ld_client, + ): + await to_lang_graph( + _make_def_promise(graph_def), + opts={"context": ctx}, + ).invoke("hi") + + graph_events = [ + (e, d) + for e, d in calls + if e + in ( + "$ld:ai:graph:invocation_success", + "$ld:ai:graph:duration:total", + "$ld:ai:graph:total_tokens", + ) + ] + assert len(graph_events) == 3 + for _, data in graph_events: + assert data["configKey"] == "test-graph" + assert data["graphKey"] == "test-graph" + + @pytest.mark.asyncio + async def test_no_span_left_open_when_graph_setup_fails(self) -> None: + """A setup error before the run must not leave a started span un-ended.""" + import launchdarkly_ai_langchain_agents.native_graph as ng_mod + + mock_span = MagicMock() + mock_trace = MagicMock() + mock_trace.get_tracer.return_value.start_span.return_value = mock_span + + mocks = _make_langgraph_mocks(_make_ai_msg("done")) + + class _BadCompileStateGraph: + def __init__(self, *a: Any, **kw: Any) -> None: + pass + + def add_node(self, *a: Any, **kw: Any) -> None: + pass + + def add_edge(self, *a: Any, **kw: Any) -> None: + pass + + def add_conditional_edges(self, *a: Any, **kw: Any) -> None: + pass + + def compile(self) -> Any: + raise ValueError("bad graph") + + mocks["langgraph.graph"].StateGraph = _BadCompileStateGraph + graph_def = _make_graph_def() + + with _patch_imports(mocks): + with patch.object(ng_mod, "trace", mock_trace): + with patch.object(ng_mod, "_HAS_OTEL", True): + with pytest.raises(ValueError, match="bad graph"): + await to_lang_graph(_make_def_promise(graph_def)).invoke("hi") + + started = mock_trace.get_tracer.return_value.start_span.call_count + assert mock_span.end.call_count == started + @pytest.mark.asyncio async def test_node_function_emits_graph_node(self) -> None: track_calls: list[tuple[str, Any, Any]] = [] diff --git a/packages/openai-agents/src/launchdarkly_ai_openai_agents/native_graph.py b/packages/openai-agents/src/launchdarkly_ai_openai_agents/native_graph.py index a8b8a4a1..6014014a 100644 --- a/packages/openai-agents/src/launchdarkly_ai_openai_agents/native_graph.py +++ b/packages/openai-agents/src/launchdarkly_ai_openai_agents/native_graph.py @@ -19,8 +19,10 @@ get_client, make_track_data, parse_template, + set_ld_span_attributes, to_ld_context, ) +from launchdarkly_ai_server.utils import make_graph_track_data from .handler import _parse_message_content, _to_openai_agent_items @@ -129,13 +131,6 @@ async def invoke( else None ) - tracer_name = "@launchdarkly/ai-openai-agents" - if _HAS_OTEL: - span = trace.get_tracer(tracer_name).start_span("launchdarkly.graph") - span.set_attribute("launchdarkly.graph.key", def_obj.key) - else: - span = None - start_time = time.monotonic() run_id = str(uuid.uuid4()) path: list[str] = [] @@ -256,6 +251,20 @@ async def on_agent_start(self, context: Any, agent: Any) -> None: ) root_prompt = _to_openai_agent_items(turns) + tracer_name = "@launchdarkly/ai-openai-agents" + if _HAS_OTEL: + span = trace.get_tracer(tracer_name).start_span("launchdarkly.graph") + else: + span = None + if span: + set_ld_span_attributes( + span, + { + "__ld": make_graph_track_data(def_obj.key, run_id), + "ldContext": raw_ld_context, + }, + ) + try: result = await Runner.run(root_agent, root_prompt, hooks=hooks) if span: @@ -266,7 +275,7 @@ async def on_agent_start(self, context: Any, agent: Any) -> None: span.set_status(SpanStatusCode.ERROR, str(exc)) span.end() if ld_context: - td = make_track_data(root, def_obj.key, run_id) + td = make_graph_track_data(def_obj.key, run_id) get_client().track("$ld:ai:graph:invocation_failure", ld_context, td, 1) raise @@ -303,11 +312,13 @@ async def on_agent_start(self, context: Any, agent: Any) -> None: span.end() if ld_context: - root_td = make_track_data(root, def_obj.key, run_id) + graph_td = make_graph_track_data(def_obj.key, run_id) client = get_client() - client.track("$ld:ai:graph:duration:total", ld_context, root_td, duration) - client.track("$ld:ai:graph:total_tokens", ld_context, root_td, total_tokens) - client.track("$ld:ai:graph:invocation_success", ld_context, root_td, 1) + client.track("$ld:ai:graph:duration:total", ld_context, graph_td, duration) + client.track( + "$ld:ai:graph:total_tokens", ld_context, graph_td, total_tokens + ) + client.track("$ld:ai:graph:invocation_success", ld_context, graph_td, 1) return {"response": final_output, "usage": total_usage} diff --git a/packages/openai-agents/tests/test_native_graph.py b/packages/openai-agents/tests/test_native_graph.py index 2df4696f..7c465b2d 100644 --- a/packages/openai-agents/tests/test_native_graph.py +++ b/packages/openai-agents/tests/test_native_graph.py @@ -573,6 +573,172 @@ async def test_otel_span_has_graph_key_attribute(self) -> None: assert "launchdarkly.graph.key" in set_attr_calls assert set_attr_calls["launchdarkly.graph.key"] == "test-graph" + @pytest.mark.asyncio + async def test_graph_span_carries_config_identity( + self, monkeypatch: pytest.MonkeyPatch + ) -> None: + """The graph span must identify the config so Monitoring can link the trace.""" + monkeypatch.setenv("LD_ENVIRONMENT_ID", "env-123") + mock_span = MagicMock() + mock_trace = MagicMock() + mock_trace.get_tracer.return_value.start_span.return_value = mock_span + + run_result = _make_run_result("done") + agents_mock = _make_agents_mock(run_result) + graph_def = _make_graph_def() + ctx = {"kind": "user", "key": "u1"} + + with patch( + "importlib.import_module", + side_effect=lambda n: agents_mock if n == "agents" else __import__(n), + ): + with patch.object(_openai_ng, "get_client", return_value=MagicMock()): + with patch.object(_openai_ng, "trace", mock_trace): + with patch.object(_openai_ng, "_HAS_OTEL", True): + await to_openai_agents( + _make_def_promise(graph_def), + opts={"context": ctx}, + ).invoke("hi") + + attrs = {c[0][0]: c[0][1] for c in mock_span.set_attribute.call_args_list} + assert attrs["launchdarkly.operation.type"] == "gen_ai" + assert attrs["launchdarkly.config.key"] == "test-graph" + assert attrs["launchdarkly.graph.key"] == "test-graph" + assert attrs["launchdarkly.run.id"] + assert "launchdarkly.variation.key" in attrs + assert attrs["context.contextKeys.user"] == "u1" + + events = {c[0][0]: c[0][1] for c in mock_span.add_event.call_args_list} + assert events["feature_flag"] == { + "feature_flag.key": "test-graph", + "feature_flag.provider.name": "LaunchDarkly", + "feature_flag.set.id": "env-123", + "feature_flag.context.id": "u1", + "feature_flag.contextKeys": '{"user":"u1"}', + } + + @pytest.mark.asyncio + async def test_tracking_events_carry_the_environment_id( + self, monkeypatch: pytest.MonkeyPatch + ) -> None: + monkeypatch.setenv("LD_ENVIRONMENT_ID", "env-123") + payloads: list[Any] = [] + mock_ld_client = MagicMock() + mock_ld_client.track = MagicMock( + side_effect=lambda evt, ctx, data, val: payloads.append(data) + ) + + run_result = _make_run_result("done") + agents_mock = _make_agents_mock(run_result) + graph_def = _make_graph_def() + ctx = {"kind": "user", "key": "u1"} + + with patch( + "importlib.import_module", + side_effect=lambda n: agents_mock if n == "agents" else __import__(n), + ): + with patch.object(_openai_ng, "get_client", return_value=mock_ld_client): + await to_openai_agents( + _make_def_promise(graph_def), + opts={"context": ctx}, + ).invoke("hi") + + assert payloads + for data in payloads: + assert data["environmentId"] == "env-123" + + @pytest.mark.asyncio + async def test_graph_events_are_keyed_to_the_graph(self) -> None: + """Graph-level events carry the graph key, like the graph span and graph().""" + calls: list[tuple[str, Any]] = [] + mock_ld_client = MagicMock() + mock_ld_client.track = MagicMock( + side_effect=lambda evt, ctx, data, val: calls.append((evt, data)) + ) + + run_result = _make_run_result("done") + agents_mock = _make_agents_mock(run_result) + graph_def = _make_graph_def() + ctx = {"kind": "user", "key": "u1"} + + with patch( + "importlib.import_module", + side_effect=lambda n: agents_mock if n == "agents" else __import__(n), + ): + with patch.object(_openai_ng, "get_client", return_value=mock_ld_client): + await to_openai_agents( + _make_def_promise(graph_def), + opts={"context": ctx}, + ).invoke("hi") + + graph_events = [ + (e, d) + for e, d in calls + if e + in ( + "$ld:ai:graph:invocation_success", + "$ld:ai:graph:duration:total", + "$ld:ai:graph:total_tokens", + ) + ] + assert len(graph_events) == 3 + for _, data in graph_events: + assert data["configKey"] == "test-graph" + assert data["graphKey"] == "test-graph" + + @pytest.mark.asyncio + async def test_invocation_failure_is_keyed_to_the_graph(self) -> None: + calls: list[tuple[str, Any]] = [] + mock_ld_client = MagicMock() + mock_ld_client.track = MagicMock( + side_effect=lambda evt, ctx, data, val: calls.append((evt, data)) + ) + + agents_mock = _make_agents_mock(_make_run_result("done")) + agents_mock.Runner.run = AsyncMock(side_effect=RuntimeError("provider error")) + graph_def = _make_graph_def() + ctx = {"kind": "user", "key": "u1"} + + with patch( + "importlib.import_module", + side_effect=lambda n: agents_mock if n == "agents" else __import__(n), + ): + with patch.object(_openai_ng, "get_client", return_value=mock_ld_client): + with pytest.raises(RuntimeError): + await to_openai_agents( + _make_def_promise(graph_def), + opts={"context": ctx}, + ).invoke("hi") + + failures = [d for e, d in calls if e == "$ld:ai:graph:invocation_failure"] + assert len(failures) == 1 + assert failures[0]["configKey"] == "test-graph" + + @pytest.mark.asyncio + async def test_no_span_left_open_when_agent_setup_fails(self) -> None: + """A setup error before the run must not leave a started span un-ended.""" + mock_span = MagicMock() + mock_trace = MagicMock() + mock_trace.get_tracer.return_value.start_span.return_value = mock_span + + agents_mock = _make_agents_mock(_make_run_result("done")) + agents_mock.Agent = MagicMock(side_effect=ValueError("bad agent")) + graph_def = _make_graph_def() + + with patch( + "importlib.import_module", + side_effect=lambda n: agents_mock if n == "agents" else __import__(n), + ): + with patch.object(_openai_ng, "trace", mock_trace): + with patch.object(_openai_ng, "_HAS_OTEL", True): + with pytest.raises(ValueError, match="bad agent"): + await to_openai_agents(_make_def_promise(graph_def)).invoke( + "hi" + ) + + started = mock_trace.get_tracer.return_value.start_span.call_count + assert mock_span.end.call_count == started + @pytest.mark.asyncio async def test_agent_end_hook_emits_generation_success(self) -> None: """agent_end hook must emit $ld:ai:generation:success for the agent's node."""