1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68

//! This module defines the `TextPrompt` struct, which represents a text-only prompt without implying the existence of a chat.
//! This is useful for non-chat models. As an added benefit, `TextPrompt` can be used in chat models as well.

use serde::{Deserialize, Serialize};

use super::chat::{ChatMessage, ChatRole};
use super::traits::Prompt;
use crate::PromptTemplate;
use std::fmt;

/// Represents a text-only prompt, without implying the existence of a chat. This is useful for non-chat models.
/// As an added benefit, `TextPrompt` can be used in chat models as well.
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct TextPrompt {
    content: PromptTemplate,
}

impl TextPrompt {
    /// Creates a new `TextPrompt` with the provided template.
    ///
    /// # Arguments
    ///
    /// * `content` - A string that represents the template for the text prompt.
    ///
    /// # Example
    ///
    /// ```
    /// use llm_chain::prompt::text::TextPrompt;
    ///
    /// let text_prompt = TextPrompt::new("Hello, {{name}}!");
    /// ```
    #[cfg(feature = "tera")]
    pub fn new<S: Into<String>>(content: S) -> Self {
        Self {
            content: PromptTemplate::tera(content.into()),
        }
    }
    #[cfg(not(feature = "tera"))]
    pub fn new<S: Into<String>>(content: S) -> Self {
        Self {
            content: PromptTemplate::legacy(content.into()),
        }
    }
}

/// Implement the `Prompt` trait for `TextPrompt`.
impl Prompt for TextPrompt {
    /// Converts a `TextPrompt` into a `Vec<ChatMessage>` where the `TextPrompt` is treated as a user message.
    fn as_chat_prompt(&self) -> Vec<ChatMessage> {
        vec![ChatMessage::from_template(
            ChatRole::User,
            self.content.clone(),
        )]
    }

    /// Returns a reference to the `PromptTemplate` for the `TextPrompt`.
    fn as_text_prompt(&self) -> Option<&PromptTemplate> {
        Some(&self.content)
    }
}

/// Implement `fmt::Display` for `TextPrompt` for pretty-printing.
impl fmt::Display for TextPrompt {
    /// Formats the `TextPrompt` for display.
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        write!(f, "{}", self.content)
    }
}