kproc-llm 0.7.0

use std::io::Write as _;

use clap::Parser;
use kproc_llm::prelude::*;

trait Example
{
  type Model: kproc_llm::LargeLanguageModel + Send + Sync;
  async fn create_model() -> Self::Model;
  async fn generate(args: Args) -> kproc_llm::Result<()>
  {
    use kproc_llm::LargeLanguageModel;
    let model = Self::create_model().await;

    let stream = if let Some(yaml_file) = args.from_yaml
    {
      let chat = serde_saphyr::from_slice(&std::fs::read(yaml_file).unwrap()).unwrap();
      model.chat_stream(chat)?.await?
    }
    else
    {
      let prompt = kproc_llm::GenerationPrompt::prompt(args.prompt).thinking(args.thinking);
      model.generate_stream(prompt)?.await?
    };
    let mut stream = if args.thinking
    {
      stream
    }
    else
    {
      ccutils::streams::filter_tag(stream, "think")
    };

    use futures::stream::StreamExt;
    print!("Answer: ```");
    while let Some(next_token) = stream.next().await
    {
      print!("{}", next_token?);
      let _ = std::io::stdout().flush();
    }
    println!("```");
    Ok(())
  }
}

#[cfg(feature = "ollama")]
struct OllamaSimple {}

#[cfg(feature = "ollama")]
impl Example for OllamaSimple
{
  type Model = kproc_llm::ollama::Ollama;
  async fn create_model() -> Self::Model
  {
    ollama::Ollama::from_model("llama3:instruct")
  }
}

#[cfg(feature = "llama.cpp")]
struct LlamaCppSimple {}

#[cfg(feature = "llama.cpp")]
impl Example for LlamaCppSimple
{
  type Model = kproc_llm::llama_cpp::LlamaCpp;
  async fn create_model() -> Self::Model
  {
    llama_cpp::LlamaCpp::from_model(
      "hf://ggml-org/Meta-Llama-3.1-8B-Instruct-Q4_0-GGUF/meta-llama-3.1-8b-instruct-q4_0.gguf",
    )
    .unwrap()
  }
}

#[cfg(feature = "simple-api")]
struct SimpleApi {}

#[cfg(feature = "simple-api")]
impl Example for SimpleApi
{
  type Model = kproc_llm::simple_api::SimpleApi<yaaral::tokio::Runtime>;
  async fn create_model() -> Self::Model
  {
    simple_api::SimpleApi::new(
      yaaral::tokio::Runtime::current(),
      "http://localhost",
      8080,
      None,
    )
    .unwrap()
  }
}

#[cfg(any(feature = "candle", feature = "candle-git"))]
struct CandleLlama;

#[cfg(any(feature = "candle", feature = "candle-git"))]
impl Example for CandleLlama
{
  type Model = kproc_llm::candle::Candle;
  async fn create_model() -> Self::Model
  {
    kproc_llm::candle::factory::llama2_7b()
      .build()
      .await
      .unwrap()
  }
}

#[cfg(feature = "candle")]
struct CandleSmol;

#[cfg(feature = "candle")]
impl Example for CandleSmol
{
  type Model = kproc_llm::candle::Candle;
  async fn create_model() -> Self::Model
  {
    kproc_llm::candle::factory::smol_lm_2_1_7b()
      .build()
      .await
      .unwrap()
  }
}

#[cfg(feature = "candle-git")]
struct CandleSmol;

#[cfg(feature = "candle-git")]
impl Example for CandleSmol
{
  type Model = kproc_llm::candle::Candle;
  async fn create_model() -> Self::Model
  {
    kproc_llm::candle::factory::smol_lm_3_q4_k_m()
      .build()
      .await
      .unwrap()
  }
}

#[derive(Parser)]
#[command(author, version, about, long_about = None)]
struct Args
{
  #[arg(long, value_name = "BACKEND", help = "Sets the backend to use.")]
  backend: String,

  #[arg(long, help = "Enable thinking.")]
  thinking: bool,

  #[arg(long, value_name = "YAML_FILE_NAME")]
  from_yaml: Option<String>,

  #[arg(value_name = "PROMPT", default_value = "")]
  prompt: String,
}

#[tokio::main]
async fn main()
{
  colog::init();
  let args = Args::parse();
  match args.backend.as_str()
  {
    #[cfg(feature = "ollama")]
    "ollama" => OllamaSimple::generate(args).await.unwrap(),
    #[cfg(feature = "llama.cpp")]
    "llama.cpp" => LlamaCppSimple::generate(args).await.unwrap(),
    #[cfg(feature = "simple-api")]
    "simple-api" => SimpleApi::generate(args).await.unwrap(),
    #[cfg(any(feature = "candle", feature = "candle-git"))]
    "candle-llama" => CandleLlama::generate(args).await.unwrap(),
    #[cfg(feature = "candle")]
    "candle-smol" => CandleSmol::generate(args).await.unwrap(),
    #[cfg(feature = "candle-git")]
    "candle-smol" => CandleSmol::generate(args).await.unwrap(),
    _ => log::error!("Unknown backend '{}'.", args.backend),
  }
}