42
42
SubscribeRequest ,
43
43
Tool ,
44
44
UnsubscribeRequest ,
45
+ TextContent ,
46
+ EmbeddedResource ,
47
+ PromptMessage ,
45
48
)
46
49
47
50
logger = logging .getLogger (__name__ )
@@ -117,8 +120,6 @@ def get_prompt(self):
117
120
GetPromptRequest ,
118
121
GetPromptResult ,
119
122
ImageContent ,
120
- SamplingMessage ,
121
- TextContent ,
122
123
)
123
124
from mcp_python .types import (
124
125
Role as Role ,
@@ -133,7 +134,7 @@ def decorator(
133
134
134
135
async def handler (req : GetPromptRequest ):
135
136
prompt_get = await func (req .params .name , req .params .arguments )
136
- messages : list [SamplingMessage ] = []
137
+ messages : list [PromptMessage ] = []
137
138
for message in prompt_get .messages :
138
139
match message .content :
139
140
case str () as text_content :
@@ -144,15 +145,20 @@ async def handler(req: GetPromptRequest):
144
145
data = img_content .data ,
145
146
mimeType = img_content .mime_type ,
146
147
)
148
+ case types .EmbeddedResource () as resource :
149
+ content = EmbeddedResource (
150
+ type = "resource" ,
151
+ resource = resource .resource
152
+ )
147
153
case _:
148
154
raise ValueError (
149
155
f"Unexpected content type: { type (message .content )} "
150
156
)
151
157
152
- sampling_message = SamplingMessage (
158
+ prompt_message = PromptMessage (
153
159
role = message .role , content = content
154
160
)
155
- messages .append (sampling_message )
161
+ messages .append (prompt_message )
156
162
157
163
return ServerResult (
158
164
GetPromptResult (description = prompt_get .desc , messages = messages )
@@ -276,14 +282,46 @@ async def handler(_: Any):
276
282
return decorator
277
283
278
284
def call_tool (self ):
279
- from mcp_python .types import CallToolResult
285
+ from mcp_python .types import CallToolResult , TextContent , ImageContent , EmbeddedResource
280
286
281
- def decorator (func : Callable [..., Awaitable [Any ]]):
287
+ def decorator (
288
+ func : Callable [..., Awaitable [list [str | types .ImageContent | types .EmbeddedResource ]]]
289
+ ):
282
290
logger .debug ("Registering handler for CallToolRequest" )
283
291
284
292
async def handler (req : CallToolRequest ):
285
- result = await func (req .params .name , (req .params .arguments or {}))
286
- return ServerResult (CallToolResult (toolResult = result ))
293
+ try :
294
+ results = await func (req .params .name , (req .params .arguments or {}))
295
+ content = []
296
+ for result in results :
297
+ match result :
298
+ case str () as text :
299
+ content .append (TextContent (type = "text" , text = text ))
300
+ case types .ImageContent () as img :
301
+ content .append (ImageContent (
302
+ type = "image" ,
303
+ data = img .data ,
304
+ mimeType = img .mime_type
305
+ ))
306
+ case types .EmbeddedResource () as resource :
307
+ content .append (EmbeddedResource (
308
+ type = "resource" ,
309
+ resource = resource .resource
310
+ ))
311
+
312
+ return ServerResult (
313
+ CallToolResult (
314
+ content = content ,
315
+ isError = False
316
+ )
317
+ )
318
+ except Exception as e :
319
+ return ServerResult (
320
+ CallToolResult (
321
+ content = [TextContent (type = "text" , text = str (e ))],
322
+ isError = True
323
+ )
324
+ )
287
325
288
326
self .request_handlers [CallToolRequest ] = handler
289
327
return func
0 commit comments