From 8c0f238ec6bb0f0c63173f40937edb6964cd28cb Mon Sep 17 00:00:00 2001 From: jsonbailey Date: Wed, 11 Mar 2026 14:42:53 -0500 Subject: [PATCH 01/10] feat: add ManagedAgentGraph, OpenAIAgentGraphRunner, LangGraphAgentGraphRunner MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Implements PR 5 — ManagedAgentGraph + create_agent_graph(): ldai: - managed_agent_graph.py: ManagedAgentGraph wrapper holding AgentGraphRunner + AIGraphTracker; exposes run(), get_agent_graph_runner(), get_tracker() - LDAIClient.create_agent_graph(key, context, tools): resolves graph via agent_graph(), delegates to RunnerFactory, returns ManagedAgentGraph - Exports ManagedAgentGraph from top-level ldai package ldai_openai: - OpenAIAgentGraphRunner(AgentGraphRunner): builds agents via reverse_traverse using the openai-agents SDK; auto-tracks path, tool calls, handoffs, latency, invocation success/failure - OpenAIRunnerFactory.create_agent_graph(graph_def, tools) -> OpenAIAgentGraphRunner ldai_langchain: - LangGraphAgentGraphRunner(AgentGraphRunner): builds a LangGraph StateGraph via traverse(); auto-tracks latency and invocation success/failure - LangChainRunnerFactory.create_agent_graph(graph_def, tools) -> LangGraphAgentGraphRunner Co-Authored-By: Claude Sonnet 4.6 --- .../src/ldai_langchain/__init__.py | 2 + .../langchain_runner_factory.py | 15 +- .../langgraph_agent_graph_runner.py | 138 ++++++++++++ .../test_langgraph_agent_graph_runner.py | 139 ++++++++++++ .../src/ldai_openai/__init__.py | 2 + .../ldai_openai/openai_agent_graph_runner.py | 210 ++++++++++++++++++ .../src/ldai_openai/openai_runner_factory.py | 14 +- .../tests/test_openai_agent_graph_runner.py | 134 +++++++++++ packages/sdk/server-ai/src/ldai/__init__.py | 2 + packages/sdk/server-ai/src/ldai/client.py | 52 +++++ .../server-ai/src/ldai/managed_agent_graph.py | 61 +++++ .../tests/test_managed_agent_graph.py | 196 ++++++++++++++++ 12 files changed, 963 insertions(+), 2 deletions(-) create mode 100644 packages/ai-providers/server-ai-langchain/src/ldai_langchain/langgraph_agent_graph_runner.py create mode 100644 packages/ai-providers/server-ai-langchain/tests/test_langgraph_agent_graph_runner.py create mode 100644 packages/ai-providers/server-ai-openai/src/ldai_openai/openai_agent_graph_runner.py create mode 100644 packages/ai-providers/server-ai-openai/tests/test_openai_agent_graph_runner.py create mode 100644 packages/sdk/server-ai/src/ldai/managed_agent_graph.py create mode 100644 packages/sdk/server-ai/tests/test_managed_agent_graph.py diff --git a/packages/ai-providers/server-ai-langchain/src/ldai_langchain/__init__.py b/packages/ai-providers/server-ai-langchain/src/ldai_langchain/__init__.py index cb455e5..04e299d 100644 --- a/packages/ai-providers/server-ai-langchain/src/ldai_langchain/__init__.py +++ b/packages/ai-providers/server-ai-langchain/src/ldai_langchain/__init__.py @@ -9,12 +9,14 @@ ) from ldai_langchain.langchain_model_runner import LangChainModelRunner from ldai_langchain.langchain_runner_factory import LangChainRunnerFactory +from ldai_langchain.langgraph_agent_graph_runner import LangGraphAgentGraphRunner __version__ = "0.1.0" __all__ = [ '__version__', 'LangChainRunnerFactory', + 'LangGraphAgentGraphRunner', 'LangChainModelRunner', 'convert_messages_to_langchain', 'create_langchain_model', diff --git a/packages/ai-providers/server-ai-langchain/src/ldai_langchain/langchain_runner_factory.py b/packages/ai-providers/server-ai-langchain/src/ldai_langchain/langchain_runner_factory.py index 402e295..f52923b 100644 --- a/packages/ai-providers/server-ai-langchain/src/ldai_langchain/langchain_runner_factory.py +++ b/packages/ai-providers/server-ai-langchain/src/ldai_langchain/langchain_runner_factory.py @@ -1,5 +1,7 @@ +from typing import Any + from ldai.models import AIConfigKind -from ldai.providers import AIProvider +from ldai.providers import AIProvider, ToolRegistry from ldai_langchain.langchain_helper import create_langchain_model from ldai_langchain.langchain_model_runner import LangChainModelRunner @@ -8,6 +10,17 @@ class LangChainRunnerFactory(AIProvider): """LangChain ``AIProvider`` implementation for the LaunchDarkly AI SDK.""" + def create_agent_graph(self, graph_def: Any, tools: ToolRegistry) -> Any: + """ + Create a configured LangGraphAgentGraphRunner for the given graph definition. + + :param graph_def: The AgentGraphDefinition to execute + :param tools: Registry mapping tool names to callables (langchain-compatible) + :return: LangGraphAgentGraphRunner ready to execute the graph + """ + from ldai_langchain.langgraph_agent_graph_runner import LangGraphAgentGraphRunner + return LangGraphAgentGraphRunner(graph_def, tools) + def create_model(self, config: AIConfigKind) -> LangChainModelRunner: """ Create a configured LangChainModelRunner for the given AI config. 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 new file mode 100644 index 0000000..99f81a1 --- /dev/null +++ b/packages/ai-providers/server-ai-langchain/src/ldai_langchain/langgraph_agent_graph_runner.py @@ -0,0 +1,138 @@ +"""LangGraph agent graph runner for LaunchDarkly AI SDK.""" + +import operator +import time +from typing import Annotated, Any, List + +from ldai.agent_graph import AgentGraphDefinition, AgentGraphNode +from ldai.providers.types import LDAIMetrics +from ldai.runners.agent_graph_runner import AgentGraphRunner +from ldai.runners.types import AgentGraphResult, ToolRegistry + + +class LangGraphAgentGraphRunner(AgentGraphRunner): + """ + AgentGraphRunner implementation for LangGraph. + + Builds a LangGraph StateGraph from an AgentGraphDefinition and + ToolRegistry via traverse(), compiles it, and executes it with + ainvoke(). Auto-tracks latency and invocation success/failure via + the graph's AIGraphTracker. + + Requires ``langgraph`` to be installed. + """ + + def __init__(self, graph: AgentGraphDefinition, tools: ToolRegistry): + """ + Initialize the runner. + + :param graph: The AgentGraphDefinition to execute + :param tools: Registry mapping tool names to callables (langchain-compatible) + """ + self._graph = graph + self._tools = tools + + async def run(self, input: Any) -> AgentGraphResult: + """ + Run the agent graph with the given input. + + Builds a LangGraph StateGraph from the AgentGraphDefinition, compiles + it, and invokes it. Tracks latency and invocation success/failure. + + :param input: The string prompt to send to the agent graph + :return: AgentGraphResult with the final output and metrics + """ + tracker = self._graph.get_tracker() + start_time = time.time() + try: + try: + from langchain.chat_models import init_chat_model + from langchain_core.messages import AnyMessage, HumanMessage + from langgraph.graph import END, START, StateGraph + from typing_extensions import TypedDict + except ImportError as exc: + raise ImportError( + "langgraph is required for LangGraphAgentGraphRunner. " + "Install it with: pip install langgraph" + ) from exc + + class WorkflowState(TypedDict): + messages: Annotated[List[AnyMessage], operator.add] + + agent_builder: StateGraph = StateGraph(WorkflowState) + root_node = self._graph.root() + root_key = root_node.get_key() if root_node else None + tools_ref = self._tools + + def handle_traversal(node: AgentGraphNode, ctx: dict) -> None: + node_config = node.get_config() + node_key = node.get_key() + + model = None + if node_config.model: + lc_model = init_chat_model(model=node_config.model.name) + tool_defs = node_config.model.get_parameter('tools') or [] + tool_fns = [ + tools_ref[t.get('name', '')] + for t in tool_defs + if t.get('name', '') in tools_ref + ] + if tool_fns: + lc_model = lc_model.bind_tools(tool_fns) + model = lc_model + + def invoke(state: WorkflowState) -> WorkflowState: + if model: + response = model.invoke(state['messages']) + return {'messages': [response]} + return state + + invoke.__name__ = node_key + + agent_builder.add_node(name=node_key, node=invoke) + + if node_key == root_key: + agent_builder.add_edge(START, node_key) + + if node.is_terminal(): + agent_builder.add_edge(node_key, END) + + for edge in node.get_edges(): + agent_builder.add_edge(node_key, edge.target_config) + + return None + + self._graph.traverse(fn=handle_traversal) + compiled = agent_builder.compile() + + result = await compiled.ainvoke( + {'messages': [HumanMessage(content=str(input))]} + ) + duration = int((time.time() - start_time) * 1000) + + output = '' + messages = result.get('messages', []) + if messages: + last = messages[-1] + if hasattr(last, 'content'): + output = str(last.content) + + if tracker: + tracker.track_latency(duration) + tracker.track_invocation_success() + + return AgentGraphResult( + output=output, + raw=result, + metrics=LDAIMetrics(success=True), + ) + except Exception: + duration = int((time.time() - start_time) * 1000) + if tracker: + tracker.track_latency(duration) + tracker.track_invocation_failure() + return AgentGraphResult( + output='', + raw=None, + metrics=LDAIMetrics(success=False), + ) diff --git a/packages/ai-providers/server-ai-langchain/tests/test_langgraph_agent_graph_runner.py b/packages/ai-providers/server-ai-langchain/tests/test_langgraph_agent_graph_runner.py new file mode 100644 index 0000000..8ad88a3 --- /dev/null +++ b/packages/ai-providers/server-ai-langchain/tests/test_langgraph_agent_graph_runner.py @@ -0,0 +1,139 @@ +"""Tests for LangGraphAgentGraphRunner and LangChainRunnerFactory.create_agent_graph().""" + +import pytest +from unittest.mock import AsyncMock, MagicMock, patch + +from ldai.agent_graph import AgentGraphDefinition +from ldai.models import AIAgentGraphConfig, AIAgentConfig, ModelConfig, ProviderConfig +from ldai.runners.types import AgentGraphResult, ToolRegistry +from ldai_langchain.langgraph_agent_graph_runner import LangGraphAgentGraphRunner +from ldai_langchain.langchain_runner_factory import LangChainRunnerFactory + + +def _make_graph(enabled: bool = True) -> AgentGraphDefinition: + root_config = AIAgentConfig( + key='root-agent', + enabled=enabled, + model=ModelConfig(name='gpt-4'), + provider=ProviderConfig(name='openai'), + instructions='You are a helpful assistant.', + tracker=MagicMock(), + ) + graph_config = AIAgentGraphConfig( + key='test-graph', + root_config_key='root-agent', + edges=[], + enabled=enabled, + ) + nodes = AgentGraphDefinition.build_nodes(graph_config, {'root-agent': root_config}) + return AgentGraphDefinition( + agent_graph=graph_config, + nodes=nodes, + context=MagicMock(), + enabled=enabled, + tracker=MagicMock(), + ) + + +# --- Factory --- + +def test_langchain_runner_factory_create_agent_graph_returns_runner(): + graph = _make_graph() + tools: ToolRegistry = {'fetch_weather': lambda loc: f'weather in {loc}'} + factory = LangChainRunnerFactory() + runner = factory.create_agent_graph(graph, tools) + assert isinstance(runner, LangGraphAgentGraphRunner) + + +def test_langchain_runner_factory_create_agent_graph_wires_graph_and_tools(): + graph = _make_graph() + tools: ToolRegistry = {} + factory = LangChainRunnerFactory() + runner = factory.create_agent_graph(graph, tools) + assert runner._graph is graph + assert runner._tools is tools + + +# --- LangGraphAgentGraphRunner --- + +def test_langgraph_runner_stores_graph_and_tools(): + graph = _make_graph() + tools: ToolRegistry = {} + runner = LangGraphAgentGraphRunner(graph, tools) + assert runner._graph is graph + assert runner._tools is tools + + +@pytest.mark.asyncio +async def test_langgraph_runner_run_raises_when_langgraph_not_installed(): + graph = _make_graph() + runner = LangGraphAgentGraphRunner(graph, {}) + + with patch.dict('sys.modules', {'langgraph': None, 'langgraph.graph': None}): + result = await runner.run("test") + assert isinstance(result, AgentGraphResult) + assert result.metrics.success is False + + +@pytest.mark.asyncio +async def test_langgraph_runner_run_tracks_failure_on_exception(): + graph = _make_graph() + tracker = graph.get_tracker() + runner = LangGraphAgentGraphRunner(graph, {}) + + with patch.dict('sys.modules', {'langgraph': None, 'langgraph.graph': None}): + result = await runner.run("fail") + + assert result.metrics.success is False + tracker.track_invocation_failure.assert_called_once() + tracker.track_latency.assert_called_once() + + +@pytest.mark.asyncio +async def test_langgraph_runner_run_success(): + graph = _make_graph() + tracker = graph.get_tracker() + + mock_message = MagicMock() + mock_message.content = "langgraph answer" + + mock_compiled = MagicMock() + mock_compiled.ainvoke = AsyncMock(return_value={'messages': [mock_message]}) + + mock_state_graph_instance = MagicMock() + mock_state_graph_instance.add_node = MagicMock() + mock_state_graph_instance.add_edge = MagicMock() + mock_state_graph_instance.compile = MagicMock(return_value=mock_compiled) + + mock_langgraph_graph = MagicMock() + mock_langgraph_graph.END = 'END' + mock_langgraph_graph.START = 'START' + mock_langgraph_graph.StateGraph = MagicMock(return_value=mock_state_graph_instance) + + mock_human_message = MagicMock() + mock_lc_core_messages = MagicMock() + mock_lc_core_messages.HumanMessage = MagicMock(return_value=mock_human_message) + mock_lc_core_messages.AnyMessage = MagicMock() + + mock_init_model = MagicMock() + mock_init_model.return_value = MagicMock() + mock_langchain_chat = MagicMock() + mock_langchain_chat.init_chat_model = mock_init_model + + with patch.dict('sys.modules', { + 'langgraph': MagicMock(), + 'langgraph.graph': mock_langgraph_graph, + 'langchain_core': MagicMock(), + 'langchain_core.messages': mock_lc_core_messages, + 'langchain': MagicMock(), + 'langchain.chat_models': mock_langchain_chat, + 'typing_extensions': __import__('typing_extensions'), + }): + runner = LangGraphAgentGraphRunner(graph, {}) + result = await runner.run("find restaurants") + + assert isinstance(result, AgentGraphResult) + assert result.output == "langgraph answer" + assert result.metrics.success is True + tracker.track_invocation_success.assert_called_once() + tracker.track_latency.assert_called_once() diff --git a/packages/ai-providers/server-ai-openai/src/ldai_openai/__init__.py b/packages/ai-providers/server-ai-openai/src/ldai_openai/__init__.py index 8a8199b..422c059 100644 --- a/packages/ai-providers/server-ai-openai/src/ldai_openai/__init__.py +++ b/packages/ai-providers/server-ai-openai/src/ldai_openai/__init__.py @@ -1,3 +1,4 @@ +from ldai_openai.openai_agent_graph_runner import OpenAIAgentGraphRunner from ldai_openai.openai_helper import ( convert_messages_to_openai, get_ai_metrics_from_response, @@ -8,6 +9,7 @@ __all__ = [ 'OpenAIRunnerFactory', + 'OpenAIAgentGraphRunner', 'OpenAIModelRunner', 'convert_messages_to_openai', 'get_ai_metrics_from_response', 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 new file mode 100644 index 0000000..447058e --- /dev/null +++ b/packages/ai-providers/server-ai-openai/src/ldai_openai/openai_agent_graph_runner.py @@ -0,0 +1,210 @@ +"""OpenAI agent graph runner for LaunchDarkly AI SDK.""" + +import time +from typing import Any, List, Optional + +from ldai.agent_graph import AgentGraphDefinition, AgentGraphNode +from ldai.providers.types import LDAIMetrics +from ldai.runners.agent_graph_runner import AgentGraphRunner +from ldai.runners.types import AgentGraphResult, ToolRegistry +from ldai.tracker import TokenUsage + + +def _to_openai_name(name: str) -> str: + """Convert a hyphenated tool/node name to an underscore-separated OpenAI function name.""" + return name.replace('-', '_') + + +class OpenAIAgentGraphRunner(AgentGraphRunner): + """ + AgentGraphRunner implementation for the OpenAI Agents SDK. + + Builds agents from an AgentGraphDefinition and a ToolRegistry via + reverse_traverse, executes them with Runner.run(), and auto-tracks + path, tool calls, handoffs, latency, and invocation success/failure + via the graph's AIGraphTracker. + + Requires ``openai-agents`` to be installed. + """ + + def __init__(self, graph: AgentGraphDefinition, tools: ToolRegistry): + """ + Initialize the runner. + + :param graph: The AgentGraphDefinition to execute + :param tools: Registry mapping OpenAI-formatted tool names to callables + """ + self._graph = graph + self._tools = tools + + async def run(self, input: Any) -> AgentGraphResult: + """ + Run the agent graph with the given input. + + Builds the agent tree via reverse_traverse, then invokes the root + agent with Runner.run(). Tracks path, latency, and invocation + success/failure. + + :param input: The string prompt to send to the agent graph + :return: AgentGraphResult with the final output and metrics + """ + tracker = self._graph.get_tracker() + path: List[str] = [] + root_node = self._graph.root() + if root_node: + path.append(root_node.get_key()) + + start_time = time.time() + try: + try: + from agents import Runner + except ImportError as exc: + raise ImportError( + "openai-agents is required for OpenAIAgentGraphRunner. " + "Install it with: pip install openai-agents" + ) from exc + + root_agent = self._build_agents(path) + result = await Runner.run(root_agent, str(input)) + duration = int((time.time() - start_time) * 1000) + + if tracker: + tracker.track_path(path) + tracker.track_latency(duration) + tracker.track_invocation_success() + + return AgentGraphResult( + output=str(result.final_output), + raw=result, + metrics=LDAIMetrics(success=True), + ) + except Exception: + duration = int((time.time() - start_time) * 1000) + if tracker: + tracker.track_latency(duration) + tracker.track_invocation_failure() + return AgentGraphResult( + output='', + raw=None, + metrics=LDAIMetrics(success=False), + ) + + def _build_agents(self, path: List[str]) -> Any: + """ + Build the agent tree from the graph definition via reverse_traverse. + + Agents are constructed from terminal nodes upward so that handoff + targets exist before the agents that hand off to them. + + :param path: Mutable list to accumulate the execution path + :return: The root Agent instance + """ + try: + from agents import Agent, FunctionTool, Handoff, RunContextWrapper, Tool, handoff + from agents.extensions.handoff_prompt import RECOMMENDED_PROMPT_PREFIX + from agents.tool_context import ToolContext + except ImportError as exc: + raise ImportError( + "openai-agents is required for OpenAIAgentGraphRunner. " + "Install it with: pip install openai-agents" + ) from exc + + tracker = self._graph.get_tracker() + + def build_node(node: AgentGraphNode, ctx: dict) -> Any: + node_config = node.get_config() + config_tracker = node_config.tracker + model = node_config.model + + if not model: + raise ValueError(f"Model not set for node '{node_config.key}'") + + tool_defs = model.get_parameter('tools') or [] + + # --- handoffs --- + agent_handoffs: List[Handoff] = [] + for edge in node.get_edges(): + target_key = edge.target_config + + def _make_on_handoff(src: str, tgt: str): + def on_handoff(run_ctx: RunContextWrapper) -> None: + path.append(tgt) + if tracker: + tracker.track_handoff_success(src, tgt) + tracker.track_node_invocation(src) + if config_tracker: + try: + usage_entry = run_ctx.usage.request_usage_entries[-1] + config_tracker.track_tokens( + TokenUsage( + total=usage_entry.total_tokens, + input=usage_entry.input_tokens, + output=usage_entry.output_tokens, + ) + ) + except Exception: + pass + config_tracker.track_success() + return on_handoff + + agent_handoffs.append( + handoff( + agent=ctx[target_key], + on_handoff=_make_on_handoff(node_config.key, target_key), + ) + ) + + # --- tools --- + agent_tools: List[Tool] = [] + for tool_def in tool_defs: + tool_name_raw = tool_def.get('name', '') + tool_name = _to_openai_name(tool_name_raw) + tool_fn = self._tools.get(tool_name) or self._tools.get(tool_name_raw) + if not tool_fn: + continue + + def _make_tool( + name: str, + raw_name: str, + fn: Any, + description: str, + params_schema: dict, + cfg_key: str, + ) -> FunctionTool: + def wrapped(tool_ctx: ToolContext, tool_args: str) -> Any: + import json + try: + args = json.loads(tool_args) + except Exception: + args = {} + path.append(raw_name) + if tracker: + tracker.track_tool_call(config_key=cfg_key, tool_key=name) + return fn(**args) + + return FunctionTool( + name=f'tool_{name}', + description=description, + params_json_schema=params_schema, + on_invoke_tool=wrapped, + ) + + agent_tools.append( + _make_tool( + tool_name, + tool_name_raw, + tool_fn, + tool_def.get('description', ''), + tool_def.get('parameters', {}), + node_config.key, + ) + ) + + return Agent( + name=_to_openai_name(node_config.key), + 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) diff --git a/packages/ai-providers/server-ai-openai/src/ldai_openai/openai_runner_factory.py b/packages/ai-providers/server-ai-openai/src/ldai_openai/openai_runner_factory.py index d80fc01..c52b5a6 100644 --- a/packages/ai-providers/server-ai-openai/src/ldai_openai/openai_runner_factory.py +++ b/packages/ai-providers/server-ai-openai/src/ldai_openai/openai_runner_factory.py @@ -1,8 +1,9 @@ import os -from typing import Optional +from typing import Any, Optional from ldai.models import AIConfigKind from ldai.providers import AIProvider +from ldai.runners.types import ToolRegistry from openai import AsyncOpenAI from ldai_openai.openai_model_runner import OpenAIModelRunner @@ -36,6 +37,17 @@ def create_model(self, config: AIConfigKind) -> OpenAIModelRunner: parameters = model_dict.get('parameters') or {} return OpenAIModelRunner(self._client, model_name, parameters) + def create_agent_graph(self, graph_def: Any, tools: ToolRegistry) -> Any: + """ + Create a configured OpenAIAgentGraphRunner for the given graph definition. + + :param graph_def: The AgentGraphDefinition to execute + :param tools: Registry mapping tool names to callables + :return: OpenAIAgentGraphRunner ready to execute the graph + """ + from ldai_openai.openai_agent_graph_runner import OpenAIAgentGraphRunner + return OpenAIAgentGraphRunner(graph_def, tools) + def get_client(self) -> AsyncOpenAI: """ Return the underlying AsyncOpenAI client. diff --git a/packages/ai-providers/server-ai-openai/tests/test_openai_agent_graph_runner.py b/packages/ai-providers/server-ai-openai/tests/test_openai_agent_graph_runner.py new file mode 100644 index 0000000..bc6fb2c --- /dev/null +++ b/packages/ai-providers/server-ai-openai/tests/test_openai_agent_graph_runner.py @@ -0,0 +1,134 @@ +"""Tests for OpenAIAgentGraphRunner and OpenAIRunnerFactory.create_agent_graph().""" + +import pytest +from unittest.mock import MagicMock, AsyncMock, patch + +from ldai.agent_graph import AgentGraphDefinition +from ldai.models import AIAgentGraphConfig, AIAgentConfig, Edge, ModelConfig, ProviderConfig +from ldai.runners.types import AgentGraphResult, ToolRegistry +from ldai_openai.openai_agent_graph_runner import OpenAIAgentGraphRunner +from ldai_openai.openai_runner_factory import OpenAIRunnerFactory + + +def _make_graph(enabled: bool = True) -> AgentGraphDefinition: + """Build a minimal single-node AgentGraphDefinition for testing.""" + root_config = AIAgentConfig( + key='root-agent', + enabled=enabled, + model=ModelConfig(name='gpt-4'), + provider=ProviderConfig(name='openai'), + instructions='You are a helpful assistant.', + tracker=MagicMock(), + ) + graph_config = AIAgentGraphConfig( + key='test-graph', + root_config_key='root-agent', + edges=[], + enabled=enabled, + ) + nodes = AgentGraphDefinition.build_nodes(graph_config, {'root-agent': root_config}) + return AgentGraphDefinition( + agent_graph=graph_config, + nodes=nodes, + context=MagicMock(), + enabled=enabled, + tracker=MagicMock(), + ) + + +# --- Factory --- + +def test_openai_runner_factory_create_agent_graph_returns_runner(): + graph = _make_graph() + tools: ToolRegistry = {'search': lambda q: q} + factory = OpenAIRunnerFactory(client=MagicMock()) + runner = factory.create_agent_graph(graph, tools) + assert isinstance(runner, OpenAIAgentGraphRunner) + + +def test_openai_runner_factory_create_agent_graph_wires_graph_and_tools(): + graph = _make_graph() + tools: ToolRegistry = {'my_tool': lambda: None} + factory = OpenAIRunnerFactory(client=MagicMock()) + runner = factory.create_agent_graph(graph, tools) + assert runner._graph is graph + assert runner._tools is tools + + +# --- OpenAIAgentGraphRunner --- + +def test_openai_agent_graph_runner_stores_graph_and_tools(): + graph = _make_graph() + tools: ToolRegistry = {} + runner = OpenAIAgentGraphRunner(graph, tools) + assert runner._graph is graph + assert runner._tools is tools + + +@pytest.mark.asyncio +async def test_openai_agent_graph_runner_run_raises_when_agents_not_installed(): + graph = _make_graph() + runner = OpenAIAgentGraphRunner(graph, {}) + + with patch.dict('sys.modules', {'agents': None}): + # The import inside run() will fail — runner should return failure result + # rather than propagate the ImportError, since it's caught by the except block + result = await runner.run("test input") + assert isinstance(result, AgentGraphResult) + assert result.metrics.success is False + + +@pytest.mark.asyncio +async def test_openai_agent_graph_runner_run_tracks_invocation_failure_on_exception(): + graph = _make_graph() + tracker = graph.get_tracker() + runner = OpenAIAgentGraphRunner(graph, {}) + + with patch.dict('sys.modules', {'agents': None}): + result = await runner.run("fail") + + assert result.metrics.success is False + tracker.track_invocation_failure.assert_called_once() + tracker.track_latency.assert_called_once() + + +@pytest.mark.asyncio +async def test_openai_agent_graph_runner_run_success(): + graph = _make_graph() + tracker = graph.get_tracker() + + mock_result = MagicMock() + mock_result.final_output = "agent answer" + + mock_runner_module = MagicMock() + mock_runner_module.run = AsyncMock(return_value=mock_result) + + mock_agents = MagicMock() + mock_agents.Runner = mock_runner_module + mock_agents.Agent = MagicMock(return_value=MagicMock()) + mock_agents.FunctionTool = MagicMock() + mock_agents.Handoff = MagicMock() + mock_agents.RunContextWrapper = MagicMock() + mock_agents.Tool = MagicMock() + mock_agents.handoff = MagicMock() + + mock_agents_ext = MagicMock() + mock_agents_ext.RECOMMENDED_PROMPT_PREFIX = '[PREFIX]' + + mock_tool_context = MagicMock() + + with patch.dict('sys.modules', { + 'agents': mock_agents, + 'agents.extensions': MagicMock(), + 'agents.extensions.handoff_prompt': mock_agents_ext, + 'agents.tool_context': mock_tool_context, + }): + runner = OpenAIAgentGraphRunner(graph, {}) + result = await runner.run("find restaurants") + + assert isinstance(result, AgentGraphResult) + assert result.output == "agent answer" + assert result.metrics.success is True + tracker.track_invocation_success.assert_called_once() + tracker.track_path.assert_called_once() + tracker.track_latency.assert_called_once() diff --git a/packages/sdk/server-ai/src/ldai/__init__.py b/packages/sdk/server-ai/src/ldai/__init__.py index 944a0cb..25295a7 100644 --- a/packages/sdk/server-ai/src/ldai/__init__.py +++ b/packages/sdk/server-ai/src/ldai/__init__.py @@ -6,6 +6,7 @@ from ldai.chat import Chat # Deprecated — use ManagedModel from ldai.client import LDAIClient from ldai.judge import Judge +from ldai.managed_agent_graph import ManagedAgentGraph from ldai.managed_model import ManagedModel from ldai.models import ( # Deprecated aliases for backward compatibility AIAgentConfig, @@ -56,6 +57,7 @@ 'AIJudgeConfig', 'AIJudgeConfigDefault', 'ManagedModel', + 'ManagedAgentGraph', 'EvalScore', 'AgentGraphDefinition', 'Judge', diff --git a/packages/sdk/server-ai/src/ldai/client.py b/packages/sdk/server-ai/src/ldai/client.py index 358f9eb..1c72ec4 100644 --- a/packages/sdk/server-ai/src/ldai/client.py +++ b/packages/sdk/server-ai/src/ldai/client.py @@ -7,6 +7,7 @@ from ldai import log from ldai.agent_graph import AgentGraphDefinition from ldai.judge import Judge +from ldai.managed_agent_graph import ManagedAgentGraph from ldai.managed_model import ManagedModel from ldai.models import ( AIAgentConfig, @@ -25,6 +26,7 @@ ProviderConfig, ) from ldai.providers.runner_factory import RunnerFactory +from ldai.runners.types import ToolRegistry from ldai.sdk_info import AI_SDK_LANGUAGE, AI_SDK_NAME, AI_SDK_VERSION from ldai.tracker import AIGraphTracker, LDAIConfigTracker @@ -35,6 +37,7 @@ _TRACK_USAGE_CREATE_JUDGE = '$ld:ai:usage:create-judge' _TRACK_USAGE_AGENT_CONFIG = '$ld:ai:usage:agent-config' _TRACK_USAGE_AGENT_CONFIGS = '$ld:ai:usage:agent-configs' +_TRACK_USAGE_CREATE_AGENT_GRAPH = '$ld:ai:usage:create-agent-graph' _INIT_TRACK_CONTEXT = Context.builder('ld-internal-tracking').kind('ld_ai').anonymous(True).build() @@ -609,6 +612,55 @@ def agent_graph( tracker=tracker, ) + async def create_agent_graph( + self, + key: str, + context: Context, + tools: Optional[ToolRegistry] = None, + default_ai_provider: Optional[str] = None, + ) -> Optional[ManagedAgentGraph]: + """ + Creates and returns a new ManagedAgentGraph for AI agent graph execution. + + Resolves the graph configuration via ``agent_graph()``, creates a + provider-specific runner, and wraps it in a ``ManagedAgentGraph``. + + :param key: The key identifying the agent graph configuration + :param context: Standard Context used when evaluating flags + :param tools: Registry mapping tool names to callables + :param default_ai_provider: Optional provider override ('openai', 'langchain', …) + :return: ManagedAgentGraph instance, or None if the graph is disabled or unsupported + + Example:: + + graph = await client.create_agent_graph( + "travel-assistant-graph", + context, + tools={ + "web_search_tool": my_search_fn, + "get_weather": my_weather_fn, + } + ) + + if graph: + result = await graph.run("Find me restaurants in Seattle") + print(result.output) + """ + self._client.track(_TRACK_USAGE_CREATE_AGENT_GRAPH, context, key, 1) + log.debug(f"Creating managed agent graph for key: {key}") + + graph = self.agent_graph(key, context) + if not graph.enabled: + return None + + runner = await RunnerFactory.create_agent_graph( + graph, tools or {}, default_ai_provider + ) + if not runner: + return None + + return ManagedAgentGraph(runner, graph.get_tracker()) + def agents( self, agent_configs: List[AIAgentConfigRequest], diff --git a/packages/sdk/server-ai/src/ldai/managed_agent_graph.py b/packages/sdk/server-ai/src/ldai/managed_agent_graph.py new file mode 100644 index 0000000..6eac9d3 --- /dev/null +++ b/packages/sdk/server-ai/src/ldai/managed_agent_graph.py @@ -0,0 +1,61 @@ +"""ManagedAgentGraph — LaunchDarkly managed wrapper for agent graph execution.""" + +from typing import Any, Optional + +from ldai.runners.agent_graph_runner import AgentGraphRunner +from ldai.runners.types import AgentGraphResult +from ldai.tracker import AIGraphTracker + + +class ManagedAgentGraph: + """ + LaunchDarkly managed wrapper for AI agent graph execution. + + Holds an AgentGraphRunner and an AIGraphTracker. Auto-tracking of path, + tool calls, handoffs, latency, and invocation success/failure is handled + by the runner implementation. + + Obtain an instance via ``LDAIClient.create_agent_graph()``. + """ + + def __init__( + self, + runner: AgentGraphRunner, + tracker: Optional[AIGraphTracker] = None, + ): + """ + Initialize ManagedAgentGraph. + + :param runner: The AgentGraphRunner to delegate execution to + :param tracker: The AIGraphTracker for this graph + """ + self._runner = runner + self._tracker = tracker + + async def run(self, input: Any) -> AgentGraphResult: + """ + Run the agent graph with the given input. + + Delegates to the underlying AgentGraphRunner, which handles + execution and all auto-tracking internally. + + :param input: The input prompt or structured input for the graph + :return: AgentGraphResult containing the output, raw response, and metrics + """ + return await self._runner.run(input) + + def get_agent_graph_runner(self) -> AgentGraphRunner: + """ + Return the underlying AgentGraphRunner for advanced use. + + :return: The AgentGraphRunner instance + """ + return self._runner + + def get_tracker(self) -> Optional[AIGraphTracker]: + """ + Return the AIGraphTracker for this graph. + + :return: The AIGraphTracker instance, or None if not available + """ + return self._tracker diff --git a/packages/sdk/server-ai/tests/test_managed_agent_graph.py b/packages/sdk/server-ai/tests/test_managed_agent_graph.py new file mode 100644 index 0000000..8ffa73b --- /dev/null +++ b/packages/sdk/server-ai/tests/test_managed_agent_graph.py @@ -0,0 +1,196 @@ +"""Tests for ManagedAgentGraph and LDAIClient.create_agent_graph().""" + +import pytest +from unittest.mock import AsyncMock, MagicMock, patch +from ldclient import Config, Context, LDClient +from ldclient.integrations.test_data import TestData + +from ldai import LDAIClient, ManagedAgentGraph +from ldai.providers.types import LDAIMetrics +from ldai.runners.agent_graph_runner import AgentGraphRunner +from ldai.runners.types import AgentGraphResult, ToolRegistry +from ldai.tracker import AIGraphTracker + + +# --- Test double --- + +class StubAgentGraphRunner(AgentGraphRunner): + def __init__(self, output: str = "stub output"): + self._output = output + + async def run(self, input) -> AgentGraphResult: + return AgentGraphResult( + output=self._output, + raw={"input": input}, + metrics=LDAIMetrics(success=True), + ) + + +# --- ManagedAgentGraph unit tests --- + +@pytest.mark.asyncio +async def test_managed_agent_graph_run_delegates_to_runner(): + runner = StubAgentGraphRunner("hello world") + managed = ManagedAgentGraph(runner) + result = await managed.run("test input") + assert result.output == "hello world" + assert result.metrics.success is True + + +def test_managed_agent_graph_get_runner(): + runner = StubAgentGraphRunner() + managed = ManagedAgentGraph(runner) + assert managed.get_agent_graph_runner() is runner + + +def test_managed_agent_graph_get_tracker_none_by_default(): + runner = StubAgentGraphRunner() + managed = ManagedAgentGraph(runner) + assert managed.get_tracker() is None + + +def test_managed_agent_graph_get_tracker_returns_tracker(): + runner = StubAgentGraphRunner() + tracker = MagicMock(spec=AIGraphTracker) + managed = ManagedAgentGraph(runner, tracker) + assert managed.get_tracker() is tracker + + +# --- LDAIClient.create_agent_graph() integration tests --- + +@pytest.fixture +def td() -> TestData: + td = TestData.data_source() + + td.update( + td.flag('travel-graph') + .variations({ + 'root': 'triage-agent', + 'edges': { + 'triage-agent': [{'key': 'specialist-agent'}], + }, + '_ldMeta': {'enabled': True, 'variationKey': 'v1', 'version': 1}, + }) + .variation_for_all(0) + ) + + td.update( + td.flag('triage-agent') + .variations({ + 'model': {'name': 'gpt-4'}, + 'provider': {'name': 'openai'}, + 'instructions': 'You are a triage agent.', + '_ldMeta': {'enabled': True, 'variationKey': 'triage-v1', 'version': 1}, + }) + .variation_for_all(0) + ) + + td.update( + td.flag('specialist-agent') + .variations({ + 'model': {'name': 'gpt-4'}, + 'provider': {'name': 'openai'}, + 'instructions': 'You are a specialist.', + '_ldMeta': {'enabled': True, 'variationKey': 'specialist-v1', 'version': 1}, + }) + .variation_for_all(0) + ) + + td.update( + td.flag('disabled-graph') + .variations({ + '_ldMeta': {'enabled': False, 'variationKey': 'disabled-v1', 'version': 1}, + }) + .variation_for_all(0) + ) + + return td + + +@pytest.fixture +def client(td: TestData) -> LDClient: + config = Config('sdk-key', update_processor_class=td, send_events=False) + return LDClient(config=config) + + +@pytest.fixture +def ldai_client(client: LDClient) -> LDAIClient: + return LDAIClient(client) + + +@pytest.mark.asyncio +async def test_create_agent_graph_returns_managed_agent_graph(ldai_client: LDAIClient): + context = Context.create('user-key') + stub_runner = StubAgentGraphRunner("result") + + with patch( + 'ldai.providers.runner_factory.RunnerFactory.create_agent_graph', + new=AsyncMock(return_value=stub_runner), + ): + managed = await ldai_client.create_agent_graph('travel-graph', context) + + assert managed is not None + assert isinstance(managed, ManagedAgentGraph) + assert managed.get_agent_graph_runner() is stub_runner + + +@pytest.mark.asyncio +async def test_create_agent_graph_returns_none_when_disabled(ldai_client: LDAIClient): + context = Context.create('user-key') + managed = await ldai_client.create_agent_graph('disabled-graph', context) + assert managed is None + + +@pytest.mark.asyncio +async def test_create_agent_graph_returns_none_when_runner_factory_fails(ldai_client: LDAIClient): + context = Context.create('user-key') + + with patch( + 'ldai.providers.runner_factory.RunnerFactory.create_agent_graph', + new=AsyncMock(return_value=None), + ): + managed = await ldai_client.create_agent_graph('travel-graph', context) + + assert managed is None + + +@pytest.mark.asyncio +async def test_create_agent_graph_passes_tools_to_factory(ldai_client: LDAIClient): + context = Context.create('user-key') + tools: ToolRegistry = {'search': lambda q: f'results for {q}'} + captured = {} + + async def fake_create_agent_graph(graph_def, tools_arg, default_ai_provider=None): + captured['tools'] = tools_arg + return StubAgentGraphRunner() + + with patch( + 'ldai.providers.runner_factory.RunnerFactory.create_agent_graph', + new=fake_create_agent_graph, + ): + await ldai_client.create_agent_graph('travel-graph', context, tools=tools) + + assert captured['tools'] is tools + + +@pytest.mark.asyncio +async def test_create_agent_graph_run_produces_result(ldai_client: LDAIClient): + context = Context.create('user-key') + + with patch( + 'ldai.providers.runner_factory.RunnerFactory.create_agent_graph', + new=AsyncMock(return_value=StubAgentGraphRunner("final answer")), + ): + managed = await ldai_client.create_agent_graph('travel-graph', context) + + assert managed is not None + result = await managed.run("find restaurants") + assert result.output == "final answer" + assert result.metrics.success is True + + +# --- Top-level export --- + +def test_managed_agent_graph_exported(): + import ldai + assert hasattr(ldai, 'ManagedAgentGraph') From be26d6d2d5ee9ed861020e64eaa1a00563dcbff0 Mon Sep 17 00:00:00 2001 From: jsonbailey Date: Wed, 25 Mar 2026 17:22:50 -0500 Subject: [PATCH 02/10] feat: improve graph runner tracking with node metrics, path, and token rollup Co-Authored-By: Claude Sonnet 4.6 --- .../langgraph_agent_graph_runner.py | 65 ++++--- .../test_langgraph_agent_graph_runner.py | 14 +- .../ldai_openai/openai_agent_graph_runner.py | 169 +++++++++++++----- .../tests/test_openai_agent_graph_runner.py | 3 + 4 files changed, 185 insertions(+), 66 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 99f81a1..5388d8f 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 @@ -4,20 +4,22 @@ import time from typing import Annotated, Any, List +from ldai import log from ldai.agent_graph import AgentGraphDefinition, AgentGraphNode from ldai.providers.types import LDAIMetrics from ldai.runners.agent_graph_runner import AgentGraphRunner from ldai.runners.types import AgentGraphResult, ToolRegistry +from ldai_langchain.langchain_helper import LangChainHelper + class LangGraphAgentGraphRunner(AgentGraphRunner): """ AgentGraphRunner implementation for LangGraph. - Builds a LangGraph StateGraph from an AgentGraphDefinition and - ToolRegistry via traverse(), compiles it, and executes it with - ainvoke(). Auto-tracks latency and invocation success/failure via - the graph's AIGraphTracker. + Compiles and runs the agent graph with LangGraph and automatically records + graph- and node-level AI metric data to the LaunchDarkly trackers on the + graph definition and each node. Requires ``langgraph`` to be installed. """ @@ -43,18 +45,12 @@ async def run(self, input: Any) -> AgentGraphResult: :return: AgentGraphResult with the final output and metrics """ tracker = self._graph.get_tracker() - start_time = time.time() + start_ns = time.perf_counter_ns() try: - try: - from langchain.chat_models import init_chat_model - from langchain_core.messages import AnyMessage, HumanMessage - from langgraph.graph import END, START, StateGraph - from typing_extensions import TypedDict - except ImportError as exc: - raise ImportError( - "langgraph is required for LangGraphAgentGraphRunner. " - "Install it with: pip install langgraph" - ) from exc + from langchain.chat_models import init_chat_model + from langchain_core.messages import AnyMessage, HumanMessage + from langgraph.graph import END, START, StateGraph + from typing_extensions import TypedDict class WorkflowState(TypedDict): messages: Annotated[List[AnyMessage], operator.add] @@ -63,10 +59,12 @@ class WorkflowState(TypedDict): root_node = self._graph.root() root_key = root_node.get_key() if root_node else None tools_ref = self._tools + exec_path: List[str] = [] def handle_traversal(node: AgentGraphNode, ctx: dict) -> None: node_config = node.get_config() node_key = node.get_key() + node_tracker = node_config.tracker model = None if node_config.model: @@ -82,10 +80,24 @@ def handle_traversal(node: AgentGraphNode, ctx: dict) -> None: model = lc_model def invoke(state: WorkflowState) -> WorkflowState: - if model: + exec_path.append(node_key) + if not model: + return state + gk = tracker.graph_key if tracker is not None else None + if node_tracker: + response = node_tracker.track_metrics_of( + lambda: model.invoke(state['messages']), + LangChainHelper.get_ai_metrics_from_response, + graph_key=gk, + ) + node_tracker.track_tool_calls( + LangChainHelper.get_tool_calls_from_response(response), + graph_key=tracker.graph_key if tracker is not None else None, + ) + else: response = model.invoke(state['messages']) - return {'messages': [response]} - return state + + return {'messages': [response]} invoke.__name__ = node_key @@ -108,7 +120,7 @@ def invoke(state: WorkflowState) -> WorkflowState: result = await compiled.ainvoke( {'messages': [HumanMessage(content=str(input))]} ) - duration = int((time.time() - start_time) * 1000) + duration = (time.perf_counter_ns() - start_ns) // 1_000_000 output = '' messages = result.get('messages', []) @@ -118,16 +130,27 @@ def invoke(state: WorkflowState) -> WorkflowState: output = str(last.content) if tracker: + tracker.track_path(exec_path) tracker.track_latency(duration) tracker.track_invocation_success() + tracker.track_total_tokens( + LangChainHelper.sum_token_usage_from_messages(messages) + ) return AgentGraphResult( output=output, raw=result, metrics=LDAIMetrics(success=True), ) - except Exception: - duration = int((time.time() - start_time) * 1000) + except Exception as exc: + if isinstance(exc, ImportError): + log.warning( + "langgraph is required for LangGraphAgentGraphRunner. " + "Install it with: pip install langgraph" + ) + else: + log.warning(f'LangGraphAgentGraphRunner run failed: {exc}') + duration = (time.perf_counter_ns() - start_ns) // 1_000_000 if tracker: tracker.track_latency(duration) tracker.track_invocation_failure() diff --git a/packages/ai-providers/server-ai-langchain/tests/test_langgraph_agent_graph_runner.py b/packages/ai-providers/server-ai-langchain/tests/test_langgraph_agent_graph_runner.py index 8ad88a3..67229d5 100644 --- a/packages/ai-providers/server-ai-langchain/tests/test_langgraph_agent_graph_runner.py +++ b/packages/ai-providers/server-ai-langchain/tests/test_langgraph_agent_graph_runner.py @@ -96,6 +96,8 @@ async def test_langgraph_runner_run_success(): mock_message = MagicMock() mock_message.content = "langgraph answer" + mock_message.usage_metadata = None + mock_message.response_metadata = None mock_compiled = MagicMock() mock_compiled.ainvoke = AsyncMock(return_value={'messages': [mock_message]}) @@ -115,8 +117,17 @@ async def test_langgraph_runner_run_success(): mock_lc_core_messages.HumanMessage = MagicMock(return_value=mock_human_message) mock_lc_core_messages.AnyMessage = MagicMock() + mock_model_response = MagicMock() + mock_model_response.content = 'langgraph answer' + mock_model_response.usage_metadata = None + mock_model_response.response_metadata = None + mock_model_response.tool_calls = None + + mock_llm = MagicMock() + mock_llm.invoke = MagicMock(return_value=mock_model_response) + mock_init_model = MagicMock() - mock_init_model.return_value = MagicMock() + mock_init_model.return_value = mock_llm mock_langchain_chat = MagicMock() mock_langchain_chat.init_chat_model = mock_init_model @@ -135,5 +146,6 @@ async def test_langgraph_runner_run_success(): assert isinstance(result, AgentGraphResult) assert result.output == "langgraph answer" assert result.metrics.success is True + tracker.track_path.assert_called_once_with([]) tracker.track_invocation_success.assert_called_once() tracker.track_latency.assert_called_once() 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 447058e..afebbcb 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 @@ -3,6 +3,7 @@ import time from typing import Any, List, Optional +from ldai import log from ldai.agent_graph import AgentGraphDefinition, AgentGraphNode from ldai.providers.types import LDAIMetrics from ldai.runners.agent_graph_runner import AgentGraphRunner @@ -15,14 +16,46 @@ def _to_openai_name(name: str) -> str: return name.replace('-', '_') +def _log_run_result_shape(result: Any) -> None: + """Print RunResult attributes (excluding final_output) for debugging.""" + attrs = [a for a in dir(result) if not a.startswith('_')] + print("RunResult public attributes:", attrs) + for name in attrs: + if name == 'final_output': + continue + try: + val = getattr(result, name) + if callable(val): + print(f" {name}: callable") + else: + print(f" {name}: {repr(val)}") + except Exception as e: + print(f" {name}: (error reading: {e})") + + +def _build_native_tool_map() -> dict: + try: + from agents import CodeInterpreterTool, FileSearchTool, ImageGenerationTool, WebSearchTool + return { + 'web_search_tool': lambda _: WebSearchTool(), + 'file_search_tool': lambda _: FileSearchTool(), + 'code_interpreter': lambda _: CodeInterpreterTool(), + 'image_generation': lambda _: ImageGenerationTool(), + } + except ImportError: + return {} + + +_NATIVE_OPENAI_TOOLS = _build_native_tool_map() + + class OpenAIAgentGraphRunner(AgentGraphRunner): """ AgentGraphRunner implementation for the OpenAI Agents SDK. - Builds agents from an AgentGraphDefinition and a ToolRegistry via - reverse_traverse, executes them with Runner.run(), and auto-tracks - path, tool calls, handoffs, latency, and invocation success/failure - via the graph's AIGraphTracker. + Runs the agent graph with the OpenAI Agents SDK and automatically records + graph- and node-level AI metric data to the LaunchDarkly trackers on the + graph definition and each node. Requires ``openai-agents`` to be installed. """ @@ -54,32 +87,41 @@ async def run(self, input: Any) -> AgentGraphResult: if root_node: path.append(root_node.get_key()) - start_time = time.time() + start_ns = time.perf_counter_ns() try: - try: - from agents import Runner - except ImportError as exc: - raise ImportError( - "openai-agents is required for OpenAIAgentGraphRunner. " - "Install it with: pip install openai-agents" - ) from exc - + from agents import Runner root_agent = self._build_agents(path) result = await Runner.run(root_agent, str(input)) - duration = int((time.time() - start_time) * 1000) + # _log_run_result_shape(result) + duration = (time.perf_counter_ns() - start_ns) // 1_000_000 if tracker: tracker.track_path(path) tracker.track_latency(duration) tracker.track_invocation_success() + usage = result.context_wrapper.usage + tracker.track_total_tokens( + TokenUsage( + total=usage.total_tokens, + input=usage.input_tokens, + output=usage.output_tokens, + ) + ) return AgentGraphResult( output=str(result.final_output), raw=result, metrics=LDAIMetrics(success=True), ) - except Exception: - duration = int((time.time() - start_time) * 1000) + except Exception as exc: + if isinstance(exc, ImportError): + log.warning( + "openai-agents is required for OpenAIAgentGraphRunner. " + "Install it with: pip install openai-agents" + ) + else: + log.warning(f'OpenAIAgentGraphRunner run failed: {exc}') + duration = (time.perf_counter_ns() - start_ns) // 1_000_000 if tracker: tracker.track_latency(duration) tracker.track_invocation_failure() @@ -89,6 +131,54 @@ async def run(self, input: Any) -> AgentGraphResult: metrics=LDAIMetrics(success=False), ) + def _handle_handoff( + self, + run_ctx: Any, + src: str, + tgt: str, + path: List[str], + tracker: Any, + config_tracker: Any, + ) -> None: + path.append(tgt) + if tracker: + tracker.track_handoff_success(src, tgt) + + usage: Optional[TokenUsage] = None + duration_ms: Optional[int] = None + try: + usage_entry = run_ctx.usage.request_usage_entries[-1] + usage = TokenUsage( + total=usage_entry.total_tokens, + input=usage_entry.input_tokens, + output=usage_entry.output_tokens, + ) + duration_ms = getattr(usage_entry, 'duration_ms', None) + if duration_ms is None: + duration_ms = getattr(usage_entry, 'latency_ms', None) + 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) + if duration_ms is not None: + config_tracker.track_duration(int(duration_ms), graph_key=gk) + config_tracker.track_success(graph_key=gk) + + def _make_on_handoff( + self, + src: str, + tgt: str, + path: List[str], + tracker: Any, + config_tracker: Any, + ): + def on_handoff(run_ctx: Any) -> None: + self._handle_handoff(run_ctx, src, tgt, path, tracker, config_tracker) + return on_handoff + def _build_agents(self, path: List[str]) -> Any: """ Build the agent tree from the graph definition via reverse_traverse. @@ -125,32 +215,16 @@ def build_node(node: AgentGraphNode, ctx: dict) -> Any: agent_handoffs: List[Handoff] = [] for edge in node.get_edges(): target_key = edge.target_config - - def _make_on_handoff(src: str, tgt: str): - def on_handoff(run_ctx: RunContextWrapper) -> None: - path.append(tgt) - if tracker: - tracker.track_handoff_success(src, tgt) - tracker.track_node_invocation(src) - if config_tracker: - try: - usage_entry = run_ctx.usage.request_usage_entries[-1] - config_tracker.track_tokens( - TokenUsage( - total=usage_entry.total_tokens, - input=usage_entry.input_tokens, - output=usage_entry.output_tokens, - ) - ) - except Exception: - pass - config_tracker.track_success() - return on_handoff - agent_handoffs.append( handoff( agent=ctx[target_key], - on_handoff=_make_on_handoff(node_config.key, target_key), + on_handoff=self._make_on_handoff( + node_config.key, + target_key, + path, + tracker, + config_tracker, + ), ) ) @@ -159,6 +233,12 @@ def on_handoff(run_ctx: RunContextWrapper) -> None: for tool_def in tool_defs: tool_name_raw = tool_def.get('name', '') tool_name = _to_openai_name(tool_name_raw) + + # Check native OpenAI tools first, then fall back to ToolRegistry + if tool_name in _NATIVE_OPENAI_TOOLS: + agent_tools.append(_NATIVE_OPENAI_TOOLS[tool_name](tool_def)) + continue + tool_fn = self._tools.get(tool_name) or self._tools.get(tool_name_raw) if not tool_fn: continue @@ -169,7 +249,6 @@ def _make_tool( fn: Any, description: str, params_schema: dict, - cfg_key: str, ) -> FunctionTool: def wrapped(tool_ctx: ToolContext, tool_args: str) -> Any: import json @@ -178,8 +257,11 @@ def wrapped(tool_ctx: ToolContext, tool_args: str) -> Any: except Exception: args = {} path.append(raw_name) - if tracker: - tracker.track_tool_call(config_key=cfg_key, tool_key=name) + if config_tracker is not None: + config_tracker.track_tool_call( + name, + graph_key=tracker.graph_key if tracker is not None else None, + ) return fn(**args) return FunctionTool( @@ -196,13 +278,12 @@ def wrapped(tool_ctx: ToolContext, tool_args: str) -> Any: tool_fn, tool_def.get('description', ''), tool_def.get('parameters', {}), - node_config.key, ) ) return Agent( name=_to_openai_name(node_config.key), - instructions=f'[RECOMMENDED_PROMPT_PREFIX] {node_config.instructions or ""}', + instructions=f'{RECOMMENDED_PROMPT_PREFIX} {node_config.instructions or ""}', handoffs=list(agent_handoffs), tools=list(agent_tools), ) diff --git a/packages/ai-providers/server-ai-openai/tests/test_openai_agent_graph_runner.py b/packages/ai-providers/server-ai-openai/tests/test_openai_agent_graph_runner.py index bc6fb2c..9cc0a67 100644 --- a/packages/ai-providers/server-ai-openai/tests/test_openai_agent_graph_runner.py +++ b/packages/ai-providers/server-ai-openai/tests/test_openai_agent_graph_runner.py @@ -99,6 +99,9 @@ async def test_openai_agent_graph_runner_run_success(): mock_result = MagicMock() mock_result.final_output = "agent answer" + mock_result.context_wrapper.usage.total_tokens = 0 + mock_result.context_wrapper.usage.input_tokens = 0 + mock_result.context_wrapper.usage.output_tokens = 0 mock_runner_module = MagicMock() mock_runner_module.run = AsyncMock(return_value=mock_result) From b32f5626a015ad757356979d9b6c865360772969 Mon Sep 17 00:00:00 2001 From: jsonbailey Date: Thu, 26 Mar 2026 14:46:53 -0500 Subject: [PATCH 03/10] fix: resolve lint errors from rebase onto main Co-Authored-By: Claude Sonnet 4.6 --- .../langchain_runner_factory.py | 4 +++- .../langgraph_agent_graph_runner.py | 20 +++++++++++-------- .../test_langgraph_agent_graph_runner.py | 2 +- .../ldai_openai/openai_agent_graph_runner.py | 19 ++++++++++++++---- .../src/ldai_openai/openai_runner_factory.py | 3 +-- .../tests/test_openai_agent_graph_runner.py | 2 +- packages/sdk/server-ai/src/ldai/client.py | 4 ++-- .../server-ai/src/ldai/managed_agent_graph.py | 3 +-- .../tests/test_managed_agent_graph.py | 11 +++++----- 9 files changed, 41 insertions(+), 27 deletions(-) diff --git a/packages/ai-providers/server-ai-langchain/src/ldai_langchain/langchain_runner_factory.py b/packages/ai-providers/server-ai-langchain/src/ldai_langchain/langchain_runner_factory.py index f52923b..43febb5 100644 --- a/packages/ai-providers/server-ai-langchain/src/ldai_langchain/langchain_runner_factory.py +++ b/packages/ai-providers/server-ai-langchain/src/ldai_langchain/langchain_runner_factory.py @@ -18,7 +18,9 @@ def create_agent_graph(self, graph_def: Any, tools: ToolRegistry) -> Any: :param tools: Registry mapping tool names to callables (langchain-compatible) :return: LangGraphAgentGraphRunner ready to execute the graph """ - from ldai_langchain.langgraph_agent_graph_runner import LangGraphAgentGraphRunner + from ldai_langchain.langgraph_agent_graph_runner import ( + LangGraphAgentGraphRunner, + ) return LangGraphAgentGraphRunner(graph_def, tools) def create_model(self, config: AIConfigKind) -> LangChainModelRunner: 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 5388d8f..123bebe 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 @@ -6,11 +6,15 @@ from ldai import log from ldai.agent_graph import AgentGraphDefinition, AgentGraphNode +from ldai.providers import AgentGraphResult, AgentGraphRunner, ToolRegistry from ldai.providers.types import LDAIMetrics -from ldai.runners.agent_graph_runner import AgentGraphRunner -from ldai.runners.types import AgentGraphResult, ToolRegistry -from ldai_langchain.langchain_helper import LangChainHelper +from ldai_langchain.langchain_helper import ( + get_ai_metrics_from_response, + get_ai_usage_from_response, + get_tool_calls_from_response, + sum_token_usage_from_messages, +) class LangGraphAgentGraphRunner(AgentGraphRunner): @@ -53,7 +57,7 @@ async def run(self, input: Any) -> AgentGraphResult: from typing_extensions import TypedDict class WorkflowState(TypedDict): - messages: Annotated[List[AnyMessage], operator.add] + messages: Annotated[List[Any], operator.add] agent_builder: StateGraph = StateGraph(WorkflowState) root_node = self._graph.root() @@ -87,11 +91,11 @@ def invoke(state: WorkflowState) -> WorkflowState: if node_tracker: response = node_tracker.track_metrics_of( lambda: model.invoke(state['messages']), - LangChainHelper.get_ai_metrics_from_response, + get_ai_metrics_from_response, graph_key=gk, ) node_tracker.track_tool_calls( - LangChainHelper.get_tool_calls_from_response(response), + get_tool_calls_from_response(response), graph_key=tracker.graph_key if tracker is not None else None, ) else: @@ -101,7 +105,7 @@ def invoke(state: WorkflowState) -> WorkflowState: invoke.__name__ = node_key - agent_builder.add_node(name=node_key, node=invoke) + agent_builder.add_node(node_key, invoke) if node_key == root_key: agent_builder.add_edge(START, node_key) @@ -134,7 +138,7 @@ def invoke(state: WorkflowState) -> WorkflowState: tracker.track_latency(duration) tracker.track_invocation_success() tracker.track_total_tokens( - LangChainHelper.sum_token_usage_from_messages(messages) + sum_token_usage_from_messages(messages) ) return AgentGraphResult( diff --git a/packages/ai-providers/server-ai-langchain/tests/test_langgraph_agent_graph_runner.py b/packages/ai-providers/server-ai-langchain/tests/test_langgraph_agent_graph_runner.py index 67229d5..de5e8d9 100644 --- a/packages/ai-providers/server-ai-langchain/tests/test_langgraph_agent_graph_runner.py +++ b/packages/ai-providers/server-ai-langchain/tests/test_langgraph_agent_graph_runner.py @@ -5,7 +5,7 @@ from ldai.agent_graph import AgentGraphDefinition from ldai.models import AIAgentGraphConfig, AIAgentConfig, ModelConfig, ProviderConfig -from ldai.runners.types import AgentGraphResult, ToolRegistry +from ldai.providers import AgentGraphResult, ToolRegistry from ldai_langchain.langgraph_agent_graph_runner import LangGraphAgentGraphRunner from ldai_langchain.langchain_runner_factory import LangChainRunnerFactory 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 afebbcb..4247310 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 @@ -5,9 +5,8 @@ from ldai import log from ldai.agent_graph import AgentGraphDefinition, AgentGraphNode +from ldai.providers import AgentGraphResult, AgentGraphRunner, ToolRegistry from ldai.providers.types import LDAIMetrics -from ldai.runners.agent_graph_runner import AgentGraphRunner -from ldai.runners.types import AgentGraphResult, ToolRegistry from ldai.tracker import TokenUsage @@ -35,7 +34,12 @@ def _log_run_result_shape(result: Any) -> None: def _build_native_tool_map() -> dict: try: - from agents import CodeInterpreterTool, FileSearchTool, ImageGenerationTool, WebSearchTool + from agents import ( + CodeInterpreterTool, + FileSearchTool, + ImageGenerationTool, + WebSearchTool, + ) return { 'web_search_tool': lambda _: WebSearchTool(), 'file_search_tool': lambda _: FileSearchTool(), @@ -190,7 +194,14 @@ def _build_agents(self, path: List[str]) -> Any: :return: The root Agent instance """ try: - from agents import Agent, FunctionTool, Handoff, RunContextWrapper, Tool, handoff + from agents import ( + Agent, + FunctionTool, + Handoff, + RunContextWrapper, + Tool, + handoff, + ) from agents.extensions.handoff_prompt import RECOMMENDED_PROMPT_PREFIX from agents.tool_context import ToolContext except ImportError as exc: diff --git a/packages/ai-providers/server-ai-openai/src/ldai_openai/openai_runner_factory.py b/packages/ai-providers/server-ai-openai/src/ldai_openai/openai_runner_factory.py index c52b5a6..ae4094c 100644 --- a/packages/ai-providers/server-ai-openai/src/ldai_openai/openai_runner_factory.py +++ b/packages/ai-providers/server-ai-openai/src/ldai_openai/openai_runner_factory.py @@ -2,8 +2,7 @@ from typing import Any, Optional from ldai.models import AIConfigKind -from ldai.providers import AIProvider -from ldai.runners.types import ToolRegistry +from ldai.providers import AIProvider, ToolRegistry from openai import AsyncOpenAI from ldai_openai.openai_model_runner import OpenAIModelRunner diff --git a/packages/ai-providers/server-ai-openai/tests/test_openai_agent_graph_runner.py b/packages/ai-providers/server-ai-openai/tests/test_openai_agent_graph_runner.py index 9cc0a67..960e353 100644 --- a/packages/ai-providers/server-ai-openai/tests/test_openai_agent_graph_runner.py +++ b/packages/ai-providers/server-ai-openai/tests/test_openai_agent_graph_runner.py @@ -5,7 +5,7 @@ from ldai.agent_graph import AgentGraphDefinition from ldai.models import AIAgentGraphConfig, AIAgentConfig, Edge, ModelConfig, ProviderConfig -from ldai.runners.types import AgentGraphResult, ToolRegistry +from ldai.providers import AgentGraphResult, ToolRegistry from ldai_openai.openai_agent_graph_runner import OpenAIAgentGraphRunner from ldai_openai.openai_runner_factory import OpenAIRunnerFactory diff --git a/packages/sdk/server-ai/src/ldai/client.py b/packages/sdk/server-ai/src/ldai/client.py index 1c72ec4..6aef7c6 100644 --- a/packages/sdk/server-ai/src/ldai/client.py +++ b/packages/sdk/server-ai/src/ldai/client.py @@ -25,8 +25,8 @@ ModelConfig, ProviderConfig, ) +from ldai.providers import ToolRegistry from ldai.providers.runner_factory import RunnerFactory -from ldai.runners.types import ToolRegistry from ldai.sdk_info import AI_SDK_LANGUAGE, AI_SDK_NAME, AI_SDK_VERSION from ldai.tracker import AIGraphTracker, LDAIConfigTracker @@ -653,7 +653,7 @@ async def create_agent_graph( if not graph.enabled: return None - runner = await RunnerFactory.create_agent_graph( + runner = RunnerFactory.create_agent_graph( graph, tools or {}, default_ai_provider ) if not runner: diff --git a/packages/sdk/server-ai/src/ldai/managed_agent_graph.py b/packages/sdk/server-ai/src/ldai/managed_agent_graph.py index 6eac9d3..bb04add 100644 --- a/packages/sdk/server-ai/src/ldai/managed_agent_graph.py +++ b/packages/sdk/server-ai/src/ldai/managed_agent_graph.py @@ -2,8 +2,7 @@ from typing import Any, Optional -from ldai.runners.agent_graph_runner import AgentGraphRunner -from ldai.runners.types import AgentGraphResult +from ldai.providers import AgentGraphResult, AgentGraphRunner from ldai.tracker import AIGraphTracker diff --git a/packages/sdk/server-ai/tests/test_managed_agent_graph.py b/packages/sdk/server-ai/tests/test_managed_agent_graph.py index 8ffa73b..476ac02 100644 --- a/packages/sdk/server-ai/tests/test_managed_agent_graph.py +++ b/packages/sdk/server-ai/tests/test_managed_agent_graph.py @@ -7,8 +7,7 @@ from ldai import LDAIClient, ManagedAgentGraph from ldai.providers.types import LDAIMetrics -from ldai.runners.agent_graph_runner import AgentGraphRunner -from ldai.runners.types import AgentGraphResult, ToolRegistry +from ldai.providers import AgentGraphResult, AgentGraphRunner, ToolRegistry from ldai.tracker import AIGraphTracker @@ -125,7 +124,7 @@ async def test_create_agent_graph_returns_managed_agent_graph(ldai_client: LDAIC with patch( 'ldai.providers.runner_factory.RunnerFactory.create_agent_graph', - new=AsyncMock(return_value=stub_runner), + new=MagicMock(return_value=stub_runner), ): managed = await ldai_client.create_agent_graph('travel-graph', context) @@ -147,7 +146,7 @@ async def test_create_agent_graph_returns_none_when_runner_factory_fails(ldai_cl with patch( 'ldai.providers.runner_factory.RunnerFactory.create_agent_graph', - new=AsyncMock(return_value=None), + new=MagicMock(return_value=None), ): managed = await ldai_client.create_agent_graph('travel-graph', context) @@ -160,7 +159,7 @@ async def test_create_agent_graph_passes_tools_to_factory(ldai_client: LDAIClien tools: ToolRegistry = {'search': lambda q: f'results for {q}'} captured = {} - async def fake_create_agent_graph(graph_def, tools_arg, default_ai_provider=None): + def fake_create_agent_graph(graph_def, tools_arg, default_ai_provider=None): captured['tools'] = tools_arg return StubAgentGraphRunner() @@ -179,7 +178,7 @@ async def test_create_agent_graph_run_produces_result(ldai_client: LDAIClient): with patch( 'ldai.providers.runner_factory.RunnerFactory.create_agent_graph', - new=AsyncMock(return_value=StubAgentGraphRunner("final answer")), + new=MagicMock(return_value=StubAgentGraphRunner("final answer")), ): managed = await ldai_client.create_agent_graph('travel-graph', context) From e23119bb2943f84ceb05e2c8b84211754ba5ff64 Mon Sep 17 00:00:00 2001 From: jsonbailey Date: Fri, 27 Mar 2026 08:52:14 -0500 Subject: [PATCH 04/10] properly track durations on nodes --- .../ldai_openai/openai_agent_graph_runner.py | 60 ++++++++++++++++--- .../tests/test_openai_agent_graph_runner.py | 5 ++ 2 files changed, 58 insertions(+), 7 deletions(-) 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 4247310..1122538 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 @@ -92,11 +92,15 @@ async def run(self, input: Any) -> AgentGraphResult: path.append(root_node.get_key()) start_ns = time.perf_counter_ns() + # Mutable cell so handoff callbacks can update time-between-handoffs without globals. + last_handoff_ns: List[int] = [start_ns] try: from agents import Runner - root_agent = self._build_agents(path) + root_agent = self._build_agents(path, last_handoff_ns) result = await Runner.run(root_agent, str(input)) # _log_run_result_shape(result) + self._flush_final_segment(path, last_handoff_ns, tracker, result) + duration = (time.perf_counter_ns() - start_ns) // 1_000_000 if tracker: @@ -135,6 +139,44 @@ async def run(self, input: Any) -> AgentGraphResult: metrics=LDAIMetrics(success=False), ) + def _flush_final_segment( + self, + path: List[str], + last_handoff_ns: List[int], + tracker: Any, + result: Any, + ) -> None: + """Record duration/tokens for the last active agent (no handoff after it).""" + if not path: + return + last_key = path[-1] + node = self._graph.get_node(last_key) + if node is None: + return + config_tracker = node.get_config().tracker + if config_tracker is None: + return + + now_ns = time.perf_counter_ns() + duration_ms = (now_ns - last_handoff_ns[0]) // 1_000_000 + + usage: Optional[TokenUsage] = None + try: + usage_entry = result.context_wrapper.usage.request_usage_entries[-1] + usage = TokenUsage( + total=usage_entry.total_tokens, + input=usage_entry.input_tokens, + output=usage_entry.output_tokens, + ) + 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) + def _handle_handoff( self, run_ctx: Any, @@ -143,13 +185,16 @@ def _handle_handoff( path: List[str], tracker: Any, config_tracker: Any, + last_handoff_ns: List[int], ) -> None: path.append(tgt) if tracker: tracker.track_handoff_success(src, tgt) usage: Optional[TokenUsage] = None - duration_ms: Optional[int] = None + now_ns = time.perf_counter_ns() + duration_ms = (now_ns - last_handoff_ns[0]) // 1_000_000 + last_handoff_ns[0] = now_ns try: usage_entry = run_ctx.usage.request_usage_entries[-1] usage = TokenUsage( @@ -157,9 +202,6 @@ def _handle_handoff( input=usage_entry.input_tokens, output=usage_entry.output_tokens, ) - duration_ms = getattr(usage_entry, 'duration_ms', None) - if duration_ms is None: - duration_ms = getattr(usage_entry, 'latency_ms', None) except Exception: pass @@ -178,12 +220,15 @@ def _make_on_handoff( path: List[str], tracker: Any, config_tracker: Any, + last_handoff_ns: List[int], ): def on_handoff(run_ctx: Any) -> None: - self._handle_handoff(run_ctx, src, tgt, path, tracker, config_tracker) + self._handle_handoff( + run_ctx, src, tgt, path, tracker, config_tracker, last_handoff_ns + ) return on_handoff - def _build_agents(self, path: List[str]) -> Any: + def _build_agents(self, path: List[str], last_handoff_ns: List[int]) -> Any: """ Build the agent tree from the graph definition via reverse_traverse. @@ -235,6 +280,7 @@ def build_node(node: AgentGraphNode, ctx: dict) -> Any: path, tracker, config_tracker, + last_handoff_ns, ), ) ) diff --git a/packages/ai-providers/server-ai-openai/tests/test_openai_agent_graph_runner.py b/packages/ai-providers/server-ai-openai/tests/test_openai_agent_graph_runner.py index 960e353..56dd0c1 100644 --- a/packages/ai-providers/server-ai-openai/tests/test_openai_agent_graph_runner.py +++ b/packages/ai-providers/server-ai-openai/tests/test_openai_agent_graph_runner.py @@ -135,3 +135,8 @@ async def test_openai_agent_graph_runner_run_success(): tracker.track_invocation_success.assert_called_once() tracker.track_path.assert_called_once() tracker.track_latency.assert_called_once() + + root_tracker = graph.get_node('root-agent').get_config().tracker + root_tracker.track_duration.assert_called_once() + root_tracker.track_tokens.assert_called_once() + root_tracker.track_success.assert_called_once() From f3c0131299901490e727bbedee7fa16ef9bf16d3 Mon Sep 17 00:00:00 2001 From: jsonbailey Date: Fri, 27 Mar 2026 08:58:15 -0500 Subject: [PATCH 05/10] address code review feedback --- .../src/ldai_langchain/langgraph_agent_graph_runner.py | 6 +++--- 1 file changed, 3 insertions(+), 3 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 123bebe..9def295 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 @@ -10,6 +10,7 @@ from ldai.providers.types import LDAIMetrics from ldai_langchain.langchain_helper import ( + create_langchain_model, get_ai_metrics_from_response, get_ai_usage_from_response, get_tool_calls_from_response, @@ -51,7 +52,6 @@ async def run(self, input: Any) -> AgentGraphResult: tracker = self._graph.get_tracker() start_ns = time.perf_counter_ns() try: - from langchain.chat_models import init_chat_model from langchain_core.messages import AnyMessage, HumanMessage from langgraph.graph import END, START, StateGraph from typing_extensions import TypedDict @@ -72,7 +72,7 @@ def handle_traversal(node: AgentGraphNode, ctx: dict) -> None: model = None if node_config.model: - lc_model = init_chat_model(model=node_config.model.name) + lc_model = create_langchain_model(node_config) tool_defs = node_config.model.get_parameter('tools') or [] tool_fns = [ tools_ref[t.get('name', '')] @@ -86,7 +86,7 @@ def handle_traversal(node: AgentGraphNode, ctx: dict) -> None: def invoke(state: WorkflowState) -> WorkflowState: exec_path.append(node_key) if not model: - return state + return {'messages': []} gk = tracker.graph_key if tracker is not None else None if node_tracker: response = node_tracker.track_metrics_of( From c2c05379a425983bbac37ab3c39cb0e6b87040c2 Mon Sep 17 00:00:00 2001 From: jsonbailey Date: Fri, 27 Mar 2026 09:00:03 -0500 Subject: [PATCH 06/10] remove debugging code --- .../ldai_openai/openai_agent_graph_runner.py | 19 ------------------- 1 file changed, 19 deletions(-) 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 1122538..09b5dfd 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 @@ -14,24 +14,6 @@ def _to_openai_name(name: str) -> str: """Convert a hyphenated tool/node name to an underscore-separated OpenAI function name.""" return name.replace('-', '_') - -def _log_run_result_shape(result: Any) -> None: - """Print RunResult attributes (excluding final_output) for debugging.""" - attrs = [a for a in dir(result) if not a.startswith('_')] - print("RunResult public attributes:", attrs) - for name in attrs: - if name == 'final_output': - continue - try: - val = getattr(result, name) - if callable(val): - print(f" {name}: callable") - else: - print(f" {name}: {repr(val)}") - except Exception as e: - print(f" {name}: (error reading: {e})") - - def _build_native_tool_map() -> dict: try: from agents import ( @@ -98,7 +80,6 @@ async def run(self, input: Any) -> AgentGraphResult: from agents import Runner root_agent = self._build_agents(path, last_handoff_ns) result = await Runner.run(root_agent, str(input)) - # _log_run_result_shape(result) self._flush_final_segment(path, last_handoff_ns, tracker, result) duration = (time.perf_counter_ns() - start_ns) // 1_000_000 From fbb03e334658eb8354da9478e70428b5cf18deaa Mon Sep 17 00:00:00 2001 From: jsonbailey Date: Fri, 27 Mar 2026 09:12:36 -0500 Subject: [PATCH 07/10] fix lint issues --- .../src/ldai_langchain/langgraph_agent_graph_runner.py | 4 +--- .../src/ldai_openai/openai_agent_graph_runner.py | 1 + 2 files changed, 2 insertions(+), 3 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 9def295..5d1e182 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 @@ -79,9 +79,7 @@ def handle_traversal(node: AgentGraphNode, ctx: dict) -> None: for t in tool_defs if t.get('name', '') in tools_ref ] - if tool_fns: - lc_model = lc_model.bind_tools(tool_fns) - model = lc_model + model = lc_model.bind_tools(tool_fns) if tool_fns else lc_model def invoke(state: WorkflowState) -> WorkflowState: exec_path.append(node_key) 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 09b5dfd..571a7fb 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 @@ -14,6 +14,7 @@ def _to_openai_name(name: str) -> str: """Convert a hyphenated tool/node name to an underscore-separated OpenAI function name.""" return name.replace('-', '_') + def _build_native_tool_map() -> dict: try: from agents import ( From 131f30e242dfa40197fc78ee7bc45946a0bd856b Mon Sep 17 00:00:00 2001 From: jsonbailey Date: Fri, 27 Mar 2026 09:19:17 -0500 Subject: [PATCH 08/10] use run state for accurate tracking --- .../ldai_openai/openai_agent_graph_runner.py | 46 +++++++++++-------- 1 file changed, 27 insertions(+), 19 deletions(-) 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 571a7fb..b16f881 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 @@ -36,6 +36,14 @@ def _build_native_tool_map() -> dict: _NATIVE_OPENAI_TOOLS = _build_native_tool_map() +class _RunState: + """Mutable state shared across handoff and tool callbacks during a single run.""" + + def __init__(self, last_handoff_ns: int, last_node_key: str) -> None: + self.last_handoff_ns = last_handoff_ns + self.last_node_key = last_node_key + + class OpenAIAgentGraphRunner(AgentGraphRunner): """ AgentGraphRunner implementation for the OpenAI Agents SDK. @@ -71,17 +79,17 @@ async def run(self, input: Any) -> AgentGraphResult: tracker = self._graph.get_tracker() path: List[str] = [] root_node = self._graph.root() - if root_node: - path.append(root_node.get_key()) + root_key = root_node.get_key() if root_node else '' + if root_key: + path.append(root_key) start_ns = time.perf_counter_ns() - # Mutable cell so handoff callbacks can update time-between-handoffs without globals. - last_handoff_ns: List[int] = [start_ns] + state = _RunState(last_handoff_ns=start_ns, last_node_key=root_key) try: from agents import Runner - root_agent = self._build_agents(path, last_handoff_ns) + root_agent = self._build_agents(path, state) result = await Runner.run(root_agent, str(input)) - self._flush_final_segment(path, last_handoff_ns, tracker, result) + self._flush_final_segment(state, tracker, result) duration = (time.perf_counter_ns() - start_ns) // 1_000_000 @@ -123,16 +131,14 @@ async def run(self, input: Any) -> AgentGraphResult: def _flush_final_segment( self, - path: List[str], - last_handoff_ns: List[int], + state: _RunState, tracker: Any, result: Any, ) -> None: """Record duration/tokens for the last active agent (no handoff after it).""" - if not path: + if not state.last_node_key: return - last_key = path[-1] - node = self._graph.get_node(last_key) + node = self._graph.get_node(state.last_node_key) if node is None: return config_tracker = node.get_config().tracker @@ -140,7 +146,7 @@ def _flush_final_segment( return now_ns = time.perf_counter_ns() - duration_ms = (now_ns - last_handoff_ns[0]) // 1_000_000 + duration_ms = (now_ns - state.last_handoff_ns) // 1_000_000 usage: Optional[TokenUsage] = None try: @@ -167,16 +173,17 @@ def _handle_handoff( path: List[str], tracker: Any, config_tracker: Any, - last_handoff_ns: List[int], + state: _RunState, ) -> None: path.append(tgt) + state.last_node_key = tgt if tracker: tracker.track_handoff_success(src, tgt) usage: Optional[TokenUsage] = None now_ns = time.perf_counter_ns() - duration_ms = (now_ns - last_handoff_ns[0]) // 1_000_000 - last_handoff_ns[0] = now_ns + duration_ms = (now_ns - state.last_handoff_ns) // 1_000_000 + state.last_handoff_ns = now_ns try: usage_entry = run_ctx.usage.request_usage_entries[-1] usage = TokenUsage( @@ -202,15 +209,15 @@ def _make_on_handoff( path: List[str], tracker: Any, config_tracker: Any, - last_handoff_ns: List[int], + state: _RunState, ): def on_handoff(run_ctx: Any) -> None: self._handle_handoff( - run_ctx, src, tgt, path, tracker, config_tracker, last_handoff_ns + run_ctx, src, tgt, path, tracker, config_tracker, state ) return on_handoff - def _build_agents(self, path: List[str], last_handoff_ns: List[int]) -> Any: + def _build_agents(self, path: List[str], state: _RunState) -> Any: """ Build the agent tree from the graph definition via reverse_traverse. @@ -218,6 +225,7 @@ def _build_agents(self, path: List[str], last_handoff_ns: List[int]) -> Any: targets exist before the agents that hand off to them. :param path: Mutable list to accumulate the execution path + :param state: Shared run state for tracking handoff timing and last node :return: The root Agent instance """ try: @@ -262,7 +270,7 @@ def build_node(node: AgentGraphNode, ctx: dict) -> Any: path, tracker, config_tracker, - last_handoff_ns, + state, ), ) ) From 37f34fa25cd30e4bbb7090967404d19c92dd15eb Mon Sep 17 00:00:00 2001 From: jsonbailey Date: Fri, 27 Mar 2026 15:33:58 -0500 Subject: [PATCH 09/10] fix lint issue --- .../src/ldai_langchain/langgraph_agent_graph_runner.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) 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 5d1e182..c0c0b5c 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 @@ -119,7 +119,7 @@ def invoke(state: WorkflowState) -> WorkflowState: self._graph.traverse(fn=handle_traversal) compiled = agent_builder.compile() - result = await compiled.ainvoke( + result = await compiled.ainvoke( # type: ignore[call-overload] {'messages': [HumanMessage(content=str(input))]} ) duration = (time.perf_counter_ns() - start_ns) // 1_000_000 From f1836a4440fd1da6a0933be48a6aec349343d221 Mon Sep 17 00:00:00 2001 From: jsonbailey Date: Fri, 27 Mar 2026 16:21:12 -0500 Subject: [PATCH 10/10] catch failures when reading token counts --- .../ldai_openai/openai_agent_graph_runner.py | 17 ++++++++++------- 1 file changed, 10 insertions(+), 7 deletions(-) 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 b16f881..df2acf6 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 @@ -97,14 +97,17 @@ async def run(self, input: Any) -> AgentGraphResult: tracker.track_path(path) tracker.track_latency(duration) tracker.track_invocation_success() - usage = result.context_wrapper.usage - tracker.track_total_tokens( - TokenUsage( - total=usage.total_tokens, - input=usage.input_tokens, - output=usage.output_tokens, + try: + usage = result.context_wrapper.usage + tracker.track_total_tokens( + TokenUsage( + total=usage.total_tokens, + input=usage.input_tokens, + output=usage.output_tokens, + ) ) - ) + except Exception: + pass return AgentGraphResult( output=str(result.final_output),