Python: Introducing Local MCP Servers (#389)

* mcp parts

* mcp parts 2

* removed structured output in favor of handling in chatresponse, mcp as AITool and running samples

* updated naming

* fixed test

python/packages/main/agent_framework/__init__.py

@@ -11,5 +11,6 @@ except importlib.metadata.PackageNotFoundError:
 from ._agents import *  # noqa: F403
 from ._clients import *  # noqa: F403
 from ._logging import *  # noqa: F403
+from ._mcp import *  # noqa: F403
 from ._tools import *  # noqa: F403
 from ._types import *  # noqa: F403

python/packages/main/agent_framework/_agents.py

@@ -2,14 +2,16 @@
 
 import sys
 from collections.abc import AsyncIterable, Callable, MutableMapping, Sequence
-from contextlib import AbstractAsyncContextManager
+from contextlib import AbstractAsyncContextManager, AsyncExitStack
 from enum import Enum
+from itertools import chain
 from typing import Any, ClassVar, Literal, Protocol, TypeVar, runtime_checkable
 from uuid import uuid4
 
-from pydantic import BaseModel, Field
+from pydantic import BaseModel, Field, PrivateAttr
 
 from ._clients import ChatClient
+from ._mcp import McpTool
 from ._pydantic import AFBaseModel
 from ._tools import AITool
 from ._types import (
@@ -315,6 +317,8 @@ class ChatClientAgent(AgentBase):
     chat_client: ChatClient
     instructions: str | None = None
     chat_options: ChatOptions
+    _local_mcp_tools: list[McpTool] = PrivateAttr(default_factory=list)  # type: ignore[reportUnknownVariableType]
+    _async_exit_stack: AsyncExitStack = PrivateAttr(default_factory=AsyncExitStack)
 
     def __init__(
         self,
@@ -383,6 +387,11 @@ class ChatClientAgent(AgentBase):
         """
         kwargs.update(additional_properties or {})
 
+        # We ignore the MCP Servers here and store them separately,
+        # we add their functions to the tools list at runtime
+        normalized_tools = [] if tools is None else tools if isinstance(tools, list) else [tools]
+        local_mcp_tools = [tool for tool in normalized_tools if isinstance(tool, McpTool)]
+        final_tools = [tool for tool in normalized_tools if not isinstance(tool, McpTool)]
         args: dict[str, Any] = {
             "chat_client": chat_client,
             "chat_options": ChatOptions(
@@ -398,7 +407,7 @@ class ChatClientAgent(AgentBase):
                 store=store,
                 temperature=temperature,
                 tool_choice=tool_choice,
-                tools=tools,  # type: ignore
+                tools=final_tools,  # type: ignore[reportArgumentType]
                 top_p=top_p,
                 user=user,
                 additional_properties=kwargs,
@@ -415,23 +424,27 @@ class ChatClientAgent(AgentBase):
 
         super().__init__(**args)
         self._update_agent_name()
+        self._local_mcp_tools = local_mcp_tools  # type: ignore[assignment]
 
     async def __aenter__(self) -> "Self":
         """Async context manager entry.
 
-        If the chat_client supports async context management, enter its context.
+        If either the chat_client or the local_mcp_tools are context managers,
+        they will be entered into the async exit stack to ensure proper cleanup.
+
+        This list might be extended in the future.
         """
-        if isinstance(self.chat_client, AbstractAsyncContextManager):
-            await self.chat_client.__aenter__()  # type: ignore[reportUnknownMemberType]
+        for context_manager in chain([self.chat_client], self._local_mcp_tools):
+            if isinstance(context_manager, AbstractAsyncContextManager):
+                await self._async_exit_stack.enter_async_context(context_manager)
         return self
 
     async def __aexit__(self, exc_type: type[BaseException] | None, exc_val: BaseException | None, exc_tb: Any) -> None:
         """Async context manager exit.
 
-        If the chat_client supports async context management, exit its context.
+        Close the async exit stack to ensure all context managers are exited properly.
         """
-        if isinstance(self.chat_client, AbstractAsyncContextManager):
-            await self.chat_client.__aexit__(exc_type, exc_val, exc_tb)  # type: ignore[reportUnknownMemberType]
+        await self._async_exit_stack.aclose()
 
     def _update_agent_name(self) -> None:
         """Update the agent name in a chat client.
@@ -506,6 +519,19 @@ class ChatClientAgent(AgentBase):
         thread, thread_messages = await self._prepare_thread_and_messages(thread=thread, input_messages=input_messages)
         agent_name = self._get_agent_name()
 
+        # Resolve final tool list (runtime provided tools + local MCP server tools)
+        final_tools: list[AITool | dict[str, Any] | Callable[..., Any]] = []
+        # Normalize tools argument to a list without mutating the original parameter
+        normalized_tools = [] if tools is None else tools if isinstance(tools, list) else [tools]
+        for tool in normalized_tools:
+            if isinstance(tool, McpTool):
+                final_tools.extend(tool.functions)  # type: ignore
+            else:
+                final_tools.append(tool)  # type: ignore
+
+        for mcp_server in self._local_mcp_tools:
+            final_tools.extend(mcp_server.functions)
+
         response = await self.chat_client.get_response(
             messages=thread_messages,
             chat_options=self.chat_options
@@ -523,7 +549,7 @@ class ChatClientAgent(AgentBase):
                 store=store,
                 temperature=temperature,
                 tool_choice=tool_choice,
-                tools=tools,  # type: ignore
+                tools=final_tools,  # type: ignore[reportArgumentType]
                 top_p=top_p,
                 user=user,
                 additional_properties=additional_properties or {},
@@ -617,6 +643,19 @@ class ChatClientAgent(AgentBase):
         agent_name = self._get_agent_name()
         response_updates: list[ChatResponseUpdate] = []
 
+        # Resolve final tool list (runtime provided tools + local MCP server tools)
+        final_tools: list[AITool | MutableMapping[str, Any] | Callable[..., Any]] = []
+        # Normalize tools argument to a list without mutating the original parameter
+        normalized_tools = [] if tools is None else tools if isinstance(tools, list) else [tools]
+        for tool in normalized_tools:
+            if isinstance(tool, McpTool):
+                final_tools.extend(tool.functions)  # type: ignore
+            else:
+                final_tools.append(tool)
+
+        for mcp_server in self._local_mcp_tools:
+            final_tools.extend(mcp_server.functions)
+
         async for update in self.chat_client.get_streaming_response(
             messages=thread_messages,
             chat_options=self.chat_options
@@ -634,7 +673,7 @@ class ChatClientAgent(AgentBase):
                 store=store,
                 temperature=temperature,
                 tool_choice=tool_choice,
-                tools=tools,  # type: ignore
+                tools=final_tools,  # type: ignore[reportArgumentType]
                 top_p=top_p,
                 user=user,
                 additional_properties=additional_properties or {},

python/packages/main/agent_framework/_clients.py

@@ -285,13 +285,13 @@ class ChatClient(Protocol):
 
         Args:
             messages: The sequence of input messages to send.
+            response_format: the format of the response.
             frequency_penalty: the frequency penalty to use.
             logit_bias: the logit bias to use.
             max_tokens: The maximum number of tokens to generate.
             metadata: additional metadata to include in the request.
             model: The model to use for the agent.
             presence_penalty: the presence penalty to use.
-            response_format: the format of the response.
             seed: the random seed to use.
             stop: the stop sequence(s) for the request.
             store: whether to store the response.

python/packages/main/agent_framework/_mcp.py

@@ -0,0 +1,784 @@
+# Copyright (c) Microsoft. All rights reserved.
+
+import json
+import logging
+import re
+import sys
+from abc import abstractmethod
+from contextlib import AsyncExitStack, _AsyncGeneratorContextManager  # type: ignore
+from datetime import timedelta
+from functools import partial
+from typing import TYPE_CHECKING, Any
+
+from mcp import types
+from mcp.client.session import ClientSession
+from mcp.client.sse import sse_client
+from mcp.client.stdio import StdioServerParameters, stdio_client
+from mcp.client.streamable_http import streamablehttp_client
+from mcp.client.websocket import websocket_client
+from mcp.shared.context import RequestContext
+from mcp.shared.exceptions import McpError
+from mcp.shared.session import RequestResponder
+from pydantic import BaseModel, create_model
+
+from ._tools import AIFunction
+from ._types import AIContents, ChatMessage, ChatRole, DataContent, TextContent, UriContent
+from .exceptions import ToolException, ToolExecutionException
+
+if sys.version_info >= (3, 11):
+    from typing import Self  # pragma: no cover
+else:
+    from typing_extensions import Self  # pragma: no cover
+
+if TYPE_CHECKING:
+    from ._clients import ChatClient
+
+logger = logging.getLogger(__name__)
+
+# region: Helpers
+
+LOG_LEVEL_MAPPING: dict[types.LoggingLevel, int] = {
+    "debug": logging.DEBUG,
+    "info": logging.INFO,
+    "notice": logging.INFO,
+    "warning": logging.WARNING,
+    "error": logging.ERROR,
+    "critical": logging.CRITICAL,
+    "alert": logging.CRITICAL,
+    "emergency": logging.CRITICAL,
+}
+
+__all__ = [
+    "McpSseTools",
+    "McpStdioTool",
+    "McpStreamableHttpTool",
+    "McpWebsocketTool",
+]
+
+
+def _mcp_prompt_message_to_chat_message(
+    mcp_type: types.PromptMessage | types.SamplingMessage,
+) -> ChatMessage:
+    """Convert a MCP container type to a Agent Framework type."""
+    return ChatMessage(
+        role=ChatRole(value=mcp_type.role),
+        contents=[_mcp_type_to_ai_content(mcp_type.content)],  # type: ignore[call-arg]
+        raw_representation=mcp_type,
+    )
+
+
+def _mcp_call_tool_result_to_ai_contents(
+    mcp_type: types.CallToolResult,
+) -> list[AIContents]:
+    """Convert a MCP container type to a Agent Framework type."""
+    return [_mcp_type_to_ai_content(item) for item in mcp_type.content]
+
+
+def _mcp_type_to_ai_content(
+    mcp_type: types.ImageContent | types.TextContent | types.AudioContent | types.EmbeddedResource | types.ResourceLink,
+) -> AIContents:
+    """Convert a MCP type to a Agent Framework type."""
+    match mcp_type:
+        case types.TextContent():
+            return TextContent(text=mcp_type.text, raw_representation=mcp_type)
+        case types.ImageContent() | types.AudioContent():
+            return DataContent(uri=mcp_type.data, media_type=mcp_type.mimeType, raw_representation=mcp_type)
+        case types.ResourceLink():
+            return UriContent(
+                uri=str(mcp_type.uri), media_type=mcp_type.mimeType or "application/json", raw_representation=mcp_type
+            )
+        case _:
+            match mcp_type.resource:
+                case types.TextResourceContents():
+                    return TextContent(
+                        text=mcp_type.resource.text,
+                        raw_representation=mcp_type,
+                        additional_properties=mcp_type.annotations.model_dump() if mcp_type.annotations else None,
+                    )
+                case types.BlobResourceContents():
+                    return DataContent(
+                        uri=mcp_type.resource.blob,
+                        media_type=mcp_type.resource.mimeType,
+                        raw_representation=mcp_type,
+                        additional_properties=mcp_type.annotations.model_dump() if mcp_type.annotations else None,
+                    )
+
+
+def _ai_content_to_mcp_types(
+    content: AIContents,
+) -> types.TextContent | types.ImageContent | types.AudioContent | types.EmbeddedResource | types.ResourceLink | None:
+    """Convert a AIContent type to a MCP type."""
+    match content:
+        case TextContent():
+            return types.TextContent(type="text", text=content.text)
+        case DataContent():
+            if content.media_type and content.media_type.startswith("image/"):
+                return types.ImageContent(type="image", data=content.uri, mimeType=content.media_type)
+            if content.media_type and content.media_type.startswith("audio/"):
+                return types.AudioContent(type="audio", data=content.uri, mimeType=content.media_type)
+            if content.media_type and content.media_type.startswith("application/"):
+                return types.EmbeddedResource(
+                    type="resource",
+                    resource=types.BlobResourceContents(
+                        blob=content.uri,
+                        mimeType=content.media_type,
+                        # uri's are not limited in MCP but they have to be set.
+                        # the uri of data content, contains the data uri, which
+                        # is not the uri meant here, UriContent would match this.
+                        uri=content.additional_properties.get("uri", "af://binary")
+                        if content.additional_properties
+                        else "af://binary",  # type: ignore[reportArgumentType]
+                    ),
+                )
+            return None
+        case UriContent():
+            return types.ResourceLink(
+                type="resource_link",
+                uri=content.uri,  # type: ignore[reportArgumentType]
+                mimeType=content.media_type,
+                name=content.additional_properties.get("name", "Unknown")
+                if content.additional_properties
+                else "Unknown",
+            )
+        case _:
+            return None
+
+
+def _chat_message_to_mcp_types(
+    content: ChatMessage,
+) -> list[types.TextContent | types.ImageContent | types.AudioContent | types.EmbeddedResource | types.ResourceLink]:
+    """Convert a ChatMessage to a list of MCP types."""
+    messages: list[
+        types.TextContent | types.ImageContent | types.AudioContent | types.EmbeddedResource | types.ResourceLink
+    ] = []
+    for item in content.contents:
+        mcp_content = _ai_content_to_mcp_types(item)
+        if mcp_content:
+            messages.append(mcp_content)
+    return messages
+
+
+def _get_input_model_from_mcp_prompt(prompt: types.Prompt) -> type[BaseModel]:
+    """Creates a Pydantic model from a prompt's parameters."""
+    # Check if 'arguments' is missing or empty
+    if not prompt.arguments:
+        return create_model(f"{prompt.name}_input")
+
+    field_definitions: dict[str, Any] = {}
+    for prompt_argument in prompt.arguments:
+        # For prompts, all arguments are typically required and string type
+        # unless specified otherwise in the prompt argument
+        python_type = str  # Default type for prompt arguments
+
+        # Create field definition for create_model
+        if prompt_argument.required:
+            field_definitions[prompt_argument.name] = (python_type, ...)
+        else:
+            field_definitions[prompt_argument.name] = (python_type, None)
+
+    return create_model(f"{prompt.name}_input", **field_definitions)
+
+
+def _get_input_model_from_mcp_tool(tool: types.Tool) -> type[BaseModel]:
+    """Creates a Pydantic model from a tools parameters."""
+    properties = tool.inputSchema.get("properties", None)
+    required = tool.inputSchema.get("required", [])
+    # Check if 'properties' is missing or not a dictionary
+    if not properties:
+        return create_model(f"{tool.name}_input")
+
+    field_definitions: dict[str, Any] = {}
+    for prop_name, prop_details in properties.items():
+        prop_details = json.loads(prop_details) if isinstance(prop_details, str) else prop_details
+
+        # Map JSON Schema types to Python types
+        json_type = prop_details.get("type", "string")
+        python_type: type = str  # default
+        if json_type == "integer":
+            python_type = int
+        elif json_type == "number":
+            python_type = float
+        elif json_type == "boolean":
+            python_type = bool
+        elif json_type == "array":
+            python_type = list
+        elif json_type == "object":
+            python_type = dict
+
+        # Create field definition for create_model
+        if prop_name in required:
+            field_definitions[prop_name] = (python_type, ...)
+        else:
+            default_value = prop_details.get("default", None)
+            field_definitions[prop_name] = (python_type, default_value)
+
+    return create_model(f"{tool.name}_input", **field_definitions)
+
+
+def _normalize_mcp_name(name: str) -> str:
+    """Normalize MCP tool/prompt names to allowed identifier pattern (A-Za-z0-9_.-)."""
+    return re.sub(r"[^A-Za-z0-9_.-]", "-", name)
+
+
+# region: MCP Plugin
+
+
+class McpTool:
+    """Base class with the MCP logic."""
+
+    def __init__(
+        self,
+        name: str,
+        description: str | None = None,
+        additional_properties: dict[str, Any] | None = None,
+        load_tools: bool = True,
+        load_prompts: bool = True,
+        session: ClientSession | None = None,
+        request_timeout: int | None = None,
+        chat_client: "ChatClient | None" = None,
+    ) -> None:
+        """Initialize the MCP Plugin Base."""
+        self.name = name
+        self.description = description or ""
+        self.additional_properties = additional_properties
+        self.load_tools_flag = load_tools
+        self.load_prompts_flag = load_prompts
+        self._exit_stack = AsyncExitStack()
+        self.session = session
+        self.request_timeout = request_timeout
+        self.chat_client = chat_client
+        self.functions: list[AIFunction[Any, Any]] = []
+
+    def __str__(self) -> str:
+        return f"McpTool(name={self.name}, description={self.description})"
+
+    async def connect(self) -> None:
+        """Connect to the MCP server."""
+        if not self.session:
+            try:
+                transport = await self._exit_stack.enter_async_context(self.get_mcp_client())
+            except Exception as ex:
+                await self._exit_stack.aclose()
+                raise ToolException(
+                    "Failed to connect to the MCP server. Please check your configuration.", inner_exception=ex
+                ) from ex
+            try:
+                session = await self._exit_stack.enter_async_context(
+                    ClientSession(
+                        read_stream=transport[0],
+                        write_stream=transport[1],
+                        read_timeout_seconds=timedelta(seconds=self.request_timeout) if self.request_timeout else None,
+                        message_handler=self.message_handler,
+                        logging_callback=self.logging_callback,
+                        sampling_callback=self.sampling_callback,
+                    )
+                )
+            except Exception as ex:
+                await self._exit_stack.aclose()
+                raise ToolException(
+                    message="Failed to create a session. Please check your configuration.", inner_exception=ex
+                ) from ex
+            await session.initialize()
+            self.session = session
+        elif self.session._request_id == 0:  # type: ignore[reportPrivateUsage]
+            # If the session is not initialized, we need to reinitialize it
+            await self.session.initialize()
+        logger.debug("Connected to MCP server: %s", self.session)
+        if self.load_tools_flag:
+            await self.load_tools()
+        if self.load_prompts_flag:
+            await self.load_prompts()
+
+        if logger.level != logging.NOTSET:
+            try:
+                await self.session.set_logging_level(
+                    next(level for level, value in LOG_LEVEL_MAPPING.items() if value == logger.level)
+                )
+            except Exception as exc:
+                logger.warning("Failed to set log level to %s", logger.level, exc_info=exc)
+
+    async def sampling_callback(
+        self, context: RequestContext[ClientSession, Any], params: types.CreateMessageRequestParams
+    ) -> types.CreateMessageResult | types.ErrorData:
+        """Callback function for sampling.
+
+        This function is called when the MCP server needs to get a message completed.
+
+        This is a simple version of this function, it can be overridden to allow more complex sampling.
+        It get's added to the session at initialization time, so overriding it is the best way to do this.
+        """
+        if not self.chat_client:
+            return types.ErrorData(
+                code=types.INTERNAL_ERROR,
+                message="No chat client available. Please set a chat client.",
+            )
+        logger.debug("Sampling callback called with params: %s", params)
+        messages: list[ChatMessage] = []
+        for msg in params.messages:
+            messages.append(_mcp_prompt_message_to_chat_message(msg))
+        try:
+            response = await self.chat_client.get_response(
+                messages,
+                temperature=params.temperature,
+                max_tokens=params.maxTokens,
+                stop=params.stopSequences,
+            )
+        except Exception as ex:
+            return types.ErrorData(
+                code=types.INTERNAL_ERROR,
+                message=f"Failed to get chat message content: {ex}",
+            )
+        if not response or not response.messages:
+            return types.ErrorData(
+                code=types.INTERNAL_ERROR,
+                message="Failed to get chat message content.",
+            )
+        mcp_contents = _chat_message_to_mcp_types(response.messages[0])
+        # grab the first content that is of type TextContent or ImageContent
+        mcp_content = next(
+            (content for content in mcp_contents if isinstance(content, (types.TextContent, types.ImageContent))),
+            None,
+        )
+        if not mcp_content:
+            return types.ErrorData(
+                code=types.INTERNAL_ERROR,
+                message="Failed to get right content types from the response.",
+            )
+        return types.CreateMessageResult(
+            role="assistant",
+            content=mcp_content,
+            model=response.ai_model_id or "unknown",
+        )
+
+    async def logging_callback(self, params: types.LoggingMessageNotificationParams) -> None:
+        """Callback function for logging.
+
+        This function is called when the MCP Server sends a log message.
+        By default it will log the message to the logger with the level set in the params.
+
+        Please subclass the MCP*Plugin and override this function if you want to adapt the behavior.
+        """
+        logger.log(LOG_LEVEL_MAPPING[params.level], params.data)
+
+    async def message_handler(
+        self,
+        message: RequestResponder[types.ServerRequest, types.ClientResult] | types.ServerNotification | Exception,
+    ) -> None:
+        """Handle messages from the MCP server.
+
+        By default this function will handle exceptions on the server, by logging those.
+
+        And it will trigger a reload of the tools and prompts when the list changed notification is received.
+
+        If you want to extend this behavior you can subclass the MCPPlugin and override this function,
+        if you want to keep the default behavior, make sure to call `super().message_handler(message)`.
+        """
+        if isinstance(message, Exception):
+            logger.error("Error from MCP server: %s", message, exc_info=message)
+            return
+        if isinstance(message, types.ServerNotification):
+            match message.root.method:
+                case "notifications/tools/list_changed":
+                    await self.load_tools()
+                case "notifications/prompts/list_changed":
+                    await self.load_prompts()
+                case _:
+                    logger.debug("Unhandled notification: %s", message.root.method)
+
+    async def load_prompts(self) -> None:
+        """Load prompts from the MCP server."""
+        if not self.session:
+            raise ToolExecutionException("MCP server not connected, please call connect() before using this method.")
+        try:
+            prompt_list = await self.session.list_prompts()
+        except Exception as exc:
+            logger.info(
+                "Prompt could not be loaded, you can exclude trying to load, by setting: load_prompts=False",
+                exc_info=exc,
+            )
+            prompt_list = None
+        for prompt in prompt_list.prompts if prompt_list else []:
+            local_name = _normalize_mcp_name(prompt.name)
+            input_model = _get_input_model_from_mcp_prompt(prompt)
+            func: AIFunction[BaseModel, list[ChatMessage]] = AIFunction(
+                func=partial(self.get_prompt, prompt.name),
+                name=local_name,
+                description=prompt.description or "",
+                input_model=input_model,
+            )
+            self.functions.append(func)
+
+    async def load_tools(self) -> None:
+        """Load tools from the MCP server."""
+        if not self.session:
+            raise ToolExecutionException("MCP server not connected, please call connect() before using this method.")
+        try:
+            tool_list = await self.session.list_tools()
+        except Exception as exc:
+            logger.info(
+                "Tools could not be loaded, you can exclude trying to load, by setting: load_tools=False",
+                exc_info=exc,
+            )
+            tool_list = None
+        for tool in tool_list.tools if tool_list else []:
+            local_name = _normalize_mcp_name(tool.name)
+            input_model = _get_input_model_from_mcp_tool(tool)
+            # Create AIFunctions out of each tool
+            func: AIFunction[BaseModel, list[AIContents]] = AIFunction(
+                func=partial(self.call_tool, tool.name),
+                name=local_name,
+                description=tool.description or "",
+                input_model=input_model,
+            )
+            self.functions.append(func)
+
+    async def close(self) -> None:
+        """Disconnect from the MCP server."""
+        await self._exit_stack.aclose()
+        self.session = None
+
+    @abstractmethod
+    def get_mcp_client(self) -> _AsyncGeneratorContextManager[Any, None]:
+        """Get an MCP client."""
+        pass
+
+    async def call_tool(self, tool_name: str, **kwargs: Any) -> list[AIContents]:
+        """Call a tool with the given arguments."""
+        if not self.session:
+            raise ToolExecutionException("MCP server not connected, please call connect() before using this method.")
+        if not self.load_tools_flag:
+            raise ToolExecutionException(
+                "Tools are not loaded for this server, please set load_tools=True in the constructor."
+            )
+        try:
+            return _mcp_call_tool_result_to_ai_contents(await self.session.call_tool(tool_name, arguments=kwargs))
+        except McpError as mcp_exc:
+            raise ToolExecutionException(mcp_exc.error.message, inner_exception=mcp_exc) from mcp_exc
+        except Exception as ex:
+            raise ToolExecutionException(f"Failed to call tool '{tool_name}'.", inner_exception=ex) from ex
+
+    async def get_prompt(self, prompt_name: str, **kwargs: Any) -> list[ChatMessage]:
+        """Call a prompt with the given arguments."""
+        if not self.session:
+            raise ToolExecutionException("MCP server not connected, please call connect() before using this method.")
+        if not self.load_prompts_flag:
+            raise ToolExecutionException(
+                "Prompts are not loaded for this server, please set load_prompts=True in the constructor."
+            )
+        try:
+            prompt_result = await self.session.get_prompt(prompt_name, arguments=kwargs)
+            return [_mcp_prompt_message_to_chat_message(message) for message in prompt_result.messages]
+        except McpError as mcp_exc:
+            raise ToolExecutionException(mcp_exc.error.message, inner_exception=mcp_exc) from mcp_exc
+        except Exception as ex:
+            raise ToolExecutionException(f"Failed to call prompt '{prompt_name}'.", inner_exception=ex) from ex
+
+    async def __aenter__(self) -> Self:
+        """Enter the context manager."""
+        try:
+            await self.connect()
+            return self
+        except ToolException:
+            raise
+        except Exception as ex:
+            await self._exit_stack.aclose()
+            raise ToolExecutionException("Failed to enter context manager.", inner_exception=ex) from ex
+
+    async def __aexit__(
+        self, exc_type: type[BaseException] | None, exc_value: BaseException | None, traceback: Any
+    ) -> None:
+        """Exit the context manager."""
+        await self.close()
+
+
+# region: MCP Plugin Implementations
+
+
+class McpStdioTool(McpTool):
+    """MCP stdio server configuration."""
+
+    def __init__(
+        self,
+        name: str,
+        command: str,
+        *,
+        load_tools: bool = True,
+        load_prompts: bool = True,
+        request_timeout: int | None = None,
+        session: ClientSession | None = None,
+        description: str | None = None,
+        additional_properties: dict[str, Any] | None = None,
+        args: list[str] | None = None,
+        env: dict[str, str] | None = None,
+        encoding: str | None = None,
+        chat_client: "ChatClient | None" = None,
+        **kwargs: Any,
+    ) -> None:
+        """Initialize the MCP stdio plugin.
+
+        The arguments are used to create a StdioServerParameters object.
+        Which is then used to create a stdio client.
+        see mcp.client.stdio.stdio_client and mcp.client.stdio.stdio_server_parameters
+        for more details.
+
+        Args:
+            name: The name of the plugin.
+            command: The command to run the MCP server.
+            load_tools: Whether to load tools from the MCP server.
+            load_prompts: Whether to load prompts from the MCP server.
+            request_timeout: The default timeout used for all requests.
+            session: The session to use for the MCP connection.
+            description: The description of the plugin.
+            additional_properties: Additional properties.
+            args: The arguments to pass to the command.
+            env: The environment variables to set for the command.
+            encoding: The encoding to use for the command output.
+            chat_client: The chat client to use for sampling.
+            kwargs: Any extra arguments to pass to the stdio client.
+
+        """
+        super().__init__(
+            name=name,
+            description=description,
+            additional_properties=additional_properties,
+            session=session,
+            chat_client=chat_client,
+            load_tools=load_tools,
+            load_prompts=load_prompts,
+            request_timeout=request_timeout,
+        )
+        self.command = command
+        self.args = args or []
+        self.env = env
+        self.encoding = encoding
+        self._client_kwargs = kwargs
+
+    def get_mcp_client(self) -> _AsyncGeneratorContextManager[Any, None]:
+        """Get an MCP stdio client."""
+        args: dict[str, Any] = {
+            "command": self.command,
+            "args": self.args,
+            "env": self.env,
+        }
+        if self.encoding:
+            args["encoding"] = self.encoding
+        if self._client_kwargs:
+            args.update(self._client_kwargs)
+        return stdio_client(server=StdioServerParameters(**args))
+
+
+class McpSseTools(McpTool):
+    """MCP sse server configuration."""
+
+    def __init__(
+        self,
+        name: str,
+        url: str,
+        *,
+        load_tools: bool = True,
+        load_prompts: bool = True,
+        request_timeout: int | None = None,
+        session: ClientSession | None = None,
+        description: str | None = None,
+        additional_properties: dict[str, Any] | None = None,
+        headers: dict[str, Any] | None = None,
+        timeout: float | None = None,
+        sse_read_timeout: float | None = None,
+        chat_client: "ChatClient | None" = None,
+        **kwargs: Any,
+    ) -> None:
+        """Initialize the MCP sse plugin.
+
+                The arguments are used to create a sse client.
+        see mcp.client.sse.sse_client for more details.
+
+        Any extra arguments passed to the constructor will be passed to the
+        sse client constructor.
+
+        Args:
+            name: The name of the plugin.
+            url: The URL of the MCP server.
+            load_tools: Whether to load tools from the MCP server.
+            load_prompts: Whether to load prompts from the MCP server.
+            request_timeout: The default timeout used for all requests.
+            session: The session to use for the MCP connection.
+            description: The description of the plugin.
+            additional_properties: Additional properties.
+            headers: The headers to send with the request.
+            timeout: The timeout for the request.
+            sse_read_timeout: The timeout for reading from the SSE stream.
+            chat_client: The chat client to use for sampling.
+            kwargs: Any extra arguments to pass to the sse client.
+
+        """
+        super().__init__(
+            name=name,
+            description=description,
+            additional_properties=additional_properties,
+            session=session,
+            chat_client=chat_client,
+            load_tools=load_tools,
+            load_prompts=load_prompts,
+            request_timeout=request_timeout,
+        )
+        self.url = url
+        self.headers = headers or {}
+        self.timeout = timeout
+        self.sse_read_timeout = sse_read_timeout
+        self._client_kwargs = kwargs
+
+    def get_mcp_client(self) -> _AsyncGeneratorContextManager[Any, None]:
+        """Get an MCP SSE client."""
+        args: dict[str, Any] = {
+            "url": self.url,
+        }
+        if self.headers:
+            args["headers"] = self.headers
+        if self.timeout is not None:
+            args["timeout"] = self.timeout
+        if self.sse_read_timeout is not None:
+            args["sse_read_timeout"] = self.sse_read_timeout
+        if self._client_kwargs:
+            args.update(self._client_kwargs)
+        return sse_client(**args)
+
+
+class McpStreamableHttpTool(McpTool):
+    """MCP streamable http server configuration."""
+
+    def __init__(
+        self,
+        name: str,
+        url: str,
+        *,
+        load_tools: bool = True,
+        load_prompts: bool = True,
+        request_timeout: int | None = None,
+        session: ClientSession | None = None,
+        description: str | None = None,
+        additional_properties: dict[str, Any] | None = None,
+        headers: dict[str, Any] | None = None,
+        timeout: float | None = None,
+        sse_read_timeout: float | None = None,
+        terminate_on_close: bool | None = None,
+        chat_client: "ChatClient | None" = None,
+        **kwargs: Any,
+    ) -> None:
+        """Initialize the MCP streamable http plugin.
+
+        The arguments are used to create a streamable http client.
+        see mcp.client.streamable_http.streamablehttp_client for more details.
+
+        Any extra arguments passed to the constructor will be passed to the
+        streamable http client constructor.
+
+        Args:
+            name: The name of the plugin.
+            url: The URL of the MCP server.
+            load_tools: Whether to load tools from the MCP server.
+            load_prompts: Whether to load prompts from the MCP server.
+            request_timeout: The default timeout used for all requests.
+            session: The session to use for the MCP connection.
+            description: The description of the plugin.
+            additional_properties: Additional properties.
+            headers: The headers to send with the request.
+            timeout: The timeout for the request.
+            sse_read_timeout: The timeout for reading from the SSE stream.
+            terminate_on_close: Close the transport when the MCP client is terminated.
+            chat_client: The chat client to use for sampling.
+            kwargs: Any extra arguments to pass to the sse client.
+        """
+        super().__init__(
+            name=name,
+            description=description,
+            additional_properties=additional_properties,
+            session=session,
+            chat_client=chat_client,
+            load_tools=load_tools,
+            load_prompts=load_prompts,
+            request_timeout=request_timeout,
+        )
+        self.url = url
+        self.headers = headers or {}
+        self.timeout = timeout
+        self.sse_read_timeout = sse_read_timeout
+        self.terminate_on_close = terminate_on_close
+        self._client_kwargs = kwargs
+
+    def get_mcp_client(self) -> _AsyncGeneratorContextManager[Any, None]:
+        """Get an MCP streamable http client."""
+        args: dict[str, Any] = {
+            "url": self.url,
+        }
+        if self.headers:
+            args["headers"] = self.headers
+        if self.timeout is not None:
+            args["timeout"] = self.timeout
+        if self.sse_read_timeout is not None:
+            args["sse_read_timeout"] = self.sse_read_timeout
+        if self.terminate_on_close is not None:
+            args["terminate_on_close"] = self.terminate_on_close
+        if self._client_kwargs:
+            args.update(self._client_kwargs)
+        return streamablehttp_client(**args)
+
+
+class McpWebsocketTool(McpTool):
+    """MCP websocket server configuration."""
+
+    def __init__(
+        self,
+        name: str,
+        url: str,
+        *,
+        load_tools: bool = True,
+        load_prompts: bool = True,
+        request_timeout: int | None = None,
+        session: ClientSession | None = None,
+        description: str | None = None,
+        additional_properties: dict[str, Any] | None = None,
+        chat_client: "ChatClient | None" = None,
+        **kwargs: Any,
+    ) -> None:
+        """Initialize the MCP websocket plugin.
+
+                The arguments are used to create a websocket client.
+        see mcp.client.websocket.websocket_client for more details.
+
+        Any extra arguments passed to the constructor will be passed to the
+        websocket client constructor.
+
+        Args:
+            name: The name of the plugin.
+            url: The URL of the MCP server.
+            load_tools: Whether to load tools from the MCP server.
+            load_prompts: Whether to load prompts from the MCP server.
+            request_timeout: The default timeout used for all requests.
+            session: The session to use for the MCP connection.
+            description: The description of the plugin.
+            additional_properties: Additional properties.
+            chat_client: The chat client to use for sampling.
+            kwargs: Any extra arguments to pass to the websocket client.
+
+        """
+        super().__init__(
+            name=name,
+            description=description,
+            additional_properties=additional_properties,
+            session=session,
+            chat_client=chat_client,
+            load_tools=load_tools,
+            load_prompts=load_prompts,
+            request_timeout=request_timeout,
+        )
+        self.url = url
+        self._client_kwargs = kwargs
+
+    def get_mcp_client(self) -> _AsyncGeneratorContextManager[Any, None]:
+        """Get an MCP websocket client."""
+        args: dict[str, Any] = {
+            "url": self.url,
+        }
+        if self._client_kwargs:
+            args.update(self._client_kwargs)
+        return websocket_client(**args)

python/packages/main/agent_framework/_types.py

@@ -26,6 +26,7 @@ from pydantic import (
     model_serializer,
 )
 
+from ._logging import get_logger
 from ._pydantic import AFBaseModel
 from ._tools import AITool, ai_function
 from .exceptions import AgentFrameworkException
@@ -35,9 +36,10 @@ if sys.version_info >= (3, 11):
 else:
     from typing_extensions import Self  # pragma: no cover
 
+logger = get_logger("agent_framework")
+
 # region Constants and types
 _T = TypeVar("_T")
-TValue = TypeVar("TValue")
 TEmbedding = TypeVar("TEmbedding")
 TChatResponse = TypeVar("TChatResponse", bound="ChatResponse")
 TChatToolMode = TypeVar("TChatToolMode", bound="ChatToolMode")
@@ -99,7 +101,6 @@ __all__ = [
     "HostedFileContent",
     "HostedVectorStoreContent",
     "SpeechToTextOptions",
-    "StructuredResponse",
     "TextContent",
     "TextReasoningContent",
     "TextSpanRegion",
@@ -1317,10 +1318,9 @@ class ChatResponse(AFBaseModel):
         created_at: A timestamp for the chat response.
         finish_reason: The reason for the chat response.
         usage_details: The usage details for the chat response.
+        structured_output: The structured output of the chat response, if applicable.
         additional_properties: Any additional properties associated with the chat response.
         raw_representation: The raw representation of the chat response from an underlying implementation.
-
-
     """
 
     messages: list[ChatMessage]
@@ -1338,6 +1338,8 @@ class ChatResponse(AFBaseModel):
     """The reason for the chat response."""
     usage_details: UsageDetails | None = None
     """The usage details for the chat response."""
+    value: Any | None = None
+    """The structured output of the chat response, if applicable."""
     additional_properties: dict[str, Any] | None = None
     """Any additional properties associated with the chat response."""
     raw_representation: Any | None = None
@@ -1354,6 +1356,8 @@ class ChatResponse(AFBaseModel):
         created_at: CreatedAtT | None = None,
         finish_reason: ChatFinishReason | None = None,
         usage_details: UsageDetails | None = None,
+        value: Any | None = None,
+        response_format: type[BaseModel] | None = None,
         additional_properties: dict[str, Any] | None = None,
         raw_representation: Any | None = None,
         **kwargs: Any,
@@ -1368,6 +1372,8 @@ class ChatResponse(AFBaseModel):
             created_at: Optional timestamp for the chat response.
             finish_reason: Optional reason for the chat response.
             usage_details: Optional usage details for the chat response.
+            value: Optional value of the structured output.
+            response_format: Optional response format for the chat response.
             messages: List of ChatMessage objects to include in the response.
             additional_properties: Optional additional properties associated with the chat response.
             raw_representation: Optional raw representation of the chat response from an underlying implementation.
@@ -1385,6 +1391,8 @@ class ChatResponse(AFBaseModel):
         created_at: CreatedAtT | None = None,
         finish_reason: ChatFinishReason | None = None,
         usage_details: UsageDetails | None = None,
+        value: Any | None = None,
+        response_format: type[BaseModel] | None = None,
         additional_properties: dict[str, Any] | None = None,
         raw_representation: Any | None = None,
         **kwargs: Any,
@@ -1399,6 +1407,8 @@ class ChatResponse(AFBaseModel):
             created_at: Optional timestamp for the chat response.
             finish_reason: Optional reason for the chat response.
             usage_details: Optional usage details for the chat response.
+            value: Optional value of the structured output.
+            response_format: Optional response format for the chat response.
             additional_properties: Optional additional properties associated with the chat response.
             raw_representation: Optional raw representation of the chat response from an underlying implementation.
             **kwargs: Any additional keyword arguments.
@@ -1416,6 +1426,8 @@ class ChatResponse(AFBaseModel):
         created_at: CreatedAtT | None = None,
         finish_reason: ChatFinishReason | None = None,
         usage_details: UsageDetails | None = None,
+        value: Any | None = None,
+        response_format: type[BaseModel] | None = None,
         additional_properties: dict[str, Any] | None = None,
         raw_representation: Any | None = None,
         **kwargs: Any,
@@ -1438,29 +1450,44 @@ class ChatResponse(AFBaseModel):
             created_at=created_at,  # type: ignore[reportCallIssue]
             finish_reason=finish_reason,  # type: ignore[reportCallIssue]
             usage_details=usage_details,  # type: ignore[reportCallIssue]
+            value=value,  # type: ignore[reportCallIssue]
             additional_properties=additional_properties,  # type: ignore[reportCallIssue]
             raw_representation=raw_representation,  # type: ignore[reportCallIssue]
             **kwargs,
         )
+        if response_format:
+            self.try_parse_value(output_format_type=response_format)
 
     @classmethod
-    def from_chat_response_updates(cls: type[TChatResponse], updates: Sequence["ChatResponseUpdate"]) -> TChatResponse:
+    def from_chat_response_updates(
+        cls: type[TChatResponse],
+        updates: Sequence["ChatResponseUpdate"],
+        *,
+        output_format_type: type[BaseModel] | None = None,
+    ) -> TChatResponse:
         """Joins multiple updates into a single ChatResponse."""
         msg = cls(messages=[])
         for update in updates:
             _process_update(msg, update)
         _finalize_response(msg)
+        if output_format_type:
+            msg.try_parse_value(output_format_type)
         return msg
 
     @classmethod
     async def from_chat_response_generator(
-        cls: type[TChatResponse], updates: AsyncIterable["ChatResponseUpdate"]
+        cls: type[TChatResponse],
+        updates: AsyncIterable["ChatResponseUpdate"],
+        *,
+        output_format_type: type[BaseModel] | None = None,
     ) -> TChatResponse:
         """Joins multiple updates into a single ChatResponse."""
         msg = cls(messages=[])
         async for update in updates:
             _process_update(msg, update)
         _finalize_response(msg)
+        if output_format_type:
+            msg.try_parse_value(output_format_type)
         return msg
 
     @property
@@ -1471,97 +1498,13 @@ class ChatResponse(AFBaseModel):
     def __str__(self) -> str:
         return self.text
 
-
-class StructuredResponse(ChatResponse, Generic[TValue]):
-    """Represents a structured response to a chat request.
-
-    Type Parameters:
-        TValue: The type of the value contained in the structured response.
-    """
-
-    value: TValue
-    """The result value of the chat response as an instance of `TValue`."""
-
-    @property
-    def text(self) -> str:
-        """Returns the concatenated text of all messages in the response."""
-        return "\n".join(message.text for message in self.messages)
-
-    @overload
-    def __init__(
-        self,
-        value: TValue,
-        *,
-        messages: ChatMessage | MutableSequence[ChatMessage],
-        response_id: str | None = None,
-        conversation_id: str | None = None,
-        model_id: str | None = None,
-        created_at: CreatedAtT | None = None,
-        finish_reason: ChatFinishReason | None = None,
-        usage_details: UsageDetails | None = None,
-        additional_properties: dict[str, Any] | None = None,
-        raw_representation: Any | None = None,
-        **kwargs: Any,
-    ) -> None:
-        """Initializes a StructuredResponse with the provided parameters."""
-
-    @overload
-    def __init__(
-        self,
-        value: TValue,
-        *,
-        text: TextContent | str,
-        response_id: str | None = None,
-        conversation_id: str | None = None,
-        model_id: str | None = None,
-        created_at: CreatedAtT | None = None,
-        finish_reason: ChatFinishReason | None = None,
-        usage_details: UsageDetails | None = None,
-        raw_representation: Any | None = None,
-        additional_properties: dict[str, Any] | None = None,
-        **kwargs: Any,
-    ) -> None:
-        """Initializes a StructuredResponse with the provided parameters."""
-
-    def __init__(
-        self,
-        value: TValue,
-        *,
-        messages: ChatMessage | MutableSequence[ChatMessage] | None = None,
-        text: TextContent | str | None = None,
-        response_id: str | None = None,
-        conversation_id: str | None = None,
-        model_id: str | None = None,
-        created_at: CreatedAtT | None = None,
-        finish_reason: ChatFinishReason | None = None,
-        usage_details: UsageDetails | None = None,
-        additional_properties: dict[str, Any] | None = None,
-        raw_representation: Any | None = None,
-        **kwargs: Any,
-    ) -> None:
-        """Initializes a StructuredResponse with the provided parameters."""
-        if messages is None:
-            messages = []
-        elif isinstance(messages, ChatMessage):
-            messages = [messages]
-        if text is not None:
-            if isinstance(text, str):
-                text = TextContent(text=text)
-            messages.append(ChatMessage(role=ChatRole.ASSISTANT, contents=[text]))
-
-        super().__init__(
-            value=value,
-            messages=messages,
-            conversation_id=conversation_id,
-            created_at=created_at,
-            finish_reason=finish_reason,
-            model_id=model_id,
-            response_id=response_id,
-            usage_details=usage_details,
-            additional_properties=additional_properties,
-            raw_representation=raw_representation,
-            **kwargs,
-        )
+    def try_parse_value(self, output_format_type: type[BaseModel]) -> None:
+        """If there is a value, does nothing, otherwise tries to parse the text into the value."""
+        if self.value is None:
+            try:
+                self.value = output_format_type.model_validate_json(self.text)  # type: ignore[reportUnknownMemberType]
+            except ValidationError as ex:
+                logger.debug("Failed to parse value from chat response text: %s", ex)
 
 
 # region ChatResponseUpdate

python/packages/main/agent_framework/openai/_assistants_client.py

@@ -144,7 +144,8 @@ class OpenAIAssistantsClient(OpenAIConfigBase, ChatClientBase):
         **kwargs: Any,
     ) -> ChatResponse:
         return await ChatResponse.from_chat_response_generator(
-            updates=self._inner_get_streaming_response(messages=messages, chat_options=chat_options, **kwargs)
+            updates=self._inner_get_streaming_response(messages=messages, chat_options=chat_options, **kwargs),
+            output_format_type=chat_options.response_format,
         )
 
     async def _inner_get_streaming_response(

python/packages/main/agent_framework/openai/_chat_client.py

@@ -61,7 +61,9 @@ class OpenAIChatClientBase(OpenAIHandler, ChatClientBase):
     ) -> ChatResponse:
         options_dict = self._prepare_options(messages, chat_options)
         try:
-            return self._create_chat_response(await self.client.chat.completions.create(stream=False, **options_dict))
+            return self._create_chat_response(
+                await self.client.chat.completions.create(stream=False, **options_dict), chat_options
+            )
         except BadRequestError as ex:
             if ex.code == "content_filter":
                 raise OpenAIContentFilterException(
@@ -143,7 +145,7 @@ class OpenAIChatClientBase(OpenAIHandler, ChatClientBase):
             options_dict["response_format"] = type_to_response_format_param(chat_options.response_format)
         return options_dict
 
-    def _create_chat_response(self, response: ChatCompletion) -> "ChatResponse":
+    def _create_chat_response(self, response: ChatCompletion, chat_options: ChatOptions) -> "ChatResponse":
         """Create a chat message content object from a choice."""
         response_metadata = self._get_metadata_from_chat_response(response)
         messages: list[ChatMessage] = []
@@ -166,6 +168,7 @@ class OpenAIChatClientBase(OpenAIHandler, ChatClientBase):
             model_id=response.model,
             additional_properties=response_metadata,
             finish_reason=finish_reason,
+            response_format=chat_options.response_format,
         )
 
     def _create_chat_response_update(

python/packages/main/agent_framework/openai/_responses_client.py

@@ -44,7 +44,6 @@ from .._types import (
     FunctionCallContent,
     FunctionResultContent,
     HostedFileContent,
-    StructuredResponse,
     TextContent,
     TextSpanRegion,
     UsageDetails,
@@ -605,7 +604,8 @@ class OpenAIResponsesClientBase(OpenAIHandler, ChatClientBase):
             args["usage_details"] = usage_details
         if structured_response:
             args["value"] = structured_response
-            return StructuredResponse(**args)
+        elif chat_options.response_format:
+            args["response_format"] = chat_options.response_format
         return ChatResponse(**args)
 
     def _create_streaming_response_content(