netget/llm/
embedded_inference.rs

1//! Embedded LLM inference using llama.cpp (via llama-cpp-2 Rust bindings)
2//!
3//! Provides local LLM inference as a fallback when Ollama is unavailable.
4//! Uses GGUF model files loaded from disk.
5
6use anyhow::{anyhow, Context, Result};
7use llama_cpp_2::context::params::LlamaContextParams;
8use llama_cpp_2::llama_backend::LlamaBackend;
9use llama_cpp_2::llama_batch::LlamaBatch;
10use llama_cpp_2::model::params::LlamaModelParams;
11use llama_cpp_2::model::LlamaModel;
12use llama_cpp_2::model::{AddBos, Special};
13use llama_cpp_2::sampling::LlamaSampler;
14use std::num::NonZeroU32;
15use std::path::{Path, PathBuf};
16use std::sync::Arc;
17use tracing::{debug, info, trace};
18
19/// Configuration for embedded LLM inference
20#[derive(Debug, Clone)]
21pub struct InferenceConfig {
22    /// Context size (number of tokens)
23    pub context_size: u32,
24    /// Maximum tokens to generate
25    pub max_tokens: usize,
26    /// Temperature (0.0 = deterministic, 1.0+ = creative)
27    pub temperature: f32,
28    /// Top-p (nucleus sampling)
29    pub top_p: f32,
30    /// Number of GPU layers to offload (0 = CPU only, u32::MAX = auto-detect)
31    pub n_gpu_layers: u32,
32    /// Number of threads (0 = auto-detect)
33    pub n_threads: u32,
34}
35
36impl Default for InferenceConfig {
37    fn default() -> Self {
38        Self {
39            context_size: 4096,
40            max_tokens: 2048,
41            temperature: 0.7,
42            top_p: 0.9,
43            n_gpu_layers: u32::MAX, // Auto-detect and use all available GPU layers
44            n_threads: 0,           // Auto-detect optimal thread count
45        }
46    }
47}
48
49/// Model information for debugging and display
50#[derive(Debug, Clone)]
51pub struct ModelInfo {
52    pub path: PathBuf,
53    pub context_size: u32,
54    pub max_tokens: usize,
55    pub backend: String,
56    pub vocab_size: Option<i32>,
57}
58
59/// Embedded LLM backend using llama.cpp
60pub struct EmbeddedLLMBackend {
61    model: Arc<LlamaModel>,
62    config: InferenceConfig,
63    model_path: PathBuf,
64    /// Global llama.cpp backend (must be initialized once)
65    _backend: Arc<LlamaBackend>,
66}
67
68impl EmbeddedLLMBackend {
69    /// Load GGUF model from disk
70    ///
71    /// # Arguments
72    /// * `model_path` - Path to GGUF model file (e.g., "mistral-7b.Q4_K_M.gguf")
73    ///
74    /// # Example
75    /// ```no_run
76    /// # use netget::llm::embedded_inference::EmbeddedLLMBackend;
77    /// # async fn example() -> anyhow::Result<()> {
78    /// let backend = EmbeddedLLMBackend::new("./models/mistral-7b.Q4_K_M.gguf").await?;
79    /// # Ok(())
80    /// # }
81    /// ```
82    pub async fn new(model_path: impl AsRef<Path>) -> Result<Self> {
83        Self::new_with_config(model_path, InferenceConfig::default()).await
84    }
85
86    /// Load model with custom configuration
87    pub async fn new_with_config(
88        model_path: impl AsRef<Path>,
89        config: InferenceConfig,
90    ) -> Result<Self> {
91        let path = model_path.as_ref();
92        if !path.exists() {
93            return Err(anyhow!(
94                "Model file not found: {}. Download a GGUF model from https://huggingface.co/TheBloke",
95                path.display()
96            ));
97        }
98
99        info!("Loading embedded LLM from: {}", path.display());
100        debug!(
101            "Config: context_size={}, max_tokens={}, temp={}, top_p={}, gpu_layers={}",
102            config.context_size,
103            config.max_tokens,
104            config.temperature,
105            config.top_p,
106            if config.n_gpu_layers == u32::MAX {
107                "auto".to_string()
108            } else {
109                config.n_gpu_layers.to_string()
110            }
111        );
112
113        // Initialize llama.cpp backend (must be done once globally)
114        // This is safe to call multiple times - llama.cpp handles initialization internally
115        let backend =
116            Arc::new(LlamaBackend::init().context("Failed to initialize llama.cpp backend")?);
117
118        // Load model (blocking operation, run in spawn_blocking)
119        let model_path_clone = path.to_path_buf();
120        let n_gpu_layers = config.n_gpu_layers;
121        let backend_clone = backend.clone();
122
123        let model = tokio::task::spawn_blocking(move || {
124            // Create model params with GPU layers
125            let model_params = LlamaModelParams::default().with_n_gpu_layers(n_gpu_layers);
126
127            // Load model from file
128            LlamaModel::load_from_file(&backend_clone, model_path_clone, &model_params)
129                .context("Failed to load GGUF model")
130        })
131        .await??;
132
133        let vocab_size = model.n_vocab();
134        info!(
135            "Embedded LLM loaded successfully (vocab_size: {}, context_size: {})",
136            vocab_size, config.context_size
137        );
138
139        Ok(Self {
140            model: Arc::new(model),
141            config,
142            model_path: path.to_path_buf(),
143            _backend: backend,
144        })
145    }
146
147    /// Generate completion from prompt (non-streaming)
148    ///
149    /// # Arguments
150    /// * `prompt` - Input text prompt
151    ///
152    /// # Returns
153    /// Generated text completion
154    ///
155    /// # Example
156    /// ```no_run
157    /// # use netget::llm::embedded_inference::EmbeddedLLMBackend;
158    /// # async fn example() -> anyhow::Result<()> {
159    /// # let backend = EmbeddedLLMBackend::new("model.gguf").await?;
160    /// let response = backend.generate(&"Hello, how are you?").await?;
161    /// println!("Response: {}", response);
162    /// # Ok(())
163    /// # }
164    /// ```
165    pub async fn generate(&self, prompt: &str) -> Result<String> {
166        let model = self.model.clone();
167        let prompt = prompt.to_string();
168        let config = self.config.clone();
169
170        // Run inference in blocking task (CPU/GPU-intensive)
171        let response =
172            tokio::task::spawn_blocking(move || Self::run_inference_sync(&model, &prompt, &config))
173                .await
174                .context("Inference task panicked")??;
175
176        Ok(response)
177    }
178
179    /// Synchronous inference (runs in blocking task)
180    fn run_inference_sync(
181        model: &LlamaModel,
182        prompt: &str,
183        config: &InferenceConfig,
184    ) -> Result<String> {
185        // Create context params
186        let n_ctx =
187            NonZeroU32::new(config.context_size).context("Context size must be non-zero")?;
188
189        let mut ctx_params = LlamaContextParams::default().with_n_ctx(Some(n_ctx));
190
191        if config.n_threads > 0 {
192            ctx_params = ctx_params.with_n_threads(config.n_threads as i32);
193        }
194
195        // Create inference context
196        let backend = LlamaBackend::init()?;
197        let mut ctx = model
198            .new_context(&backend, ctx_params)
199            .context("Failed to create inference context")?;
200
201        // Create sampler (greedy for now, can be enhanced later)
202        let mut sampler = LlamaSampler::greedy();
203
204        // Tokenize prompt
205        let tokens = model
206            .str_to_token(prompt, AddBos::Always)
207            .context("Failed to tokenize prompt")?;
208
209        debug!("Tokenized prompt: {} tokens", tokens.len());
210
211        if tokens.len() >= config.context_size as usize {
212            return Err(anyhow!(
213                "Prompt too long: {} tokens (max: {})",
214                tokens.len(),
215                config.context_size
216            ));
217        }
218
219        // Prepare batch
220        let mut batch = LlamaBatch::new(config.context_size as usize, 1);
221
222        // Add prompt tokens to batch
223        for (i, token) in tokens.iter().enumerate() {
224            let is_last = i == tokens.len() - 1;
225            batch.add(*token, i as i32, &[0], is_last)?;
226        }
227
228        // Decode prompt
229        ctx.decode(&mut batch).context("Failed to decode prompt")?;
230
231        // Generate tokens
232        let mut generated_text = String::new();
233        let mut n_cur = tokens.len();
234        let mut n_decode = 0;
235
236        while n_cur < config.context_size as usize && n_decode < config.max_tokens {
237            // Sample next token using sampler
238            let new_token_id = sampler.sample(&ctx, batch.n_tokens() - 1);
239
240            // Check for end of generation
241            if model.is_eog_token(new_token_id) {
242                trace!("End of generation token encountered");
243                break;
244            }
245
246            // Decode token to text (Special::Tokenize means render special tokens)
247            let output_bytes = model
248                .token_to_bytes(new_token_id, Special::Tokenize)
249                .context("Failed to decode token")?;
250            let output_str = String::from_utf8_lossy(&output_bytes);
251            generated_text.push_str(&output_str);
252
253            trace!("Generated token: {:?}", output_str);
254
255            // Prepare next batch
256            batch.clear();
257            batch.add(new_token_id, n_cur as i32, &[0], true)?;
258
259            // Decode
260            ctx.decode(&mut batch).context("Failed to decode batch")?;
261
262            n_cur += 1;
263            n_decode += 1;
264        }
265
266        debug!("Generation complete: {} tokens generated", n_decode);
267
268        Ok(generated_text.trim().to_string())
269    }
270
271    /// Get model information (useful for debugging and display)
272    pub fn get_model_info(&self) -> ModelInfo {
273        ModelInfo {
274            path: self.model_path.clone(),
275            context_size: self.config.context_size,
276            max_tokens: self.config.max_tokens,
277            backend: "llama.cpp".to_string(),
278            vocab_size: Some(self.model.n_vocab()),
279        }
280    }
281
282    /// Get model path
283    pub fn model_path(&self) -> &Path {
284        &self.model_path
285    }
286
287    /// Check if model is loaded and ready
288    pub fn is_ready(&self) -> bool {
289        true // If we constructed successfully, we're ready
290    }
291}
netget/llm/embedded_inference.rs

netget/llm/
embedded_inference.rs
