from __future__ import annotations
import base64
import json
import time
from dataclasses import dataclass
from enum import Enum
from typing import TYPE_CHECKING, Any, Callable, Dict, Iterable, List, Optional
from ldclient import Context, LDClient, Result
from ldai import log
if TYPE_CHECKING:
from ldai.providers.types import AIGraphMetrics, AIGraphMetricSummary, LDAIMetrics
[docs]
class FeedbackKind(Enum):
"""
Types of feedback that can be provided for AI operations.
"""
Positive = "positive"
Negative = "negative"
[docs]
@dataclass
class TokenUsage:
"""
Tracks token usage for AI operations.
:param total: Total number of tokens used.
:param input: Number of tokens in the prompt.
:param output: Number of tokens in the completion.
"""
total: int
input: int
output: int
[docs]
class LDAIMetricSummary:
"""
Summary of metrics which have been tracked.
"""
[docs]
def __init__(self):
self._duration_ms: Optional[int] = None
self._success: Optional[bool] = None
self._feedback: Optional[Dict[str, FeedbackKind]] = None
self._tokens: Optional[TokenUsage] = None
self._time_to_first_token: Optional[int] = None
self._tool_calls: List[str] = []
self._resumption_token: Optional[str] = None
@property
def duration_ms(self) -> Optional[int]:
"""Duration of the AI operation in milliseconds."""
return self._duration_ms
@property
def success(self) -> Optional[bool]:
return self._success
@property
def feedback(self) -> Optional[Dict[str, FeedbackKind]]:
return self._feedback
@property
def tokens(self) -> Optional[TokenUsage]:
return self._tokens
@property
def time_to_first_token(self) -> Optional[int]:
return self._time_to_first_token
@property
def tool_calls(self) -> List[str]:
"""List of tool keys that were invoked during this operation."""
return self._tool_calls
@property
def resumption_token(self) -> Optional[str]:
"""
URL-safe Base64-encoded resumption token captured at tracker
instantiation. Useful for deferred feedback flows where a downstream
process needs to associate events with the original AI run.
"""
return self._resumption_token
[docs]
class LDAIConfigTracker:
"""
Records metrics for a single AI run.
All events a tracker emits share a runId (a UUIDv4) so LaunchDarkly can correlate
them in metrics views. See individual track methods for their specific semantics.
Call ``create_tracker`` on the AI Config to start a new run. A resumption token
preserves the runId, so events emitted by a tracker reconstructed in another
process correlate with the original run.
"""
[docs]
def __init__(
self,
ld_client: LDClient,
run_id: str,
config_key: str,
variation_key: str,
version: int,
context: Context,
model_name: str,
provider_name: str,
graph_key: Optional[str] = None,
):
"""
Initialize an AI Config tracker.
:param ld_client: LaunchDarkly client instance.
:param run_id: Unique identifier for this AI run.
:param config_key: Configuration key for tracking.
:param variation_key: Variation key for tracking.
:param version: Version of the variation.
:param context: Context for evaluation.
:param model_name: Name of the model used.
:param provider_name: Name of the provider used.
:param graph_key: When set, include ``graphKey`` in all event payloads
(e.g. config-level metrics inside a graph).
"""
self._ld_client = ld_client
self._variation_key = variation_key
self._config_key = config_key
self._version = version
self._model_name = model_name
self._provider_name = provider_name
self._context = context
self._graph_key = graph_key
self._run_id = run_id
self._summary = LDAIMetricSummary()
# Capture resumption_token immediately so it's available on the summary at instantiation.
self._summary._resumption_token = self.resumption_token
@property
def resumption_token(self) -> str:
"""
A URL-safe Base64-encoded JSON string that can be used to reconstruct
a tracker in a different process (e.g. for deferred feedback).
The token contains ``runId``, ``configKey``, ``version``, and
optionally ``variationKey`` and ``graphKey`` (omitted when empty).
``modelName`` and ``providerName`` are **not** included.
"""
data: dict = {
"runId": self._run_id,
"configKey": self._config_key,
}
if self._variation_key:
data["variationKey"] = self._variation_key
data["version"] = self._version
if self._graph_key:
data["graphKey"] = self._graph_key
payload = json.dumps(data)
return base64.urlsafe_b64encode(payload.encode("utf-8")).rstrip(b"=").decode("utf-8")
[docs]
@classmethod
def from_resumption_token(cls, token: str, ld_client: LDClient, context: Context) -> Result:
"""
Reconstruct a tracker from a resumption token.
This is used for cross-process scenarios such as deferred feedback,
where a different service needs to associate tracking events with the
original tracker's ``runId``.
:param token: A URL-safe Base64-encoded resumption token obtained from
:attr:`resumption_token`.
:param ld_client: LaunchDarkly client instance.
:param context: The context to use for track events.
:return: A :class:`Result` whose ``value`` is a new
:class:`LDAIConfigTracker` bound to the original ``runId`` from the
token on success, or whose ``error`` describes the problem on failure.
"""
try:
padded = token + "=" * (-len(token) % 4)
payload = json.loads(
base64.urlsafe_b64decode(padded.encode("utf-8")).decode("utf-8")
)
except Exception as e:
return Result.fail(f"Invalid resumption token: {e}", e)
for field in ("runId", "configKey", "version"):
if field not in payload:
return Result.fail(
f"Invalid resumption token: missing required field '{field}'"
)
return Result.success(cls(
ld_client=ld_client,
run_id=payload["runId"],
config_key=payload["configKey"],
variation_key=payload.get("variationKey") or "",
version=payload["version"],
context=context,
model_name="",
provider_name="",
graph_key=payload.get("graphKey"),
))
def __get_track_data(self) -> dict:
"""
Get tracking data for events.
:return: Dictionary containing variation and config keys.
"""
data = {
"runId": self._run_id,
"configKey": self._config_key,
"version": self._version,
"modelName": self._model_name,
"providerName": self._provider_name,
}
if self._variation_key:
data["variationKey"] = self._variation_key
if self._graph_key:
data['graphKey'] = self._graph_key
return data
[docs]
def track_duration(self, duration: int) -> None:
"""
Manually track the duration of an AI run.
Records at most once per Tracker; further calls are ignored.
:param duration: Duration in milliseconds.
"""
if self._summary.duration_ms is not None:
log.warning(
"Skipping track_duration: duration already recorded on this tracker. "
"Call create_tracker on the AI Config for a new run. %s",
self.__get_track_data(),
)
return
self._summary._duration_ms = duration
self._ld_client.track(
"$ld:ai:duration:total", self._context, self.__get_track_data(), duration
)
[docs]
def track_time_to_first_token(self, time_to_first_token: int) -> None:
"""
Manually track the time to first token of an AI run.
Records at most once per Tracker; further calls are ignored.
:param time_to_first_token: Time to first token in milliseconds.
"""
if self._summary.time_to_first_token is not None:
log.warning(
"Skipping track_time_to_first_token: time-to-first-token already recorded on this tracker. "
"Call create_tracker on the AI Config for a new run. %s",
self.__get_track_data(),
)
return
self._summary._time_to_first_token = time_to_first_token
self._ld_client.track(
"$ld:ai:tokens:ttf",
self._context,
self.__get_track_data(),
time_to_first_token,
)
[docs]
def track_duration_of(self, func):
"""
Automatically track the duration of an AI run.
An exception raised while the function runs will still record the
duration. The exception will be re-thrown.
:param func: Function to track (synchronous only).
:return: Result of the tracked function.
"""
start_ns = time.perf_counter_ns()
try:
result = func()
finally:
duration = (time.perf_counter_ns() - start_ns) // 1_000_000 # duration in milliseconds
self.track_duration(duration)
return result
def _track_from_metrics_extractor(
self,
result: Any,
metrics_extractor: Callable[[Any], Optional[LDAIMetrics]],
elapsed_ms: int,
) -> None:
metrics = None
try:
metrics = metrics_extractor(result)
except Exception as exc:
log.warning("Failed to extract metrics: %s", exc)
if metrics is None:
self.track_duration(elapsed_ms)
return
self.track_duration(metrics.duration_ms if metrics.duration_ms is not None else elapsed_ms)
if metrics.success:
self.track_success()
else:
self.track_error()
if metrics.tokens:
self.track_tokens(metrics.tokens)
if metrics.tool_calls is not None:
self.track_tool_calls(metrics.tool_calls)
[docs]
def track_metrics_of(
self,
metrics_extractor: Callable[[Any], Optional[LDAIMetrics]],
func: Callable[[], Any],
) -> Any:
"""
Track metrics for a synchronous AI operation.
This function will track the duration of the operation, extract metrics using the provided
metrics extractor function, and track success or error status accordingly.
If the provided function throws, then this method will also throw.
In the case the provided function throws, this function will record the duration and an error.
A failed operation will not have any token usage data.
For async operations, use :meth:`track_metrics_of_async`.
When the extracted :class:`~ldai.providers.types.LDAIMetrics` object has a
non-``None`` ``duration_ms`` field, that value is used as the measured duration
instead of the wall-clock elapsed time.
Because each inner metric is at-most-once per Tracker, calling this twice
on the same Tracker will run the inner block again but produce no
additional metric events.
:param metrics_extractor: Function that extracts LDAIMetrics from the operation result
:param func: Synchronous callable that runs the operation
:return: The result of the operation
"""
start_ns = time.perf_counter_ns()
try:
result = func()
except Exception as err:
duration = (time.perf_counter_ns() - start_ns) // 1_000_000
self.track_duration(duration)
self.track_error()
raise err
elapsed_ms = (time.perf_counter_ns() - start_ns) // 1_000_000
self._track_from_metrics_extractor(result, metrics_extractor, elapsed_ms)
return result
[docs]
async def track_metrics_of_async(
self,
metrics_extractor: Callable[[Any], Optional[LDAIMetrics]],
func: Callable[[], Any],
) -> Any:
"""
Track metrics for an async AI operation (``func`` is awaited).
Same event semantics as :meth:`track_metrics_of`.
When the extracted :class:`~ldai.providers.types.LDAIMetrics` object has a
non-``None`` ``duration_ms`` field, that value is used as the measured duration
instead of the wall-clock elapsed time.
Because each inner metric is at-most-once per Tracker, calling this twice
on the same Tracker will run the inner block again but produce no
additional metric events.
:param metrics_extractor: Function that extracts LDAIMetrics from the operation result
:param func: Async callable or zero-arg callable that returns an awaitable when called
:return: The result of the operation
"""
start_ns = time.perf_counter_ns()
result = None
try:
result = await func()
except Exception as err:
duration = (time.perf_counter_ns() - start_ns) // 1_000_000
self.track_duration(duration)
self.track_error()
raise err
elapsed_ms = (time.perf_counter_ns() - start_ns) // 1_000_000
self._track_from_metrics_extractor(result, metrics_extractor, elapsed_ms)
return result
[docs]
def track_judge_result(self, judge_result: Any) -> None:
"""
Track a judge result, including the evaluation score with judge config key.
May be called multiple times per Tracker; each call records the
provided judge result.
:param judge_result: JudgeResult object containing score, metric key, and success status
"""
if not judge_result.sampled:
return
if judge_result.success and judge_result.metric_key:
track_data = self.__get_track_data()
if judge_result.judge_config_key:
track_data = {**track_data, 'judgeConfigKey': judge_result.judge_config_key}
self._ld_client.track(
judge_result.metric_key,
self._context,
track_data,
judge_result.score,
)
[docs]
def track_feedback(self, feedback: Dict[str, FeedbackKind]) -> None:
"""
Track user feedback for an AI run.
Records at most once per Tracker; further calls are ignored.
:param feedback: Dictionary containing feedback kind.
"""
if self._summary.feedback is not None:
log.warning(
"Skipping track_feedback: feedback already recorded on this tracker. "
"Call create_tracker on the AI Config for a new run. %s",
self.__get_track_data(),
)
return
self._summary._feedback = feedback
if feedback["kind"] == FeedbackKind.Positive:
self._ld_client.track(
"$ld:ai:feedback:user:positive",
self._context,
self.__get_track_data(),
1,
)
elif feedback["kind"] == FeedbackKind.Negative:
self._ld_client.track(
"$ld:ai:feedback:user:negative",
self._context,
self.__get_track_data(),
1,
)
[docs]
def track_success(self) -> None:
"""
Track a successful AI generation.
Records at most once per Tracker. track_success and track_error share
state; only one of the two can record per Tracker, and subsequent calls
are ignored.
"""
if self._summary.success is not None:
log.warning(
"Skipping track_success: success/error already recorded on this tracker. "
"Call create_tracker on the AI Config for a new run. %s",
self.__get_track_data(),
)
return
self._summary._success = True
self._ld_client.track(
"$ld:ai:generation:success", self._context, self.__get_track_data(), 1
)
[docs]
def track_error(self) -> None:
"""
Track an unsuccessful AI generation attempt.
Records at most once per Tracker. track_success and track_error share
state; only one of the two can record per Tracker, and subsequent calls
are ignored.
"""
if self._summary.success is not None:
log.warning(
"Skipping track_error: success/error already recorded on this tracker. "
"Call create_tracker on the AI Config for a new run. %s",
self.__get_track_data(),
)
return
self._summary._success = False
self._ld_client.track(
"$ld:ai:generation:error", self._context, self.__get_track_data(), 1
)
[docs]
def track_tokens(self, tokens: TokenUsage) -> None:
"""
Track token usage metrics.
Records at most once per Tracker; further calls are ignored.
:param tokens: Token usage data.
"""
if self._summary.tokens is not None:
log.warning(
"Skipping track_tokens: token usage already recorded on this tracker. "
"Call create_tracker on the AI Config for a new run. %s",
self.__get_track_data(),
)
return
self._summary._tokens = tokens
td = self.__get_track_data()
if tokens.total > 0:
self._ld_client.track(
"$ld:ai:tokens:total",
self._context,
td,
tokens.total,
)
if tokens.input > 0:
self._ld_client.track(
"$ld:ai:tokens:input",
self._context,
td,
tokens.input,
)
if tokens.output > 0:
self._ld_client.track(
"$ld:ai:tokens:output",
self._context,
td,
tokens.output,
)
[docs]
def get_summary(self) -> LDAIMetricSummary:
"""
Get the current summary of AI metrics.
:return: Summary of AI metrics.
"""
return self._summary
[docs]
class AIGraphTracker:
"""
Tracks graph-level metrics for AI agent graph operations.
Maintains an internal :class:`~ldai.providers.types.AIGraphMetricSummary`
that is updated as tracking methods are called. Retrieve it via
:meth:`get_summary`.
"""
[docs]
def __init__(
self,
ld_client: LDClient,
variation_key: str,
graph_key: str,
version: int,
context: Context,
):
"""
Initialize an AI Graph tracker.
:param ld_client: LaunchDarkly client instance.
:param variation_key: Variation key for tracking.
:param graph_key: Graph configuration key for tracking.
:param version: Version of the variation.
:param context: Context for evaluation.
"""
self._ld_client = ld_client
self._variation_key = variation_key
self._graph_key = graph_key
self._version = version
self._context = context
from ldai.providers.types import AIGraphMetricSummary
self._summary = AIGraphMetricSummary()
@property
def graph_key(self) -> str:
"""Graph configuration key used in tracking payloads."""
return self._graph_key
[docs]
def get_summary(self) -> AIGraphMetricSummary:
"""
Get the current summary of graph-level metrics.
:return: Summary of graph metrics tracked so far.
"""
return self._summary
def __get_track_data(self):
"""
Get tracking data for events.
:return: Dictionary containing variation, graph key, and version.
"""
track_data = {
"variationKey": self._variation_key,
"graphKey": self._graph_key,
"version": self._version,
}
return track_data
[docs]
def track_invocation_success(self) -> None:
"""
Track a successful graph run.
Records at most once per graph tracker. track_invocation_success and
track_invocation_failure share state; only one of the two can record
per graph tracker, and subsequent calls are ignored.
"""
if self._summary.success is not None:
log.warning(
"Skipping track_invocation_success: invocation result already recorded on this graph tracker. "
"Call create_tracker on the agent graph for a new run. %s",
self.__get_track_data(),
)
return
self._summary.success = True
self._ld_client.track(
"$ld:ai:graph:invocation_success",
self._context,
self.__get_track_data(),
1,
)
[docs]
def track_invocation_failure(self) -> None:
"""
Track an unsuccessful graph run.
Records at most once per graph tracker. track_invocation_success and
track_invocation_failure share state; only one of the two can record
per graph tracker, and subsequent calls are ignored.
"""
if self._summary.success is not None:
log.warning(
"Skipping track_invocation_failure: invocation result already recorded on this graph tracker. "
"Call create_tracker on the agent graph for a new run. %s",
self.__get_track_data(),
)
return
self._summary.success = False
self._ld_client.track(
"$ld:ai:graph:invocation_failure",
self._context,
self.__get_track_data(),
1,
)
[docs]
def track_duration(self, duration: int) -> None:
"""
Track the total duration of a graph run.
Records at most once per graph tracker; further calls are ignored.
:param duration: Duration in milliseconds.
"""
if self._summary.duration_ms is not None:
log.warning(
"Skipping track_duration: duration already recorded on this graph tracker. "
"Call create_tracker on the agent graph for a new run. %s",
self.__get_track_data(),
)
return
self._summary.duration_ms = duration
self._ld_client.track(
"$ld:ai:graph:duration:total",
self._context,
self.__get_track_data(),
duration,
)
[docs]
def track_total_tokens(self, tokens: Optional[TokenUsage] = None) -> None:
"""
Track aggregated token usage across the entire graph run.
Records at most once per graph tracker; further calls are ignored.
:param tokens: Token usage data, or ``None`` when usage is unknown.
"""
if tokens is None or tokens.total <= 0:
return
if self._summary.tokens is not None:
log.warning(
"Skipping track_total_tokens: tokens already recorded on this graph tracker. "
"Call create_tracker on the agent graph for a new run. %s",
self.__get_track_data(),
)
return
self._summary.tokens = tokens
self._ld_client.track(
"$ld:ai:graph:total_tokens",
self._context,
self.__get_track_data(),
tokens.total,
)
[docs]
def track_path(self, path: List[str]) -> None:
"""
Track the path traversed through the graph during a graph run.
Appends to the summary's path list and fires a ``$ld:ai:graph:path``
event.
May be called multiple times per Tracker; each call records the
provided path segment and appends it to the summary so the full
path can be built incrementally.
:param path: An array of configuration keys representing the sequence of nodes executed during graph traversal.
"""
self._summary.path.extend(path)
track_data = {**self.__get_track_data(), "path": path}
self._ld_client.track(
"$ld:ai:graph:path",
self._context,
track_data,
1,
)
[docs]
def track_redirect(self, source_key: str, redirected_target: str) -> None:
"""
Track when a node redirects to a different target than originally specified.
:param source_key: The configuration key of the source node.
:param redirected_target: The configuration key of the target node that was redirected to.
"""
track_data = {
**self.__get_track_data(),
"sourceKey": source_key,
"redirectedTarget": redirected_target,
}
self._ld_client.track(
"$ld:ai:graph:redirect",
self._context,
track_data,
1,
)
[docs]
def track_handoff_success(self, source_key: str, target_key: str) -> None:
"""
Track successful handoffs between nodes.
:param source_key: The configuration key of the source node.
:param target_key: The configuration key of the target node.
"""
track_data = {
**self.__get_track_data(),
"sourceKey": source_key,
"targetKey": target_key,
}
self._ld_client.track(
"$ld:ai:graph:handoff_success",
self._context,
track_data,
1,
)
[docs]
def track_handoff_failure(self, source_key: str, target_key: str) -> None:
"""
Track failed handoffs between nodes.
:param source_key: The configuration key of the source node.
:param target_key: The configuration key of the target node.
"""
track_data = {
**self.__get_track_data(),
"sourceKey": source_key,
"targetKey": target_key,
}
self._ld_client.track(
"$ld:ai:graph:handoff_failure",
self._context,
track_data,
1,
)
def _track_from_graph_metrics(
self,
result: Any,
metrics_extractor: Callable[[Any], Optional[AIGraphMetrics]],
elapsed_ms: int,
) -> None:
metrics: Optional[AIGraphMetrics] = None
try:
metrics = metrics_extractor(result)
except Exception as exc:
log.warning("Failed to extract graph metrics: %s", exc)
if metrics is None:
self.track_duration(elapsed_ms)
return
self.track_duration(metrics.duration_ms if metrics.duration_ms is not None else elapsed_ms)
if metrics.success:
self.track_invocation_success()
else:
self.track_invocation_failure()
if metrics.path:
self.track_path(metrics.path)
if metrics.tokens is not None:
self.track_total_tokens(metrics.tokens)
[docs]
def track_graph_metrics_of(
self,
metrics_extractor: Callable[[Any], Optional[AIGraphMetrics]],
func: Callable[[], Any],
) -> Any:
"""
Track graph-level metrics for a synchronous graph operation.
Times the operation, extracts :class:`~ldai.providers.types.AIGraphMetrics`
via the provided extractor, and fires graph-level tracking events
(path, duration, success/failure, total tokens).
If the extracted ``AIGraphMetrics`` has a non-``None`` ``duration_ms``,
that value is used instead of the wall-clock elapsed time.
Node-level metrics are not tracked by this method.
For async operations, use :meth:`track_graph_metrics_of_async`.
:param metrics_extractor: Function that extracts AIGraphMetrics from the result
:param func: Synchronous callable that runs the graph operation
:return: The result of the operation
"""
start_ns = time.perf_counter_ns()
try:
result = func()
except Exception as err:
duration = (time.perf_counter_ns() - start_ns) // 1_000_000
self.track_duration(duration)
self.track_invocation_failure()
raise err
elapsed_ms = (time.perf_counter_ns() - start_ns) // 1_000_000
self._track_from_graph_metrics(result, metrics_extractor, elapsed_ms)
return result
[docs]
async def track_graph_metrics_of_async(
self,
metrics_extractor: Callable[[Any], Optional[AIGraphMetrics]],
func: Callable[[], Any],
) -> Any:
"""
Track graph-level metrics for an async graph operation (``func`` is awaited).
Same event semantics as :meth:`track_graph_metrics_of`.
:param metrics_extractor: Function that extracts AIGraphMetrics from the result
:param func: Async callable that runs the graph operation
:return: The result of the operation
"""
start_ns = time.perf_counter_ns()
try:
result = await func()
except Exception as err:
duration = (time.perf_counter_ns() - start_ns) // 1_000_000
self.track_duration(duration)
self.track_invocation_failure()
raise err
elapsed_ms = (time.perf_counter_ns() - start_ns) // 1_000_000
self._track_from_graph_metrics(result, metrics_extractor, elapsed_ms)
return result