Skip to content

Commit 3c287c8

Browse files
ochafikochafik
committed
add sampling backfill example
1 parent 856d9ec commit 3c287c8

File tree

3 files changed

+259
-0
lines changed

3 files changed

+259
-0
lines changed

‎package-lock.json‎

Lines changed: 22 additions & 0 deletions
Some generated files are not rendered by default. Learn more about customizing how changed files appear on GitHub.

‎package.json‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -76,6 +76,7 @@
7676
"zod-to-json-schema": "^3.24.1"
7777
},
7878
"devDependencies": {
79+
"@anthropic-ai/sdk": "^0.65.0",
7980
"@eslint/js": "^9.8.0",
8081
"@jest-mock/express": "^3.0.0",
8182
"@types/content-type": "^1.1.8",

‎src/examples/backfill/backfillSampling.ts‎

Lines changed: 236 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,236 @@
1+
/*
2+
This example implements an stdio MCP proxy that backfills sampling requests using the Claude API.
3+
4+
Usage:
5+
npx -y @modelcontextprotocol/inspector \
6+
npx -y --silent tsx src/examples/backfill/backfillSampling.ts \
7+
npx -y --silent @modelcontextprotocol/server-everything
8+
*/
9+
10+
import { Anthropic } from "@anthropic-ai/sdk";
11+
import { Base64ImageSource, ContentBlock, ContentBlockParam, MessageParam } from "@anthropic-ai/sdk/resources/messages.js";
12+
import { StdioServerTransport } from '../../server/stdio.js';
13+
import { StdioClientTransport } from '../../client/stdio.js';
14+
import {
15+
CancelledNotification,
16+
CancelledNotificationSchema,
17+
isInitializeRequest,
18+
isJSONRPCRequest,
19+
ElicitRequest,
20+
ElicitRequestSchema,
21+
CreateMessageRequest,
22+
CreateMessageRequestSchema,
23+
CreateMessageResult,
24+
JSONRPCResponse,
25+
isInitializedNotification,
26+
CallToolRequest,
27+
CallToolRequestSchema,
28+
isJSONRPCNotification,
29+
} from "../../types.js";
30+
import { Transport } from "../../shared/transport.js";
31+
32+
// TODO: move to SDK
33+
34+
const isCancelledNotification: (value: unknown) => value is CancelledNotification =
35+
((value: any) => CancelledNotificationSchema.safeParse(value).success) as any;
36+
37+
const isCallToolRequest: (value: unknown) => value is CallToolRequest =
38+
((value: any) => CallToolRequestSchema.safeParse(value).success) as any;
39+
40+
const isElicitRequest: (value: unknown) => value is ElicitRequest =
41+
((value: any) => ElicitRequestSchema.safeParse(value).success) as any;
42+
43+
const isCreateMessageRequest: (value: unknown) => value is CreateMessageRequest =
44+
((value: any) => CreateMessageRequestSchema.safeParse(value).success) as any;
45+
46+
47+
function contentToMcp(content: ContentBlock): CreateMessageResult['content'][number] {
48+
switch (content.type) {
49+
case 'text':
50+
return {type: 'text', text: content.text};
51+
default:
52+
throw new Error(`Unsupported content type: ${content.type}`);
53+
}
54+
}
55+
56+
function contentFromMcp(content: CreateMessageRequest['params']['messages'][number]['content']): ContentBlockParam {
57+
switch (content.type) {
58+
case 'text':
59+
return {type: 'text', text: content.text};
60+
case 'image':
61+
return {
62+
type: 'image',
63+
source: {
64+
data: content.data,
65+
media_type: content.mimeType as Base64ImageSource['media_type'],
66+
type: 'base64',
67+
},
68+
};
69+
case 'audio':
70+
default:
71+
throw new Error(`Unsupported content type: ${content.type}`);
72+
}
73+
}
74+
75+
export type NamedTransport<T extends Transport = Transport> = {
76+
name: 'client' | 'server',
77+
transport: T,
78+
}
79+
80+
export async function setupBackfill(client: NamedTransport, server: NamedTransport, api: Anthropic) {
81+
const backfillMeta = await (async () => {
82+
const models = new Set<string>();
83+
let defaultModel: string | undefined;
84+
for await (const info of api.models.list()) {
85+
models.add(info.id);
86+
if (info.id.indexOf('sonnet') >= 0 && defaultModel === undefined) {
87+
defaultModel = info.id;
88+
}
89+
}
90+
if (defaultModel === undefined) {
91+
if (models.size === 0) {
92+
throw new Error("No models available from the API");
93+
}
94+
defaultModel = models.values().next().value;
95+
}
96+
return {
97+
sampling_models: Array.from(models),
98+
sampling_default_model: defaultModel,
99+
};
100+
})();
101+
102+
function pickModel(preferences: CreateMessageRequest['params']['modelPreferences'] | undefined): string {
103+
if (preferences?.hints) {
104+
for (const hint of Object.values(preferences.hints)) {
105+
if (hint.name !== undefined && backfillMeta.sampling_models.includes(hint.name)) {
106+
return hint.name;
107+
}
108+
}
109+
}
110+
// TODO: linear model on preferences?.{intelligencePriority, speedPriority, costPriority} to pick betwen haiku, sonnet, opus.
111+
return backfillMeta.sampling_default_model!;
112+
}
113+
114+
let clientSupportsSampling: boolean | undefined;
115+
// let clientSupportsElicitation: boolean | undefined;
116+
117+
const propagateMessage = (source: NamedTransport, target: NamedTransport) => {
118+
source.transport.onmessage = async (message, extra) => {
119+
console.error(`[proxy]: Message from ${source.name} transport: ${JSON.stringify(message)}; extra: ${JSON.stringify(extra)}`);
120+
121+
if (isJSONRPCRequest(message)) {
122+
if (isInitializeRequest(message)) {
123+
if (!(clientSupportsSampling = !!message.params.capabilities.sampling)) {
124+
message.params.capabilities.sampling = {}
125+
message.params._meta = {...(message.params._meta ?? {}), ...backfillMeta};
126+
}
127+
} else if (isCreateMessageRequest(message) && !clientSupportsSampling) {
128+
if (message.params.includeContext !== 'none') {
129+
source.transport.send({
130+
jsonrpc: "2.0",
131+
id: message.id,
132+
error: {
133+
code: -32601, // Method not found
134+
message: "includeContext != none not supported by MCP sampling backfill",
135+
},
136+
}, {relatedRequestId: message.id});
137+
return;
138+
}
139+
140+
message.params.metadata;
141+
message.params.modelPreferences;
142+
143+
try {
144+
// message.params.
145+
const msg = await api.messages.create({
146+
model: pickModel(message.params.modelPreferences),
147+
system: message.params.systemPrompt === undefined ? undefined : [
148+
{
149+
type: "text",
150+
text: message.params.systemPrompt
151+
},
152+
],
153+
messages: message.params.messages.map(({role, content}) => (<MessageParam>{
154+
role,
155+
content: [contentFromMcp(content)]
156+
})),
157+
max_tokens: message.params.maxTokens,
158+
temperature: message.params.temperature,
159+
stop_sequences: message.params.stopSequences,
160+
});
161+
162+
if (msg.content.length !== 1) {
163+
throw new Error(`Expected exactly one content item in the response, got ${msg.content.length}`);
164+
}
165+
166+
source.transport.send(<JSONRPCResponse>{
167+
jsonrpc: "2.0",
168+
id: message.id,
169+
result: <CreateMessageResult>{
170+
model: msg.model,
171+
stopReason: msg.stop_reason,
172+
role: msg.role,
173+
content: contentToMcp(msg.content[0]),
174+
},
175+
});
176+
} catch (error) {
177+
source.transport.send({
178+
jsonrpc: "2.0",
179+
id: message.id,
180+
error: {
181+
code: -32601, // Method not found
182+
message: `Error processing message: ${(error as Error).message}`,
183+
},
184+
}, {relatedRequestId: message.id});
185+
}
186+
return;
187+
// } else if (isElicitRequest(message) && !clientSupportsElicitation) {
188+
// // TODO: form
189+
// return;
190+
}
191+
} else if (isJSONRPCNotification(message)) {
192+
if (isInitializedNotification(message) && source.name === 'server') {
193+
if (!clientSupportsSampling) {
194+
message.params = {...(message.params ?? {}), _meta: {...(message.params?._meta ?? {}), ...backfillMeta}};
195+
}
196+
}
197+
}
198+
199+
try {
200+
const relatedRequestId = isCancelledNotification(message)? message.params.requestId : undefined;
201+
await target.transport.send(message, {relatedRequestId});
202+
} catch (error) {
203+
console.error(`[proxy]: Error sending message to ${target.name}:`, error);
204+
}
205+
};
206+
};
207+
propagateMessage(server, client);
208+
propagateMessage(client, server);
209+
210+
const addErrorHandler = (transport: NamedTransport) => {
211+
transport.transport.onerror = async (error: Error) => {
212+
console.error(`[proxy]: Error from ${transport.name} transport:`, error);
213+
};
214+
};
215+
216+
addErrorHandler(client);
217+
addErrorHandler(server);
218+
219+
await server.transport.start();
220+
await client.transport.start();
221+
}
222+
223+
async function main() {
224+
const args = process.argv.slice(2);
225+
const client: NamedTransport = {name: 'client', transport: new StdioClientTransport({command: args[0], args: args.slice(1)})};
226+
const server: NamedTransport = {name: 'server', transport: new StdioServerTransport()};
227+
228+
const api = new Anthropic();
229+
await setupBackfill(client, server, api);
230+
console.error("[proxy]: Transports started.");
231+
}
232+
233+
main().catch((error) => {
234+
console.error("[proxy]: Fatal error:", error);
235+
process.exit(1);
236+
});

0 commit comments

Comments
 (0)