diff --git a/packages/client/src/launchdarkly_ai_server/graph.py b/packages/client/src/launchdarkly_ai_server/graph.py index 46be8969..8b67f6d8 100644 --- a/packages/client/src/launchdarkly_ai_server/graph.py +++ b/packages/client/src/launchdarkly_ai_server/graph.py @@ -2,7 +2,6 @@ import asyncio import inspect as _inspect -import json import logging import re as _re import time @@ -39,11 +38,10 @@ logger = logging.getLogger(__name__) MAX_TRAVERSAL_DEPTH = 100 -MAX_GRAPH_CACHE_SIZE = 512 -# graph() caches a built graph per context, track data included. Each run passes its own copy +# The built graph's track data is minted before the run's runId. Each run passes its own copy # with a fresh runId through route/stream_route opts under this key, so the handoff events -# emitted inside the cached build carry the run's id instead of the build's. +# emitted inside the build carry the run's id instead of the build's. _RUN_TRACK_DATA = "_run_track_data" @@ -878,9 +876,6 @@ def __init__( ) -> None: self._key = key self._options = options - self._cache: dict[ - str, tuple[GraphDefinition, TrackData, Callable[..., Any]] - ] = {} async def invoke( self, @@ -927,19 +922,7 @@ async def _invoke( "use resolve_graph() with a framework-native runner." ) - try: - cache_key: str | None = json.dumps(context, sort_keys=True) - except (TypeError, ValueError): - cache_key = None - if cache_key is not None and cache_key in self._cache: - built = self._cache[cache_key] - else: - built = await _build_graph(self._key, context, resolved_options) - if cache_key is not None: - if len(self._cache) >= MAX_GRAPH_CACHE_SIZE: - # Evict an arbitrary entry to keep the cache bounded. - self._cache.pop(next(iter(self._cache))) - self._cache[cache_key] = built + built = await _build_graph(self._key, context, resolved_options) graph_def, built_track_data, _ = built if not graph_def.enabled: @@ -1165,18 +1148,7 @@ async def _stream_events( "use resolve_graph() with a framework-native runner." ) - try: - cache_key: str | None = json.dumps(context, sort_keys=True) - except (TypeError, ValueError): - cache_key = None - if cache_key is not None and cache_key in self._cache: - built = self._cache[cache_key] - else: - built = await _build_graph(self._key, context, resolved_options) - if cache_key is not None: - if len(self._cache) >= MAX_GRAPH_CACHE_SIZE: - self._cache.pop(next(iter(self._cache))) - self._cache[cache_key] = built + built = await _build_graph(self._key, context, resolved_options) graph_def, built_track_data, stream_route = built if not graph_def.enabled: diff --git a/packages/client/tests/test_graph.py b/packages/client/tests/test_graph.py index c3379682..e1653ee2 100644 --- a/packages/client/tests/test_graph.py +++ b/packages/client/tests/test_graph.py @@ -324,14 +324,80 @@ async def test_graph_node_tracked_per_visited_node( assert node_events[0][2]["graphKey"] == "graph-key" assert "path" not in node_events[0][2] - async def test_node_variations_resolved_only_once( + async def test_node_variations_resolved_on_every_call( self, mock_ld_client: MagicMock ) -> None: + """Each invoke() re-resolves the graph and node flags. See TESTING.md §3.11.""" g = graph("graph-key", handlers=[_make_handler()]) await g.invoke("hi", CONTEXT) + after_first = mock_ld_client.variation.call_count await g.invoke("hi again", CONTEXT) - # Second call should use the cache — variation calls should not double - assert mock_ld_client.variation.call_count <= 4 # graph + 2 nodes × 1 cache hit + assert mock_ld_client.variation.call_count == after_first * 2 + + async def test_flag_change_reaches_reused_graph( + self, mock_ld_client: MagicMock + ) -> None: + """A node flag change between calls is seen on the next call. See TESTING.md §3.11.""" + state = {"instructions": "v1", "graph_enabled": True} + + async def variation(key: str, ctx: dict, default: Any) -> Any: + if key == "graph-key": + return ( + {"root": "only-node", "edges": {}} if state["graph_enabled"] else {} + ) + return { + "model": {"name": "gpt-4"}, + "provider": {"name": "TestProvider"}, + "instructions": state["instructions"], + "_ldMeta": {"enabled": True, "variationKey": "v1", "version": 1}, + } + + mock_ld_client.variation = AsyncMock(side_effect=variation) + seen: list[str] = [] + + async def fn( + config, user_input, tool_handlers, variables, history=None + ) -> dict: # type: ignore[no-untyped-def] + seen.append(config["instructions"]) + return {"output": "ok", "usage": {"input_tokens": 1, "output_tokens": 1}} + + g = graph( + "graph-key", + handlers=[ + ProviderHandler(fn=fn, provides_for=("TestProvider", "messages")) + ], # type: ignore[arg-type] + ) + await g.invoke("first", CONTEXT) + state["instructions"] = "v2" + await g.invoke("second", CONTEXT) + assert seen == ["v1", "v2"] + + state["graph_enabled"] = False + with pytest.raises(ValueError, match="disabled"): + await g.invoke("third", CONTEXT) + with pytest.raises(ValueError, match="disabled"): + async for _ in g.stream("fourth", CONTEXT): + pass + + async def test_failed_resolution_is_not_reused( + self, mock_ld_client: MagicMock + ) -> None: + """A transient variation failure on one call does not stick to the next.""" + original = mock_ld_client.variation.side_effect + failures = {"left": 1} + + async def flaky(key: str, ctx: dict, default: Any) -> Any: + if key != "graph-key" and failures["left"]: + failures["left"] -= 1 + raise RuntimeError("transient") + return await original(key, ctx, default) + + mock_ld_client.variation = AsyncMock(side_effect=flaky) + g = graph("graph-key", handlers=[_make_handler()]) + with pytest.raises(ValueError, match="disabled"): + await g.invoke("first", CONTEXT) + result = await g.invoke("second", CONTEXT) + assert result.response == "ok" async def test_handoff_success_tracked_on_each_handoff( self, mock_ld_client: MagicMock @@ -464,35 +530,6 @@ async def test_disabled_graph_reverse_traverse_key_is_snake_case( assert hasattr(gd, "reverse_traverse") assert not hasattr(gd, "reverseTraverse") - async def test_cache_is_bounded(self, mock_ld_client: MagicMock) -> None: - """GraphInstance._cache must not grow beyond MAX_GRAPH_CACHE_SIZE. - See TESTING.md §3.11.""" - from launchdarkly_ai_server.graph import MAX_GRAPH_CACHE_SIZE - - g = graph("graph-key", handlers=[_make_handler()]) - # Invoke with more distinct contexts than the max - limit = MAX_GRAPH_CACHE_SIZE + 10 - for i in range(limit): - await g.invoke("hi", {"kind": "user", "key": f"u{i}"}) - assert len(g._cache) <= MAX_GRAPH_CACHE_SIZE - - async def test_equal_content_contexts_share_cache_entry( - self, mock_ld_client: MagicMock - ) -> None: - """Two distinct dict objects with identical JSON content must share one - cache entry. The cache key must be json.dumps(context, sort_keys=True), - not id(context). See TESTING.md §3.11.""" - g = graph("graph-key", handlers=[_make_handler()]) - ctx_a = {"kind": "user", "key": "u1"} - ctx_b = {"kind": "user", "key": "u1"} # distinct object, same content - assert ctx_a is not ctx_b, "test requires two different dict objects" - await g.invoke("first call", ctx_a) - await g.invoke("second call", ctx_b) - assert len(g._cache) == 1, ( - "Equal-content contexts must share a single cache entry — " - "the cache key must be json.dumps(context), not id(context)." - ) - async def test_error_logging_when_node_variation_fails( self, mock_ld_client: MagicMock ) -> None: