Jeremiah Lowin commited on
Commit
262b9ea
·
1 Parent(s): c673fcb

Update prompt fn signature

Browse files
Files changed (1) hide show
  1. src/fastmcp/prompts/base.py +16 -4
src/fastmcp/prompts/base.py CHANGED
@@ -1,7 +1,7 @@
1
  """Base classes for FastMCP prompts."""
2
 
3
  import json
4
- from typing import Any, Callable, Dict, Literal, Optional, Sequence, Union
5
  import inspect
6
 
7
  from pydantic import BaseModel, Field, TypeAdapter, validate_call
@@ -41,7 +41,12 @@ class AssistantMessage(Message):
41
  super().__init__(content=content, **kwargs)
42
 
43
 
44
- message_validator = TypeAdapter(Union[UserMessage, AssistantMessage])
 
 
 
 
 
45
 
46
 
47
  class PromptArgument(BaseModel):
@@ -71,11 +76,18 @@ class Prompt(BaseModel):
71
  @classmethod
72
  def from_function(
73
  cls,
74
- fn: Callable[..., Sequence[Message]],
75
  name: Optional[str] = None,
76
  description: Optional[str] = None,
77
  ) -> "Prompt":
78
- """Create a Prompt from a function."""
 
 
 
 
 
 
 
79
  func_name = name or fn.__name__
80
 
81
  if func_name == "<lambda>":
 
1
  """Base classes for FastMCP prompts."""
2
 
3
  import json
4
+ from typing import Any, Callable, Dict, Literal, Optional, Sequence, Awaitable
5
  import inspect
6
 
7
  from pydantic import BaseModel, Field, TypeAdapter, validate_call
 
41
  super().__init__(content=content, **kwargs)
42
 
43
 
44
+ message_validator = TypeAdapter(UserMessage | AssistantMessage)
45
+
46
+ SyncPromptResult = (
47
+ str | Message | dict[str, Any] | Sequence[str | Message | dict[str, Any]]
48
+ )
49
+ PromptResult = SyncPromptResult | Awaitable[SyncPromptResult]
50
 
51
 
52
  class PromptArgument(BaseModel):
 
76
  @classmethod
77
  def from_function(
78
  cls,
79
+ fn: Callable[..., PromptResult],
80
  name: Optional[str] = None,
81
  description: Optional[str] = None,
82
  ) -> "Prompt":
83
+ """Create a Prompt from a function.
84
+
85
+ The function can return:
86
+ - A string (converted to a message)
87
+ - A Message object
88
+ - A dict (converted to a message)
89
+ - A sequence of any of the above
90
+ """
91
  func_name = name or fn.__name__
92
 
93
  if func_name == "<lambda>":