kproc-llm 0.7.0

//! Module for using `ollama`.

use ccutils::streams::pin_stream;
use futures::StreamExt;
use ollama_rs::{
  error::OllamaError,
  generation::{
    chat::{request::ChatMessageRequest, ChatMessage, MessageRole},
    completion::request::GenerationRequest,
    parameters::FormatType,
  },
};
use std::future::Future;

use crate::prelude::*;

/// Interface to Ollama
pub struct Ollama
{
  ollama: ollama_rs::Ollama,
  model: String,
}

impl Default for Ollama
{
  fn default() -> Self
  {
    Self {
      ollama: Default::default(),
      model: "llama3:instruct".to_string(),
    }
  }
}

impl Ollama
{
  /// Create a new LLM using ollama and the given model.
  pub fn new(ollama: ollama_rs::Ollama, model: impl Into<String>) -> Self
  {
    let model = model.into();
    Self { ollama, model }
  }
  /// Create a new LLM using ollama default configuration and the given model
  pub fn from_model(model: impl Into<String>) -> Self
  {
    Self::new(ollama_rs::Ollama::default(), model)
  }
}

impl TryInto<Option<FormatType>> for crate::Format
{
  type Error = Error;
  fn try_into(self) -> std::result::Result<Option<FormatType>, Self::Error>
  {
    match self
    {
      Self::Text => Ok(None),
      Self::Json => Ok(Some(FormatType::Json)),
    }
  }
}

impl LargeLanguageModel for Ollama
{
  fn chat_stream(
    &self,
    prompt: ChatPrompt,
  ) -> Result<impl Future<Output = Result<StringStream>> + Send>
  {
    let req = ChatMessageRequest::new(
      self.model.to_owned(),
      prompt
        .messages
        .into_iter()
        .map(|msg| {
          let role = match msg.role
          {
            Role::User | Role::Custom(_) => MessageRole::User,
            Role::System => MessageRole::System,
            Role::Assistant => MessageRole::Assistant,
          };
          ChatMessage {
            role,
            content: msg.content,
            tool_calls: Default::default(),
            images: Default::default(),
            thinking: None,
          }
        })
        .collect(),
    );
    Ok(async {
      let stream = self.ollama.send_chat_messages_stream(req).await?;
      Ok(pin_stream(stream.map(|r| {
        r.map(|x| x.message.content)
          .map_err(|_| Error::OllamaError(OllamaError::Other("Internal error".into())))
      })))
    })
  }
  fn generate_stream(
    &self,
    prompt: GenerationPrompt,
  ) -> Result<impl Future<Output = Result<StringStream>> + Send>
  {
    let req = GenerationRequest::new(self.model.to_owned(), prompt.user);
    let req = if let Some(system) = prompt.system
    {
      req.system(system)
    }
    else
    {
      req
    };
    let req = if let Some(format) = prompt.options.format.try_into()?
    {
      req.format(format)
    }
    else
    {
      req
    };
    #[cfg(feature = "image")]
    let req = if let Some(image) = prompt.image
    {
      req.images([image.to_ollama()?].into())
    }
    else
    {
      req
    };
    let req = req.think(prompt.options.thinking);
    Ok(async {
      let stream = self.ollama.generate_stream(req).await?;
      Ok(pin_stream(stream.flat_map(|r| {
        futures::stream::iter(match r
        {
          Ok(v) => v.into_iter().map(|x| Ok(x.response)).collect(),
          Err(e) => vec![Err(e.into())],
        })
      })))
    })
  }
}