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/composition/sub_workflow_basics.py

@@ -8,7 +8,6 @@ from agent_framework import (
     Executor,
     WorkflowBuilder,
     WorkflowContext,
-    WorkflowEvent,
     WorkflowExecutor,
     handler,
 )
@@ -46,13 +45,6 @@ class TextProcessingResult:
     char_count: int
 
 
-class AllTasksCompleted(WorkflowEvent):
-    """Event triggered when all processing tasks are complete."""
-
-    def __init__(self, results: list[TextProcessingResult]):
-        super().__init__(results)
-
-
 # Sub-workflow executor
 class TextProcessor(Executor):
     """Processes text strings - counts words and characters."""
@@ -113,7 +105,11 @@ class TextProcessingOrchestrator(Executor):
             await ctx.send_message(request, target_id="text_processor_workflow")
 
     @handler
-    async def collect_result(self, result: TextProcessingResult, ctx: WorkflowContext) -> None:
+    async def collect_result(
+        self,
+        result: TextProcessingResult,
+        ctx: WorkflowContext[Never, list[TextProcessingResult]],
+    ) -> None:
         """Collect results from sub-workflows."""
         print(f"📥 Collected result from {result.task_id}")
         self.results.append(result)
@@ -121,48 +117,54 @@ class TextProcessingOrchestrator(Executor):
         # Check if all results are collected
         if len(self.results) == self.expected_count:
             print("\n🎉 All tasks completed!")
-            await ctx.add_event(AllTasksCompleted(self.results))
+            await ctx.yield_output(self.results)
 
-    def get_summary(self) -> dict[str, Any]:
-        """Get a summary of all processing results."""
-        total_words = sum(result.word_count for result in self.results)
-        total_chars = sum(result.char_count for result in self.results)
-        avg_words = total_words / len(self.results) if self.results else 0
-        avg_chars = total_chars / len(self.results) if self.results else 0
 
-        return {
-            "total_texts": len(self.results),
-            "total_words": total_words,
-            "total_characters": total_chars,
-            "average_words_per_text": round(avg_words, 2),
-            "average_characters_per_text": round(avg_chars, 2),
-        }
+def get_result_summary(results: list[TextProcessingResult]) -> dict[str, Any]:
+    """Get a summary of all processing results."""
+    total_words = sum(result.word_count for result in results)
+    total_chars = sum(result.char_count for result in results)
+    avg_words = total_words / len(results) if results else 0
+    avg_chars = total_chars / len(results) if results else 0
+
+    return {
+        "total_texts": len(results),
+        "total_words": total_words,
+        "total_characters": total_chars,
+        "average_words_per_text": round(avg_words, 2),
+        "average_characters_per_text": round(avg_chars, 2),
+    }
+
+
+def create_sub_workflow() -> WorkflowExecutor:
+    """Create the text processing sub-workflow."""
+    print("🚀 Setting up sub-workflow...")
+
+    processing_workflow = (
+        WorkflowBuilder()
+        .register_executor(TextProcessor, name="text_processor")
+        .set_start_executor("text_processor")
+        .build()
+    )
+
+    return WorkflowExecutor(processing_workflow, id="text_processor_workflow")
 
 
 async def main():
     """Main function to run the basic sub-workflow example."""
-    print("🚀 Setting up sub-workflow...")
-
-    # Step 1: Create the text processing sub-workflow
-    text_processor = TextProcessor()
-
-    processing_workflow = WorkflowBuilder().set_start_executor(text_processor).build()
-
     print("🔧 Setting up parent workflow...")
-
-    # Step 2: Create the parent workflow
-    orchestrator = TextProcessingOrchestrator()
-    workflow_executor = WorkflowExecutor(processing_workflow, id="text_processor_workflow")
-
+    # Step 1: Create the parent workflow
     main_workflow = (
         WorkflowBuilder()
-        .set_start_executor(orchestrator)
-        .add_edge(orchestrator, workflow_executor)
-        .add_edge(workflow_executor, orchestrator)
+        .register_executor(TextProcessingOrchestrator, name="text_orchestrator")
+        .register_executor(create_sub_workflow, name="text_processor_workflow")
+        .set_start_executor("text_orchestrator")
+        .add_edge("text_orchestrator", "text_processor_workflow")
+        .add_edge("text_processor_workflow", "text_orchestrator")
         .build()
     )
 
-    # Step 3: Test data - various text strings
+    # Step 2: Test data - various text strings
     test_texts = [
         "Hello world! This is a simple test.",
         "Python is a powerful programming language used for many applications.",
@@ -175,15 +177,17 @@ async def main():
     print(f"\n🧪 Testing with {len(test_texts)} text strings")
     print("=" * 60)
 
-    # Step 4: Run the workflow
-    await main_workflow.run(test_texts)
+    # Step 3: Run the workflow
+    result = await main_workflow.run(test_texts)
 
-    # Step 5: Display results
+    # Step 4: Display results
     print("\n📊 Processing Results:")
     print("=" * 60)
 
     # Sort results by task_id for consistent display
-    sorted_results = sorted(orchestrator.results, key=lambda r: r.task_id)
+    task_results = result.get_outputs()
+    assert len(task_results) == 1
+    sorted_results = sorted(task_results[0], key=lambda r: r.task_id)
 
     for result in sorted_results:
         preview = result.text[:30] + "..." if len(result.text) > 30 else result.text
@@ -191,7 +195,7 @@ async def main():
         print(f"✅ {result.task_id}: '{preview}' -> {result.word_count} words, {result.char_count} chars")
 
     # Step 6: Display summary
-    summary = orchestrator.get_summary()
+    summary = get_result_summary(sorted_results)
     print("\n📈 Summary:")
     print("=" * 60)
     print(f"📄 Total texts processed: {summary['total_texts']}")

python/samples/getting_started/workflows/composition/sub_workflow_parallel_requests.py

@@ -169,19 +169,18 @@ def build_resource_request_distribution_workflow() -> Workflow:
             elif len(self._responses) > self._request_count:
                 raise ValueError("Received more responses than expected")
 
-    orchestrator = RequestDistribution("orchestrator")
-    resource_requester = ResourceRequester("resource_requester")
-    policy_checker = PolicyChecker("policy_checker")
-    result_collector = ResultCollector("result_collector")
-
     return (
         WorkflowBuilder()
-        .set_start_executor(orchestrator)
-        .add_edge(orchestrator, resource_requester)
-        .add_edge(orchestrator, policy_checker)
-        .add_edge(resource_requester, result_collector)
-        .add_edge(policy_checker, result_collector)
-        .add_edge(orchestrator, result_collector)  # For request count
+        .register_executor(lambda: RequestDistribution("orchestrator"), name="orchestrator")
+        .register_executor(lambda: ResourceRequester("resource_requester"), name="resource_requester")
+        .register_executor(lambda: PolicyChecker("policy_checker"), name="policy_checker")
+        .register_executor(lambda: ResultCollector("result_collector"), name="result_collector")
+        .set_start_executor("orchestrator")
+        .add_edge("orchestrator", "resource_requester")
+        .add_edge("orchestrator", "policy_checker")
+        .add_edge("resource_requester", "result_collector")
+        .add_edge("policy_checker", "result_collector")
+        .add_edge("orchestrator", "result_collector")  # For request count
         .build()
     )
 
@@ -288,29 +287,27 @@ class PolicyEngine(Executor):
 
 
 async def main() -> None:
-    # Create executors in the main workflow
-    sub_workflow = build_resource_request_distribution_workflow()
-    resource_allocator = ResourceAllocator("resource_allocator")
-    policy_engine = PolicyEngine("policy_engine")
-
-    # Create the WorkflowExecutor for the sub-workflow
-    # Setting allow_direct_output=True to let the sub-workflow output directly.
-    # This is because the sub-workflow is the both the entry point and the exit
-    # point of the main workflow.
-    sub_workflow_executor = WorkflowExecutor(
-        sub_workflow,
-        "sub_workflow_executor",
-        allow_direct_output=True,
-    )
-
     # Build the main workflow
     main_workflow = (
         WorkflowBuilder()
-        .set_start_executor(sub_workflow_executor)
-        .add_edge(sub_workflow_executor, resource_allocator)
-        .add_edge(resource_allocator, sub_workflow_executor)
-        .add_edge(sub_workflow_executor, policy_engine)
-        .add_edge(policy_engine, sub_workflow_executor)
+        .register_executor(lambda: ResourceAllocator("resource_allocator"), name="resource_allocator")
+        .register_executor(lambda: PolicyEngine("policy_engine"), name="policy_engine")
+        .register_executor(
+            lambda: WorkflowExecutor(
+                build_resource_request_distribution_workflow(),
+                "sub_workflow_executor",
+                # Setting allow_direct_output=True to let the sub-workflow output directly.
+                # This is because the sub-workflow is the both the entry point and the exit
+                # point of the main workflow.
+                allow_direct_output=True,
+            ),
+            name="sub_workflow_executor",
+        )
+        .set_start_executor("sub_workflow_executor")
+        .add_edge("sub_workflow_executor", "resource_allocator")
+        .add_edge("resource_allocator", "sub_workflow_executor")
+        .add_edge("sub_workflow_executor", "policy_engine")
+        .add_edge("policy_engine", "sub_workflow_executor")
         .build()
     )
 

python/samples/getting_started/workflows/composition/sub_workflow_request_interception.py

@@ -154,15 +154,14 @@ def build_email_address_validation_workflow() -> Workflow:
                 )
 
     # Build the workflow
-    sanitizer = EmailSanitizer(id="email_sanitizer")
-    format_validator = EmailFormatValidator(id="email_format_validator")
-    domain_validator = DomainValidator(id="domain_validator")
-
     return (
         WorkflowBuilder()
-        .set_start_executor(sanitizer)
-        .add_edge(sanitizer, format_validator)
-        .add_edge(format_validator, domain_validator)
+        .register_executor(lambda: EmailSanitizer(id="email_sanitizer"), name="email_sanitizer")
+        .register_executor(lambda: EmailFormatValidator(id="email_format_validator"), name="email_format_validator")
+        .register_executor(lambda: DomainValidator(id="domain_validator"), name="domain_validator")
+        .set_start_executor("email_sanitizer")
+        .add_edge("email_sanitizer", "email_format_validator")
+        .add_edge("email_format_validator", "domain_validator")
         .build()
     )
 
@@ -270,21 +269,22 @@ async def main() -> None:
     # A list of approved domains
     approved_domains = {"example.com", "company.com"}
 
-    # Create executors in the main workflow
-    orchestrator = SmartEmailOrchestrator(id="smart_email_orchestrator", approved_domains=approved_domains)
-    email_delivery = EmailDelivery(id="email_delivery")
-
-    # Create the sub-workflow for email address validation
-    validation_workflow = build_email_address_validation_workflow()
-    validation_workflow_executor = WorkflowExecutor(validation_workflow, id="email_validation_workflow")
-
     # Build the main workflow
     workflow = (
         WorkflowBuilder()
-        .set_start_executor(orchestrator)
-        .add_edge(orchestrator, validation_workflow_executor)
-        .add_edge(validation_workflow_executor, orchestrator)
-        .add_edge(orchestrator, email_delivery)
+        .register_executor(
+            lambda: SmartEmailOrchestrator(id="smart_email_orchestrator", approved_domains=approved_domains),
+            name="smart_email_orchestrator",
+        )
+        .register_executor(lambda: EmailDelivery(id="email_delivery"), name="email_delivery")
+        .register_executor(
+            lambda: WorkflowExecutor(build_email_address_validation_workflow(), id="email_validation_workflow"),
+            name="email_validation_workflow",
+        )
+        .set_start_executor("smart_email_orchestrator")
+        .add_edge("smart_email_orchestrator", "email_validation_workflow")
+        .add_edge("email_validation_workflow", "smart_email_orchestrator")
+        .add_edge("smart_email_orchestrator", "email_delivery")
         .build()
     )