77import pydantic_core
88from pydantic import BaseModel , Field , TypeAdapter , validate_call
99
10- from mcp .types import AudioContent , EmbeddedResource , ImageContent , TextContent
11-
12- CONTENT_TYPES = TextContent | ImageContent | AudioContent | EmbeddedResource
10+ from mcp .types import Content , TextContent
1311
1412
1513class Message (BaseModel ):
1614 """Base class for all prompt messages."""
1715
1816 role : Literal ["user" , "assistant" ]
19- content : CONTENT_TYPES
17+ content : Content
2018
21- def __init__ (self , content : str | CONTENT_TYPES , ** kwargs : Any ):
19+ def __init__ (self , content : str | Content , ** kwargs : Any ):
2220 if isinstance (content , str ):
2321 content = TextContent (type = "text" , text = content )
2422 super ().__init__ (content = content , ** kwargs )
@@ -29,7 +27,7 @@ class UserMessage(Message):
2927
3028 role : Literal ["user" , "assistant" ] = "user"
3129
32- def __init__ (self , content : str | CONTENT_TYPES , ** kwargs : Any ):
30+ def __init__ (self , content : str | Content , ** kwargs : Any ):
3331 super ().__init__ (content = content , ** kwargs )
3432
3533
@@ -38,7 +36,7 @@ class AssistantMessage(Message):
3836
3937 role : Literal ["user" , "assistant" ] = "assistant"
4038
41- def __init__ (self , content : str | CONTENT_TYPES , ** kwargs : Any ):
39+ def __init__ (self , content : str | Content , ** kwargs : Any ):
4240 super ().__init__ (content = content , ** kwargs )
4341
4442
0 commit comments