7
7
import pydantic_core
8
8
from pydantic import BaseModel , Field , TypeAdapter , validate_call
9
9
10
- from mcp .types import AudioContent , EmbeddedResource , ImageContent , TextContent
11
-
12
- CONTENT_TYPES = TextContent | ImageContent | AudioContent | EmbeddedResource
10
+ from mcp .types import Content , TextContent
13
11
14
12
15
13
class Message (BaseModel ):
16
14
"""Base class for all prompt messages."""
17
15
18
16
role : Literal ["user" , "assistant" ]
19
- content : CONTENT_TYPES
17
+ content : Content
20
18
21
- def __init__ (self , content : str | CONTENT_TYPES , ** kwargs : Any ):
19
+ def __init__ (self , content : str | Content , ** kwargs : Any ):
22
20
if isinstance (content , str ):
23
21
content = TextContent (type = "text" , text = content )
24
22
super ().__init__ (content = content , ** kwargs )
@@ -29,7 +27,7 @@ class UserMessage(Message):
29
27
30
28
role : Literal ["user" , "assistant" ] = "user"
31
29
32
- def __init__ (self , content : str | CONTENT_TYPES , ** kwargs : Any ):
30
+ def __init__ (self , content : str | Content , ** kwargs : Any ):
33
31
super ().__init__ (content = content , ** kwargs )
34
32
35
33
@@ -38,7 +36,7 @@ class AssistantMessage(Message):
38
36
39
37
role : Literal ["user" , "assistant" ] = "assistant"
40
38
41
- def __init__ (self , content : str | CONTENT_TYPES , ** kwargs : Any ):
39
+ def __init__ (self , content : str | Content , ** kwargs : Any ):
42
40
super ().__init__ (content = content , ** kwargs )
43
41
44
42
0 commit comments