77import pydantic_core
88from pydantic import BaseModel , Field , TypeAdapter , validate_call
99
10- from mcp .types import Content , TextContent
10+ from mcp .types import ContentBlock , TextContent
1111
1212
1313class Message (BaseModel ):
1414 """Base class for all prompt messages."""
1515
1616 role : Literal ["user" , "assistant" ]
17- content : Content
17+ content : ContentBlock
1818
19- def __init__ (self , content : str | Content , ** kwargs : Any ):
19+ def __init__ (self , content : str | ContentBlock , ** kwargs : Any ):
2020 if isinstance (content , str ):
2121 content = TextContent (type = "text" , text = content )
2222 super ().__init__ (content = content , ** kwargs )
@@ -27,7 +27,7 @@ class UserMessage(Message):
2727
2828 role : Literal ["user" , "assistant" ] = "user"
2929
30- def __init__ (self , content : str | Content , ** kwargs : Any ):
30+ def __init__ (self , content : str | ContentBlock , ** kwargs : Any ):
3131 super ().__init__ (content = content , ** kwargs )
3232
3333
@@ -36,7 +36,7 @@ class AssistantMessage(Message):
3636
3737 role : Literal ["user" , "assistant" ] = "assistant"
3838
39- def __init__ (self , content : str | Content , ** kwargs : Any ):
39+ def __init__ (self , content : str | ContentBlock , ** kwargs : Any ):
4040 super ().__init__ (content = content , ** kwargs )
4141
4242
0 commit comments