mirror of
https://github.com/microsoft/agent-framework.git
synced 2026-06-16 21:04:09 +08:00
Python: Update getting started with workflows sample structure and README (#653)
* Update getting started with workflows sample structure and README * Small updates * Adjust getting started samples. Fix agent executor bug. Add workflow tests to unit test file. * Fix resource links
This commit is contained in:
Evan Mattson
committed by
GitHub
Unverified
GitHub
parent
2472f35876
commit
2699b85285
@@ -0,0 +1,167 @@
|
||||
# Copyright (c) Microsoft. All rights reserved.
|
||||
|
||||
import asyncio
|
||||
from dataclasses import dataclass
|
||||
from typing import Any
|
||||
|
||||
from agent_framework import ChatMessage, Role # Core chat primitives to build LLM requests
|
||||
from agent_framework.azure import AzureChatClient # Client wrapper for Azure OpenAI chat models
|
||||
from agent_framework.workflow import (
|
||||
AgentExecutor, # Wraps an LLM agent for use inside a workflow
|
||||
AgentExecutorRequest, # The message bundle sent to an AgentExecutor
|
||||
AgentExecutorResponse, # The structured result returned by an AgentExecutor
|
||||
AgentRunEvent, # Tracing event for agent execution steps
|
||||
Executor, # Base class for custom Python executors
|
||||
WorkflowBuilder, # Fluent builder for wiring the workflow graph
|
||||
WorkflowCompletedEvent, # Terminal event carrying the final result
|
||||
WorkflowContext, # Per run context and event bus
|
||||
handler, # Decorator to mark an Executor method as invokable
|
||||
)
|
||||
from azure.identity import AzureCliCredential # Uses your az CLI login for credentials
|
||||
|
||||
"""
|
||||
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.
|
||||
- Simple tracing using AgentRunEvent to observe execution order and progress.
|
||||
|
||||
Prerequisites:
|
||||
- Familiarity with WorkflowBuilder, executors, edges, events, and streaming runs.
|
||||
- Azure OpenAI access configured for AzureChatClient. Log in with Azure CLI and set any required environment variables.
|
||||
- Comfort reading AgentExecutorResponse.agent_run_response.text for assistant output aggregation.
|
||||
"""
|
||||
|
||||
|
||||
class DispatchToExperts(Executor):
|
||||
"""Dispatches the incoming prompt to all expert agent executors for parallel processing (fan out)."""
|
||||
|
||||
def __init__(self, expert_ids: list[str], id: str | None = None):
|
||||
super().__init__(id)
|
||||
self._expert_ids = expert_ids
|
||||
|
||||
@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.
|
||||
# Each send_message targets a different AgentExecutor by id so that branches run in parallel.
|
||||
initial_message = ChatMessage(Role.USER, text=prompt)
|
||||
for expert_id in self._expert_ids:
|
||||
await ctx.send_message(
|
||||
AgentExecutorRequest(messages=[initial_message], should_respond=True),
|
||||
target_id=expert_id,
|
||||
)
|
||||
|
||||
|
||||
@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)."""
|
||||
|
||||
def __init__(self, expert_ids: list[str], id: str | None = None):
|
||||
super().__init__(id)
|
||||
self._expert_ids = expert_ids
|
||||
|
||||
@handler
|
||||
async def aggregate(self, results: list[AgentExecutorResponse], ctx: WorkflowContext[Any]) -> None:
|
||||
# Map responses to text by executor id for a simple, predictable demo.
|
||||
by_id: dict[str, str] = {}
|
||||
for r in results:
|
||||
# AgentExecutorResponse.agent_run_response.text is the assistant text produced by the agent.
|
||||
by_id[r.executor_id] = r.agent_run_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.add_event(WorkflowCompletedEvent(data=consolidated))
|
||||
|
||||
|
||||
async def main() -> None:
|
||||
# 1) Create agent executors for domain experts
|
||||
chat_client = AzureChatClient(credential=AzureCliCredential())
|
||||
|
||||
researcher = AgentExecutor(
|
||||
chat_client.create_agent(
|
||||
instructions=(
|
||||
"You're an expert market and product researcher. Given a prompt, provide concise, factual insights,"
|
||||
" opportunities, and risks."
|
||||
),
|
||||
),
|
||||
id="researcher",
|
||||
)
|
||||
marketer = AgentExecutor(
|
||||
chat_client.create_agent(
|
||||
instructions=(
|
||||
"You're a creative marketing strategist. Craft compelling value propositions and target messaging"
|
||||
" aligned to the prompt."
|
||||
),
|
||||
),
|
||||
id="marketer",
|
||||
)
|
||||
legal = AgentExecutor(
|
||||
chat_client.create_agent(
|
||||
instructions=(
|
||||
"You're a cautious legal/compliance reviewer. Highlight constraints, disclaimers, and policy concerns"
|
||||
" based on the prompt."
|
||||
),
|
||||
),
|
||||
id="legal",
|
||||
)
|
||||
|
||||
expert_ids = [researcher.id, marketer.id, legal.id]
|
||||
|
||||
dispatcher = DispatchToExperts(expert_ids=expert_ids, id="dispatcher")
|
||||
aggregator = AggregateInsights(expert_ids=expert_ids, id="aggregator")
|
||||
|
||||
# 2) Build a simple fan out and fan in workflow
|
||||
workflow = (
|
||||
WorkflowBuilder()
|
||||
.set_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 print progress plus the final consolidated output
|
||||
completion: WorkflowCompletedEvent | None = None
|
||||
async for event in workflow.run_stream("We are launching a new budget-friendly electric bike for urban commuters."):
|
||||
if isinstance(event, AgentRunEvent):
|
||||
# Show which agent ran and what step completed for lightweight observability.
|
||||
print(event)
|
||||
if isinstance(event, WorkflowCompletedEvent):
|
||||
completion = event
|
||||
|
||||
if completion:
|
||||
print("===== Final Aggregated Output =====")
|
||||
print(completion.data)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
asyncio.run(main())
|
||||
@@ -0,0 +1,317 @@
|
||||
# Copyright (c) Microsoft. All rights reserved.
|
||||
|
||||
import ast
|
||||
import asyncio
|
||||
import os
|
||||
from collections import defaultdict
|
||||
from dataclasses import dataclass
|
||||
from typing import Any
|
||||
|
||||
import aiofiles
|
||||
from agent_framework.workflow import (
|
||||
Executor, # Base class for custom workflow steps
|
||||
WorkflowBuilder, # Fluent graph builder for executors and edges
|
||||
WorkflowCompletedEvent, # Terminal event that carries final output
|
||||
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
|
||||
)
|
||||
|
||||
"""
|
||||
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 shared 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.
|
||||
- aiofiles installed for async file I/O.
|
||||
- Write access to a tmp directory next to this script.
|
||||
- A source text at resources/long_text.txt.
|
||||
- Optional for SVG export: install the viz extra for agent framework workflow.
|
||||
"""
|
||||
|
||||
# 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 shared state to store the data to be processed
|
||||
SHARED_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)
|
||||
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 shared state.
|
||||
|
||||
Args:
|
||||
data: The raw text to process.
|
||||
ctx: Workflow context to persist shared 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.
|
||||
await ctx.set_shared_state(SHARED_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 shared state keyed by its own executor id.
|
||||
await ctx.set_shared_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 shared state access and messaging.
|
||||
"""
|
||||
# Retrieve tokens and our assigned slice.
|
||||
data_to_be_processed: list[str] = await ctx.get_shared_state(SHARED_STATE_DATA_KEY)
|
||||
chunk_start, chunk_end = await ctx.get_shared_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")
|
||||
async with aiofiles.open(file_path, "w") as f:
|
||||
await 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)
|
||||
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")
|
||||
async with aiofiles.open(file_path, "w") as f:
|
||||
await 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:
|
||||
async with aiofiles.open(result.file_path, "r") as f:
|
||||
map_results.extend([
|
||||
(line.strip().split(": ")[0], int(line.strip().split(": ")[1])) for line in await 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.
|
||||
async with aiofiles.open(data.file_path, "r") as f:
|
||||
lines = await 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")
|
||||
async with aiofiles.open(file_path, "w") as f:
|
||||
await 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 emits the final completion event."""
|
||||
|
||||
@handler
|
||||
async def complete(self, data: list[ReduceCompleted], ctx: WorkflowContext[Any]) -> None:
|
||||
"""Collect reducer output file paths and publish a terminal event."""
|
||||
await ctx.add_event(WorkflowCompletedEvent(data=[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 the executors.
|
||||
map_operations = [Map(id=f"map_executor_{i}") for i in range(3)]
|
||||
split_operation = Split(
|
||||
[map_operation.id for map_operation in map_operations],
|
||||
id="split_data_executor",
|
||||
)
|
||||
reduce_operations = [Reduce(id=f"reduce_executor_{i}") for i in range(4)]
|
||||
shuffle_operation = Shuffle(
|
||||
[reduce_operation.id for reduce_operation in reduce_operations],
|
||||
id="shuffle_executor",
|
||||
)
|
||||
completion_executor = CompletionExecutor(id="completion_executor")
|
||||
|
||||
# Step 2: Build the workflow graph using fan out and fan in edges.
|
||||
workflow = (
|
||||
WorkflowBuilder()
|
||||
.set_start_executor(split_operation)
|
||||
.add_fan_out_edges(split_operation, map_operations) # Split -> many mappers
|
||||
.add_fan_in_edges(map_operations, shuffle_operation) # All mappers -> shuffle
|
||||
.add_fan_out_edges(shuffle_operation, reduce_operations) # Shuffle -> many reducers
|
||||
.add_fan_in_edges(reduce_operations, 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-workflow[viz]")
|
||||
|
||||
# Step 3: Open the text file and read its content.
|
||||
async with aiofiles.open(os.path.join(DIR, "resources", "long_text.txt"), "r") as f:
|
||||
raw_text = await f.read()
|
||||
|
||||
# Step 4: Run the workflow with the raw text as input.
|
||||
completion_event = None
|
||||
async for event in workflow.run_stream(raw_text):
|
||||
print(f"Event: {event}")
|
||||
if isinstance(event, WorkflowCompletedEvent):
|
||||
completion_event = event
|
||||
|
||||
if completion_event:
|
||||
print(f"Completion Event: {completion_event}")
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
asyncio.run(main())
|
||||
Reference in New Issue
Block a user