kproc-llm 0.6.1

use smart_default::SmartDefault as Default;

use crate::Message;

use crate::prelude::*;

/// Prompt for a chat session
#[derive(Debug, Default)]
pub struct ChatPrompt
{
  /// Messages
  pub messages: Vec<Message>,
  /// Requested output format
  #[default(Format::Text)]
  pub format: Format,
}

impl ChatPrompt
{
  /// Start an empty chat prompt.
  pub fn new() -> Self
  {
    Self {
      messages: Default::default(),
      format: Format::Text,
    }
  }
  /// Create a new message, from the given role and content.
  pub fn message(mut self, role: Role, content: impl Into<String>) -> Self
  {
    self.messages.push(Message {
      role,
      content: content.into(),
    });
    self
  }
  /// Create a new message, from the given role and content. Ignore if content is null.
  pub fn message_opt(mut self, role: Role, content: Option<String>) -> Self
  {
    if let Some(content) = content
    {
      self.messages.push(Message { role, content });
    }
    self
  }
  /// Create a new prompt, from the given string.
  pub fn user(self, content: impl Into<String>) -> Self
  {
    self.message(Role::User, content)
  }
  /// Set the system hint.
  pub fn system(self, content: impl Into<String>) -> Self
  {
    self.message(Role::System, content)
  }
  /// Set the system hint.
  pub fn system_opt(self, content: Option<String>) -> Self
  {
    self.message_opt(Role::System, content)
  }
  /// Set the system hint.
  pub fn assistant(self, content: impl Into<String>) -> Self
  {
    self.message(Role::Assistant, content)
  }
  /// Set the system hint.
  pub fn assistant_opt(self, content: Option<String>) -> Self
  {
    self.message_opt(Role::Assistant, content)
  }
  /// Set the result format
  pub fn format(mut self, format: impl Into<Format>) -> Self
  {
    self.format = format.into();
    self
  }
}

/// Prompt
#[derive(Debug)]
pub struct GenerationPrompt
{
  /// Messages
  pub user: String,
  /// Messages
  pub system: Option<String>,
  /// Messages
  pub assistant: Option<String>,
  /// Requested output format
  pub format: Format,
  /// Optional prompt image (for multi modal models)
  #[cfg(feature = "image")]
  pub image: Option<kproc_values::Image>,
}

impl GenerationPrompt
{
  /// Start creation of generation prompt, with user prompt
  pub fn prompt(user: impl Into<String>) -> Self
  {
    Self {
      user: user.into(),
      system: Default::default(),
      assistant: Default::default(),
      format: Format::Text,
      #[cfg(feature = "image")]
      image: None,
    }
  }
  /// Set the system hint.
  pub fn system(mut self, content: impl Into<String>) -> Self
  {
    self.system = Some(content.into());
    self
  }
  /// Set the assistant hint.
  pub fn assistant(mut self, content: impl Into<String>) -> Self
  {
    self.assistant = Some(content.into());
    self
  }
  /// Set the result format
  pub fn format(mut self, format: impl Into<Format>) -> Self
  {
    self.format = format.into();
    self
  }
  /// Set the image, for use in multi-modal models
  #[cfg(feature = "image")]
  pub fn image(mut self, image: impl Into<kproc_values::Image>) -> Self
  {
    self.image = Some(image.into());
    self
  }
}

impl From<GenerationPrompt> for Vec<Message>
{
  fn from(value: GenerationPrompt) -> Self
  {
    let mut vec = Self::default();
    if let Some(system) = value.system
    {
      vec.push(Message {
        role: Role::System,
        content: system,
      });
    }
    if let Some(assistant) = value.assistant
    {
      vec.push(Message {
        role: Role::Assistant,
        content: assistant,
      });
    }
    vec.push(Message {
      role: Role::User,
      content: value.user,
    });
    vec
  }
}