adds model preference support for mcp server sampling requests (microsoft#373)

Example:

from mcp.types import ModelPreferences
from mcp_extensions import send_sampling_request

    # Set model preferences
    model_preferences = ModelPreferences(
        # Can use hints to prefer models with specific names
        # hints=[
        #     ModelHint(
# # Prefer models where name starts with `name` value (so `o3` would
_include_ `o3-mini`)')
        #         name="gpt-4o",
        #     )
        # ],
# Setting speed priority to 1 to choose a faster model, like gpt-4o
        speedPriority=1,
# If needing a reasoning model, set intelligence priority to 1 instead
        # intelligencePriority=1,
    )
    
    sampling_result = await send_sampling_request(
            fastmcp_server_context=ctx,
            system_prompt="<your prompt>",
            messages=messages,
            model_preferences=model_preferences,
            max_tokens=1024,
        )

Author: bkrabach

assistants/codespace-assistant/assistant/response/response.py

@@ -52,16 +52,27 @@ async def error_handler(server_config: MCPServerConfig, error: Exception) -> Non
                 )
             )
 
+        # Get the AI client configurations for this assistant
+        generative_ai_client_config = get_ai_client_configs(config, "generative")
+        reasoning_ai_client_config = get_ai_client_configs(config, "reasoning")
+
         # TODO: This is a temporary hack to allow directing the request to the reasoning model
+        # Currently we will only use the requested AI client configuration for the turn
         request_type = "reasoning" if message.content.startswith("reason:") else "generative"
-
-        # Get the AI client configuration based on the request type
-        request_config, service_config = get_ai_client_configs(config, request_type)
+        # Set a default AI client configuration based on the request type
+        default_ai_client_config = (
+            reasoning_ai_client_config if request_type == "reasoning" else generative_ai_client_config
+        )
+        # Set the service and request configurations for the AI client
+        service_config = default_ai_client_config.service_config
+        request_config = default_ai_client_config.request_config
 
         # Create a sampling handler for handling requests from the MCP servers
         sampling_handler = OpenAISamplingHandler(
-            service_config=service_config,
-            request_config=request_config,
+            ai_client_configs=[
+                generative_ai_client_config,
+                reasoning_ai_client_config,
+            ]
         )
 
         mcp_sessions = await establish_mcp_sessions(

assistants/codespace-assistant/assistant/response/utils/openai_utils.py

@@ -2,8 +2,9 @@
 
 import logging
 from textwrap import dedent
-from typing import List, Literal, Tuple
+from typing import List, Literal, Tuple, Union
 
+from assistant_extensions.ai_clients.config import AzureOpenAIClientConfigModel, OpenAIClientConfigModel
 from assistant_extensions.mcp import (
     ExtendedCallToolRequestParams,
     MCPSession,
@@ -28,11 +29,32 @@
 
 def get_ai_client_configs(
     config: AssistantConfigModel, request_type: Literal["generative", "reasoning"] = "generative"
-) -> tuple[OpenAIRequestConfig, AzureOpenAIServiceConfig | OpenAIServiceConfig]:
-    if request_type == "reasoning":
-        return config.reasoning_ai_client_config.request_config, config.reasoning_ai_client_config.service_config
+) -> Union[AzureOpenAIClientConfigModel, OpenAIClientConfigModel]:
+    def create_ai_client_config(
+        service_config: AzureOpenAIServiceConfig | OpenAIServiceConfig,
+        request_config: OpenAIRequestConfig,
+    ) -> AzureOpenAIClientConfigModel | OpenAIClientConfigModel:
+        if isinstance(service_config, AzureOpenAIServiceConfig):
+            return AzureOpenAIClientConfigModel(
+                service_config=service_config,
+                request_config=request_config,
+            )
+
+        return OpenAIClientConfigModel(
+            service_config=service_config,
+            request_config=request_config,
+        )
 
-    return config.generative_ai_client_config.request_config, config.generative_ai_client_config.service_config
+    if request_type == "reasoning":
+        return create_ai_client_config(
+            config.reasoning_ai_client_config.service_config,
+            config.reasoning_ai_client_config.request_config,
+        )
+
+    return create_ai_client_config(
+        config.generative_ai_client_config.service_config,
+        config.generative_ai_client_config.request_config,
+    )
 
 
 async def get_completion(

libraries/python/assistant-extensions/assistant_extensions/mcp/_openai_utils.py

@@ -4,7 +4,13 @@
 import deepmerge
 from mcp import ClientSession, CreateMessageResult, SamplingMessage
 from mcp.shared.context import RequestContext
-from mcp.types import CreateMessageRequestParams, ErrorData, ImageContent, TextContent
+from mcp.types import (
+    CreateMessageRequestParams,
+    ErrorData,
+    ImageContent,
+    ModelPreferences,
+    TextContent,
+)
 from openai.types.chat import (
     ChatCompletion,
     ChatCompletionAssistantMessageParam,
@@ -14,8 +20,9 @@
     ChatCompletionToolParam,
     ChatCompletionUserMessageParam,
 )
-from openai_client import OpenAIRequestConfig, ServiceConfig, create_client
+from openai_client import OpenAIRequestConfig, create_client
 
+from ..ai_clients.config import AzureOpenAIClientConfigModel, OpenAIClientConfigModel
 from ._model import MCPSamplingMessageHandler
 from ._sampling_handler import SamplingHandler
 
@@ -40,14 +47,14 @@ def message_handler(self) -> MCPSamplingMessageHandler:
 
     def __init__(
         self,
-        service_config: ServiceConfig | None = None,
-        request_config: OpenAIRequestConfig | None = None,
+        ai_client_configs: list[
+            Union[AzureOpenAIClientConfigModel, OpenAIClientConfigModel]
+        ],
         assistant_mcp_tools: list[ChatCompletionToolParam] | None = None,
         message_processor: OpenAIMessageProcessor | None = None,
         handler: MCPSamplingMessageHandler | None = None,
     ) -> None:
-        self.service_config = service_config
-        self.request_config = request_config
+        self.ai_client_configs = ai_client_configs
         self.assistant_mcp_tools = assistant_mcp_tools
 
         # set a default message processor that converts sampling messages to
@@ -82,27 +89,21 @@ async def _default_message_handler(
     ) -> CreateMessageResult | ErrorData:
         logger.info(f"Sampling handler invoked with context: {context}")
 
-        if not self.service_config or not self.request_config:
-            raise ValueError(
-                "Service config and request config must be set before handling messages."
-            )
+        ai_client_config = self._ai_client_config_from_model_preferences(
+            params.modelPreferences
+        )
 
-        try:
-            completion_args = await self._create_completion_request(
-                request=params,
-                request_config=self.request_config,
-                template_processor=self.message_processor,
-            )
-        except Exception as e:
-            logger.exception(f"Error creating completion request: {e}")
-            return ErrorData(
-                code=500,
-                message="Error creating completion request.",
-                data=e,
-            )
+        if not ai_client_config:
+            raise ValueError("No AI client configs defined for sampling requests.")
+
+        completion_args = await self._create_completion_request(
+            request=params,
+            request_config=ai_client_config.request_config,
+            template_processor=self.message_processor,
+        )
 
         completion: ChatCompletion | None = None
-        async with create_client(self.service_config) as client:
+        async with create_client(ai_client_config.service_config) as client:
             completion = await client.chat.completions.create(**completion_args)
 
         if completion is None:
@@ -112,12 +113,6 @@ async def _default_message_handler(
             )
 
         choice = completion.choices[0]
-        if choice.message.content is None:
-            return ErrorData(
-                code=500,
-                message="No content returned from completion choice.",
-            )
-
         content = choice.message.content
         if content is None:
             content = "[no content]"
@@ -141,7 +136,61 @@ async def handle_message(
         context: RequestContext[ClientSession, Any],
         params: CreateMessageRequestParams,
     ) -> CreateMessageResult | ErrorData:
-        return await self._message_handler(context, params)
+        try:
+            return await self._message_handler(context, params)
+        except Exception as e:
+            logger.error(f"Error handling sampling request: {e}")
+            code = getattr(e, "status_code", 500)
+            message = getattr(e, "message", "Error handling sampling request.")
+            data = str(e)
+            return ErrorData(code=code, message=message, data=data)
+
+    def _ai_client_config_from_model_preferences(
+        self, model_preferences: ModelPreferences | None
+    ) -> Union[AzureOpenAIClientConfigModel, OpenAIClientConfigModel] | None:
+        """
+        Returns an AI client config from model preferences.
+        """
+
+        # if no configs are provided, return None
+        if not self.ai_client_configs or len(self.ai_client_configs) == 0:
+            return None
+
+        # if not provided, return the first config
+        if not model_preferences:
+            return self.ai_client_configs[0]
+
+        # if hints are provided, return the first hint where the name value matches
+        # the start of the model name
+        if model_preferences.hints:
+            for hint in model_preferences.hints:
+                if not hint.name:
+                    continue
+                for ai_client_config in self.ai_client_configs:
+                    if ai_client_config.request_config.model.startswith(hint.name):
+                        return ai_client_config
+
+        # if any of the priority values are set, return the first config that matches our
+        # criteria: speedPriority equates to non-reasoning models, intelligencePriority
+        # equates to reasoning models for now
+        # note: we are ignoring costPriority for now
+        speed_priority = model_preferences.speedPriority or 0
+        intelligence_priority = model_preferences.intelligencePriority or 0
+        # cost_priority = 0 # ignored for now
+
+        # later we will support more than just reasoning or non-reasoning choices, but
+        # for now we can keep it simple
+        use_reasoning_model = intelligence_priority > speed_priority
+
+        for ai_client_config in self.ai_client_configs:
+            if (
+                ai_client_config.request_config.is_reasoning_model
+                == use_reasoning_model
+            ):
+                return ai_client_config
+
+        # failing to find a config via preferences, return first config
+        return self.ai_client_configs[0]
 
     async def _create_completion_request(
         self,

mcp-servers/mcp-server-giphy/mcp_server/sampling.py

@@ -8,7 +8,7 @@
 from typing import Any, Dict, List, Union
 
 from mcp.server.fastmcp import Context
-from mcp.types import ImageContent, SamplingMessage, TextContent
+from mcp.types import ImageContent, ModelPreferences, SamplingMessage, TextContent
 from mcp_extensions import send_sampling_request, send_tool_call_progress
 
 from .utils import fetch_url
@@ -74,8 +74,8 @@ async def generate_sampling_messages(search_results: List[Dict]) -> List[Samplin
         text_content = get_text_content(result)
         if text_content is not None:
             messages.append(SamplingMessage(role="user", content=text_content))
-        if image_content is not None:
-            messages.append(SamplingMessage(role="user", content=image_content))
+        # if image_content is not None:
+        #     messages.append(SamplingMessage(role="user", content=image_content))
     return messages
 
 
@@ -108,19 +108,39 @@ async def perform_sampling(
     # Generate sampling messages
     messages += await generate_sampling_messages(search_results)
 
+    # Set model preferences
+    model_preferences = ModelPreferences(
+        # Can use hints to prefer models with specific names
+        # hints=[
+        #     ModelHint(
+        #         # Prefer models where name starts with `name` value (so `o3` would _include_ `o3-mini`)')
+        #         name="gpt-4o",
+        #     )
+        # ],
+        # Setting speed priority to 1 to choose a faster model, like gpt-4o
+        speedPriority=1,
+        # If needing a reasoning model, set intelligence priority to 1 instead
+        # intelligencePriority=1,
+    )
+
     await send_tool_call_progress(ctx, "choosing image...")
 
     # FIXME add support for structured output to enforce image selection
     # Send sampling request to FastMCP server
-    sampling_result = await send_sampling_request(
-        fastmcp_server_context=ctx,
-        system_prompt=dedent(f"""
-            Analyze these images and choose the best choice based on provided context.
-            Context: {context}
-            Return the url for the chosen image.
-        """).strip(),
-        messages=messages,
-        max_tokens=100,
-    )
-
-    return sampling_result.content
+    try:
+        sampling_result = await send_sampling_request(
+            fastmcp_server_context=ctx,
+            system_prompt=dedent(f"""
+                Analyze these images and choose the best choice based on provided context.
+                Context: {context}
+                Return the url for the chosen image.
+            """).strip(),
+            messages=messages,
+            model_preferences=model_preferences,
+            max_tokens=100,
+        )
+
+        return sampling_result.content
+    except Exception as e:
+        logger.error(f"Failed to perform sampling: {str(e)}")
+        raise e