feat!: Move graph_key to AIConfigTracker instantiation (#134)

jsonbailey · web-flow · commit 20fff24fcd02 · 2026-04-16T10:50:22.000-05:00
diff --git a/packages/ai-providers/server-ai-langchain/src/ldai_langchain/langgraph_agent_graph_runner.py b/packages/ai-providers/server-ai-langchain/src/ldai_langchain/langgraph_agent_graph_runner.py
@@ -300,7 +300,7 @@ async def run(self, input: Any) -> AgentGraphResult:
             output = extract_last_message_content(messages)
 
             # Flush per-node metrics to LD trackers
-            handler.flush(self._graph, tracker)
+            handler.flush(self._graph)
 
             # Graph-level metrics
             if tracker:
diff --git a/packages/ai-providers/server-ai-langchain/src/ldai_langchain/langgraph_callback_handler.py b/packages/ai-providers/server-ai-langchain/src/ldai_langchain/langgraph_callback_handler.py
@@ -188,16 +188,14 @@ def on_tool_end(
     # Flush
     # ------------------------------------------------------------------
 
-    def flush(self, graph: AgentGraphDefinition, graph_tracker: Any) -> None:
+    def flush(self, graph: AgentGraphDefinition) -> None:
         """
         Emit all collected per-node metrics to the LaunchDarkly trackers.
 
         Call this once after the graph run completes.
 
         :param graph: The AgentGraphDefinition whose nodes hold the LD config trackers.
-        :param graph_tracker: The AIGraphTracker for the overall graph (may be None).
         """
-        gk = graph_tracker.graph_key if graph_tracker is not None else None
         for node_key in self._path:
             node = graph.get_node(node_key)
             if not node:
@@ -208,13 +206,13 @@ def flush(self, graph: AgentGraphDefinition, graph_tracker: Any) -> None:
 
             usage = self._node_tokens.get(node_key)
             if usage:
-                config_tracker.track_tokens(usage, graph_key=gk)
+                config_tracker.track_tokens(usage)
 
             duration = self._node_duration_ms.get(node_key)
             if duration is not None:
-                config_tracker.track_duration(duration, graph_key=gk)
+                config_tracker.track_duration(duration)
 
-            config_tracker.track_success(graph_key=gk)
+            config_tracker.track_success()
 
             for tool_key in self._node_tool_calls.get(node_key, []):
-                config_tracker.track_tool_call(tool_key, graph_key=gk)
+                config_tracker.track_tool_call(tool_key)
diff --git a/packages/ai-providers/server-ai-langchain/tests/test_langgraph_callback_handler.py b/packages/ai-providers/server-ai-langchain/tests/test_langgraph_callback_handler.py
@@ -35,6 +35,7 @@ def _make_graph(mock_ld_client: MagicMock, node_key: str = 'root-agent', graph_k
         model_name='gpt-4',
         provider_name='openai',
         context=context,
+        graph_key=graph_key,
     )
     graph_tracker = AIGraphTracker(
         ld_client=mock_ld_client,
@@ -325,7 +326,7 @@ def test_flush_emits_token_events_to_ld_tracker():
     node_run_id = uuid4()
     handler.on_chain_start({}, {}, run_id=node_run_id, name='root-agent')
     handler.on_llm_end(_llm_result(15, 10, 5), run_id=uuid4(), parent_run_id=node_run_id)
-    handler.flush(graph, tracker)
+    handler.flush(graph)
 
     ev = _events(mock_ld_client)
     assert ev['$ld:ai:tokens:total'][0][1] == 15
@@ -344,7 +345,7 @@ def test_flush_emits_duration():
     run_id = uuid4()
     handler.on_chain_start({}, {}, run_id=run_id, name='root-agent')
     handler.on_chain_end({}, run_id=run_id)
-    handler.flush(graph, tracker)
+    handler.flush(graph)
 
     ev = _events(mock_ld_client)
     assert '$ld:ai:duration:total' in ev
@@ -364,7 +365,7 @@ def test_flush_emits_tool_calls():
     tools_run_id = uuid4()
     handler.on_chain_start({}, {}, run_id=tools_run_id, name='root-agent__tools')
     handler.on_tool_end('r', run_id=uuid4(), parent_run_id=tools_run_id, name='fn_search')
-    handler.flush(graph, tracker)
+    handler.flush(graph)
 
     ev = _events(mock_ld_client)
     tool_events = ev.get('$ld:ai:tool_call', [])
@@ -382,23 +383,54 @@ def test_flush_includes_graph_key_in_node_events():
     node_run_id = uuid4()
     handler.on_chain_start({}, {}, run_id=node_run_id, name='root-agent')
     handler.on_llm_end(_llm_result(5, 3, 2), run_id=uuid4(), parent_run_id=node_run_id)
-    handler.flush(graph, tracker)
+    handler.flush(graph)
 
     ev = _events(mock_ld_client)
     token_data = ev['$ld:ai:tokens:total'][0][0]
     assert token_data.get('graphKey') == 'my-graph'
 
 
-def test_flush_with_none_tracker_uses_no_graph_key():
-    """flush() with graph_tracker=None does not fail and omits graphKey."""
+def test_flush_with_no_graph_key_on_node_tracker():
+    """When node tracker has no graph_key, events omit graphKey."""
     mock_ld_client = MagicMock()
-    graph = _make_graph(mock_ld_client)
+    context = MagicMock()
+    node_tracker = LDAIConfigTracker(
+        ld_client=mock_ld_client,
+        variation_key='v1',
+        config_key='root-agent',
+        version=1,
+        model_name='gpt-4',
+        provider_name='openai',
+        context=context,
+    )
+    node_config = AIAgentConfig(
+        key='root-agent',
+        enabled=True,
+        model=ModelConfig(name='gpt-4', parameters={}),
+        provider=ProviderConfig(name='openai'),
+        instructions='Be helpful.',
+        tracker=node_tracker,
+    )
+    graph_config = AIAgentGraphConfig(
+        key='test-graph',
+        root_config_key='root-agent',
+        edges=[],
+        enabled=True,
+    )
+    nodes = AgentGraphDefinition.build_nodes(graph_config, {'root-agent': node_config})
+    graph = AgentGraphDefinition(
+        agent_graph=graph_config,
+        nodes=nodes,
+        context=context,
+        enabled=True,
+        tracker=None,
+    )
 
     handler = LDMetricsCallbackHandler({'root-agent'}, {})
     node_run_id = uuid4()
     handler.on_chain_start({}, {}, run_id=node_run_id, name='root-agent')
     handler.on_llm_end(_llm_result(5, 3, 2), run_id=uuid4(), parent_run_id=node_run_id)
-    handler.flush(graph, None)  # graph_tracker=None
+    handler.flush(graph)
 
     ev = _events(mock_ld_client)
     token_data = ev['$ld:ai:tokens:total'][0][0]
@@ -413,7 +445,7 @@ def test_flush_skips_nodes_not_in_path():
 
     # Handler with 'root-agent' in node_keys but never started
     handler = LDMetricsCallbackHandler({'root-agent'}, {})
-    handler.flush(graph, tracker)
+    handler.flush(graph)
 
     ev = _events(mock_ld_client)
     assert '$ld:ai:tokens:total' not in ev
@@ -449,7 +481,7 @@ def test_flush_skips_node_without_tracker():
     node_run_id = uuid4()
     handler.on_chain_start({}, {}, run_id=node_run_id, name='no-track')
     handler.on_llm_end(_llm_result(5, 3, 2), run_id=uuid4(), parent_run_id=node_run_id)
-    handler.flush(graph, None)  # should not raise
+    handler.flush(graph)  # should not raise
 
     mock_ld_client.track.assert_not_called()
 
diff --git a/packages/ai-providers/server-ai-langchain/tests/test_tracking_langgraph.py b/packages/ai-providers/server-ai-langchain/tests/test_tracking_langgraph.py
@@ -45,6 +45,7 @@ def _make_graph(
         model_name='gpt-4',
         provider_name='openai',
         context=context,
+        graph_key=graph_key,
     )
 
     graph_tracker = AIGraphTracker(
@@ -141,6 +142,7 @@ def _make_two_node_graph(mock_ld_client: MagicMock) -> 'AgentGraphDefinition':
         model_name='gpt-4',
         provider_name='openai',
         context=context,
+        graph_key='two-node-graph',
     )
     child_tracker = LDAIConfigTracker(
         ld_client=mock_ld_client,
@@ -150,6 +152,7 @@ def _make_two_node_graph(mock_ld_client: MagicMock) -> 'AgentGraphDefinition':
         model_name='gpt-4',
         provider_name='openai',
         context=context,
+        graph_key='two-node-graph',
     )
     graph_tracker = AIGraphTracker(
         ld_client=mock_ld_client,
@@ -240,7 +243,7 @@ async def test_tracks_node_and_graph_tokens_on_success():
     )
     handler.on_llm_end(llm_result, run_id=uuid4(), parent_run_id=node_run_id)
     handler.on_chain_end({}, run_id=node_run_id)
-    handler.flush(graph2, tracker2)
+    handler.flush(graph2)
 
     ev2 = _events(mock_ld_client2)
     assert ev2['$ld:ai:tokens:total'][0][1] == 15
@@ -314,7 +317,7 @@ def get_weather(location: str = 'NYC') -> str:
     tools_run_id = uuid4()
     handler.on_chain_start({}, {}, run_id=tools_run_id, name='root-agent__tools')
     handler.on_tool_end('sunny', run_id=uuid4(), parent_run_id=tools_run_id, name='get_weather')
-    handler.flush(graph2, tracker2)
+    handler.flush(graph2)
 
     ev2 = _events(mock_ld_client2)
     tool_events = ev2.get('$ld:ai:tool_call', [])
@@ -368,7 +371,7 @@ def summarize(text: str = '') -> str:
     handler.on_chain_start({}, {}, run_id=tools_run_id, name='root-agent__tools')
     handler.on_tool_end('result', run_id=uuid4(), parent_run_id=tools_run_id, name='search')
     handler.on_tool_end('summary', run_id=uuid4(), parent_run_id=tools_run_id, name='summarize')
-    handler.flush(graph2, tracker2)
+    handler.flush(graph2)
 
     ev2 = _events(mock_ld_client2)
     tool_keys = [data['toolKey'] for data, _ in ev2.get('$ld:ai:tool_call', [])]
@@ -399,7 +402,7 @@ async def test_tracks_graph_key_on_node_events():
         llm_output={},
     )
     handler.on_llm_end(llm_result, run_id=uuid4(), parent_run_id=node_run_id)
-    handler.flush(graph, tracker)
+    handler.flush(graph)
 
     ev = _events(mock_ld_client)
     token_data = ev['$ld:ai:tokens:total'][0][0]
@@ -484,7 +487,7 @@ def model_factory(node_config, **kwargs):
     )
     handler.on_llm_end(child_llm_result, run_id=uuid4(), parent_run_id=child_run_id)
 
-    handler.flush(graph2, tracker2)
+    handler.flush(graph2)
 
     ev2 = _events(mock_ld_client2)
 
@@ -517,6 +520,7 @@ def _node_tracker(key: str) -> LDAIConfigTracker:
             model_name='gpt-4',
             provider_name='openai',
             context=context,
+            graph_key='multi-child-graph',
         )
 
     graph_tracker = AIGraphTracker(
@@ -627,6 +631,7 @@ def _node_tracker(key: str) -> LDAIConfigTracker:
             model_name='gpt-4',
             provider_name='openai',
             context=context,
+            graph_key='multi-child-tools-graph',
         )
 
     graph_tracker = AIGraphTracker(
diff --git a/packages/ai-providers/server-ai-openai/src/ldai_openai/openai_agent_graph_runner.py b/packages/ai-providers/server-ai-openai/src/ldai_openai/openai_agent_graph_runner.py
@@ -82,8 +82,8 @@ async def run(self, input: Any) -> AgentGraphResult:
             from agents import Runner
             root_agent = self._build_agents(path, state)
             result = await Runner.run(root_agent, str(input))
-            self._flush_final_segment(state, tracker, result)
-            self._track_tool_calls(result, tracker)
+            self._flush_final_segment(state, result)
+            self._track_tool_calls(result)
 
             duration = (time.perf_counter_ns() - start_ns) // 1_000_000
 
@@ -248,18 +248,16 @@ def _handle_handoff(
         except Exception:
             pass
 
-        gk = tracker.graph_key if tracker is not None else None
         if config_tracker is not None:
             if usage is not None:
-                config_tracker.track_tokens(usage, graph_key=gk)
+                config_tracker.track_tokens(usage)
             if duration_ms is not None:
-                config_tracker.track_duration(int(duration_ms), graph_key=gk)
-            config_tracker.track_success(graph_key=gk)
+                config_tracker.track_duration(int(duration_ms))
+            config_tracker.track_success()
 
     def _flush_final_segment(
         self,
         state: _RunState,
-        tracker: Any,
         result: Any,
     ) -> None:
         """Record duration/tokens for the last active agent (no handoff after it)."""
@@ -283,15 +281,13 @@ def _flush_final_segment(
         except Exception:
             pass
 
-        gk = tracker.graph_key if tracker is not None else None
         if usage is not None:
-            config_tracker.track_tokens(usage, graph_key=gk)
-        config_tracker.track_duration(int(duration_ms), graph_key=gk)
-        config_tracker.track_success(graph_key=gk)
+            config_tracker.track_tokens(usage)
+        config_tracker.track_duration(int(duration_ms))
+        config_tracker.track_success()
 
-    def _track_tool_calls(self, result: Any, tracker: Any) -> None:
+    def _track_tool_calls(self, result: Any) -> None:
         """Track all tool calls from the run result, attributed to the node that called them."""
-        gk = tracker.graph_key if tracker is not None else None
         for agent_name, tool_fn_name in get_tool_calls_from_run_items(result.new_items):
             agent_key = self._agent_name_map.get(agent_name, agent_name)
             tool_name = self._tool_name_map.get(tool_fn_name)
@@ -302,4 +298,4 @@ def _track_tool_calls(self, result: Any, tracker: Any) -> None:
                 continue
             config_tracker = node.get_config().tracker
             if config_tracker is not None:
-                config_tracker.track_tool_call(tool_name, graph_key=gk)
+                config_tracker.track_tool_call(tool_name)
diff --git a/packages/ai-providers/server-ai-openai/tests/test_tracking_openai_agents.py b/packages/ai-providers/server-ai-openai/tests/test_tracking_openai_agents.py
@@ -40,6 +40,7 @@ def _make_graph(
         model_name='gpt-4',
         provider_name='openai',
         context=context,
+        graph_key=graph_key,
     )
 
     graph_tracker = AIGraphTracker(
@@ -178,6 +179,7 @@ def _make_two_node_graph(mock_ld_client: MagicMock) -> AgentGraphDefinition:
         model_name='gpt-4',
         provider_name='openai',
         context=context,
+        graph_key='two-node-graph',
     )
     child_tracker = LDAIConfigTracker(
         ld_client=mock_ld_client,
@@ -187,6 +189,7 @@ def _make_two_node_graph(mock_ld_client: MagicMock) -> AgentGraphDefinition:
         model_name='gpt-4',
         provider_name='openai',
         context=context,
+        graph_key='two-node-graph',
     )
     graph_tracker = AIGraphTracker(
         ld_client=mock_ld_client,
diff --git a/packages/sdk/server-ai/src/ldai/client.py b/packages/sdk/server-ai/src/ldai/client.py
@@ -606,9 +606,12 @@ def agent_graph(
             for single_edge in variation.get("edges", {}).get(edge_key, []):
                 all_agent_keys.add(single_edge.get("key", ""))
 
+        graph_key_value = key
         agent_configs = {
-            key: self.agent_config(key, context, AIAgentConfigDefault(enabled=False))
-            for key in all_agent_keys
+            agent_key: self.__evaluate_agent(
+                agent_key, context, AIAgentConfigDefault(enabled=False), graph_key=graph_key_value
+            )
+            for agent_key in all_agent_keys
         }
 
         if not all(config.enabled for config in agent_configs.values()):
@@ -748,6 +751,7 @@ def __evaluate(
         context: Context,
         default_dict: Dict[str, Any],
         variables: Optional[Dict[str, Any]] = None,
+        graph_key: Optional[str] = None,
     ) -> Tuple[
         Optional[ModelConfig], Optional[ProviderConfig], Optional[List[LDMessage]],
         Optional[str], LDAIConfigTracker, bool, Optional[Any], Dict[str, Any]
@@ -759,6 +763,7 @@ def __evaluate(
         :param context: The evaluation context.
         :param default_dict: Default configuration as dictionary.
         :param variables: Variables for interpolation.
+        :param graph_key: When set, passed to the tracker so all events include ``graphKey``.
         :return: Tuple of (model, provider, messages, instructions, tracker, enabled, judge_configuration, variation).
         """
         variation = self._client.variation(key, context, default_dict)
@@ -809,6 +814,7 @@ def __evaluate(
             model.name if model else '',
             provider_config.name if provider_config else '',
             context,
+            graph_key=graph_key,
         )
 
         enabled = variation.get('_ldMeta', {}).get('enabled', False)
@@ -836,6 +842,7 @@ def __evaluate_agent(
         context: Context,
         default: AIAgentConfigDefault,
         variables: Optional[Dict[str, Any]] = None,
+        graph_key: Optional[str] = None,
     ) -> AIAgentConfig:
         """
         Internal method to evaluate an agent configuration.
@@ -844,10 +851,11 @@ def __evaluate_agent(
         :param context: The evaluation context.
         :param default: Default agent values.
         :param variables: Variables for interpolation.
+        :param graph_key: When set, passed to the tracker so all events include ``graphKey``.
         :return: Configured AIAgentConfig instance.
         """
         model, provider, messages, instructions, tracker, enabled, judge_configuration, _ = self.__evaluate(
-            key, context, default.to_dict(), variables
+            key, context, default.to_dict(), variables, graph_key=graph_key
         )
 
         # For agents, prioritize instructions over messages
diff --git a/packages/sdk/server-ai/src/ldai/tracker.py b/packages/sdk/server-ai/src/ldai/tracker.py
diff --git a/packages/sdk/server-ai/tests/test_agent_graph.py b/packages/sdk/server-ai/tests/test_agent_graph.py
diff --git a/packages/sdk/server-ai/tests/test_tracker.py b/packages/sdk/server-ai/tests/test_tracker.py
