Expose model preferences in ctx.sample

This commit is contained in:
davenpi 2025-05-21 19:34:45 -04:00
commit 9122833998
3 changed files with 76 additions and 3 deletions

View file

@ -228,8 +228,8 @@ async def analyze_sentiment(text: str, ctx: Context) -> dict:
# Create a sampling prompt asking for sentiment analysis
prompt = f"Analyze the sentiment of the following text as positive, negative, or neutral. Just output a single word - 'positive', 'negative', or 'neutral'. Text to analyze: {text}"
# Send the sampling request to the client's LLM
response = await ctx.sample(prompt)
# Send the sampling request to the clients LLM (provide a hint for the model you want to use)
response = await ctx.sample(prompt, model_preferences="claude-3-sonnet")
# Process the LLM's response
sentiment = response.text.strip().lower()
@ -247,11 +247,12 @@ async def analyze_sentiment(text: str, ctx: Context) -> dict:
**Method signature:**
- **`ctx.sample(messages: str | list[str | SamplingMessage], system_prompt: str | None = None, temperature: float | None = None, max_tokens: int | None = None) -> TextContent | ImageContent`**
- **`ctx.sample(messages: str | list[str | SamplingMessage], system_prompt: str | None = None, temperature: float | None = None, max_tokens: int | None = None, model_preferences: ModelPreferences | str | list[str] | None = None) -> TextContent | ImageContent`**
- `messages`: A string or list of strings/message objects to send to the LLM
- `system_prompt`: Optional system prompt to guide the LLM's behavior
- `temperature`: Optional sampling temperature (controls randomness)
- `max_tokens`: Optional maximum number of tokens to generate (defaults to 512)
- `model_preferences`: Optional model selection preferences (e.g., a model hint string, list of hints, or a ModelPreferences object)
- Returns the LLM's response as TextContent or ImageContent
When providing a simple string, it's treated as a user message. For more complex scenarios, you can provide a list of messages with different roles.

View file

@ -12,6 +12,8 @@ from mcp.shared.context import RequestContext
from mcp.types import (
CreateMessageResult,
ImageContent,
ModelHint,
ModelPreferences,
Root,
SamplingMessage,
TextContent,
@ -200,6 +202,7 @@ class Context:
system_prompt: str | None = None,
temperature: float | None = None,
max_tokens: int | None = None,
model_preferences: ModelPreferences | str | list[str] | None = None,
) -> TextContent | ImageContent:
"""
Send a sampling request to the client and await the response.
@ -231,6 +234,7 @@ class Context:
system_prompt=system_prompt,
temperature=temperature,
max_tokens=max_tokens,
model_preferences=self._parse_model_preferences(model_preferences),
)
return result.content
@ -248,3 +252,42 @@ class Context:
)
return fastmcp.server.dependencies.get_http_request()
def _parse_model_preferences(self, model_preferences) -> ModelPreferences | None:
"""
Validates and converts user input for model_preferences into a ModelPreferences object.
Args:
model_preferences (ModelPreferences | str | list[str] | None):
The model preferences to use. Accepts:
- ModelPreferences (returns as-is)
- str (single model hint)
- list[str] (multiple model hints)
- None (no preferences)
Returns:
ModelPreferences | None: The parsed ModelPreferences object, or None if not provided.
Raises:
ValueError: If the input is not a supported type or contains invalid values.
"""
if model_preferences is None:
return None
if isinstance(model_preferences, ModelPreferences):
return model_preferences
if isinstance(model_preferences, str):
# Single model hint
return ModelPreferences(hints=[ModelHint(name=model_preferences)])
if isinstance(model_preferences, list):
# List of model hints (strings)
if not all(isinstance(h, str) for h in model_preferences):
raise ValueError(
"All elements of model_preferences list must be"
" strings (model name hints)."
)
return ModelPreferences(
hints=[ModelHint(name=h) for h in model_preferences]
)
raise ValueError(
"model_preferences must be one of: ModelPreferences, str, list[str], or None."
)

View file

@ -2,9 +2,11 @@ import warnings
from unittest.mock import MagicMock, patch
import pytest
from mcp.types import ModelPreferences
from starlette.requests import Request
from fastmcp.server.context import Context
from fastmcp.server.server import FastMCP
class TestContextDeprecations:
@ -57,3 +59,30 @@ class TestContextDeprecations:
assert "https://gofastmcp.com/patterns/http-requests" in str(
warning.message
)
@pytest.fixture
def context():
return Context(fastmcp=FastMCP())
class TestParseModelPreferences:
def test_parse_model_preferences_string(self, context):
mp = context._parse_model_preferences("claude-3-sonnet")
assert isinstance(mp, ModelPreferences)
assert mp.hints is not None
assert mp.hints[0].name == "claude-3-sonnet"
def test_parse_model_preferences_list(self, context):
mp = context._parse_model_preferences(["claude-3-sonnet", "claude"])
assert isinstance(mp, ModelPreferences)
assert mp.hints is not None
assert [h.name for h in mp.hints] == ["claude-3-sonnet", "claude"]
def test_parse_model_preferences_object(self, context):
obj = ModelPreferences(hints=[])
assert context._parse_model_preferences(obj) is obj
def test_parse_model_preferences_invalid_type(self, context):
with pytest.raises(ValueError):
context._parse_model_preferences(123)