chore: import upstream snapshot with attribution

python/samples/02-agents/response_stream.py

@@ -0,0 +1,358 @@
+# Copyright (c) Microsoft. All rights reserved.
+
+import asyncio
+from collections.abc import AsyncIterable, Sequence
+
+from agent_framework import ChatResponse, ChatResponseUpdate, Content, Message, ResponseStream
+
+"""ResponseStream: A Deep Dive
+
+This sample explores the ResponseStream class - a powerful abstraction for working with
+streaming responses in the Agent Framework.
+
+=== Why ResponseStream Exists ===
+
+When working with AI models, responses can be delivered in two ways:
+1. **Non-streaming**: Wait for the complete response, then return it all at once
+2. **Streaming**: Receive incremental updates as they're generated
+
+Streaming provides a better user experience (faster time-to-first-token, progressive rendering)
+but introduces complexity:
+- How do you process updates as they arrive?
+- How do you also get a final, complete response?
+- How do you ensure the underlying stream is only consumed once?
+- How do you add custom logic (hooks) at different stages?
+
+ResponseStream solves all these problems by wrapping an async iterable and providing:
+- Multiple consumption patterns (iteration OR direct finalization)
+- Hook points for transformation, cleanup, finalization, and result processing
+- The `wrap()` API to layer behavior without double-consuming the stream
+
+=== The Four Hook Types ===
+
+ResponseStream provides four ways to inject custom logic. All can be passed via constructor
+or added later via fluent methods:
+
+1. **Transform Hooks** (`transform_hooks=[]` or `.with_transform_hook()`)
+   - Called for EACH update as it's yielded during iteration
+   - Can transform updates before they're returned to the consumer
+   - Multiple hooks are called in order, each receiving the previous hook's output
+   - Only triggered during iteration (not when calling get_final_response directly)
+
+2. **Cleanup Hooks** (`cleanup_hooks=[]` or `.with_cleanup_hook()`)
+   - Called ONCE when iteration completes (stream fully consumed), BEFORE finalizer
+   - Used for cleanup: closing connections, releasing resources, logging
+   - Cannot modify the stream or response
+   - Triggered regardless of how the stream ends (normal completion or exception)
+
+3. **Finalizer** (`finalizer=` constructor parameter)
+   - Called ONCE when `get_final_response()` is invoked
+   - Receives the list of collected updates and converts to the final type
+   - There is only ONE finalizer per stream (set at construction)
+
+4. **Result Hooks** (`result_hooks=[]` or `.with_result_hook()`)
+   - Called ONCE after the finalizer produces its result
+   - Transform the final response before returning
+   - Multiple result hooks are called in order, each receiving the previous result
+   - Can return None to keep the previous value unchanged
+
+=== Two Consumption Patterns ===
+
+**Pattern 1: Async Iteration**
+```python
+async for update in response_stream:
+    print(update.text)  # Process each update
+# Stream is now consumed; updates are stored internally
+```
+- Transform hooks are called for each yielded item
+- Cleanup hooks are called after the last item
+- The stream collects all updates internally for later finalization
+- Does not run the finalizer automatically
+
+**Pattern 2: Direct Finalization**
+```python
+final = await response_stream.get_final_response()
+```
+- If the stream hasn't been iterated, it auto-iterates (consuming all updates)
+- The finalizer converts collected updates to a final response
+- Result hooks transform the response
+- You get the complete response without ever seeing individual updates
+
+** Pattern 3: Combined Usage **
+
+When you first iterate the stream and then call `get_final_response()`, the following occurs:
+- Iteration yields updates with transform hooks applied
+- Cleanup hooks run after iteration completes
+- Calling `get_final_response()` uses the already collected updates to produce the final response
+- Note that it does not re-iterate the stream since it's already been consumed
+
+```python
+async for update in response_stream:
+    print(update.text)  # See each update
+final = await response_stream.get_final_response()  # Get the aggregated result
+```
+
+=== Chaining with .map() and .with_finalizer() ===
+
+When building a Agent on top of a ChatClient, we face a challenge:
+- The ChatClient returns a ResponseStream[ChatResponseUpdate, ChatResponse]
+- The Agent needs to return a ResponseStream[AgentResponseUpdate, AgentResponse]
+- We can't iterate the ChatClient's stream twice!
+
+The `.map()` and `.with_finalizer()` methods solve this by creating new ResponseStreams that:
+- Delegate iteration to the inner stream (only consuming it once)
+- Maintain their OWN separate transform hooks, result hooks, and cleanup hooks
+- Allow type-safe transformation of updates and final responses
+
+**`.map(transform)`**: Creates a new stream that transforms each update.
+- Returns a new ResponseStream with the transformed update type
+- Falls back to the inner stream's finalizer if no new finalizer is set
+
+**`.with_finalizer(finalizer)`**: Creates a new stream with a different finalizer.
+- Returns a new ResponseStream with the new final type
+- The inner stream's finalizer and result_hooks ARE still called (see below)
+
+**IMPORTANT**: When chaining these methods via `get_final_response()`:
+1. The inner stream's finalizer runs first (on the original updates)
+2. The inner stream's result_hooks run (on the inner final result)
+3. The outer stream's finalizer runs (on the transformed updates)
+4. The outer stream's result_hooks run (on the outer final result)
+
+This ensures that post-processing hooks registered on the inner stream (e.g., context
+provider notifications, telemetry, thread updates) are still executed even when the
+stream is wrapped/mapped.
+
+```python
+# Agent does something like this internally:
+chat_stream = client.get_response(messages, stream=True)
+agent_stream = (
+    chat_stream
+    .map(_to_agent_update, _to_agent_response)
+    .with_result_hook(_notify_thread)  # Outer hook runs AFTER inner hooks
+)
+```
+
+This ensures:
+- The underlying ChatClient stream is only consumed once
+- The agent can add its own transform hooks, result hooks, and cleanup logic
+- Each layer (ChatClient, Agent, middleware) can add independent behavior
+- Inner stream post-processing (like context provider notification) still runs
+- Types flow naturally through the chain
+"""
+
+
+async def main() -> None:
+    """Demonstrate the various ResponseStream patterns and capabilities."""
+
+    # =========================================================================
+    # Example 1: Basic ResponseStream with iteration
+    # =========================================================================
+    print("=== Example 1: Basic Iteration ===\n")
+
+    async def generate_updates() -> AsyncIterable[ChatResponseUpdate]:
+        """Simulate a streaming response from an AI model."""
+        words = ["Hello", " ", "from", " ", "the", " ", "streaming", " ", "response", "!"]
+        for word in words:
+            await asyncio.sleep(0.05)  # Simulate network delay
+            yield ChatResponseUpdate(contents=[Content.from_text(word)], role="assistant")
+
+    def combine_updates(updates: Sequence[ChatResponseUpdate]) -> ChatResponse:
+        """Finalizer that combines all updates into a single response."""
+        return ChatResponse.from_updates(updates)
+
+    stream = ResponseStream(generate_updates(), finalizer=combine_updates)
+
+    print("Iterating through updates:")
+    async for update in stream:
+        print(f"  Update: '{update.text}'")
+
+    # After iteration, we can still get the final response
+    final = await stream.get_final_response()
+    print(f"\nFinal response: '{final.text}'")
+
+    # =========================================================================
+    # Example 2: Using get_final_response() without iteration
+    # =========================================================================
+    print("\n=== Example 2: Direct Finalization (No Iteration) ===\n")
+
+    # Create a fresh stream (streams can only be consumed once)
+    stream2 = ResponseStream(generate_updates(), finalizer=combine_updates)
+
+    # Skip iteration entirely - get_final_response() auto-consumes the stream
+    final2 = await stream2.get_final_response()
+    print(f"Got final response directly: '{final2.text}'")
+    print(f"Number of updates collected internally: {len(stream2.updates)}")
+
+    # =========================================================================
+    # Example 3: Transform hooks - transform updates during iteration
+    # =========================================================================
+    print("\n=== Example 3: Transform Hooks ===\n")
+
+    update_count = {"value": 0}
+
+    def counting_hook(update: ChatResponseUpdate) -> ChatResponseUpdate:
+        """Hook that counts and annotates each update."""
+        update_count["value"] += 1
+        # Return the update (or a modified version)
+        return update
+
+    def uppercase_hook(update: ChatResponseUpdate) -> ChatResponseUpdate:
+        """Hook that converts text to uppercase."""
+        if update.text:
+            return ChatResponseUpdate(
+                contents=[Content.from_text(update.text.upper())], role=None, response_id=update.response_id
+            )
+        return update
+
+    # Pass transform_hooks directly to constructor
+    stream3: ResponseStream[ChatResponseUpdate, ChatResponse] = ResponseStream(
+        generate_updates(),
+        finalizer=combine_updates,
+        transform_hooks=[counting_hook, uppercase_hook],  # First counts, then uppercases
+    )
+
+    print("Iterating with hooks applied:")
+    async for update in stream3:
+        print(f"  Received: '{update.text}'")  # Will be uppercase
+
+    print(f"\nTotal updates processed: {update_count['value']}")
+
+    # =========================================================================
+    # Example 4: Cleanup hooks - cleanup after stream consumption
+    # =========================================================================
+    print("\n=== Example 4: Cleanup Hooks ===\n")
+
+    cleanup_performed = {"value": False}
+
+    async def cleanup_hook() -> None:
+        """Cleanup hook for releasing resources after stream consumption."""
+        print("  [Cleanup] Cleaning up resources...")
+        cleanup_performed["value"] = True
+
+    # Pass cleanup_hooks directly to constructor
+    stream4 = ResponseStream(
+        generate_updates(),
+        finalizer=combine_updates,
+        cleanup_hooks=[cleanup_hook],
+    )
+
+    print("Starting iteration (cleanup happens after):")
+    async for _update in stream4:
+        pass  # Just consume the stream
+    print(f"Cleanup was performed: {cleanup_performed['value']}")
+
+    # =========================================================================
+    # Example 5: Result hooks - transform the final response
+    # =========================================================================
+    print("\n=== Example 5: Result Hooks ===\n")
+
+    def add_metadata_hook(response: ChatResponse) -> ChatResponse:
+        """Result hook that adds metadata to the response."""
+        response.additional_properties["processed"] = True
+        response.additional_properties["word_count"] = len((response.text or "").split())
+        return response
+
+    def wrap_in_quotes_hook(response: ChatResponse) -> ChatResponse:
+        """Result hook that wraps the response text in quotes."""
+        if response.text:
+            return ChatResponse(
+                messages=[Message(contents=[f'"{response.text}"'], role="assistant")],
+                additional_properties=response.additional_properties,
+            )
+        return response
+
+    # Finalizer converts updates to response, then result hooks transform it
+    stream5: ResponseStream[ChatResponseUpdate, ChatResponse] = ResponseStream(
+        generate_updates(),
+        finalizer=combine_updates,
+        result_hooks=[add_metadata_hook, wrap_in_quotes_hook],  # First adds metadata, then wraps in quotes
+    )
+
+    final5 = await stream5.get_final_response()
+    print(f"Final text: {final5.text}")
+    print(f"Metadata: {final5.additional_properties}")
+
+    # =========================================================================
+    # Example 6: The wrap() API - layering without double-consumption
+    # =========================================================================
+    print("\n=== Example 6: wrap() API for Layering ===\n")
+
+    # Simulate what ChatClient returns
+    inner_stream = ResponseStream(generate_updates(), finalizer=combine_updates)
+
+    # Simulate what Agent does: wrap the inner stream
+    def to_agent_format(update: ChatResponseUpdate) -> ChatResponseUpdate:
+        """Map ChatResponseUpdate to agent format (simulated transformation)."""
+        # In real code, this would convert to AgentResponseUpdate
+        return ChatResponseUpdate(
+            contents=[Content.from_text(f"[AGENT] {update.text}")], role=None, response_id=update.response_id
+        )
+
+    def to_agent_response(updates: Sequence[ChatResponseUpdate]) -> ChatResponse:
+        """Finalizer that converts updates to agent response (simulated)."""
+        # In real code, this would create an AgentResponse
+        text = "".join(u.text or "" for u in updates)
+        return ChatResponse(
+            messages=[Message(contents=[f"[AGENT FINAL] {text}"], role="assistant")],
+            additional_properties={"layer": "agent"},
+        )
+
+    # .map() creates a new stream that:
+    # 1. Delegates iteration to inner_stream (only consuming it once)
+    # 2. Transforms each update via the transform function
+    # 3. Uses the provided finalizer (required since update type may change)
+    outer_stream = inner_stream.map(to_agent_format, to_agent_response)
+
+    print("Iterating the mapped stream:")
+    async for update in outer_stream:
+        print(f"  {update.text}")
+
+    final_outer = await outer_stream.get_final_response()
+    print(f"\nMapped final: {final_outer.text}")
+    print(f"Mapped metadata: {final_outer.additional_properties}")
+
+    # Important: the inner stream was only consumed once!
+    print(f"Inner stream consumed: {inner_stream._consumed}")
+
+    # =========================================================================
+    # Example 7: Combining all patterns
+    # =========================================================================
+    print("\n=== Example 7: Full Integration ===\n")
+
+    stats = {"updates": 0, "characters": 0}
+
+    def track_stats(update: ChatResponseUpdate) -> ChatResponseUpdate:
+        """Track statistics as updates flow through."""
+        stats["updates"] += 1
+        stats["characters"] += len(update.text or "")
+        return update
+
+    def log_cleanup() -> None:
+        """Log when stream consumption completes."""
+        print(f"  [Cleanup] Stream complete: {stats['updates']} updates, {stats['characters']} chars")
+
+    def add_stats_to_response(response: ChatResponse) -> ChatResponse:
+        """Result hook to include the statistics in the final response."""
+        response.additional_properties["stats"] = stats.copy()
+        return response
+
+    # All hooks can be passed via constructor
+    full_stream: ResponseStream[ChatResponseUpdate, ChatResponse] = ResponseStream(
+        generate_updates(),
+        finalizer=combine_updates,
+        transform_hooks=[track_stats],
+        result_hooks=[add_stats_to_response],
+        cleanup_hooks=[log_cleanup],
+    )
+
+    print("Processing with all hooks active:")
+    async for update in full_stream:
+        print(f"  -> '{update.text}'")
+
+    final_full = await full_stream.get_final_response()
+    print(f"\nFinal: '{final_full.text}'")
+    print(f"Stats: {final_full.additional_properties['stats']}")
+
+
+if __name__ == "__main__":
+    asyncio.run(main())