Add typing for content

.github/workflows/run-static.yml

@@ -28,3 +28,9 @@ jobs:
           python-version: "3.12"
       - name: Run pre-commit
         uses: pre-commit/action@v3.0.1
+      - name: Install dependencies
+        run: |
+          python -m pip install --upgrade pip
+          pip install ".[tests]"
+      - name: Run pyright
+        run: pyright src tests

.pre-commit-config.yaml

@@ -18,8 +18,3 @@ repos:
       - id: ruff-format
       - id: ruff
         args: [--fix, --exit-non-zero-on-fix]
-
-  - repo: https://github.com/RobertCraigie/pyright-python
-    rev: v1.1.352
-    hooks:
-      - id: pyright

src/fastmcp/prompts/base.py

@@ -4,38 +4,42 @@ import json
 from typing import Any, Callable, Dict, Literal, Optional, Sequence, Union
 import inspect
 
-from pydantic import BaseModel, Field, TypeAdapter, field_validator, validate_call
+from pydantic import BaseModel, Field, TypeAdapter, validate_call
 from mcp.types import TextContent, ImageContent, EmbeddedResource
 import pydantic_core
 
+CONTENT_TYPES = TextContent | ImageContent | EmbeddedResource
+
 
 class Message(BaseModel):
     """Base class for all prompt messages."""
 
     role: Literal["user", "assistant"]
-    content: Union[TextContent, ImageContent, EmbeddedResource]
+    content: CONTENT_TYPES
 
-    def __init__(self, content, **kwargs):
+    def __init__(self, content: str | CONTENT_TYPES, **kwargs):
+        if isinstance(content, str):
+            content = TextContent(type="text", text=content)
         super().__init__(content=content, **kwargs)
 
-    @field_validator("content", mode="before")
-    def validate_content(cls, v):
-        if isinstance(v, str):
-            return TextContent(type="text", text=v)
-        return v
-
 
 class UserMessage(Message):
     """A message from the user."""
 
     role: Literal["user"] = "user"
 
+    def __init__(self, content: str | CONTENT_TYPES, **kwargs):
+        super().__init__(content=content, **kwargs)
+
 
 class AssistantMessage(Message):
     """A message from the assistant."""
 
     role: Literal["assistant"] = "assistant"
 
+    def __init__(self, content: str | CONTENT_TYPES, **kwargs):
+        super().__init__(content=content, **kwargs)
+
 
 message_validator = TypeAdapter(Union[UserMessage, AssistantMessage])