From b6d188c92a96482c11a4b3cda4bd790ca1b2ed8f Mon Sep 17 00:00:00 2001 From: jsonbailey Date: Wed, 15 Apr 2026 11:16:32 -0500 Subject: [PATCH 1/3] feat!: Move graph_key to AIConfigTracker instantiation graph_key is now set once at tracker construction time rather than passed as an optional parameter to every tracking method. When agent_graph() builds nodes it uses the internal __evaluate_agent() path so each node tracker is born with graph_key set to the parent graph's key. Public tracking methods no longer accept graph_key. BREAKING CHANGE: The graph_key parameter has been removed from all LDAIConfigTracker tracking methods (track_duration, track_success, track_error, track_tokens, track_tool_call, track_tool_calls, track_feedback, track_eval_scores, track_judge_response, track_time_to_first_token, track_metrics_of, track_metrics_of_async). Pass graph_key to the LDAIConfigTracker constructor instead. Co-Authored-By: Claude Sonnet 4.6 --- packages/sdk/server-ai/src/ldai/client.py | 12 ++- packages/sdk/server-ai/src/ldai/tracker.py | 102 +++++++----------- .../sdk/server-ai/tests/test_agent_graph.py | 13 +++ packages/sdk/server-ai/tests/test_tracker.py | 24 ++--- 4 files changed, 74 insertions(+), 77 deletions(-) diff --git a/packages/sdk/server-ai/src/ldai/client.py b/packages/sdk/server-ai/src/ldai/client.py index 9c87ee8a..190a4f2a 100644 --- a/packages/sdk/server-ai/src/ldai/client.py +++ b/packages/sdk/server-ai/src/ldai/client.py @@ -606,9 +606,10 @@ 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 +749,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 +761,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 +812,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 +840,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 +849,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 index b1894f40..ea2cfc15 100644 --- a/packages/sdk/server-ai/src/ldai/tracker.py +++ b/packages/sdk/server-ai/src/ldai/tracker.py @@ -77,6 +77,7 @@ def __init__( model_name: str, provider_name: str, context: Context, + graph_key: Optional[str] = None, ): """ Initialize an AI Config tracker. @@ -88,6 +89,8 @@ def __init__( :param model_name: Name of the model used. :param provider_name: Name of the provider used. :param context: Context for evaluation. + :param graph_key: When set, include ``graphKey`` in all event payloads + (e.g. config-level metrics inside a graph). """ self._ld_client = ld_client self._variation_key = variation_key @@ -96,13 +99,13 @@ def __init__( self._model_name = model_name self._provider_name = provider_name self._context = context + self._graph_key = graph_key self._summary = LDAIMetricSummary() - def __get_track_data(self, graph_key: Optional[str] = None) -> dict: + def __get_track_data(self) -> dict: """ Get tracking data for events. - :param graph_key: When set, include ``graphKey`` in the payload. :return: Dictionary containing variation and config keys. """ data = { @@ -112,41 +115,36 @@ def __get_track_data(self, graph_key: Optional[str] = None) -> dict: "modelName": self._model_name, "providerName": self._provider_name, } - if graph_key is not None: - data['graphKey'] = graph_key + if self._graph_key is not None: + data['graphKey'] = self._graph_key return data - def track_duration(self, duration: int, *, graph_key: Optional[str] = None) -> None: + def track_duration(self, duration: int) -> None: """ Manually track the duration of an AI operation. :param duration: Duration in milliseconds. - :param graph_key: When set, include ``graphKey`` in the event payload - (e.g. config-level metrics inside a graph). """ self._summary._duration = duration self._ld_client.track( - "$ld:ai:duration:total", self._context, self.__get_track_data(graph_key), duration + "$ld:ai:duration:total", self._context, self.__get_track_data(), duration ) - def track_time_to_first_token( - self, time_to_first_token: int, *, graph_key: Optional[str] = None - ) -> None: + def track_time_to_first_token(self, time_to_first_token: int) -> None: """ Manually track the time to first token of an AI operation. :param time_to_first_token: Time to first token in milliseconds. - :param graph_key: When set, include ``graphKey`` in the event payload. """ self._summary._time_to_first_token = time_to_first_token self._ld_client.track( "$ld:ai:tokens:ttf", self._context, - self.__get_track_data(graph_key), + self.__get_track_data(), time_to_first_token, ) - def track_duration_of(self, func, *, graph_key: Optional[str] = None): + def track_duration_of(self, func): """ Automatically track the duration of an AI operation. @@ -154,7 +152,6 @@ def track_duration_of(self, func, *, graph_key: Optional[str] = None): track the duration. The exception will be re-thrown. :param func: Function to track (synchronous only). - :param graph_key: When set, passed through to :meth:`track_duration`. :return: Result of the tracked function. """ start_ns = time.perf_counter_ns() @@ -162,7 +159,7 @@ def track_duration_of(self, func, *, graph_key: Optional[str] = None): result = func() finally: duration = (time.perf_counter_ns() - start_ns) // 1_000_000 # duration in milliseconds - self.track_duration(duration, graph_key=graph_key) + self.track_duration(duration) return result @@ -170,24 +167,20 @@ def _track_from_metrics_extractor( self, result: Any, metrics_extractor: Callable[[Any], Any], - *, - graph_key: Optional[str] = None, ) -> Any: metrics = metrics_extractor(result) if metrics.success: - self.track_success(graph_key=graph_key) + self.track_success() else: - self.track_error(graph_key=graph_key) + self.track_error() if metrics.usage: - self.track_tokens(metrics.usage, graph_key=graph_key) + self.track_tokens(metrics.usage) return result def track_metrics_of( self, func: Callable[[], Any], metrics_extractor: Callable[[Any], Any], - *, - graph_key: Optional[str] = None, ) -> Any: """ Track metrics for a synchronous AI operation. @@ -203,7 +196,6 @@ def track_metrics_of( :param func: Synchronous callable that runs the operation :param metrics_extractor: Function that extracts LDAIMetrics from the operation result - :param graph_key: When set, include ``graphKey`` on emitted config-level events. :return: The result of the operation """ start_ns = time.perf_counter_ns() @@ -211,17 +203,15 @@ def track_metrics_of( result = func() except Exception as err: duration = (time.perf_counter_ns() - start_ns) // 1_000_000 - self.track_duration(duration, graph_key=graph_key) - self.track_error(graph_key=graph_key) + self.track_duration(duration) + self.track_error() raise err duration = (time.perf_counter_ns() - start_ns) // 1_000_000 - self.track_duration(duration, graph_key=graph_key) - return self._track_from_metrics_extractor(result, metrics_extractor, graph_key=graph_key) + self.track_duration(duration) + return self._track_from_metrics_extractor(result, metrics_extractor) - async def track_metrics_of_async( - self, func, metrics_extractor, *, graph_key: Optional[str] = None - ): + async def track_metrics_of_async(self, func, metrics_extractor): """ Track metrics for an async AI operation (``func`` is awaited). @@ -229,7 +219,6 @@ async def track_metrics_of_async( :param func: Async callable or zero-arg callable that returns an awaitable when called :param metrics_extractor: Function that extracts LDAIMetrics from the operation result - :param graph_key: When set, include ``graphKey`` on emitted config-level events. :return: The result of the operation """ start_ns = time.perf_counter_ns() @@ -238,23 +227,22 @@ async def track_metrics_of_async( result = await func() except Exception as err: duration = (time.perf_counter_ns() - start_ns) // 1_000_000 - self.track_duration(duration, graph_key=graph_key) - self.track_error(graph_key=graph_key) + self.track_duration(duration) + self.track_error() raise err duration = (time.perf_counter_ns() - start_ns) // 1_000_000 - self.track_duration(duration, graph_key=graph_key) - return self._track_from_metrics_extractor(result, metrics_extractor, graph_key=graph_key) + self.track_duration(duration) + return self._track_from_metrics_extractor(result, metrics_extractor) - def track_judge_result(self, judge_result: Any, *, graph_key: Optional[str] = None) -> None: + def track_judge_result(self, judge_result: Any) -> None: """ Track a judge result, including the evaluation score with judge config key. :param judge_result: JudgeResult object containing score, metric key, and success status - :param graph_key: When set, include ``graphKey`` in the event payload. """ if judge_result.success and judge_result.metric_key: - track_data = self.__get_track_data(graph_key=graph_key) + track_data = self.__get_track_data() if judge_result.judge_config_key: track_data = {**track_data, 'judgeConfigKey': judge_result.judge_config_key} self._ld_client.track( @@ -264,49 +252,44 @@ def track_judge_result(self, judge_result: Any, *, graph_key: Optional[str] = No judge_result.score, ) - def track_feedback(self, feedback: Dict[str, FeedbackKind], *, graph_key: Optional[str] = None) -> None: + def track_feedback(self, feedback: Dict[str, FeedbackKind]) -> None: """ Track user feedback for an AI operation. :param feedback: Dictionary containing feedback kind. - :param graph_key: When set, include ``graphKey`` in the event payload. """ self._summary._feedback = feedback if feedback["kind"] == FeedbackKind.Positive: self._ld_client.track( "$ld:ai:feedback:user:positive", self._context, - self.__get_track_data(graph_key=graph_key), + self.__get_track_data(), 1, ) elif feedback["kind"] == FeedbackKind.Negative: self._ld_client.track( "$ld:ai:feedback:user:negative", self._context, - self.__get_track_data(graph_key=graph_key), + self.__get_track_data(), 1, ) - def track_success(self, *, graph_key: Optional[str] = None) -> None: + def track_success(self) -> None: """ Track a successful AI generation. - - :param graph_key: When set, include ``graphKey`` in the event payload. """ self._summary._success = True self._ld_client.track( - "$ld:ai:generation:success", self._context, self.__get_track_data(graph_key=graph_key), 1 + "$ld:ai:generation:success", self._context, self.__get_track_data(), 1 ) - def track_error(self, *, graph_key: Optional[str] = None) -> None: + def track_error(self) -> None: """ Track an unsuccessful AI generation attempt. - - :param graph_key: When set, include ``graphKey`` in the event payload. """ self._summary._success = False self._ld_client.track( - "$ld:ai:generation:error", self._context, self.__get_track_data(graph_key=graph_key), 1 + "$ld:ai:generation:error", self._context, self.__get_track_data(), 1 ) def track_openai_metrics(self, func): @@ -364,15 +347,14 @@ def track_bedrock_converse_metrics(self, res: dict) -> dict: self.track_tokens(_bedrock_to_token_usage(res["usage"])) return res - def track_tokens(self, tokens: TokenUsage, *, graph_key: Optional[str] = None) -> None: + def track_tokens(self, tokens: TokenUsage) -> None: """ Track token usage metrics. :param tokens: Token usage data from either custom, OpenAI, or Bedrock sources. - :param graph_key: When set, include ``graphKey`` in the event payload. """ self._summary._usage = tokens - td = self.__get_track_data(graph_key=graph_key) + td = self.__get_track_data() if tokens.total > 0: self._ld_client.track( "$ld:ai:tokens:total", @@ -395,14 +377,13 @@ def track_tokens(self, tokens: TokenUsage, *, graph_key: Optional[str] = None) - tokens.output, ) - def track_tool_call(self, tool_key: str, *, graph_key: Optional[str] = None) -> None: + def track_tool_call(self, tool_key: str) -> None: """ Track a tool invocation for this configuration (standalone or within a graph). :param tool_key: Identifier of the tool that was invoked. - :param graph_key: When set, include ``graphKey`` in the event payload. """ - track_data = {**self.__get_track_data(graph_key=graph_key), "toolKey": tool_key} + track_data = {**self.__get_track_data(), "toolKey": tool_key} self._ld_client.track( "$ld:ai:tool_call", self._context, @@ -410,17 +391,14 @@ def track_tool_call(self, tool_key: str, *, graph_key: Optional[str] = None) -> 1, ) - def track_tool_calls( - self, tool_keys: Iterable[str], *, graph_key: Optional[str] = None - ) -> None: + def track_tool_calls(self, tool_keys: Iterable[str]) -> None: """ Track multiple tool invocations for this configuration. :param tool_keys: Tool identifiers (e.g. from a model response). - :param graph_key: When set, include ``graphKey`` on each event. """ for tool_key in tool_keys: - self.track_tool_call(tool_key, graph_key=graph_key) + self.track_tool_call(tool_key) def get_summary(self) -> LDAIMetricSummary: """ diff --git a/packages/sdk/server-ai/tests/test_agent_graph.py b/packages/sdk/server-ai/tests/test_agent_graph.py index 130f766b..20d2308b 100644 --- a/packages/sdk/server-ai/tests/test_agent_graph.py +++ b/packages/sdk/server-ai/tests/test_agent_graph.py @@ -387,6 +387,19 @@ def handle_reverse_traverse(node, context): ] +def test_agent_graph_node_trackers_have_graph_key(ldai_client: LDAIClient): + graph = ldai_client.agent_graph("test-agent-graph", Context.create("user-key")) + + assert graph.enabled is True + for node in [graph.get_node("customer-support-agent"), + graph.get_node("personalized-agent"), + graph.get_node("multi-context-agent"), + graph.get_node("minimal-agent")]: + config = node.get_config() + assert config.tracker is not None + assert config.tracker._graph_key == "test-agent-graph" + + def test_agent_graph_handoff(ldai_client: LDAIClient): graph = ldai_client.agent_graph( "test-agent-graph-depth-3", Context.create("user-key") diff --git a/packages/sdk/server-ai/tests/test_tracker.py b/packages/sdk/server-ai/tests/test_tracker.py index 5fea61a8..32e018e9 100644 --- a/packages/sdk/server-ai/tests/test_tracker.py +++ b/packages/sdk/server-ai/tests/test_tracker.py @@ -456,31 +456,31 @@ def _base_td() -> dict: def test_config_tracker_includes_graph_key_when_provided(client: LDClient): context = Context.create("user-key") tracker = LDAIConfigTracker( - client, "variation-key", "config-key", 3, "fakeModel", "fakeProvider", context + client, "variation-key", "config-key", 3, "fakeModel", "fakeProvider", context, graph_key="my-graph" ) expected = {**_base_td(), "graphKey": "my-graph"} - tracker.track_success(graph_key="my-graph") + tracker.track_success() client.track.assert_called_with("$ld:ai:generation:success", context, expected, 1) # type: ignore def test_config_tracker_track_tokens_with_graph_key(client: LDClient): context = Context.create("user-key") tracker = LDAIConfigTracker( - client, "variation-key", "config-key", 3, "fakeModel", "fakeProvider", context + client, "variation-key", "config-key", 3, "fakeModel", "fakeProvider", context, graph_key="g1" ) tokens = TokenUsage(10, 4, 6) expected = {**_base_td(), "graphKey": "g1"} - tracker.track_tokens(tokens, graph_key="g1") + tracker.track_tokens(tokens) client.track.assert_any_call("$ld:ai:tokens:total", context, expected, 10) # type: ignore def test_config_tracker_track_feedback_with_graph_key(client: LDClient): context = Context.create("user-key") tracker = LDAIConfigTracker( - client, "variation-key", "config-key", 3, "fakeModel", "fakeProvider", context + client, "variation-key", "config-key", 3, "fakeModel", "fakeProvider", context, graph_key="gx" ) expected = {**_base_td(), "graphKey": "gx"} - tracker.track_feedback({"kind": FeedbackKind.Positive}, graph_key="gx") + tracker.track_feedback({"kind": FeedbackKind.Positive}) client.track.assert_called_with( "$ld:ai:feedback:user:positive", context, expected, 1 ) # type: ignore @@ -499,19 +499,19 @@ def test_config_tracker_track_tool_call(client: LDClient): def test_config_tracker_track_tool_call_with_graph_key(client: LDClient): context = Context.create("user-key") tracker = LDAIConfigTracker( - client, "variation-key", "config-key", 3, "fakeModel", "fakeProvider", context + client, "variation-key", "config-key", 3, "fakeModel", "fakeProvider", context, graph_key="my-graph" ) expected = {**_base_td(), "graphKey": "my-graph", "toolKey": "calc"} - tracker.track_tool_call("calc", graph_key="my-graph") + tracker.track_tool_call("calc") client.track.assert_called_with("$ld:ai:tool_call", context, expected, 1) # type: ignore def test_config_tracker_track_tool_calls(client: LDClient): context = Context.create("user-key") tracker = LDAIConfigTracker( - client, "variation-key", "config-key", 3, "fakeModel", "fakeProvider", context + client, "variation-key", "config-key", 3, "fakeModel", "fakeProvider", context, graph_key="g" ) - tracker.track_tool_calls(["a", "b"], graph_key="g") + tracker.track_tool_calls(["a", "b"]) assert client.track.call_count == 2 # type: ignore client.track.assert_any_call( "$ld:ai:tool_call", @@ -550,7 +550,7 @@ def extract(r): async def test_config_tracker_track_metrics_of_async_passes_graph_key(client: LDClient): context = Context.create("user-key") tracker = LDAIConfigTracker( - client, "variation-key", "config-key", 3, "fakeModel", "fakeProvider", context + client, "variation-key", "config-key", 3, "fakeModel", "fakeProvider", context, graph_key="gg" ) async def fn(): @@ -559,7 +559,7 @@ async def fn(): def extract(r): return LDAIMetrics(success=True, usage=TokenUsage(5, 2, 3)) - await tracker.track_metrics_of_async(fn, extract, graph_key="gg") + await tracker.track_metrics_of_async(fn, extract) gk_td = {**_base_td(), "graphKey": "gg"} calls = client.track.mock_calls # type: ignore assert any( From 7624cf3e73853e8c2b3af1cd638040c086433e5b Mon Sep 17 00:00:00 2001 From: jsonbailey Date: Wed, 15 Apr 2026 15:53:02 -0500 Subject: [PATCH 2/3] fix: Update provider packages for graph_key-on-init refactor Remove graph_key kwarg from tracker method calls in langchain and openai provider packages. Fix line length in client.py. Update test helpers to pass graph_key at node tracker construction. Co-Authored-By: Claude Sonnet 4.6 --- .../langgraph_callback_handler.py | 9 ++--- .../tests/test_langgraph_callback_handler.py | 40 +++++++++++++++++-- .../tests/test_tracking_langgraph.py | 5 +++ .../ldai_openai/openai_agent_graph_runner.py | 17 ++++---- .../tests/test_tracking_openai_agents.py | 3 ++ packages/sdk/server-ai/src/ldai/client.py | 4 +- 6 files changed, 58 insertions(+), 20 deletions(-) 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 index 026b8d6f..4d495e77 100644 --- 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 @@ -197,7 +197,6 @@ def flush(self, graph: AgentGraphDefinition, graph_tracker: Any) -> None: :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 +207,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 index 79a4a213..6326479c 100644 --- 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, @@ -389,16 +390,47 @@ def test_flush_includes_graph_key_in_node_events(): 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, None) ev = _events(mock_ld_client) token_data = ev['$ld:ai:tokens:total'][0][0] 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 index f8f0649a..5fd98e4e 100644 --- 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, @@ -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 index 9729dbb6..f32abd61 100644 --- 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 @@ -248,13 +248,12 @@ 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, @@ -283,15 +282,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: """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 +299,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 index 931ee59f..39c75034 100644 --- 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 index 190a4f2a..d47e00bf 100644 --- a/packages/sdk/server-ai/src/ldai/client.py +++ b/packages/sdk/server-ai/src/ldai/client.py @@ -608,7 +608,9 @@ def agent_graph( graph_key_value = key agent_configs = { - agent_key: self.__evaluate_agent(agent_key, context, AIAgentConfigDefault(enabled=False), graph_key=graph_key_value) + agent_key: self.__evaluate_agent( + agent_key, context, AIAgentConfigDefault(enabled=False), graph_key=graph_key_value + ) for agent_key in all_agent_keys } From 4549a4722f3ac4c56f01ea35da63a39b1a9c1f23 Mon Sep 17 00:00:00 2001 From: jsonbailey Date: Thu, 16 Apr 2026 09:25:32 -0500 Subject: [PATCH 3/3] chore: remove unused tracker params from openai and langchain provider runners Co-Authored-By: Claude Opus 4.6 --- .../ldai_langchain/langgraph_agent_graph_runner.py | 2 +- .../ldai_langchain/langgraph_callback_handler.py | 3 +-- .../tests/test_langgraph_callback_handler.py | 14 +++++++------- .../tests/test_tracking_langgraph.py | 10 +++++----- .../src/ldai_openai/openai_agent_graph_runner.py | 7 +++---- 5 files changed, 17 insertions(+), 19 deletions(-) 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 index a82b2029..158af79c 100644 --- 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 index 4d495e77..c4b1d8c9 100644 --- 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,14 +188,13 @@ 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). """ for node_key in self._path: node = graph.get_node(node_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 index 6326479c..8b644995 100644 --- 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 @@ -326,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 @@ -345,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 @@ -365,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', []) @@ -383,7 +383,7 @@ 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] @@ -430,7 +430,7 @@ def test_flush_with_no_graph_key_on_node_tracker(): 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) + handler.flush(graph) ev = _events(mock_ld_client) token_data = ev['$ld:ai:tokens:total'][0][0] @@ -445,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 @@ -481,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 index 5fd98e4e..948d58d6 100644 --- a/packages/ai-providers/server-ai-langchain/tests/test_tracking_langgraph.py +++ b/packages/ai-providers/server-ai-langchain/tests/test_tracking_langgraph.py @@ -243,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 @@ -317,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', []) @@ -371,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', [])] @@ -402,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] @@ -487,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) 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 index f32abd61..3cfd595b 100644 --- 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 @@ -258,7 +258,6 @@ def _handle_handoff( def _flush_final_segment( self, state: _RunState, - tracker: Any, result: Any, ) -> None: """Record duration/tokens for the last active agent (no handoff after it).""" @@ -287,7 +286,7 @@ def _flush_final_segment( 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.""" 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)