add multi node tests and sanitize agent names for openai

jsonbailey · jsonbailey · commit 69ba99f82bf8 · 2026-04-01T11:02:03.000-05:00
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
@@ -11,7 +11,7 @@
 from unittest.mock import MagicMock, patch
 
 from ldai.agent_graph import AgentGraphDefinition
-from ldai.models import AIAgentGraphConfig, AIAgentConfig, ModelConfig, ProviderConfig
+from ldai.models import AIAgentGraphConfig, AIAgentConfig, Edge, ModelConfig, ProviderConfig
 from ldai.tracker import AIGraphTracker, LDAIConfigTracker
 from ldai_langchain.langgraph_agent_graph_runner import LangGraphAgentGraphRunner
 
@@ -129,6 +129,74 @@ def _mock_model(response):
     return model
 
 
+def _make_two_node_graph(mock_ld_client: MagicMock) -> 'AgentGraphDefinition':
+    """Build a two-node AgentGraphDefinition (root-agent → child-agent)."""
+    context = MagicMock()
+
+    root_tracker = LDAIConfigTracker(
+        ld_client=mock_ld_client,
+        variation_key='test-variation',
+        config_key='root-agent',
+        version=1,
+        model_name='gpt-4',
+        provider_name='openai',
+        context=context,
+    )
+    child_tracker = LDAIConfigTracker(
+        ld_client=mock_ld_client,
+        variation_key='test-variation',
+        config_key='child-agent',
+        version=1,
+        model_name='gpt-4',
+        provider_name='openai',
+        context=context,
+    )
+    graph_tracker = AIGraphTracker(
+        ld_client=mock_ld_client,
+        variation_key='test-variation',
+        graph_key='two-node-graph',
+        version=1,
+        context=context,
+    )
+
+    root_config = AIAgentConfig(
+        key='root-agent',
+        enabled=True,
+        model=ModelConfig(name='gpt-4', parameters={}),
+        provider=ProviderConfig(name='openai'),
+        instructions='You are root.',
+        tracker=root_tracker,
+    )
+    child_config = AIAgentConfig(
+        key='child-agent',
+        enabled=True,
+        model=ModelConfig(name='gpt-4', parameters={}),
+        provider=ProviderConfig(name='openai'),
+        instructions='You are child.',
+        tracker=child_tracker,
+    )
+
+    edge = Edge(key='root-to-child', source_config='root-agent', target_config='child-agent')
+    graph_config = AIAgentGraphConfig(
+        key='two-node-graph',
+        root_config_key='root-agent',
+        edges=[edge],
+        enabled=True,
+    )
+
+    nodes = AgentGraphDefinition.build_nodes(graph_config, {
+        'root-agent': root_config,
+        'child-agent': child_config,
+    })
+    return AgentGraphDefinition(
+        agent_graph=graph_config,
+        nodes=nodes,
+        context=context,
+        enabled=True,
+        tracker=graph_tracker,
+    )
+
+
 # ---------------------------------------------------------------------------
 # Tests
 # ---------------------------------------------------------------------------
@@ -258,3 +326,41 @@ async def test_tracks_failure_and_latency_on_model_error():
     assert '$ld:ai:graph:invocation_failure' in ev
     assert '$ld:ai:graph:latency' in ev
     assert '$ld:ai:graph:invocation_success' not in ev
+
+
+@pytest.mark.asyncio
+async def test_multi_node_tracks_per_node_tokens_and_path():
+    """Each node emits its own token events; path and graph total cover both nodes."""
+    mock_ld_client = MagicMock()
+    graph = _make_two_node_graph(mock_ld_client)
+
+    root_response = _make_fake_response('Root done.', input_tokens=10, output_tokens=5)
+    child_response = _make_fake_response('Child done.', input_tokens=3, output_tokens=2)
+
+    def model_factory(node_config):
+        if node_config.key == 'root-agent':
+            return _mock_model(root_response)
+        return _mock_model(child_response)
+
+    with patch('ldai_langchain.langgraph_agent_graph_runner.create_langchain_model',
+               side_effect=model_factory):
+        runner = LangGraphAgentGraphRunner(graph, {})
+        result = await runner.run('hello')
+
+    assert result.metrics.success is True
+
+    ev = _events(mock_ld_client)
+
+    # Per-node token events identified by configKey
+    root_tokens = [(d, v) for d, v in ev.get('$ld:ai:tokens:total', []) if d.get('configKey') == 'root-agent']
+    child_tokens = [(d, v) for d, v in ev.get('$ld:ai:tokens:total', []) if d.get('configKey') == 'child-agent']
+    assert root_tokens[0][1] == 15
+    assert child_tokens[0][1] == 5
+
+    # Graph-level total accumulates both nodes (10+3 in, 5+2 out)
+    assert ev['$ld:ai:graph:total_tokens'][0][1] == 20
+
+    # Execution path includes both node keys
+    path_data = ev['$ld:ai:graph:path'][0][0]
+    assert 'root-agent' in path_data['path']
+    assert 'child-agent' in path_data['path']
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
@@ -1,7 +1,8 @@
 """OpenAI agent graph runner for LaunchDarkly AI SDK."""
 
+import re
 import time
-from typing import Any, List, Optional
+from typing import Any, Dict, List, Optional
 
 from ldai import log
 from ldai.agent_graph import AgentGraphDefinition, AgentGraphNode
@@ -16,6 +17,11 @@
 )
 
 
+def _sanitize_agent_name(key: str) -> str:
+    """Replace characters invalid for OpenAI function names with underscores."""
+    return re.sub(r'[^a-zA-Z0-9_]', '_', key)
+
+
 class _RunState:
     """Mutable state shared across handoff and tool callbacks during a single run."""
 
@@ -44,6 +50,7 @@ def __init__(self, graph: AgentGraphDefinition, tools: ToolRegistry):
         """
         self._graph = graph
         self._tools = tools
+        self._agent_name_map: Dict[str, str] = {}
 
     async def run(self, input: Any) -> AgentGraphResult:
         """
@@ -132,6 +139,7 @@ def _build_agents(self, path: List[str], state: _RunState) -> Any:
             ) from exc
 
         tracker = self._graph.get_tracker()
+        name_map: Dict[str, str] = {}
 
         def build_node(node: AgentGraphNode, ctx: dict) -> Any:
             node_config = node.get_config()
@@ -142,6 +150,8 @@ def build_node(node: AgentGraphNode, ctx: dict) -> Any:
                 raise ValueError(f"Model not set for node '{node_config.key}'")
 
             tool_defs = model.get_parameter('tools') or []
+            sanitized_name = _sanitize_agent_name(node_config.key)
+            name_map[sanitized_name] = node_config.key
 
             # --- handoffs ---
             agent_handoffs: List[Handoff] = []
@@ -173,14 +183,16 @@ def build_node(node: AgentGraphNode, ctx: dict) -> Any:
                 agent_tools.append(function_tool(tool_fn))
 
             return Agent(
-                name=node_config.key,
+                name=sanitized_name,
                 model=model.name,
                 instructions=f'{RECOMMENDED_PROMPT_PREFIX} {node_config.instructions or ""}',
                 handoffs=list(agent_handoffs),
                 tools=list(agent_tools),
             )
 
-        return self._graph.reverse_traverse(fn=build_node)
+        root = self._graph.reverse_traverse(fn=build_node)
+        self._agent_name_map = name_map
+        return root
 
     def _make_on_handoff(
         self,
@@ -269,7 +281,8 @@ 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_name in get_tool_calls_from_run_items(result.new_items):
-            node = self._graph.get_node(agent_name)
+            original_key = self._agent_name_map.get(agent_name, agent_name)
+            node = self._graph.get_node(original_key)
             if node is None:
                 continue
             config_tracker = node.get_config().tracker
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
@@ -11,7 +11,7 @@
 from unittest.mock import AsyncMock, MagicMock, patch
 
 from ldai.agent_graph import AgentGraphDefinition
-from ldai.models import AIAgentGraphConfig, AIAgentConfig, ModelConfig, ProviderConfig
+from ldai.models import AIAgentGraphConfig, AIAgentConfig, Edge, ModelConfig, ProviderConfig
 from ldai.tracker import AIGraphTracker, LDAIConfigTracker
 from ldai_openai.openai_agent_graph_runner import OpenAIAgentGraphRunner
 
@@ -153,6 +153,74 @@ def _make_agents_modules(run_result: MagicMock) -> dict:
     }
 
 
+def _make_two_node_graph(mock_ld_client: MagicMock) -> AgentGraphDefinition:
+    """Build a two-node AgentGraphDefinition (root-agent → child-agent)."""
+    context = MagicMock()
+
+    root_tracker = LDAIConfigTracker(
+        ld_client=mock_ld_client,
+        variation_key='test-variation',
+        config_key='root-agent',
+        version=1,
+        model_name='gpt-4',
+        provider_name='openai',
+        context=context,
+    )
+    child_tracker = LDAIConfigTracker(
+        ld_client=mock_ld_client,
+        variation_key='test-variation',
+        config_key='child-agent',
+        version=1,
+        model_name='gpt-4',
+        provider_name='openai',
+        context=context,
+    )
+    graph_tracker = AIGraphTracker(
+        ld_client=mock_ld_client,
+        variation_key='test-variation',
+        graph_key='two-node-graph',
+        version=1,
+        context=context,
+    )
+
+    root_config = AIAgentConfig(
+        key='root-agent',
+        enabled=True,
+        model=ModelConfig(name='gpt-4', parameters={}),
+        provider=ProviderConfig(name='openai'),
+        instructions='You are root.',
+        tracker=root_tracker,
+    )
+    child_config = AIAgentConfig(
+        key='child-agent',
+        enabled=True,
+        model=ModelConfig(name='gpt-4', parameters={}),
+        provider=ProviderConfig(name='openai'),
+        instructions='You are child.',
+        tracker=child_tracker,
+    )
+
+    edge = Edge(key='root-to-child', source_config='root-agent', target_config='child-agent')
+    graph_config = AIAgentGraphConfig(
+        key='two-node-graph',
+        root_config_key='root-agent',
+        edges=[edge],
+        enabled=True,
+    )
+
+    nodes = AgentGraphDefinition.build_nodes(graph_config, {
+        'root-agent': root_config,
+        'child-agent': child_config,
+    })
+    return AgentGraphDefinition(
+        agent_graph=graph_config,
+        nodes=nodes,
+        context=context,
+        enabled=True,
+        tracker=graph_tracker,
+    )
+
+
 def _events(mock_ld_client: MagicMock) -> dict:
     """Return dict of event_name -> list of (data, value) from all track() calls."""
     result = defaultdict(list)
@@ -303,3 +371,89 @@ async def test_tracks_failure_and_latency_on_runner_error():
     assert '$ld:ai:graph:invocation_failure' in ev
     assert '$ld:ai:graph:latency' in ev
     assert '$ld:ai:graph:invocation_success' not in ev
+
+
+@pytest.mark.asyncio
+async def test_multi_node_tracks_per_node_tokens_and_handoff():
+    """Each node emits its own token events; handoff event fires between them."""
+    mock_ld_client = MagicMock()
+    graph = _make_two_node_graph(mock_ld_client)
+
+    root_entry = MagicMock()
+    root_entry.total_tokens = 15
+    root_entry.input_tokens = 10
+    root_entry.output_tokens = 5
+
+    child_entry = MagicMock()
+    child_entry.total_tokens = 9
+    child_entry.input_tokens = 6
+    child_entry.output_tokens = 3
+
+    run_result = MagicMock()
+    run_result.final_output = 'child answer'
+    run_result.new_items = []
+    run_result.usage = None
+    run_result.context_wrapper.usage.total_tokens = 24
+    run_result.context_wrapper.usage.input_tokens = 16
+    run_result.context_wrapper.usage.output_tokens = 8
+    run_result.context_wrapper.usage.request_usage_entries = [root_entry, child_entry]
+
+    on_handoff_callbacks = []
+
+    def capture_handoff(**kwargs):
+        cb = kwargs.get('on_handoff')
+        if cb:
+            on_handoff_callbacks.append(cb)
+        return MagicMock()
+
+    async def mock_run(agent, input_str, **kwargs):
+        # Simulate the root→child handoff before returning
+        if on_handoff_callbacks:
+            run_ctx = MagicMock()
+            run_ctx.usage.request_usage_entries = [root_entry]
+            on_handoff_callbacks[0](run_ctx)
+        return run_result
+
+    mock_runner_cls = MagicMock()
+    mock_runner_cls.run = mock_run
+
+    mock_agents = MagicMock()
+    mock_agents.Runner = mock_runner_cls
+    mock_agents.Agent = MagicMock(return_value=MagicMock())
+    mock_agents.Handoff = MagicMock()
+    mock_agents.Tool = MagicMock()
+    mock_agents.function_tool = lambda fn: MagicMock()
+    mock_agents.handoff = capture_handoff
+
+    mock_ext = MagicMock()
+    mock_ext.RECOMMENDED_PROMPT_PREFIX = '[PREFIX]'
+
+    with patch.dict('sys.modules', {
+        'agents': mock_agents,
+        'agents.extensions': MagicMock(),
+        'agents.extensions.handoff_prompt': mock_ext,
+        'agents.tool_context': MagicMock(),
+    }):
+        runner = OpenAIAgentGraphRunner(graph, {})
+        result = await runner.run('hello')
+
+    assert result.metrics.success is True
+
+    ev = _events(mock_ld_client)
+
+    # Per-node token events identified by configKey
+    root_tokens = [(d, v) for d, v in ev.get('$ld:ai:tokens:total', []) if d.get('configKey') == 'root-agent']
+    child_tokens = [(d, v) for d, v in ev.get('$ld:ai:tokens:total', []) if d.get('configKey') == 'child-agent']
+    assert root_tokens[0][1] == 15
+    assert child_tokens[0][1] == 9
+
+    # Execution path includes both node keys
+    path_data = ev['$ld:ai:graph:path'][0][0]
+    assert 'root-agent' in path_data['path']
+    assert 'child-agent' in path_data['path']
+
+    # Handoff event fires with correct source and target
+    handoff_events = ev.get('$ld:ai:graph:handoff_success', [])
+    assert len(handoff_events) == 1
+    assert handoff_events[0][0]['sourceKey'] == 'root-agent'
+    assert handoff_events[0][0]['targetKey'] == 'child-agent'
