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            // NOLINT: empty vec is correct when schema has no 'required' field
58            .unwrap_or_default();
59        if let Some(props) = obj.get("properties").and_then(|p| p.as_object()) {
60            for (name, prop) in props {
61                let description = prop
62                    .get("description")
63                    .and_then(|d| d.as_str())
64                    .unwrap_or("")
65                    .to_string();
66                let required = required_fields.contains(&name.as_str());
67                arguments.push(PromptArgumentDefinition {
68                    name: name.clone(),
69                    description,
70                    required,
71                });
72            }
73        }
74    }
75
76    PromptDefinition {
77        name: T::NAME.to_string(),
78        description: prompt.description().into_owned(),
79        arguments,
80    }
81}
82
83/// Type-erased future returned by [`ErasedPrompt::render_erased`].
84pub type BoxPromptFuture<'a> =
85    Pin<Box<dyn Future<Output = Result<PromptOutput, ToolError>> + Send + 'a>>;
86
87/// Type-erased wrapper enabling heterogeneous prompt storage.
88pub trait ErasedPrompt: Send + Sync {
89    /// Return the prompt definition.
90    fn definition(&self) -> PromptDefinition;
91
92    /// Deserialize arguments and render the prompt template.
93    fn render_erased(&self, args: serde_json::Value) -> BoxPromptFuture<'_>;
94}
95
96impl<T: RustPrompt> ErasedPrompt for T {
97    fn definition(&self) -> PromptDefinition {
98        definition_of_prompt(self)
99    }
100
101    fn render_erased(&self, args: serde_json::Value) -> BoxPromptFuture<'_> {
102        Box::pin(async move {
103            let params: T::Params = serde_json::from_value(args).map_err(|e| {
104                ToolError::new(format!("Failed to deserialize prompt parameters: {e}"))
105            })?;
106            self.render(params).await
107        })
108    }
109}