quaynor 2.2.0

Lightweight local AI inference engine: load GGUF models and chat on-device with streaming, tool calling, embeddings, and reranking
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95

pub mod chat;
pub mod crossencoder;
pub mod encoder;
pub mod errors;
pub mod llm;
pub mod memory;
pub mod sampler_config;
pub mod template;
pub mod tokenizer;
pub mod tool_calling;

pub fn send_llamacpp_logs_to_tracing() {
    llama_cpp_2::send_logs_to_tracing(llama_cpp_2::LogOptions::default().with_logs_enabled(true));
}

#[cfg(test)]
pub mod test_utils {
    use crate::llm::{get_model, Model};
    use std::sync::{Arc, Once};

    static INIT: Once = Once::new();

    /// Initialize tracing for tests
    pub fn init_test_tracing() {
        INIT.call_once(|| {
            tracing_subscriber::fmt()
                .with_max_level(tracing::Level::TRACE)
                .with_timer(tracing_subscriber::fmt::time::uptime())
                .with_span_events(tracing_subscriber::fmt::format::FmtSpan::CLOSE)
                .try_init()
                .ok();
        });
    }

    /// Get path to test model from TEST_MODEL env var or default to "model.gguf"
    pub fn test_model_path() -> String {
        std::env::var("TEST_MODEL").unwrap_or_else(|_| "model.gguf".to_string())
    }

    /// Get path to test embeddings model from TEST_EMBEDDINGS_MODEL env var
    pub fn test_embeddings_model_path() -> String {
        std::env::var("TEST_EMBEDDINGS_MODEL").unwrap_or_else(|_| "embeddings.gguf".to_string())
    }

    /// Get path to test crossencoder model from TEST_CROSSENCODER_MODEL env var
    pub fn test_crossencoder_model_path() -> String {
        std::env::var("TEST_CROSSENCODER_MODEL").unwrap_or_else(|_| "crossencoder.gguf".to_string())
    }

    /// Get path to test multimodal projector from TEST_MMPROJ env var
    pub fn test_mmproj_path() -> String {
        std::env::var("TEST_MMPROJ").unwrap_or_else(|_| "mmproj.gguf".to_string())
    }

    /// Get path to test vision model from TEST_VISION_MODEL env var
    pub fn test_vision_model_path() -> String {
        std::env::var("TEST_VISION_MODEL").unwrap_or_else(|_| "vision-model.gguf".to_string())
    }

    /// Load the test model with GPU acceleration if available
    pub fn load_test_model() -> Arc<Model> {
        let path = test_model_path();
        Arc::new(
            get_model(&path, true, None)
                .unwrap_or_else(|e| panic!("Failed to load test model from {}: {:?}", path, e)),
        )
    }

    /// Load the embeddings model with GPU acceleration if available
    pub fn load_embeddings_model() -> Arc<Model> {
        let path = test_embeddings_model_path();
        // XXX: loading the embeddings model for unit tests without GPU offloading
        //      because it otherwise caused a segfault specifically with the llvmpipe vulkan driver.
        //      (which is used in the nix sandbox, since we don't have access to the host GPU)
        //      llvmpipe is very rare in the wild, so it shouldn't cause any problems in general
        //      this segfault doesn't happen on quaynor commit 94d51c5.
        //      it's most likely related to an upstream change in llama.cpp
        Arc::new(
            get_model(&path, false, None).unwrap_or_else(|e| {
                panic!("Failed to load embeddings model from {}: {:?}", path, e)
            }),
        )
    }

    /// Load the crossencoder model with GPU acceleration if available
    pub fn load_crossencoder_model() -> Arc<Model> {
        let path = test_crossencoder_model_path();
        // Same GPU offloading note as embeddings model
        Arc::new(
            get_model(&path, false, None).unwrap_or_else(|e| {
                panic!("Failed to load crossencoder model from {}: {:?}", path, e)
            }),
        )
    }
}