flymcprust_core/
prompt.rs

1//! Prompt definitions, arguments and messages.
2//!
3//! Prompts are usually declared with the `#[prompt]` macro, which implements
4//! [`Prompt`] for a marker type. [`PromptDef`] can also be built by hand.
5
6use std::collections::BTreeMap;
7use std::future::Future;
8use std::pin::Pin;
9use std::str::FromStr;
10
11use crate::completion::{CompleteFn, Completer};
12use crate::context::Ctx;
13use crate::spec::{
14    self, AudioContent, ContentBlock, EmbeddedResource, GetPromptResult, ImageContent,
15    PromptMessage, ResourceLink, Role, TextContent,
16};
17use crate::tool::{IntoToolError, StateRequirement, ToolError};
18
19/// The future returned by a prompt.
20pub type PromptFuture =
21    Pin<Box<dyn Future<Output = Result<GetPromptResult, ToolError>> + Send + 'static>>;
22
23/// Entry point of a prompt: receives the (string) arguments of `prompts/get`.
24pub type GetFn = fn(arguments: &PromptArguments, ctx: Ctx) -> Result<PromptFuture, ToolError>;
25
26/// A prompt that can be registered with
27/// [`ServerBuilder::prompt`](crate::ServerBuilder::prompt).
28pub trait Prompt: 'static {
29    /// Builds the prompt definition. Called once, at registration.
30    fn definition() -> PromptDef;
31}
32
33/// A registered prompt: its protocol description plus its entry point.
34#[derive(Debug, Clone)]
35pub struct PromptDef {
36    pub(crate) prompt: spec::Prompt,
37    pub(crate) get: GetFn,
38    pub(crate) required_state: Vec<StateRequirement>,
39    pub(crate) completers: Vec<Completer>,
40    pub(crate) scopes: Vec<String>,
41}
42
43impl PromptDef {
44    /// Creates a definition from the protocol description and entry point.
45    #[must_use]
46    pub fn new(prompt: spec::Prompt, get: GetFn) -> Self {
47        Self {
48            prompt,
49            get,
50            required_state: Vec::new(),
51            completers: Vec::new(),
52            scopes: Vec::new(),
53        }
54    }
55
56    /// Declares the OAuth scopes a caller needs (all of them). Enforced when
57    /// the HTTP transport requires OAuth; the client gets `403` with an
58    /// `insufficient_scope` challenge listing these scopes.
59    #[must_use]
60    pub fn scopes<I, S>(mut self, scopes: I) -> Self
61    where
62        I: IntoIterator<Item = S>,
63        S: Into<String>,
64    {
65        self.scopes = scopes.into_iter().map(Into::into).collect();
66        self
67    }
68
69    /// The OAuth scopes declared with [`scopes`](Self::scopes).
70    #[must_use]
71    pub fn required_scopes(&self) -> &[String] {
72        &self.scopes
73    }
74
75    /// Declares that the prompt needs shared state of type `T`.
76    #[must_use]
77    pub fn require_state<T: Send + Sync + 'static>(mut self) -> Self {
78        self.required_state.push(StateRequirement::of::<T>());
79        self
80    }
81
82    /// Attaches a completion provider to an argument.
83    #[must_use]
84    pub fn complete(mut self, argument: impl Into<String>, complete: CompleteFn) -> Self {
85        self.completers.push(Completer {
86            argument: argument.into(),
87            complete,
88        });
89        self
90    }
91
92    /// The protocol description.
93    #[must_use]
94    pub fn prompt(&self) -> &spec::Prompt {
95        &self.prompt
96    }
97}
98
99/// The arguments of a `prompts/get` request.
100#[derive(Debug, Clone, Default, PartialEq, Eq)]
101pub struct PromptArguments(BTreeMap<String, String>);
102
103impl PromptArguments {
104    pub(crate) fn new(arguments: BTreeMap<String, String>) -> Self {
105        Self(arguments)
106    }
107
108    /// The raw value of an argument.
109    #[must_use]
110    pub fn get(&self, name: &str) -> Option<&str> {
111        self.0.get(name).map(String::as_str)
112    }
113
114    /// Parses a required argument.
115    ///
116    /// # Errors
117    ///
118    /// Returns an invalid-argument error (`-32602`) if the argument is missing
119    /// or does not parse.
120    pub fn required<T: FromStr>(&self, name: &str) -> Result<T, ToolError> {
121        let value = self
122            .get(name)
123            .ok_or_else(|| ToolError::invalid_argument(format!("missing argument `{name}`")))?;
124        parse(name, value)
125    }
126
127    /// Parses an optional argument.
128    ///
129    /// # Errors
130    ///
131    /// Returns an invalid-argument error (`-32602`) if the argument is present
132    /// but does not parse.
133    pub fn optional<T: FromStr>(&self, name: &str) -> Result<Option<T>, ToolError> {
134        self.get(name).map(|v| parse(name, v)).transpose()
135    }
136
137    /// All arguments.
138    #[must_use]
139    pub fn as_map(&self) -> &BTreeMap<String, String> {
140        &self.0
141    }
142}
143
144fn parse<T: FromStr>(name: &str, value: &str) -> Result<T, ToolError> {
145    value
146        .parse()
147        .map_err(|_| ToolError::invalid_argument(format!("argument `{name}` is not valid")))
148}
149
150/// Content that can appear in a message: `&str`, `String` (as text) and the
151/// specification's content types.
152pub trait IntoContent {
153    /// Performs the conversion.
154    fn into_content(self) -> ContentBlock;
155}
156
157impl IntoContent for &str {
158    fn into_content(self) -> ContentBlock {
159        TextContent::new(self.to_owned()).into()
160    }
161}
162
163impl IntoContent for String {
164    fn into_content(self) -> ContentBlock {
165        TextContent::new(self).into()
166    }
167}
168
169macro_rules! content_types {
170    ($($t:ty),* $(,)?) => {
171        $(impl IntoContent for $t {
172            fn into_content(self) -> ContentBlock {
173                self.into()
174            }
175        })*
176    };
177}
178
179content_types!(
180    TextContent,
181    ImageContent,
182    AudioContent,
183    EmbeddedResource,
184    ResourceLink
185);
186
187impl IntoContent for ContentBlock {
188    fn into_content(self) -> ContentBlock {
189        self
190    }
191}
192
193/// Builders for prompt messages.
194pub mod message {
195    use super::{IntoContent, PromptMessage, Role};
196
197    /// A message from the user.
198    pub fn user(content: impl IntoContent) -> PromptMessage {
199        PromptMessage::new(content.into_content(), Role::User)
200    }
201
202    /// A message from the assistant.
203    pub fn assistant(content: impl IntoContent) -> PromptMessage {
204        PromptMessage::new(content.into_content(), Role::Assistant)
205    }
206}
207
208/// Conversion of a prompt's return value into a result.
209///
210/// | Return type | Result |
211/// | --- | --- |
212/// | `String`, `&'static str` | One user text message |
213/// | `PromptMessage`, `Vec<PromptMessage>` | The messages |
214/// | `GetPromptResult` | As given |
215/// | `Result<T, E>` | `T` as above; `E` as for tools |
216pub trait IntoPromptResult {
217    /// Performs the conversion.
218    ///
219    /// # Errors
220    ///
221    /// Returns the handler's error.
222    fn into_prompt_result(self) -> Result<GetPromptResult, ToolError>;
223}
224
225fn messages(messages: Vec<PromptMessage>) -> GetPromptResult {
226    GetPromptResult::new(messages, "complete".to_owned())
227}
228
229impl IntoPromptResult for String {
230    fn into_prompt_result(self) -> Result<GetPromptResult, ToolError> {
231        Ok(messages(vec![message::user(self)]))
232    }
233}
234
235impl IntoPromptResult for &'static str {
236    fn into_prompt_result(self) -> Result<GetPromptResult, ToolError> {
237        Ok(messages(vec![message::user(self)]))
238    }
239}
240
241impl IntoPromptResult for PromptMessage {
242    fn into_prompt_result(self) -> Result<GetPromptResult, ToolError> {
243        Ok(messages(vec![self]))
244    }
245}
246
247impl IntoPromptResult for Vec<PromptMessage> {
248    fn into_prompt_result(self) -> Result<GetPromptResult, ToolError> {
249        Ok(messages(self))
250    }
251}
252
253impl IntoPromptResult for GetPromptResult {
254    fn into_prompt_result(self) -> Result<GetPromptResult, ToolError> {
255        Ok(self)
256    }
257}
258
259impl<T: IntoPromptResult, E: IntoToolError> IntoPromptResult for Result<T, E> {
260    fn into_prompt_result(self) -> Result<GetPromptResult, ToolError> {
261        match self {
262            Ok(value) => value.into_prompt_result(),
263            Err(error) => Err(error.into_tool_error()),
264        }
265    }
266}
267
268#[cfg(test)]
269mod tests {
270    use super::*;
271
272    #[test]
273    fn arguments_parse() {
274        let args = PromptArguments::new(BTreeMap::from([
275            ("n".to_owned(), "3".to_owned()),
276            ("s".to_owned(), "x".to_owned()),
277        ]));
278        assert_eq!(args.required::<u8>("n").unwrap(), 3);
279        assert_eq!(args.optional::<u8>("missing").unwrap(), None);
280        assert!(args.required::<u8>("s").is_err());
281        assert!(args.required::<String>("missing").is_err());
282    }
283
284    #[test]
285    fn string_becomes_a_user_message() {
286        let r = "hello".into_prompt_result().unwrap();
287        assert_eq!(
288            serde_json::to_value(&r.messages).unwrap(),
289            serde_json::json!([{ "role": "user", "content": { "type": "text", "text": "hello" } }])
290        );
291    }
292}
flymcprust_core/prompt.rs

flymcprust_core/
prompt.rs
