llama_cpp_v3/

lib.rs

use llama_cpp_sys_v3::{LlamaLib, LoadError};
use std::path::{Path, PathBuf};
use std::sync::Arc;

pub mod backend;
pub mod downloader;

pub use backend::Backend;

#[derive(Debug, thiserror::Error)]
pub enum LlamaError {
    #[error("Failed to load DLL: {0}")]
    DllLoad(#[from] LoadError),
    #[error("Failed to download backend: {0}")]
    Download(#[from] downloader::DownloadError),
    #[error("Failed to initialize backend")]
    BackendInit,
    #[error("Failed to load model from file")]
    ModelLoad,
    #[error("Failed to create context")]
    ContextCreate,
    #[error("Decode error with status code {0}")]
    Decode(i32),
}

pub struct LoadOptions<'a> {
    pub backend: Backend,
    pub app_name: &'a str,
    pub version: Option<&'a str>,        // None = latest
    pub explicit_path: Option<&'a Path>, // Absolute bypass of all resolution
    pub cache_dir: Option<PathBuf>,      // Exact directory to save/load downloaded DLLs
}

/// The initialized Llama capabilities backend.
/// Holds the DLL handle alive.
pub struct LlamaBackend {
    pub lib: Arc<LlamaLib>,
}

impl Drop for LlamaBackend {
    fn drop(&mut self) {
        // Technically backend free is global in llama.cpp
        // We only free if we are the very last reference to the library object.
        if Arc::strong_count(&self.lib) == 1 {
            unsafe {
                (self.lib.symbols.llama_backend_free)();
            }
        }
    }
}

impl LlamaBackend {
    /// Load the specified backend DLL, downloading it from GitHub releases if necessary.
    pub fn load(options: LoadOptions<'_>) -> Result<Self, LlamaError> {
        let dll_path = if let Some(path) = options.explicit_path {
            path.to_path_buf()
        } else if let Ok(env_path) = std::env::var("LLAMA_DLL_PATH") {
            PathBuf::from(env_path)
        } else {
            // Attempt auto-download
            downloader::Downloader::ensure_dll(
                options.backend,
                options.app_name,
                options.version,
                options.cache_dir,
            )?
        };

        if let Some(parent) = dll_path.parent() {
            if let Some(path_ext) = std::env::var_os("PATH") {
                let mut paths = std::env::split_paths(&path_ext).collect::<Vec<_>>();
                let parent_buf = parent.to_path_buf();
                if !paths.contains(&parent_buf) {
                    paths.insert(0, parent_buf);
                    if let Ok(new_path) = std::env::join_paths(paths) {
                        std::env::set_var("PATH", new_path);
                    }
                }
            }
        }

        let lib = LlamaLib::open(&dll_path)?;

        if let Some(parent) = dll_path.parent() {
            let parent_str = parent.to_string_lossy().to_string();
            let c_parent = std::ffi::CString::new(parent_str).unwrap();
            unsafe {
                (lib.symbols.ggml_backend_load_all_from_path)(c_parent.as_ptr());
            }
        } else {
            unsafe {
                (lib.symbols.ggml_backend_load_all)();
            }
        }

        unsafe {
            (lib.symbols.llama_backend_init)();
        }

        Ok(Self { lib: Arc::new(lib) })
    }
}

/// A loaded GGUF model
pub struct LlamaModel {
    pub backend: Arc<LlamaLib>,
    pub handle: *mut llama_cpp_sys_v3::llama_model,
}

impl Drop for LlamaModel {
    fn drop(&mut self) {
        unsafe {
            (self.backend.symbols.llama_model_free)(self.handle);
        }
    }
}

unsafe impl Send for LlamaModel {}
unsafe impl Sync for LlamaModel {}

impl LlamaModel {
    pub fn load_from_file(
        backend: &LlamaBackend,
        path: &str,
        params: llama_cpp_sys_v3::llama_model_params,
    ) -> Result<Self, LlamaError> {
        let c_path = std::ffi::CString::new(path).map_err(|_| LlamaError::ModelLoad)?;
        let handle =
            unsafe { (backend.lib.symbols.llama_model_load_from_file)(c_path.as_ptr(), params) };

        if handle.is_null() {
            return Err(LlamaError::ModelLoad);
        }

        Ok(Self {
            backend: backend.lib.clone(),
            handle,
        })
    }

    pub fn default_params(backend: &LlamaBackend) -> llama_cpp_sys_v3::llama_model_params {
        unsafe { (backend.lib.symbols.llama_model_default_params)() }
    }

    pub fn get_vocab(&self) -> LlamaVocab {
        let handle = unsafe { (self.backend.symbols.llama_model_get_vocab)(self.handle) };
        LlamaVocab {
            backend: self.backend.clone(),
            handle,
        }
    }

    pub fn tokenize(
        &self,
        text: &str,
        add_special: bool,
        parse_special: bool,
    ) -> Result<Vec<llama_cpp_sys_v3::llama_token>, LlamaError> {
        let vocab = self.get_vocab();
        let c_text = std::ffi::CString::new(text).map_err(|_| LlamaError::ModelLoad)?;

        // First call to get required size
        let n_tokens = unsafe {
            (self.backend.symbols.llama_tokenize)(
                vocab.handle,
                c_text.as_ptr(),
                text.len() as i32,
                std::ptr::null_mut(),
                0,
                add_special,
                parse_special,
            )
        };

        if n_tokens < 0 {
            let mut tokens = vec![0; (-n_tokens) as usize];
            let actual_tokens = unsafe {
                (self.backend.symbols.llama_tokenize)(
                    vocab.handle,
                    c_text.as_ptr(),
                    text.len() as i32,
                    tokens.as_mut_ptr(),
                    tokens.len() as i32,
                    add_special,
                    parse_special,
                )
            };
            if actual_tokens < 0 {
                return Err(LlamaError::Decode(actual_tokens));
            }
            tokens.truncate(actual_tokens as usize);
            Ok(tokens)
        } else {
            let mut tokens = vec![0; n_tokens as usize];
            let actual_tokens = unsafe {
                (self.backend.symbols.llama_tokenize)(
                    vocab.handle,
                    c_text.as_ptr(),
                    text.len() as i32,
                    tokens.as_mut_ptr(),
                    tokens.len() as i32,
                    add_special,
                    parse_special,
                )
            };
            if actual_tokens < 0 {
                return Err(LlamaError::Decode(actual_tokens));
            }
            tokens.truncate(actual_tokens as usize);
            Ok(tokens)
        }
    }

    pub fn token_to_piece(&self, token: llama_cpp_sys_v3::llama_token) -> String {
        let vocab = self.get_vocab();
        let mut buf = vec![0u8; 128];
        let n = unsafe {
            (self.backend.symbols.llama_token_to_piece)(
                vocab.handle,
                token,
                buf.as_mut_ptr() as *mut std::ffi::c_char,
                buf.len() as i32,
                0,
                true,
            )
        };

        if n < 0 {
            buf.resize((-n) as usize, 0);
            unsafe {
                (self.backend.symbols.llama_token_to_piece)(
                    vocab.handle,
                    token,
                    buf.as_mut_ptr() as *mut std::ffi::c_char,
                    buf.len() as i32,
                    0,
                    true,
                );
            }
        } else {
            buf.truncate(n as usize);
        }

        String::from_utf8_lossy(&buf).to_string()
    }
}

pub struct LlamaVocab {
    pub backend: Arc<LlamaLib>,
    pub handle: *const llama_cpp_sys_v3::llama_vocab,
}

impl LlamaVocab {
    pub fn bos(&self) -> llama_cpp_sys_v3::llama_token {
        unsafe { (self.backend.symbols.llama_vocab_bos)(self.handle) }
    }

    pub fn eos(&self) -> llama_cpp_sys_v3::llama_token {
        unsafe { (self.backend.symbols.llama_vocab_eos)(self.handle) }
    }

    pub fn is_eog(&self, token: llama_cpp_sys_v3::llama_token) -> bool {
        unsafe { (self.backend.symbols.llama_vocab_is_eog)(self.handle, token) }
    }
}

pub struct LlamaSampler {
    pub backend: Arc<LlamaLib>,
    pub handle: *mut llama_cpp_sys_v3::llama_sampler,
}

impl Drop for LlamaSampler {
    fn drop(&mut self) {
        unsafe {
            (self.backend.symbols.llama_sampler_free)(self.handle);
        }
    }
}

impl LlamaSampler {
    pub fn new_greedy(backend: Arc<LlamaLib>) -> Self {
        let handle = unsafe { (backend.symbols.llama_sampler_init_greedy)() };
        Self { backend, handle }
    }

    pub fn sample(&self, ctx: &LlamaContext, idx: i32) -> llama_cpp_sys_v3::llama_token {
        unsafe { (self.backend.symbols.llama_sampler_sample)(self.handle, ctx.handle, idx) }
    }
}

/// Inference context attached to a model
pub struct LlamaContext {
    pub backend: Arc<LlamaLib>,
    pub handle: *mut llama_cpp_sys_v3::llama_context,
}

impl Drop for LlamaContext {
    fn drop(&mut self) {
        unsafe {
            (self.backend.symbols.llama_free)(self.handle);
        }
    }
}

impl LlamaContext {
    pub fn new(
        model: &LlamaModel,
        params: llama_cpp_sys_v3::llama_context_params,
    ) -> Result<Self, LlamaError> {
        let handle = unsafe { (model.backend.symbols.llama_init_from_model)(model.handle, params) };

        if handle.is_null() {
            return Err(LlamaError::ContextCreate);
        }

        Ok(Self {
            backend: model.backend.clone(),
            handle,
        })
    }

    pub fn default_params(model: &LlamaModel) -> llama_cpp_sys_v3::llama_context_params {
        unsafe { (model.backend.symbols.llama_context_default_params)() }
    }

    pub fn decode(&mut self, batch: llama_cpp_sys_v3::llama_batch) -> Result<(), LlamaError> {
        let res = unsafe { (self.backend.symbols.llama_decode)(self.handle, batch) };
        if res != 0 {
            Err(LlamaError::Decode(res))
        } else {
            Ok(())
        }
    }
}