Added an optional sources parameter to CodeExecutorAgent for more control on which agents can contribute code. Also moved the _extract_markdown_code_blocks to an member method so that it can be overriden by subclasses.

Author: afourney

python/packages/autogen-agentchat/src/autogen_agentchat/agents/_code_executor_agent.py

@@ -9,17 +9,6 @@
 from ._base_chat_agent import BaseChatAgent
 
 
-def _extract_markdown_code_blocks(markdown_text: str) -> List[CodeBlock]:
-    pattern = re.compile(r"```(?:\s*([\w\+\-]+))?\n([\s\S]*?)```")
-    matches = pattern.findall(markdown_text)
-    code_blocks: List[CodeBlock] = []
-    for match in matches:
-        language = match[0].strip() if match[0] else ""
-        code_content = match[1]
-        code_blocks.append(CodeBlock(code=code_content, language=language))
-    return code_blocks
-
-
 class CodeExecutorAgent(BaseChatAgent):
     """An agent that extracts and executes code snippets found in received messages and returns the output.
 
@@ -75,9 +64,11 @@ def __init__(
         code_executor: CodeExecutor,
         *,
         description: str = "A computer terminal that performs no other action than running Python scripts (provided to it quoted in ```python code blocks), or sh shell scripts (provided to it quoted in ```sh code blocks).",
+        sources: Sequence[str] | None = None,
     ) -> None:
         super().__init__(name=name, description=description)
         self._code_executor = code_executor
+        self._sources = sources
 
     @property
     def produced_message_types(self) -> Sequence[type[ChatMessage]]:
@@ -89,7 +80,9 @@ async def on_messages(self, messages: Sequence[ChatMessage], cancellation_token:
         code_blocks: List[CodeBlock] = []
         for msg in messages:
             if isinstance(msg, TextMessage):
-                code_blocks.extend(_extract_markdown_code_blocks(msg.content))
+                if self._sources is not None and msg.source not in self._sources:
+                    continue
+                code_blocks.extend(self._extract_markdown_code_blocks(msg.content))
         if code_blocks:
             # Execute the code blocks.
             result = await self._code_executor.execute_code_blocks(code_blocks, cancellation_token=cancellation_token)
@@ -114,3 +107,13 @@ async def on_messages(self, messages: Sequence[ChatMessage], cancellation_token:
     async def on_reset(self, cancellation_token: CancellationToken) -> None:
         """It it's a no-op as the code executor agent has no mutable state."""
         pass
+
+    def _extract_markdown_code_blocks(self, markdown_text: str) -> List[CodeBlock]:
+        pattern = re.compile(r"```(?:\s*([\w\+\-]+))?\n([\s\S]*?)```")
+        matches = pattern.findall(markdown_text)
+        code_blocks: List[CodeBlock] = []
+        for match in matches:
+            language = match[0].strip() if match[0] else ""
+            code_content = match[1]
+            code_blocks.append(CodeBlock(code=code_content, language=language))
+        return code_blocks