diff --git a/examples/graph_streaming.py b/examples/graph_streaming.py index 10f1c6cc..8d36f36a 100644 --- a/examples/graph_streaming.py +++ b/examples/graph_streaming.py @@ -42,8 +42,14 @@ async def run(key: str, user_input: str) -> None: elif event["type"] == "node_start": print(f"\n[node_start] {event['nodeKey']}", file=sys.stderr) elif event["type"] == "node_done": + # A node's own judge results, when one of its judges produced a score. + judged = ( + f" judgeResults={json.dumps(event['judgeResults'], default=str)}" + if event.get("judgeResults") + else "" + ) print( - f"\n[node_done] {event['nodeKey']} usage={json.dumps(event.get('usage'))}", + f"\n[node_done] {event['nodeKey']} usage={json.dumps(event.get('usage'))}{judged}", file=sys.stderr, ) elif event["type"] == "handoff": diff --git a/packages/client/src/launchdarkly_ai_server/graph.py b/packages/client/src/launchdarkly_ai_server/graph.py index 4f8677b6..8b3f0f20 100644 --- a/packages/client/src/launchdarkly_ai_server/graph.py +++ b/packages/client/src/launchdarkly_ai_server/graph.py @@ -577,12 +577,16 @@ async def stream_node( 1, ) - yield { + # The node's own judges; omitted when none produced a result, as on done. + node_done: GraphStreamEvent = { "type": "node_done", "nodeKey": node.key, "response": response, "usage": usage, } + if judge_results: + node_done["judgeResults"] = judge_results + yield node_done if outcome is not None: outcome.clear() outcome.update( @@ -721,12 +725,16 @@ async def stream_route( 1, ) - yield { + # The node's own judges; omitted when none produced a result, as on done. + node_done: GraphStreamEvent = { "type": "node_done", "nodeKey": node.key, "response": response, "usage": usage, } + if judge_results: + node_done["judgeResults"] = judge_results + yield node_done if outcome is not None: outcome.clear() outcome.update( diff --git a/packages/client/src/launchdarkly_ai_server/types.py b/packages/client/src/launchdarkly_ai_server/types.py index c3f17341..3143ba64 100644 --- a/packages/client/src/launchdarkly_ai_server/types.py +++ b/packages/client/src/launchdarkly_ai_server/types.py @@ -389,7 +389,7 @@ class JudgeRunResult: # Graph stream events (camelCase public fields, matching the TypeScript SDK) # node_start: {type, nodeKey} # chunk: {type, text, nodeKey} -# node_done: {type, nodeKey, response, usage} +# node_done: {type, nodeKey, response, usage, judgeResults?} (the node's own judges) # handoff: {type, sourceKey, targetKey} # done: {type, response, usage, judgeResults?} GraphStreamEvent = dict[str, Any] diff --git a/packages/client/tests/test_graph_stream.py b/packages/client/tests/test_graph_stream.py index f3f6d624..a38ad0a6 100644 --- a/packages/client/tests/test_graph_stream.py +++ b/packages/client/tests/test_graph_stream.py @@ -493,6 +493,82 @@ async def test_omits_judge_results_when_empty( ) assert "judgeResults" not in events[-1] + async def test_node_done_carries_the_nodes_own_judge_results( + self, mock_ld_client: MagicMock + ) -> None: + leaf_judge = { + "leaf-judge": { + "usage": {"input": 1, "output": 1, "total": 2}, + "response": "good", + "score": 0.9, + } + } + + async def judges_for(**kwargs: Any) -> dict[str, Any]: + return leaf_judge if kwargs["config"]["instructions"] == "I am leaf" else {} + + with patch( + "launchdarkly_ai_server.judges.run_judges", + new_callable=AsyncMock, + side_effect=judges_for, + ): + events = await _collect( + graph("graph-key", handlers=[_make_streaming_handler(["ok"])]).stream( + "hi", CONTEXT + ) + ) + dones = [e for e in events if e["type"] == "node_done"] + assert [d["nodeKey"] for d in dones] == ["root-node", "leaf-node"] + assert "judgeResults" not in dones[0] + assert dones[1]["judgeResults"] == leaf_judge + # done holds only the graph-level judge; per-node results are not merged into it. + assert "judgeResults" not in events[-1] + + async def test_node_judges_stay_on_node_done_when_a_graph_judge_runs( + self, mock_ld_client: MagicMock + ) -> None: + leaf_judge = { + "leaf-judge": { + "usage": {"input": 1, "output": 1, "total": 2}, + "response": "good", + "score": 0.9, + } + } + graph_judge = { + "graph-judge": { + "usage": {"input": 1, "output": 1, "total": 2}, + "response": "whole", + "score": 0.5, + } + } + + async def judges_for(**kwargs: Any) -> dict[str, Any]: + config = kwargs["config"] + judges = (config.get("judgeConfiguration") or {}).get("judges") or [] + if judges and judges[0]["key"] == "graph-judge": + return graph_judge + return leaf_judge if config.get("instructions") == "I am leaf" else {} + + with patch( + "launchdarkly_ai_server.judges.run_judges", + new_callable=AsyncMock, + side_effect=judges_for, + ): + events = await _collect( + graph( + "graph-key", + handlers=[_make_streaming_handler(["ok"])], + graph_judge="graph-judge", + ).stream("hi", CONTEXT) + ) + leaf_done = next( + e + for e in events + if e["type"] == "node_done" and e["nodeKey"] == "leaf-node" + ) + assert leaf_done["judgeResults"] == leaf_judge + assert events[-1]["judgeResults"] == graph_judge + async def test_per_node_judge_receives_the_trajectory( self, mock_ld_client: MagicMock ) -> None: @@ -652,6 +728,34 @@ async def test_judges_receive_original_config( assert judged["instructions"] == "I am root" assert not any(k.startswith("__handoff_") for k in (judged.get("tools") or {})) + async def test_routing_node_judge_results_on_its_node_done( + self, mock_ld_client: MagicMock + ) -> None: + root_judge = { + "root-judge": { + "usage": {"input": 1, "output": 1, "total": 2}, + "response": "routed", + "score": 1.0, + } + } + + async def judges_for(**kwargs: Any) -> dict[str, Any]: + return root_judge if kwargs["config"]["instructions"] == "I am root" else {} + + h = _make_branch_picking_stream_handler("agent-b") + with patch( + "launchdarkly_ai_server.judges.run_judges", + new_callable=AsyncMock, + side_effect=judges_for, + ): + events = await _collect( + graph("graph-key", handlers=[h]).stream("hi", CONTEXT) + ) + dones = [e for e in events if e["type"] == "node_done"] + assert [d["nodeKey"] for d in dones] == ["root-node", "agent-b"] + assert dones[0]["judgeResults"] == root_judge + assert "judgeResults" not in dones[1] + async def test_routed_judge_receives_the_trajectory( self, mock_ld_client: MagicMock ) -> None: