kproc-llm 0.7.0

use std::collections::HashMap;

use serde::{Deserialize, Serialize};
use smart_default::SmartDefault as Default;

use crate::Message;

use crate::prelude::*;

/// Format
#[derive(Debug, Clone, Copy, Default, Deserialize, Serialize)]
pub enum Format
{
  /// Output text
  #[default]
  Text,
  /// Output Json
  Json,
}

/// Generation options
#[derive(Debug, Default, Deserialize, Serialize)]
pub struct Options
{
  /// Requested output format
  #[default(Format::Text)]
  pub(crate) format: Format,
  /// Enable thinking mode (if the model support)
  #[default(false)]
  pub(crate) thinking: bool,
}

/// Prompt for a chat session
#[derive(Debug, Default, Deserialize, Serialize)]
pub struct ChatPrompt
{
  /// Messages
  pub(crate) messages: Vec<Message>,
  /// Generation options
  #[serde(default)]
  pub(crate) options: Options,
  /// Extraxt context passed to the template
  pub(crate) template_context: HashMap<String, minijinja::Value>,
}

impl ChatPrompt
{
  /// 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.options.format = format.into();
    self
  }
  /// Enable thinking
  pub fn thiking(mut self, thinking: bool) -> Self
  {
    self.options.thinking = thinking;
    self
  }
  /// Set options
  pub fn options(mut self, options: Options) -> Self
  {
    self.options = options;
    self
  }
  /// Add to template context
  pub fn template_context(mut self, key: String, value: impl Into<minijinja::Value>) -> Self
  {
    self.template_context.insert(key, value.into());
    self
  }
}

/// Prompt
#[derive(Debug)]
pub struct GenerationPrompt
{
  /// Messages
  pub(crate) user: String,
  /// Messages
  pub(crate) system: Option<String>,
  /// Messages
  pub(crate) assistant: Option<String>,
  /// Requested generation options
  pub(crate) options: Options,
  /// Extraxt context passed to the template
  pub(crate) template_context: HashMap<String, minijinja::Value>,
  /// Optional prompt image (for multi modal models)
  #[cfg(feature = "image")]
  pub(crate) 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(),
      options: Default::default(),
      template_context: Default::default(),
      #[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.options.format = format.into();
    self
  }
  /// Enable thinking
  pub fn thinking(mut self, thinking: bool) -> Self
  {
    self.options.thinking = thinking;
    self
  }
  /// Add to template context
  pub fn template_context(mut self, key: String, value: impl Into<minijinja::Value>) -> Self
  {
    self.template_context.insert(key, value.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
  }
}