feat(AIC-3210): support graph().stream() in the Python AI SDK

examples/graph_streaming.py

@@ -0,0 +1,63 @@
+"""
+Example: ``graph().stream()`` — the streaming counterpart to ``examples/graph_example.py``.
+
+Model text goes to stdout; node boundaries go to stderr so stdout stays a clean transcript.
+
+Like ``examples/streaming.py``, the generator is built inside ``conversation_id`` and iterated
+*outside* it. That is the shape a server produces when it hands a stream to a transport, and it
+is what exercises call-time binding: an async generator body does not run until the first
+``__anext__``, so both the conversation id and the ``ld.ai.graph`` span's OTel parent have to be
+captured when ``stream()`` is called, not when iteration starts.
+
+Writes no JSON output file, same as the single-config streaming example.
+
+Usage (via main.py):
+    python main.py graph-streaming <flag-key> "<user input>"
+"""
+
+from __future__ import annotations
+
+import json
+import sys
+
+import examples.register  # noqa: F401 – side-effect: populate global_registry
+from examples.utils import new_context, new_conversation_id
+from launchdarkly_ai_server import conversation_id, global_registry, graph
+
+
+async def run(key: str, user_input: str) -> None:
+    conversation = new_conversation_id("graph-streaming-example")
+    print(f"[conversation] {conversation}", file=sys.stderr)
+
+    with conversation_id(conversation):
+        stream = graph(
+            key,
+            registry=global_registry,
+        ).stream(user_input, new_context(), {"user_id": "user-123"})
+
+    async for event in stream:
+        if event["type"] == "chunk":
+            sys.stdout.write(event.get("text", ""))
+            sys.stdout.flush()
+        elif event["type"] == "node_start":
+            print(f"\n[node_start] {event['nodeKey']}", file=sys.stderr)
+        elif event["type"] == "node_done":
+            print(
+                f"\n[node_done] {event['nodeKey']} usage={json.dumps(event.get('usage'))}",
+                file=sys.stderr,
+            )
+        elif event["type"] == "handoff":
+            print(
+                f"[handoff] {event['sourceKey']} -> {event['targetKey']}",
+                file=sys.stderr,
+            )
+        else:
+            # Final event — usage aggregated across nodes, plus graph judge results when configured.
+            sys.stdout.write("\n\n")
+            print("Usage:", json.dumps(event.get("usage"), indent=2, default=str))
+            if event.get("judgeResults"):
+                print(
+                    "Judge results:",
+                    json.dumps(event["judgeResults"], indent=2, default=str),
+                )
+            sys.stdout.write("\n")

main.py

@@ -10,6 +10,7 @@
     python main.py judge                    launch-darkly-documentation-summarizer "What is the LaunchDarkly AI SDK?"
     python main.py graph                    my-agent-graph "What is the LaunchDarkly AI SDK?"
     python main.py graph-history            my-agent-graph ""
+    python main.py graph-streaming          travel-agent-flow "I was double charged for my flight"
     python main.py openai-only              my-openai-flag "Tell me about feature flags"
     python main.py langchain                my-langchain-flag "Tell me about feature flags"
     python main.py claude-agents            launch-darkly-documentation-summarizer "What is the LaunchDarkly AI SDK?"
@@ -46,6 +47,7 @@
     "streaming": "examples.streaming",
     "graph": "examples.graph_example",
     "graph-history": "examples.graph_history",
+    "graph-streaming": "examples.graph_streaming",
     "conversation": "examples.conversation",
     "history": "examples.history",
     "judge": "examples.judge_example",

packages/client/src/launchdarkly_ai_server/__init__.py

@@ -73,6 +73,7 @@
     GraphEdge,
     GraphNode,
     GraphOptions,
+    GraphStreamEvent,
     GraphTopology,
     HandlerResult,
     HandlerStreamEvent,
@@ -136,6 +137,7 @@
     "GraphEdge",
     "GraphNode",
     "GraphOptions",
+    "GraphStreamEvent",
     "GraphTopology",
     "HandlerResult",
     "HandlerStreamEvent",

packages/client/src/launchdarkly_ai_server/conversation.py

@@ -200,6 +200,34 @@ def bind_conversation_id(
     return _stream_with_bound_id(generator, conversation)
 
 
+async def bind_span_context(
+    generator: AsyncGenerator[Any, None],
+    ctx: otel_context.Context,
+) -> AsyncGenerator[Any, None]:
+    """Re-enter ``ctx`` around every step of ``generator``.
+
+    Sibling of :func:`bind_conversation_id`, which deliberately carries only the conversation id
+    and leaves span parenting alone. A generator body suspends at each ``yield``, so the context
+    has to be re-applied on every ``__anext__`` — wrapping the body once is not enough.
+    """
+    try:
+        while True:
+            token = otel_context.attach(ctx)
+            try:
+                item = await generator.__anext__()
+            except StopAsyncIteration:
+                return
+            finally:
+                otel_context.detach(token)
+            yield item
+    finally:
+        token = otel_context.attach(ctx)
+        try:
+            await generator.aclose()
+        finally:
+            otel_context.detach(token)
+
+
 @asynccontextmanager
 async def with_judge_evaluation(name: str) -> AsyncIterator[RecordEvaluation]:
     """Hold the judge ``invoke_agent`` span open until ``record`` runs.