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

llm_tool/
rust_prompt.rs
