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 = match serde_json::to_value(&schema) {
36        Ok(v) => v,
37        Err(e) => {
38            tracing::error!(
39                prompt = T::NAME,
40                error = %e,
41                "failed to serialize prompt schema, producing empty definition"
42            );
43            return PromptDefinition {
44                name: T::NAME.to_string(),
45                description: prompt.description().into_owned(),
46                arguments: Vec::new(),
47            };
48        }
49    };
50    let mut arguments = Vec::new();
51
52    if let Some(obj) = val.as_object() {
53        let required_fields: Vec<&str> = obj
54            .get("required")
55            .and_then(|r| r.as_array())
56            .map(|arr| arr.iter().filter_map(|v| v.as_str()).collect())
57            .unwrap_or_default();
58        if let Some(props) = obj.get("properties").and_then(|p| p.as_object()) {
59            for (name, prop) in props {
60                let description = prop
61                    .get("description")
62                    .and_then(|d| d.as_str())
63                    .unwrap_or("")
64                    .to_string();
65                let required = required_fields.contains(&name.as_str());
66                arguments.push(PromptArgumentDefinition {
67                    name: name.clone(),
68                    description,
69                    required,
70                });
71            }
72        }
73    }
74
75    PromptDefinition {
76        name: T::NAME.to_string(),
77        description: prompt.description().into_owned(),
78        arguments,
79    }
80}
81
82/// Type-erased future returned by [`ErasedPrompt::render_erased`].
83pub type BoxPromptFuture<'a> =
84    Pin<Box<dyn Future<Output = Result<PromptOutput, ToolError>> + Send + 'a>>;
85
86/// Type-erased wrapper enabling heterogeneous prompt storage.
87pub trait ErasedPrompt: Send + Sync {
88    /// Return the prompt definition.
89    fn definition(&self) -> PromptDefinition;
90
91    /// Deserialize arguments and render the prompt template.
92    fn render_erased(&self, args: serde_json::Value) -> BoxPromptFuture<'_>;
93}
94
95impl<T: RustPrompt> ErasedPrompt for T {
96    fn definition(&self) -> PromptDefinition {
97        definition_of_prompt(self)
98    }
99
100    fn render_erased(&self, args: serde_json::Value) -> BoxPromptFuture<'_> {
101        Box::pin(async move {
102            let params: T::Params = serde_json::from_value(args).map_err(|e| {
103                ToolError::new(format!("Failed to deserialize prompt parameters: {e}"))
104            })?;
105            self.render(params).await
106        })
107    }
108}