From 296ef93b1929bd88a0ec90890894e2b09a2d2376 Mon Sep 17 00:00:00 2001
From: Jeremiah Lowin <153965+jlowin@users.noreply.github.com>
Date: Tue, 3 Dec 2024 13:49:02 -0500
Subject: [PATCH] Update prompt fn signature

---
 src/fastmcp/prompts/base.py | 20 ++++++++++++++++----
 1 file changed, 16 insertions(+), 4 deletions(-)

diff --git a/src/fastmcp/prompts/base.py b/src/fastmcp/prompts/base.py
index 2aaaba4d8..d44fc1823 100644
--- a/src/fastmcp/prompts/base.py
+++ b/src/fastmcp/prompts/base.py
@@ -1,7 +1,7 @@
 """Base classes for FastMCP prompts."""
 
 import json
-from typing import Any, Callable, Dict, Literal, Optional, Sequence, Union
+from typing import Any, Callable, Dict, Literal, Optional, Sequence, Awaitable
 import inspect
 
 from pydantic import BaseModel, Field, TypeAdapter, validate_call
@@ -41,7 +41,12 @@ class AssistantMessage(Message):
         super().__init__(content=content, **kwargs)
 
 
-message_validator = TypeAdapter(Union[UserMessage, AssistantMessage])
+message_validator = TypeAdapter(UserMessage | AssistantMessage)
+
+SyncPromptResult = (
+    str | Message | dict[str, Any] | Sequence[str | Message | dict[str, Any]]
+)
+PromptResult = SyncPromptResult | Awaitable[SyncPromptResult]
 
 
 class PromptArgument(BaseModel):
@@ -71,11 +76,18 @@ class Prompt(BaseModel):
     @classmethod
     def from_function(
         cls,
-        fn: Callable[..., Sequence[Message]],
+        fn: Callable[..., PromptResult],
         name: Optional[str] = None,
         description: Optional[str] = None,
     ) -> "Prompt":
-        """Create a Prompt from a function."""
+        """Create a Prompt from a function.
+
+        The function can return:
+        - A string (converted to a message)
+        - A Message object
+        - A dict (converted to a message)
+        - A sequence of any of the above
+        """
         func_name = name or fn.__name__
 
         if func_name == "<lambda>":