From 60ebd3c4c1e03f1035b5f28743528c943da4d546 Mon Sep 17 00:00:00 2001 From: Harsh Sadhvani Date: Wed, 28 Oct 2026 13:34:15 -0700 Subject: [PATCH 1/4] feat(shadow-requests): run Shadow Requests from the experimental entry point --- packages/client/README.md | 28 + ...ai_server.experimental.shadow_requests.txt | 22 + .../api-report/launchdarkly_ai_server.txt | 4 +- .../src/launchdarkly_ai_server/client.py | 65 ++ .../experimental/shadow_requests.py | 319 +++++++++ .../invocation_hooks.py | 64 ++ .../src/launchdarkly_ai_server/lifecycle.py | 33 +- .../src/launchdarkly_ai_server/tracking.py | 127 ++-- .../src/launchdarkly_ai_server/utils.py | 4 + packages/client/tests/test_shadow_requests.py | 652 ++++++++++++++++++ release-please-config.json | 3 + 11 files changed, 1263 insertions(+), 58 deletions(-) create mode 100644 packages/client/api-report/launchdarkly_ai_server.experimental.shadow_requests.txt create mode 100644 packages/client/src/launchdarkly_ai_server/experimental/shadow_requests.py create mode 100644 packages/client/src/launchdarkly_ai_server/invocation_hooks.py create mode 100644 packages/client/tests/test_shadow_requests.py diff --git a/packages/client/README.md b/packages/client/README.md index ac0ebf49..c113a14d 100644 --- a/packages/client/README.md +++ b/packages/client/README.md @@ -308,6 +308,34 @@ result = await evals.run( **The list is checked before any network I/O.** A blank key, an uppercase key, a `schema` that is not a JSON object, a schema that is not JSON-serializable (including a `NaN` or `Infinity` value), and a non-callable implementation each fail with zero requests issued. A repeated key fails too, and keys are compared without case. A `NativeTool` is valid only for a library tool, because the provider supplies its schema. +### Shadow Requests (experimental) + +A Shadow Request runs another variation of an AI Config in the background on the same input, and reports it on the normal AI metrics with `isShadow: true`. LaunchDarkly attaches shadow requests to a targeting rule or the fallthrough. The shadow's response is discarded, and your response, latency and events are unchanged. Shadow runs use your provider account, so they are billed like real calls. + +Nothing runs until you enable it: + +```python +from launchdarkly_ai_server.experimental.shadow_requests import enable_shadow_requests + +enable_shadow_requests() # run sampled shadows in the background, 30 s timeout +enable_shadow_requests(timeout_s=10) # or pick the timeout +``` + +Shadows run for `config().invoke()` and `config().stream()`, not for graphs or judges. A shadow makes no tool calls and runs no judges. Call `disable_shadow_requests()` to turn them off. + +To run shadows on your own queue or worker instead, pass `on_task`. It receives a picklable `ShadowTask` before your handler runs and must only enqueue it. Run the task later with `run_shadow`: + +```python +from launchdarkly_ai_server.experimental.shadow_requests import enable_shadow_requests, run_shadow + +enable_shadow_requests(on_task=queue.put_nowait) + +# in the worker +result = await run_shadow(task, handlers=[create_openai_messages_handler()]) +``` + +These names may change in a minor release. + ### `config(**args)` The primary entry point for AI config invocations. Accepts either a single handler or a list of handlers and routes to the correct one at invoke-time based on the flag variation's provider and mode. diff --git a/packages/client/api-report/launchdarkly_ai_server.experimental.shadow_requests.txt b/packages/client/api-report/launchdarkly_ai_server.experimental.shadow_requests.txt new file mode 100644 index 00000000..938940c3 --- /dev/null +++ b/packages/client/api-report/launchdarkly_ai_server.experimental.shadow_requests.txt @@ -0,0 +1,22 @@ +# Public API of launchdarkly_ai_server.experimental.shadow_requests. Generated by `make api-report`; do not edit. +disable_shadow_requests() -> None +enable_shadow_requests(*, timeout_s: float = 30.0, on_task: Callable[[ShadowTask], None] | None = None) -> None +run_shadow(task: ShadowTask, handlers: list[ProviderHandler]) -> ShadowResult [async] +ShadowResult [class, @dataclass(frozen=True)] + __init__(usage: dict[str, int], track_data: dict[str, Any], error: str | None = None) + error: str | None = None + track_data: dict[str, Any] + usage: dict[str, int] +ShadowTask [class, @dataclass(frozen=True)] + __init__(config_key: str, shadow_request_id: str, shadow_config: AiConfigRep, shadow_meta: VariationMeta, user_input: str | None, history: list[dict[str, Any]] | None, variables: dict[str, Any] | None, user_context: LDContext, timeout_s: float, track_data: dict[str, Any], span_attributes: dict[str, Any] = dict()) + config_key: str + history: list[dict[str, Any]] | None + shadow_config: AiConfigRep + shadow_meta: VariationMeta + shadow_request_id: str + span_attributes: dict[str, Any] = field(default_factory=dict) + timeout_s: float + track_data: dict[str, Any] + user_context: LDContext + user_input: str | None + variables: dict[str, Any] | None diff --git a/packages/client/api-report/launchdarkly_ai_server.txt b/packages/client/api-report/launchdarkly_ai_server.txt index 4ed2c00e..8d52226b 100644 --- a/packages/client/api-report/launchdarkly_ai_server.txt +++ b/packages/client/api-report/launchdarkly_ai_server.txt @@ -24,8 +24,8 @@ create_handler(provides_for: tuple[str, Literal['agent', 'messages']], fn: _Hand create_run_usage() -> RunUsage end_span_once(span: Any, tracker: set[int], abandoned: bool = False, cancelled: bool = False) -> None end_unfinished_spans(*spans: Any) -> None -execute_and_stream(*, config_key: str, config: AiConfigRep, meta: VariationMeta, user_context: LDContext, handler: ProviderHandler, user_input: str | None = None, tool_handlers: dict[str, Callable[..., Any] | NativeTool] | None = None, variables: dict[str, Any] | None = None, graph_key: str | None = None, history: list[dict[str, Any]] | None = None) -> AsyncGenerator[ExecuteStreamEvent, None] [async] -execute_and_track(*, config_key: str, config: AiConfigRep, meta: VariationMeta, user_context: LDContext, handler: ProviderHandler, user_input: str | None = None, tool_handlers: dict[str, Callable[..., Any] | NativeTool] | None = None, variables: dict[str, Any] | None = None, graph_key: str | None = None, history: list[dict[str, Any]] | None = None) -> dict[str, Any] [async] +execute_and_stream(*, config_key: str, config: AiConfigRep, meta: VariationMeta, user_context: LDContext, handler: ProviderHandler, user_input: str | None = None, tool_handlers: dict[str, Callable[..., Any] | NativeTool] | None = None, variables: dict[str, Any] | None = None, graph_key: str | None = None, history: list[dict[str, Any]] | None = None, run_id: str | None = None, extra_track_data: dict[str, Any] | None = None, span_attributes: dict[str, Any] | None = None) -> AsyncGenerator[ExecuteStreamEvent, None] [async] +execute_and_track(*, config_key: str, config: AiConfigRep, meta: VariationMeta, user_context: LDContext, handler: ProviderHandler, user_input: str | None = None, tool_handlers: dict[str, Callable[..., Any] | NativeTool] | None = None, variables: dict[str, Any] | None = None, graph_key: str | None = None, history: list[dict[str, Any]] | None = None, run_id: str | None = None, extra_track_data: dict[str, Any] | None = None, span_attributes: dict[str, Any] | None = None) -> dict[str, Any] [async] ExecuteStreamDoneEvent = dict[str, Any] ExecuteStreamEvent = dict[str, Any] extract_variation(config_key: str, context: dict[str, Any]) -> dict[str, Any] [async] diff --git a/packages/client/src/launchdarkly_ai_server/client.py b/packages/client/src/launchdarkly_ai_server/client.py index be124d87..5d3ed1f4 100644 --- a/packages/client/src/launchdarkly_ai_server/client.py +++ b/packages/client/src/launchdarkly_ai_server/client.py @@ -1,11 +1,17 @@ from __future__ import annotations +import uuid from collections.abc import AsyncGenerator, Callable from contextlib import aclosing from typing import Any from .agui import StableProjector, judge_completed_stream from .conversation import bind_conversation_id +from .invocation_hooks import ( + InvocationEvent, + has_invocation_hooks, + run_invocation_hooks, +) from .judge_scoring import judge_text from .judges import _build_judge_tasks, run_judges from .lifecycle import extract_variation @@ -45,6 +51,39 @@ def _resolve_output_format_response( return parsed if parsed is not None else raw_response +def _run_hooks( + *, + config_key: str, + variation: dict[str, Any], + meta: VariationMeta, + context: LDContext, + user_input: str | None, + history: list[dict[str, Any]] | None, + variables: dict[str, Any] | None, + handlers: list[ProviderHandler], + run_id: str, +) -> None: + if not has_invocation_hooks(): + return + raw = variation.get("raw") + run_invocation_hooks( + InvocationEvent( + config_key=config_key, + value=raw if isinstance(raw, dict) else {}, + reason=variation.get("reason"), + context=context, + user_input=user_input, + history=history, + variables=variables, + handlers=handlers, + run_id=run_id, + variation_key=str(meta.get("variationKey", "")) + if isinstance(meta, dict) + else "", + ) + ) + + class ConfigInstance: """Return type of ``config()``.""" @@ -98,6 +137,18 @@ async def _invoke( if not resolved_handler_list: raise ValueError("No handlers provided to config()") handler = select_handler(config, meta, resolved_handler_list) + run_id = str(uuid.uuid4()) + _run_hooks( + config_key=self._key, + variation=variation, + meta=meta, + context=context, + user_input=user_input, + history=history, + variables=variables, + handlers=resolved_handler_list, + run_id=run_id, + ) result = await execute_and_track( config_key=self._key, @@ -109,6 +160,7 @@ async def _invoke( tool_handlers=resolved_tools, variables=variables, history=history, + run_id=run_id, ) raw_response = result["response"] @@ -228,6 +280,18 @@ async def _execute_stream( sink["tools"] = resolved_tools sink["provider"] = provider sink["model"] = model + run_id = str(uuid.uuid4()) + _run_hooks( + config_key=self._key, + variation=variation, + meta=meta, + context=context, + user_input=user_input, + history=history, + variables=variables, + handlers=resolved_handler_list, + run_id=run_id, + ) async with aclosing( execute_and_stream( @@ -240,6 +304,7 @@ async def _execute_stream( tool_handlers=resolved_tools, variables=variables, history=history, + run_id=run_id, ) ) as source: async for event in source: diff --git a/packages/client/src/launchdarkly_ai_server/experimental/shadow_requests.py b/packages/client/src/launchdarkly_ai_server/experimental/shadow_requests.py new file mode 100644 index 00000000..ffb61dd1 --- /dev/null +++ b/packages/client/src/launchdarkly_ai_server/experimental/shadow_requests.py @@ -0,0 +1,319 @@ +""" +Shadow Requests (experimental). + +Call ``enable_shadow_requests()`` to run the shadow variations LaunchDarkly attaches +to a targeting rule or the fallthrough. Each sampled shadow runs in the background on +the same input, and is tracked on the normal ``$ld:ai:*`` events with ``isShadow``. +Its response is discarded, and the caller's response, latency and events are +unchanged. These names may change in a minor release, so review the changelog when +you upgrade. +""" + +from __future__ import annotations + +import asyncio +import hashlib +import logging +import time +from collections.abc import Callable +from dataclasses import dataclass, field +from typing import Any + +from ..invocation_hooks import InvocationEvent, register_invocation_hook +from ..sdk_usage import report_usage +from ..shutdown_hooks import register_shutdown_hook +from ..tracking import build_track_data, execute_and_track +from ..types import AiConfigRep, LDContext, ProviderHandler, VariationMeta +from ..types_validation import parse_ai_config +from ..utils import ( + _context_identity_from_ld_context, + model_stamps_from_meta, + parse_usage, + select_handler, + to_ld_context, +) + +__all__ = [ + "ShadowResult", + "ShadowTask", + "disable_shadow_requests", + "enable_shadow_requests", + "run_shadow", +] + +logger = logging.getLogger("launchdarkly_ai_server.experimental.shadow_requests") + +_FEATURE = "Shadow Requests" +_FALLTHROUGH = "fallthrough" +_STRIPPED_KEYS = ("_ldMeta", "tools", "judgeConfiguration", "shadowConfiguration") +_BUCKET_SCALE = float(0xFFFFFFFFFFFFFFF) + + +@dataclass(frozen=True) +class ShadowTask: + """ + Everything needed to run and track one sampled shadow, with no flag evaluation. + Plain data, so it survives ``pickle``. Run it with ``run_shadow(task, handlers)``. + """ + + config_key: str + shadow_request_id: str + shadow_config: AiConfigRep + shadow_meta: VariationMeta + user_input: str | None + history: list[dict[str, Any]] | None + variables: dict[str, Any] | None + user_context: LDContext + timeout_s: float + track_data: dict[str, Any] + span_attributes: dict[str, Any] = field(default_factory=dict) + + +@dataclass(frozen=True) +class ShadowResult: + """The outcome of one shadow run. There is no response text: it is discarded.""" + + usage: dict[str, int] + track_data: dict[str, Any] + error: str | None = None + + +@dataclass(frozen=True) +class _Settings: + timeout_s: float + on_task: Callable[[ShadowTask], None] | None + + +_settings: _Settings | None = None +_pending: set[asyncio.Task[ShadowResult]] = set() + + +def enable_shadow_requests( + *, + timeout_s: float = 30.0, + on_task: Callable[[ShadowTask], None] | None = None, +) -> None: + """ + Turns Shadow Requests on for this process. Calling it again replaces the settings. + + By default each sampled shadow runs in the background with *timeout_s*. Pass + *on_task* to run them yourself instead: it receives a ``ShadowTask`` for each + sampled shadow, before the real handler runs, and must only enqueue it. + """ + report_usage("experimental.enableShadowRequests") + if not timeout_s > 0: + raise ValueError("timeout_s must be greater than 0") + global _settings + _settings = _Settings(timeout_s=timeout_s, on_task=on_task) + register_invocation_hook(_FEATURE, _on_invocation) + + +def disable_shadow_requests() -> None: + """Turns Shadow Requests off. Shadows already running finish.""" + global _settings + _settings = None + register_invocation_hook(_FEATURE, None) + + +async def run_shadow(task: ShadowTask, handlers: list[ProviderHandler]) -> ShadowResult: + """Runs one ``ShadowTask``. Never raises; a failure is returned in ``error``.""" + report_usage("experimental.runShadow") + return await _run(task, handlers) + + +def _reset() -> None: + for pending in list(_pending): + pending.cancel() + _pending.clear() + disable_shadow_requests() + + +register_shutdown_hook(_FEATURE, _reset) + + +async def _wait_for_shadows_for_testing() -> None: + while _pending: + await asyncio.gather(*list(_pending), return_exceptions=True) + + +def _in_sample(shadow_id: str, context: LDContext, sampling_rate: float) -> bool: + identity = _context_identity_from_ld_context(context) + if identity is None: + return False + canonical, _ = identity + digest = hashlib.sha1(f"{shadow_id}.{canonical}".encode(), usedforsecurity=False) + bucket = int(digest.hexdigest()[:15], 16) / _BUCKET_SCALE + return bucket < sampling_rate + + +def _slot(reason: Any) -> str | None: + if not isinstance(reason, dict): + return None + kind = reason.get("kind") + if kind == "FALLTHROUGH": + return _FALLTHROUGH + if kind == "RULE_MATCH" and isinstance(reason.get("ruleId"), str): + return str(reason["ruleId"]) + return None + + +def _valid_entry(entry: Any) -> bool: + if not isinstance(entry, dict): + return False + numbers_ok = all( + isinstance(entry.get(k), (int, float)) and not isinstance(entry.get(k), bool) + for k in ("samplingRate", "endsAt") + ) + strings_ok = all( + isinstance(entry.get(k), str) and entry.get(k) for k in ("id", "variationKey") + ) + return numbers_ok and strings_ok + + +def _parse_variation(raw: Any) -> tuple[AiConfigRep, VariationMeta] | None: + if not isinstance(raw, dict): + return None + ld_meta = raw.get("_ldMeta") + if not isinstance(ld_meta, dict) or not ld_meta.get("enabled"): + return None + parsed = parse_ai_config({k: v for k, v in raw.items() if k not in _STRIPPED_KEYS}) + if not parsed.success: + return None + meta: VariationMeta = { + "enabled": True, + "variationKey": ld_meta.get("variationKey", ""), + "version": ld_meta.get("version", 1), + "mode": ld_meta.get("mode"), + **model_stamps_from_meta(ld_meta), + } + return parsed.data, meta + + +def _on_invocation(event: InvocationEvent) -> None: + settings = _settings + if settings is None: + return + slot = _slot(event.reason) + configuration = event.value.get("shadowConfiguration") + if slot is None or not isinstance(configuration, dict): + return + shadows = configuration.get("shadows") + variations = configuration.get("variations") + if not isinstance(shadows, dict) or not isinstance(variations, dict): + return + entries = shadows.get(slot) + if not isinstance(entries, list): + return + + now_ms = time.time() * 1000 + for entry in entries: + if not _valid_entry(entry) or entry["endsAt"] <= now_ms: + continue + if not _in_sample(entry["id"], event.context, float(entry["samplingRate"])): + continue + parsed = _parse_variation(variations.get(entry["variationKey"])) + if parsed is None: + continue + shadow_config, shadow_meta = parsed + try: + select_handler(shadow_config, shadow_meta, event.handlers) + except Exception: + logger.debug("No handler for shadow request %s", entry["id"]) + continue + task = ShadowTask( + config_key=event.config_key, + shadow_request_id=entry["id"], + shadow_config=shadow_config, + shadow_meta=shadow_meta, + user_input=event.user_input, + history=list(event.history) if event.history else None, + variables=dict(event.variables) if event.variables else None, + user_context=dict(event.context), + timeout_s=settings.timeout_s, + track_data={ + "runId": event.run_id, + "isShadow": True, + "primaryVariationKey": event.variation_key, + "shadowRequestId": entry["id"], + "ownerKind": entry.get("ownerKind", ""), + "ownerId": entry.get("ownerId", ""), + }, + span_attributes={ + "launchdarkly.shadow": True, + "launchdarkly.shadow.request.id": entry["id"], + "launchdarkly.shadow.primary.variation.key": event.variation_key, + }, + ) + if settings.on_task is not None: + try: + settings.on_task(task) + except Exception: + logger.warning( + "The Shadow Requests on_task callback failed", exc_info=True + ) + continue + running = asyncio.get_running_loop().create_task(_run(task, event.handlers)) + _pending.add(running) + running.add_done_callback(_pending.discard) + + +async def _run(task: ShadowTask, handlers: list[ProviderHandler]) -> ShadowResult: + extra = {k: v for k, v in task.track_data.items() if k != "runId"} + run_id = task.track_data.get("runId") + track_data = build_track_data( + config_key=task.config_key, + config=task.shadow_config, + meta=task.shadow_meta, + run_id=run_id, + extra_track_data=extra, + ) + try: + handler = select_handler(task.shadow_config, task.shadow_meta, handlers) + except Exception as err: + return ShadowResult( + usage=parse_usage({}), track_data=track_data, error=str(err) + ) + + start = time.monotonic() + try: + result = await asyncio.wait_for( + execute_and_track( + config_key=task.config_key, + config=task.shadow_config, + meta=task.shadow_meta, + user_context=task.user_context, + handler=handler, + user_input=task.user_input, + tool_handlers=None, + variables=task.variables, + history=list(task.history) if task.history else None, + run_id=run_id, + extra_track_data=extra, + span_attributes=task.span_attributes, + ), + timeout=task.timeout_s, + ) + except TimeoutError: + _track_timeout(task, track_data, int((time.monotonic() - start) * 1000)) + return ShadowResult( + usage=parse_usage({}), track_data=track_data, error="timeout" + ) + except Exception as err: + return ShadowResult( + usage=parse_usage({}), track_data=track_data, error=str(err) + ) + return ShadowResult(usage=result["usage"], track_data=result["track_data"]) + + +def _track_timeout( + task: ShadowTask, track_data: dict[str, Any], elapsed_ms: int +) -> None: + from ..lifecycle import get_client + + try: + client = get_client() + ld_ctx = to_ld_context(client, task.user_context) + client.track("$ld:ai:duration:total", ld_ctx, track_data, elapsed_ms) + client.track("$ld:ai:generation:error", ld_ctx, track_data, 1) + except Exception: + logger.debug("Could not track a shadow timeout", exc_info=True) diff --git a/packages/client/src/launchdarkly_ai_server/invocation_hooks.py b/packages/client/src/launchdarkly_ai_server/invocation_hooks.py new file mode 100644 index 00000000..0534f404 --- /dev/null +++ b/packages/client/src/launchdarkly_ai_server/invocation_hooks.py @@ -0,0 +1,64 @@ +""" +Hooks that ``config().invoke()`` and ``.stream()`` run before the handler. + +This is how a feature reacts to an invocation without core importing it: a +feature module registers a hook, and core runs whatever is registered. Core +never imports, names or calls a feature module, so a hook's *name* comes from +the module that registers it. Same shape as ``shutdown_hooks``. +""" + +from __future__ import annotations + +import logging +from collections.abc import Callable +from dataclasses import dataclass +from typing import Any + +from .types import LDContext, ProviderHandler + +logger = logging.getLogger("launchdarkly_ai_server.client") + + +@dataclass(frozen=True) +class InvocationEvent: + """What a hook sees: one ``config()`` call, after routing and before the handler runs.""" + + config_key: str + value: dict[str, Any] + reason: Any + context: LDContext + user_input: str | None + history: list[dict[str, Any]] | None + variables: dict[str, Any] | None + handlers: list[ProviderHandler] + run_id: str + variation_key: str + + +_hooks: dict[str, Callable[[InvocationEvent], None]] = {} + + +def register_invocation_hook( + name: str, hook: Callable[[InvocationEvent], None] | None +) -> None: + """ + Registers *hook* under *name*, replacing any earlier hook with that name. + Passing ``None`` removes it. + """ + if hook is None: + _hooks.pop(name, None) + else: + _hooks[name] = hook + + +def has_invocation_hooks() -> bool: + return bool(_hooks) + + +def run_invocation_hooks(event: InvocationEvent) -> None: + """Runs every hook. A hook that raises is logged and must not break the call.""" + for name, hook in list(_hooks.items()): + try: + hook(event) + except Exception: + logger.warning("The %s invocation hook failed", name, exc_info=True) diff --git a/packages/client/src/launchdarkly_ai_server/lifecycle.py b/packages/client/src/launchdarkly_ai_server/lifecycle.py index e7bb8fc7..d0249749 100644 --- a/packages/client/src/launchdarkly_ai_server/lifecycle.py +++ b/packages/client/src/launchdarkly_ai_server/lifecycle.py @@ -601,13 +601,26 @@ async def extract_variation( # The real ldclient.variation() expects an ldclient.Context, not a plain dict. # Convert when running against the real SDK; test mocks accept dicts directly. ld_context = to_ld_context(client, context) - # variation() is synchronous in the real Python LD SDK; AsyncMock in tests. - variation_result = client.variation(config_key, ld_context, None) - raw = ( - await variation_result - if inspect.isawaitable(variation_result) - else variation_result - ) + # The reason is only read while an invocation hook is registered, so an app + # with none keeps calling variation() exactly as before (spec §9.6). + from .invocation_hooks import has_invocation_hooks + + reason: Any = None + if has_invocation_hooks(): + detail_result = client.variation_detail(config_key, ld_context, None) + detail = ( + await detail_result if inspect.isawaitable(detail_result) else detail_result + ) + raw = getattr(detail, "value", None) + reason = getattr(detail, "reason", None) + else: + # variation() is synchronous in the real Python LD SDK; AsyncMock in tests. + variation_result = client.variation(config_key, ld_context, None) + raw = ( + await variation_result + if inspect.isawaitable(variation_result) + else variation_result + ) if raw is None: raise RuntimeError( @@ -629,9 +642,9 @@ async def extract_variation( **model_stamps_from_meta(ld_meta), } - # Strip _ldMeta for config parsing + # Strip _ldMeta and shadowConfiguration for config parsing config_raw = ( - {k: v for k, v in raw.items() if k != "_ldMeta"} + {k: v for k, v in raw.items() if k not in ("_ldMeta", "shadowConfiguration")} if isinstance(raw, dict) else raw ) @@ -642,4 +655,4 @@ async def extract_variation( f"Invalid AI config variation for '{config_key}': {result.error['message']}" ) - return {"config": result.data, "meta": meta} + return {"config": result.data, "meta": meta, "raw": raw, "reason": reason} diff --git a/packages/client/src/launchdarkly_ai_server/tracking.py b/packages/client/src/launchdarkly_ai_server/tracking.py index acc4f8b5..f2f2d631 100644 --- a/packages/client/src/launchdarkly_ai_server/tracking.py +++ b/packages/client/src/launchdarkly_ai_server/tracking.py @@ -52,6 +52,59 @@ def _try_get_environment_id() -> str | None: return os.environ.get("LD_ENVIRONMENT_ID") or None +SPAN_ATTRIBUTES_KEY = "__ld_span" +"""Variables key for extra root-span attributes. Set only by the SDK, never from template variables.""" + + +def build_track_data( + *, + config_key: str, + config: AiConfigRep, + meta: VariationMeta, + graph_key: str | None = None, + run_id: str | None = None, + extra_track_data: dict[str, Any] | None = None, +) -> TrackData: + """The tracking payload for one invocation. *extra_track_data* is added last.""" + environment_id = _try_get_environment_id() + track_data: TrackData = { + "runId": run_id or str(uuid.uuid4()), + "configKey": config_key, + "variationKey": meta.get("variationKey", "") if isinstance(meta, dict) else "", + "version": meta.get("version", 1) if isinstance(meta, dict) else 1, + "modelName": config.get("model", {}).get("name", "") + if isinstance(config, dict) + else "", + "providerName": config.get("provider", {}).get("name", "") + if isinstance(config, dict) + else "", + **model_stamps_from_meta(meta), + } + if graph_key: + track_data["graphKey"] = graph_key + if environment_id: + track_data["environmentId"] = environment_id + if extra_track_data: + track_data.update(extra_track_data) + return track_data + + +def _merge_variables( + variables: dict[str, Any] | None, + user_context: LDContext, + track_data: TrackData, + span_attributes: dict[str, Any] | None, +) -> dict[str, Any]: + merged: dict[str, Any] = { + **{k: v for k, v in (variables or {}).items() if k != SPAN_ATTRIBUTES_KEY}, + "ldContext": {**user_context}, + "__ld": track_data, + } + if span_attributes: + merged[SPAN_ATTRIBUTES_KEY] = dict(span_attributes) + return merged + + def wrap_tool_handlers( tool_handlers: dict[str, Callable[..., Any] | NativeTool] | None, user_context: LDContext, @@ -167,6 +220,9 @@ async def execute_and_track( variables: dict[str, Any] | None = None, graph_key: str | None = None, history: list[dict[str, Any]] | None = None, + run_id: str | None = None, + extra_track_data: dict[str, Any] | None = None, + span_attributes: dict[str, Any] | None = None, ) -> dict[str, Any]: """ Calls *handler*, emits LD telemetry events, and returns @@ -174,24 +230,14 @@ async def execute_and_track( """ from .lifecycle import get_client - environment_id = _try_get_environment_id() - track_data: TrackData = { - "runId": str(uuid.uuid4()), - "configKey": config_key, - "variationKey": meta.get("variationKey", "") if isinstance(meta, dict) else "", - "version": meta.get("version", 1) if isinstance(meta, dict) else 1, - "modelName": config.get("model", {}).get("name", "") - if isinstance(config, dict) - else "", - "providerName": config.get("provider", {}).get("name", "") - if isinstance(config, dict) - else "", - **model_stamps_from_meta(meta), - } - if graph_key: - track_data["graphKey"] = graph_key - if environment_id: - track_data["environmentId"] = environment_id + track_data = build_track_data( + config_key=config_key, + config=config, + meta=meta, + graph_key=graph_key, + run_id=run_id, + extra_track_data=extra_track_data, + ) client = get_client() ld_ctx = to_ld_context(client, user_context) @@ -207,11 +253,9 @@ async def execute_and_track( ld_ctx, track_data, ) - merged_variables: dict[str, Any] = { - **(variables or {}), - "ldContext": {**user_context}, - "__ld": track_data, - } + merged_variables = _merge_variables( + variables, user_context, track_data, span_attributes + ) start_time = time.monotonic() @@ -266,6 +310,9 @@ async def execute_and_stream( variables: dict[str, Any] | None = None, graph_key: str | None = None, history: list[dict[str, Any]] | None = None, + run_id: str | None = None, + extra_track_data: dict[str, Any] | None = None, + span_attributes: dict[str, Any] | None = None, ) -> AsyncGenerator[ExecuteStreamEvent, None]: """ Streaming counterpart to ``execute_and_track``. Yields ``chunk`` events, @@ -278,24 +325,14 @@ async def execute_and_stream( """ from .lifecycle import get_client - environment_id = _try_get_environment_id() - track_data: TrackData = { - "runId": str(uuid.uuid4()), - "configKey": config_key, - "variationKey": meta.get("variationKey", "") if isinstance(meta, dict) else "", - "version": meta.get("version", 1) if isinstance(meta, dict) else 1, - "modelName": config.get("model", {}).get("name", "") - if isinstance(config, dict) - else "", - "providerName": config.get("provider", {}).get("name", "") - if isinstance(config, dict) - else "", - **model_stamps_from_meta(meta), - } - if graph_key: - track_data["graphKey"] = graph_key - if environment_id: - track_data["environmentId"] = environment_id + track_data = build_track_data( + config_key=config_key, + config=config, + meta=meta, + graph_key=graph_key, + run_id=run_id, + extra_track_data=extra_track_data, + ) client = get_client() ld_ctx = to_ld_context(client, user_context) @@ -307,11 +344,9 @@ async def execute_and_stream( ld_ctx, track_data, ) - merged_variables: dict[str, Any] = { - **(variables or {}), - "ldContext": {**user_context}, - "__ld": track_data, - } + merged_variables = _merge_variables( + variables, user_context, track_data, span_attributes + ) start_time = time.monotonic() full_text = "" diff --git a/packages/client/src/launchdarkly_ai_server/utils.py b/packages/client/src/launchdarkly_ai_server/utils.py index 5e1fabba..8fa5c1c6 100644 --- a/packages/client/src/launchdarkly_ai_server/utils.py +++ b/packages/client/src/launchdarkly_ai_server/utils.py @@ -806,6 +806,10 @@ def set_ld_span_attributes(span: Any, variables: dict[str, Any] | None) -> None: span.set_attribute("launchdarkly.run.id", ld.get("runId", "")) if ld.get("graphKey"): span.set_attribute("launchdarkly.graph.key", ld["graphKey"]) + extra = variables.get("__ld_span") + if isinstance(extra, dict): + for name, value in extra.items(): + span.set_attribute(name, value) feature_flag_attrs: dict[str, str] = { "feature_flag.key": ld.get("configKey", ""), diff --git a/packages/client/tests/test_shadow_requests.py b/packages/client/tests/test_shadow_requests.py new file mode 100644 index 00000000..ac9ba20d --- /dev/null +++ b/packages/client/tests/test_shadow_requests.py @@ -0,0 +1,652 @@ +""" +Tests for §9 Shadow Requests (experimental). +Reference: TESTING.md §9, A.22 +""" + +import asyncio +import logging +import pickle +import time +from collections.abc import Iterator +from types import SimpleNamespace +from typing import Any +from unittest.mock import AsyncMock, MagicMock + +import pytest + +import launchdarkly_ai_server.lifecycle as lifecycle_module +from launchdarkly_ai_server import ProviderHandler, config, shutdown +from launchdarkly_ai_server.experimental import shadow_requests as sr +from launchdarkly_ai_server.invocation_hooks import register_invocation_hook +from launchdarkly_ai_server.sdk_usage import reset_sdk_usage +from launchdarkly_ai_server.utils import set_ld_span_attributes + +CONTEXT = {"kind": "user", "key": "u1"} +FUTURE = int(time.time() * 1000) + 3_600_000 +PAST = int(time.time() * 1000) - 1 +_SHADOW_FIELDS = { + "isShadow", + "primaryVariationKey", + "shadowRequestId", + "ownerKind", + "ownerId", +} +_BOOKKEEPING = {"$ld:ai:sdk:usage", "$ld:ai:sdk:info"} + + +def _entry(**overrides: Any) -> dict[str, Any]: + entry = { + "id": "shadow-1", + "variationKey": "b", + "samplingRate": 1.0, + "endsAt": FUTURE, + "ownerKind": "member", + "ownerId": "member-1", + } + entry.update(overrides) + return entry + + +def _variation_b() -> dict[str, Any]: + return { + "_ldMeta": { + "enabled": True, + "variationKey": "b", + "version": 7, + "mode": "messages", + "modelKey": "claude", + "modelVersion": 2, + }, + "model": {"name": "claude-x"}, + "provider": {"name": "ShadowProvider"}, + "instructions": "Shadow.", + "tools": {"search": {"name": "search", "type": "function", "parameters": {}}}, + "judgeConfiguration": {"judges": [{"key": "j", "samplingRate": 1}]}, + } + + +def _variation_a( + shadows: dict[str, Any] | None = None, variations: dict[str, Any] | None = None +) -> dict[str, Any]: + return { + "_ldMeta": { + "enabled": True, + "variationKey": "a", + "version": 1, + "mode": "messages", + }, + "model": {"name": "gpt-4"}, + "provider": {"name": "TestProvider"}, + "instructions": "Be helpful.", + "shadowConfiguration": { + "shadows": shadows if shadows is not None else {"fallthrough": [_entry()]}, + "variations": variations + if variations is not None + else {"b": _variation_b()}, + }, + } + + +def _detail(value: dict[str, Any], reason: dict[str, Any] | None) -> SimpleNamespace: + return SimpleNamespace(value=value, reason=reason) + + +FALLTHROUGH = {"kind": "FALLTHROUGH"} + + +def _make_client( + value: dict[str, Any], reason: dict[str, Any] | None = FALLTHROUGH +) -> MagicMock: + c = MagicMock() + c.track = MagicMock() + c.flush = AsyncMock() + c.close = AsyncMock() + c.variation = AsyncMock(return_value=value) + c.variation_detail = AsyncMock(return_value=_detail(value, reason)) + return c + + +class _Recorder: + def __init__( + self, + provider: str, + output: str = "ok", + *, + raises: bool = False, + sleep_s: float = 0, + ) -> None: + self.calls: list[dict[str, Any]] = [] + self._output = output + self._raises = raises + self._sleep_s = sleep_s + self.handler = ProviderHandler(fn=self._fn, provides_for=(provider, "messages")) # type: ignore[arg-type] + + async def _fn( + self, + cfg: Any, + user_input: Any, + tool_handlers: Any, + variables: Any, + history: Any = None, + ) -> dict[str, Any]: + self.calls.append( + { + "config": cfg, + "input": user_input, + "tools": tool_handlers, + "variables": variables, + "history": history, + } + ) + if self._sleep_s: + await asyncio.sleep(self._sleep_s) + if self._raises: + raise RuntimeError("shadow boom") + return { + "output": self._output, + "usage": {"input_tokens": 3, "output_tokens": 2}, + } + + +def _events(client: MagicMock) -> list[tuple[str, dict[str, Any], Any]]: + return [ + (c[0][0], c[0][2], c[0][3]) + for c in client.track.call_args_list + if c[0][0] not in _BOOKKEEPING + ] + + +def _shadow_events(client: MagicMock) -> list[tuple[str, dict[str, Any], Any]]: + return [e for e in _events(client) if e[1].get("isShadow")] + + +def _real_events(client: MagicMock) -> list[tuple[str, dict[str, Any], Any]]: + return [e for e in _events(client) if not e[1].get("isShadow")] + + +@pytest.fixture(autouse=True) +def _reset() -> Iterator[None]: + lifecycle_module._reset_for_testing() + reset_sdk_usage() + sr.disable_shadow_requests() + yield + sr.disable_shadow_requests() + lifecycle_module._reset_for_testing() + + +def _use(client: MagicMock) -> MagicMock: + lifecycle_module._set_client_for_testing(client) + return client + + +async def _invoke( + client: MagicMock, + *, + user_input: str = "hi", + history: list[dict[str, Any]] | None = None, +) -> tuple[Any, _Recorder, _Recorder]: + _use(client) + a = _Recorder("TestProvider", "real answer") + b = _Recorder("ShadowProvider", "shadow answer") + response = await config(key="flag", handler=[a.handler, b.handler]).invoke( + user_input, CONTEXT, {"topic": "x"}, history + ) + await sr._wait_for_shadows_for_testing() + return response, a, b + + +# --------------------------------------------------------------------------- +# §9.1 Turning Shadow Requests on and off +# --------------------------------------------------------------------------- + + +class TestOnOff: + async def test_nothing_runs_until_enabled(self) -> None: + client = _make_client(_variation_a()) + _, a, b = await _invoke(client) + assert len(a.calls) == 1 + assert b.calls == [] + assert _shadow_events(client) == [] + client.variation_detail.assert_not_called() + client.variation.assert_called_once() + + async def test_enabled_runs_the_fallthrough_shadow(self) -> None: + sr.enable_shadow_requests() + client = _make_client(_variation_a()) + _, a, b = await _invoke(client) + assert len(a.calls) == 1 + assert len(b.calls) == 1 + client.variation_detail.assert_called_once() + client.variation.assert_not_called() + + def test_timeout_must_be_positive(self) -> None: + with pytest.raises(ValueError): + sr.enable_shadow_requests(timeout_s=0) + + async def test_disable_is_the_kill_switch(self) -> None: + sr.enable_shadow_requests() + sr.disable_shadow_requests() + client = _make_client(_variation_a()) + _, _, b = await _invoke(client) + assert b.calls == [] + + async def test_disable_stops_deferred_mode(self) -> None: + tasks: list[sr.ShadowTask] = [] + sr.enable_shadow_requests(on_task=tasks.append) + sr.disable_shadow_requests() + await _invoke(_make_client(_variation_a())) + assert tasks == [] + + async def test_enabling_again_replaces_settings(self) -> None: + tasks: list[sr.ShadowTask] = [] + sr.enable_shadow_requests(on_task=tasks.append) + sr.enable_shadow_requests() + _, _, b = await _invoke(_make_client(_variation_a())) + assert tasks == [] + assert len(b.calls) == 1 + + async def test_shutdown_disables(self) -> None: + sr.enable_shadow_requests() + _use(_make_client(_variation_a())) + await shutdown() + client = _make_client(_variation_a()) + _, _, b = await _invoke(client) + assert b.calls == [] + + async def test_enable_reports_usage(self) -> None: + client = _use(_make_client(_variation_a())) + sr.enable_shadow_requests() + helpers = [ + c[0][2].get("helper") + for c in client.track.call_args_list + if c[0][0] == "$ld:ai:sdk:usage" + ] + assert "experimental.enableShadowRequests" in helpers + + +# --------------------------------------------------------------------------- +# §9.2 Payload shape +# --------------------------------------------------------------------------- + + +class TestPayload: + async def test_shadow_configuration_never_reaches_the_handler(self) -> None: + sr.enable_shadow_requests() + _, a, b = await _invoke(_make_client(_variation_a())) + assert "shadowConfiguration" not in a.calls[0]["config"] + assert "shadowConfiguration" not in b.calls[0]["config"] + + @pytest.mark.parametrize( + "shadow_configuration", + [ + "nope", + {"shadows": {"fallthrough": "nope"}, "variations": {"b": _variation_b()}}, + { + "shadows": {"fallthrough": [{"id": "x"}]}, + "variations": {"b": _variation_b()}, + }, + { + "shadows": {"fallthrough": [_entry(samplingRate="high")]}, + "variations": {"b": _variation_b()}, + }, + ], + ) + async def test_malformed_configuration_is_ignored( + self, shadow_configuration: Any + ) -> None: + sr.enable_shadow_requests() + value = _variation_a() + value["shadowConfiguration"] = shadow_configuration + response, _, b = await _invoke(_make_client(value)) + assert response.response == "real answer" + assert b.calls == [] + + +# --------------------------------------------------------------------------- +# §9.3 Selecting shadows +# --------------------------------------------------------------------------- + + +class TestSelection: + async def test_rule_match_uses_only_that_rule(self) -> None: + sr.enable_shadow_requests() + shadows = { + "r1": [_entry(id="rule-shadow")], + "fallthrough": [_entry(id="ft-shadow")], + } + client = _make_client( + _variation_a(shadows), + {"kind": "RULE_MATCH", "ruleId": "r1", "ruleIndex": 0}, + ) + await _invoke(client) + ids = {e[1]["shadowRequestId"] for e in _shadow_events(client)} + assert ids == {"rule-shadow"} + + @pytest.mark.parametrize( + "reason", + [ + {"kind": "TARGET_MATCH"}, + {"kind": "OFF"}, + {"kind": "PREREQUISITE_FAILED"}, + {"kind": "ERROR"}, + None, + ], + ) + async def test_other_reasons_run_nothing( + self, reason: dict[str, Any] | None + ) -> None: + sr.enable_shadow_requests() + _, _, b = await _invoke(_make_client(_variation_a(), reason)) + assert b.calls == [] + + async def test_entry_past_its_end_is_skipped(self) -> None: + sr.enable_shadow_requests() + shadows = {"fallthrough": [_entry(id="ended", endsAt=PAST), _entry(id="live")]} + client = _make_client(_variation_a(shadows)) + await _invoke(client) + ids = {e[1]["shadowRequestId"] for e in _shadow_events(client)} + assert ids == {"live"} + + async def test_missing_variation_is_skipped(self) -> None: + sr.enable_shadow_requests() + _, _, b = await _invoke( + _make_client(_variation_a({"fallthrough": [_entry(variationKey="gone")]})) + ) + assert b.calls == [] + + async def test_no_matching_handler_is_skipped(self) -> None: + sr.enable_shadow_requests() + client = _use(_make_client(_variation_a())) + a = _Recorder("TestProvider", "real answer") + response = await config(key="flag", handler=a.handler).invoke("hi", CONTEXT) + await sr._wait_for_shadows_for_testing() + assert response.response == "real answer" + assert _shadow_events(client) == [] + + def test_sampling_is_sticky(self) -> None: + assert sr._in_sample("s1", CONTEXT, 0.5) == sr._in_sample("s1", CONTEXT, 0.5) + + def test_sampling_honors_the_rate(self) -> None: + hits = sum( + sr._in_sample("s1", {"kind": "user", "key": f"k{i}"}, 0.2) + for i in range(10_000) + ) + assert 1800 <= hits <= 2200 + + def test_sampling_edges(self) -> None: + contexts = [{"kind": "user", "key": f"k{i}"} for i in range(200)] + assert all(sr._in_sample("s1", c, 1.0) for c in contexts) + assert not any(sr._in_sample("s1", c, 0.0) for c in contexts) + + def test_context_without_key_is_not_sampled(self) -> None: + assert not sr._in_sample("s1", {"kind": "user"}, 1.0) + + async def test_stream_runs_the_shadow_as_invoke(self) -> None: + sr.enable_shadow_requests() + _use(_make_client(_variation_a())) + a = _Recorder("TestProvider", "real answer") + b = _Recorder("ShadowProvider", "shadow answer") + events = [ + e + async for e in config(key="flag", handler=[a.handler, b.handler]).stream( + "hi", CONTEXT + ) + ] + await sr._wait_for_shadows_for_testing() + assert events[-1]["type"] == "done" + assert len(b.calls) == 1 + + +# --------------------------------------------------------------------------- +# §9.4 Running a shadow +# --------------------------------------------------------------------------- + + +class TestRunning: + async def test_same_input_and_stripped_config(self) -> None: + sr.enable_shadow_requests() + history = [{"role": "user", "content": "earlier"}] + _, _, b = await _invoke( + _make_client(_variation_a()), user_input="question", history=history + ) + call = b.calls[0] + assert call["input"] == "question" + assert call["history"] == history + assert call["variables"]["topic"] == "x" + assert "tools" not in call["config"] + assert "judgeConfiguration" not in call["config"] + assert not call["tools"] + + async def test_shadow_answer_never_joins_history(self) -> None: + sr.enable_shadow_requests() + history = [{"role": "user", "content": "earlier"}] + await _invoke(_make_client(_variation_a()), history=history) + assert history == [{"role": "user", "content": "earlier"}] + + async def test_failing_shadow_never_reaches_the_caller(self) -> None: + sr.enable_shadow_requests() + client = _use(_make_client(_variation_a())) + a = _Recorder("TestProvider", "real answer") + b = _Recorder("ShadowProvider", raises=True) + response = await config(key="flag", handler=[a.handler, b.handler]).invoke( + "hi", CONTEXT + ) + await sr._wait_for_shadows_for_testing() + assert response.response == "real answer" + assert "$ld:ai:generation:success" in [e[0] for e in _real_events(client)] + assert "$ld:ai:generation:error" in [e[0] for e in _shadow_events(client)] + + async def test_timeout_records_an_error(self) -> None: + sr.enable_shadow_requests(timeout_s=0.05) + client = _use(_make_client(_variation_a())) + a = _Recorder("TestProvider", "real answer") + b = _Recorder("ShadowProvider", sleep_s=1) + response = await config(key="flag", handler=[a.handler, b.handler]).invoke( + "hi", CONTEXT + ) + await sr._wait_for_shadows_for_testing() + assert response.response == "real answer" + names = [e[0] for e in _shadow_events(client)] + assert "$ld:ai:generation:error" in names + assert "$ld:ai:duration:total" in names + + async def test_real_run_is_unchanged(self) -> None: + baseline_client = _make_client(_variation_a()) + baseline, _, _ = await _invoke(baseline_client) + lifecycle_module._reset_for_testing() + sr.enable_shadow_requests() + client = _make_client(_variation_a()) + response, _, _ = await _invoke(client) + assert response.response == baseline.response + assert response.usage == baseline.usage + assert [e[0] for e in _real_events(client)] == [ + e[0] for e in _real_events(baseline_client) + ] + + +# --------------------------------------------------------------------------- +# §9.5 Automatic and deferred modes +# --------------------------------------------------------------------------- + + +class TestDeferred: + async def test_deferred_runs_nothing_inline(self) -> None: + tasks: list[sr.ShadowTask] = [] + sr.enable_shadow_requests(on_task=tasks.append) + client = _make_client(_variation_a()) + _, a, b = await _invoke(client) + assert len(a.calls) == 1 + assert b.calls == [] + assert len(tasks) == 1 + assert _shadow_events(client) == [] + + async def test_task_survives_pickle(self) -> None: + tasks: list[sr.ShadowTask] = [] + sr.enable_shadow_requests(on_task=tasks.append) + await _invoke(_make_client(_variation_a())) + restored = pickle.loads(pickle.dumps(tasks[0])) + assert restored == tasks[0] + assert restored.shadow_request_id == "shadow-1" + + async def test_run_shadow_matches_automatic_mode(self) -> None: + tasks: list[sr.ShadowTask] = [] + sr.enable_shadow_requests(on_task=tasks.append) + deferred_client = _make_client(_variation_a()) + await _invoke(deferred_client) + b = _Recorder("ShadowProvider", "shadow answer") + result = await sr.run_shadow(tasks[0], [b.handler]) + deferred = [ + (e[0], {k: v for k, v in e[1].items() if k != "runId"}) + for e in _shadow_events(deferred_client) + ] + assert result.error is None + assert result.track_data["isShadow"] is True + + lifecycle_module._reset_for_testing() + sr.enable_shadow_requests() + auto_client = _make_client(_variation_a()) + await _invoke(auto_client) + automatic = [ + (e[0], {k: v for k, v in e[1].items() if k != "runId"}) + for e in _shadow_events(auto_client) + ] + assert deferred == automatic + + async def test_raising_callback_does_not_fail_the_call( + self, caplog: pytest.LogCaptureFixture + ) -> None: + def boom(task: sr.ShadowTask) -> None: + raise RuntimeError("queue full") + + sr.enable_shadow_requests(on_task=boom) + with caplog.at_level(logging.WARNING): + response, a, _ = await _invoke(_make_client(_variation_a())) + assert response.response == "real answer" + assert len(a.calls) == 1 + + async def test_run_shadow_never_raises(self) -> None: + tasks: list[sr.ShadowTask] = [] + sr.enable_shadow_requests(on_task=tasks.append) + await _invoke(_make_client(_variation_a())) + result = await sr.run_shadow( + tasks[0], [_Recorder("ShadowProvider", raises=True).handler] + ) + assert result.error + + async def test_run_shadow_reports_usage(self) -> None: + tasks: list[sr.ShadowTask] = [] + sr.enable_shadow_requests(on_task=tasks.append) + client = _make_client(_variation_a()) + await _invoke(client) + await sr.run_shadow(tasks[0], [_Recorder("ShadowProvider").handler]) + helpers = [ + c[0][2].get("helper") + for c in client.track.call_args_list + if c[0][0] == "$ld:ai:sdk:usage" + ] + assert "experimental.runShadow" in helpers + + +# --------------------------------------------------------------------------- +# §9.6 Core invocation hook +# --------------------------------------------------------------------------- + + +class TestInvocationHook: + async def test_raising_hook_does_not_fail_the_call( + self, caplog: pytest.LogCaptureFixture + ) -> None: + def bad_hook(event: Any) -> None: + raise RuntimeError("hook boom") + + register_invocation_hook("test-bad-hook", bad_hook) + try: + with caplog.at_level(logging.WARNING): + response, a, _ = await _invoke(_make_client(_variation_a())) + assert response.response == "real answer" + assert len(a.calls) == 1 + assert any( + "hook boom" in r.getMessage() or r.exc_info for r in caplog.records + ) + finally: + register_invocation_hook("test-bad-hook", None) + + async def test_hook_sees_the_run_id_of_the_real_run(self) -> None: + seen: list[Any] = [] + register_invocation_hook("test-seen", seen.append) + try: + response, _, _ = await _invoke(_make_client(_variation_a())) + assert seen[0].run_id == response.track_data["runId"] + assert seen[0].reason == FALLTHROUGH + finally: + register_invocation_hook("test-seen", None) + + +# --------------------------------------------------------------------------- +# §9.7 Shadow events and spans +# --------------------------------------------------------------------------- + + +class TestEvents: + async def test_shadow_events_describe_b(self) -> None: + sr.enable_shadow_requests() + client = _make_client(_variation_a()) + response, _, _ = await _invoke(client) + events = _shadow_events(client) + assert {e[0] for e in events} >= { + "$ld:ai:duration:total", + "$ld:ai:generation:success", + "$ld:ai:tokens:total", + } + data = events[0][1] + assert data["configKey"] == "flag" + assert data["variationKey"] == "b" + assert data["version"] == 7 + assert data["modelName"] == "claude-x" + assert data["providerName"] == "ShadowProvider" + assert data["modelKey"] == "claude" + assert data["modelVersion"] == 2 + assert data["runId"] == response.track_data["runId"] + assert data["isShadow"] is True + assert data["primaryVariationKey"] == "a" + assert data["shadowRequestId"] == "shadow-1" + assert data["ownerKind"] == "member" + assert data["ownerId"] == "member-1" + + async def test_nothing_is_added_to_real_events(self) -> None: + sr.enable_shadow_requests() + client = _make_client(_variation_a()) + response, _, _ = await _invoke(client) + assert not (_SHADOW_FIELDS & set(response.track_data)) + for _, data, _ in _real_events(client): + assert not (_SHADOW_FIELDS & set(data)) + + def test_span_attributes_are_set_on_the_root_span( + self, mock_span: MagicMock + ) -> None: + set_ld_span_attributes( + mock_span, + { + "__ld": {"configKey": "flag", "variationKey": "b", "runId": "r"}, + "__ld_span": {"launchdarkly.shadow": True}, + }, + ) + mock_span.set_attribute.assert_any_call("launchdarkly.shadow", True) + + async def test_shadow_span_attributes_reach_the_handler(self) -> None: + sr.enable_shadow_requests() + _, _, b = await _invoke(_make_client(_variation_a())) + span_attrs = b.calls[0]["variables"]["__ld_span"] + assert span_attrs == { + "launchdarkly.shadow": True, + "launchdarkly.shadow.request.id": "shadow-1", + "launchdarkly.shadow.primary.variation.key": "a", + } + + async def test_template_variables_cannot_set_span_attributes(self) -> None: + _use(_make_client(_variation_a())) + a = _Recorder("TestProvider", "real answer") + await config(key="flag", handler=a.handler).invoke( + "hi", CONTEXT, {"__ld_span": {"x": 1}} + ) + assert "__ld_span" not in a.calls[0]["variables"] diff --git a/release-please-config.json b/release-please-config.json index 9c94b155..ba490b9a 100644 --- a/release-please-config.json +++ b/release-please-config.json @@ -161,6 +161,9 @@ { "type": "feat", "scope": "evaluations", "section": "Experimental" }, { "type": "fix", "scope": "evaluations", "section": "Experimental" }, { "type": "perf", "scope": "evaluations", "section": "Experimental" }, + { "type": "feat", "scope": "shadow-requests", "section": "Experimental" }, + { "type": "fix", "scope": "shadow-requests", "section": "Experimental" }, + { "type": "perf", "scope": "shadow-requests", "section": "Experimental" }, { "type": "feat", "section": "Features" }, { "type": "fix", "section": "Bug Fixes" }, { "type": "perf", "section": "Performance Improvements" }, From c9642201e87c0b1952bea0f4905236994f1c0378 Mon Sep 17 00:00:00 2001 From: Harsh Sadhvani Date: Wed, 28 Oct 2026 13:56:29 -0700 Subject: [PATCH 2/4] fix(shadow-requests): name the shadow span attribute launchdarkly.is_shadow --- .../launchdarkly_ai_server/experimental/shadow_requests.py | 2 +- packages/client/tests/test_shadow_requests.py | 6 +++--- 2 files changed, 4 insertions(+), 4 deletions(-) diff --git a/packages/client/src/launchdarkly_ai_server/experimental/shadow_requests.py b/packages/client/src/launchdarkly_ai_server/experimental/shadow_requests.py index ffb61dd1..78ef0d1d 100644 --- a/packages/client/src/launchdarkly_ai_server/experimental/shadow_requests.py +++ b/packages/client/src/launchdarkly_ai_server/experimental/shadow_requests.py @@ -239,7 +239,7 @@ def _on_invocation(event: InvocationEvent) -> None: "ownerId": entry.get("ownerId", ""), }, span_attributes={ - "launchdarkly.shadow": True, + "launchdarkly.is_shadow": True, "launchdarkly.shadow.request.id": entry["id"], "launchdarkly.shadow.primary.variation.key": event.variation_key, }, diff --git a/packages/client/tests/test_shadow_requests.py b/packages/client/tests/test_shadow_requests.py index ac9ba20d..6c1837de 100644 --- a/packages/client/tests/test_shadow_requests.py +++ b/packages/client/tests/test_shadow_requests.py @@ -628,17 +628,17 @@ def test_span_attributes_are_set_on_the_root_span( mock_span, { "__ld": {"configKey": "flag", "variationKey": "b", "runId": "r"}, - "__ld_span": {"launchdarkly.shadow": True}, + "__ld_span": {"launchdarkly.is_shadow": True}, }, ) - mock_span.set_attribute.assert_any_call("launchdarkly.shadow", True) + mock_span.set_attribute.assert_any_call("launchdarkly.is_shadow", True) async def test_shadow_span_attributes_reach_the_handler(self) -> None: sr.enable_shadow_requests() _, _, b = await _invoke(_make_client(_variation_a())) span_attrs = b.calls[0]["variables"]["__ld_span"] assert span_attrs == { - "launchdarkly.shadow": True, + "launchdarkly.is_shadow": True, "launchdarkly.shadow.request.id": "shadow-1", "launchdarkly.shadow.primary.variation.key": "a", } From 5f6258c0693b12a3c8b101ae3743b921c16b442a Mon Sep 17 00:00:00 2001 From: Harsh Sadhvani Date: Wed, 28 Oct 2026 14:36:27 -0700 Subject: [PATCH 3/4] fix(shadow-requests): track a shadow timeout once and reject non-finite timeouts --- .../experimental/shadow_requests.py | 69 ++++++++++--------- packages/client/tests/test_shadow_requests.py | 32 +++++++++ 2 files changed, 68 insertions(+), 33 deletions(-) diff --git a/packages/client/src/launchdarkly_ai_server/experimental/shadow_requests.py b/packages/client/src/launchdarkly_ai_server/experimental/shadow_requests.py index 78ef0d1d..c3080a21 100644 --- a/packages/client/src/launchdarkly_ai_server/experimental/shadow_requests.py +++ b/packages/client/src/launchdarkly_ai_server/experimental/shadow_requests.py @@ -14,6 +14,7 @@ import asyncio import hashlib import logging +import math import time from collections.abc import Callable from dataclasses import dataclass, field @@ -30,7 +31,6 @@ model_stamps_from_meta, parse_usage, select_handler, - to_ld_context, ) __all__ = [ @@ -101,8 +101,8 @@ def enable_shadow_requests( sampled shadow, before the real handler runs, and must only enqueue it. """ report_usage("experimental.enableShadowRequests") - if not timeout_s > 0: - raise ValueError("timeout_s must be greater than 0") + if not (math.isfinite(timeout_s) and timeout_s > 0): + raise ValueError("timeout_s must be a finite number greater than 0") global _settings _settings = _Settings(timeout_s=timeout_s, on_task=on_task) register_invocation_hook(_FEATURE, _on_invocation) @@ -274,27 +274,22 @@ async def _run(task: ShadowTask, handlers: list[ProviderHandler]) -> ShadowResul usage=parse_usage({}), track_data=track_data, error=str(err) ) - start = time.monotonic() try: - result = await asyncio.wait_for( - execute_and_track( - config_key=task.config_key, - config=task.shadow_config, - meta=task.shadow_meta, - user_context=task.user_context, - handler=handler, - user_input=task.user_input, - tool_handlers=None, - variables=task.variables, - history=list(task.history) if task.history else None, - run_id=run_id, - extra_track_data=extra, - span_attributes=task.span_attributes, - ), - timeout=task.timeout_s, + result = await execute_and_track( + config_key=task.config_key, + config=task.shadow_config, + meta=task.shadow_meta, + user_context=task.user_context, + handler=_with_deadline(handler, task.timeout_s), + user_input=task.user_input, + tool_handlers=None, + variables=task.variables, + history=list(task.history) if task.history else None, + run_id=run_id, + extra_track_data=extra, + span_attributes=task.span_attributes, ) - except TimeoutError: - _track_timeout(task, track_data, int((time.monotonic() - start) * 1000)) + except _ShadowTimeout: return ShadowResult( usage=parse_usage({}), track_data=track_data, error="timeout" ) @@ -305,15 +300,23 @@ async def _run(task: ShadowTask, handlers: list[ProviderHandler]) -> ShadowResul return ShadowResult(usage=result["usage"], track_data=result["track_data"]) -def _track_timeout( - task: ShadowTask, track_data: dict[str, Any], elapsed_ms: int -) -> None: - from ..lifecycle import get_client +class _ShadowTimeout(Exception): + pass - try: - client = get_client() - ld_ctx = to_ld_context(client, task.user_context) - client.track("$ld:ai:duration:total", ld_ctx, track_data, elapsed_ms) - client.track("$ld:ai:generation:error", ld_ctx, track_data, 1) - except Exception: - logger.debug("Could not track a shadow timeout", exc_info=True) + +def _with_deadline(handler: ProviderHandler, timeout_s: float) -> ProviderHandler: + async def timed(*args: Any) -> Any: + deadline = asyncio.timeout(timeout_s) + try: + async with deadline: + return await handler(*args) + except TimeoutError: + if deadline.expired(): + raise _ShadowTimeout("timeout") from None + raise + + return ProviderHandler( + fn=timed, + provides_for=handler.provides_for, + capture_content=handler.capture_content, + ) diff --git a/packages/client/tests/test_shadow_requests.py b/packages/client/tests/test_shadow_requests.py index 6c1837de..4ddb669f 100644 --- a/packages/client/tests/test_shadow_requests.py +++ b/packages/client/tests/test_shadow_requests.py @@ -114,11 +114,13 @@ def __init__( *, raises: bool = False, sleep_s: float = 0, + error: Exception | None = None, ) -> None: self.calls: list[dict[str, Any]] = [] self._output = output self._raises = raises self._sleep_s = sleep_s + self._error = error self.handler = ProviderHandler(fn=self._fn, provides_for=(provider, "messages")) # type: ignore[arg-type] async def _fn( @@ -140,6 +142,8 @@ async def _fn( ) if self._sleep_s: await asyncio.sleep(self._sleep_s) + if self._error is not None: + raise self._error if self._raises: raise RuntimeError("shadow boom") return { @@ -450,6 +454,34 @@ async def test_timeout_records_an_error(self) -> None: assert "$ld:ai:generation:error" in names assert "$ld:ai:duration:total" in names + async def test_timeout_is_tracked_once(self) -> None: + sr.enable_shadow_requests(timeout_s=0.05) + client = _use(_make_client(_variation_a())) + a = _Recorder("TestProvider", "real answer") + b = _Recorder("ShadowProvider", sleep_s=1) + await config(key="flag", handler=[a.handler, b.handler]).invoke("hi", CONTEXT) + await sr._wait_for_shadows_for_testing() + names = [e[0] for e in _shadow_events(client)] + assert names.count("$ld:ai:generation:error") == 1 + assert names.count("$ld:ai:duration:total") == 1 + + async def test_handler_timeout_error_is_tracked_once(self) -> None: + tasks: list[sr.ShadowTask] = [] + sr.enable_shadow_requests(on_task=tasks.append) + client = _make_client(_variation_a()) + await _invoke(client) + b = _Recorder("ShadowProvider", error=TimeoutError("provider timeout")) + result = await sr.run_shadow(tasks[0], [b.handler]) + names = [e[0] for e in _shadow_events(client)] + assert result.error == "provider timeout" + assert names.count("$ld:ai:generation:error") == 1 + assert names.count("$ld:ai:duration:total") == 1 + + @pytest.mark.parametrize("timeout_s", [0, -1, float("inf"), float("nan")]) + def test_timeout_must_be_positive_and_finite(self, timeout_s: float) -> None: + with pytest.raises(ValueError): + sr.enable_shadow_requests(timeout_s=timeout_s) + async def test_real_run_is_unchanged(self) -> None: baseline_client = _make_client(_variation_a()) baseline, _, _ = await _invoke(baseline_client) From 4c5efa1670079098a6d6aa94091d4ebd39f3ac59 Mon Sep 17 00:00:00 2001 From: Harsh Sadhvani Date: Thu, 29 Oct 2026 10:27:37 -0700 Subject: [PATCH 4/4] refactor(client): run shadow requests from core like judges --- packages/claude-agents/tests/test_skills.py | 1 + packages/client/README.md | 25 +- ...ai_server.experimental.shadow_requests.txt | 22 - .../api-report/launchdarkly_ai_server.txt | 28 +- .../src/launchdarkly_ai_server/__init__.py | 6 + .../src/launchdarkly_ai_server/client.py | 84 ++-- .../experimental/__init__.py | 2 +- .../experimental/shadow_requests.py | 322 ------------ .../invocation_hooks.py | 64 --- .../src/launchdarkly_ai_server/lifecycle.py | 11 +- .../src/launchdarkly_ai_server/shadows.py | 283 +++++++++++ .../src/launchdarkly_ai_server/types.py | 39 ++ packages/client/tests/conftest.py | 1 + packages/client/tests/test_agui.py | 1 + packages/client/tests/test_client.py | 1 + packages/client/tests/test_public_surface.py | 3 + packages/client/tests/test_sdk_usage.py | 1 + ...est_shadow_requests.py => test_shadows.py} | 460 +++++++----------- .../client/tests/test_stream_conversation.py | 1 + .../opencode-agents/tests/test_handler.py | 1 + release-please-config.json | 3 - 21 files changed, 597 insertions(+), 762 deletions(-) delete mode 100644 packages/client/api-report/launchdarkly_ai_server.experimental.shadow_requests.txt delete mode 100644 packages/client/src/launchdarkly_ai_server/experimental/shadow_requests.py delete mode 100644 packages/client/src/launchdarkly_ai_server/invocation_hooks.py create mode 100644 packages/client/src/launchdarkly_ai_server/shadows.py rename packages/client/tests/{test_shadow_requests.py => test_shadows.py} (62%) diff --git a/packages/claude-agents/tests/test_skills.py b/packages/claude-agents/tests/test_skills.py index 5fdb836d..b2a8a2bd 100644 --- a/packages/claude-agents/tests/test_skills.py +++ b/packages/claude-agents/tests/test_skills.py @@ -672,6 +672,7 @@ async def during(kwargs: dict[str, Any]) -> None: async def ld_client() -> Any: await shutdown() client = MagicMock() + del client.variation_detail client.track = MagicMock() client.flush = MagicMock() client.close = MagicMock() diff --git a/packages/client/README.md b/packages/client/README.md index ffa378e0..46dfafad 100644 --- a/packages/client/README.md +++ b/packages/client/README.md @@ -305,34 +305,29 @@ result = await evals.run( **The list is checked before any network I/O.** A blank key, an uppercase key, a `schema` that is not a JSON object, a schema that is not JSON-serializable (including a `NaN` or `Infinity` value), and a non-callable implementation each fail with zero requests issued. A repeated key fails too, and keys are compared without case. A `NativeTool` is valid only for a library tool, because the provider supplies its schema. -### Shadow Requests (experimental) +### Shadow requests -A Shadow Request runs another variation of an AI Config in the background on the same input, and reports it on the normal AI metrics with `isShadow: true`. LaunchDarkly attaches shadow requests to a targeting rule or the fallthrough. The shadow's response is discarded, and your response, latency and events are unchanged. Shadow runs use your provider account, so they are billed like real calls. +A shadow request runs another variation of an AI Config in the background on the same input, and reports it on the normal AI metrics with `isShadow: true`. You set shadow requests up on a targeting rule or the fallthrough in LaunchDarkly. The shadow's response is discarded, and your response, latency and events are unchanged. Shadow runs use your provider account, so they are billed like real calls. -Nothing runs until you enable it: +`config().invoke()` and `config().stream()` run them for you in the background, with a 30 second timeout. Shadows don't run for graphs or judges, and a shadow makes no tool calls and runs no judges. ```python -from launchdarkly_ai_server.experimental.shadow_requests import enable_shadow_requests - -enable_shadow_requests() # run sampled shadows in the background, 30 s timeout -enable_shadow_requests(timeout_s=10) # or pick the timeout +caller = config(key="my-ai-config", handler=handlers, shadow_timeout_s=10) ``` -Shadows run for `config().invoke()` and `config().stream()`, not for graphs or judges. A shadow makes no tool calls and runs no judges. Call `disable_shadow_requests()` to turn them off. - -To run shadows on your own queue or worker instead, pass `on_task`. It receives a picklable `ShadowTask` before your handler runs and must only enqueue it. Run the task later with `run_shadow`: +To run shadows on your own queue or worker instead, pass `skip_shadows=True`. Each response then carries `shadow_tasks`, which you can pickle and run later with `run_shadow`, the same way `skip_judges` and `run_judge` work: ```python -from launchdarkly_ai_server.experimental.shadow_requests import enable_shadow_requests, run_shadow +from launchdarkly_ai_server import config, run_shadow -enable_shadow_requests(on_task=queue.put_nowait) +result = await config(key="my-ai-config", handler=handlers, skip_shadows=True).invoke(...) +for task in result.shadow_tasks or []: + queue.put_nowait(task) # in the worker -result = await run_shadow(task, handlers=[create_openai_messages_handler()]) +outcome = await run_shadow(task, handlers) ``` -These names may change in a minor release. - ### `config(**args)` The primary entry point for AI config invocations. Accepts either a single handler or a list of handlers and routes to the correct one at invoke-time based on the flag variation's provider and mode. diff --git a/packages/client/api-report/launchdarkly_ai_server.experimental.shadow_requests.txt b/packages/client/api-report/launchdarkly_ai_server.experimental.shadow_requests.txt deleted file mode 100644 index 938940c3..00000000 --- a/packages/client/api-report/launchdarkly_ai_server.experimental.shadow_requests.txt +++ /dev/null @@ -1,22 +0,0 @@ -# Public API of launchdarkly_ai_server.experimental.shadow_requests. Generated by `make api-report`; do not edit. -disable_shadow_requests() -> None -enable_shadow_requests(*, timeout_s: float = 30.0, on_task: Callable[[ShadowTask], None] | None = None) -> None -run_shadow(task: ShadowTask, handlers: list[ProviderHandler]) -> ShadowResult [async] -ShadowResult [class, @dataclass(frozen=True)] - __init__(usage: dict[str, int], track_data: dict[str, Any], error: str | None = None) - error: str | None = None - track_data: dict[str, Any] - usage: dict[str, int] -ShadowTask [class, @dataclass(frozen=True)] - __init__(config_key: str, shadow_request_id: str, shadow_config: AiConfigRep, shadow_meta: VariationMeta, user_input: str | None, history: list[dict[str, Any]] | None, variables: dict[str, Any] | None, user_context: LDContext, timeout_s: float, track_data: dict[str, Any], span_attributes: dict[str, Any] = dict()) - config_key: str - history: list[dict[str, Any]] | None - shadow_config: AiConfigRep - shadow_meta: VariationMeta - shadow_request_id: str - span_attributes: dict[str, Any] = field(default_factory=dict) - timeout_s: float - track_data: dict[str, Any] - user_context: LDContext - user_input: str | None - variables: dict[str, Any] | None diff --git a/packages/client/api-report/launchdarkly_ai_server.txt b/packages/client/api-report/launchdarkly_ai_server.txt index 08c06615..0dd5940a 100644 --- a/packages/client/api-report/launchdarkly_ai_server.txt +++ b/packages/client/api-report/launchdarkly_ai_server.txt @@ -10,10 +10,10 @@ any_multimodal(turns: list[dict[str, Any]]) -> bool build_judge_tasks(*, config: AiConfigRep, user_context: LDContext, handler: ProviderHandler, handlers: list[ProviderHandler] | None = None, llm_response: str, base_track_data: TrackData, user_input: str | None = None, trajectory: str = '') -> list[JudgeTask] [async] compose(a: Registry, b: Registry) -> Registry compose_history(*, history: list[dict[str, Any]], user_input: str | None = None, config_messages: list[dict[str, Any]] | None = None) -> list[dict[str, Any]] -config(*, key: str, handler: ProviderHandler | list[ProviderHandler] | None = None, tool_handlers: Mapping[str, Callable[..., Any] | NativeTool] | None = None, registry: Any = None, skip_judges: bool = False) -> ConfigInstance +config(*, key: str, handler: ProviderHandler | list[ProviderHandler] | None = None, tool_handlers: Mapping[str, Callable[..., Any] | NativeTool] | None = None, registry: Any = None, skip_judges: bool = False, skip_shadows: bool = False, shadow_timeout_s: float = DEFAULT_SHADOW_TIMEOUT_S) -> ConfigInstance config_skills(key: str, context: LDContext) -> SkillSource ConfigInstance [class] - __init__(key: str, handler: ProviderHandler | list[ProviderHandler] | None, tool_handlers: Mapping[str, Callable[..., Any] | NativeTool] | None, registry: Any, skip_judges: bool = False) + __init__(key: str, handler: ProviderHandler | list[ProviderHandler] | None, tool_handlers: Mapping[str, Callable[..., Any] | NativeTool] | None, registry: Any, skip_judges: bool = False, skip_shadows: bool = False, shadow_timeout_s: float = DEFAULT_SHADOW_TIMEOUT_S) invoke(user_input: str | None, context: LDContext, variables: dict[str, Any] | None = None, history: list[dict[str, Any]] | None = None) -> ProviderResponse[Any] [async] stream(user_input: str | None, context: LDContext, variables: dict[str, Any] | None = None, history: list[dict[str, Any]] | None = None) -> AsyncGenerator[StreamEvent, None] configure(*, options: SDKOptions) -> None @@ -65,7 +65,7 @@ execute_and_stream(*, config_key: str, config: AiConfigRep, meta: VariationMeta, execute_and_track(*, config_key: str, config: AiConfigRep, meta: VariationMeta, user_context: LDContext, handler: ProviderHandler, user_input: str | None = None, tool_handlers: Mapping[str, Callable[..., Any] | NativeTool] | None = None, variables: dict[str, Any] | None = None, graph_key: str | None = None, history: list[dict[str, Any]] | None = None, run_id: str | None = None, extra_track_data: dict[str, Any] | None = None, span_attributes: dict[str, Any] | None = None) -> dict[str, Any] [async] ExecuteStreamDoneEvent = dict[str, Any] ExecuteStreamEvent = dict[str, Any] -extract_variation(config_key: str, context: dict[str, Any]) -> dict[str, Any] [async] +extract_variation(config_key: str, context: dict[str, Any], *, with_reason: bool = False) -> dict[str, Any] [async] FDv2SkillStore = LaunchDarklySkillStore GenerationConfig(TypedDict, total=False) [class] instructions: str @@ -255,10 +255,11 @@ ProviderHandler [class] provides_for: tuple[str, Literal['agent', 'messages']] | None = provides_for stream(config: AiConfigRep, user_input: str | None = None, tool_handlers: Mapping[str, Callable[..., Any] | NativeTool] | None = None, variables: dict[str, Any] | None = None, history: list[dict[str, Any]] | None = None) -> AsyncGenerator[dict[str, Any], None] [async] ProviderResponse(Generic[T]) [class, @dataclass] - __init__(response: T, usage: UsageDict, judge_results: dict[str, JudgeResult] | None = None, judge_tasks: list[JudgeTask] | None = None, track_data: TrackData | None = None) + __init__(response: T, usage: UsageDict, judge_results: dict[str, JudgeResult] | None = None, judge_tasks: list[JudgeTask] | None = None, shadow_tasks: list[ShadowTask] | None = None, track_data: TrackData | None = None) judge_results: dict[str, JudgeResult] | None = None judge_tasks: list[JudgeTask] | None = None response: T + shadow_tasks: list[ShadowTask] | None = None track_data: TrackData | None = None usage: UsageDict ReconcileAction [class, @dataclass(frozen=True)] @@ -286,6 +287,7 @@ resolve_handlers(registry: Registry | None, local_handlers: list[ProviderHandler resolve_tools(registry: Registry | None, local_tools: Mapping[str, Callable[..., Any] | NativeTool] | None) -> Mapping[str, Callable[..., Any] | NativeTool] | None run_judge(task: JudgeTask, handlers: list[ProviderHandler]) -> JudgeRunResult | None [async] run_judges(*, config: AiConfigRep, user_context: LDContext, handler: ProviderHandler, handlers: list[ProviderHandler] | None = None, user_input: str | None, llm_response: str, base_track_data: TrackData, tool_handlers: Mapping[str, Callable[..., Any] | NativeTool] | None = None, graph_key: str | None = None, trajectory: str = '') -> dict[str, JudgeResult] [async] +run_shadow(task: ShadowTask, handlers: list[ProviderHandler]) -> ShadowResult [async] RunSummary [class, @dataclass] __init__(total_rows: int = 0, passed_rows: int = 0, failed_rows: int = 0, error_rows: int = 0, pending_rows: int = 0) error_rows: int = 0 @@ -324,6 +326,24 @@ set_skill_store(store: SkillStore | None) -> None set_tool_call_content_attributes(span: Any, capture: bool, *, arguments: Any = None, result: Any = None) -> None set_tool_definition_attributes(span: Any, capture: bool, tools: list[ToolDefinitionInput]) -> None set_usage_span_attributes(span: Any, usage: SpanUsage) -> None +ShadowResult [class, @dataclass(frozen=True)] + __init__(usage: dict[str, int], track_data: TrackData, error: str | None = None) + error: str | None = None + track_data: TrackData + usage: dict[str, int] +ShadowTask [class, @dataclass(frozen=True)] + __init__(config_key: str, shadow_request_id: str, shadow_config: AiConfigRep, shadow_meta: VariationMeta, user_input: str | None, history: list[dict[str, Any]] | None, variables: dict[str, Any] | None, user_context: LDContext, timeout_s: float, track_data: TrackData, span_attributes: dict[str, Any] = dict()) + config_key: str + history: list[dict[str, Any]] | None + shadow_config: AiConfigRep + shadow_meta: VariationMeta + shadow_request_id: str + span_attributes: dict[str, Any] = field(default_factory=dict) + timeout_s: float + track_data: TrackData + user_context: LDContext + user_input: str | None + variables: dict[str, Any] | None shutdown() -> None [async] Skill [class, @dataclass(frozen=True)] __init__(key: str, version: int, content: bytes, content_hash: str, name: str | None = None, description: str | None = None) diff --git a/packages/client/src/launchdarkly_ai_server/__init__.py b/packages/client/src/launchdarkly_ai_server/__init__.py index 48355f8c..84c70c96 100644 --- a/packages/client/src/launchdarkly_ai_server/__init__.py +++ b/packages/client/src/launchdarkly_ai_server/__init__.py @@ -71,6 +71,7 @@ ) from .sdk_info import SDK_INFO_CONTEXT, SDK_INFO_EVENT, register_ai_sdk_package from .sdk_usage import report_usage +from .shadows import run_shadow from .skills import ( InMemorySkillStore, all_skills, @@ -128,6 +129,8 @@ ReconcileActionKind, ReconcileReport, SDKOptions, + ShadowResult, + ShadowTask, Skill, SkillOutcome, SkillOutcomeReason, @@ -200,6 +203,8 @@ "ProviderGraphResponse", "ProviderHandler", "ProviderResponse", + "ShadowResult", + "ShadowTask", "StreamChunkEvent", "StreamDoneEvent", "StreamEvent", @@ -275,6 +280,7 @@ # judges "build_judge_tasks", "run_judge", + "run_shadow", "run_judges", # client "config", diff --git a/packages/client/src/launchdarkly_ai_server/client.py b/packages/client/src/launchdarkly_ai_server/client.py index 23fd8bf6..618a5f17 100644 --- a/packages/client/src/launchdarkly_ai_server/client.py +++ b/packages/client/src/launchdarkly_ai_server/client.py @@ -1,5 +1,6 @@ from __future__ import annotations +import math import uuid from collections.abc import AsyncGenerator, Callable, Mapping from contextlib import aclosing @@ -7,16 +8,12 @@ from .agui import StableProjector, judge_completed_stream from .conversation import bind_conversation_id -from .invocation_hooks import ( - InvocationEvent, - has_invocation_hooks, - run_invocation_hooks, -) from .judge_scoring import judge_text from .judges import _build_judge_tasks, run_judges from .lifecycle import extract_variation from .registry import resolve_handlers, resolve_tools from .sdk_usage import report_usage +from .shadows import DEFAULT_SHADOW_TIMEOUT_S, build_shadow_tasks, start_shadows from .tracking import execute_and_stream, execute_and_track from .types import ( AiConfigRep, @@ -24,6 +21,7 @@ NativeTool, ProviderHandler, ProviderResponse, + ShadowTask, StreamEvent, VariationMeta, ) @@ -51,7 +49,7 @@ def _resolve_output_format_response( return parsed if parsed is not None else raw_response -def _run_hooks( +def _shadow_tasks( *, config_key: str, variation: dict[str, Any], @@ -62,25 +60,20 @@ def _run_hooks( variables: dict[str, Any] | None, handlers: list[ProviderHandler], run_id: str, -) -> None: - if not has_invocation_hooks(): - return - raw = variation.get("raw") - run_invocation_hooks( - InvocationEvent( - config_key=config_key, - value=raw if isinstance(raw, dict) else {}, - reason=variation.get("reason"), - context=context, - user_input=user_input, - history=history, - variables=variables, - handlers=handlers, - run_id=run_id, - variation_key=str(meta.get("variationKey", "")) - if isinstance(meta, dict) - else "", - ) + timeout_s: float, +) -> list[ShadowTask]: + return build_shadow_tasks( + config_key=config_key, + raw_variation=variation.get("raw"), + reason=variation.get("reason"), + primary_variation_key=str(meta.get("variationKey", "")), + context=context, + user_input=user_input, + history=history, + variables=variables, + handlers=handlers, + run_id=run_id, + timeout_s=timeout_s, ) @@ -94,12 +87,18 @@ def __init__( tool_handlers: Mapping[str, Callable[..., Any] | NativeTool] | None, registry: Any, # Registry | None skip_judges: bool = False, + skip_shadows: bool = False, + shadow_timeout_s: float = DEFAULT_SHADOW_TIMEOUT_S, ) -> None: + if not (math.isfinite(shadow_timeout_s) and shadow_timeout_s > 0): + raise ValueError("shadow_timeout_s must be a finite number greater than 0") self._key = key self._handler = handler self._tool_handlers = tool_handlers self._registry = registry self._skip_judges = skip_judges + self._skip_shadows = skip_shadows + self._shadow_timeout_s = shadow_timeout_s def _normalize_handlers(self) -> list[ProviderHandler] | None: if self._handler is None: @@ -130,7 +129,7 @@ async def _invoke( ) resolved_tools = resolve_tools(self._registry, self._tool_handlers) - variation = await extract_variation(self._key, context) + variation = await extract_variation(self._key, context, with_reason=True) config: AiConfigRep = variation["config"] meta: VariationMeta = variation["meta"] @@ -138,7 +137,7 @@ async def _invoke( raise ValueError("No handlers provided to config()") handler = select_handler(config, meta, resolved_handler_list) run_id = str(uuid.uuid4()) - _run_hooks( + shadow_tasks = _shadow_tasks( config_key=self._key, variation=variation, meta=meta, @@ -148,7 +147,10 @@ async def _invoke( variables=variables, handlers=resolved_handler_list, run_id=run_id, + timeout_s=self._shadow_timeout_s, ) + if not self._skip_shadows: + start_shadows(shadow_tasks, resolved_handler_list) result = await execute_and_track( config_key=self._key, @@ -192,6 +194,7 @@ async def _invoke( response=parsed_response, usage=usage_obj, judge_tasks=judge_tasks, + shadow_tasks=shadow_tasks if self._skip_shadows else None, track_data=track_data, ) @@ -210,6 +213,7 @@ async def _invoke( response=parsed_response, usage=usage_obj, judge_results=judge_results if judge_results else None, + shadow_tasks=shadow_tasks if self._skip_shadows else None, track_data=track_data, ) @@ -262,7 +266,7 @@ async def _execute_stream( ) resolved_tools = resolve_tools(self._registry, self._tool_handlers) - variation = await extract_variation(self._key, context) + variation = await extract_variation(self._key, context, with_reason=True) config: AiConfigRep = variation["config"] meta: VariationMeta = variation["meta"] @@ -281,7 +285,7 @@ async def _execute_stream( sink["provider"] = provider sink["model"] = model run_id = str(uuid.uuid4()) - _run_hooks( + shadow_tasks = _shadow_tasks( config_key=self._key, variation=variation, meta=meta, @@ -291,7 +295,11 @@ async def _execute_stream( variables=variables, handlers=resolved_handler_list, run_id=run_id, + timeout_s=self._shadow_timeout_s, ) + if not self._skip_shadows: + start_shadows(shadow_tasks, resolved_handler_list) + sink["shadow_tasks"] = shadow_tasks if self._skip_shadows else None async with aclosing( execute_and_stream( @@ -331,12 +339,15 @@ async def _stream_events( judge_results = await judge_completed_stream( self, sink, done_event, user_input, context ) - yield { + done: StreamEvent = { "type": "done", "response": done_event.get("response", ""), "usage": done_event.get("usage"), "judge_results": judge_results if judge_results else None, } + if sink.get("shadow_tasks") is not None: + done["shadow_tasks"] = sink["shadow_tasks"] + yield done class _InternalConfigInstance(ConfigInstance): @@ -372,6 +383,8 @@ def config( tool_handlers: Mapping[str, Callable[..., Any] | NativeTool] | None = None, registry: Any = None, skip_judges: bool = False, + skip_shadows: bool = False, + shadow_timeout_s: float = DEFAULT_SHADOW_TIMEOUT_S, ) -> ConfigInstance: """ Creates a ``ConfigInstance`` bound to *key*. Accepts a single handler or a @@ -384,6 +397,11 @@ def config( ``.invoke()`` / ``.stream()``. When set, ``invoke()`` returns ``judge_tasks: list[JudgeTask]`` — pre-packaged tasks ready for a background thread calling ``run_judge(task, handlers)``. + + Shadow requests attached to the matched rule or fallthrough run in the + background with *shadow_timeout_s*. Pass ``skip_shadows=True`` to get them + back as ``shadow_tasks`` instead, and run each with + ``run_shadow(task, handlers)``. """ return ConfigInstance( key=key, @@ -391,6 +409,8 @@ def config( tool_handlers=tool_handlers, registry=registry, skip_judges=skip_judges, + skip_shadows=skip_shadows, + shadow_timeout_s=shadow_timeout_s, ) @@ -401,6 +421,8 @@ def _config( tool_handlers: Mapping[str, Callable[..., Any] | NativeTool] | None = None, registry: Any = None, skip_judges: bool = False, + skip_shadows: bool = False, + shadow_timeout_s: float = DEFAULT_SHADOW_TIMEOUT_S, ) -> ConfigInstance: """Non-reporting :func:`config` for SDK wrappers.""" return _InternalConfigInstance( @@ -409,4 +431,6 @@ def _config( tool_handlers=tool_handlers, registry=registry, skip_judges=skip_judges, + skip_shadows=skip_shadows, + shadow_timeout_s=shadow_timeout_s, ) diff --git a/packages/client/src/launchdarkly_ai_server/experimental/__init__.py b/packages/client/src/launchdarkly_ai_server/experimental/__init__.py index fbdde28f..5fdaa0fb 100644 --- a/packages/client/src/launchdarkly_ai_server/experimental/__init__.py +++ b/packages/client/src/launchdarkly_ai_server/experimental/__init__.py @@ -1,7 +1,7 @@ """ Experimental features of the LaunchDarkly AI SDK. -Each feature is a submodule of this package: today, `shadow_requests`. +Each feature is a submodule of this package. No feature is experimental now. Names imported from here may change in a minor release, so review the changelog when you upgrade. Names from the package root change only in a major release. """ diff --git a/packages/client/src/launchdarkly_ai_server/experimental/shadow_requests.py b/packages/client/src/launchdarkly_ai_server/experimental/shadow_requests.py deleted file mode 100644 index c3080a21..00000000 --- a/packages/client/src/launchdarkly_ai_server/experimental/shadow_requests.py +++ /dev/null @@ -1,322 +0,0 @@ -""" -Shadow Requests (experimental). - -Call ``enable_shadow_requests()`` to run the shadow variations LaunchDarkly attaches -to a targeting rule or the fallthrough. Each sampled shadow runs in the background on -the same input, and is tracked on the normal ``$ld:ai:*`` events with ``isShadow``. -Its response is discarded, and the caller's response, latency and events are -unchanged. These names may change in a minor release, so review the changelog when -you upgrade. -""" - -from __future__ import annotations - -import asyncio -import hashlib -import logging -import math -import time -from collections.abc import Callable -from dataclasses import dataclass, field -from typing import Any - -from ..invocation_hooks import InvocationEvent, register_invocation_hook -from ..sdk_usage import report_usage -from ..shutdown_hooks import register_shutdown_hook -from ..tracking import build_track_data, execute_and_track -from ..types import AiConfigRep, LDContext, ProviderHandler, VariationMeta -from ..types_validation import parse_ai_config -from ..utils import ( - _context_identity_from_ld_context, - model_stamps_from_meta, - parse_usage, - select_handler, -) - -__all__ = [ - "ShadowResult", - "ShadowTask", - "disable_shadow_requests", - "enable_shadow_requests", - "run_shadow", -] - -logger = logging.getLogger("launchdarkly_ai_server.experimental.shadow_requests") - -_FEATURE = "Shadow Requests" -_FALLTHROUGH = "fallthrough" -_STRIPPED_KEYS = ("_ldMeta", "tools", "judgeConfiguration", "shadowConfiguration") -_BUCKET_SCALE = float(0xFFFFFFFFFFFFFFF) - - -@dataclass(frozen=True) -class ShadowTask: - """ - Everything needed to run and track one sampled shadow, with no flag evaluation. - Plain data, so it survives ``pickle``. Run it with ``run_shadow(task, handlers)``. - """ - - config_key: str - shadow_request_id: str - shadow_config: AiConfigRep - shadow_meta: VariationMeta - user_input: str | None - history: list[dict[str, Any]] | None - variables: dict[str, Any] | None - user_context: LDContext - timeout_s: float - track_data: dict[str, Any] - span_attributes: dict[str, Any] = field(default_factory=dict) - - -@dataclass(frozen=True) -class ShadowResult: - """The outcome of one shadow run. There is no response text: it is discarded.""" - - usage: dict[str, int] - track_data: dict[str, Any] - error: str | None = None - - -@dataclass(frozen=True) -class _Settings: - timeout_s: float - on_task: Callable[[ShadowTask], None] | None - - -_settings: _Settings | None = None -_pending: set[asyncio.Task[ShadowResult]] = set() - - -def enable_shadow_requests( - *, - timeout_s: float = 30.0, - on_task: Callable[[ShadowTask], None] | None = None, -) -> None: - """ - Turns Shadow Requests on for this process. Calling it again replaces the settings. - - By default each sampled shadow runs in the background with *timeout_s*. Pass - *on_task* to run them yourself instead: it receives a ``ShadowTask`` for each - sampled shadow, before the real handler runs, and must only enqueue it. - """ - report_usage("experimental.enableShadowRequests") - if not (math.isfinite(timeout_s) and timeout_s > 0): - raise ValueError("timeout_s must be a finite number greater than 0") - global _settings - _settings = _Settings(timeout_s=timeout_s, on_task=on_task) - register_invocation_hook(_FEATURE, _on_invocation) - - -def disable_shadow_requests() -> None: - """Turns Shadow Requests off. Shadows already running finish.""" - global _settings - _settings = None - register_invocation_hook(_FEATURE, None) - - -async def run_shadow(task: ShadowTask, handlers: list[ProviderHandler]) -> ShadowResult: - """Runs one ``ShadowTask``. Never raises; a failure is returned in ``error``.""" - report_usage("experimental.runShadow") - return await _run(task, handlers) - - -def _reset() -> None: - for pending in list(_pending): - pending.cancel() - _pending.clear() - disable_shadow_requests() - - -register_shutdown_hook(_FEATURE, _reset) - - -async def _wait_for_shadows_for_testing() -> None: - while _pending: - await asyncio.gather(*list(_pending), return_exceptions=True) - - -def _in_sample(shadow_id: str, context: LDContext, sampling_rate: float) -> bool: - identity = _context_identity_from_ld_context(context) - if identity is None: - return False - canonical, _ = identity - digest = hashlib.sha1(f"{shadow_id}.{canonical}".encode(), usedforsecurity=False) - bucket = int(digest.hexdigest()[:15], 16) / _BUCKET_SCALE - return bucket < sampling_rate - - -def _slot(reason: Any) -> str | None: - if not isinstance(reason, dict): - return None - kind = reason.get("kind") - if kind == "FALLTHROUGH": - return _FALLTHROUGH - if kind == "RULE_MATCH" and isinstance(reason.get("ruleId"), str): - return str(reason["ruleId"]) - return None - - -def _valid_entry(entry: Any) -> bool: - if not isinstance(entry, dict): - return False - numbers_ok = all( - isinstance(entry.get(k), (int, float)) and not isinstance(entry.get(k), bool) - for k in ("samplingRate", "endsAt") - ) - strings_ok = all( - isinstance(entry.get(k), str) and entry.get(k) for k in ("id", "variationKey") - ) - return numbers_ok and strings_ok - - -def _parse_variation(raw: Any) -> tuple[AiConfigRep, VariationMeta] | None: - if not isinstance(raw, dict): - return None - ld_meta = raw.get("_ldMeta") - if not isinstance(ld_meta, dict) or not ld_meta.get("enabled"): - return None - parsed = parse_ai_config({k: v for k, v in raw.items() if k not in _STRIPPED_KEYS}) - if not parsed.success: - return None - meta: VariationMeta = { - "enabled": True, - "variationKey": ld_meta.get("variationKey", ""), - "version": ld_meta.get("version", 1), - "mode": ld_meta.get("mode"), - **model_stamps_from_meta(ld_meta), - } - return parsed.data, meta - - -def _on_invocation(event: InvocationEvent) -> None: - settings = _settings - if settings is None: - return - slot = _slot(event.reason) - configuration = event.value.get("shadowConfiguration") - if slot is None or not isinstance(configuration, dict): - return - shadows = configuration.get("shadows") - variations = configuration.get("variations") - if not isinstance(shadows, dict) or not isinstance(variations, dict): - return - entries = shadows.get(slot) - if not isinstance(entries, list): - return - - now_ms = time.time() * 1000 - for entry in entries: - if not _valid_entry(entry) or entry["endsAt"] <= now_ms: - continue - if not _in_sample(entry["id"], event.context, float(entry["samplingRate"])): - continue - parsed = _parse_variation(variations.get(entry["variationKey"])) - if parsed is None: - continue - shadow_config, shadow_meta = parsed - try: - select_handler(shadow_config, shadow_meta, event.handlers) - except Exception: - logger.debug("No handler for shadow request %s", entry["id"]) - continue - task = ShadowTask( - config_key=event.config_key, - shadow_request_id=entry["id"], - shadow_config=shadow_config, - shadow_meta=shadow_meta, - user_input=event.user_input, - history=list(event.history) if event.history else None, - variables=dict(event.variables) if event.variables else None, - user_context=dict(event.context), - timeout_s=settings.timeout_s, - track_data={ - "runId": event.run_id, - "isShadow": True, - "primaryVariationKey": event.variation_key, - "shadowRequestId": entry["id"], - "ownerKind": entry.get("ownerKind", ""), - "ownerId": entry.get("ownerId", ""), - }, - span_attributes={ - "launchdarkly.is_shadow": True, - "launchdarkly.shadow.request.id": entry["id"], - "launchdarkly.shadow.primary.variation.key": event.variation_key, - }, - ) - if settings.on_task is not None: - try: - settings.on_task(task) - except Exception: - logger.warning( - "The Shadow Requests on_task callback failed", exc_info=True - ) - continue - running = asyncio.get_running_loop().create_task(_run(task, event.handlers)) - _pending.add(running) - running.add_done_callback(_pending.discard) - - -async def _run(task: ShadowTask, handlers: list[ProviderHandler]) -> ShadowResult: - extra = {k: v for k, v in task.track_data.items() if k != "runId"} - run_id = task.track_data.get("runId") - track_data = build_track_data( - config_key=task.config_key, - config=task.shadow_config, - meta=task.shadow_meta, - run_id=run_id, - extra_track_data=extra, - ) - try: - handler = select_handler(task.shadow_config, task.shadow_meta, handlers) - except Exception as err: - return ShadowResult( - usage=parse_usage({}), track_data=track_data, error=str(err) - ) - - try: - result = await execute_and_track( - config_key=task.config_key, - config=task.shadow_config, - meta=task.shadow_meta, - user_context=task.user_context, - handler=_with_deadline(handler, task.timeout_s), - user_input=task.user_input, - tool_handlers=None, - variables=task.variables, - history=list(task.history) if task.history else None, - run_id=run_id, - extra_track_data=extra, - span_attributes=task.span_attributes, - ) - except _ShadowTimeout: - return ShadowResult( - usage=parse_usage({}), track_data=track_data, error="timeout" - ) - except Exception as err: - return ShadowResult( - usage=parse_usage({}), track_data=track_data, error=str(err) - ) - return ShadowResult(usage=result["usage"], track_data=result["track_data"]) - - -class _ShadowTimeout(Exception): - pass - - -def _with_deadline(handler: ProviderHandler, timeout_s: float) -> ProviderHandler: - async def timed(*args: Any) -> Any: - deadline = asyncio.timeout(timeout_s) - try: - async with deadline: - return await handler(*args) - except TimeoutError: - if deadline.expired(): - raise _ShadowTimeout("timeout") from None - raise - - return ProviderHandler( - fn=timed, - provides_for=handler.provides_for, - capture_content=handler.capture_content, - ) diff --git a/packages/client/src/launchdarkly_ai_server/invocation_hooks.py b/packages/client/src/launchdarkly_ai_server/invocation_hooks.py deleted file mode 100644 index 0534f404..00000000 --- a/packages/client/src/launchdarkly_ai_server/invocation_hooks.py +++ /dev/null @@ -1,64 +0,0 @@ -""" -Hooks that ``config().invoke()`` and ``.stream()`` run before the handler. - -This is how a feature reacts to an invocation without core importing it: a -feature module registers a hook, and core runs whatever is registered. Core -never imports, names or calls a feature module, so a hook's *name* comes from -the module that registers it. Same shape as ``shutdown_hooks``. -""" - -from __future__ import annotations - -import logging -from collections.abc import Callable -from dataclasses import dataclass -from typing import Any - -from .types import LDContext, ProviderHandler - -logger = logging.getLogger("launchdarkly_ai_server.client") - - -@dataclass(frozen=True) -class InvocationEvent: - """What a hook sees: one ``config()`` call, after routing and before the handler runs.""" - - config_key: str - value: dict[str, Any] - reason: Any - context: LDContext - user_input: str | None - history: list[dict[str, Any]] | None - variables: dict[str, Any] | None - handlers: list[ProviderHandler] - run_id: str - variation_key: str - - -_hooks: dict[str, Callable[[InvocationEvent], None]] = {} - - -def register_invocation_hook( - name: str, hook: Callable[[InvocationEvent], None] | None -) -> None: - """ - Registers *hook* under *name*, replacing any earlier hook with that name. - Passing ``None`` removes it. - """ - if hook is None: - _hooks.pop(name, None) - else: - _hooks[name] = hook - - -def has_invocation_hooks() -> bool: - return bool(_hooks) - - -def run_invocation_hooks(event: InvocationEvent) -> None: - """Runs every hook. A hook that raises is logged and must not break the call.""" - for name, hook in list(_hooks.items()): - try: - hook(event) - except Exception: - logger.warning("The %s invocation hook failed", name, exc_info=True) diff --git a/packages/client/src/launchdarkly_ai_server/lifecycle.py b/packages/client/src/launchdarkly_ai_server/lifecycle.py index 42d60994..0501424a 100644 --- a/packages/client/src/launchdarkly_ai_server/lifecycle.py +++ b/packages/client/src/launchdarkly_ai_server/lifecycle.py @@ -586,6 +586,8 @@ async def inspect_config( async def extract_variation( config_key: str, context: dict[str, Any], + *, + with_reason: bool = False, ) -> dict[str, Any]: """ Fetches an AI config variation from the LD client and parses it. @@ -603,11 +605,12 @@ async def extract_variation( # The real ldclient.variation() expects an ldclient.Context, not a plain dict. # Convert when running against the real SDK; test mocks accept dicts directly. ld_context = to_ld_context(client, context) - from .invocation_hooks import has_invocation_hooks - + # The evaluation reason picks the shadow requests for the matched rule. reason: Any = None - variation_detail = getattr(client, "variation_detail", None) - if has_invocation_hooks() and variation_detail is not None: + variation_detail = ( + getattr(client, "variation_detail", None) if with_reason else None + ) + if variation_detail is not None: detail_result = variation_detail(config_key, ld_context, None) detail = ( await detail_result if inspect.isawaitable(detail_result) else detail_result diff --git a/packages/client/src/launchdarkly_ai_server/shadows.py b/packages/client/src/launchdarkly_ai_server/shadows.py new file mode 100644 index 00000000..00abfaa9 --- /dev/null +++ b/packages/client/src/launchdarkly_ai_server/shadows.py @@ -0,0 +1,283 @@ +""" +Shadow requests: run another variation of the AI Config in the background on +the same input, and track it on the normal ``$ld:ai:*`` events with +``isShadow``. The shadow's response is discarded, and the caller's response, +latency and events are unchanged. + +LaunchDarkly attaches shadow requests to a targeting rule or the fallthrough +and sends them in the served variation as ``shadowConfiguration``. +``config().invoke()`` and ``.stream()`` call :func:`build_shadow_tasks`, then +either :func:`start_shadows` or hand the tasks back when ``skip_shadows=True``. +""" + +from __future__ import annotations + +import asyncio +import hashlib +import logging +import time +from typing import Any + +from .sdk_usage import report_usage +from .shutdown_hooks import register_shutdown_hook +from .tracking import build_track_data, execute_and_track +from .types import ( + AiConfigRep, + LDContext, + ProviderHandler, + ShadowResult, + ShadowTask, + VariationMeta, +) +from .types_validation import parse_ai_config +from .utils import ( + _context_identity_from_ld_context, + model_stamps_from_meta, + parse_usage, + select_handler, +) + +logger = logging.getLogger("launchdarkly_ai_server.shadows") + +DEFAULT_SHADOW_TIMEOUT_S = 30.0 +_FALLTHROUGH = "fallthrough" +_STRIPPED_KEYS = ("_ldMeta", "tools", "judgeConfiguration", "shadowConfiguration") +_BUCKET_SCALE = float(0xFFFFFFFFFFFFFFF) + +_running: set[asyncio.Task[ShadowResult]] = set() + + +def build_shadow_tasks( + *, + config_key: str, + raw_variation: Any, + reason: Any, + primary_variation_key: str, + context: LDContext, + user_input: str | None, + history: list[dict[str, Any]] | None, + variables: dict[str, Any] | None, + handlers: list[ProviderHandler], + run_id: str, + timeout_s: float, +) -> list[ShadowTask]: + """ + Returns a task for each shadow request on the rule or fallthrough that + served this call, skipping entries that have ended, are not sampled for this + context, or have no matching handler. Never raises. + """ + try: + entries, variations = _entries_for(raw_variation, reason) + except Exception: + logger.debug("Ignoring a malformed shadowConfiguration", exc_info=True) + return [] + + tasks: list[ShadowTask] = [] + now_ms = time.time() * 1000 + for entry in entries: + if not _valid_entry(entry) or entry["endsAt"] <= now_ms: + continue + if not _in_sample(entry["id"], context, float(entry["samplingRate"])): + continue + parsed = _parse_variation(variations.get(entry["variationKey"])) + if parsed is None: + continue + shadow_config, shadow_meta = parsed + try: + select_handler(shadow_config, shadow_meta, handlers) + except Exception: + logger.debug("No handler for shadow request %s", entry["id"]) + continue + tasks.append( + ShadowTask( + config_key=config_key, + shadow_request_id=entry["id"], + shadow_config=shadow_config, + shadow_meta=shadow_meta, + user_input=user_input, + history=list(history) if history else None, + variables=dict(variables) if variables else None, + user_context=dict(context), + timeout_s=timeout_s, + track_data={ + "runId": run_id, + "isShadow": True, + "primaryVariationKey": primary_variation_key, + "shadowRequestId": entry["id"], + "ownerKind": entry.get("ownerKind", ""), + "ownerId": entry.get("ownerId", ""), + }, + span_attributes={ + "launchdarkly.is_shadow": True, + "launchdarkly.shadow.request.id": entry["id"], + "launchdarkly.shadow.primary.variation.key": primary_variation_key, + }, + ) + ) + return tasks + + +def start_shadows(tasks: list[ShadowTask], handlers: list[ProviderHandler]) -> None: + """Runs each task in the background. The caller does not wait for them.""" + if not tasks: + return + loop = asyncio.get_running_loop() + for task in tasks: + running = loop.create_task(_run(task, handlers)) + _running.add(running) + running.add_done_callback(_running.discard) + + +async def run_shadow(task: ShadowTask, handlers: list[ProviderHandler]) -> ShadowResult: + """Runs one :class:`ShadowTask`. Never raises; a failure is in ``error``.""" + report_usage("client.runShadow") + return await _run(task, handlers) + + +def _cancel_running() -> None: + for running in list(_running): + running.cancel() + _running.clear() + + +register_shutdown_hook("Shadow requests", _cancel_running) + + +async def _wait_for_shadows_for_testing() -> None: + while _running: + await asyncio.gather(*list(_running), return_exceptions=True) + + +def _entries_for(raw_variation: Any, reason: Any) -> tuple[list[Any], dict[str, Any]]: + configuration = ( + raw_variation.get("shadowConfiguration") + if isinstance(raw_variation, dict) + else None + ) + slot = _slot(reason) + if slot is None or not isinstance(configuration, dict): + return [], {} + shadows = configuration.get("shadows") + variations = configuration.get("variations") + if not isinstance(shadows, dict) or not isinstance(variations, dict): + return [], {} + entries = shadows.get(slot) + return (entries if isinstance(entries, list) else []), variations + + +def _slot(reason: Any) -> str | None: + if not isinstance(reason, dict): + return None + kind = reason.get("kind") + if kind == "FALLTHROUGH": + return _FALLTHROUGH + if kind == "RULE_MATCH" and isinstance(reason.get("ruleId"), str): + return str(reason["ruleId"]) + return None + + +def _in_sample(shadow_id: str, context: LDContext, sampling_rate: float) -> bool: + identity = _context_identity_from_ld_context(context) + if identity is None: + return False + canonical, _ = identity + digest = hashlib.sha1(f"{shadow_id}.{canonical}".encode(), usedforsecurity=False) + bucket = int(digest.hexdigest()[:15], 16) / _BUCKET_SCALE + return bucket < sampling_rate + + +def _valid_entry(entry: Any) -> bool: + if not isinstance(entry, dict): + return False + numbers_ok = all( + isinstance(entry.get(k), (int, float)) and not isinstance(entry.get(k), bool) + for k in ("samplingRate", "endsAt") + ) + strings_ok = all( + isinstance(entry.get(k), str) and entry.get(k) for k in ("id", "variationKey") + ) + return numbers_ok and strings_ok + + +def _parse_variation(raw: Any) -> tuple[AiConfigRep, VariationMeta] | None: + if not isinstance(raw, dict): + return None + ld_meta = raw.get("_ldMeta") + if not isinstance(ld_meta, dict) or not ld_meta.get("enabled"): + return None + parsed = parse_ai_config({k: v for k, v in raw.items() if k not in _STRIPPED_KEYS}) + if not parsed.success: + return None + meta: VariationMeta = { + "enabled": True, + "variationKey": ld_meta.get("variationKey", ""), + "version": ld_meta.get("version", 1), + "mode": ld_meta.get("mode"), + **model_stamps_from_meta(ld_meta), + } + return parsed.data, meta + + +async def _run(task: ShadowTask, handlers: list[ProviderHandler]) -> ShadowResult: + extra = {k: v for k, v in task.track_data.items() if k != "runId"} + run_id = task.track_data.get("runId") + track_data = build_track_data( + config_key=task.config_key, + config=task.shadow_config, + meta=task.shadow_meta, + run_id=run_id, + extra_track_data=extra, + ) + try: + handler = select_handler(task.shadow_config, task.shadow_meta, handlers) + except Exception as err: + return ShadowResult( + usage=parse_usage({}), track_data=track_data, error=str(err) + ) + + try: + result = await execute_and_track( + config_key=task.config_key, + config=task.shadow_config, + meta=task.shadow_meta, + user_context=task.user_context, + handler=_with_deadline(handler, task.timeout_s), + user_input=task.user_input, + tool_handlers=None, + variables=task.variables, + history=list(task.history) if task.history else None, + run_id=run_id, + extra_track_data=extra, + span_attributes=task.span_attributes, + ) + except _ShadowTimeout: + return ShadowResult( + usage=parse_usage({}), track_data=track_data, error="timeout" + ) + except Exception as err: + return ShadowResult( + usage=parse_usage({}), track_data=track_data, error=str(err) + ) + return ShadowResult(usage=result["usage"], track_data=result["track_data"]) + + +class _ShadowTimeout(Exception): + pass + + +def _with_deadline(handler: ProviderHandler, timeout_s: float) -> ProviderHandler: + async def timed(*args: Any) -> Any: + deadline = asyncio.timeout(timeout_s) + try: + async with deadline: + return await handler(*args) + except TimeoutError: + if deadline.expired(): + raise _ShadowTimeout("timeout") from None + raise + + return ProviderHandler( + fn=timed, + provides_for=handler.provides_for, + capture_content=handler.capture_content, + ) diff --git a/packages/client/src/launchdarkly_ai_server/types.py b/packages/client/src/launchdarkly_ai_server/types.py index 079f9d2e..99161147 100644 --- a/packages/client/src/launchdarkly_ai_server/types.py +++ b/packages/client/src/launchdarkly_ai_server/types.py @@ -288,6 +288,12 @@ class ProviderResponse(Generic[T]): ``None`` when ``skip_judges=False`` (judges ran inline). """ + shadow_tasks: list[ShadowTask] | None = None + """ + Shadow requests sampled for this call, produced when ``skip_shadows=True``. + Run each with ``run_shadow(task, handlers)``. ``None`` when shadows ran in + the background. + """ track_data: TrackData | None = None """ Tracking payload from this invocation. Carried inside each :class:`JudgeTask` @@ -357,6 +363,39 @@ class JudgeTask: """LD metric key to track the score against.""" +@dataclass(frozen=True) +class ShadowTask: + """ + Everything needed to run one shadow request later, without evaluating the + flag again. Returned in ``ProviderResponse.shadow_tasks`` when + ``skip_shadows=True``. Plain data, so it can be pickled and queued. Run it + with ``run_shadow(task, handlers)``. + """ + + config_key: str + shadow_request_id: str + shadow_config: AiConfigRep + """The shadow variation, with tools, judges and shadow configuration removed.""" + shadow_meta: VariationMeta + user_input: str | None + history: list[dict[str, Any]] | None + variables: dict[str, Any] | None + user_context: LDContext + timeout_s: float + track_data: TrackData + """Fields added to every event of the shadow run, including the real run's ``runId``.""" + span_attributes: dict[str, Any] = field(default_factory=dict) + + +@dataclass(frozen=True) +class ShadowResult: + """The outcome of one shadow run. The shadow's response is discarded.""" + + usage: dict[str, int] + track_data: TrackData + error: str | None = None + + @dataclass class JudgeRunResult: """ diff --git a/packages/client/tests/conftest.py b/packages/client/tests/conftest.py index de69af0c..5c40d4c7 100644 --- a/packages/client/tests/conftest.py +++ b/packages/client/tests/conftest.py @@ -14,6 +14,7 @@ def mock_ld_client() -> MagicMock: """Stub LaunchDarkly client with variation/track/flush/close.""" client = MagicMock() + del client.variation_detail client.variation = AsyncMock(return_value=None) client.track = MagicMock() client.flush = AsyncMock() diff --git a/packages/client/tests/test_agui.py b/packages/client/tests/test_agui.py index 68613a08..16c7f3d1 100644 --- a/packages/client/tests/test_agui.py +++ b/packages/client/tests/test_agui.py @@ -32,6 +32,7 @@ def _client() -> MagicMock: client = MagicMock() + del client.variation_detail client.track = MagicMock() client.flush = AsyncMock() client.close = AsyncMock() diff --git a/packages/client/tests/test_client.py b/packages/client/tests/test_client.py index 051f8d62..231d4344 100644 --- a/packages/client/tests/test_client.py +++ b/packages/client/tests/test_client.py @@ -34,6 +34,7 @@ def _generation_calls(client: MagicMock) -> list[Any]: def _make_client(response: Any = "Hello", usage: dict | None = None) -> MagicMock: c = MagicMock() + del c.variation_detail c.track = MagicMock() c.flush = AsyncMock() c.close = AsyncMock() diff --git a/packages/client/tests/test_public_surface.py b/packages/client/tests/test_public_surface.py index 29177a3a..a22e7070 100644 --- a/packages/client/tests/test_public_surface.py +++ b/packages/client/tests/test_public_surface.py @@ -44,6 +44,8 @@ "ProviderGraphResponse", "ProviderHandler", "ProviderResponse", + "ShadowResult", + "ShadowTask", "Registry", "RunUsage", "SDK_INFO_CONTEXT", @@ -101,6 +103,7 @@ "resolve_handlers", "resolve_tools", "run_judge", + "run_shadow", "run_judges", "set_conversation_id_if_absent", "set_input_content_attributes", diff --git a/packages/client/tests/test_sdk_usage.py b/packages/client/tests/test_sdk_usage.py index 9f2a8d8e..0a649706 100644 --- a/packages/client/tests/test_sdk_usage.py +++ b/packages/client/tests/test_sdk_usage.py @@ -43,6 +43,7 @@ def _client(variation: Any = None) -> MagicMock: client = MagicMock() + del client.variation_detail client.track = MagicMock() client.variation = AsyncMock(return_value=variation) client.flush = AsyncMock() diff --git a/packages/client/tests/test_shadow_requests.py b/packages/client/tests/test_shadows.py similarity index 62% rename from packages/client/tests/test_shadow_requests.py rename to packages/client/tests/test_shadows.py index 4ddb669f..d318113b 100644 --- a/packages/client/tests/test_shadow_requests.py +++ b/packages/client/tests/test_shadows.py @@ -1,10 +1,6 @@ -""" -Tests for §9 Shadow Requests (experimental). -Reference: TESTING.md §9, A.22 -""" +"""Tests for shadow requests.""" import asyncio -import logging import pickle import time from collections.abc import Iterator @@ -15,9 +11,14 @@ import pytest import launchdarkly_ai_server.lifecycle as lifecycle_module -from launchdarkly_ai_server import ProviderHandler, config, shutdown -from launchdarkly_ai_server.experimental import shadow_requests as sr -from launchdarkly_ai_server.invocation_hooks import register_invocation_hook +from launchdarkly_ai_server import ( + ProviderHandler, + ShadowTask, + config, + run_shadow, + shutdown, +) +from launchdarkly_ai_server import shadows as sh from launchdarkly_ai_server.sdk_usage import reset_sdk_usage from launchdarkly_ai_server.utils import set_ld_span_attributes @@ -172,9 +173,7 @@ def _real_events(client: MagicMock) -> list[tuple[str, dict[str, Any], Any]]: def _reset() -> Iterator[None]: lifecycle_module._reset_for_testing() reset_sdk_usage() - sr.disable_shadow_requests() yield - sr.disable_shadow_requests() lifecycle_module._reset_for_testing() @@ -188,34 +187,26 @@ async def _invoke( *, user_input: str = "hi", history: list[dict[str, Any]] | None = None, + **config_args: Any, ) -> tuple[Any, _Recorder, _Recorder]: _use(client) a = _Recorder("TestProvider", "real answer") b = _Recorder("ShadowProvider", "shadow answer") - response = await config(key="flag", handler=[a.handler, b.handler]).invoke( - user_input, CONTEXT, {"topic": "x"}, history - ) - await sr._wait_for_shadows_for_testing() + response = await config( + key="flag", handler=[a.handler, b.handler], **config_args + ).invoke(user_input, CONTEXT, {"topic": "x"}, history) + await sh._wait_for_shadows_for_testing() return response, a, b -# --------------------------------------------------------------------------- -# §9.1 Turning Shadow Requests on and off -# --------------------------------------------------------------------------- - +async def _deferred_task(client: MagicMock) -> ShadowTask: + response, _, _ = await _invoke(client, skip_shadows=True) + assert response.shadow_tasks is not None + return response.shadow_tasks[0] -class TestOnOff: - async def test_nothing_runs_until_enabled(self) -> None: - client = _make_client(_variation_a()) - _, a, b = await _invoke(client) - assert len(a.calls) == 1 - assert b.calls == [] - assert _shadow_events(client) == [] - client.variation_detail.assert_not_called() - client.variation.assert_called_once() - async def test_enabled_runs_the_fallthrough_shadow(self) -> None: - sr.enable_shadow_requests() +class TestRunning: + async def test_runs_the_fallthrough_shadow(self) -> None: client = _make_client(_variation_a()) _, a, b = await _invoke(client) assert len(a.calls) == 1 @@ -223,63 +214,111 @@ async def test_enabled_runs_the_fallthrough_shadow(self) -> None: client.variation_detail.assert_called_once() client.variation.assert_not_called() - def test_timeout_must_be_positive(self) -> None: - with pytest.raises(ValueError): - sr.enable_shadow_requests(timeout_s=0) - - async def test_disable_is_the_kill_switch(self) -> None: - sr.enable_shadow_requests() - sr.disable_shadow_requests() + async def test_client_without_variation_detail_runs_no_shadows(self) -> None: client = _make_client(_variation_a()) - _, _, b = await _invoke(client) + del client.variation_detail + response, a, b = await _invoke(client) + assert response.response == "real answer" + assert len(a.calls) == 1 assert b.calls == [] - async def test_disable_stops_deferred_mode(self) -> None: - tasks: list[sr.ShadowTask] = [] - sr.enable_shadow_requests(on_task=tasks.append) - sr.disable_shadow_requests() - await _invoke(_make_client(_variation_a())) - assert tasks == [] - - async def test_enabling_again_replaces_settings(self) -> None: - tasks: list[sr.ShadowTask] = [] - sr.enable_shadow_requests(on_task=tasks.append) - sr.enable_shadow_requests() - _, _, b = await _invoke(_make_client(_variation_a())) - assert tasks == [] - assert len(b.calls) == 1 + @pytest.mark.parametrize("timeout_s", [0, -1, float("inf"), float("nan")]) + def test_timeout_must_be_positive_and_finite(self, timeout_s: float) -> None: + with pytest.raises(ValueError): + config(key="flag", shadow_timeout_s=timeout_s) - async def test_shutdown_disables(self) -> None: - sr.enable_shadow_requests() + async def test_shutdown_cancels_running_shadows(self) -> None: _use(_make_client(_variation_a())) + a = _Recorder("TestProvider", "real answer") + b = _Recorder("ShadowProvider", sleep_s=5) + await config(key="flag", handler=[a.handler, b.handler]).invoke("hi", CONTEXT) + assert sh._running await shutdown() - client = _make_client(_variation_a()) - _, _, b = await _invoke(client) - assert b.calls == [] + assert not sh._running + + async def test_same_input_and_stripped_config(self) -> None: + history = [{"role": "user", "content": "earlier"}] + _, a, b = await _invoke( + _make_client(_variation_a()), user_input="question", history=history + ) + call = b.calls[0] + assert call["input"] == "question" + assert call["history"] == history + assert call["variables"]["topic"] == "x" + assert "tools" not in call["config"] + assert "judgeConfiguration" not in call["config"] + assert "shadowConfiguration" not in a.calls[0]["config"] + assert "shadowConfiguration" not in call["config"] + assert not call["tools"] - async def test_enable_reports_usage(self) -> None: + async def test_shadow_answer_never_joins_history(self) -> None: + history = [{"role": "user", "content": "earlier"}] + await _invoke(_make_client(_variation_a()), history=history) + assert history == [{"role": "user", "content": "earlier"}] + + async def test_failing_shadow_never_reaches_the_caller(self) -> None: client = _use(_make_client(_variation_a())) - sr.enable_shadow_requests() - helpers = [ - c[0][2].get("helper") - for c in client.track.call_args_list - if c[0][0] == "$ld:ai:sdk:usage" - ] - assert "experimental.enableShadowRequests" in helpers + a = _Recorder("TestProvider", "real answer") + b = _Recorder("ShadowProvider", raises=True) + response = await config(key="flag", handler=[a.handler, b.handler]).invoke( + "hi", CONTEXT + ) + await sh._wait_for_shadows_for_testing() + assert response.response == "real answer" + assert "$ld:ai:generation:success" in [e[0] for e in _real_events(client)] + assert "$ld:ai:generation:error" in [e[0] for e in _shadow_events(client)] + + async def test_timeout_is_tracked_once(self) -> None: + client = _use(_make_client(_variation_a())) + a = _Recorder("TestProvider", "real answer") + b = _Recorder("ShadowProvider", sleep_s=1) + response = await config( + key="flag", handler=[a.handler, b.handler], shadow_timeout_s=0.05 + ).invoke("hi", CONTEXT) + await sh._wait_for_shadows_for_testing() + assert response.response == "real answer" + names = [e[0] for e in _shadow_events(client)] + assert names.count("$ld:ai:generation:error") == 1 + assert names.count("$ld:ai:duration:total") == 1 + async def test_handler_timeout_error_is_tracked_once(self) -> None: + client = _make_client(_variation_a()) + task = await _deferred_task(client) + b = _Recorder("ShadowProvider", error=TimeoutError("provider timeout")) + result = await run_shadow(task, [b.handler]) + names = [e[0] for e in _shadow_events(client)] + assert result.error == "provider timeout" + assert names.count("$ld:ai:generation:error") == 1 -# --------------------------------------------------------------------------- -# §9.2 Payload shape -# --------------------------------------------------------------------------- + async def test_real_run_is_unchanged(self) -> None: + baseline_client = _make_client(_variation_a(shadows={})) + baseline, _, _ = await _invoke(baseline_client) + lifecycle_module._reset_for_testing() + client = _make_client(_variation_a()) + response, _, _ = await _invoke(client) + assert response.response == baseline.response + assert response.usage == baseline.usage + assert [e[0] for e in _real_events(client)] == [ + e[0] for e in _real_events(baseline_client) + ] + + async def test_stream_runs_the_shadow_as_invoke(self) -> None: + _use(_make_client(_variation_a())) + a = _Recorder("TestProvider", "real answer") + b = _Recorder("ShadowProvider", "shadow answer") + events = [ + e + async for e in config(key="flag", handler=[a.handler, b.handler]).stream( + "hi", CONTEXT + ) + ] + await sh._wait_for_shadows_for_testing() + assert events[-1]["type"] == "done" + assert "shadow_tasks" not in events[-1] + assert len(b.calls) == 1 class TestPayload: - async def test_shadow_configuration_never_reaches_the_handler(self) -> None: - sr.enable_shadow_requests() - _, a, b = await _invoke(_make_client(_variation_a())) - assert "shadowConfiguration" not in a.calls[0]["config"] - assert "shadowConfiguration" not in b.calls[0]["config"] - @pytest.mark.parametrize( "shadow_configuration", [ @@ -298,7 +337,6 @@ async def test_shadow_configuration_never_reaches_the_handler(self) -> None: async def test_malformed_configuration_is_ignored( self, shadow_configuration: Any ) -> None: - sr.enable_shadow_requests() value = _variation_a() value["shadowConfiguration"] = shadow_configuration response, _, b = await _invoke(_make_client(value)) @@ -306,14 +344,8 @@ async def test_malformed_configuration_is_ignored( assert b.calls == [] -# --------------------------------------------------------------------------- -# §9.3 Selecting shadows -# --------------------------------------------------------------------------- - - class TestSelection: async def test_rule_match_uses_only_that_rule(self) -> None: - sr.enable_shadow_requests() shadows = { "r1": [_entry(id="rule-shadow")], "fallthrough": [_entry(id="ft-shadow")], @@ -339,12 +371,10 @@ async def test_rule_match_uses_only_that_rule(self) -> None: async def test_other_reasons_run_nothing( self, reason: dict[str, Any] | None ) -> None: - sr.enable_shadow_requests() _, _, b = await _invoke(_make_client(_variation_a(), reason)) assert b.calls == [] async def test_entry_past_its_end_is_skipped(self) -> None: - sr.enable_shadow_requests() shadows = {"fallthrough": [_entry(id="ended", endsAt=PAST), _entry(id="live")]} client = _make_client(_variation_a(shadows)) await _invoke(client) @@ -352,276 +382,115 @@ async def test_entry_past_its_end_is_skipped(self) -> None: assert ids == {"live"} async def test_missing_variation_is_skipped(self) -> None: - sr.enable_shadow_requests() _, _, b = await _invoke( _make_client(_variation_a({"fallthrough": [_entry(variationKey="gone")]})) ) assert b.calls == [] async def test_no_matching_handler_is_skipped(self) -> None: - sr.enable_shadow_requests() client = _use(_make_client(_variation_a())) a = _Recorder("TestProvider", "real answer") response = await config(key="flag", handler=a.handler).invoke("hi", CONTEXT) - await sr._wait_for_shadows_for_testing() + await sh._wait_for_shadows_for_testing() assert response.response == "real answer" assert _shadow_events(client) == [] def test_sampling_is_sticky(self) -> None: - assert sr._in_sample("s1", CONTEXT, 0.5) == sr._in_sample("s1", CONTEXT, 0.5) + assert sh._in_sample("s1", CONTEXT, 0.5) == sh._in_sample("s1", CONTEXT, 0.5) def test_sampling_honors_the_rate(self) -> None: hits = sum( - sr._in_sample("s1", {"kind": "user", "key": f"k{i}"}, 0.2) + sh._in_sample("s1", {"kind": "user", "key": f"k{i}"}, 0.2) for i in range(10_000) ) assert 1800 <= hits <= 2200 def test_sampling_edges(self) -> None: contexts = [{"kind": "user", "key": f"k{i}"} for i in range(200)] - assert all(sr._in_sample("s1", c, 1.0) for c in contexts) - assert not any(sr._in_sample("s1", c, 0.0) for c in contexts) + assert all(sh._in_sample("s1", c, 1.0) for c in contexts) + assert not any(sh._in_sample("s1", c, 0.0) for c in contexts) - def test_context_without_key_is_not_sampled(self) -> None: - assert not sr._in_sample("s1", {"kind": "user"}, 1.0) + def test_sampling_matches_typescript(self) -> None: + assert sh._in_sample("s1", CONTEXT, 0.732) + assert not sh._in_sample("s1", CONTEXT, 0.731) - async def test_stream_runs_the_shadow_as_invoke(self) -> None: - sr.enable_shadow_requests() - _use(_make_client(_variation_a())) - a = _Recorder("TestProvider", "real answer") - b = _Recorder("ShadowProvider", "shadow answer") - events = [ - e - async for e in config(key="flag", handler=[a.handler, b.handler]).stream( - "hi", CONTEXT - ) - ] - await sr._wait_for_shadows_for_testing() - assert events[-1]["type"] == "done" - assert len(b.calls) == 1 - - -# --------------------------------------------------------------------------- -# §9.4 Running a shadow -# --------------------------------------------------------------------------- - - -class TestRunning: - async def test_same_input_and_stripped_config(self) -> None: - sr.enable_shadow_requests() - history = [{"role": "user", "content": "earlier"}] - _, _, b = await _invoke( - _make_client(_variation_a()), user_input="question", history=history - ) - call = b.calls[0] - assert call["input"] == "question" - assert call["history"] == history - assert call["variables"]["topic"] == "x" - assert "tools" not in call["config"] - assert "judgeConfiguration" not in call["config"] - assert not call["tools"] - - async def test_shadow_answer_never_joins_history(self) -> None: - sr.enable_shadow_requests() - history = [{"role": "user", "content": "earlier"}] - await _invoke(_make_client(_variation_a()), history=history) - assert history == [{"role": "user", "content": "earlier"}] - - async def test_failing_shadow_never_reaches_the_caller(self) -> None: - sr.enable_shadow_requests() - client = _use(_make_client(_variation_a())) - a = _Recorder("TestProvider", "real answer") - b = _Recorder("ShadowProvider", raises=True) - response = await config(key="flag", handler=[a.handler, b.handler]).invoke( - "hi", CONTEXT - ) - await sr._wait_for_shadows_for_testing() - assert response.response == "real answer" - assert "$ld:ai:generation:success" in [e[0] for e in _real_events(client)] - assert "$ld:ai:generation:error" in [e[0] for e in _shadow_events(client)] - - async def test_timeout_records_an_error(self) -> None: - sr.enable_shadow_requests(timeout_s=0.05) - client = _use(_make_client(_variation_a())) - a = _Recorder("TestProvider", "real answer") - b = _Recorder("ShadowProvider", sleep_s=1) - response = await config(key="flag", handler=[a.handler, b.handler]).invoke( - "hi", CONTEXT - ) - await sr._wait_for_shadows_for_testing() - assert response.response == "real answer" - names = [e[0] for e in _shadow_events(client)] - assert "$ld:ai:generation:error" in names - assert "$ld:ai:duration:total" in names - - async def test_timeout_is_tracked_once(self) -> None: - sr.enable_shadow_requests(timeout_s=0.05) - client = _use(_make_client(_variation_a())) - a = _Recorder("TestProvider", "real answer") - b = _Recorder("ShadowProvider", sleep_s=1) - await config(key="flag", handler=[a.handler, b.handler]).invoke("hi", CONTEXT) - await sr._wait_for_shadows_for_testing() - names = [e[0] for e in _shadow_events(client)] - assert names.count("$ld:ai:generation:error") == 1 - assert names.count("$ld:ai:duration:total") == 1 - - async def test_handler_timeout_error_is_tracked_once(self) -> None: - tasks: list[sr.ShadowTask] = [] - sr.enable_shadow_requests(on_task=tasks.append) - client = _make_client(_variation_a()) - await _invoke(client) - b = _Recorder("ShadowProvider", error=TimeoutError("provider timeout")) - result = await sr.run_shadow(tasks[0], [b.handler]) - names = [e[0] for e in _shadow_events(client)] - assert result.error == "provider timeout" - assert names.count("$ld:ai:generation:error") == 1 - assert names.count("$ld:ai:duration:total") == 1 - - @pytest.mark.parametrize("timeout_s", [0, -1, float("inf"), float("nan")]) - def test_timeout_must_be_positive_and_finite(self, timeout_s: float) -> None: - with pytest.raises(ValueError): - sr.enable_shadow_requests(timeout_s=timeout_s) - - async def test_real_run_is_unchanged(self) -> None: - baseline_client = _make_client(_variation_a()) - baseline, _, _ = await _invoke(baseline_client) - lifecycle_module._reset_for_testing() - sr.enable_shadow_requests() - client = _make_client(_variation_a()) - response, _, _ = await _invoke(client) - assert response.response == baseline.response - assert response.usage == baseline.usage - assert [e[0] for e in _real_events(client)] == [ - e[0] for e in _real_events(baseline_client) - ] - - -# --------------------------------------------------------------------------- -# §9.5 Automatic and deferred modes -# --------------------------------------------------------------------------- + def test_context_without_key_is_not_sampled(self) -> None: + assert not sh._in_sample("s1", {"kind": "user"}, 1.0) class TestDeferred: - async def test_deferred_runs_nothing_inline(self) -> None: - tasks: list[sr.ShadowTask] = [] - sr.enable_shadow_requests(on_task=tasks.append) + async def test_skip_shadows_returns_tasks_and_runs_nothing(self) -> None: client = _make_client(_variation_a()) - _, a, b = await _invoke(client) + response, a, b = await _invoke(client, skip_shadows=True) assert len(a.calls) == 1 assert b.calls == [] - assert len(tasks) == 1 + assert response.shadow_tasks is not None + assert len(response.shadow_tasks) == 1 assert _shadow_events(client) == [] + async def test_tasks_are_none_when_shadows_run_in_the_background(self) -> None: + response, _, _ = await _invoke(_make_client(_variation_a())) + assert response.shadow_tasks is None + + async def test_stream_done_event_carries_tasks(self) -> None: + _use(_make_client(_variation_a())) + a = _Recorder("TestProvider", "real answer") + events = [ + e + async for e in config( + key="flag", handler=a.handler, skip_shadows=True + ).stream("hi", CONTEXT) + ] + assert events[-1]["type"] == "done" + assert events[-1]["shadow_tasks"] == [] + async def test_task_survives_pickle(self) -> None: - tasks: list[sr.ShadowTask] = [] - sr.enable_shadow_requests(on_task=tasks.append) - await _invoke(_make_client(_variation_a())) - restored = pickle.loads(pickle.dumps(tasks[0])) - assert restored == tasks[0] + task = await _deferred_task(_make_client(_variation_a())) + restored = pickle.loads(pickle.dumps(task)) + assert restored == task assert restored.shadow_request_id == "shadow-1" - async def test_run_shadow_matches_automatic_mode(self) -> None: - tasks: list[sr.ShadowTask] = [] - sr.enable_shadow_requests(on_task=tasks.append) + async def test_run_shadow_matches_background_mode(self) -> None: deferred_client = _make_client(_variation_a()) - await _invoke(deferred_client) - b = _Recorder("ShadowProvider", "shadow answer") - result = await sr.run_shadow(tasks[0], [b.handler]) + task = await _deferred_task(deferred_client) + result = await run_shadow(task, [_Recorder("ShadowProvider").handler]) + assert result.error is None + assert result.track_data["isShadow"] is True deferred = [ (e[0], {k: v for k, v in e[1].items() if k != "runId"}) for e in _shadow_events(deferred_client) ] - assert result.error is None - assert result.track_data["isShadow"] is True lifecycle_module._reset_for_testing() - sr.enable_shadow_requests() - auto_client = _make_client(_variation_a()) - await _invoke(auto_client) - automatic = [ + background_client = _make_client(_variation_a()) + await _invoke(background_client) + background = [ (e[0], {k: v for k, v in e[1].items() if k != "runId"}) - for e in _shadow_events(auto_client) + for e in _shadow_events(background_client) ] - assert deferred == automatic - - async def test_raising_callback_does_not_fail_the_call( - self, caplog: pytest.LogCaptureFixture - ) -> None: - def boom(task: sr.ShadowTask) -> None: - raise RuntimeError("queue full") - - sr.enable_shadow_requests(on_task=boom) - with caplog.at_level(logging.WARNING): - response, a, _ = await _invoke(_make_client(_variation_a())) - assert response.response == "real answer" - assert len(a.calls) == 1 + assert deferred == background async def test_run_shadow_never_raises(self) -> None: - tasks: list[sr.ShadowTask] = [] - sr.enable_shadow_requests(on_task=tasks.append) - await _invoke(_make_client(_variation_a())) - result = await sr.run_shadow( - tasks[0], [_Recorder("ShadowProvider", raises=True).handler] + task = await _deferred_task(_make_client(_variation_a())) + result = await run_shadow( + task, [_Recorder("ShadowProvider", raises=True).handler] ) assert result.error async def test_run_shadow_reports_usage(self) -> None: - tasks: list[sr.ShadowTask] = [] - sr.enable_shadow_requests(on_task=tasks.append) client = _make_client(_variation_a()) - await _invoke(client) - await sr.run_shadow(tasks[0], [_Recorder("ShadowProvider").handler]) + task = await _deferred_task(client) + await run_shadow(task, [_Recorder("ShadowProvider").handler]) helpers = [ c[0][2].get("helper") for c in client.track.call_args_list if c[0][0] == "$ld:ai:sdk:usage" ] - assert "experimental.runShadow" in helpers - - -# --------------------------------------------------------------------------- -# §9.6 Core invocation hook -# --------------------------------------------------------------------------- - - -class TestInvocationHook: - async def test_raising_hook_does_not_fail_the_call( - self, caplog: pytest.LogCaptureFixture - ) -> None: - def bad_hook(event: Any) -> None: - raise RuntimeError("hook boom") - - register_invocation_hook("test-bad-hook", bad_hook) - try: - with caplog.at_level(logging.WARNING): - response, a, _ = await _invoke(_make_client(_variation_a())) - assert response.response == "real answer" - assert len(a.calls) == 1 - assert any( - "hook boom" in r.getMessage() or r.exc_info for r in caplog.records - ) - finally: - register_invocation_hook("test-bad-hook", None) - - async def test_hook_sees_the_run_id_of_the_real_run(self) -> None: - seen: list[Any] = [] - register_invocation_hook("test-seen", seen.append) - try: - response, _, _ = await _invoke(_make_client(_variation_a())) - assert seen[0].run_id == response.track_data["runId"] - assert seen[0].reason == FALLTHROUGH - finally: - register_invocation_hook("test-seen", None) - - -# --------------------------------------------------------------------------- -# §9.7 Shadow events and spans -# --------------------------------------------------------------------------- + assert "client.runShadow" in helpers class TestEvents: - async def test_shadow_events_describe_b(self) -> None: - sr.enable_shadow_requests() + async def test_shadow_events_describe_the_shadow_variation(self) -> None: client = _make_client(_variation_a()) response, _, _ = await _invoke(client) events = _shadow_events(client) @@ -646,7 +515,6 @@ async def test_shadow_events_describe_b(self) -> None: assert data["ownerId"] == "member-1" async def test_nothing_is_added_to_real_events(self) -> None: - sr.enable_shadow_requests() client = _make_client(_variation_a()) response, _, _ = await _invoke(client) assert not (_SHADOW_FIELDS & set(response.track_data)) @@ -666,10 +534,8 @@ def test_span_attributes_are_set_on_the_root_span( mock_span.set_attribute.assert_any_call("launchdarkly.is_shadow", True) async def test_shadow_span_attributes_reach_the_handler(self) -> None: - sr.enable_shadow_requests() _, _, b = await _invoke(_make_client(_variation_a())) - span_attrs = b.calls[0]["variables"]["__ld_span"] - assert span_attrs == { + assert b.calls[0]["variables"]["__ld_span"] == { "launchdarkly.is_shadow": True, "launchdarkly.shadow.request.id": "shadow-1", "launchdarkly.shadow.primary.variation.key": "a", diff --git a/packages/client/tests/test_stream_conversation.py b/packages/client/tests/test_stream_conversation.py index 198f993c..a94dc57c 100644 --- a/packages/client/tests/test_stream_conversation.py +++ b/packages/client/tests/test_stream_conversation.py @@ -47,6 +47,7 @@ def _reset_exporter() -> Iterator[None]: @pytest.fixture def mock_ld_client() -> Iterator[MagicMock]: c = MagicMock() + del c.variation_detail c.track = MagicMock() c.flush = AsyncMock() c.close = AsyncMock() diff --git a/packages/opencode-agents/tests/test_handler.py b/packages/opencode-agents/tests/test_handler.py index 6adc2f20..d321f9b7 100644 --- a/packages/opencode-agents/tests/test_handler.py +++ b/packages/opencode-agents/tests/test_handler.py @@ -277,6 +277,7 @@ def _reset_sdk() -> Iterator[None]: @pytest.fixture def mock_ld_client() -> Iterator[MagicMock]: client = MagicMock() + del client.variation_detail client.track = MagicMock() client.flush = AsyncMock() client.close = AsyncMock() diff --git a/release-please-config.json b/release-please-config.json index 73cfe3de..79ae254a 100644 --- a/release-please-config.json +++ b/release-please-config.json @@ -158,9 +158,6 @@ } }, "changelog-sections": [ - { "type": "feat", "scope": "shadow-requests", "section": "Experimental" }, - { "type": "fix", "scope": "shadow-requests", "section": "Experimental" }, - { "type": "perf", "scope": "shadow-requests", "section": "Experimental" }, { "type": "feat", "section": "Features" }, { "type": "fix", "section": "Bug Fixes" }, { "type": "perf", "section": "Performance Improvements" },