Skip to main content

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};
7
8/// A custom prompt template implemented in Rust with strongly-typed parameters.
9pub trait RustPrompt: Send + Sync {
10    /// The strongly-typed parameters struct deriving `Deserialize` + `JsonSchema`.
11    type Params: serde::de::DeserializeOwned + schemars::JsonSchema + Send;
12
13    /// Unique prompt name.
14    const NAME: &'static str;
15
16    /// Human-readable description.
17    const DESCRIPTION: &'static str;
18
19    /// Return the prompt description.
20    fn description(&self) -> Cow<'static, str> {
21        Cow::Borrowed(Self::DESCRIPTION)
22    }
23
24    /// Render the prompt template with typed parameters.
25    fn render(
26        &self,
27        params: Self::Params,
28    ) -> impl Future<Output = Result<PromptOutput, ToolError>> + Send;
29}
30
31/// Build a [`PromptDefinition`] from any [`RustPrompt`] implementor.
32#[must_use]
33pub fn definition_of_prompt<T: RustPrompt>(prompt: &T) -> PromptDefinition {
34    let schema = schemars::schema_for!(T::Params);
35    let val = serde_json::to_value(&schema).unwrap_or(serde_json::Value::Null);
36    let mut arguments = Vec::new();
37
38    if let Some(obj) = val.as_object() {
39        let required_fields: Vec<&str> = obj
40            .get("required")
41            .and_then(|r| r.as_array())
42            .map(|arr| arr.iter().filter_map(|v| v.as_str()).collect())
43            .unwrap_or_default();
44        if let Some(props) = obj.get("properties").and_then(|p| p.as_object()) {
45            for (name, prop) in props {
46                let description = prop
47                    .get("description")
48                    .and_then(|d| d.as_str())
49                    .unwrap_or("")
50                    .to_string();
51                let required = required_fields.contains(&name.as_str());
52                arguments.push(PromptArgumentDefinition {
53                    name: name.clone(),
54                    description,
55                    required,
56                });
57            }
58        }
59    }
60
61    PromptDefinition {
62        name: T::NAME.to_string(),
63        description: prompt.description().into_owned(),
64        arguments,
65    }
66}
67
68/// Type-erased future returned by [`ErasedPrompt::render_erased`].
69pub type BoxPromptFuture<'a> =
70    Pin<Box<dyn Future<Output = Result<PromptOutput, ToolError>> + Send + 'a>>;
71
72/// Type-erased wrapper enabling heterogeneous prompt storage.
73pub trait ErasedPrompt: Send + Sync {
74    /// Return the prompt definition.
75    fn definition(&self) -> PromptDefinition;
76
77    /// Deserialize arguments and render the prompt template.
78    fn render_erased(&self, args: serde_json::Value) -> BoxPromptFuture<'_>;
79}
80
81impl<T: RustPrompt> ErasedPrompt for T {
82    fn definition(&self) -> PromptDefinition {
83        definition_of_prompt(self)
84    }
85
86    fn render_erased(&self, args: serde_json::Value) -> BoxPromptFuture<'_> {
87        Box::pin(async move {
88            let params: T::Params = serde_json::from_value(args).map_err(|e| {
89                ToolError::new(format!("Failed to deserialize prompt parameters: {e}"))
90            })?;
91            self.render(params).await
92        })
93    }
94}