fim_engine/
lib.rs

1//! `fim-engine` — embedded fill-in-the-middle code completion.
2//!
3//! A self-contained local code-completion engine: it downloads a small
4//! quantized [qwen2.5-coder] model on first use, caches it, and runs
5//! inference in-process via [candle] — no external daemon, no API key,
6//! no network after the one-time download.
7//!
8//! Shared by mnml + tmnl. Typical use:
9//!
10//! ```no_run
11//! use fim_engine::{FimEngine, ModelChoice};
12//!
13//! // Blocking — do this on a worker thread.
14//! let cache = fim_engine::default_cache_dir();
15//! let mut engine = FimEngine::load(&cache, ModelChoice::Qwen1_5B, &|p| {
16//!     eprintln!("{}: {}/{:?}", p.label, p.received, p.total);
17//! })?;
18//! let completion = engine.complete("fn add(a: i32, b: i32) -> i32 {\n    ", "\n}", 64)?;
19//! # Ok::<(), String>(())
20//! ```
21//!
22//! [qwen2.5-coder]: https://huggingface.co/Qwen/Qwen2.5-Coder-1.5B
23//! [candle]: https://github.com/huggingface/candle
24
25mod download;
26mod infer;
27
28pub use download::{DownloadProgress, ModelChoice, ModelPaths, is_model_cached};
29
30use std::path::{Path, PathBuf};
31
32/// The canonical, host-agnostic model cache directory — every consumer
33/// (mnml, tmnl, …) should pass this to [`FimEngine::load`] so the
34/// ~1 GB download is shared, not duplicated per app.
35///
36/// `$XDG_CACHE_HOME/fim-engine` when set, else `~/.cache/fim-engine`,
37/// else `./.fim-engine-cache` as a last resort.
38pub fn default_cache_dir() -> PathBuf {
39    if let Ok(xdg) = std::env::var("XDG_CACHE_HOME")
40        && !xdg.is_empty()
41    {
42        return PathBuf::from(xdg).join("fim-engine");
43    }
44    if let Ok(home) = std::env::var("HOME")
45        && !home.is_empty()
46    {
47        return PathBuf::from(home).join(".cache").join("fim-engine");
48    }
49    PathBuf::from(".fim-engine-cache")
50}
51
52/// A loaded local FIM completion engine. Holds the model in memory;
53/// keep one alive and call [`FimEngine::complete`] repeatedly.
54pub struct FimEngine {
55    model: infer::Model,
56}
57
58impl FimEngine {
59    /// Download (if needed) + load the `choice` model. `cache_dir` is
60    /// where the GGUF weights + tokenizer are cached (see
61    /// [`default_cache_dir`]). `progress` fires periodically while
62    /// files download.
63    ///
64    /// Blocking and slow on the first call (a ~1 GB download); fast
65    /// afterwards (just the load). Run it on a worker thread.
66    pub fn load(
67        cache_dir: &Path,
68        choice: ModelChoice,
69        progress: &(dyn Fn(DownloadProgress) + Sync),
70    ) -> Result<Self, String> {
71        let paths = download::ensure_model(cache_dir, choice, progress)?;
72        let model = infer::Model::load(&paths.gguf, &paths.tokenizer)?;
73        Ok(FimEngine { model })
74    }
75
76    /// Generate a completion for the cursor sitting between `prefix`
77    /// (code before) and `suffix` (code after). Returns the text to
78    /// insert — never includes the surrounding code. `max_tokens`
79    /// bounds the length (≈ 64 is a good inline default).
80    ///
81    /// Blocking + CPU-bound (~100–400 ms for the 1.5B model). Call on
82    /// a worker thread; never on the UI thread.
83    pub fn complete(
84        &mut self,
85        prefix: &str,
86        suffix: &str,
87        max_tokens: usize,
88    ) -> Result<String, String> {
89        self.model.complete(prefix, suffix, max_tokens)
90    }
91}
fim_engine/lib.rs

fim_engine/
lib.rs
