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/// Build a [`PromptDefinition`] from any [`RustPrompt`] implementor.
35///
36/// # Errors
37///
38/// Returns `Err` if the prompt's `Params` type fails to produce a JSON schema.
39pub fn definition_of_prompt<T: RustPrompt>(prompt: &T) -> Result<PromptDefinition, ToolError> {
40    let schema = schemars::schema_for!(T::Params);
41    let val = serde_json::to_value(&schema).map_err(|e| {
42        ToolError::new(format!(
43            "Failed to serialize schema for prompt '{}': {e}",
44            T::NAME
45        ))
46    })?;
47    let mut arguments = Vec::new();
48
49    if let Some(obj) = val.as_object() {
50        let required_fields: Vec<&str> = match obj.get("required").and_then(|r| r.as_array()) {
51            Some(arr) => arr.iter().filter_map(|v| v.as_str()).collect(),
52            // An absent `required` field correctly yields no required arguments.
53            None => Vec::new(),
54        };
55        if let Some(props) = obj.get("properties").and_then(|p| p.as_object()) {
56            for (name, prop) in props {
57                let description = prop
58                    .get("description")
59                    .and_then(|d| d.as_str())
60                    .unwrap_or("")
61                    .to_string();
62                let required = required_fields.contains(&name.as_str());
63                arguments.push(PromptArgumentDefinition {
64                    name: name.clone(),
65                    description,
66                    required,
67                });
68            }
69        }
70    }
71
72    Ok(PromptDefinition {
73        name: T::NAME.to_string(),
74        description: prompt.description().into_owned(),
75        arguments,
76    })
77}
78
79/// Type-erased future returned by [`ErasedPrompt::render_erased`].
80pub(crate) type BoxPromptFuture<'a> =
81    Pin<Box<dyn Future<Output = Result<PromptOutput, ToolError>> + Send + 'a>>;
82
83/// Type-erased wrapper enabling heterogeneous prompt storage.
84///
85/// This is an internal implementation detail of [`PromptRegistry`]; callers
86/// interact with prompts through the registry rather than this trait.
87pub(crate) trait ErasedPrompt: Send + Sync {
88    /// Deserialize arguments and render the prompt template.
89    fn render_erased(&self, args: serde_json::Value) -> BoxPromptFuture<'_>;
90}
91
92impl<T: RustPrompt> ErasedPrompt for T {
93    fn render_erased(&self, args: serde_json::Value) -> BoxPromptFuture<'_> {
94        Box::pin(async move {
95            let params: T::Params = serde_json::from_value(args).map_err(|e| {
96                ToolError::new(format!("Failed to deserialize prompt parameters: {e}"))
97            })?;
98            self.render(params).await
99        })
100    }
101}
102
103/// A registered prompt: its cached definition plus the type-erased handler.
104struct RegisteredPrompt {
105    definition: PromptDefinition,
106    erased: Box<dyn ErasedPrompt>,
107}
108
109/// A registry of named prompt templates for dynamic dispatch.
110///
111/// Mirrors [`ToolRegistry`](crate::ToolRegistry) for prompts: it stores
112/// type-erased [`RustPrompt`] implementations keyed by name, caches each
113/// [`PromptDefinition`] at registration time, and renders them on demand,
114/// keeping the type-erasure machinery a private implementation detail.
115#[derive(Default)]
116pub struct PromptRegistry {
117    prompts: HashMap<&'static str, RegisteredPrompt>,
118}
119
120impl core::fmt::Debug for PromptRegistry {
121    fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {
122        let names: Vec<&str> = self.prompts.keys().copied().collect();
123        f.debug_struct("PromptRegistry")
124            .field("prompt_count", &self.prompts.len())
125            .field("prompt_names", &names)
126            .finish()
127    }
128}
129
130impl PromptRegistry {
131    /// Create an empty prompt registry.
132    #[must_use]
133    pub fn new() -> Self {
134        Self {
135            prompts: HashMap::new(),
136        }
137    }
138
139    /// Register a [`RustPrompt`], replacing any existing prompt of the same name.
140    ///
141    /// Returns `&mut Self` for chaining.
142    ///
143    /// # Panics
144    ///
145    /// Panics if the prompt's JSON schema cannot be serialized. This indicates
146    /// a bug in the prompt's `Params` type. Use
147    /// [`try_register`](Self::try_register) for the non-panicking variant.
148    pub fn register<P: RustPrompt + 'static>(&mut self, prompt: P) -> &mut Self {
149        if let Err(e) = self.try_register(prompt) {
150            panic!("Failed to build definition for prompt '{}': {e}", P::NAME);
151        }
152        self
153    }
154
155    /// Register a [`RustPrompt`], returning an error instead of panicking if
156    /// the prompt's JSON schema cannot be built.
157    ///
158    /// This is the fallible counterpart to [`register`](Self::register). If a
159    /// prompt with the same name was already registered, it is replaced.
160    ///
161    /// # Errors
162    ///
163    /// Returns [`ToolError`] if the prompt's `Params` type fails to produce a
164    /// JSON schema.
165    pub fn try_register<P: RustPrompt + 'static>(
166        &mut self,
167        prompt: P,
168    ) -> Result<&mut Self, ToolError> {
169        let definition = definition_of_prompt(&prompt)?;
170        self.prompts.insert(
171            P::NAME,
172            RegisteredPrompt {
173                definition,
174                erased: Box::new(prompt),
175            },
176        );
177        Ok(self)
178    }
179
180    /// Register a [`RustPrompt`], consuming and returning `Self` for chaining.
181    ///
182    /// # Panics
183    ///
184    /// Panics if the prompt's JSON schema cannot be serialized; see
185    /// [`register`](Self::register).
186    #[must_use]
187    pub fn with_prompt<P: RustPrompt + 'static>(mut self, prompt: P) -> Self {
188        self.register(prompt);
189        self
190    }
191
192    /// Collect [`PromptDefinition`]s for all registered prompts.
193    ///
194    /// Returns clones of the cached definitions computed at registration time.
195    #[must_use]
196    pub fn definitions(&self) -> Vec<PromptDefinition> {
197        self.prompts
198            .values()
199            .map(|entry| entry.definition.clone())
200            .collect()
201    }
202
203    /// Number of registered prompts.
204    #[must_use]
205    pub fn len(&self) -> usize {
206        self.prompts.len()
207    }
208
209    /// Whether the registry has no registered prompts.
210    #[must_use]
211    pub fn is_empty(&self) -> bool {
212        self.prompts.is_empty()
213    }
214
215    /// Whether a prompt with the given name is registered.
216    #[must_use]
217    pub fn contains(&self, name: &str) -> bool {
218        self.prompts.contains_key(name)
219    }
220
221    /// Borrow the cached [`PromptDefinition`] for a registered prompt by name.
222    ///
223    /// Returns `None` if no prompt named `name` is registered. Unlike
224    /// [`definitions`](Self::definitions), this clones nothing.
225    #[must_use]
226    pub fn definition(&self, name: &str) -> Option<&PromptDefinition> {
227        self.prompts.get(name).map(|entry| &entry.definition)
228    }
229
230    /// Iterate over `(name, definition)` pairs for every registered prompt.
231    ///
232    /// Yields clones of the cached definitions computed at registration time.
233    #[must_use]
234    pub fn iter(&self) -> PromptDefinitions<'_> {
235        PromptDefinitions {
236            inner: self.prompts.iter(),
237        }
238    }
239
240    /// Render a registered prompt by name with raw JSON arguments.
241    ///
242    /// # Errors
243    ///
244    /// Returns [`ToolError::not_found`] if no prompt named `name` is registered
245    /// (carrying `error_kind = "not_registered"` metadata), or a render error
246    /// if argument deserialization or rendering fails.
247    pub async fn render(
248        &self,
249        name: &str,
250        args: serde_json::Value,
251    ) -> Result<PromptOutput, ToolError> {
252        let Some(entry) = self.prompts.get(name) else {
253            return Err(ToolError::not_found(RegistryItem::Prompt, name));
254        };
255        entry.erased.render_erased(args).await
256    }
257}
258
259/// Borrowing iterator over `(name, definition)` pairs, yielded by
260/// [`PromptRegistry::iter`] and by `&PromptRegistry`'s [`IntoIterator`] impl.
261///
262/// Each cached [`PromptDefinition`] is cloned lazily as it is yielded.
263pub struct PromptDefinitions<'a> {
264    inner: HashMapIter<'a, &'static str, RegisteredPrompt>,
265}
266
267impl Iterator for PromptDefinitions<'_> {
268    type Item = (&'static str, PromptDefinition);
269
270    fn next(&mut self) -> Option<Self::Item> {
271        self.inner
272            .next()
273            .map(|(name, entry)| (*name, entry.definition.clone()))
274    }
275
276    fn size_hint(&self) -> (usize, Option<usize>) {
277        self.inner.size_hint()
278    }
279}
280
281impl ExactSizeIterator for PromptDefinitions<'_> {
282    fn len(&self) -> usize {
283        self.inner.len()
284    }
285}
286
287/// Iterate over `(name, definition)` pairs for every registered prompt.
288impl<'a> IntoIterator for &'a PromptRegistry {
289    type Item = (&'static str, PromptDefinition);
290    type IntoIter = PromptDefinitions<'a>;
291
292    fn into_iter(self) -> Self::IntoIter {
293        self.iter()
294    }
295}
llm_tool/rust_prompt.rs

llm_tool/
rust_prompt.rs
