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..6fb6e6b7 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 @@ -236,20 +238,30 @@ async def invoke( ) raw_handlers: dict[str, Any] = _opts.get("tool_handlers") or {} + start_time = time.monotonic() + run_id = str(uuid.uuid4()) 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()) path: list[str] = [] total_usage = {"input": 0, "output": 0, "total": 0} subagent_tool_ctx: dict[str, Any] = {} + # One try covers setup and the run, so a setup error (a sub-agent tool + # that fails to build) is recorded on the span and tracked as an + # invocation failure, like a run error. The span ends exactly once, in + # the finally. try: + if span: + set_ld_span_attributes( + span, + { + "__ld": make_graph_track_data(def_obj.key, run_id), + "ldContext": raw_ld_context, + }, + ) async def _build_node(node: GraphNode) -> None: if node.key == root.key: @@ -351,29 +363,17 @@ async def _subagent_execute( ) root_start = time.monotonic() - try: - result = await _run_query( - root, - input_text, - vs, - raw_handlers, - ld_context, - def_obj.key, - run_id, - root_child_tools, - history, - ) - except Exception as exc: - if span: - span.record_exception(exc) - span.set_status(SpanStatusCode.ERROR, str(exc)) - span.end() - if ld_context: - td = make_track_data(root, def_obj.key, run_id) - get_client().track( - "$ld:ai:graph:invocation_failure", ld_context, td, 1 - ) - raise + result = await _run_query( + root, + input_text, + vs, + raw_handlers, + ld_context, + def_obj.key, + run_id, + root_child_tools, + history, + ) final_output = result["output"] root_usage = result["usage"] @@ -407,30 +407,38 @@ async def _subagent_execute( span.set_attribute("gen_ai.usage.input_tokens", total_usage["input"]) span.set_attribute("gen_ai.usage.output_tokens", total_usage["output"]) span.set_attribute("gen_ai.usage.total_tokens", total_usage["total"]) - span.set_status(SpanStatusCode.OK) - 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) + if span: + span.set_status(SpanStatusCode.OK) return {"response": final_output, "usage": total_usage} - except Exception as exc: if span: span.record_exception(exc) span.set_status(SpanStatusCode.ERROR, str(exc)) - span.end() + if ld_context: + get_client().track( + "$ld:ai:graph:invocation_failure", + ld_context, + make_graph_track_data(def_obj.key, run_id), + 1, + ) raise + finally: + if span: + span.end() return types.SimpleNamespace(invoke=invoke) diff --git a/packages/claude-agents/tests/test_native_graph.py b/packages/claude-agents/tests/test_native_graph.py index e01900f4..7a55aba4 100644 --- a/packages/claude-agents/tests/test_native_graph.py +++ b/packages/claude-agents/tests/test_native_graph.py @@ -133,6 +133,55 @@ async def _query(**kwargs: Any) -> AsyncIterator[Any]: # --------------------------------------------------------------------------- +def _recording_trace() -> tuple[Any, Any, list[Any]]: + """A real tracer behind the module's ``trace`` lookup, with every span's ``end`` spied. + + Returns the stand-in ``trace`` module, the exporter holding finished spans, and the + ``end`` spies, so a test can check both what was exported and how often each span ended. + """ + from opentelemetry.sdk.trace import TracerProvider + from opentelemetry.sdk.trace.export import SimpleSpanProcessor + from opentelemetry.sdk.trace.export.in_memory_span_exporter import ( + InMemorySpanExporter, + ) + + exporter = InMemorySpanExporter() + provider = TracerProvider() + provider.add_span_processor(SimpleSpanProcessor(exporter)) + tracer = provider.get_tracer("test") + end_spies: list[Any] = [] + + def _start_span(name: str, *args: Any, **kwargs: Any) -> Any: + span = tracer.start_span(name, *args, **kwargs) + span.end = MagicMock(wraps=span.end) # type: ignore[method-assign] + end_spies.append(span.end) + return span + + fake_trace = MagicMock() + fake_trace.get_tracer.return_value.start_span.side_effect = _start_span + return fake_trace, exporter, end_spies + + +def _assert_one_failed_graph_span( + exporter: Any, end_spies: list[Any], message: str, track_calls: list[Any] +) -> None: + from opentelemetry.trace import StatusCode + + assert [spy.call_count for spy in end_spies] == [1] + spans = [s for s in exporter.get_finished_spans() if s.name == "launchdarkly.graph"] + assert len(spans) == 1 + span = spans[0] + assert span.status.status_code is StatusCode.ERROR + assert span.status.description == message + assert [e.name for e in span.events].count("exception") == 1 + assert (span.attributes or {}).get("launchdarkly.graph.key") == "test-graph" + + failures = [d for e, d in track_calls if e == "$ld:ai:graph:invocation_failure"] + assert len(failures) == 1 + assert failures[0]["configKey"] == "test-graph" + assert failures[0]["graphKey"] == "test-graph" + + class TestToClaudeAgentsTopology: @pytest.mark.asyncio async def test_each_graph_node_translated(self) -> None: @@ -587,6 +636,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] = [] @@ -698,7 +865,44 @@ async def _bad_query(**kwargs: Any) -> AsyncIterator[Any]: "hi" ) - mock_span.end.assert_called() + mock_span.end.assert_called_once() + mock_span.record_exception.assert_called_once() + + @pytest.mark.asyncio + async def test_subagent_setup_error_fails_the_graph_span(self) -> None: + """A sub-agent build error before the run ends one ERROR span and tracks the failure.""" + fake_trace, exporter, end_spies = _recording_trace() + track_calls: list[tuple[str, Any]] = [] + mock_ld_client = MagicMock() + mock_ld_client.track = MagicMock( + side_effect=lambda evt, ctx, data, val: track_calls.append((evt, data)) + ) + mock_sdk = _make_sdk_mock("done") + graph_def = _make_graph_def() + + with patch( + "importlib.import_module", + side_effect=lambda n: ( + mock_sdk if n == "claude_agent_sdk" else __import__(n) + ), + ): + with patch.object(_claude_ng, "trace", fake_trace): + with patch.object(_claude_ng, "_HAS_OTEL", True): + with patch.object( + _claude_ng, "get_client", return_value=mock_ld_client + ): + with patch.object( + _claude_ng, + "_reverse_traverse", + AsyncMock(side_effect=ValueError("bad sub-agent")), + ): + with pytest.raises(ValueError, match="bad sub-agent"): + await to_claude_agents( + _make_def_promise(graph_def), + opts={"context": {"kind": "user", "key": "u1"}}, + ).invoke("hi") + + _assert_one_failed_graph_span(exporter, end_spies, "bad sub-agent", track_calls) @pytest.mark.asyncio async def test_build_tool_mcp_throws_when_tool_not_in_handlers(self) -> None: 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..ffe80c4c 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,259 +169,283 @@ async def invoke( else None ) + start_time = time.monotonic() + run_id = str(uuid.uuid4()) 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] = [] - total_usage = {"input": 0, "output": 0, "total": 0} - edges_from = def_obj.edges_from - - # WorkflowState must reference add_messages from module-level scope. - # With `from __future__ import annotations`, LangGraph resolves annotations - # via get_type_hints() in the *module* global namespace — a local variable - # would cause NameError at StateGraph(WorkflowState) time (AIC-2948). - class WorkflowState(TypedDict): - messages: Annotated[list[Any], add_messages] - - builder = StateGraph(WorkflowState) - - async def _traverse_node(node: GraphNode) -> None: - node_key = _sanitize_name(node.key) - outgoing = edges_from(node.key) - is_terminal = node.is_terminal - is_multi_child = len(outgoing) > 1 - - # Get chat model - if model_factory: - chat_model = model_factory(node) - else: - lc_openai = importlib.import_module("langchain_openai") - model_cfg = node.config.get("model") or {} - raw = model_cfg.get("parameters") - kwargs = dict(raw) if isinstance(raw, dict) else {} - kwargs["model"] = model_cfg.get("name") or "gpt-4o" - chat_model = lc_openai.ChatOpenAI(**kwargs) - - regular_tools = _build_node_tools(node, tool_handlers) - - # Handoff tools (return Command to route) - lc_tools = importlib.import_module("langchain_core.tools") - tool_fn = lc_tools.tool - langgraph_types = importlib.import_module("langgraph.types") - Command = langgraph_types.Command - - handoff_tools = [] - for edge in outgoing: - target_key = _sanitize_name(edge.target_key) - - async def _handoff_exec( - _target: str = target_key, - _node: GraphNode = node, - ) -> Any: - if ld_context: - td = make_track_data(_node, def_obj.key, run_id) - get_client().track( - "$ld:ai:graph:handoff_success", ld_context, td, 1 - ) - return Command(goto=_target) - - ht = tool_fn( - f"transfer_to_{_sanitize_name(edge.target_key)}", - _handoff_exec, - description=f"Transfer control to the {edge.target_key} agent", - args_schema={}, + # One try covers setup and the run, so a setup error (a graph that + # fails to compile) is recorded on the span and tracked as an invocation + # failure, like a run error. The span ends exactly once, in the finally. + try: + if span: + set_ld_span_attributes( + span, + { + "__ld": make_graph_track_data(def_obj.key, run_id), + "ldContext": raw_ld_context, + }, ) - handoff_tools.append(ht) - - all_tools = regular_tools + handoff_tools + path: list[str] = [] + total_usage = {"input": 0, "output": 0, "total": 0} + edges_from = def_obj.edges_from + + # WorkflowState must reference add_messages from module-level scope. + # With `from __future__ import annotations`, LangGraph resolves annotations + # via get_type_hints() in the *module* global namespace — a local variable + # would cause NameError at StateGraph(WorkflowState) time (AIC-2948). + class WorkflowState(TypedDict): + messages: Annotated[list[Any], add_messages] + + builder = StateGraph(WorkflowState) + + async def _traverse_node(node: GraphNode) -> None: + node_key = _sanitize_name(node.key) + outgoing = edges_from(node.key) + is_terminal = node.is_terminal + is_multi_child = len(outgoing) > 1 + + # Get chat model + if model_factory: + chat_model = model_factory(node) + else: + lc_openai = importlib.import_module("langchain_openai") + model_cfg = node.config.get("model") or {} + raw = model_cfg.get("parameters") + kwargs = dict(raw) if isinstance(raw, dict) else {} + kwargs["model"] = model_cfg.get("name") or "gpt-4o" + chat_model = lc_openai.ChatOpenAI(**kwargs) + + regular_tools = _build_node_tools(node, tool_handlers) + + # Handoff tools (return Command to route) + lc_tools = importlib.import_module("langchain_core.tools") + tool_fn = lc_tools.tool + langgraph_types = importlib.import_module("langgraph.types") + Command = langgraph_types.Command + + handoff_tools = [] + for edge in outgoing: + target_key = _sanitize_name(edge.target_key) + + async def _handoff_exec( + _target: str = target_key, + _node: GraphNode = node, + ) -> Any: + if ld_context: + td = make_track_data(_node, def_obj.key, run_id) + get_client().track( + "$ld:ai:graph:handoff_success", ld_context, td, 1 + ) + return Command(goto=_target) + + ht = tool_fn( + f"transfer_to_{_sanitize_name(edge.target_key)}", + _handoff_exec, + description=f"Transfer control to the {edge.target_key} agent", + args_schema={}, + ) + handoff_tools.append(ht) + + all_tools = regular_tools + handoff_tools + + async def _node_fn( + state: WorkflowState, _node: GraphNode = node + ) -> dict[str, Any]: + if _node.key not in path: + index = len(path) + path.append(_node.key) + if ld_context: + node_td = make_track_data(_node, def_obj.key, run_id) + get_client().track( + "$ld:ai:graph:node", + ld_context, + {**node_td, "nodeKey": _node.key, "index": index}, + 1, + ) + node_start = time.monotonic() + + system_prompt = _build_system_prompt(_node, vs) + conv_messages: list[Any] = state.get("messages", []) + full_messages = ( + [SystemMessage(system_prompt), *conv_messages] + if system_prompt + else list(conv_messages) + ) - async def _node_fn( - state: WorkflowState, _node: GraphNode = node - ) -> dict[str, Any]: - if _node.key not in path: - index = len(path) - path.append(_node.key) - if ld_context: - node_td = make_track_data(_node, def_obj.key, run_id) - get_client().track( - "$ld:ai:graph:node", - ld_context, - {**node_td, "nodeKey": _node.key, "index": index}, - 1, + bound = ( + chat_model.bind_tools( + all_tools, + **( + {"parallel_tool_calls": False} if is_multi_child else {} + ), ) - node_start = time.monotonic() - - system_prompt = _build_system_prompt(_node, vs) - conv_messages: list[Any] = state.get("messages", []) - full_messages = ( - [SystemMessage(system_prompt), *conv_messages] - if system_prompt - else list(conv_messages) - ) - - bound = ( - chat_model.bind_tools( - all_tools, - **({"parallel_tool_calls": False} if is_multi_child else {}), + if all_tools + else chat_model ) - if all_tools - else chat_model - ) - result_msg = await bound.ainvoke(full_messages) - usage = _extract_usage(result_msg) - total_usage["input"] += usage["input"] - total_usage["output"] += usage["output"] - total_usage["total"] += usage["total"] - - if ld_context: - td = make_track_data(_node, def_obj.key, run_id) - dur = int((time.monotonic() - node_start) * 1000) - client = get_client() - client.track("$ld:ai:duration:total", ld_context, td, dur) - client.track("$ld:ai:generation:success", ld_context, td, 1) - if usage["total"] > 0: - client.track( - "$ld:ai:tokens:total", ld_context, td, usage["total"] + result_msg = await bound.ainvoke(full_messages) + usage = _extract_usage(result_msg) + total_usage["input"] += usage["input"] + total_usage["output"] += usage["output"] + total_usage["total"] += usage["total"] + + if ld_context: + td = make_track_data(_node, def_obj.key, run_id) + dur = int((time.monotonic() - node_start) * 1000) + client = get_client() + client.track("$ld:ai:duration:total", ld_context, td, dur) + client.track("$ld:ai:generation:success", ld_context, td, 1) + if usage["total"] > 0: + client.track( + "$ld:ai:tokens:total", ld_context, td, usage["total"] + ) + if usage["input"] > 0: + client.track( + "$ld:ai:tokens:input", ld_context, td, usage["input"] + ) + if usage["output"] > 0: + client.track( + "$ld:ai:tokens:output", ld_context, td, usage["output"] + ) + + return {"messages": [result_msg]} + + builder.add_node(node_key, _node_fn) + + if all_tools: + builder.add_node(f"{node_key}_tools", ToolNode(all_tools)) + + # Edge wiring + if node.key == root.key: + builder.add_edge(START, node_key) + + if is_terminal: + if all_tools: + builder.add_conditional_edges( + node_key, + tools_condition, + {"tools": f"{node_key}_tools", "__end__": END}, ) - if usage["input"] > 0: - client.track( - "$ld:ai:tokens:input", ld_context, td, usage["input"] + builder.add_edge(f"{node_key}_tools", node_key) + else: + builder.add_edge(node_key, END) + elif is_multi_child: + if all_tools: + builder.add_conditional_edges( + node_key, + tools_condition, + {"tools": f"{node_key}_tools", "__end__": END}, ) - if usage["output"] > 0: - client.track( - "$ld:ai:tokens:output", ld_context, td, usage["output"] + builder.add_edge(f"{node_key}_tools", node_key) + else: + builder.add_edge(node_key, END) + else: + child_key = _sanitize_name(outgoing[0].target_key) + if all_tools: + builder.add_conditional_edges( + node_key, + tools_condition, + {"tools": f"{node_key}_tools", "__end__": child_key}, ) + builder.add_edge(f"{node_key}_tools", node_key) + else: + builder.add_edge(node_key, child_key) + + # Pre-order traversal (root first, to mirror TS traverse) + pre_visited: set[str] = set() + + async def _pre_visit(node_key: str) -> None: + if node_key in pre_visited: + return + pre_visited.add(node_key) + node = def_obj.get_node(node_key) + if node: + await _traverse_node(node) + for edge in edges_from(node_key): + await _pre_visit(edge.target_key) + + await _pre_visit(root.key) + + compiled = builder.compile() + + # History is a root-only concern: it seeds the initial message state the + # entry node reads. Downstream nodes are reached through handoffs and see + # the accumulated graph state, never the original `history` array. + initial_messages = ( + to_lang_chain_messages( + compose_history(history=history, user_input=input_text) + ) + if history + else [HumanMessage(input_text)] + ) - return {"messages": [result_msg]} + result = await compiled.ainvoke({"messages": initial_messages}) - builder.add_node(node_key, _node_fn) + duration = int((time.monotonic() - start_time) * 1000) - if all_tools: - builder.add_node(f"{node_key}_tools", ToolNode(all_tools)) + # Extract final output from last AI message + result_messages = ( + result.get("messages", []) if isinstance(result, dict) else [] + ) + last_msg = result_messages[-1] if result_messages else None + + def _content_str(msg: Any) -> str: + if msg is None: + return "" + c = msg.content + if isinstance(c, str): + return c + if isinstance(c, list): + return "".join( + part.get("text", "") if isinstance(part, dict) else "" + for part in c + if isinstance(part, dict) and part.get("type") == "text" + ) + return "" - # Edge wiring - if node.key == root.key: - builder.add_edge(START, node_key) + final_output = _content_str(last_msg) - if is_terminal: - if all_tools: - builder.add_conditional_edges( - node_key, - tools_condition, - {"tools": f"{node_key}_tools", "__end__": END}, - ) - builder.add_edge(f"{node_key}_tools", node_key) - else: - builder.add_edge(node_key, END) - elif is_multi_child: - if all_tools: - builder.add_conditional_edges( - node_key, - tools_condition, - {"tools": f"{node_key}_tools", "__end__": END}, - ) - builder.add_edge(f"{node_key}_tools", node_key) - else: - builder.add_edge(node_key, END) - else: - child_key = _sanitize_name(outgoing[0].target_key) - if all_tools: - builder.add_conditional_edges( - node_key, - tools_condition, - {"tools": f"{node_key}_tools", "__end__": child_key}, - ) - builder.add_edge(f"{node_key}_tools", node_key) - else: - builder.add_edge(node_key, child_key) - - # Pre-order traversal (root first, to mirror TS traverse) - pre_visited: set[str] = set() - - async def _pre_visit(node_key: str) -> None: - if node_key in pre_visited: - return - pre_visited.add(node_key) - node = def_obj.get_node(node_key) - if node: - await _traverse_node(node) - for edge in edges_from(node_key): - await _pre_visit(edge.target_key) - - await _pre_visit(root.key) - - compiled = builder.compile() - - # History is a root-only concern: it seeds the initial message state the - # entry node reads. Downstream nodes are reached through handoffs and see - # the accumulated graph state, never the original `history` array. - initial_messages = ( - to_lang_chain_messages( - compose_history(history=history, user_input=input_text) - ) - if history - else [HumanMessage(input_text)] - ) + if span: + span.set_attribute("launchdarkly.graph.path", "->".join(path)) + span.set_attribute("gen_ai.usage.input_tokens", total_usage["input"]) + span.set_attribute("gen_ai.usage.output_tokens", total_usage["output"]) + span.set_attribute("gen_ai.usage.total_tokens", total_usage["total"]) + + if ld_context: + graph_td = make_graph_track_data(def_obj.key, run_id) + client = get_client() + client.track( + "$ld:ai:graph:duration:total", ld_context, graph_td, duration + ) + client.track( + "$ld:ai:graph:total_tokens", + ld_context, + graph_td, + total_usage["total"], + ) + client.track("$ld:ai:graph:invocation_success", ld_context, graph_td, 1) - try: - result = await compiled.ainvoke({"messages": initial_messages}) if span: span.set_status(SpanStatusCode.OK) + return {"response": final_output, "usage": total_usage} except Exception as exc: if span: span.record_exception(exc) span.set_status(SpanStatusCode.ERROR, str(exc)) - span.end() if ld_context: - td = make_track_data(root, def_obj.key, run_id) - get_client().track("$ld:ai:graph:invocation_failure", ld_context, td, 1) - raise - - duration = int((time.monotonic() - start_time) * 1000) - - # Extract final output from last AI message - result_messages = result.get("messages", []) if isinstance(result, dict) else [] - last_msg = result_messages[-1] if result_messages else None - - def _content_str(msg: Any) -> str: - if msg is None: - return "" - c = msg.content - if isinstance(c, str): - return c - if isinstance(c, list): - return "".join( - part.get("text", "") if isinstance(part, dict) else "" - for part in c - if isinstance(part, dict) and part.get("type") == "text" + get_client().track( + "$ld:ai:graph:invocation_failure", + ld_context, + make_graph_track_data(def_obj.key, run_id), + 1, ) - return "" - - final_output = _content_str(last_msg) - - if span: - span.set_attribute("launchdarkly.graph.path", "->".join(path)) - span.set_attribute("gen_ai.usage.input_tokens", total_usage["input"]) - span.set_attribute("gen_ai.usage.output_tokens", total_usage["output"]) - span.set_attribute("gen_ai.usage.total_tokens", total_usage["total"]) - span.end() - - if ld_context: - root_td = make_track_data(root, 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_usage["total"] - ) - client.track("$ld:ai:graph:invocation_success", ld_context, root_td, 1) - - return {"response": final_output, "usage": total_usage} + raise + finally: + if span: + span.end() return types.SimpleNamespace(invoke=invoke) diff --git a/packages/langchain-agents/tests/test_native_graph.py b/packages/langchain-agents/tests/test_native_graph.py index 53457ddc..b0934ba8 100644 --- a/packages/langchain-agents/tests/test_native_graph.py +++ b/packages/langchain-agents/tests/test_native_graph.py @@ -236,6 +236,55 @@ def _patch_imports(mocks: dict[str, Any]) -> Any: # --------------------------------------------------------------------------- +def _recording_trace() -> tuple[Any, Any, list[Any]]: + """A real tracer behind the module's ``trace`` lookup, with every span's ``end`` spied. + + Returns the stand-in ``trace`` module, the exporter holding finished spans, and the + ``end`` spies, so a test can check both what was exported and how often each span ended. + """ + from opentelemetry.sdk.trace import TracerProvider + from opentelemetry.sdk.trace.export import SimpleSpanProcessor + from opentelemetry.sdk.trace.export.in_memory_span_exporter import ( + InMemorySpanExporter, + ) + + exporter = InMemorySpanExporter() + provider = TracerProvider() + provider.add_span_processor(SimpleSpanProcessor(exporter)) + tracer = provider.get_tracer("test") + end_spies: list[Any] = [] + + def _start_span(name: str, *args: Any, **kwargs: Any) -> Any: + span = tracer.start_span(name, *args, **kwargs) + span.end = MagicMock(wraps=span.end) # type: ignore[method-assign] + end_spies.append(span.end) + return span + + fake_trace = MagicMock() + fake_trace.get_tracer.return_value.start_span.side_effect = _start_span + return fake_trace, exporter, end_spies + + +def _assert_one_failed_graph_span( + exporter: Any, end_spies: list[Any], message: str, track_calls: list[Any] +) -> None: + from opentelemetry.trace import StatusCode + + assert [spy.call_count for spy in end_spies] == [1] + spans = [s for s in exporter.get_finished_spans() if s.name == "launchdarkly.graph"] + assert len(spans) == 1 + span = spans[0] + assert span.status.status_code is StatusCode.ERROR + assert span.status.description == message + assert [e.name for e in span.events].count("exception") == 1 + assert (span.attributes or {}).get("launchdarkly.graph.key") == "test-graph" + + failures = [d for e, d in track_calls if e == "$ld:ai:graph:invocation_failure"] + assert len(failures) == 1 + assert failures[0]["configKey"] == "test-graph" + assert failures[0]["graphKey"] == "test-graph" + + class TestToLangGraphTopology: @pytest.mark.asyncio async def test_each_graph_node_translated(self) -> None: @@ -436,6 +485,166 @@ 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_graph_setup_error_fails_the_graph_span(self) -> None: + """A compile error before the run ends one ERROR span and tracks the failure.""" + import launchdarkly_ai_langchain_agents.native_graph as ng_mod + + fake_trace, exporter, end_spies = _recording_trace() + track_calls: list[tuple[str, Any]] = [] + mock_ld_client = MagicMock() + mock_ld_client.track = MagicMock( + side_effect=lambda evt, ctx, data, val: track_calls.append((evt, data)) + ) + + 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", fake_trace): + with patch.object(ng_mod, "_HAS_OTEL", True): + with patch.object( + ng_mod, "get_client", return_value=mock_ld_client + ): + with pytest.raises(ValueError, match="bad graph"): + await to_lang_graph( + _make_def_promise(graph_def), + opts={"context": {"kind": "user", "key": "u1"}}, + ).invoke("hi") + + _assert_one_failed_graph_span(exporter, end_spies, "bad graph", track_calls) + @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..7f38a597 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,186 +131,207 @@ async def invoke( else None ) + start_time = time.monotonic() + run_id = str(uuid.uuid4()) 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] = [] - agent_name_to_key: dict[str, str] = {} - agent_ctx: dict[str, Any] = {} - edges_from = def_obj.edges_from - - # Post-order traversal (leaves first) using the edges_from function - visited: set[str] = set() - - async def _visit(node_key: str) -> None: - if node_key in visited: - return - visited.add(node_key) - for edge in edges_from(node_key): - await _visit(edge.target_key) - - node = def_obj.get_node(node_key) - if node is None: - return - - child_handoffs = [] - for edge in edges_from(node_key): - child_agent = agent_ctx.get(edge.target_key) - if child_agent is None: - raise ValueError( - f'Child agent "{edge.target_key}" not built before parent "{node_key}"' - ) - child_handoffs.append(handoff_fn(child_agent)) - - instructions = _build_instructions(node, vs) - tools = _build_node_tools(node, tool_handlers) - agent_name = _sanitize_name(node.key) - agent_name_to_key[agent_name] = node.key - - agent = Agent( - name=agent_name, - model=node.config.get("model", {}).get("name", "gpt-4o"), - **({"instructions": instructions} if instructions else {}), - **({"tools": tools} if tools else {}), - **({"handoffs": child_handoffs} if child_handoffs else {}), - ) - agent_ctx[node.key] = agent - - await _visit(root.key) + # One try covers setup and the run, so a setup error (an agent that + # fails to build) is recorded on the span and tracked as an invocation + # failure, like a run error. The span ends exactly once, in the finally. + try: + if span: + set_ld_span_attributes( + span, + { + "__ld": make_graph_track_data(def_obj.key, run_id), + "ldContext": raw_ld_context, + }, + ) + path: list[str] = [] + agent_name_to_key: dict[str, str] = {} + agent_ctx: dict[str, Any] = {} + edges_from = def_obj.edges_from + + # Post-order traversal (leaves first) using the edges_from function + visited: set[str] = set() + + async def _visit(node_key: str) -> None: + if node_key in visited: + return + visited.add(node_key) + for edge in edges_from(node_key): + await _visit(edge.target_key) - root_agent = agent_ctx.get(root.key) - if root_agent is None: - raise ValueError(f'Root agent "{root.key}" was not built') + node = def_obj.get_node(node_key) + if node is None: + return - # Lifecycle hooks for LD tracking - class _LDHooks(RunHooks): # type: ignore[misc, valid-type] - async def on_agent_end(self, context: Any, agent: Any, output: Any) -> None: - node_key = agent_name_to_key.get(agent.name) - if node_key and ld_context: + child_handoffs = [] + for edge in edges_from(node_key): + child_agent = agent_ctx.get(edge.target_key) + if child_agent is None: + raise ValueError( + f'Child agent "{edge.target_key}" not built before parent "{node_key}"' + ) + child_handoffs.append(handoff_fn(child_agent)) + + instructions = _build_instructions(node, vs) + tools = _build_node_tools(node, tool_handlers) + agent_name = _sanitize_name(node.key) + agent_name_to_key[agent_name] = node.key + + agent = Agent( + name=agent_name, + model=node.config.get("model", {}).get("name", "gpt-4o"), + **({"instructions": instructions} if instructions else {}), + **({"tools": tools} if tools else {}), + **({"handoffs": child_handoffs} if child_handoffs else {}), + ) + agent_ctx[node.key] = agent + + await _visit(root.key) + + root_agent = agent_ctx.get(root.key) + if root_agent is None: + raise ValueError(f'Root agent "{root.key}" was not built') + + # Lifecycle hooks for LD tracking + class _LDHooks(RunHooks): # type: ignore[misc, valid-type] + async def on_agent_end( + self, context: Any, agent: Any, output: Any + ) -> None: + node_key = agent_name_to_key.get(agent.name) + if node_key and ld_context: + node = def_obj.get_node(node_key) + if node: + td = make_track_data(node, def_obj.key, run_id) + get_client().track( + "$ld:ai:generation:success", ld_context, td, 1 + ) + + async def on_handoff( + self, context: Any, from_agent: Any, to_agent: Any + ) -> None: + from_key = agent_name_to_key.get(from_agent.name) + if from_key and ld_context: + from_node = def_obj.get_node(from_key) + if from_node: + td = make_track_data(from_node, def_obj.key, run_id) + get_client().track( + "$ld:ai:graph:handoff_success", ld_context, td, 1 + ) + + async def on_agent_start(self, context: Any, agent: Any) -> None: + node_key = agent_name_to_key.get(agent.name) + if not node_key or node_key in path: + return + index = len(path) + path.append(node_key) + if not ld_context: + return node = def_obj.get_node(node_key) if node: td = make_track_data(node, def_obj.key, run_id) get_client().track( - "$ld:ai:generation:success", ld_context, td, 1 + "$ld:ai:graph:node", + ld_context, + {**td, "nodeKey": node_key, "index": index}, + 1, ) - async def on_handoff( - self, context: Any, from_agent: Any, to_agent: Any - ) -> None: - from_key = agent_name_to_key.get(from_agent.name) - if from_key and ld_context: - from_node = def_obj.get_node(from_key) - if from_node: - td = make_track_data(from_node, def_obj.key, run_id) - get_client().track( - "$ld:ai:graph:handoff_success", ld_context, td, 1 - ) + hooks = _LDHooks() + + root_prompt: str | list[dict[str, Any]] = input_text + if history: + # config.instructions takes priority over config.messages, so skip + # config conversation turns when instructions are set (parity with the + # single-node handler and TESTING.md §1.11 composition order). + config_messages = ( + [] + if root.config.get("instructions") + else [ + { + **message, + "content": _parse_message_content( + message.get("content", ""), vs + ), + } + for message in (root.config.get("messages") or []) + if message.get("role") != "system" + ] + ) + turns = compose_history( + history=history, + user_input=input_text, + config_messages=config_messages, + ) + root_prompt = _to_openai_agent_items(turns) - async def on_agent_start(self, context: Any, agent: Any) -> None: - node_key = agent_name_to_key.get(agent.name) - if not node_key or node_key in path: - return - index = len(path) - path.append(node_key) - if not ld_context: - return - node = def_obj.get_node(node_key) - if node: - td = make_track_data(node, def_obj.key, run_id) - get_client().track( - "$ld:ai:graph:node", - ld_context, - {**td, "nodeKey": node_key, "index": index}, - 1, - ) - - hooks = _LDHooks() - - root_prompt: str | list[dict[str, Any]] = input_text - if history: - # config.instructions takes priority over config.messages, so skip - # config conversation turns when instructions are set (parity with the - # single-node handler and TESTING.md §1.11 composition order). - config_messages = ( - [] - if root.config.get("instructions") - else [ - { - **message, - "content": _parse_message_content( - message.get("content", ""), vs - ), - } - for message in (root.config.get("messages") or []) - if message.get("role") != "system" - ] + result = await Runner.run(root_agent, root_prompt, hooks=hooks) + + final_output = str(result.final_output or "") + # Sum usage across all raw_responses + input_tokens = sum( + getattr(r.usage, "input_tokens", 0) + for r in result.raw_responses + if hasattr(r, "usage") + ) + output_tokens = sum( + getattr(r.usage, "output_tokens", 0) + for r in result.raw_responses + if hasattr(r, "usage") ) - turns = compose_history( - history=history, - user_input=input_text, - config_messages=config_messages, + total_tokens = sum( + getattr(r.usage, "total_tokens", 0) + for r in result.raw_responses + if hasattr(r, "usage") ) - root_prompt = _to_openai_agent_items(turns) - try: - result = await Runner.run(root_agent, root_prompt, hooks=hooks) + total_usage = { + "input": input_tokens, + "output": output_tokens, + "total": total_tokens, + } + duration = int((time.monotonic() - start_time) * 1000) + + if span: + span.set_attribute("launchdarkly.graph.path", "->".join(path)) + span.set_attribute("gen_ai.usage.input_tokens", input_tokens) + span.set_attribute("gen_ai.usage.output_tokens", output_tokens) + span.set_attribute("gen_ai.usage.total_tokens", total_tokens) + + if ld_context: + graph_td = make_graph_track_data(def_obj.key, run_id) + client = get_client() + 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) + if span: span.set_status(SpanStatusCode.OK) + return {"response": final_output, "usage": total_usage} except Exception as exc: if span: span.record_exception(exc) span.set_status(SpanStatusCode.ERROR, str(exc)) - span.end() if ld_context: - td = make_track_data(root, def_obj.key, run_id) - get_client().track("$ld:ai:graph:invocation_failure", ld_context, td, 1) + get_client().track( + "$ld:ai:graph:invocation_failure", + ld_context, + make_graph_track_data(def_obj.key, run_id), + 1, + ) raise - - final_output = str(result.final_output or "") - # Sum usage across all raw_responses - input_tokens = sum( - getattr(r.usage, "input_tokens", 0) - for r in result.raw_responses - if hasattr(r, "usage") - ) - output_tokens = sum( - getattr(r.usage, "output_tokens", 0) - for r in result.raw_responses - if hasattr(r, "usage") - ) - total_tokens = sum( - getattr(r.usage, "total_tokens", 0) - for r in result.raw_responses - if hasattr(r, "usage") - ) - - total_usage = { - "input": input_tokens, - "output": output_tokens, - "total": total_tokens, - } - duration = int((time.monotonic() - start_time) * 1000) - - if span: - span.set_attribute("launchdarkly.graph.path", "->".join(path)) - span.set_attribute("gen_ai.usage.input_tokens", input_tokens) - span.set_attribute("gen_ai.usage.output_tokens", output_tokens) - span.set_attribute("gen_ai.usage.total_tokens", total_tokens) - span.end() - - if ld_context: - root_td = make_track_data(root, 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) - - return {"response": final_output, "usage": total_usage} + finally: + if span: + span.end() return types.SimpleNamespace(invoke=invoke) diff --git a/packages/openai-agents/tests/test_native_graph.py b/packages/openai-agents/tests/test_native_graph.py index 2df4696f..73b20f3a 100644 --- a/packages/openai-agents/tests/test_native_graph.py +++ b/packages/openai-agents/tests/test_native_graph.py @@ -154,6 +154,55 @@ async def on_agent_start(self, context: Any, agent: Any) -> None: # --------------------------------------------------------------------------- +def _recording_trace() -> tuple[Any, Any, list[Any]]: + """A real tracer behind the module's ``trace`` lookup, with every span's ``end`` spied. + + Returns the stand-in ``trace`` module, the exporter holding finished spans, and the + ``end`` spies, so a test can check both what was exported and how often each span ended. + """ + from opentelemetry.sdk.trace import TracerProvider + from opentelemetry.sdk.trace.export import SimpleSpanProcessor + from opentelemetry.sdk.trace.export.in_memory_span_exporter import ( + InMemorySpanExporter, + ) + + exporter = InMemorySpanExporter() + provider = TracerProvider() + provider.add_span_processor(SimpleSpanProcessor(exporter)) + tracer = provider.get_tracer("test") + end_spies: list[Any] = [] + + def _start_span(name: str, *args: Any, **kwargs: Any) -> Any: + span = tracer.start_span(name, *args, **kwargs) + span.end = MagicMock(wraps=span.end) # type: ignore[method-assign] + end_spies.append(span.end) + return span + + fake_trace = MagicMock() + fake_trace.get_tracer.return_value.start_span.side_effect = _start_span + return fake_trace, exporter, end_spies + + +def _assert_one_failed_graph_span( + exporter: Any, end_spies: list[Any], message: str, track_calls: list[Any] +) -> None: + from opentelemetry.trace import StatusCode + + assert [spy.call_count for spy in end_spies] == [1] + spans = [s for s in exporter.get_finished_spans() if s.name == "launchdarkly.graph"] + assert len(spans) == 1 + span = spans[0] + assert span.status.status_code is StatusCode.ERROR + assert span.status.description == message + assert [e.name for e in span.events].count("exception") == 1 + assert (span.attributes or {}).get("launchdarkly.graph.key") == "test-graph" + + failures = [d for e, d in track_calls if e == "$ld:ai:graph:invocation_failure"] + assert len(failures) == 1 + assert failures[0]["configKey"] == "test-graph" + assert failures[0]["graphKey"] == "test-graph" + + class TestToOpenAIAgentsTopology: @pytest.mark.asyncio async def test_each_graph_node_translated(self) -> None: @@ -573,6 +622,179 @@ 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_agent_setup_error_fails_the_graph_span(self) -> None: + """A setup error before the run ends one ERROR span and tracks the failure.""" + fake_trace, exporter, end_spies = _recording_trace() + track_calls: list[tuple[str, Any]] = [] + mock_ld_client = MagicMock() + mock_ld_client.track = MagicMock( + side_effect=lambda evt, ctx, data, val: track_calls.append((evt, data)) + ) + + 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", fake_trace): + with patch.object(_openai_ng, "_HAS_OTEL", True): + with patch.object( + _openai_ng, "get_client", return_value=mock_ld_client + ): + with pytest.raises(ValueError, match="bad agent"): + await to_openai_agents( + _make_def_promise(graph_def), + opts={"context": {"kind": "user", "key": "u1"}}, + ).invoke("hi") + + agents_mock.Runner.run.assert_not_called() + _assert_one_failed_graph_span(exporter, end_spies, "bad agent", track_calls) + @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."""