ag2 0.9.3__tar.gz → 0.9.4__tar.gz

{ag2-0.9.3 → ag2-0.9.4}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: ag2
-Version: 0.9.3
+Version: 0.9.4
 Summary: A programming framework for agentic AI
 Project-URL: Homepage, https://ag2.ai/
 Project-URL: Documentation, https://docs.ag2.ai

{ag2-0.9.3 → ag2-0.9.4}/autogen/agentchat/contrib/agent_optimizer.py

@@ -140,7 +140,7 @@ History:
 
 According to the information I provide, please take one of four actions to manipulate list B using the functions you know.
 Instead of returning TERMINATE directly or taking no action, you should try your best to optimize the function list. Only take no action if you really think the current list is optimal, as more actions will harm performance in future tasks.
-Even adding a general function that can substitute the assistant’s repeated suggestions of Python code with the same functionality could also be helpful.
+Even adding a general function that can substitute the assistant's repeated suggestions of Python code with the same functionality could also be helpful.
 """
 
 
@@ -176,14 +176,15 @@ class AgentOptimizer:
     def __init__(
         self,
         max_actions_per_step: int,
-        llm_config: Union[LLMConfig, dict[str, Any]],
+        llm_config: Optional[Union[LLMConfig, dict[str, Any]]] = None,
         optimizer_model: Optional[str] = "gpt-4-1106-preview",
     ):
         """(These APIs are experimental and may change in the future.)
 
         Args:
             max_actions_per_step (int): the maximum number of actions that the optimizer can take in one step.
-            llm_config (LLMConfig or dict): llm inference configuration.
+            llm_config (LLMConfig or dict or None): llm inference configuration.
+                If None, the current LLMConfig from context is used.
                 Please refer to [OpenAIWrapper.create](https://docs.ag2.ai/latest/docs/api-reference/autogen/OpenAIWrapper/#autogen.OpenAIWrapper.create) for available options.
                 When using OpenAI or Azure OpenAI endpoints, please specify a non-empty 'model' either in `llm_config` or in each config of 'config_list' in `llm_config`.
             optimizer_model: the model used for the optimizer.
@@ -203,6 +204,8 @@ class AgentOptimizer:
         self._failure_functions_performance = []
         self._best_performance = -1
 
+        if llm_config is None:
+            llm_config = LLMConfig.current
         assert isinstance(llm_config, (dict, LLMConfig)), "llm_config must be a dict or LLMConfig"
         llm_config = copy.deepcopy(llm_config)
         self.llm_config = llm_config

{ag2-0.9.3 → ag2-0.9.4}/autogen/agentchat/conversable_agent.py

@@ -77,13 +77,13 @@ from .chat import (
     initiate_chats,
 )
 from .group.context_variables import ContextVariables
+from .group.guardrails import Guardrail
 from .group.handoffs import Handoffs
 from .utils import consolidate_chat_info, gather_usage_summary
 
 if TYPE_CHECKING:
     from .group.on_condition import OnCondition
     from .group.on_context_condition import OnContextCondition
-
 __all__ = ("ConversableAgent",)
 
 logger = logging.getLogger(__name__)
@@ -224,6 +224,8 @@ class ConversableAgent(LLMAgent):
             handoffs (Handoffs): Handoffs object containing all handoff transition conditions.
         """
         self.handoffs = handoffs if handoffs is not None else Handoffs()
+        self.input_guardrails: list["Guardrail"] = []
+        self.output_guardrails: list["Guardrail"] = []
 
         # we change code_execution_config below and we have to make sure we don't change the input
         # in case of UserProxyAgent, without this we could even change the default value {}
@@ -1471,7 +1473,11 @@ class ConversableAgent(LLMAgent):
                 self.send(msg2send, recipient, request_reply=True, silent=silent)
 
             else:  # No breaks in the for loop, so we have reached max turns
-                iostream.send(TerminationEvent(termination_reason=f"Maximum turns ({max_turns}) reached"))
+                iostream.send(
+                    TerminationEvent(
+                        termination_reason=f"Maximum turns ({max_turns}) reached", sender=self, recipient=recipient
+                    )
+                )
         else:
             self._prepare_chat(recipient, clear_history)
             if isinstance(message, Callable):
@@ -1651,7 +1657,11 @@ class ConversableAgent(LLMAgent):
                     break
                 await self.a_send(msg2send, recipient, request_reply=True, silent=silent)
             else:  # No breaks in the for loop, so we have reached max turns
-                iostream.send(TerminationEvent(termination_reason=f"Maximum turns ({max_turns}) reached"))
+                iostream.send(
+                    TerminationEvent(
+                        termination_reason=f"Maximum turns ({max_turns}) reached", sender=self, recipient=recipient
+                    )
+                )
         else:
             self._prepare_chat(recipient, clear_history)
             if isinstance(message, Callable):
@@ -2587,7 +2597,7 @@ class ConversableAgent(LLMAgent):
             self._consecutive_auto_reply_counter[sender] = 0
 
             if termination_reason:
-                iostream.send(TerminationEvent(termination_reason=termination_reason))
+                iostream.send(TerminationEvent(termination_reason=termination_reason, sender=self, recipient=sender))
 
             return True, None
 
@@ -2727,7 +2737,7 @@ class ConversableAgent(LLMAgent):
             self._consecutive_auto_reply_counter[sender] = 0
 
             if termination_reason:
-                iostream.send(TerminationEvent(termination_reason=termination_reason))
+                iostream.send(TerminationEvent(termination_reason=termination_reason, sender=self, recipient=sender))
 
             return True, None
 
@@ -3994,6 +4004,42 @@ class ConversableAgent(LLMAgent):
         """
         self.handoffs.add_many(conditions)
 
+    def register_input_guardrail(self, guardrail: "Guardrail") -> None:
+        """
+        Register a guardrail to be used for input validation.
+
+        Args:
+            guardrail: The guardrail to register.
+        """
+        self.input_guardrails.append(guardrail)
+
+    def register_input_guardrails(self, guardrails: list["Guardrail"]) -> None:
+        """
+        Register multiple guardrails to be used for input validation.
+
+        Args:
+            guardrails: List of guardrails to register.
+        """
+        self.input_guardrails.extend(guardrails)
+
+    def register_output_guardrail(self, guardrail: "Guardrail") -> None:
+        """
+        Register a guardrail to be used for output validation.
+
+        Args:
+            guardrail: The guardrail to register.
+        """
+        self.output_guardrails.append(guardrail)
+
+    def register_output_guardrails(self, guardrails: list["Guardrail"]) -> None:
+        """
+        Register multiple guardrails to be used for output validation.
+
+        Args:
+            guardrails: List of guardrails to register.
+        """
+        self.output_guardrails.extend(guardrails)
+
 
 @export_module("autogen")
 def register_function(

{ag2-0.9.3 → ag2-0.9.4}/autogen/agentchat/group/group_utils.py

@@ -137,13 +137,7 @@ def _run_oncontextconditions(
         if is_available and (
             on_condition.condition is None or on_condition.condition.evaluate(agent.context_variables)
         ):
-            # Condition has been met, we'll set the Tool Executor's next target
-            # attribute and that will be picked up on the next iteration when
-            # _determine_next_agent is called
-            for agent in agent._group_manager.groupchat.agents:  # type: ignore[attr-defined]
-                if isinstance(agent, GroupToolExecutor):
-                    agent.set_next_target(on_condition.target)
-                    break
+            on_condition.target.activate_target(agent._group_manager.groupchat)  # type: ignore[attr-defined]
 
             transfer_name = on_condition.target.display_name()
 
@@ -212,6 +206,18 @@ def ensure_handoff_agents_in_group(agents: list["ConversableAgent"]) -> None:
                 raise ValueError("Agent in after work target Hand-offs must be in the agents list")
 
 
+def ensure_guardrail_agents_in_group(agents: list["ConversableAgent"]) -> None:
+    """Ensure the agents in handoffs are in the group chat."""
+    agent_names = [agent.name for agent in agents]
+    for agent in agents:
+        for guardrail in agent.input_guardrails + agent.output_guardrails:
+            if (
+                isinstance(guardrail.target, (AgentTarget, AgentNameTarget))
+                and guardrail.target.agent_name not in agent_names
+            ):
+                raise ValueError("Agent in guardrail's target must be in the agents list")
+
+
 def prepare_exclude_transit_messages(agents: list["ConversableAgent"]) -> None:
     """Preparation for excluding transit messages by getting all tool names and registering a hook on agents to remove those messages."""
     # get all transit functions names
@@ -254,6 +260,9 @@ def prepare_group_agents(
     # Ensure all agents in hand-off after-works are in the passed in agents list
     ensure_handoff_agents_in_group(agents)
 
+    # Ensure all agents in guardrails are in the passed in agents list
+    ensure_guardrail_agents_in_group(agents)
+
     # Create Tool Executor for the group
     tool_execution = GroupToolExecutor()
 

ag2-0.9.4/autogen/agentchat/group/guardrails.py

@@ -0,0 +1,171 @@
+# Copyright (c) 2023 - 2025, AG2ai, Inc., AG2ai open-source projects maintainers and core contributors
+#
+# SPDX-License-Identifier: Apache-2.0
+
+import json
+import re
+from abc import ABC, abstractmethod
+from typing import TYPE_CHECKING, Any, Optional, Union
+
+from pydantic import BaseModel, Field
+
+from ...oai.client import OpenAIWrapper
+
+if TYPE_CHECKING:
+    from ...llm_config import LLMConfig
+    from .targets.transition_target import TransitionTarget
+
+
+class GuardrailResult(BaseModel):
+    """Represents the outcome of a guardrail check."""
+
+    activated: bool
+    justification: str = Field(default="No justification provided")
+
+    def __str__(self) -> str:
+        return f"Guardrail Result: {self.activated}\nJustification: {self.justification}"
+
+    @staticmethod
+    def parse(text: str) -> "GuardrailResult":
+        """Parses a JSON string into a GuardrailResult object.
+
+        Args:
+            text (str): The JSON string to parse.
+
+        Returns:
+            GuardrailResult: The parsed GuardrailResult object.
+        """
+        try:
+            data = json.loads(text)
+            return GuardrailResult(**data)
+        except (json.JSONDecodeError, ValueError) as e:
+            raise ValueError(f"Failed to parse GuardrailResult from text: {text}") from e
+
+
+class Guardrail(ABC):
+    """Abstract base class for guardrails."""
+
+    def __init__(
+        self, name: str, condition: str, target: "TransitionTarget", activation_message: Optional[str] = None
+    ) -> None:
+        self.name = name
+        self.condition = condition
+        self.target = target
+        self.activation_message = (
+            activation_message if activation_message else f"Guardrail '{name}' has been activated."
+        )
+
+    @abstractmethod
+    def check(
+        self,
+        context: Union[str, list[dict[str, Any]]],
+    ) -> GuardrailResult:
+        """Checks the text against the guardrail and returns a GuardrailResult.
+
+        Args:
+            context (Union[str, list[dict[str, Any]]]): The context to check against the guardrail.
+
+        Returns:
+            GuardrailResult: The result of the guardrail check.
+        """
+        pass
+
+
+class LLMGuardrail(Guardrail):
+    """Guardrail that uses an LLM to check the context."""
+
+    def __init__(
+        self,
+        name: str,
+        condition: str,
+        target: "TransitionTarget",
+        llm_config: "LLMConfig",
+        activation_message: Optional[str] = None,
+    ) -> None:
+        super().__init__(name, condition, target, activation_message)
+
+        if not llm_config:
+            raise ValueError("LLMConfig is required.")
+
+        self.llm_config = llm_config.deepcopy()
+        setattr(self.llm_config, "response_format", GuardrailResult)
+        self.client = OpenAIWrapper(**self.llm_config.model_dump())
+
+        self.check_prompt = f"""You are a guardrail that checks if a condition is met in the conversation you are given.
+You will activate the guardrail only if the condition is met.
+
+**Condition: {self.condition}**"""
+
+    def check(
+        self,
+        context: Union[str, list[dict[str, Any]]],
+    ) -> GuardrailResult:
+        """Checks the context against the guardrail using an LLM.
+
+        Args:
+            context (Union[str, list[dict[str, Any]]]): The context to check against the guardrail.
+
+        Returns:
+            GuardrailResult: The result of the guardrail check.
+        """
+        # Set the check prompt as the system message
+        check_messages = [{"role": "system", "content": self.check_prompt}]
+        # If context is a string, wrap it in a user message and append it
+        if isinstance(context, str):
+            check_messages.append({"role": "user", "content": context})
+        # If context is a list of messages, append them
+        elif isinstance(context, list):
+            check_messages.extend(context)
+        else:
+            raise ValueError("Context must be a string or a list of messages.")
+        # Call the LLM with the check messages
+        response = self.client.create(messages=check_messages)
+        return GuardrailResult.parse(response.choices[0].message.content)  # type: ignore
+
+
+class RegexGuardrail(Guardrail):
+    """Guardrail that checks the context against a regular expression."""
+
+    def __init__(
+        self,
+        name: str,
+        condition: str,
+        target: "TransitionTarget",
+        activation_message: Optional[str] = None,
+    ) -> None:
+        super().__init__(name, condition, target, activation_message)
+        # Compile the regular expression condition
+        try:
+            self.regex = re.compile(condition)
+        except re.error as e:
+            raise ValueError(f"Invalid regex pattern '{condition}': {str(e)}")
+
+    def check(
+        self,
+        context: Union[str, list[dict[str, Any]]],
+    ) -> GuardrailResult:
+        """Checks the context against the guardrail using a regular expression.
+
+        Args:
+            context (Union[str, list[dict[str, Any]]]): The context to check against the guardrail.
+
+        Returns:
+            GuardrailResult: The result of the guardrail check.
+        """
+        # Create a list of the messages to check
+        if isinstance(context, str):
+            messages = [context]
+        elif isinstance(context, list):
+            messages = [message.get("content", "") for message in context]
+        else:
+            raise ValueError("Context must be a string or a list of messages.")
+
+        # Check each message against the regex
+        for message in messages:
+            match = self.regex.search(message)
+            # If a match is found, activate the guardrail and return the result
+            if match:
+                activated = True
+                justification = f"Match found -> {match.group(0)}"
+                return GuardrailResult(activated=activated, justification=justification)
+        return GuardrailResult(activated=False, justification="No match found in the context.")

{ag2-0.9.3 → ag2-0.9.4}/autogen/agentchat/group/targets/transition_target.py

@@ -68,6 +68,16 @@ class TransitionTarget(BaseModel):
         """Create a wrapper agent for the target if needed."""
         raise NotImplementedError("Requires subclasses to implement.")
 
+    def activate_target(self, groupchat: "GroupChat") -> None:
+        """Activate the target in the groupchat, setting the next target for GroupToolExecutor.
+
+        The Tool Executor's next target attribute will be picked up on the next iteration when _determine_next_agent is called"""
+        for agent in groupchat.agents:  # type: ignore[attr-defined]
+            # get the GroupToolExecutor agent
+            if type(agent).__name__ == "GroupToolExecutor":
+                agent.set_next_target(self)  # type: ignore[attr-defined]
+                return
+
 
 class AgentTarget(TransitionTarget):
     """Target that represents a direct agent reference."""

{ag2-0.9.3 → ag2-0.9.4}/autogen/agentchat/groupchat.py

@@ -1023,6 +1023,38 @@ class GroupChat:
                 mentions[agent.name] = count
         return mentions
 
+    def _run_input_guardrails(
+        self,
+        agent: "ConversableAgent",
+        messages: Optional[list[dict[str, Any]]] = None,
+    ) -> Optional[str]:
+        """Run input guardrails for an agent before the reply is generated.
+        Args:
+            agent (ConversableAgent): The agent whose input guardrails to run.
+            messages (Optional[list[dict[str, Any]]]): The messages to check against the guardrails.
+        """
+        for guardrail in agent.input_guardrails:
+            guardrail_result = guardrail.check(context=messages)
+
+            if guardrail_result.activated:
+                guardrail.target.activate_target(self)
+                return f"{guardrail.activation_message}\nJustification: {guardrail_result.justification}"
+        return None
+
+    def _run_output_guardrails(self, agent: "ConversableAgent", reply: str) -> None:
+        """Run output guardrails for an agent after the reply is generated.
+        Args:
+            agent (ConversableAgent): The agent whose output guardrails to run.
+            reply (str): The reply generated by the agent.
+        """
+        for guardrail in agent.output_guardrails:
+            guardrail_result = guardrail.check(context=reply)
+
+            if guardrail_result.activated:
+                guardrail.target.activate_target(self)
+                return f"{guardrail.activation_message}\nJustification: {guardrail_result.justification}"
+        return None
+
 
 @export_module("autogen")
 class GroupChatManager(ConversableAgent):
@@ -1194,8 +1226,17 @@ class GroupChatManager(ConversableAgent):
                 if not silent:
                     iostream = IOStream.get_default()
                     iostream.send(GroupChatRunChatEvent(speaker=speaker, silent=silent))
-                # let the speaker speak
-                reply = speaker.generate_reply(sender=self)
+
+                guardrails_activated = False
+                guardrails_reply = groupchat._run_input_guardrails(speaker, speaker._oai_messages[self])
+
+                if guardrails_reply is not None:
+                    # if a guardrail has been activated, then the next target has been set and the guardrail reply will be sent
+                    guardrails_activated = True
+                    reply = guardrails_reply
+                else:
+                    # let the speaker speak
+                    reply = speaker.generate_reply(sender=self)
             except KeyboardInterrupt:
                 # let the admin agent speak if interrupted
                 if groupchat.admin_name in groupchat.agent_names:
@@ -1215,6 +1256,15 @@ class GroupChatManager(ConversableAgent):
                 termination_reason = "No reply generated"
                 break
 
+            if not guardrails_activated:
+                # if the input guardrails were not activated, and the agent returned a reply
+                guardrails_reply = groupchat._run_output_guardrails(speaker, reply)
+
+                if guardrails_reply is not None:
+                    # if a guardrail has been activated, then the next target has been set and the guardrail reply will be sent
+                    guardrails_activated = True
+                    reply = guardrails_reply
+
             # check for "clear history" phrase in reply and activate clear history function if found
             if (
                 groupchat.enable_clear_history
@@ -1233,7 +1283,11 @@ class GroupChatManager(ConversableAgent):
                 a.previous_cache = None
 
         if termination_reason:
-            iostream.send(TerminationEvent(termination_reason=termination_reason))
+            iostream.send(
+                TerminationEvent(
+                    termination_reason=termination_reason, sender=self, recipient=speaker if speaker else None
+                )
+            )
 
         return True, None
 
@@ -1287,8 +1341,19 @@ class GroupChatManager(ConversableAgent):
             try:
                 # select the next speaker
                 speaker = await groupchat.a_select_speaker(speaker, self)
-                # let the speaker speak
-                reply = await speaker.a_generate_reply(sender=self)
+                if not silent:
+                    iostream.send(GroupChatRunChatEvent(speaker=speaker, silent=silent))
+
+                guardrails_activated = False
+                guardrails_reply = groupchat._run_input_guardrails(speaker, speaker._oai_messages[self])
+
+                if guardrails_reply is not None:
+                    # if a guardrail has been activated, then the next target has been set and the guardrail reply will be sent
+                    guardrails_activated = True
+                    reply = guardrails_reply
+                else:
+                    # let the speaker speak
+                    reply = await speaker.a_generate_reply(sender=self)
             except KeyboardInterrupt:
                 # let the admin agent speak if interrupted
                 if groupchat.admin_name in groupchat.agent_names:
@@ -1308,6 +1373,24 @@ class GroupChatManager(ConversableAgent):
                 termination_reason = "No reply generated"
                 break
 
+            if not guardrails_activated:
+                # if the input guardrails were not activated, and the agent returned a reply
+                guardrails_reply = groupchat._run_output_guardrails(speaker, reply)
+
+                if guardrails_reply is not None:
+                    # if a guardrail has been activated, then the next target has been set and the guardrail reply will be sent
+                    guardrails_activated = True
+                    reply = guardrails_reply
+
+            # check for "clear history" phrase in reply and activate clear history function if found
+            if (
+                groupchat.enable_clear_history
+                and isinstance(reply, dict)
+                and reply["content"]
+                and "CLEAR HISTORY" in reply["content"].upper()
+            ):
+                reply["content"] = self.clear_agents_history(reply, groupchat)
+
             # The speaker sends the message without requesting a reply
             await speaker.a_send(reply, self, request_reply=False, silent=silent)
             message = self.last_message(speaker)
@@ -1317,7 +1400,11 @@ class GroupChatManager(ConversableAgent):
                 a.previous_cache = None
 
         if termination_reason:
-            iostream.send(TerminationEvent(termination_reason=termination_reason))
+            iostream.send(
+                TerminationEvent(
+                    termination_reason=termination_reason, sender=self, recipient=speaker if speaker else None
+                )
+            )
 
         return True, None
 

{ag2-0.9.3 → ag2-0.9.4}/autogen/agents/experimental/websurfer/websurfer.py

@@ -12,7 +12,9 @@ from ....tools.experimental import (
     BrowserUseTool,
     Crawl4AITool,
     DuckDuckGoSearchTool,
+    FirecrawlTool,
     PerplexitySearchTool,
+    SearxngSearchTool,
     TavilySearchTool,
 )
 
@@ -28,7 +30,9 @@ class WebSurferAgent(ConversableAgent):
         *,
         llm_config: Optional[Union[LLMConfig, dict[str, Any]]] = None,
         web_tool_llm_config: Optional[Union[LLMConfig, dict[str, Any]]] = None,
-        web_tool: Literal["browser_use", "crawl4ai", "duckduckgo", "perplexity", "tavily"] = "browser_use",
+        web_tool: Literal[
+            "browser_use", "crawl4ai", "duckduckgo", "firecrawl", "perplexity", "tavily", "searxng"
+        ] = "browser_use",
         web_tool_kwargs: Optional[dict[str, Any]] = None,
         **kwargs: Any,
     ) -> None:
@@ -48,12 +52,16 @@ class WebSurferAgent(ConversableAgent):
             self.tool: Tool = BrowserUseTool(llm_config=web_tool_llm_config, **web_tool_kwargs)  # type: ignore[arg-type]
         elif web_tool == "crawl4ai":
             self.tool = Crawl4AITool(llm_config=web_tool_llm_config, **web_tool_kwargs)
+        elif web_tool == "firecrawl":
+            self.tool = FirecrawlTool(llm_config=web_tool_llm_config, **web_tool_kwargs)
         elif web_tool == "perplexity":
             self.tool = PerplexitySearchTool(**web_tool_kwargs)
         elif web_tool == "tavily":
             self.tool = TavilySearchTool(llm_config=web_tool_llm_config, **web_tool_kwargs)
         elif web_tool == "duckduckgo":
             self.tool = DuckDuckGoSearchTool(**web_tool_kwargs)
+        elif web_tool == "searxng":
+            self.tool = SearxngSearchTool(**web_tool_kwargs)
         else:
             raise ValueError(f"Unsupported {web_tool=}.")
 

{ag2-0.9.3 → ag2-0.9.4}/autogen/events/agent_events.py

@@ -669,16 +669,22 @@ class TerminationEvent(BaseEvent):
     """When a workflow termination condition is met"""
 
     termination_reason: str
+    sender: str
+    recipient: Optional[str] = None
 
     def __init__(
         self,
         *,
         uuid: Optional[UUID] = None,
+        sender: Union["Agent", str],
+        recipient: Optional[Union["Agent", str]] = None,
         termination_reason: str,
     ):
         super().__init__(
             uuid=uuid,
             termination_reason=termination_reason,
+            sender=sender.name if hasattr(sender, "name") else sender,
+            recipient=recipient.name if hasattr(recipient, "name") else recipient if recipient else None,
         )
 
     def print(self, f: Optional[Callable[..., Any]] = None) -> None:

{ag2-0.9.3 → ag2-0.9.4}/autogen/events/helpers.py

@@ -13,12 +13,15 @@ logger = logging.getLogger(__name__)
 def deprecated_by(
     new_class: type[BaseModel],
     param_mapping: dict[str, str] = None,
+    default_params: dict[str, any] = None,
 ) -> Callable[[type[BaseModel]], Callable[..., BaseModel]]:
     param_mapping = param_mapping or {}
+    default_params = default_params or {}
 
     def decorator(
         old_class: type[BaseModel],
         param_mapping: dict[str, str] = param_mapping,
+        default_params: dict[str, any] = default_params,
     ) -> Callable[..., BaseModel]:
         @wraps(old_class)
         def wrapper(*args, **kwargs) -> BaseModel:
@@ -28,6 +31,11 @@ def deprecated_by(
             # Translate old parameters to new parameters
             new_kwargs = {param_mapping.get(k, k): v for k, v in kwargs.items()}
 
+            # Add default parameters if not already present
+            for key, value in default_params.items():
+                if key not in new_kwargs:
+                    new_kwargs[key] = value
+
             # Pass the translated parameters to the new class
             return new_class(*args, **new_kwargs)
 

ag2-0.9.4/autogen/mcp/helpers.py

@@ -0,0 +1,45 @@
+# Copyright (c) 2023 - 2025, AG2ai, Inc., AG2ai open-source projects maintainers and core contributors
+#
+# SPDX-License-Identifier: Apache-2.0
+import asyncio
+import os
+import signal
+from asyncio.subprocess import PIPE, Process, create_subprocess_exec
+from contextlib import asynccontextmanager
+from typing import AsyncGenerator, Dict, Optional
+
+
+@asynccontextmanager
+async def run_streamable_http_client(
+    *, mcp_server_path: str, env_vars: Optional[Dict[str, str]] = None, startup_wait_secs: float = 5.0
+) -> AsyncGenerator[Process, None]:
+    """
+    Async context manager to run a Python subprocess for streamable-http with custom env vars.
+
+    Args:
+        mcp_server_path: Path to the Python script to run.
+        env_vars: Environment variables to export to the subprocess.
+        startup_wait_secs: Time to wait for the server to start (in seconds).
+    Yields:
+        An asyncio.subprocess.Process object.
+    """
+    env = os.environ.copy()
+    if env_vars:
+        env.update(env_vars)
+
+    process = await create_subprocess_exec(
+        "python", mcp_server_path, "streamable-http", env=env, stdout=PIPE, stderr=PIPE
+    )
+
+    # Optional startup delay to let the server initialize
+    await asyncio.sleep(startup_wait_secs)
+
+    try:
+        yield process
+    finally:
+        if process.returncode is None:
+            process.send_signal(signal.SIGINT)
+            try:
+                await asyncio.wait_for(process.wait(), timeout=5.0)
+            except asyncio.TimeoutError:
+                process.kill()