From ee8f6f0cf655c53acc60dedef33c06735cb3627a Mon Sep 17 00:00:00 2001 From: Andrew Klatzke Date: Fri, 2 Oct 2026 15:08:29 -0800 Subject: [PATCH] feat(AIC-3439): let wildcard handlers accept only named providers A LangChain adapter used for one provider no longer has to match every provider, and two scoped wildcards of the same mode can be registered together. Co-authored-by: Cursor --- AGENTS.md | 6 +- packages/client/README.md | 2 +- .../evaluations/runner.py | 62 +++++--- .../src/launchdarkly_ai_server/judges.py | 88 +++-------- .../src/launchdarkly_ai_server/registry.py | 20 ++- .../src/launchdarkly_ai_server/types.py | 4 + .../src/launchdarkly_ai_server/utils.py | 137 ++++++++++++++---- packages/client/tests/test_client.py | 103 ++++++++++++- packages/client/tests/test_evaluations_run.py | 74 ++++++++++ packages/client/tests/test_graph.py | 12 ++ packages/client/tests/test_judges.py | 54 +++++++ packages/client/tests/test_registry.py | 55 ++++++- packages/client/tests/test_utils.py | 53 +++++++ packages/langchain-agents/README.md | 2 + .../launchdarkly_ai_langchain_agents/graph.py | 10 +- .../handler.py | 22 ++- packages/langchain-agents/tests/test_graph.py | 9 ++ .../langchain-agents/tests/test_handler.py | 28 ++++ packages/langchain-messages/README.md | 2 + .../handler.py | 20 ++- .../langchain-messages/tests/test_handler.py | 39 +++++ 21 files changed, 669 insertions(+), 133 deletions(-) diff --git a/AGENTS.md b/AGENTS.md index 305613a9..381a6be8 100644 --- a/AGENTS.md +++ b/AGENTS.md @@ -355,7 +355,7 @@ Returns a `ConfigInstance` with: **Behavior when `.invoke()` is called:** 1. Fetches and validates the `AiConfigRep` variation from LaunchDarkly using `key` and the supplied `context`. Raises if the variation is disabled or invalid. -2. Selects the handler by matching on `[config.provider.name, normalized mode]`. Selection priority: (a) exact provider match, (b) wildcard `['*', mode]` fallback for multi-provider adapters (e.g. LangChain). Raises if no matching handler is found. +2. Selects the handler by matching on `[config.provider.name, normalized mode]`. Selection priority: (a) exact provider match, ignoring any `providers` list; (b) a wildcard `['*', mode]` whose `providers` list contains the name, shortest list first and earlier registration on a tie; (c) an unscoped wildcard `['*', mode]`. Raises if no matching handler is found. 3. Invokes the selected handler with the config, user input, tool handlers, variables, and history. The `context` passed to `.invoke()` is automatically merged into `variables` under the key `ldContext`, so templates can reference `{{ldContext.key}}`, `{{ldContext.email}}`, etc. If `history` is provided, it is passed to the handler as the 5th positional argument — messages-mode handlers splice it into the messages array; agent-mode handlers append it to the system prompt. 4. Emits LaunchDarkly telemetry events: duration (`$ld:ai:duration:total`), outcome (`$ld:ai:generation:success` / `$ld:ai:generation:error`), and token counts (`$ld:ai:tokens:*`). 5. If `judgeConfiguration` is present: @@ -390,7 +390,7 @@ registry = Registry( ) ``` -`.register(handlers=[], tools={})` can be called multiple times to add more handlers or tools. Duplicate `provides_for` keys or tool names produce a warning and the last registration wins. +`.register(handlers=[], tools={})` can be called multiple times to add more handlers or tools. Handler identity is `provides_for` plus the provider set: two wildcards of the same mode with different `providers` lists both stay registered, while the same set (order-independent; repeated names collapse) or two unscoped wildcards warn and the last registration wins. Duplicate tool names do the same. `global_registry` is a pre-constructed singleton `Registry` instance. @@ -441,7 +441,7 @@ async def handler( handler.provides_for = [provider_name: str, mode: Literal["agent", "messages"]] ``` -The `provides_for` list is how `config()` routes to the correct handler at runtime. The mode element must exactly match the normalized `meta.mode`. The provider element must either exactly match `config.provider.name` **or** be the wildcard `'*'`. A wildcard handler is chosen only when no handler with an exact provider name matches — it acts as a fallback for multi-provider adapters like LangChain. **Always attach `provides_for` using `create_handler` rather than direct attribute assignment.** +The `provides_for` list is how `config()` routes to the correct handler at runtime. The mode element must exactly match the normalized `meta.mode`. The provider element must either exactly match `config.provider.name` **or** be the wildcard `'*'`. A wildcard handler is chosen only when no handler with an exact provider name matches — it acts as a fallback for multi-provider adapters like LangChain. Pass `providers=["Bedrock"]` to `create_handler` to limit which `config.provider.name` values that wildcard accepts; names must match exactly, and an empty list or a blank name is rejected. Omit `providers` to accept every provider. **Always attach `provides_for` using `create_handler` rather than direct attribute assignment.** ### Factory Function diff --git a/packages/client/README.md b/packages/client/README.md index b2607d4e..57f0d826 100644 --- a/packages/client/README.md +++ b/packages/client/README.md @@ -161,7 +161,7 @@ Two limits keep a trajectory from spending the judge's context window: at most 5 A tool result is now judge-prompt input. It stays literal for the same reason the generated output does: the judge config is handed to the handler unrendered and the handler makes exactly one template pass, so a `{{...}}` sequence coming back from a tool is never expanded into the judge prompt. -**Judges are independent AI Configs, so handlers are routed per judge.** A judge may resolve to a different provider or mode than `generation`, and a handler built for one provider cannot execute another's config. `handler` runs a judge when it provides for that judge's provider; pass handlers for any other providers in `judge_handlers`. Selection prefers a handler naming the judge's provider outright over a wildcard multi-provider adapter, and an agent-mode handler can serve a messages-mode judge with its messages collapsed into one instructions block. A plain callable that declares no `provides_for` routes itself, exactly as it already does for the generation config. +**Judges are independent AI Configs, so handlers are routed per judge.** A judge may resolve to a different provider or mode than `generation`, and a handler built for one provider cannot execute another's config. `handler` runs a judge when it provides for that judge's provider; pass handlers for any other providers in `judge_handlers`. Selection prefers a handler naming the judge's provider outright over a wildcard multi-provider adapter, and an agent-mode handler can serve a messages-mode judge with its messages collapsed into one instructions block. A wildcard's `providers` list, when set, accepts only those names. A plain callable that declares no `provides_for` routes itself, exactly as it already does for the generation config. Judges are resolved through flag delivery, and handlers are matched to them, **before** any evaluation records are created — a missing judge or one no handler covers fails the run up front rather than after the generation spend. After that point a criterion failure never aborts the run: an unparseable judge response, an out-of-range score, a raising handler or scorer, and a row whose generation errored each become a per-criterion `ERROR` event with a cause code (`invalid_judge_output`, `invalid_score`, `handler_raised`, `scorer_raised`, `generation_incomplete`) and a top-level `errorMessage`. Event *delivery* is different: the backend needs one result per `(row, criterion)` to finish row accounting, so if tracking a criterion event fails, every remaining result is still attempted and flushed and then `run()` raises — rather than polling to its timeout with the cause hidden. diff --git a/packages/client/src/launchdarkly_ai_server/evaluations/runner.py b/packages/client/src/launchdarkly_ai_server/evaluations/runner.py index 6240e980..e728c092 100644 --- a/packages/client/src/launchdarkly_ai_server/evaluations/runner.py +++ b/packages/client/src/launchdarkly_ai_server/evaluations/runner.py @@ -85,11 +85,19 @@ def _provides_for( return None -def _covers_provider( - provides_for: tuple[str, Literal["agent", "messages"]], - provider: str | None, -) -> bool: - return provides_for[0] == provider or provides_for[0] == "*" +def _covers_provider(handler: EvalHandler, provider: str | None) -> bool: + """Exact provider name, or a wildcard whose allowlist contains it.""" + provides_for = _provides_for(handler) + if provides_for is None or provider is None: + return False + if provides_for[0] == provider: + return True + if provides_for[0] != "*": + return False + allow = getattr(handler, "providers", None) + if allow is None: + return True + return provider in tuple(allow) def _find_judge_handler( @@ -104,19 +112,37 @@ def _find_judge_handler( ``config()`` already applies to a generation config. Searching in one pass would instead let the order the caller happened to list its handlers in decide, sending an OpenAI judge through a LangChain adapter that was merely - listed first. + listed first. Among wildcards that accept the name, the shortest allowlist + wins and an unscoped ``*`` is the last resort. """ - for exact in (True, False): - for candidate in judge_handlers: - provides_for = _provides_for(candidate) - if provides_for is None or provides_for[1] != mode: - continue - if exact: - if provides_for[0] == provider: - return candidate - elif provides_for[0] == "*": - return candidate - return None + for candidate in judge_handlers: + provides_for = _provides_for(candidate) + if ( + provides_for is not None + and provider is not None + and provides_for[0] == provider + and provides_for[1] == mode + ): + return candidate + + scoped: list[tuple[int, int, EvalHandler]] = [] + unscoped: EvalHandler | None = None + for index, candidate in enumerate(judge_handlers): + provides_for = _provides_for(candidate) + if provides_for is None or provides_for[0] != "*" or provides_for[1] != mode: + continue + allow = getattr(candidate, "providers", None) + if allow is None: + if unscoped is None: + unscoped = candidate + continue + names = tuple(allow) + if provider is not None and provider in names: + scoped.append((len(names), index, candidate)) + if scoped: + scoped.sort(key=lambda item: (item[0], item[1])) + return scoped[0][2] + return unscoped def _select_judge_handler( @@ -153,7 +179,7 @@ def _select_judge_handler( generation_provides_for = _provides_for(handler) if generation_provides_for is None: return JudgeExecution(resolved=resolved, handler=handler) - if _covers_provider(generation_provides_for, resolved.provider): + if _covers_provider(handler, resolved.provider): return JudgeExecution( resolved=resolved, handler=handler, diff --git a/packages/client/src/launchdarkly_ai_server/judges.py b/packages/client/src/launchdarkly_ai_server/judges.py index 56d7e032..5dab8985 100644 --- a/packages/client/src/launchdarkly_ai_server/judges.py +++ b/packages/client/src/launchdarkly_ai_server/judges.py @@ -27,17 +27,20 @@ from .utils import ( normalize_mode, omit_model_stamps, + select_mode_handler, to_ld_context, to_usage_dict, + wildcard_covers, ) def _provider_matches(handler: ProviderHandler, provider: str | None) -> bool: - """Returns True when the handler covers the given provider or is a wildcard.""" - return bool( - handler.provides_for - and (handler.provides_for[0] == provider or handler.provides_for[0] == "*") - ) + """True when the handler names the provider or its wildcard allowlist contains it.""" + if not handler.provides_for: + return False + if provider is not None and handler.provides_for[0] == provider: + return True + return wildcard_covers(handler, provider) logger = logging.getLogger(__name__) @@ -109,28 +112,10 @@ async def run_judges( judge_handler: ProviderHandler = handler collapse_messages = False if handlers: - exact = next( - ( - h - for h in handlers - if _provider_matches(h, judge_provider) - and h.provides_for - and h.provides_for[1] == judge_mode - ), - None, - ) + exact = select_mode_handler(handlers, judge_provider, judge_mode) agent_fallback = ( - next( - ( - h - for h in handlers - if _provider_matches(h, judge_provider) - and h.provides_for - and h.provides_for[1] == "agent" - ), - None, - ) - if not exact and judge_mode == "messages" + select_mode_handler(handlers, judge_provider, "agent") + if exact is None and judge_mode == "messages" else None ) if exact: @@ -145,6 +130,10 @@ async def run_judges( and handler.provides_for is not None and handler.provides_for[1] == "agent" ) + else: + # No compatible handler — skip rather than calling a provider + # whose allowlist does not include this judge. + continue effective_judge_config = ( _collapse_messages_to_instructions(judge_ai_config) @@ -267,28 +256,10 @@ async def build_judge_tasks( collapse_messages = False if handlers: - exact = next( - ( - h - for h in handlers - if _provider_matches(h, judge_provider) - and h.provides_for - and h.provides_for[1] == judge_mode - ), - None, - ) + exact = select_mode_handler(handlers, judge_provider, judge_mode) agent_fallback = ( - next( - ( - h - for h in handlers - if _provider_matches(h, judge_provider) - and h.provides_for - and h.provides_for[1] == "agent" - ), - None, - ) - if not exact and judge_mode == "messages" + select_mode_handler(handlers, judge_provider, "agent") + if exact is None and judge_mode == "messages" else None ) if exact: @@ -353,27 +324,10 @@ async def run_judge( """ from .tracking import execute_and_track - def _matches(h: ProviderHandler) -> bool: - return _provider_matches(h, task.judge_provider) - - exact = next( - ( - h - for h in handlers - if _matches(h) and h.provides_for and h.provides_for[1] == task.judge_mode - ), - None, - ) + exact = select_mode_handler(handlers, task.judge_provider, task.judge_mode) agent_fallback = ( - next( - ( - h - for h in handlers - if _matches(h) and h.provides_for and h.provides_for[1] == "agent" - ), - None, - ) - if task.judge_mode == "messages" and not exact + select_mode_handler(handlers, task.judge_provider, "agent") + if task.judge_mode == "messages" and exact is None else None ) diff --git a/packages/client/src/launchdarkly_ai_server/registry.py b/packages/client/src/launchdarkly_ai_server/registry.py index bceb8388..536f97d7 100644 --- a/packages/client/src/launchdarkly_ai_server/registry.py +++ b/packages/client/src/launchdarkly_ai_server/registry.py @@ -9,6 +9,18 @@ logger = logging.getLogger(__name__) +def _registration_key( + handler: ProviderHandler, +) -> tuple[tuple[str, str], frozenset[str] | None]: + """``provides_for`` plus the allowlist as a set. Order and repeated names do not make a new identity.""" + provides_for = handler.provides_for + if provides_for is None: + raise ValueError("handler is missing provides_for") + providers = handler.providers + scope = None if providers is None else frozenset(providers) + return (provides_for, scope) + + class Registry: """ Manages handlers and tools for ``routed_model`` and ``graph``. @@ -42,13 +54,15 @@ def register( if handlers: for handler in handlers: if handler.provides_for is not None: - # Check for duplicate by providesFor key - key = handler.provides_for + # Identity is provides_for plus the provider set. Two scoped + # wildcards of the same mode stay registered when their lists differ. + key = _registration_key(handler) existing_idx = next( ( i for i, h in enumerate(self._handlers) - if h.provides_for == key + if h.provides_for is not None + and _registration_key(h) == key ), None, ) diff --git a/packages/client/src/launchdarkly_ai_server/types.py b/packages/client/src/launchdarkly_ai_server/types.py index 761faf5b..e0eec8c2 100644 --- a/packages/client/src/launchdarkly_ai_server/types.py +++ b/packages/client/src/launchdarkly_ai_server/types.py @@ -148,6 +148,7 @@ class ProviderHandler: - ``__call__`` — blocking invocation - ``stream`` — optional async-generator streaming (may be ``None``) - ``provides_for`` — ``(provider_name, mode)`` tuple or ``None`` + - ``providers`` — optional allowlist consulted only when ``provides_for`` names ``"*"`` - ``capture_content`` — whether this handler was built with content capture on ``capture_content`` is declared here so the client core can apply the handler's own content @@ -156,6 +157,7 @@ class ProviderHandler: """ provides_for: tuple[str, Literal["agent", "messages"]] | None + providers: tuple[str, ...] | None capture_content: bool def __init__( @@ -164,11 +166,13 @@ def __init__( provides_for: tuple[str, Literal["agent", "messages"]] | None = None, stream_fn: _StreamFn | None = None, capture_content: bool = False, + providers: tuple[str, ...] | None = None, ) -> None: self._fn = fn self.provides_for = provides_for self._stream_fn = stream_fn self.capture_content = capture_content + self.providers = providers async def __call__( self, diff --git a/packages/client/src/launchdarkly_ai_server/utils.py b/packages/client/src/launchdarkly_ai_server/utils.py index 2958de6e..8d357863 100644 --- a/packages/client/src/launchdarkly_ai_server/utils.py +++ b/packages/client/src/launchdarkly_ai_server/utils.py @@ -2,7 +2,7 @@ import json import re -from collections.abc import Mapping +from collections.abc import Mapping, Sequence from dataclasses import dataclass, field from typing import Any, Literal @@ -18,24 +18,117 @@ ) +def normalize_providers(providers: Sequence[str] | None) -> tuple[str, ...] | None: + """ + Validates an optional wildcard allowlist. + + ``None`` means the handler is unscoped. A list must contain at least one + non-blank name and is stored in the caller's order. Names are compared to + ``config.provider.name`` exactly, so surrounding whitespace is rejected + rather than trimmed. + """ + if providers is None: + return None + if ( + isinstance(providers, (str, bytes)) + or len(providers) == 0 + or any(not isinstance(name, str) or name.strip() == "" for name in providers) + ): + raise ValueError("providers must be a non-empty list of provider names") + return tuple(providers) + + def create_handler( provides_for: tuple[str, Literal["agent", "messages"]], fn: _HandlerFn, stream_fn: _StreamFn | None = None, capture_content: bool = False, + *, + providers: Sequence[str] | None = None, ) -> ProviderHandler: """ Wraps a plain async callable in a :class:`ProviderHandler` with the given ``provides_for`` metadata and optional streaming implementation. + + ``providers`` is consulted only when ``provides_for`` starts with ``"*"``. + It limits which ``config.provider.name`` values that wildcard accepts. """ return ProviderHandler( fn=fn, provides_for=provides_for, stream_fn=stream_fn, capture_content=capture_content, + providers=normalize_providers(providers), ) +def wildcard_covers(handler: Any, provider: str | None) -> bool: + """Whether a ``"*"`` handler accepts *provider*. An allowlist is a filter; an absent one accepts every name.""" + provides_for = getattr(handler, "provides_for", None) + if not provides_for or provides_for[0] != "*": + return False + allow = getattr(handler, "providers", None) + if allow is None: + return True + return provider is not None and provider in tuple(allow) + + +def best_wildcard(handlers: list[Any], provider: str | None, mode: str) -> Any | None: + """ + The ``"*"`` handler of *mode* that accepts *provider*. + + A scoped list beats an unscoped wildcard. Among scoped lists that contain + the name, the shorter list wins, and equal lengths keep the earlier + registration. + """ + scoped: list[tuple[int, int, Any]] = [] + unscoped: Any | None = None + for index, handler in enumerate(handlers): + provides_for = getattr(handler, "provides_for", None) + if not provides_for or provides_for[1] != mode: + continue + if not wildcard_covers(handler, provider): + continue + allow = getattr(handler, "providers", None) + if allow is None: + if unscoped is None: + unscoped = handler + continue + scoped.append((len(tuple(allow)), index, handler)) + if scoped: + scoped.sort(key=lambda item: (item[0], item[1])) + return scoped[0][2] + return unscoped + + +def select_mode_handler( + handlers: list[Any], provider: str | None, mode: str +) -> Any | None: + """Exact ``(provider, mode)`` first, then the best wildcard of that mode.""" + if provider: + exact = next( + ( + handler + for handler in handlers + if getattr(handler, "provides_for", None) == (provider, mode) + ), + None, + ) + if exact: + return exact + return best_wildcard(handlers, provider, mode) + + +def covers_provider_name(handler: Any, provider: str | None) -> bool: + """True when the handler names *provider* or its wildcard allowlist contains it, in any mode.""" + provides_for = getattr(handler, "provides_for", None) + if not provides_for or provider is None: + return False + if provides_for[0] == provider: + return True + return wildcard_covers(handler, provider) + + def collapse_messages_to_instructions(config: AiConfigRep) -> AiConfigRep: """ When only an agent handler is available for a messages-mode config, collapse @@ -483,15 +576,18 @@ def select_handler( Selects a handler from *handlers* based on the provider and mode in *config*/*meta*. Resolution order: - 1. Exact ``(provider, mode)`` match. - 2. Wildcard ``("*", mode)`` match — for multi-provider adapters like LangChain. - 3. (Non-strict only) Provider-only match, then single-handler fallback. + 1. Exact ``(provider, mode)`` match. A ``providers`` list on an exact handler is ignored. + 2. Wildcard ``("*", mode)`` whose ``providers`` list contains the name. The + shortest list wins; equal lengths keep the earlier registration. + 3. Unscoped wildcard ``("*", mode)`` (``providers`` is ``None``). + 4. (Non-strict only) Provider-only match, then single-handler fallback. + A sole handler whose allowlist omits the name is not used. - When *strict* is ``True`` (default, used by ``config()``), steps 1–2 are tried - and a descriptive error is raised if neither matches. + When *strict* is ``True`` (default, used by ``config()``), steps 1–3 are tried + and a descriptive error is raised if none match. When *strict* is ``False`` (used by ``graph()``), resolution falls back - progressively: exact match → wildcard match → provider-only match → single-handler fallback. + progressively through the provider-only match and the single-handler fallback. """ provider = ( (config.get("provider") or {}).get("name") if isinstance(config, dict) else None @@ -501,26 +597,13 @@ def select_handler( mode = normalize_mode(meta.get("mode") if isinstance(meta, dict) else None) - exact = next((h for h in handlers if h.provides_for == (provider, mode)), None) - if exact: - return exact - - wildcard = next( - ( - h - for h in handlers - if h.provides_for and h.provides_for[0] == "*" and h.provides_for[1] == mode - ), - None, - ) - if wildcard: - return wildcard + chosen = select_mode_handler(handlers, provider, mode) + if isinstance(chosen, ProviderHandler): + return chosen if strict: has_coverage = any( - h.provides_for - and (h.provides_for[0] == provider or h.provides_for[0] == "*") - for h in handlers + covers_provider_name(handler, provider) for handler in handlers ) if not has_coverage: raise ValueError(f"Handler for provider {provider} not found") @@ -534,7 +617,11 @@ def select_handler( return by_provider if len(handlers) == 1: - return handlers[0] + only = handlers[0] + allow = getattr(only, "providers", None) + if allow is not None and provider not in tuple(allow): + raise ValueError(f"Handler for provider {provider} not found") + return only raise ValueError(f"Handler for provider {provider} not found") diff --git a/packages/client/tests/test_client.py b/packages/client/tests/test_client.py index d7357308..d61357f6 100644 --- a/packages/client/tests/test_client.py +++ b/packages/client/tests/test_client.py @@ -407,11 +407,33 @@ async def test_parse_failure_returns_raw_string_when_output_format_set( # --------------------------------------------------------------------------- -def _handler_for(provider: str, mode: str, response: str = "ok") -> ProviderHandler: +def _handler_for( + provider: str, + mode: str, + response: str = "ok", + providers: list[str] | None = None, +) -> ProviderHandler: async def fn(cfg, user_input, tool_handlers, variables, history=None) -> dict: # type: ignore[override] return {"output": response, "usage": {"input_tokens": 1, "output_tokens": 1}} - return ProviderHandler(fn=fn, provides_for=(provider, mode)) # type: ignore[arg-type] + handler = ProviderHandler(fn=fn, provides_for=(provider, mode)) # type: ignore[arg-type] + if providers is not None: + handler.providers = tuple(providers) # type: ignore[attr-defined] + return handler + + +def _variation(provider: str, mode: str = "messages") -> dict[str, Any]: + return { + "model": {"name": "gpt-4"}, + "provider": {"name": provider}, + "instructions": "hi", + "_ldMeta": { + "enabled": True, + "variationKey": "v1", + "version": 1, + "mode": mode, + }, + } class TestConfigMultiHandler: @@ -509,6 +531,83 @@ async def test_wildcard_handler_does_not_match_different_mode( with pytest.raises((ValueError, RuntimeError)): await rm.invoke("hi", CONTEXT) + async def test_scoped_wildcard_matches_a_listed_provider( + self, mock_ld_client: MagicMock + ) -> None: + mock_ld_client.variation = AsyncMock(return_value=_variation("Bedrock")) + h = _handler_for("*", "messages", "from-bedrock", providers=["Bedrock"]) + result = await config(key="flag", handler=[h]).invoke("hi", CONTEXT) + assert result.response == "from-bedrock" + + async def test_scoped_wildcard_rejects_an_unlisted_provider( + self, mock_ld_client: MagicMock + ) -> None: + mock_ld_client.variation = AsyncMock(return_value=_variation("OpenAI")) + h = _handler_for("*", "messages", "from-bedrock", providers=["Bedrock"]) + with pytest.raises(ValueError, match="OpenAI"): + await config(key="flag", handler=[h]).invoke("hi", CONTEXT) + + async def test_two_scoped_wildcards_divide_providers( + self, mock_ld_client: MagicMock + ) -> None: + bedrock = _handler_for("*", "messages", "bedrock", providers=["Bedrock"]) + anthropic = _handler_for("*", "messages", "anthropic", providers=["Anthropic"]) + mock_ld_client.variation = AsyncMock(return_value=_variation("Anthropic")) + result = await config(key="flag", handler=[bedrock, anthropic]).invoke( + "hi", CONTEXT + ) + assert result.response == "anthropic" + + async def test_exact_handler_wins_over_scoped_wildcard( + self, mock_ld_client: MagicMock + ) -> None: + scoped = _handler_for("*", "messages", "scoped", providers=["TestProvider"]) + exact = _handler_for("TestProvider", "messages", "exact") + result = await config(key="flag", handler=[scoped, exact]).invoke("hi", CONTEXT) + assert result.response == "exact" + + async def test_unscoped_wildcard_is_the_fallback( + self, mock_ld_client: MagicMock + ) -> None: + scoped = _handler_for("*", "messages", "bedrock", providers=["Bedrock"]) + unscoped = _handler_for("*", "messages", "any") + result = await config(key="flag", handler=[scoped, unscoped]).invoke( + "hi", CONTEXT + ) + assert result.response == "any" + + async def test_shorter_provider_list_wins(self, mock_ld_client: MagicMock) -> None: + mock_ld_client.variation = AsyncMock(return_value=_variation("Bedrock")) + wide = _handler_for("*", "messages", "wide", providers=["Bedrock", "Anthropic"]) + narrow = _handler_for("*", "messages", "narrow", providers=["Bedrock"]) + result = await config(key="flag", handler=[wide, narrow]).invoke("hi", CONTEXT) + assert result.response == "narrow" + + async def test_equal_length_overlap_keeps_the_earlier_handler( + self, mock_ld_client: MagicMock + ) -> None: + mock_ld_client.variation = AsyncMock(return_value=_variation("Bedrock")) + earlier = _handler_for( + "*", "messages", "earlier", providers=["Bedrock", "Gemini"] + ) + later = _handler_for( + "*", "messages", "later", providers=["Bedrock", "Anthropic"] + ) + result = await config(key="flag", handler=[earlier, later]).invoke( + "hi", CONTEXT + ) + assert result.response == "earlier" + + async def test_scoped_wildcard_does_not_match_a_different_mode( + self, mock_ld_client: MagicMock + ) -> None: + mock_ld_client.variation = AsyncMock( + return_value=_variation("Bedrock", mode="agent") + ) + h = _handler_for("*", "messages", "bedrock", providers=["Bedrock"]) + with pytest.raises(ValueError, match="Bedrock"): + await config(key="flag", handler=[h]).invoke("hi", CONTEXT) + async def test_resolves_handlers_from_registry( self, mock_ld_client: MagicMock ) -> None: diff --git a/packages/client/tests/test_evaluations_run.py b/packages/client/tests/test_evaluations_run.py index 4a4b8571..ba1b79e2 100644 --- a/packages/client/tests/test_evaluations_run.py +++ b/packages/client/tests/test_evaluations_run.py @@ -2091,6 +2091,80 @@ async def wildcard_judge( assert [config["provider"]["name"] for config in judged] == ["Anthropic"] +@pytest.mark.asyncio +async def test_scoped_wildcard_does_not_run_a_judge_outside_its_list( + monkeypatch: pytest.MonkeyPatch, +) -> None: + transport = judge_run_transport() + judge_variation(monkeypatch, provider="Anthropic") + evals = init_evaluations(api_token="token", sdk_key="sdk-key", transport=transport) + + async def bedrock_judge( + config: dict[str, Any], + user_input: str | None = None, + tool_handlers: dict[str, Callable[..., Any]] | None = None, + variables: dict[str, Any] | None = None, + history: list[dict[str, Any]] | None = None, + ) -> dict[str, Any]: + return {"output": '{"score": 1, "reasoning": "ok"}'} + + scoped = create_handler(("*", "messages"), bedrock_judge) + scoped.providers = ("Bedrock",) # type: ignore[attr-defined] + + with pytest.raises(EvaluationsError, match="Anthropic"): + await evals.run( + project_key="proj", + key="support-qa", + dataset="golden", + handler=create_handler(("OpenAI", "messages"), _generation_only), + generation={"provider": "OpenAI", "model": "gpt-4o"}, + criteria=[Judge(key="$ld:ai:judge:accuracy")], + judge_handlers=[scoped], + ) + + +@pytest.mark.asyncio +async def test_two_scoped_judge_handlers_divide_providers( + monkeypatch: pytest.MonkeyPatch, + stub_sdk_client: MagicMock, +) -> None: + transport = judge_run_transport() + judge_variation(monkeypatch, provider="Anthropic") + evals = init_evaluations(api_token="token", sdk_key="sdk-key", transport=transport) + chosen: list[str] = [] + + def judge_handler(name: str) -> Any: + async def run( + config: dict[str, Any], + user_input: str | None = None, + tool_handlers: dict[str, Callable[..., Any]] | None = None, + variables: dict[str, Any] | None = None, + history: list[dict[str, Any]] | None = None, + ) -> dict[str, Any]: + chosen.append(name) + return {"output": '{"score": 1, "reasoning": "ok"}'} + + return run + + bedrock = create_handler(("*", "messages"), judge_handler("bedrock")) + bedrock.providers = ("Bedrock",) # type: ignore[attr-defined] + anthropic = create_handler(("*", "messages"), judge_handler("anthropic")) + anthropic.providers = ("Anthropic",) # type: ignore[attr-defined] + + result = await evals.run( + project_key="proj", + key="support-qa", + dataset="golden", + handler=create_handler(("OpenAI", "messages"), _generation_only), + generation={"provider": "OpenAI", "model": "gpt-4o"}, + criteria=[Judge(key="$ld:ai:judge:accuracy")], + judge_handlers=[bedrock, anthropic], + ) + + assert result.passed is True + assert chosen == ["anthropic"] + + @pytest.mark.asyncio async def test_agent_handler_runs_a_messages_mode_judge_with_collapsed_messages( monkeypatch: pytest.MonkeyPatch, diff --git a/packages/client/tests/test_graph.py b/packages/client/tests/test_graph.py index a172e021..3abb6c5d 100644 --- a/packages/client/tests/test_graph.py +++ b/packages/client/tests/test_graph.py @@ -170,6 +170,18 @@ async def failing_variation(key: str, ctx: dict, default: Any) -> Any: # --------------------------------------------------------------------------- +class TestScopedWildcard: + async def test_sole_scoped_handler_does_not_run_an_unlisted_provider( + self, mock_ld_client: MagicMock + ) -> None: + handler = _make_handler() + handler.provides_for = ("*", "messages") # type: ignore[assignment] + handler.providers = ("Bedrock",) # type: ignore[attr-defined] + g = graph("graph-key", handlers=[handler]) + with pytest.raises(ValueError, match="TestProvider"): + await g.invoke("hi", CONTEXT) + + class TestGraphInvoke: async def test_throws_when_graph_disabled(self, mock_ld_client: MagicMock) -> None: mock_ld_client.variation = AsyncMock(return_value={"edges": {}}) diff --git a/packages/client/tests/test_judges.py b/packages/client/tests/test_judges.py index 5d560948..309af74a 100644 --- a/packages/client/tests/test_judges.py +++ b/packages/client/tests/test_judges.py @@ -194,6 +194,60 @@ async def recording_fn( assert effective.get("instructions") is not None assert effective.get("messages") == [] + async def test_scoped_wildcard_skips_a_judge_outside_its_list( + self, mock_ld_client: MagicMock + ) -> None: + called = False + + async def recording_fn( + config, user_input, tool_handlers, variables, history=None + ) -> dict: # type: ignore[override] + nonlocal called + called = True + return {"output": '{"score": 0.8, "reasoning": "ok"}', "usage": {}} + + scoped = ProviderHandler( + fn=recording_fn, + provides_for=("*", "agent"), # type: ignore[arg-type] + ) + scoped.providers = ("Bedrock",) # type: ignore[attr-defined] + parent = ProviderHandler( + fn=recording_fn, + provides_for=("OpenAI", "messages"), # type: ignore[arg-type] + ) + mock_ld_client.variation = AsyncMock( + return_value={ + "model": {"name": "claude-3-5-sonnet"}, + "provider": {"name": "Anthropic"}, + "messages": [{"role": "user", "content": "Evaluate this."}], + "_ldMeta": { + "enabled": True, + "variationKey": "j1", + "version": 1, + "mode": "judge", + }, + } + ) + config = { + "model": {"name": "gpt-4"}, + "provider": {"name": "OpenAI"}, + "instructions": "hi", + "judgeConfiguration": {"judges": [{"key": "judge-1", "samplingRate": 1.0}]}, + } + import random + + with patch.object(random, "random", return_value=0.0): + await run_judges( + config=config, + user_context=CONTEXT, + handler=parent, + handlers=[scoped], + user_input="q", + llm_response="response", + base_track_data={"runId": "x"}, + ) + assert called is False + async def test_exact_agent_handler_fallback_collapses_messages( self, mock_ld_client: MagicMock ) -> None: diff --git a/packages/client/tests/test_registry.py b/packages/client/tests/test_registry.py index 52a9fedd..f178bc06 100644 --- a/packages/client/tests/test_registry.py +++ b/packages/client/tests/test_registry.py @@ -19,11 +19,15 @@ def _make_handler( provides_for: tuple[str, str] | None = ("Test", "messages"), + providers: list[str] | None = None, ) -> ProviderHandler: async def fn(config, user_input, tool_handlers, variables): # type: ignore[override] return {"output": "ok"} - return ProviderHandler(fn=fn, provides_for=provides_for) # type: ignore[arg-type] + handler = ProviderHandler(fn=fn, provides_for=provides_for) # type: ignore[arg-type] + if providers is not None: + handler.providers = tuple(providers) # type: ignore[attr-defined] + return handler # --------------------------------------------------------------------------- @@ -63,6 +67,43 @@ def test_duplicate_handler_warns_and_replaces( assert r.handlers[0] is h2 assert any("already registered" in m for m in caplog.messages) + def test_scoped_wildcards_with_different_lists_both_stay(self) -> None: + r = Registry() + bedrock = _make_handler(("*", "agent"), providers=["Bedrock"]) + anthropic = _make_handler(("*", "agent"), providers=["Anthropic"]) + r.register(handlers=[bedrock, anthropic]) + assert r.handlers == [bedrock, anthropic] + + def test_unscoped_wildcard_stays_beside_a_scoped_one(self) -> None: + r = Registry() + unscoped = _make_handler(("*", "agent")) + scoped = _make_handler(("*", "agent"), providers=["Bedrock"]) + r.register(handlers=[unscoped, scoped]) + assert r.handlers == [unscoped, scoped] + + def test_same_provider_set_warns_and_replaces_regardless_of_order( + self, caplog: pytest.LogCaptureFixture + ) -> None: + r = Registry() + first = _make_handler(("*", "agent"), providers=["Anthropic", "Bedrock"]) + second = _make_handler(("*", "agent"), providers=["Bedrock", "Anthropic"]) + with caplog.at_level(logging.WARNING): + r.register(handlers=[first]) + r.register(handlers=[second]) + assert r.handlers == [second] + assert any("already registered" in message for message in caplog.messages) + + def test_repeated_provider_name_is_the_same_identity( + self, caplog: pytest.LogCaptureFixture + ) -> None: + r = Registry() + first = _make_handler(("*", "agent"), providers=["Bedrock"]) + second = _make_handler(("*", "agent"), providers=["Bedrock", "Bedrock"]) + with caplog.at_level(logging.WARNING): + r.register(handlers=[first, second]) + assert r.handlers == [second] + assert any("already registered" in message for message in caplog.messages) + def test_handler_without_provides_for_always_appended(self) -> None: r = Registry() h1 = _make_handler(None) @@ -117,6 +158,18 @@ def test_b_overrides_a_on_handler_conflict(self) -> None: assert len(result.handlers) == 1 assert result.handlers[0] is h_b + def test_compose_keeps_scoped_wildcards_with_different_lists(self) -> None: + bedrock = _make_handler(("*", "agent"), providers=["Bedrock"]) + anthropic = _make_handler(("*", "agent"), providers=["Anthropic"]) + result = compose(Registry(handlers=[bedrock]), Registry(handlers=[anthropic])) + assert result.handlers == [bedrock, anthropic] + + def test_compose_replaces_the_same_provider_set(self) -> None: + first = _make_handler(("*", "agent"), providers=["Bedrock", "Anthropic"]) + second = _make_handler(("*", "agent"), providers=["Anthropic", "Bedrock"]) + result = compose(Registry(handlers=[first]), Registry(handlers=[second])) + assert result.handlers == [second] + def test_b_overrides_a_on_tool_conflict(self) -> None: fn_a, fn_b = AsyncMock(), AsyncMock() a = Registry(tools={"t": fn_a}) diff --git a/packages/client/tests/test_utils.py b/packages/client/tests/test_utils.py index 88f05298..c0861f3f 100644 --- a/packages/client/tests/test_utils.py +++ b/packages/client/tests/test_utils.py @@ -195,6 +195,59 @@ async def fn( h = create_handler(("MyProvider", "messages"), fn) assert h.provides_for == ("MyProvider", "messages") + def test_providers_defaults_to_none(self) -> None: + async def fn( + config: object, + user_input: object, + tool_handlers: object, + variables: object, + history: object = None, + ) -> dict: # type: ignore[override] + return {"output": "ok"} + + h = create_handler(("*", "messages"), fn) + assert h.providers is None + + def test_attaches_providers(self) -> None: + async def fn( + config: object, + user_input: object, + tool_handlers: object, + variables: object, + history: object = None, + ) -> dict: # type: ignore[override] + return {"output": "ok"} + + h = create_handler(("*", "messages"), fn, providers=["Bedrock", "Anthropic"]) + assert h.provides_for == ("*", "messages") + assert h.providers == ("Bedrock", "Anthropic") + + def test_rejects_an_empty_provider_list(self) -> None: + async def fn( + config: object, + user_input: object, + tool_handlers: object, + variables: object, + history: object = None, + ) -> dict: # type: ignore[override] + return {"output": "ok"} + + with pytest.raises(ValueError, match="providers"): + create_handler(("*", "messages"), fn, providers=[]) + + def test_rejects_a_blank_provider_name(self) -> None: + async def fn( + config: object, + user_input: object, + tool_handlers: object, + variables: object, + history: object = None, + ) -> dict: # type: ignore[override] + return {"output": "ok"} + + with pytest.raises(ValueError, match="providers"): + create_handler(("*", "messages"), fn, providers=["Bedrock", ""]) + def test_returns_callable(self) -> None: async def fn( config: object, diff --git a/packages/langchain-agents/README.md b/packages/langchain-agents/README.md index 71599fed..20c90f9c 100644 --- a/packages/langchain-agents/README.md +++ b/packages/langchain-agents/README.md @@ -4,6 +4,8 @@ LangChain handler for `launchdarkly-ai-server` using **LangGraph's `StateGraph`* **`provides_for`:** `['*', 'agent']` — matches any flag variation where `meta.mode` is `"agent"` and no more-specific handler is registered. LangChain is a framework adapter, not a provider: it routes through `langchain-anthropic`, `langchain-openai`, and others at runtime based on `config.provider.name`. Use `'*'` so that flags configured with `provider.name = "Anthropic"` or `"OpenAI"` are automatically handled without requiring a separate native handler. +Pass `providers=["Bedrock"]` to `create_langchain_agents_handler`, `langchain_agents`, or `langchain_graph` to accept only those `config.provider.name` values. Names must match the config exactly. `langchain_graph` forwards `providers` to the handler and not to `graph()`. Several scoped handlers can be registered together; the shortest list that contains the name wins. Omit `providers` to keep matching every provider. + ## Installation ```bash diff --git a/packages/langchain-agents/src/launchdarkly_ai_langchain_agents/graph.py b/packages/langchain-agents/src/launchdarkly_ai_langchain_agents/graph.py index 95dc0c11..41a5e625 100644 --- a/packages/langchain-agents/src/launchdarkly_ai_langchain_agents/graph.py +++ b/packages/langchain-agents/src/launchdarkly_ai_langchain_agents/graph.py @@ -14,5 +14,13 @@ def langchain_graph(key: str, llm: Any = None, **options: Any) -> Any: Equivalent to ``graph(key, handlers=[create_langchain_agents_handler(llm)], **options)``. Use the base ``graph()`` directly for multi-provider graphs. + + ``providers`` is forwarded to the handler, not to ``graph()``. It limits which + ``config.provider.name`` values the wildcard accepts. """ - return graph(key, handlers=[create_langchain_agents_handler(llm)], **options) + providers = options.pop("providers", None) + return graph( + key, + handlers=[create_langchain_agents_handler(llm, providers=providers)], + **options, + ) diff --git a/packages/langchain-agents/src/launchdarkly_ai_langchain_agents/handler.py b/packages/langchain-agents/src/launchdarkly_ai_langchain_agents/handler.py index bc72beb3..a45112aa 100644 --- a/packages/langchain-agents/src/launchdarkly_ai_langchain_agents/handler.py +++ b/packages/langchain-agents/src/launchdarkly_ai_langchain_agents/handler.py @@ -7,7 +7,7 @@ import asyncio import json -from collections.abc import AsyncGenerator +from collections.abc import AsyncGenerator, Sequence from typing import Any from launchdarkly_ai_server import ( @@ -248,7 +248,10 @@ def _run_usage_from_messages(messages: list[Any]) -> Any: def create_langchain_agents_handler( - llm: Any = None, *, capture_content: bool = False + llm: Any = None, + *, + providers: Sequence[str] | None = None, + capture_content: bool = False, ) -> ProviderHandler: """Creates a ``ProviderHandler`` for LangChain via ``create_react_agent``. @@ -258,6 +261,10 @@ def create_langchain_agents_handler( Set *capture_content* to put prompts, model output, tool arguments and tool results on the emitted spans. It defaults to off. Conversation content is PII, so a run emits only metadata, meaning models, token counts, timings and tool names, until a caller asks for more. + + *providers* limits which ``config.provider.name`` values this wildcard accepts. + Omit it to match every provider. Names must match the config exactly, for example + ``["Bedrock"]``. """ async def _call_impl( @@ -411,6 +418,7 @@ def _stream_impl( _call_impl, # type: ignore[arg-type] _stream_impl, # type: ignore[arg-type] capture_content=capture_content, + providers=providers, ) @@ -561,13 +569,15 @@ def langchain_agents( **kwargs: Any, ) -> Any: """Convenience wrapper: creates a handler and calls config(...).invoke().""" - # Both are lifted out of kwargs: capture_content configures the handler, variables belong to - # the invocation. Leaving either in would pass it to config(), which takes neither, so a caller - # asking for content on spans got a TypeError instead of content. + # Lifted out of kwargs: capture_content and providers configure the handler, variables belong + # to the invocation. Leaving them in would pass them to config(), which takes none of them. variables = kwargs.pop("variables", None) capture_content = kwargs.pop("capture_content", False) + providers = kwargs.pop("providers", None) return config( key=config_key, - handler=create_langchain_agents_handler(capture_content=capture_content), + handler=create_langchain_agents_handler( + providers=providers, capture_content=capture_content + ), **kwargs, ).invoke(user_input, context, variables=variables) diff --git a/packages/langchain-agents/tests/test_graph.py b/packages/langchain-agents/tests/test_graph.py index ce5e455a..4bd3acf6 100644 --- a/packages/langchain-agents/tests/test_graph.py +++ b/packages/langchain-agents/tests/test_graph.py @@ -40,6 +40,15 @@ def test_user_cannot_override_handlers(self) -> None: assert "handlers" in kw assert len(kw["handlers"]) == 1 + def test_providers_forwarded_to_the_handler(self) -> None: + with patch("launchdarkly_ai_langchain_agents.graph.graph") as mock_graph: + mock_graph.return_value = MagicMock() + langchain_graph("key", providers=["Bedrock"]) + kw = mock_graph.call_args[1] + assert "providers" not in kw + assert kw["handlers"][0].providers == ("Bedrock",) + assert kw["handlers"][0].provides_for == ("*", "agent") + def test_llm_option_forwarded(self) -> None: mock_llm = MagicMock() with patch("launchdarkly_ai_langchain_agents.graph.graph") as mock_graph: diff --git a/packages/langchain-agents/tests/test_handler.py b/packages/langchain-agents/tests/test_handler.py index 1861a1ed..ea587cda 100644 --- a/packages/langchain-agents/tests/test_handler.py +++ b/packages/langchain-agents/tests/test_handler.py @@ -177,6 +177,19 @@ def test_provides_for_values_are_correct(self) -> None: h = create_langchain_agents_handler() assert h.provides_for == ("*", "agent") + def test_providers_is_unset_by_default(self) -> None: + h = create_langchain_agents_handler() + assert h.providers is None + + def test_providers_scopes_the_wildcard(self) -> None: + h = create_langchain_agents_handler(providers=["Bedrock"]) + assert h.provides_for == ("*", "agent") + assert h.providers == ("Bedrock",) + + def test_empty_providers_is_rejected(self) -> None: + with pytest.raises(ValueError, match="providers"): + create_langchain_agents_handler(providers=[]) + def test_multiple_calls_return_independent_instances(self) -> None: h1 = create_langchain_agents_handler() h2 = create_langchain_agents_handler() @@ -1189,6 +1202,21 @@ def test_config_key_forwarded_as_key(self) -> None: "hello", ctx, variables=None ) + def test_providers_is_forwarded_to_the_handler(self) -> None: + mock_config_instance = MagicMock() + mock_config_fn = MagicMock(return_value=mock_config_instance) + mock_config_instance.invoke = MagicMock(return_value="result") + + with patch.object(handler_mod, "config", mock_config_fn): + from launchdarkly_ai_langchain_agents.handler import langchain_agents + + ctx = {"kind": "user", "key": "u1"} + langchain_agents("my-flag", "hello", ctx, providers=["Bedrock"]) + + call_kwargs = mock_config_fn.call_args.kwargs + assert "providers" not in call_kwargs + assert call_kwargs["handler"].providers == ("Bedrock",) + def test_callable_without_extra_kwargs(self) -> None: mock_config_instance = MagicMock() diff --git a/packages/langchain-messages/README.md b/packages/langchain-messages/README.md index abe5e2b1..724ef0e9 100644 --- a/packages/langchain-messages/README.md +++ b/packages/langchain-messages/README.md @@ -4,6 +4,8 @@ LangChain handler for `launchdarkly-ai-server` using **LangChain chat models** ( **`provides_for`:** `['*', 'messages']` — matches any flag variation where `meta.mode` is `"messages"` and no more-specific handler is registered. LangChain is a framework adapter, not a provider: it routes through `langchain-anthropic`, `langchain-openai`, and others at runtime based on `config.provider.name`. Use `'*'` so that flags configured with `provider.name = "Anthropic"` or `"OpenAI"` are automatically handled without requiring a separate native handler. +Pass `providers=["Bedrock"]` to `create_langchain_messages_handler` (and to `langchain_messages`) to accept only those `config.provider.name` values. Names must match the config exactly. Several scoped handlers can be registered together; the shortest list that contains the name wins. Omit `providers` to keep matching every provider. + ## Installation ```bash diff --git a/packages/langchain-messages/src/launchdarkly_ai_langchain_messages/handler.py b/packages/langchain-messages/src/launchdarkly_ai_langchain_messages/handler.py index b30e8b80..947eae2c 100644 --- a/packages/langchain-messages/src/launchdarkly_ai_langchain_messages/handler.py +++ b/packages/langchain-messages/src/launchdarkly_ai_langchain_messages/handler.py @@ -2,7 +2,7 @@ import asyncio import json -from collections.abc import AsyncGenerator +from collections.abc import AsyncGenerator, Sequence from types import SimpleNamespace from typing import Any @@ -370,7 +370,10 @@ async def _run_structured_turn( def create_langchain_messages_handler( - llm: Any = None, *, capture_content: bool = False + llm: Any = None, + *, + providers: Sequence[str] | None = None, + capture_content: bool = False, ) -> ProviderHandler: """ Creates a ``ProviderHandler`` for LangChain (chat models). @@ -383,6 +386,10 @@ def create_langchain_messages_handler( Set *capture_content* to put prompts, model output, tool arguments and tool results on the emitted spans. It defaults to off. Conversation content is PII, so a run emits only metadata, meaning models, token counts, timings and tool names, until a caller asks for more. + + *providers* limits which ``config.provider.name`` values this wildcard accepts. + Omit it to match every provider. Names must match the config exactly, for example + ``["Bedrock"]``. """ async def _call_impl( @@ -699,6 +706,7 @@ def _stream_impl( _call_impl, # type: ignore[arg-type] _stream_impl, # type: ignore[arg-type] capture_content=capture_content, + providers=providers, ) @@ -1010,15 +1018,15 @@ def langchain_messages( **kwargs: Any, ) -> Any: """Convenience wrapper: creates a handler and calls config(...).invoke().""" - # Both are lifted out of kwargs: capture_content configures the handler, variables belong to - # the invocation. Leaving either in would pass it to config(), which takes neither, so a caller - # asking for content on spans got a TypeError instead of content. + # Lifted out of kwargs: capture_content and providers configure the handler, variables belong + # to the invocation. Leaving them in would pass them to config(), which takes none of them. variables = kwargs.pop("variables", None) capture_content = kwargs.pop("capture_content", False) + providers = kwargs.pop("providers", None) return config( key=config_key, handler=create_langchain_messages_handler( - llm=llm, capture_content=capture_content + llm=llm, providers=providers, capture_content=capture_content ), **kwargs, ).invoke(user_input, context, variables=variables) diff --git a/packages/langchain-messages/tests/test_handler.py b/packages/langchain-messages/tests/test_handler.py index d00e43fe..4a19c947 100644 --- a/packages/langchain-messages/tests/test_handler.py +++ b/packages/langchain-messages/tests/test_handler.py @@ -186,6 +186,27 @@ def test_provides_for_is_the_wildcard_provider(self) -> None: h = create_langchain_messages_handler(llm=_make_llm()) assert h.provides_for == ("*", "messages") + def test_providers_is_unset_by_default(self) -> None: + from launchdarkly_ai_langchain_messages import create_langchain_messages_handler + + h = create_langchain_messages_handler(llm=_make_llm()) + assert h.providers is None + + def test_providers_scopes_the_wildcard(self) -> None: + from launchdarkly_ai_langchain_messages import create_langchain_messages_handler + + h = create_langchain_messages_handler( + llm=_make_llm(), providers=["Bedrock", "Anthropic"] + ) + assert h.provides_for == ("*", "messages") + assert h.providers == ("Bedrock", "Anthropic") + + def test_empty_providers_is_rejected(self) -> None: + from launchdarkly_ai_langchain_messages import create_langchain_messages_handler + + with pytest.raises(ValueError, match="providers"): + create_langchain_messages_handler(llm=_make_llm(), providers=[]) + def test_multiple_calls_return_independent_instances(self) -> None: from launchdarkly_ai_langchain_messages import create_langchain_messages_handler @@ -1125,6 +1146,24 @@ def test_calls_through_to_model_call(self) -> None: "hello", ctx, variables=None ) + def test_providers_is_forwarded_to_the_handler(self) -> None: + import launchdarkly_ai_langchain_messages.handler as handler_mod + + mock_config_instance = MagicMock() + mock_config_fn = MagicMock(return_value=mock_config_instance) + mock_config_instance.invoke = MagicMock(return_value="result") + + with patch.object(handler_mod, "config", mock_config_fn): + from launchdarkly_ai_langchain_messages.handler import langchain_messages + + ctx = {"kind": "user", "key": "u1"} + langchain_messages("my-flag", "hello", ctx, providers=["Bedrock"]) + + call_kwargs = mock_config_fn.call_args.kwargs + assert "providers" not in call_kwargs + assert call_kwargs["handler"].providers == ("Bedrock",) + assert call_kwargs["handler"].provides_for == ("*", "messages") + def test_callable_without_extra_kwargs(self) -> None: import launchdarkly_ai_langchain_messages.handler as handler_mod