Python: WorkflowBuilder registry (#2486)

* Add workflow builder factory pattern

* Add internal edge groups to registered executors; next samples

* Update samples: Part 1

* register -> register_executor

* update hil samples

* Update other samples

* Update agent  samples

* Update doc string

* Add new sample

* Fix mypy

* Address comments

* Fix mypy

python/samples/getting_started/workflows/control-flow/edge_condition.py

@@ -5,9 +5,9 @@ import os
 from typing import Any
 
 from agent_framework import (  # Core chat primitives used to build requests
-    AgentExecutor,  # Wraps an LLM agent that can be invoked inside a workflow
     AgentExecutorRequest,  # Input message bundle for an AgentExecutor
-    AgentExecutorResponse,  # Output from an AgentExecutor
+    AgentExecutorResponse,
+    ChatAgent,  # Output from an AgentExecutor
     ChatMessage,
     Role,
     WorkflowBuilder,  # Fluent builder for wiring executors and edges
@@ -128,38 +128,35 @@ async def to_email_assistant_request(
     await ctx.send_message(AgentExecutorRequest(messages=[user_msg], should_respond=True))
 
 
-async def main() -> None:
-    # Create agents
+def create_spam_detector_agent() -> ChatAgent:
+    """Helper to create a spam detection agent."""
     # AzureCliCredential uses your current az login. This avoids embedding secrets in code.
-    chat_client = AzureOpenAIChatClient(credential=AzureCliCredential())
-
-    # Agent 1. Classifies spam and returns a DetectionResult object.
-    # response_format enforces that the LLM returns parsable JSON for the Pydantic model.
-    spam_detection_agent = AgentExecutor(
-        chat_client.create_agent(
-            instructions=(
-                "You are a spam detection assistant that identifies spam emails. "
-                "Always return JSON with fields is_spam (bool), reason (string), and email_content (string). "
-                "Include the original email content in email_content."
-            ),
-            response_format=DetectionResult,
+    return AzureOpenAIChatClient(credential=AzureCliCredential()).create_agent(
+        instructions=(
+            "You are a spam detection assistant that identifies spam emails. "
+            "Always return JSON with fields is_spam (bool), reason (string), and email_content (string). "
+            "Include the original email content in email_content."
         ),
-        id="spam_detection_agent",
+        name="spam_detection_agent",
+        response_format=DetectionResult,
     )
 
-    # Agent 2. Drafts a professional reply. Also uses structured JSON output for reliability.
-    email_assistant_agent = AgentExecutor(
-        chat_client.create_agent(
-            instructions=(
-                "You are an email assistant that helps users draft professional responses to emails. "
-                "Your input may be a JSON object that includes 'email_content'; base your reply on that content. "
-                "Return JSON with a single field 'response' containing the drafted reply."
-            ),
-            response_format=EmailResponse,
+
+def create_email_assistant_agent() -> ChatAgent:
+    """Helper to create an email assistant agent."""
+    # AzureCliCredential uses your current az login. This avoids embedding secrets in code.
+    return AzureOpenAIChatClient(credential=AzureCliCredential()).create_agent(
+        instructions=(
+            "You are an email assistant that helps users draft professional responses to emails. "
+            "Your input may be a JSON object that includes 'email_content'; base your reply on that content. "
+            "Return JSON with a single field 'response' containing the drafted reply."
         ),
-        id="email_assistant_agent",
+        name="email_assistant_agent",
+        response_format=EmailResponse,
     )
 
+
+async def main() -> None:
     # Build the workflow graph.
     # Start at the spam detector.
     # If not spam, hop to a transformer that creates a new AgentExecutorRequest,
@@ -167,13 +164,18 @@ async def main() -> None:
     # If spam, go directly to the spam handler and finalize.
     workflow = (
         WorkflowBuilder()
-        .set_start_executor(spam_detection_agent)
+        .register_agent(create_spam_detector_agent, name="spam_detection_agent")
+        .register_agent(create_email_assistant_agent, name="email_assistant_agent")
+        .register_executor(lambda: to_email_assistant_request, name="to_email_assistant_request")
+        .register_executor(lambda: handle_email_response, name="send_email")
+        .register_executor(lambda: handle_spam_classifier_response, name="handle_spam")
+        .set_start_executor("spam_detection_agent")
         # Not spam path: transform response -> request for assistant -> assistant -> send email
-        .add_edge(spam_detection_agent, to_email_assistant_request, condition=get_condition(False))
-        .add_edge(to_email_assistant_request, email_assistant_agent)
-        .add_edge(email_assistant_agent, handle_email_response)
+        .add_edge("spam_detection_agent", "to_email_assistant_request", condition=get_condition(False))
+        .add_edge("to_email_assistant_request", "email_assistant_agent")
+        .add_edge("email_assistant_agent", "send_email")
         # Spam path: send to spam handler
-        .add_edge(spam_detection_agent, handle_spam_classifier_response, condition=get_condition(True))
+        .add_edge("spam_detection_agent", "handle_spam", condition=get_condition(True))
         .build()
     )
 

python/samples/getting_started/workflows/control-flow/multi_selection_edge_group.py

@@ -9,9 +9,9 @@ from typing import Literal
 from uuid import uuid4
 
 from agent_framework import (
-    AgentExecutor,
     AgentExecutorRequest,
     AgentExecutorResponse,
+    ChatAgent,
     ChatMessage,
     Role,
     WorkflowBuilder,
@@ -181,40 +181,38 @@ async def database_access(analysis: AnalysisResult, ctx: WorkflowContext[Never,
     await ctx.add_event(DatabaseEvent(f"Email {analysis.email_id} saved to database."))
 
 
+def create_email_analysis_agent() -> ChatAgent:
+    """Creates the email analysis agent."""
+    return AzureOpenAIChatClient(credential=AzureCliCredential()).create_agent(
+        instructions=(
+            "You are a spam detection assistant that identifies spam emails. "
+            "Always return JSON with fields 'spam_decision' (one of NotSpam, Spam, Uncertain) "
+            "and 'reason' (string)."
+        ),
+        name="email_analysis_agent",
+        response_format=AnalysisResultAgent,
+    )
+
+
+def create_email_assistant_agent() -> ChatAgent:
+    """Creates the email assistant agent."""
+    return AzureOpenAIChatClient(credential=AzureCliCredential()).create_agent(
+        instructions=("You are an email assistant that helps users draft responses to emails with professionalism."),
+        name="email_assistant_agent",
+        response_format=EmailResponse,
+    )
+
+
+def create_email_summary_agent() -> ChatAgent:
+    """Creates the email summary agent."""
+    return AzureOpenAIChatClient(credential=AzureCliCredential()).create_agent(
+        instructions=("You are an assistant that helps users summarize emails."),
+        name="email_summary_agent",
+        response_format=EmailSummaryModel,
+    )
+
+
 async def main() -> None:
-    # Agents
-    chat_client = AzureOpenAIChatClient(credential=AzureCliCredential())
-
-    email_analysis_agent = AgentExecutor(
-        chat_client.create_agent(
-            instructions=(
-                "You are a spam detection assistant that identifies spam emails. "
-                "Always return JSON with fields 'spam_decision' (one of NotSpam, Spam, Uncertain) "
-                "and 'reason' (string)."
-            ),
-            response_format=AnalysisResultAgent,
-        ),
-        id="email_analysis_agent",
-    )
-
-    email_assistant_agent = AgentExecutor(
-        chat_client.create_agent(
-            instructions=(
-                "You are an email assistant that helps users draft responses to emails with professionalism."
-            ),
-            response_format=EmailResponse,
-        ),
-        id="email_assistant_agent",
-    )
-
-    email_summary_agent = AgentExecutor(
-        chat_client.create_agent(
-            instructions=("You are an assistant that helps users summarize emails."),
-            response_format=EmailSummaryModel,
-        ),
-        id="email_summary_agent",
-    )
-
     # Build the workflow
     def select_targets(analysis: AnalysisResult, target_ids: list[str]) -> list[str]:
         # Order: [handle_spam, submit_to_email_assistant, summarize_email, handle_uncertain]
@@ -228,24 +226,39 @@ async def main() -> None:
             return targets
         return [handle_uncertain_id]
 
-    workflow = (
+    workflow_builder = (
         WorkflowBuilder()
-        .set_start_executor(store_email)
-        .add_edge(store_email, email_analysis_agent)
-        .add_edge(email_analysis_agent, to_analysis_result)
+        .register_agent(create_email_analysis_agent, name="email_analysis_agent")
+        .register_agent(create_email_assistant_agent, name="email_assistant_agent")
+        .register_agent(create_email_summary_agent, name="email_summary_agent")
+        .register_executor(lambda: store_email, name="store_email")
+        .register_executor(lambda: to_analysis_result, name="to_analysis_result")
+        .register_executor(lambda: submit_to_email_assistant, name="submit_to_email_assistant")
+        .register_executor(lambda: finalize_and_send, name="finalize_and_send")
+        .register_executor(lambda: summarize_email, name="summarize_email")
+        .register_executor(lambda: merge_summary, name="merge_summary")
+        .register_executor(lambda: handle_spam, name="handle_spam")
+        .register_executor(lambda: handle_uncertain, name="handle_uncertain")
+        .register_executor(lambda: database_access, name="database_access")
+    )
+
+    workflow = (
+        workflow_builder.set_start_executor("store_email")
+        .add_edge("store_email", "email_analysis_agent")
+        .add_edge("email_analysis_agent", "to_analysis_result")
         .add_multi_selection_edge_group(
-            to_analysis_result,
-            [handle_spam, submit_to_email_assistant, summarize_email, handle_uncertain],
+            "to_analysis_result",
+            ["handle_spam", "submit_to_email_assistant", "summarize_email", "handle_uncertain"],
             selection_func=select_targets,
         )
-        .add_edge(submit_to_email_assistant, email_assistant_agent)
-        .add_edge(email_assistant_agent, finalize_and_send)
-        .add_edge(summarize_email, email_summary_agent)
-        .add_edge(email_summary_agent, merge_summary)
+        .add_edge("submit_to_email_assistant", "email_assistant_agent")
+        .add_edge("email_assistant_agent", "finalize_and_send")
+        .add_edge("summarize_email", "email_summary_agent")
+        .add_edge("email_summary_agent", "merge_summary")
         # Save to DB if short (no summary path)
-        .add_edge(to_analysis_result, database_access, condition=lambda r: r.email_length <= LONG_EMAIL_THRESHOLD)
+        .add_edge("to_analysis_result", "database_access", condition=lambda r: r.email_length <= LONG_EMAIL_THRESHOLD)
         # Save to DB with summary when long
-        .add_edge(merge_summary, database_access)
+        .add_edge("merge_summary", "database_access")
         .build()
     )
 

python/samples/getting_started/workflows/control-flow/sequential_executors.py

@@ -61,20 +61,18 @@ class ReverseTextExecutor(Executor):
 
 async def main() -> None:
     """Build a two step sequential workflow and run it with streaming to observe events."""
-    # Step 1: Create executor instances.
-    upper_case_executor = UpperCaseExecutor(id="upper_case_executor")
-    reverse_text_executor = ReverseTextExecutor(id="reverse_text_executor")
-
-    # Step 2: Build the workflow graph.
+    # Step 1: Build the workflow graph.
     # Order matters. We connect upper_case_executor -> reverse_text_executor and set the start.
     workflow = (
         WorkflowBuilder()
-        .add_edge(upper_case_executor, reverse_text_executor)
-        .set_start_executor(upper_case_executor)
+        .register_executor(lambda: UpperCaseExecutor(id="upper_case_executor"), name="upper_case_executor")
+        .register_executor(lambda: ReverseTextExecutor(id="reverse_text_executor"), name="reverse_text_executor")
+        .add_edge("upper_case_executor", "reverse_text_executor")
+        .set_start_executor("upper_case_executor")
         .build()
     )
 
-    # Step 3: Stream events for a single input.
+    # Step 2: Stream events for a single input.
     # The stream will include executor invoke and completion events, plus workflow outputs.
     outputs: list[str] = []
     async for event in workflow.run_stream("hello world"):

python/samples/getting_started/workflows/control-flow/sequential_streaming.py

@@ -52,11 +52,18 @@ async def reverse_text(text: str, ctx: WorkflowContext[Never, str]) -> None:
 
 async def main():
     """Build a two-step sequential workflow and run it with streaming to observe events."""
-    # Step 2: Build the workflow with the defined edges.
+    # Step 1: Build the workflow with the defined edges.
     # Order matters. upper_case_executor runs first, then reverse_text_executor.
-    workflow = WorkflowBuilder().add_edge(to_upper_case, reverse_text).set_start_executor(to_upper_case).build()
+    workflow = (
+        WorkflowBuilder()
+        .register_executor(lambda: to_upper_case, name="upper_case_executor")
+        .register_executor(lambda: reverse_text, name="reverse_text_executor")
+        .add_edge("upper_case_executor", "reverse_text_executor")
+        .set_start_executor("upper_case_executor")
+        .build()
+    )
 
-    # Step 3: Run the workflow and stream events in real time.
+    # Step 2: Run the workflow and stream events in real time.
     async for event in workflow.run_stream("hello world"):
         # You will see executor invoke and completion events as the workflow progresses.
         print(f"Event: {event}")

python/samples/getting_started/workflows/control-flow/simple_loop.py

@@ -4,16 +4,15 @@ import asyncio
 from enum import Enum
 
 from agent_framework import (
-    AgentExecutor,
     AgentExecutorRequest,
     AgentExecutorResponse,
+    ChatAgent,
     ChatMessage,
     Executor,
     ExecutorCompletedEvent,
     Role,
     WorkflowBuilder,
     WorkflowContext,
-    WorkflowOutputEvent,
     handler,
 )
 from agent_framework.azure import AzureOpenAIChatClient
@@ -49,9 +48,9 @@ class NumberSignal(Enum):
 class GuessNumberExecutor(Executor):
     """An executor that guesses a number."""
 
-    def __init__(self, bound: tuple[int, int], id: str | None = None):
+    def __init__(self, bound: tuple[int, int], id: str):
         """Initialize the executor with a target number."""
-        super().__init__(id=id or "guess_number")
+        super().__init__(id=id)
         self._lower = bound[0]
         self._upper = bound[1]
 
@@ -116,43 +115,37 @@ class ParseJudgeResponse(Executor):
             await ctx.send_message(NumberSignal.BELOW)
 
 
+def create_judge_agent() -> ChatAgent:
+    """Create a judge agent that evaluates guesses."""
+    return AzureOpenAIChatClient(credential=AzureCliCredential()).create_agent(
+        instructions=("You strictly respond with one of: MATCHED, ABOVE, BELOW based on the given target and guess."),
+        name="judge_agent",
+    )
+
+
 async def main():
     """Main function to run the workflow."""
-    # Step 1: Create the executors.
-    guess_number_executor = GuessNumberExecutor((1, 100))
-
-    # Agent judge setup
-    chat_client = AzureOpenAIChatClient(credential=AzureCliCredential())
-    judge_agent = AgentExecutor(
-        chat_client.create_agent(
-            instructions=(
-                "You strictly respond with one of: MATCHED, ABOVE, BELOW based on the given target and guess."
-            )
-        ),
-        id="judge_agent",
-    )
-    submit_to_judge = SubmitToJudgeAgent(judge_agent_id=judge_agent.id, target=30, id="submit_judge")
-    parse_judge = ParseJudgeResponse(id="parse_judge")
-
-    # Step 2: Build the workflow with the defined edges.
+    # Step 1: Build the workflow with the defined edges.
     # This time we are creating a loop in the workflow.
     workflow = (
         WorkflowBuilder()
-        .add_edge(guess_number_executor, submit_to_judge)
-        .add_edge(submit_to_judge, judge_agent)
-        .add_edge(judge_agent, parse_judge)
-        .add_edge(parse_judge, guess_number_executor)
-        .set_start_executor(guess_number_executor)
+        .register_executor(lambda: GuessNumberExecutor((1, 100), "guess_number"), name="guess_number")
+        .register_agent(create_judge_agent, name="judge_agent")
+        .register_executor(lambda: SubmitToJudgeAgent(judge_agent_id="judge_agent", target=30), name="submit_judge")
+        .register_executor(lambda: ParseJudgeResponse(id="parse_judge"), name="parse_judge")
+        .add_edge("guess_number", "submit_judge")
+        .add_edge("submit_judge", "judge_agent")
+        .add_edge("judge_agent", "parse_judge")
+        .add_edge("parse_judge", "guess_number")
+        .set_start_executor("guess_number")
         .build()
     )
 
-    # Step 3: Run the workflow and print the events.
+    # Step 2: Run the workflow and print the events.
     iterations = 0
     async for event in workflow.run_stream(NumberSignal.INIT):
-        if isinstance(event, ExecutorCompletedEvent) and event.executor_id == guess_number_executor.id:
+        if isinstance(event, ExecutorCompletedEvent) and event.executor_id == "guess_number":
             iterations += 1
-        elif isinstance(event, WorkflowOutputEvent):
-            print(f"Final result: {event.data}")
         print(f"Event: {event}")
 
     # This is essentially a binary search, so the number of iterations should be logarithmic.

python/samples/getting_started/workflows/control-flow/switch_case_edge_group.py

@@ -7,10 +7,10 @@ from typing import Any, Literal
 from uuid import uuid4
 
 from agent_framework import (  # Core chat primitives used to form LLM requests
-    AgentExecutor,  # Wraps an agent so it can run inside a workflow
     AgentExecutorRequest,  # Message bundle sent to an AgentExecutor
     AgentExecutorResponse,  # Result returned by an AgentExecutor
-    Case,  # Case entry for a switch-case edge group
+    Case,
+    ChatAgent,  # Case entry for a switch-case edge group
     ChatMessage,
     Default,  # Default branch when no cases match
     Role,
@@ -152,51 +152,56 @@ async def handle_uncertain(detection: DetectionResult, ctx: WorkflowContext[Neve
         raise RuntimeError("This executor should only handle Uncertain messages.")
 
 
+def create_spam_detection_agent() -> ChatAgent:
+    """Create and return the spam detection agent."""
+    return AzureOpenAIChatClient(credential=AzureCliCredential()).create_agent(
+        instructions=(
+            "You are a spam detection assistant that identifies spam emails. "
+            "Be less confident in your assessments. "
+            "Always return JSON with fields 'spam_decision' (one of NotSpam, Spam, Uncertain) "
+            "and 'reason' (string)."
+        ),
+        name="spam_detection_agent",
+        response_format=DetectionResultAgent,
+    )
+
+
+def create_email_assistant_agent() -> ChatAgent:
+    """Create and return the email assistant agent."""
+    return AzureOpenAIChatClient(credential=AzureCliCredential()).create_agent(
+        instructions=("You are an email assistant that helps users draft responses to emails with professionalism."),
+        name="email_assistant_agent",
+        response_format=EmailResponse,
+    )
+
+
 async def main():
     """Main function to run the workflow."""
-    chat_client = AzureOpenAIChatClient(credential=AzureCliCredential())
-
-    # Agents. response_format enforces that the LLM returns JSON that Pydantic can validate.
-    spam_detection_agent = AgentExecutor(
-        chat_client.create_agent(
-            instructions=(
-                "You are a spam detection assistant that identifies spam emails. "
-                "Be less confident in your assessments. "
-                "Always return JSON with fields 'spam_decision' (one of NotSpam, Spam, Uncertain) "
-                "and 'reason' (string)."
-            ),
-            response_format=DetectionResultAgent,
-        ),
-        id="spam_detection_agent",
-    )
-
-    email_assistant_agent = AgentExecutor(
-        chat_client.create_agent(
-            instructions=(
-                "You are an email assistant that helps users draft responses to emails with professionalism."
-            ),
-            response_format=EmailResponse,
-        ),
-        id="email_assistant_agent",
-    )
-
     # Build workflow: store -> detection agent -> to_detection_result -> switch (NotSpam or Spam or Default).
     # The switch-case group evaluates cases in order, then falls back to Default when none match.
     workflow = (
         WorkflowBuilder()
-        .set_start_executor(store_email)
-        .add_edge(store_email, spam_detection_agent)
-        .add_edge(spam_detection_agent, to_detection_result)
+        .register_agent(create_spam_detection_agent, name="spam_detection_agent")
+        .register_agent(create_email_assistant_agent, name="email_assistant_agent")
+        .register_executor(lambda: store_email, name="store_email")
+        .register_executor(lambda: to_detection_result, name="to_detection_result")
+        .register_executor(lambda: submit_to_email_assistant, name="submit_to_email_assistant")
+        .register_executor(lambda: finalize_and_send, name="finalize_and_send")
+        .register_executor(lambda: handle_spam, name="handle_spam")
+        .register_executor(lambda: handle_uncertain, name="handle_uncertain")
+        .set_start_executor("store_email")
+        .add_edge("store_email", "spam_detection_agent")
+        .add_edge("spam_detection_agent", "to_detection_result")
         .add_switch_case_edge_group(
-            to_detection_result,
+            "to_detection_result",
             [
-                Case(condition=get_case("NotSpam"), target=submit_to_email_assistant),
-                Case(condition=get_case("Spam"), target=handle_spam),
-                Default(target=handle_uncertain),
+                Case(condition=get_case("NotSpam"), target="submit_to_email_assistant"),
+                Case(condition=get_case("Spam"), target="handle_spam"),
+                Default(target="handle_uncertain"),
             ],
         )
-        .add_edge(submit_to_email_assistant, email_assistant_agent)
-        .add_edge(email_assistant_agent, finalize_and_send)
+        .add_edge("submit_to_email_assistant", "email_assistant_agent")
+        .add_edge("email_assistant_agent", "finalize_and_send")
         .build()
     )