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
4 changes: 4 additions & 0 deletions AGENTS.md
Original file line number Diff line number Diff line change
Expand Up @@ -193,6 +193,8 @@ LaunchDarkly metadata attached to a flag variation.
| `variationKey` | `str?` | Identifier for the specific variation. |
| `version` | `int?` | Variation version number. |
| `mode` | `"agent" \| "completion" \| "judge"` | Execution mode, used alongside `provider.name` to select a handler. |
| `modelKey` | `str?` | Stable key of the pinned model config, from `_ldMeta.modelKey`. Absent when the variation has no linked model config. Copied onto `TrackData`. |
| `modelVersion` | `int?` | Pinned model config version, from `_ldMeta.modelVersion`. Copied onto `TrackData`. |

#### `ProviderResponse`

Expand Down Expand Up @@ -240,6 +242,8 @@ Payload attached to every LaunchDarkly tracking event.
| `version` | `int` | Variation version number. |
| `modelName` | `str` | Model name from the config. |
| `providerName` | `str` | Provider name from the config. |
| `modelKey` | `str?` | Stable key of the pinned model config, read from `_ldMeta.modelKey`. Omitted when the variation has no pinned model config. |
| `modelVersion` | `int?` | Pinned model config version, read from `_ldMeta.modelVersion`. Omitted when absent. |
| `graphKey` | `str?` | Present when the event was produced inside an agent graph. |
| `toolKey` | `str?` | Present when the event is for a tool call. |
| `judgeConfigKey` | `str?` | Present when the event is from a judge execution. |
Expand Down
4 changes: 4 additions & 0 deletions packages/client/src/launchdarkly_ai_server/__init__.py
Original file line number Diff line number Diff line change
Expand Up @@ -109,8 +109,10 @@
end_unfinished_spans,
lang_chain_span_usage,
make_track_data,
model_stamps_from_meta,
normalize_mode,
number_or_zero,
omit_model_stamps,
parse_json_with_possible_fences,
parse_template,
parse_usage,
Expand Down Expand Up @@ -195,6 +197,8 @@
# utils
"create_handler",
"make_track_data",
"model_stamps_from_meta",
"omit_model_stamps",
"normalize_mode",
"parse_json_with_possible_fences",
"parse_template",
Expand Down
3 changes: 2 additions & 1 deletion packages/client/src/launchdarkly_ai_server/graph.py
Original file line number Diff line number Diff line change
Expand Up @@ -24,7 +24,7 @@
UsageDict,
VariationMeta,
)
from .utils import select_handler, to_ld_context
from .utils import model_stamps_from_meta, select_handler, to_ld_context

logger = logging.getLogger(__name__)

Expand Down Expand Up @@ -108,6 +108,7 @@ async def _build_graph(
"version": meta.get("version", 1) if isinstance(meta, dict) else 1,
"modelName": "",
"providerName": "",
**model_stamps_from_meta(meta),

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

🟡 Graph model stamps never refresh

After a graph variation changes, graph_track_data keeps cached model stamps for the same context. Later graph events remain attributed to the old pinned model version.

Learn more

GraphInstance.invoke caches both the resolved GraphDefinition and graph_track_data by context in the context cache. There is no expiry or LaunchDarkly-driven invalidation. The cached definition also contains every node's metadata, so both graph-level and node-level model stamps remain fixed for the lifetime of that cache entry. Model reassignment or version changes therefore never reach later tracking events from the same GraphInstance and context.

Example: A graph first resolves with modelVersion=5. LaunchDarkly updates the pin to version 6, but the next invocation with the same context reuses version 5 in all affected events.

Recommended fix: Do not indefinitely cache variation-derived metadata. Reevaluate graph and node variations for each invocation, or add a bounded freshness/invalidation mechanism that refreshes the cached definition and tracking data together.

Devin Review


Was this helpful? React with 👍 or 👎 to provide feedback.

"graphKey": key,
}

Expand Down
6 changes: 5 additions & 1 deletion packages/client/src/launchdarkly_ai_server/judges.py
Original file line number Diff line number Diff line change
Expand Up @@ -26,6 +26,7 @@
)
from .utils import (
normalize_mode,
omit_model_stamps,
to_ld_context,
to_usage_dict,
)
Expand Down Expand Up @@ -413,8 +414,11 @@ def _matches(h: ProviderHandler) -> bool:

usage = to_usage_dict(raw_usage)

# A judge without a pinned model config must not inherit the parent's
# modelKey / modelVersion; every other parent-only key (graphKey, ...)
# is still carried over.
merged_track_data: TrackData = {
**task.parent_track_data,
**omit_model_stamps(task.parent_track_data),
**result["track_data"],
"judgeConfigKey": task.config_key,
}
Expand Down
4 changes: 4 additions & 0 deletions packages/client/src/launchdarkly_ai_server/lifecycle.py
Original file line number Diff line number Diff line change
Expand Up @@ -8,6 +8,7 @@

from .sdk_info import flush_ai_sdk_info, reset_ai_sdk_info
from .types import InitClientOptions
from .utils import model_stamps_from_meta

logger = logging.getLogger(__name__)

Expand Down Expand Up @@ -353,6 +354,9 @@ async def extract_variation(
"variationKey": ld_meta.get("variationKey", ""),
"version": ld_meta.get("version", 1),
"mode": ld_meta.get("mode"),
# Pinned model-config identity; keys are omitted when absent so that
# tracking payloads never carry ``None`` values.
**model_stamps_from_meta(ld_meta),
}

# Strip _ldMeta for config parsing
Expand Down
4 changes: 3 additions & 1 deletion packages/client/src/launchdarkly_ai_server/tracking.py
Original file line number Diff line number Diff line change
Expand Up @@ -17,7 +17,7 @@
TrackData,
VariationMeta,
)
from .utils import parse_usage, to_ld_context
from .utils import model_stamps_from_meta, parse_usage, to_ld_context


def _try_get_environment_id() -> str | None:
Expand Down Expand Up @@ -129,6 +129,7 @@ async def execute_and_track(
"providerName": config.get("provider", {}).get("name", "")
if isinstance(config, dict)
else "",
**model_stamps_from_meta(meta),
Comment thread
atornsii marked this conversation as resolved.
}
if graph_key:
track_data["graphKey"] = graph_key
Expand Down Expand Up @@ -210,6 +211,7 @@ async def execute_and_stream(
"providerName": config.get("provider", {}).get("name", "")
if isinstance(config, dict)
else "",
**model_stamps_from_meta(meta),
}
if graph_key:
track_data["graphKey"] = graph_key
Expand Down
4 changes: 3 additions & 1 deletion packages/client/src/launchdarkly_ai_server/types.py
Original file line number Diff line number Diff line change
Expand Up @@ -97,7 +97,9 @@ class Message:

VariationMeta = dict[str, Any]
"""
Variation metadata: ``enabled``, ``variation_key``, ``version``, ``mode``.
Variation metadata: ``enabled``, ``variationKey``, ``version``, ``mode``, and
the pinned model-config identity ``modelKey`` / ``modelVersion`` (delivered in
``_ldMeta``; absent when no model config is pinned).
"""

# ---------------------------------------------------------------------------
Expand Down
55 changes: 55 additions & 0 deletions packages/client/src/launchdarkly_ai_server/utils.py
Original file line number Diff line number Diff line change
Expand Up @@ -2,6 +2,7 @@

import json
import re
from collections.abc import Mapping
from dataclasses import dataclass, field
from typing import Any, Literal

Expand Down Expand Up @@ -538,6 +539,59 @@ def select_handler(
raise ValueError(f"Handler for provider {provider} not found")


def model_stamps_from_meta(meta: Any) -> dict[str, Any]:
"""
Copies the pinned model-config identity (``modelKey``, ``modelVersion``)
from a variation's ``_ldMeta`` into a dict that can be merged into
``TrackData``. Keys are omitted (never set to ``None``) when absent; an
empty ``modelKey`` is treated as absent and ``modelVersion`` is coerced to
``int``. Gonfalon's cost attribution reads these two fields from every
``$ld:ai:*`` event payload.
"""
if not isinstance(meta, dict):
return {}
stamps: dict[str, Any] = {}
model_key = meta.get("modelKey")
if isinstance(model_key, str) and model_key:
stamps["modelKey"] = model_key
model_version = _coerce_model_version(meta.get("modelVersion"))
if model_version is not None:
stamps["modelVersion"] = model_version
return stamps


def _coerce_model_version(value: Any) -> int | None:
"""
Coerces an ``_ldMeta.modelVersion`` value to ``int``. Accepts ``int``
(but not ``bool``), integral ``float`` and integer-looking ``str``;
returns ``None`` for anything else. ``_ldMeta`` is an untyped flag payload,
so a malformed value must be dropped rather than abort the invocation.
"""
if isinstance(value, bool):
return None
if isinstance(value, int):
return value
if isinstance(value, float):
return int(value) if value.is_integer() else None
if isinstance(value, str):
try:
return int(value.strip())
except ValueError:
return None
return None


def omit_model_stamps(track_data: Mapping[str, Any]) -> dict[str, Any]:
"""
Returns a copy of ``track_data`` without ``modelKey`` / ``modelVersion``.
Used when overlaying a judge's ``track_data`` on its parent's so a judge
without a pinned model config does not inherit the parent's identity.
"""
return {
k: v for k, v in track_data.items() if k not in ("modelKey", "modelVersion")
}


def make_track_data(node: GraphNode, graph_key: str, run_id: str) -> dict[str, Any]:
"""
Builds the standard tracking payload for a graph node event.
Expand All @@ -552,6 +606,7 @@ def make_track_data(node: GraphNode, graph_key: str, run_id: str) -> dict[str, A
"version": meta.get("version", 1),
"modelName": config.get("model", {}).get("name", ""),
"providerName": config.get("provider", {}).get("name", ""),
**model_stamps_from_meta(meta),
"graphKey": graph_key,
}

Expand Down
68 changes: 68 additions & 0 deletions packages/client/tests/test_client.py
Original file line number Diff line number Diff line change
Expand Up @@ -112,6 +112,74 @@ async def test_returns_provider_response(self, mock_ld_client: MagicMock) -> Non
result = await m.invoke("q", CONTEXT)
assert result.response == "answer"

async def test_track_data_copies_model_key_and_version_from_ld_meta(
self, mock_ld_client: MagicMock
) -> None:
raw = await mock_ld_client.variation("flag", CONTEXT, None)
raw = {
**raw,
"_ldMeta": {**raw["_ldMeta"], "modelKey": "my-model", "modelVersion": 3},
}
mock_ld_client.variation = AsyncMock(return_value=raw)
m = config(key="flag", handler=_make_handler())
result = await m.invoke("q", CONTEXT)
assert result.track_data is not None
assert result.track_data["modelKey"] == "my-model"
assert result.track_data["modelVersion"] == 3
assert isinstance(result.track_data["modelVersion"], int)
assert mock_ld_client.track.call_args_list
for call in mock_ld_client.track.call_args_list:
payload = call[0][2]
assert payload["modelKey"] == "my-model"
assert payload["modelVersion"] == 3

async def test_track_data_omits_model_key_and_version_when_absent(
self, mock_ld_client: MagicMock
) -> None:
m = config(key="flag", handler=_make_handler())
result = await m.invoke("q", CONTEXT)
assert result.track_data is not None
assert "modelKey" not in result.track_data
assert "modelVersion" not in result.track_data
for call in mock_ld_client.track.call_args_list:
payload = call[0][2]
assert "modelKey" not in payload
assert "modelVersion" not in payload

async def test_track_data_treats_empty_model_key_as_absent(
self, mock_ld_client: MagicMock
) -> None:
raw = await mock_ld_client.variation("flag", CONTEXT, None)
raw = {
**raw,
"_ldMeta": {**raw["_ldMeta"], "modelKey": "", "modelVersion": "2"},
}
mock_ld_client.variation = AsyncMock(return_value=raw)
m = config(key="flag", handler=_make_handler())
result = await m.invoke("q", CONTEXT)
assert result.track_data is not None
assert "modelKey" not in result.track_data
assert result.track_data["modelVersion"] == 2
assert isinstance(result.track_data["modelVersion"], int)

async def test_stream_track_data_copies_model_key_and_version(
self, mock_ld_client: MagicMock
) -> None:
raw = await mock_ld_client.variation("flag", CONTEXT, None)
raw = {
**raw,
"_ldMeta": {**raw["_ldMeta"], "modelKey": "my-model", "modelVersion": 3},
}
mock_ld_client.variation = AsyncMock(return_value=raw)
m = config(key="flag", handler=_make_handler(stream_chunks=["a", "b"]))
async for _ in m.stream("q", CONTEXT):
pass
assert mock_ld_client.track.call_args_list
for call in mock_ld_client.track.call_args_list:
payload = call[0][2]
assert payload["modelKey"] == "my-model"
assert payload["modelVersion"] == 3

async def test_generation_success_on_success(
self, mock_ld_client: MagicMock
) -> None:
Expand Down
49 changes: 49 additions & 0 deletions packages/client/tests/test_graph.py
Original file line number Diff line number Diff line change
Expand Up @@ -192,6 +192,55 @@ async def test_traverses_root_leaf_returns_aggregated_result(
result = await g.invoke("hi", CONTEXT)
assert result.response is not None

async def test_graph_events_copy_model_key_and_version_from_ld_meta(
self, mock_ld_client: MagicMock
) -> None:
graph_var = {
"_ldMeta": {
"enabled": True,
"variationKey": "gv1",
"version": 2,
"modelKey": "graph-model",
"modelVersion": 5,
},
"root": "root-node",
"edges": {"root-node": [{"key": "leaf-node"}]},
}
original = mock_ld_client.variation

async def side_effect(key: str, ctx: dict, default: Any) -> Any:
if key == "graph-key":
return graph_var
return await original(key, ctx, default)

mock_ld_client.variation = AsyncMock(side_effect=side_effect)
g = graph("graph-key", handlers=[_make_handler()])
await g.invoke("hi", CONTEXT)
graph_calls = [
c
for c in mock_ld_client.track.call_args_list
if str(c[0][0]).startswith("$ld:ai:graph:")
]
assert graph_calls
for c in graph_calls:
assert c[0][2]["modelKey"] == "graph-model"
assert c[0][2]["modelVersion"] == 5

async def test_graph_events_omit_model_key_and_version_when_absent(
self, mock_ld_client: MagicMock
) -> None:
g = graph("graph-key", handlers=[_make_handler()])
await g.invoke("hi", CONTEXT)
graph_calls = [
c
for c in mock_ld_client.track.call_args_list
if str(c[0][0]).startswith("$ld:ai:graph:")
]
assert graph_calls
for c in graph_calls:
assert "modelKey" not in c[0][2]
assert "modelVersion" not in c[0][2]

async def test_graph_duration_total_tracked(
self, mock_ld_client: MagicMock
) -> None:
Expand Down
Loading
Loading