rig_core/providers/
llamafile.rs

1//! Llamafile API client and Rig integration
2//!
3//! [Llamafile](https://github.com/Mozilla-Ocho/llamafile) is a Mozilla Builders project
4//! that distributes LLMs as single-file executables. When started, it exposes an
5//! OpenAI-compatible API at `http://localhost:8080/v1`.
6//!
7//! # Example
8//! ```no_run
9//! use rig_core::{
10//!     client::CompletionClient,
11//!     completion::CompletionModel,
12//!     providers::llamafile,
13//! };
14//!
15//! # async fn run() -> Result<(), Box<dyn std::error::Error>> {
16//! // Create a new Llamafile client (defaults to http://localhost:8080)
17//! let client = llamafile::Client::from_url("http://localhost:8080")?;
18//!
19//! // Send a completion request with a preamble.
20//! let model = client.completion_model(llamafile::LLAMA_CPP);
21//! let request = model
22//!     .completion_request("Hello!")
23//!     .preamble("You are a helpful assistant.".to_string())
24//!     .build();
25//! let response = model.completion(request).await?;
26//! println!("{:?}", response.choice);
27//! # Ok(())
28//! # }
29//! ```
30
31use crate::client::{
32    self, Capabilities, Capable, DebugExt, Nothing, Provider, ProviderBuilder, ProviderClient,
33    Transport,
34};
35use crate::http_client::{self, HttpClientExt};
36use crate::providers::openai;
37
38// ================================================================
39// Main Llamafile Client
40// ================================================================
41const LLAMAFILE_API_BASE_URL: &str = "http://localhost:8080";
42
43/// The default model identifier reported by llamafile.
44pub const LLAMA_CPP: &str = "LLaMA_CPP";
45
46#[derive(Debug, Default, Clone, Copy)]
47pub struct LlamafileExt;
48
49#[derive(Debug, Default, Clone, Copy)]
50pub struct LlamafileBuilder;
51
52impl Provider for LlamafileExt {
53    type Builder = LlamafileBuilder;
54    const VERIFY_PATH: &'static str = "/models";
55
56    // Llamafile clients are constructed from a bare host URL
57    // (e.g. `http://localhost:8080`) while the shared OpenAI-compatible
58    // endpoints are relative to `/v1`.
59    fn build_uri(&self, base_url: &str, path: &str, _transport: Transport) -> String {
60        let base_url = base_url.trim_end_matches('/');
61        format!("{base_url}/v1/{}", path.trim_start_matches('/'))
62    }
63}
64
65impl openai::completion::OpenAICompatibleProvider for LlamafileExt {
66    const PROVIDER_NAME: &'static str = "llamafile";
67
68    type StreamingUsage = openai::Usage;
69
70    // llama.cpp-based servers can emit a whole tool call in one streaming chunk.
71    const EMITS_COMPLETE_SINGLE_CHUNK_TOOL_CALLS: bool = true;
72
73    type Response = openai::CompletionResponse;
74}
75
76impl openai::embedding::OpenAIEmbeddingsCompatible for LlamafileExt {
77    const PROVIDER_NAME: &'static str = "llamafile";
78}
79
80impl<H> Capabilities<H> for LlamafileExt {
81    type Completion = Capable<openai::completion::GenericCompletionModel<LlamafileExt, H>>;
82    type Embeddings = Capable<openai::embedding::GenericEmbeddingModel<LlamafileExt, H>>;
83    type Transcription = Nothing;
84    type ModelListing = Nothing;
85    #[cfg(feature = "image")]
86    type ImageGeneration = Nothing;
87    #[cfg(feature = "audio")]
88    type AudioGeneration = Nothing;
89    type Rerank = Nothing;
90}
91
92impl DebugExt for LlamafileExt {}
93
94impl ProviderBuilder for LlamafileBuilder {
95    type Extension<H>
96        = LlamafileExt
97    where
98        H: HttpClientExt;
99    type ApiKey = Nothing;
100
101    const BASE_URL: &'static str = LLAMAFILE_API_BASE_URL;
102
103    fn build<H>(
104        _builder: &client::ClientBuilder<Self, Self::ApiKey, H>,
105    ) -> http_client::Result<Self::Extension<H>>
106    where
107        H: HttpClientExt,
108    {
109        Ok(LlamafileExt)
110    }
111}
112
113pub type Client<H = reqwest::Client> = client::Client<LlamafileExt, H>;
114pub type ClientBuilder<H = crate::markers::Missing> =
115    client::ClientBuilder<LlamafileBuilder, Nothing, H>;
116
117/// Llamafile completion model, driven by the shared OpenAI Chat Completions path.
118pub type CompletionModel<H = reqwest::Client> =
119    openai::completion::GenericCompletionModel<LlamafileExt, H>;
120
121/// Llamafile embedding model, driven by the shared OpenAI embeddings path.
122pub type EmbeddingModel<H = reqwest::Client> =
123    openai::embedding::GenericEmbeddingModel<LlamafileExt, H>;
124
125impl Client {
126    /// Create a client pointing at the given llamafile base URL
127    /// (e.g. `http://localhost:8080`).
128    pub fn from_url(base_url: &str) -> crate::client::ProviderClientResult<Self> {
129        Self::builder()
130            .api_key(Nothing)
131            .base_url(base_url)
132            .build()
133            .map_err(Into::into)
134    }
135}
136
137impl ProviderClient for Client {
138    type Input = Nothing;
139    type Error = crate::client::ProviderClientError;
140
141    fn from_env() -> Result<Self, Self::Error> {
142        let api_base = crate::client::required_env_var("LLAMAFILE_API_BASE_URL")?;
143        Self::from_url(&api_base)
144    }
145
146    fn from_val(_: Self::Input) -> Result<Self, Self::Error> {
147        Self::builder().api_key(Nothing).build().map_err(Into::into)
148    }
149}
150
151// ================================================================
152// Tests
153// ================================================================
154#[cfg(test)]
155mod tests {
156    use super::*;
157    use crate::client::{EmbeddingsClient, Nothing};
158    use crate::embeddings::EmbeddingModel as _;
159    use crate::providers::openai::embedding::EncodingFormat;
160    use crate::test_utils::RecordingHttpClient;
161
162    #[test]
163    fn test_client_initialization() {
164        let _client =
165            crate::providers::llamafile::Client::new(Nothing).expect("Client::new() failed");
166        let _client_from_builder = crate::providers::llamafile::Client::builder()
167            .api_key(Nothing)
168            .build()
169            .expect("Client::builder() failed");
170    }
171
172    #[test]
173    fn test_client_from_url() {
174        let _client = crate::providers::llamafile::Client::from_url("http://localhost:8080");
175    }
176
177    #[test]
178    fn test_build_uri_routes_through_v1() {
179        let ext = LlamafileExt;
180        assert_eq!(
181            ext.build_uri(
182                "http://localhost:8080",
183                "/chat/completions",
184                Transport::Http
185            ),
186            "http://localhost:8080/v1/chat/completions"
187        );
188        assert_eq!(
189            ext.build_uri("http://localhost:8080/", "/embeddings", Transport::Http),
190            "http://localhost:8080/v1/embeddings"
191        );
192        assert_eq!(
193            ext.build_uri(
194                "http://localhost:8080",
195                LlamafileExt::VERIFY_PATH,
196                Transport::Http
197            ),
198            "http://localhost:8080/v1/models"
199        );
200    }
201
202    #[tokio::test]
203    async fn embedding_model_preserves_v1_path_and_usage() {
204        let response = r#"{
205            "object": "list",
206            "model": "LLaMA_CPP",
207            "usage": { "prompt_tokens": 2, "total_tokens": 2 },
208            "data": [{ "object": "embedding", "index": 0, "embedding": [0.1, 0.2] }]
209        }"#;
210        let http_client = RecordingHttpClient::new(response);
211        let client = Client::builder()
212            .api_key(Nothing)
213            .http_client(http_client.clone())
214            .build()
215            .expect("client should build");
216        let model = client.embedding_model(LLAMA_CPP);
217
218        let response = model
219            .embed_texts_with_usage(["hello".to_string()])
220            .await
221            .expect("embedding request should succeed");
222
223        assert_eq!(response.usage.total_tokens, 2);
224        assert_eq!(
225            http_client.requests()[0].uri,
226            "http://localhost:8080/v1/embeddings"
227        );
228    }
229
230    #[tokio::test]
231    async fn embedding_model_rejects_base64_before_sending() {
232        let http_client = RecordingHttpClient::new("{}");
233        let client = Client::builder()
234            .api_key(Nothing)
235            .http_client(http_client.clone())
236            .build()
237            .expect("client should build");
238        let model = client
239            .embedding_model(LLAMA_CPP)
240            .encoding_format(EncodingFormat::Base64);
241
242        let error = model
243            .embed_texts(["hello".to_string()])
244            .await
245            .expect_err("numeric response parser should reject base64");
246
247        assert!(matches!(
248            error,
249            crate::embeddings::EmbeddingError::UnsupportedResponseEncoding {
250                provider: "llamafile",
251                encoding_format: "base64"
252            }
253        ));
254        assert!(http_client.requests().is_empty());
255    }
256}
rig_core/providers/llamafile.rs

rig_core/providers/
llamafile.rs
