Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension


Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
1,178 changes: 1,178 additions & 0 deletions packages/sdk/server-ai/poetry.lock

Large diffs are not rendered by default.

2 changes: 1 addition & 1 deletion packages/sdk/server-ai/src/ldai/client.py
Original file line number Diff line number Diff line change
Expand Up @@ -799,7 +799,7 @@ async def create_agent_graph(
if not runner:
return None

return ManagedAgentGraph(runner)
return ManagedAgentGraph(graph, runner)

def agents(
self,
Expand Down
80 changes: 60 additions & 20 deletions packages/sdk/server-ai/src/ldai/managed_agent_graph.py
Original file line number Diff line number Diff line change
@@ -1,61 +1,101 @@
"""ManagedAgentGraph — LaunchDarkly managed wrapper for agent graph execution."""

from typing import Any
from typing import Any, Dict

from ldai.providers import AgentGraphResult, AgentGraphRunner
from ldai.providers.types import GraphMetricSummary, ManagedGraphResult
from ldai.agent_graph import AgentGraphDefinition
from ldai.providers import AgentGraphRunner
from ldai.providers.types import (
AgentGraphRunnerResult,
LDAIMetrics,
ManagedGraphResult,
)
from ldai.tracker import LDAIMetricSummary


class ManagedAgentGraph:
"""
LaunchDarkly managed wrapper for AI agent graph execution.

Holds an AgentGraphRunner. Wraps the runner result in a
:class:`~ldai.providers.types.ManagedGraphResult` and builds a
:class:`~ldai.providers.types.GraphMetricSummary` from the runner's metrics.
Holds an AgentGraphRunner and an AgentGraphDefinition. Delegates execution
to the runner, then drives all graph-level and per-node tracking from the
returned :class:`~ldai.providers.types.AgentGraphRunnerResult`.

Obtain an instance via ``LDAIClient.create_agent_graph()``.
"""

def __init__(
self,
graph: AgentGraphDefinition,
runner: AgentGraphRunner,
):
"""
Initialize ManagedAgentGraph.

:param graph: The AgentGraphDefinition used to drive graph-level and
per-node tracking from the runner result metrics.
:param runner: The AgentGraphRunner to delegate execution to
"""
self._graph = graph
self._runner = runner

async def run(self, input: Any) -> ManagedGraphResult:
"""
Run the agent graph with the given input.

Delegates to the underlying AgentGraphRunner, then drives all
LaunchDarkly tracking from ``result.metrics``.

:param input: The input prompt or structured input for the graph
:return: ManagedGraphResult containing the content, metric summary, raw response,
and an optional evaluations task (currently always ``None`` for graphs —
per-graph evaluations will be added in a future PR).
:return: ManagedGraphResult containing the content, metric summary,
and raw response.
"""
result: AgentGraphResult = await self._runner.run(input)

# Build a GraphMetricSummary from the runner result's LDAIMetrics.
# path and node_metrics will be populated once graph runners are migrated
# to return AgentGraphRunnerResult with GraphMetrics (PR 11).
metrics = result.metrics
summary = GraphMetricSummary(
success=metrics.success,
usage=metrics.usage,
duration_ms=getattr(metrics, 'duration_ms', None),
graph_tracker = self._graph.create_tracker()
result = await graph_tracker.track_graph_metrics_of_async(
lambda r: r.metrics,
lambda: self._runner.run(input),
)

summary = graph_tracker.get_summary()
summary.node_metrics = self._track_node_metrics(result.metrics.node_metrics)

return ManagedGraphResult(
content=result.output,
content=result.content,
metrics=summary,
raw=result.raw,
evaluations=None,
)

def _track_node_metrics(
self, node_metrics: Dict[str, LDAIMetrics]
) -> Dict[str, LDAIMetricSummary]:
"""
Drive per-node LaunchDarkly tracking events and collect node metric summaries.

For each node key present in ``node_metrics``, obtains the node's
config tracker via the graph definition, fires tracking events, and
returns a map of node key to the tracker's metric summary.
"""
node_summaries: Dict[str, LDAIMetricSummary] = {}
for node_key, node_ldai_metrics in node_metrics.items():
node = self._graph.get_node(node_key)
if node is None:
continue
node_tracker = node.get_config().create_tracker()

if node_ldai_metrics.usage is not None:
node_tracker.track_tokens(node_ldai_metrics.usage)
if node_ldai_metrics.duration_ms is not None:
node_tracker.track_duration(node_ldai_metrics.duration_ms)
if node_ldai_metrics.tool_calls:
node_tracker.track_tool_calls(node_ldai_metrics.tool_calls)
if node_ldai_metrics.success:
node_tracker.track_success()
else:
node_tracker.track_error()

node_summaries[node_key] = node_tracker.get_summary()
return node_summaries

def get_agent_graph_runner(self) -> AgentGraphRunner:
"""
Return the underlying AgentGraphRunner for advanced use.
Expand Down
Original file line number Diff line number Diff line change
@@ -1,6 +1,6 @@
from typing import Any, Protocol, runtime_checkable

from ldai.providers.types import AgentGraphResult
from ldai.providers.types import AgentGraphRunnerResult


@runtime_checkable
Expand All @@ -18,11 +18,11 @@ class AgentGraphRunner(Protocol):
the caller just passes input.
"""

async def run(self, input: Any) -> AgentGraphResult:
async def run(self, input: Any) -> AgentGraphRunnerResult:
"""
Run the agent graph with the given input.

:param input: The input to the agent graph (string prompt or structured input)
:return: AgentGraphResult containing the output, raw response, and metrics
:return: AgentGraphRunnerResult containing the content, raw response, and GraphMetrics
"""
...
8 changes: 4 additions & 4 deletions packages/sdk/server-ai/src/ldai/providers/types.py
Original file line number Diff line number Diff line change
Expand Up @@ -110,8 +110,8 @@ class GraphMetrics:
class GraphMetricSummary:
"""Contains a summary of metrics for an agent graph run."""

success: bool
"""Whether the graph run succeeded."""
success: Optional[bool] = None
"""Whether the graph run succeeded. Absent if invocation status has not been tracked."""

path: List[str] = field(default_factory=list)
"""Ordered list of node keys visited during the run."""
Expand All @@ -122,8 +122,8 @@ class GraphMetricSummary:
usage: Optional[TokenUsage] = None
"""Optional aggregate token usage information across all nodes in the graph run."""

node_metrics: Dict[str, LDAIMetrics] = field(default_factory=dict)
"""Per-node metrics keyed by node key."""
node_metrics: Dict[str, LDAIMetricSummary] = field(default_factory=dict)
"""Per-node metric summaries keyed by node key."""

resumption_token: Optional[str] = None
"""Optional resumption token from the graph tracker for cross-process resumption."""
Expand Down
Loading
Loading