llm_tool/
rust_prompt.rs

1//! Strongly-typed Rust prompt trait and type-erasure machinery.
2
3use alloc::{borrow::Cow, boxed::Box, format, string::ToString, vec::Vec};
4use core::{future::Future, pin::Pin};
5
6use super::types::{
7    PromptArgumentDefinition, PromptDefinition, PromptOutput, RegistryItem, ToolError,
8};
9use crate::compat::{HashMap, HashMapIter};
10
11/// A custom prompt template implemented in Rust with strongly-typed parameters.
12pub trait RustPrompt: Send + Sync {
13    /// The strongly-typed parameters struct deriving `Deserialize` + `JsonSchema`.
14    type Params: serde::de::DeserializeOwned + schemars::JsonSchema + Send;
15
16    /// Unique prompt name.
17    const NAME: &'static str;
18
19    /// Human-readable description.
20    const DESCRIPTION: &'static str;
21
22    /// Return the prompt description.
23    fn description(&self) -> Cow<'static, str> {
24        Cow::Borrowed(Self::DESCRIPTION)
25    }
26
27    /// Render the prompt template with typed parameters.
28    fn render(
29        &self,
30        params: Self::Params,
31    ) -> impl Future<Output = Result<PromptOutput, ToolError>> + Send;
32}
33
34/// Strongly-typed representation of a JSON Schema property for prompt arguments.
35#[derive(serde::Deserialize)]
36struct PromptPropertySchema {
37    /// Optional human-readable description from field doc comments.
38    description: Option<alloc::string::String>,
39}
40
41/// Strongly-typed representation of an object JSON Schema for [`RustPrompt::Params`].
42#[derive(serde::Deserialize)]
43struct PromptParamsSchema {
44    /// JSON Schema `"type"` discriminator (must be `"object"` for prompt parameter structs).
45    #[serde(rename = "type")]
46    schema_type: Option<alloc::string::String>,
47    /// Required parameter names.
48    required: Option<Vec<alloc::string::String>>,
49    /// Declared object properties in deterministic key order.
50    properties: Option<alloc::collections::BTreeMap<alloc::string::String, PromptPropertySchema>>,
51}
52
53/// Build a [`PromptDefinition`] from any [`RustPrompt`] implementor.
54///
55/// # Errors
56///
57/// Returns `Err` if the prompt's `Params` type fails to produce a valid
58/// `"type": "object"` JSON schema.
59pub fn definition_of_prompt<T: RustPrompt>(prompt: &T) -> Result<PromptDefinition, ToolError> {
60    let schema = schemars::schema_for!(T::Params);
61    let val = serde_json::to_value(&schema).map_err(|e| {
62        ToolError::new(format!(
63            "Failed to serialize schema for prompt '{}': {e}",
64            T::NAME
65        ))
66    })?;
67    let parsed: PromptParamsSchema = serde_json::from_value(val).map_err(|e| {
68        ToolError::new(format!(
69            "Failed to parse object schema for prompt '{}': {e}",
70            T::NAME
71        ))
72    })?;
73    if parsed.schema_type.as_deref() != Some("object") {
74        return Err(ToolError::new(format!(
75            "Prompt '{}' Params schema must have type \"object\", got {:?}",
76            T::NAME,
77            parsed.schema_type
78        )));
79    }
80
81    let required_fields = parsed.required.as_deref().unwrap_or(&[]);
82    let arguments = parsed
83        .properties
84        .into_iter()
85        .flatten()
86        .map(|(name, prop)| {
87            let required = required_fields.iter().any(|req| req == &name);
88            PromptArgumentDefinition {
89                name,
90                description: prop.description.unwrap_or_else(alloc::string::String::new),
91                required,
92            }
93        })
94        .collect();
95
96    Ok(PromptDefinition {
97        name: T::NAME.to_string(),
98        description: prompt.description().into_owned(),
99        arguments,
100    })
101}
102
103/// Type-erased future returned by [`ErasedPrompt::render_erased`].
104pub type BoxPromptFuture<'a> =
105    Pin<Box<dyn Future<Output = Result<PromptOutput, ToolError>> + Send + 'a>>;
106
107/// Type-erased wrapper enabling heterogeneous prompt storage.
108///
109/// This is an internal implementation detail of [`PromptRegistry`]; callers
110/// interact with prompts through the registry rather than this trait.
111pub trait ErasedPrompt: Send + Sync {
112    /// Deserialize arguments and render the prompt template.
113    fn render_erased(&self, args: serde_json::Value) -> BoxPromptFuture<'_>;
114}
115
116impl<T: RustPrompt> ErasedPrompt for T {
117    fn render_erased(&self, args: serde_json::Value) -> BoxPromptFuture<'_> {
118        Box::pin(async move {
119            let params: T::Params = serde_json::from_value(args).map_err(|e| {
120                ToolError::new(format!("Failed to deserialize prompt parameters: {e}"))
121            })?;
122            self.render(params).await
123        })
124    }
125}
126
127/// A registered prompt: its cached definition plus the type-erased handler.
128struct RegisteredPrompt {
129    definition: PromptDefinition,
130    erased: Box<dyn ErasedPrompt>,
131}
132
133/// A registry of named prompt templates for dynamic dispatch.
134///
135/// Mirrors [`ToolRegistry`](crate::ToolRegistry) for prompts: it stores
136/// type-erased [`RustPrompt`] implementations keyed by name, caches each
137/// [`PromptDefinition`] at registration time, and renders them on demand,
138/// keeping the type-erasure machinery a private implementation detail.
139#[derive(Default)]
140pub struct PromptRegistry {
141    prompts: HashMap<&'static str, RegisteredPrompt>,
142}
143
144impl core::fmt::Debug for PromptRegistry {
145    fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {
146        let names: Vec<&str> = self.prompts.keys().copied().collect();
147        f.debug_struct("PromptRegistry")
148            .field("prompt_count", &self.prompts.len())
149            .field("prompt_names", &names)
150            .finish()
151    }
152}
153
154impl PromptRegistry {
155    /// Create an empty prompt registry.
156    #[must_use]
157    pub fn new() -> Self {
158        Self {
159            prompts: HashMap::new(),
160        }
161    }
162
163    /// Register a [`RustPrompt`], replacing any existing prompt of the same name.
164    ///
165    /// Returns `&mut Self` for chaining.
166    ///
167    /// # Panics
168    ///
169    /// Panics if the prompt's JSON schema cannot be serialized. This indicates
170    /// a bug in the prompt's `Params` type. Use
171    /// [`try_register`](Self::try_register) for the non-panicking variant.
172    pub fn register<P: RustPrompt + 'static>(&mut self, prompt: P) -> &mut Self {
173        if let Err(e) = self.try_register(prompt) {
174            panic!("Failed to build definition for prompt '{}': {e}", P::NAME);
175        }
176        self
177    }
178
179    /// Register a [`RustPrompt`], returning an error instead of panicking if
180    /// the prompt's JSON schema cannot be built.
181    ///
182    /// This is the fallible counterpart to [`register`](Self::register). If a
183    /// prompt with the same name was already registered, it is replaced.
184    ///
185    /// # Errors
186    ///
187    /// Returns [`ToolError`] if the prompt's `Params` type fails to produce a
188    /// JSON schema.
189    pub fn try_register<P: RustPrompt + 'static>(
190        &mut self,
191        prompt: P,
192    ) -> Result<&mut Self, ToolError> {
193        let definition = definition_of_prompt(&prompt)?;
194        self.prompts.insert(
195            P::NAME,
196            RegisteredPrompt {
197                definition,
198                erased: Box::new(prompt),
199            },
200        );
201        Ok(self)
202    }
203
204    /// Register a [`RustPrompt`], consuming and returning `Self` for chaining.
205    ///
206    /// # Panics
207    ///
208    /// Panics if the prompt's JSON schema cannot be serialized; see
209    /// [`register`](Self::register).
210    #[must_use]
211    pub fn with_prompt<P: RustPrompt + 'static>(mut self, prompt: P) -> Self {
212        self.register(prompt);
213        self
214    }
215
216    /// Collect [`PromptDefinition`]s for all registered prompts.
217    ///
218    /// Returns clones of the cached definitions computed at registration time.
219    #[must_use]
220    pub fn definitions(&self) -> Vec<PromptDefinition> {
221        self.prompts
222            .values()
223            .map(|entry| entry.definition.clone())
224            .collect()
225    }
226
227    /// Remove a prompt by name, returning `true` if it was present.
228    pub fn remove(&mut self, name: &str) -> bool {
229        self.prompts.remove(name).is_some()
230    }
231
232    /// Clear all registered prompts.
233    pub fn clear(&mut self) {
234        self.prompts.clear();
235    }
236
237    /// Number of registered prompts.
238    #[must_use]
239    pub fn len(&self) -> usize {
240        self.prompts.len()
241    }
242
243    /// Whether the registry has no registered prompts.
244    #[must_use]
245    pub fn is_empty(&self) -> bool {
246        self.prompts.is_empty()
247    }
248
249    /// Whether a prompt with the given name is registered.
250    #[must_use]
251    pub fn contains(&self, name: &str) -> bool {
252        self.prompts.contains_key(name)
253    }
254
255    /// Borrow the cached [`PromptDefinition`] for a registered prompt by name.
256    ///
257    /// Returns `None` if no prompt named `name` is registered. Unlike
258    /// [`definitions`](Self::definitions), this clones nothing.
259    #[must_use]
260    pub fn definition(&self, name: &str) -> Option<&PromptDefinition> {
261        self.prompts.get(name).map(|entry| &entry.definition)
262    }
263
264    /// Iterate over `(name, definition)` pairs for every registered prompt.
265    ///
266    /// Yields clones of the cached definitions computed at registration time.
267    #[must_use]
268    pub fn iter(&self) -> PromptDefinitions<'_> {
269        PromptDefinitions {
270            inner: self.prompts.iter(),
271        }
272    }
273
274    /// Render a registered prompt by name with raw JSON arguments.
275    ///
276    /// # Errors
277    ///
278    /// Returns [`ToolError::not_found`] if no prompt named `name` is registered
279    /// (carrying `error_kind = "not_registered"` metadata), or a render error
280    /// if argument deserialization or rendering fails.
281    pub async fn render(
282        &self,
283        name: &str,
284        args: serde_json::Value,
285    ) -> Result<PromptOutput, ToolError> {
286        let Some(entry) = self.prompts.get(name) else {
287            return Err(ToolError::not_found(RegistryItem::Prompt, name));
288        };
289        entry.erased.render_erased(args).await
290    }
291}
292
293/// Borrowing iterator over `(name, definition)` pairs, yielded by
294/// [`PromptRegistry::iter`] and by `&PromptRegistry`'s [`IntoIterator`] impl.
295///
296/// Each cached [`PromptDefinition`] is cloned lazily as it is yielded.
297pub struct PromptDefinitions<'a> {
298    inner: HashMapIter<'a, &'static str, RegisteredPrompt>,
299}
300
301impl Iterator for PromptDefinitions<'_> {
302    type Item = (&'static str, PromptDefinition);
303
304    fn next(&mut self) -> Option<Self::Item> {
305        self.inner
306            .next()
307            .map(|(name, entry)| (*name, entry.definition.clone()))
308    }
309
310    fn size_hint(&self) -> (usize, Option<usize>) {
311        self.inner.size_hint()
312    }
313}
314
315impl ExactSizeIterator for PromptDefinitions<'_> {
316    fn len(&self) -> usize {
317        self.inner.len()
318    }
319}
320
321/// Iterate over `(name, definition)` pairs for every registered prompt.
322impl<'a> IntoIterator for &'a PromptRegistry {
323    type Item = (&'static str, PromptDefinition);
324    type IntoIter = PromptDefinitions<'a>;
325
326    fn into_iter(self) -> Self::IntoIter {
327        self.iter()
328    }
329}
330
331#[cfg(test)]
332mod tests {
333    use super::*;
334
335    struct NonObjectPrompt;
336
337    impl RustPrompt for NonObjectPrompt {
338        type Params = alloc::string::String;
339        const NAME: &'static str = "non_object_prompt";
340        const DESCRIPTION: &'static str = "Prompt with non-object params";
341
342        async fn render(&self, params: Self::Params) -> Result<PromptOutput, ToolError> {
343            Ok(PromptOutput::user(params))
344        }
345    }
346
347    #[test]
348    fn definition_of_prompt_rejects_non_object_schema() {
349        let err = definition_of_prompt(&NonObjectPrompt)
350            .expect_err("non-object prompt schema must fail loudly");
351        assert!(
352            err.message.contains("must have type \"object\""),
353            "unexpected error: {}",
354            err.message
355        );
356    }
357}
llm_tool/rust_prompt.rs

llm_tool/
rust_prompt.rs
