Skip to main content

agenterra_rmcp/model/
prompt.rs

1use base64::engine::{Engine, general_purpose::STANDARD as BASE64_STANDARD};
2use serde::{Deserialize, Serialize};
3
4use super::{
5    AnnotateAble, Annotations, RawEmbeddedResource, RawImageContent,
6    content::{EmbeddedResource, ImageContent},
7    resource::ResourceContents,
8};
9
10/// A prompt that can be used to generate text from a model
11#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
12#[serde(rename_all = "camelCase")]
13#[cfg_attr(feature = "schemars", derive(schemars::JsonSchema))]
14pub struct Prompt {
15    /// The name of the prompt
16    pub name: String,
17    /// Optional description of what the prompt does
18    #[serde(skip_serializing_if = "Option::is_none")]
19    pub description: Option<String>,
20    /// Optional arguments that can be passed to customize the prompt
21    #[serde(skip_serializing_if = "Option::is_none")]
22    pub arguments: Option<Vec<PromptArgument>>,
23}
24
25impl Prompt {
26    /// Create a new prompt with the given name, description and arguments
27    pub fn new<N, D>(
28        name: N,
29        description: Option<D>,
30        arguments: Option<Vec<PromptArgument>>,
31    ) -> Self
32    where
33        N: Into<String>,
34        D: Into<String>,
35    {
36        Prompt {
37            name: name.into(),
38            description: description.map(Into::into),
39            arguments,
40        }
41    }
42}
43
44/// Represents a prompt argument that can be passed to customize the prompt
45#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
46#[cfg_attr(feature = "schemars", derive(schemars::JsonSchema))]
47pub struct PromptArgument {
48    /// The name of the argument
49    pub name: String,
50    /// A description of what the argument is used for
51    #[serde(skip_serializing_if = "Option::is_none")]
52    pub description: Option<String>,
53    /// Whether this argument is required
54    #[serde(skip_serializing_if = "Option::is_none")]
55    pub required: Option<bool>,
56}
57
58/// Represents the role of a message sender in a prompt conversation
59#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
60#[serde(rename_all = "camelCase")]
61#[cfg_attr(feature = "schemars", derive(schemars::JsonSchema))]
62pub enum PromptMessageRole {
63    User,
64    Assistant,
65}
66
67/// Content types that can be included in prompt messages
68#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
69#[serde(tag = "type", rename_all = "camelCase")]
70#[cfg_attr(feature = "schemars", derive(schemars::JsonSchema))]
71pub enum PromptMessageContent {
72    /// Plain text content
73    Text { text: String },
74    /// Image content with base64-encoded data
75    Image {
76        #[serde(flatten)]
77        image: ImageContent,
78    },
79    /// Embedded server-side resource
80    Resource { resource: EmbeddedResource },
81}
82
83impl PromptMessageContent {
84    pub fn text(text: impl Into<String>) -> Self {
85        Self::Text { text: text.into() }
86    }
87}
88
89/// A message in a prompt conversation
90#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
91#[cfg_attr(feature = "schemars", derive(schemars::JsonSchema))]
92pub struct PromptMessage {
93    /// The role of the message sender
94    pub role: PromptMessageRole,
95    /// The content of the message
96    pub content: PromptMessageContent,
97}
98
99impl PromptMessage {
100    /// Create a new text message with the given role and text content
101    pub fn new_text<S: Into<String>>(role: PromptMessageRole, text: S) -> Self {
102        Self {
103            role,
104            content: PromptMessageContent::Text { text: text.into() },
105        }
106    }
107    #[cfg(feature = "base64")]
108    pub fn new_image(
109        role: PromptMessageRole,
110        data: &[u8],
111        mime_type: &str,
112        annotations: Option<Annotations>,
113    ) -> Self {
114        let mime_type = mime_type.into();
115
116        let base64 = BASE64_STANDARD.encode(data);
117
118        Self {
119            role,
120            content: PromptMessageContent::Image {
121                image: RawImageContent {
122                    data: base64,
123                    mime_type,
124                }
125                .optional_annotate(annotations),
126            },
127        }
128    }
129
130    /// Create a new resource message
131    pub fn new_resource(
132        role: PromptMessageRole,
133        uri: String,
134        mime_type: String,
135        text: Option<String>,
136        annotations: Option<Annotations>,
137    ) -> Self {
138        let resource_contents = ResourceContents::TextResourceContents {
139            uri,
140            mime_type: Some(mime_type),
141            text: text.unwrap_or_default(),
142        };
143
144        Self {
145            role,
146            content: PromptMessageContent::Resource {
147                resource: RawEmbeddedResource {
148                    resource: resource_contents,
149                }
150                .optional_annotate(annotations),
151            },
152        }
153    }
154}
155
156/// A template for a prompt
157#[derive(Debug, Serialize, Deserialize)]
158pub struct PromptTemplate {
159    pub id: String,
160    pub template: String,
161    pub arguments: Vec<PromptArgumentTemplate>,
162}
163
164/// A template for a prompt argument, this should be identical to PromptArgument
165#[derive(Debug, Serialize, Deserialize)]
166pub struct PromptArgumentTemplate {
167    pub name: String,
168    pub description: Option<String>,
169    pub required: Option<bool>,
170}