-
Notifications
You must be signed in to change notification settings - Fork 2.8k
Implement SEP-1577 - Sampling With Tools #1594
New issue
Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.
By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.
Already on GitHub? Sign in to your account
Open
ochafik
wants to merge
9
commits into
main
Choose a base branch
from
ochafik/sep-1577-sampling-tools
base: main
Could not load branches
Branch not found: {{ refName }}
Loading
Could not load tags
Nothing to show
Loading
Are you sure you want to change the base?
Some commits from the old base branch may be removed from the timeline,
and old review comments may become outdated.
+412
−21
Open
Changes from all commits
Commits
Show all changes
9 commits
Select commit
Hold shift + click to select a range
4ded25a
Implement SEP-1577: Sampling With Tools
ochafik b74bff2
Rename sampling capability types for clarity
ochafik 224fc22
Simplify sampling types to match draft spec
ochafik 03de21d
Merge branch 'main' into ochafik/sep-1577-sampling-tools
ochafik 475df1f
Merge branch 'main' into ochafik/sep-1577-sampling-tools
ochafik 2d24c3a
Merge branch 'main' into ochafik/sep-1577-sampling-tools
ochafik fd3e999
Fix ruff formatting for SamplingMessageContentBlock type alias
felixweinberger 3459880
Fix pyright type errors for CreateMessageResult.content
felixweinberger 607976e
Update README snippets
felixweinberger File filter
Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
There are no files selected for viewing
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -1,5 +1,5 @@ | ||
| from collections.abc import Callable | ||
| from typing import Annotated, Any, Generic, Literal, TypeAlias, TypeVar | ||
| from typing import Annotated, Any, Generic, Literal, TypeAlias, TypeVar, Union | ||
|
|
||
| from pydantic import BaseModel, ConfigDict, Field, FileUrl, RootModel | ||
| from pydantic.networks import AnyUrl, UrlConstraints | ||
|
|
@@ -250,8 +250,24 @@ class RootsCapability(BaseModel): | |
| model_config = ConfigDict(extra="allow") | ||
|
|
||
|
|
||
| class SamplingCapability(BaseModel): | ||
| """Capability for sampling operations.""" | ||
| class SamplingContextCapability(BaseModel): | ||
| """ | ||
| Capability for context inclusion during sampling. | ||
|
|
||
| Indicates support for non-'none' values in the includeContext parameter. | ||
| SOFT-DEPRECATED: New implementations should use tools parameter instead. | ||
| """ | ||
|
|
||
| model_config = ConfigDict(extra="allow") | ||
|
|
||
|
|
||
| class SamplingToolsCapability(BaseModel): | ||
| """ | ||
| Capability indicating support for tool calling during sampling. | ||
|
|
||
| When present in ClientCapabilities.sampling, indicates that the client | ||
| supports the tools and toolChoice parameters in sampling requests. | ||
| """ | ||
|
|
||
| model_config = ConfigDict(extra="allow") | ||
|
|
||
|
|
@@ -262,13 +278,34 @@ class ElicitationCapability(BaseModel): | |
| model_config = ConfigDict(extra="allow") | ||
|
|
||
|
|
||
| class SamplingCapability(BaseModel): | ||
| """ | ||
| Sampling capability structure, allowing fine-grained capability advertisement. | ||
| """ | ||
|
|
||
| context: SamplingContextCapability | None = None | ||
| """ | ||
| Present if the client supports non-'none' values for includeContext parameter. | ||
| SOFT-DEPRECATED: New implementations should use tools parameter instead. | ||
| """ | ||
| tools: SamplingToolsCapability | None = None | ||
| """ | ||
| Present if the client supports tools and toolChoice parameters in sampling requests. | ||
| Presence indicates full tool calling support during sampling. | ||
| """ | ||
| model_config = ConfigDict(extra="allow") | ||
|
|
||
|
|
||
| class ClientCapabilities(BaseModel): | ||
| """Capabilities a client may support.""" | ||
|
|
||
| experimental: dict[str, dict[str, Any]] | None = None | ||
| """Experimental, non-standard capabilities that the client supports.""" | ||
| sampling: SamplingCapability | None = None | ||
| """Present if the client supports sampling from an LLM.""" | ||
| """ | ||
| Present if the client supports sampling from an LLM. | ||
| Can contain fine-grained capabilities like context and tools support. | ||
| """ | ||
| elicitation: ElicitationCapability | None = None | ||
| """Present if the client supports elicitation from the user.""" | ||
| roots: RootsCapability | None = None | ||
|
|
@@ -742,11 +779,89 @@ class AudioContent(BaseModel): | |
| model_config = ConfigDict(extra="allow") | ||
|
|
||
|
|
||
| class ToolUseContent(BaseModel): | ||
| """ | ||
| Content representing an assistant's request to invoke a tool. | ||
|
|
||
| This content type appears in assistant messages when the LLM wants to call a tool | ||
| during sampling. The server should execute the tool and return a ToolResultContent | ||
| in the next user message. | ||
| """ | ||
|
|
||
| type: Literal["tool_use"] | ||
| """Discriminator for tool use content.""" | ||
|
|
||
| name: str | ||
| """The name of the tool to invoke. Must match a tool name from the request's tools array.""" | ||
|
|
||
| id: str | ||
| """Unique identifier for this tool call, used to correlate with ToolResultContent.""" | ||
|
|
||
| input: dict[str, Any] | ||
| """Arguments to pass to the tool. Must conform to the tool's inputSchema.""" | ||
|
|
||
| meta: dict[str, Any] | None = Field(alias="_meta", default=None) | ||
| """ | ||
| See [MCP specification](https://github.com/modelcontextprotocol/modelcontextprotocol/blob/47339c03c143bb4ec01a26e721a1b8fe66634ebe/docs/specification/draft/basic/index.mdx#general-fields) | ||
| for notes on _meta usage. | ||
| """ | ||
| model_config = ConfigDict(extra="allow") | ||
|
|
||
|
|
||
| class ToolResultContent(BaseModel): | ||
| """ | ||
| Content representing the result of a tool execution. | ||
|
|
||
| This content type appears in user messages as a response to a ToolUseContent | ||
| from the assistant. It contains the output of executing the requested tool. | ||
| """ | ||
|
|
||
| type: Literal["tool_result"] | ||
| """Discriminator for tool result content.""" | ||
|
|
||
| toolUseId: str | ||
| """The unique identifier that corresponds to the tool call's id field.""" | ||
|
|
||
| content: list[Union[TextContent, ImageContent, AudioContent, "ResourceLink", "EmbeddedResource"]] = [] | ||
|
Contributor
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. I think we can just use |
||
| """ | ||
| A list of content objects representing the tool result. | ||
| Defaults to empty list if not provided. | ||
| """ | ||
|
|
||
| structuredContent: dict[str, Any] | None = None | ||
| """ | ||
| Optional structured tool output that matches the tool's outputSchema (if defined). | ||
| """ | ||
|
|
||
| isError: bool | None = None | ||
| """Whether the tool execution resulted in an error.""" | ||
|
|
||
| meta: dict[str, Any] | None = Field(alias="_meta", default=None) | ||
| """ | ||
| See [MCP specification](https://github.com/modelcontextprotocol/modelcontextprotocol/blob/47339c03c143bb4ec01a26e721a1b8fe66634ebe/docs/specification/draft/basic/index.mdx#general-fields) | ||
| for notes on _meta usage. | ||
| """ | ||
| model_config = ConfigDict(extra="allow") | ||
|
|
||
|
|
||
| SamplingMessageContentBlock: TypeAlias = TextContent | ImageContent | AudioContent | ToolUseContent | ToolResultContent | ||
| """Content block types allowed in sampling messages.""" | ||
|
|
||
|
|
||
| class SamplingMessage(BaseModel): | ||
| """Describes a message issued to or received from an LLM API.""" | ||
|
|
||
| role: Role | ||
| content: TextContent | ImageContent | AudioContent | ||
| content: SamplingMessageContentBlock | list[SamplingMessageContentBlock] | ||
| """ | ||
| Message content. Can be a single content block or an array of content blocks | ||
| for multi-modal messages and tool interactions. | ||
| """ | ||
| meta: dict[str, Any] | None = Field(alias="_meta", default=None) | ||
| """ | ||
| See [MCP specification](https://github.com/modelcontextprotocol/modelcontextprotocol/blob/47339c03c143bb4ec01a26e721a1b8fe66634ebe/docs/specification/draft/basic/index.mdx#general-fields) | ||
| for notes on _meta usage. | ||
| """ | ||
| model_config = ConfigDict(extra="allow") | ||
|
|
||
|
|
||
|
|
@@ -1035,6 +1150,25 @@ class ModelPreferences(BaseModel): | |
| model_config = ConfigDict(extra="allow") | ||
|
|
||
|
|
||
| class ToolChoice(BaseModel): | ||
| """ | ||
| Controls tool usage behavior during sampling. | ||
|
|
||
| Allows the server to specify whether and how the LLM should use tools | ||
| in its response. | ||
| """ | ||
|
|
||
| mode: Literal["auto", "required", "none"] | None = None | ||
| """ | ||
| Controls when tools are used: | ||
| - "auto": Model decides whether to use tools (default) | ||
| - "required": Model MUST use at least one tool before completing | ||
| - "none": Model should not use tools | ||
| """ | ||
|
|
||
| model_config = ConfigDict(extra="allow") | ||
|
|
||
|
|
||
| class CreateMessageRequestParams(RequestParams): | ||
| """Parameters for creating a message.""" | ||
|
|
||
|
|
@@ -1057,6 +1191,16 @@ class CreateMessageRequestParams(RequestParams): | |
| stopSequences: list[str] | None = None | ||
| metadata: dict[str, Any] | None = None | ||
| """Optional metadata to pass through to the LLM provider.""" | ||
| tools: list["Tool"] | None = None | ||
| """ | ||
| Tool definitions for the LLM to use during sampling. | ||
| Requires clientCapabilities.sampling.tools to be present. | ||
| """ | ||
| toolChoice: ToolChoice | None = None | ||
| """ | ||
| Controls tool usage behavior. | ||
| Requires clientCapabilities.sampling.tools and the tools parameter to be present. | ||
| """ | ||
| model_config = ConfigDict(extra="allow") | ||
|
|
||
|
|
||
|
|
@@ -1067,18 +1211,26 @@ class CreateMessageRequest(Request[CreateMessageRequestParams, Literal["sampling | |
| params: CreateMessageRequestParams | ||
|
|
||
|
|
||
| StopReason = Literal["endTurn", "stopSequence", "maxTokens"] | str | ||
| StopReason = Literal["endTurn", "stopSequence", "maxTokens", "toolUse"] | str | ||
|
|
||
|
|
||
| class CreateMessageResult(Result): | ||
| """The client's response to a sampling/create_message request from the server.""" | ||
|
|
||
| role: Role | ||
| content: TextContent | ImageContent | AudioContent | ||
| """The role of the message sender (typically 'assistant' for LLM responses).""" | ||
| content: SamplingMessageContentBlock | list[SamplingMessageContentBlock] | ||
| """ | ||
| Response content. May be a single content block or an array. | ||
| May include ToolUseContent if stopReason is 'toolUse'. | ||
| """ | ||
| model: str | ||
| """The name of the model that generated the message.""" | ||
| stopReason: StopReason | None = None | ||
| """The reason why sampling stopped, if known.""" | ||
| """ | ||
| The reason why sampling stopped, if known. | ||
| 'toolUse' indicates the model wants to use a tool. | ||
| """ | ||
|
|
||
|
|
||
| class ResourceTemplateReference(BaseModel): | ||
|
|
||
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Oops, something went wrong.
Add this suggestion to a batch that can be applied as a single commit.
This suggestion is invalid because no changes were made to the code.
Suggestions cannot be applied while the pull request is closed.
Suggestions cannot be applied while viewing a subset of changes.
Only one suggestion per line can be applied in a batch.
Add this suggestion to a batch that can be applied as a single commit.
Applying suggestions on deleted lines is not supported.
You must change the existing code in this line in order to create a valid suggestion.
Outdated suggestions cannot be applied.
This suggestion has been applied or marked resolved.
Suggestions cannot be applied from pending reviews.
Suggestions cannot be applied on multi-line comments.
Suggestions cannot be applied while the pull request is queued to merge.
Suggestion cannot be applied right now. Please check back later.
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
I don't think we want this as the example of how to use the new multiple content blocks :)
(same comment on other places where this pattern shows up)