chore: import upstream snapshot with attribution

python/samples/03-workflows/parallelism/aggregate_results_of_different_types.py

@@ -0,0 +1,98 @@
+# Copyright (c) Microsoft. All rights reserved.
+
+import asyncio
+import random
+
+from agent_framework import Executor, WorkflowBuilder, WorkflowContext, handler
+from typing_extensions import Never
+
+"""
+Sample: Concurrent fan out and fan in with two different tasks that output results of different types.
+
+Purpose:
+Show how to construct a parallel branch pattern in workflows. Demonstrate:
+- Fan out by targeting multiple executors from one dispatcher.
+- Fan in by collecting a list of results from the executors.
+
+Prerequisites:
+- Familiarity with WorkflowBuilder, executors, edges, events, and streaming runs.
+"""
+
+
+class Dispatcher(Executor):
+    """
+    The sole purpose of this decorator is to dispatch the input of the workflow to
+    other executors.
+    """
+
+    @handler
+    async def handle(self, numbers: list[int], ctx: WorkflowContext[list[int]]):
+        if not numbers:
+            raise RuntimeError("Input must be a valid list of integers.")
+
+        await ctx.send_message(numbers)
+
+
+class Average(Executor):
+    """Calculate the average of a list of integers."""
+
+    @handler
+    async def handle(self, numbers: list[int], ctx: WorkflowContext[float]):
+        average: float = sum(numbers) / len(numbers)
+        await ctx.send_message(average)
+
+
+class Sum(Executor):
+    """Calculate the sum of a list of integers."""
+
+    @handler
+    async def handle(self, numbers: list[int], ctx: WorkflowContext[int]):
+        total: int = sum(numbers)
+        await ctx.send_message(total)
+
+
+class Aggregator(Executor):
+    """Aggregate the results from the different tasks and yield the final output."""
+
+    @handler
+    async def handle(self, results: list[int | float], ctx: WorkflowContext[Never, list[int | float]]):
+        """Receive the results from the source executors.
+
+        The framework will automatically collect messages from the source executors
+        and deliver them as a list.
+
+        Args:
+            results (list[int | float]): execution results from upstream executors.
+                The type annotation must be a list of union types that the upstream
+                executors will produce.
+            ctx (WorkflowContext[Never, list[int | float]]): A workflow context that can yield the final output.
+        """
+        await ctx.yield_output(results)
+
+
+async def main() -> None:
+    # 1) Build a simple fan out and fan in workflow
+    dispatcher = Dispatcher(id="dispatcher")
+    average = Average(id="average")
+    summation = Sum(id="summation")
+    aggregator = Aggregator(id="aggregator")
+
+    workflow = (
+        WorkflowBuilder(start_executor=dispatcher)
+        .add_fan_out_edges(dispatcher, [average, summation])
+        .add_fan_in_edges([average, summation], aggregator)
+        .build()
+    )
+
+    # 2) Run the workflow
+    output: list[int | float] | None = None
+    async for event in workflow.run([random.randint(1, 100) for _ in range(10)], stream=True):
+        if event.type == "output":
+            output = event.data
+
+    if output is not None:
+        print(output)
+
+
+if __name__ == "__main__":
+    asyncio.run(main())

python/samples/03-workflows/parallelism/fan_out_fan_in_edges.py

@@ -0,0 +1,197 @@
+# Copyright (c) Microsoft. All rights reserved.
+
+import asyncio
+import os
+from dataclasses import dataclass
+
+from agent_framework import (
+    Agent,
+    AgentExecutor,  # Wraps a ChatAgent as an Executor for use in workflows
+    AgentExecutorRequest,  # The message bundle sent to an AgentExecutor
+    AgentExecutorResponse,  # The structured result returned by an AgentExecutor
+    AgentResponseUpdate,
+    Executor,  # Base class for custom Python executors
+    Message,  # Chat message structure
+    WorkflowBuilder,  # Fluent builder for wiring the workflow graph
+    WorkflowContext,  # Per run context and event bus
+    handler,  # Decorator to mark an Executor method as invokable
+)
+from agent_framework.foundry import FoundryChatClient
+from azure.identity import AzureCliCredential  # Uses your az CLI login for credentials
+from dotenv import load_dotenv
+from typing_extensions import Never
+
+# Load environment variables from .env file
+load_dotenv()
+
+"""
+Sample: Concurrent fan out and fan in with three domain agents
+
+A dispatcher fans out the same user prompt to research, marketing, and legal AgentExecutor nodes.
+An aggregator then fans in their responses and produces a single consolidated report.
+
+Purpose:
+Show how to construct a parallel branch pattern in workflows. Demonstrate:
+- Fan out by targeting multiple AgentExecutor nodes from one dispatcher.
+- Fan in by collecting a list of AgentExecutorResponse objects and reducing them to a single result.
+
+Prerequisites:
+- FOUNDRY_PROJECT_ENDPOINT must be your Azure AI Foundry Agent Service (V2) project endpoint.
+- FOUNDRY_MODEL must be set to your Azure OpenAI model deployment name.
+- Familiarity with WorkflowBuilder, executors, edges, events, and streaming runs.
+- Comfort reading AgentExecutorResponse.agent_response.text for assistant output aggregation.
+"""
+
+
+class DispatchToExperts(Executor):
+    """Dispatches the incoming prompt to all expert agent executors for parallel processing (fan out)."""
+
+    @handler
+    async def dispatch(self, prompt: str, ctx: WorkflowContext[AgentExecutorRequest]) -> None:
+        # Wrap the incoming prompt as a user message for each expert and request a response.
+        initial_message = Message("user", contents=[prompt])
+        await ctx.send_message(AgentExecutorRequest(messages=[initial_message], should_respond=True))
+
+
+@dataclass
+class AggregatedInsights:
+    """Typed container for the aggregator to hold per domain strings before formatting."""
+
+    research: str
+    marketing: str
+    legal: str
+
+
+class AggregateInsights(Executor):
+    """Aggregates expert agent responses into a single consolidated result (fan in)."""
+
+    @handler
+    async def aggregate(self, results: list[AgentExecutorResponse], ctx: WorkflowContext[Never, str]) -> None:
+        # Map responses to text by executor id for a simple, predictable demo.
+        by_id: dict[str, str] = {}
+        for r in results:
+            # AgentExecutorResponse.agent_response.text is the assistant text produced by the agent.
+            by_id[r.executor_id] = r.agent_response.text
+
+        research_text = by_id.get("researcher", "")
+        marketing_text = by_id.get("marketer", "")
+        legal_text = by_id.get("legal", "")
+
+        aggregated = AggregatedInsights(
+            research=research_text,
+            marketing=marketing_text,
+            legal=legal_text,
+        )
+
+        # Provide a readable, consolidated string as the final workflow result.
+        consolidated = (
+            "Consolidated Insights\n"
+            "====================\n\n"
+            f"Research Findings:\n{aggregated.research}\n\n"
+            f"Marketing Angle:\n{aggregated.marketing}\n\n"
+            f"Legal/Compliance Notes:\n{aggregated.legal}\n"
+        )
+
+        await ctx.yield_output(consolidated)
+
+
+def render_live_streams(buffers: dict[str, str], order: list[str], completed: set[str]) -> None:
+    """Render concurrent agent streams in separate sections."""
+    # Clear terminal and move cursor to top-left for a live dashboard effect.
+    print("\033[2J\033[H", end="")
+    print("=== Expert Streams (Live) ===")
+    print("Concurrent agent updates are shown below as they stream.\n")
+    for agent_id in order:
+        state = "completed" if agent_id in completed else "streaming"
+        print(f"[{agent_id}] ({state})")
+        print(buffers.get(agent_id, ""))
+        print("-" * 80)
+    print("", end="", flush=True)
+
+
+async def main() -> None:
+    # 1) Create executor and agent instances
+    dispatcher = DispatchToExperts(id="dispatcher")
+    aggregator = AggregateInsights(id="aggregator")
+
+    researcher = AgentExecutor(
+        Agent(
+            client=FoundryChatClient(
+                project_endpoint=os.environ["FOUNDRY_PROJECT_ENDPOINT"],
+                model=os.environ["FOUNDRY_MODEL"],
+                credential=AzureCliCredential(),
+            ),
+            instructions=(
+                "You're an expert market and product researcher. Given a prompt, provide concise, factual insights,"
+                " opportunities, and risks."
+            ),
+            name="researcher",
+        )
+    )
+    marketer = AgentExecutor(
+        Agent(
+            client=FoundryChatClient(
+                project_endpoint=os.environ["FOUNDRY_PROJECT_ENDPOINT"],
+                model=os.environ["FOUNDRY_MODEL"],
+                credential=AzureCliCredential(),
+            ),
+            instructions=(
+                "You're a creative marketing strategist. Craft compelling value propositions and target messaging"
+                " aligned to the prompt."
+            ),
+            name="marketer",
+        )
+    )
+    legal = AgentExecutor(
+        Agent(
+            client=FoundryChatClient(
+                project_endpoint=os.environ["FOUNDRY_PROJECT_ENDPOINT"],
+                model=os.environ["FOUNDRY_MODEL"],
+                credential=AzureCliCredential(),
+            ),
+            instructions=(
+                "You're a cautious legal/compliance reviewer. Highlight constraints, disclaimers, and policy concerns"
+                " based on the prompt."
+            ),
+            name="legal",
+        )
+    )
+
+    # 2) Build a simple fan out and fan in workflow
+    workflow = (
+        WorkflowBuilder(start_executor=dispatcher)
+        .add_fan_out_edges(dispatcher, [researcher, marketer, legal])  # Parallel branches
+        .add_fan_in_edges([researcher, marketer, legal], aggregator)  # Join at the aggregator
+        .build()
+    )
+
+    # 3) Run with a single prompt and render live expert streams plus final consolidated output.
+    expert_order = ["researcher", "marketer", "legal"]
+    expert_buffers: dict[str, str] = {expert_id: "" for expert_id in expert_order}
+    completed_experts: set[str] = set()
+    final_output: str | None = None
+
+    async for event in workflow.run(
+        "We are launching a new budget-friendly electric bike for urban commuters.", stream=True
+    ):
+        if event.type == "executor_completed" and event.executor_id in expert_buffers:
+            completed_experts.add(event.executor_id)
+            render_live_streams(expert_buffers, expert_order, completed_experts)
+        elif event.type == "output":
+            if isinstance(event.data, AgentResponseUpdate):
+                executor_id = event.executor_id or ""
+                if executor_id in expert_buffers:
+                    expert_buffers[executor_id] += event.data.text
+                    render_live_streams(expert_buffers, expert_order, completed_experts)
+                continue
+
+            if event.executor_id == "aggregator":
+                final_output = str(event.data)
+
+    if final_output:
+        print("\n=== Final Consolidated Output ===\n")
+        print(final_output)
+
+
+if __name__ == "__main__":
+    asyncio.run(main())

python/samples/03-workflows/parallelism/map_reduce_and_visualization.py

@@ -0,0 +1,318 @@
+# Copyright (c) Microsoft. All rights reserved.
+
+import ast
+import asyncio
+import os
+from collections import defaultdict
+from dataclasses import dataclass
+
+from agent_framework import (
+    Executor,  # Base class for custom workflow steps
+    WorkflowBuilder,  # Fluent builder for executors and edges
+    WorkflowContext,  # Per run context with shared state and messaging
+    WorkflowViz,  # Utility to visualize a workflow graph
+    handler,  # Decorator to expose an Executor method as a step
+)
+from typing_extensions import Never
+
+"""
+Sample: Map reduce word count with fan out and fan in over file backed intermediate results
+
+The workflow splits a large text into chunks, maps words to counts in parallel,
+shuffles intermediate pairs to reducers, then reduces to per word totals.
+It also demonstrates WorkflowViz for graph visualization.
+
+Purpose:
+Show how to:
+- Partition input once and coordinate parallel mappers with workflow state.
+- Implement map, shuffle, and reduce executors that pass file paths instead of large payloads.
+- Use fan out and fan in edges to express parallelism and joins.
+- Persist intermediate results to disk to bound memory usage for large inputs.
+- Visualize the workflow graph using WorkflowViz and export to SVG with the optional viz extra.
+
+Prerequisites:
+- Familiarity with WorkflowBuilder, executors, fan out and fan in edges, events, and streaming runs.
+- Write access to a tmp directory next to this script.
+- A source text at resources/long_text.txt.
+- Optional for SVG export: install graphviz.
+
+Installation:
+    pip install agent-framework graphviz
+"""
+
+# Define the temporary directory for storing intermediate results
+DIR = os.path.dirname(__file__)
+TEMP_DIR = os.path.join(DIR, "tmp")
+# Ensure the temporary directory exists
+os.makedirs(TEMP_DIR, exist_ok=True)
+
+# Define a key for the workflow state to store the data to be processed
+STATE_DATA_KEY = "data_to_be_processed"
+
+
+class SplitCompleted:
+    """Marker type published when splitting finishes. Triggers map executors."""
+
+    ...
+
+
+class Split(Executor):
+    """Splits data into roughly equal chunks based on the number of mapper nodes."""
+
+    def __init__(self, map_executor_ids: list[str], id: str | None = None):
+        """Store mapper ids so we can assign non overlapping ranges per mapper."""
+        super().__init__(id=id or "split")
+        self._map_executor_ids = map_executor_ids
+
+    @handler
+    async def split(self, data: str, ctx: WorkflowContext[SplitCompleted]) -> None:
+        """Tokenize input and assign contiguous index ranges to each mapper via workflow state.
+
+        Args:
+            data: The raw text to process.
+            ctx: Workflow context to persist state and send messages.
+        """
+        # Process data into a list of words and remove empty lines or words.
+        word_list = self._preprocess(data)
+
+        # Store tokenized words once so all mappers can read by index.
+        ctx.set_state(STATE_DATA_KEY, word_list)
+
+        # Divide indices into contiguous slices for each mapper.
+        map_executor_count = len(self._map_executor_ids)
+        chunk_size = len(word_list) // map_executor_count  # Assumes count > 0.
+
+        async def _process_chunk(i: int) -> None:
+            """Assign the slice for mapper i, then signal that splitting is done."""
+            start_index = i * chunk_size
+            end_index = start_index + chunk_size if i < map_executor_count - 1 else len(word_list)
+
+            # The mapper reads its slice from workflow state keyed by its own executor id.
+            ctx.set_state(self._map_executor_ids[i], (start_index, end_index))
+            await ctx.send_message(SplitCompleted(), self._map_executor_ids[i])
+
+        tasks = [asyncio.create_task(_process_chunk(i)) for i in range(map_executor_count)]
+        await asyncio.gather(*tasks)
+
+    def _preprocess(self, data: str) -> list[str]:
+        """Normalize lines and split on whitespace. Return a flat list of tokens."""
+        line_list = [line.strip() for line in data.splitlines() if line.strip()]
+        return [word for line in line_list for word in line.split() if word]
+
+
+@dataclass
+class MapCompleted:
+    """Signal that a mapper wrote its intermediate pairs to file."""
+
+    file_path: str
+
+
+class Map(Executor):
+    """Maps each token to a count of 1 and writes pairs to a per mapper file."""
+
+    @handler
+    async def map(self, _: SplitCompleted, ctx: WorkflowContext[MapCompleted]) -> None:
+        """Read the assigned slice, emit (word, 1) pairs, and persist to disk.
+
+        Args:
+            _: SplitCompleted marker indicating maps can begin.
+            ctx: Workflow context for workflow state access and messaging.
+        """
+        # Retrieve tokens and our assigned slice.
+        data_to_be_processed: list[str] = ctx.get_state(STATE_DATA_KEY)
+        chunk_start, chunk_end = ctx.get_state(self.id)
+
+        results = [(item, 1) for item in data_to_be_processed[chunk_start:chunk_end]]
+
+        # Write this mapper's results as simple text lines for easy debugging.
+        file_path = os.path.join(TEMP_DIR, f"map_results_{self.id}.txt")
+        with open(file_path, "w") as f:
+            f.writelines([f"{item}: {count}\n" for item, count in results])
+
+        await ctx.send_message(MapCompleted(file_path))
+
+
+@dataclass
+class ShuffleCompleted:
+    """Signal that a shuffle partition file is ready for a specific reducer."""
+
+    file_path: str
+    reducer_id: str
+
+
+class Shuffle(Executor):
+    """Groups intermediate pairs by key and partitions them across reducers."""
+
+    def __init__(self, reducer_ids: list[str], id: str | None = None):
+        """Remember reducer ids so we can partition work deterministically."""
+        super().__init__(id=id or "shuffle")
+        self._reducer_ids = reducer_ids
+
+    @handler
+    async def shuffle(self, data: list[MapCompleted], ctx: WorkflowContext[ShuffleCompleted]) -> None:
+        """Aggregate mapper outputs and write one partition file per reducer.
+
+        Args:
+            data: MapCompleted records with file paths for each mapper output.
+            ctx: Workflow context to emit per reducer ShuffleCompleted messages.
+        """
+        chunks = await self._preprocess(data)
+
+        async def _process_chunk(chunk: list[tuple[str, list[int]]], index: int) -> None:
+            """Write one grouped partition for reducer index and notify that reducer."""
+            file_path = os.path.join(TEMP_DIR, f"shuffle_results_{index}.txt")
+            with open(file_path, "w") as f:
+                f.writelines([f"{key}: {value}\n" for key, value in chunk])
+            await ctx.send_message(ShuffleCompleted(file_path, self._reducer_ids[index]))
+
+        tasks = [asyncio.create_task(_process_chunk(chunk, i)) for i, chunk in enumerate(chunks)]
+        await asyncio.gather(*tasks)
+
+    async def _preprocess(self, data: list[MapCompleted]) -> list[list[tuple[str, list[int]]]]:
+        """Load all mapper files, group by key, sort keys, and partition for reducers.
+
+        Returns:
+            List of partitions. Each partition is a list of (key, [1, 1, ...]) tuples.
+        """
+        # Load all intermediate pairs.
+        map_results: list[tuple[str, int]] = []
+        for result in data:
+            with open(result.file_path) as f:
+                map_results.extend([
+                    (line.strip().split(": ")[0], int(line.strip().split(": ")[1])) for line in f.readlines()
+                ])
+
+        # Group values by token.
+        intermediate_results: defaultdict[str, list[int]] = defaultdict(list[int])
+        for key, value in map_results:
+            intermediate_results[key].append(value)
+
+        # Deterministic ordering helps with debugging and test stability.
+        aggregated_results = [(key, values) for key, values in intermediate_results.items()]
+        aggregated_results.sort(key=lambda x: x[0])
+
+        # Partition keys across reducers as evenly as possible.
+        reduce_executor_count = len(self._reducer_ids)
+        chunk_size = len(aggregated_results) // reduce_executor_count
+        remaining = len(aggregated_results) % reduce_executor_count
+
+        chunks = [
+            aggregated_results[i : i + chunk_size] for i in range(0, len(aggregated_results) - remaining, chunk_size)
+        ]
+        if remaining > 0:
+            chunks[-1].extend(aggregated_results[-remaining:])
+
+        return chunks
+
+
+@dataclass
+class ReduceCompleted:
+    """Signal that a reducer wrote final counts for its partition."""
+
+    file_path: str
+
+
+class Reduce(Executor):
+    """Sums grouped counts per key for its assigned partition."""
+
+    @handler
+    async def _execute(self, data: ShuffleCompleted, ctx: WorkflowContext[ReduceCompleted]) -> None:
+        """Read one shuffle partition and reduce it to totals.
+
+        Args:
+            data: ShuffleCompleted with the partition file path and target reducer id.
+            ctx: Workflow context used to emit ReduceCompleted with our output file path.
+        """
+        if data.reducer_id != self.id:
+            # This partition belongs to a different reducer. Skip.
+            return
+
+        # Read grouped values from the shuffle output.
+        with open(data.file_path) as f:
+            lines = f.readlines()
+
+        # Sum values per key. Values are serialized Python lists like [1, 1, ...].
+        reduced_results: dict[str, int] = defaultdict(int)
+        for line in lines:
+            key, value = line.split(": ")
+            reduced_results[key] = sum(ast.literal_eval(value))
+
+        # Persist our partition totals.
+        file_path = os.path.join(TEMP_DIR, f"reduced_results_{self.id}.txt")
+        with open(file_path, "w") as f:
+            f.writelines([f"{key}: {value}\n" for key, value in reduced_results.items()])
+
+        await ctx.send_message(ReduceCompleted(file_path))
+
+
+class CompletionExecutor(Executor):
+    """Joins all reducer outputs and yields the final output."""
+
+    @handler
+    async def complete(self, data: list[ReduceCompleted], ctx: WorkflowContext[Never, list[str]]) -> None:
+        """Collect reducer output file paths and yield final output."""
+        await ctx.yield_output([result.file_path for result in data])
+
+
+async def main():
+    """Construct the map reduce workflow, visualize it, then run it over a sample file."""
+
+    # Step 1: Create executor instances.
+    map_executor_0 = Map(id="map_executor_0")
+    map_executor_1 = Map(id="map_executor_1")
+    map_executor_2 = Map(id="map_executor_2")
+    split_data_executor = Split(["map_executor_0", "map_executor_1", "map_executor_2"], id="split_data_executor")
+    reduce_executor_0 = Reduce(id="reduce_executor_0")
+    reduce_executor_1 = Reduce(id="reduce_executor_1")
+    reduce_executor_2 = Reduce(id="reduce_executor_2")
+    reduce_executor_3 = Reduce(id="reduce_executor_3")
+    shuffle_executor = Shuffle(
+        ["reduce_executor_0", "reduce_executor_1", "reduce_executor_2", "reduce_executor_3"],
+        id="shuffle_executor",
+    )
+    completion_executor = CompletionExecutor(id="completion_executor")
+
+    mappers = [map_executor_0, map_executor_1, map_executor_2]
+    reducers = [reduce_executor_0, reduce_executor_1, reduce_executor_2, reduce_executor_3]
+
+    # Step 2: Build the workflow graph using fan out and fan in edges.
+    workflow = (
+        WorkflowBuilder(start_executor=split_data_executor)
+        .add_fan_out_edges(split_data_executor, mappers)  # Split -> many mappers
+        .add_fan_in_edges(mappers, shuffle_executor)  # All mappers -> shuffle
+        .add_fan_out_edges(shuffle_executor, reducers)  # Shuffle -> many reducers
+        .add_fan_in_edges(reducers, completion_executor)  # All reducers -> completion
+        .build()
+    )
+
+    # Step 2.5: Visualize the workflow (optional)
+    print("Generating workflow visualization...")
+    viz = WorkflowViz(workflow)
+    # Print out the Mermaid string.
+    print("Mermaid string: \n=======")
+    print(viz.to_mermaid())
+    print("=======")
+    # Print out the DiGraph string.
+    print("DiGraph string: \n=======")
+    print(viz.to_digraph())
+    print("=======")
+    try:
+        # Export the DiGraph visualization as SVG.
+        svg_file = viz.export(format="svg")
+        print(f"SVG file saved to: {svg_file}")
+    except ImportError:
+        print("Tip: Install 'viz' extra to export workflow visualization: pip install agent-framework[viz] --pre")
+
+    # Step 3: Open the text file and read its content.
+    with open(os.path.join(DIR, "../resources", "long_text.txt")) as f:
+        raw_text = f.read()
+
+    # Step 4: Run the workflow with the raw text as input.
+    async for event in workflow.run(raw_text, stream=True):
+        print(f"Event: {event}")
+        if event.type == "output":
+            print(f"Final Output: {event.data}")
+
+
+if __name__ == "__main__":
+    asyncio.run(main())