Skip to main content

rig_core/providers/
galadriel.rs

1//! Galadriel API client and Rig integration
2//!
3//! # Example
4//! ```no_run
5//! use rig_core::{client::CompletionClient, providers::galadriel};
6//!
7//! # fn run() -> Result<(), Box<dyn std::error::Error>> {
8//! let client = galadriel::Client::new("YOUR_API_KEY")?;
9//! // to use a fine-tuned model
10//! // let client = galadriel::Client::builder()
11//! //     .api_key("YOUR_API_KEY")
12//! //     .fine_tune_api_key("FINE_TUNE_API_KEY")
13//! //     .build()?;
14//!
15//! let gpt4o = client.completion_model(galadriel::GPT_4O);
16//! # Ok(())
17//! # }
18//! ```
19use super::openai;
20use crate::client::{
21    self, BearerAuth, Capabilities, Capable, DebugExt, Nothing, Provider, ProviderBuilder,
22    ProviderClient,
23};
24use crate::http_client::{self, HttpClientExt};
25use crate::message::MessageError;
26use crate::providers::openai::send_compatible_streaming_request;
27use crate::streaming::StreamingCompletionResponse;
28use crate::{
29    OneOrMany,
30    completion::{self, CompletionError, CompletionRequest},
31    json_utils, message,
32};
33use serde::{Deserialize, Serialize};
34use tracing::{Instrument, enabled, info_span};
35
36// ================================================================
37// Main Galadriel Client
38// ================================================================
39const GALADRIEL_API_BASE_URL: &str = "https://api.galadriel.com/v1/verified";
40
41#[derive(Debug, Default, Clone)]
42pub struct GaladrielExt {
43    fine_tune_api_key: Option<String>,
44}
45
46#[derive(Debug, Default, Clone)]
47pub struct GaladrielBuilder {
48    fine_tune_api_key: Option<String>,
49}
50
51type GaladrielApiKey = BearerAuth;
52
53impl Provider for GaladrielExt {
54    type Builder = GaladrielBuilder;
55
56    /// There is currently no way to verify a Galadriel api key without consuming tokens
57    const VERIFY_PATH: &'static str = "";
58}
59
60impl<H> Capabilities<H> for GaladrielExt {
61    type Completion = Capable<CompletionModel<H>>;
62    type Embeddings = Nothing;
63    type Transcription = Nothing;
64    type ModelListing = Nothing;
65    #[cfg(feature = "image")]
66    type ImageGeneration = Nothing;
67    #[cfg(feature = "audio")]
68    type AudioGeneration = Nothing;
69}
70
71impl DebugExt for GaladrielExt {
72    fn fields(&self) -> impl Iterator<Item = (&'static str, &dyn std::fmt::Debug)> {
73        std::iter::once((
74            "fine_tune_api_key",
75            (&self.fine_tune_api_key as &dyn std::fmt::Debug),
76        ))
77    }
78}
79
80impl ProviderBuilder for GaladrielBuilder {
81    type Extension<H>
82        = GaladrielExt
83    where
84        H: HttpClientExt;
85    type ApiKey = GaladrielApiKey;
86
87    const BASE_URL: &'static str = GALADRIEL_API_BASE_URL;
88
89    fn build<H>(
90        builder: &crate::client::ClientBuilder<Self, Self::ApiKey, H>,
91    ) -> http_client::Result<Self::Extension<H>>
92    where
93        H: HttpClientExt,
94    {
95        let GaladrielBuilder { fine_tune_api_key } = builder.ext().clone();
96
97        Ok(GaladrielExt { fine_tune_api_key })
98    }
99}
100
101pub type Client<H = reqwest::Client> = client::Client<GaladrielExt, H>;
102pub type ClientBuilder<H = crate::markers::Missing> =
103    client::ClientBuilder<GaladrielBuilder, GaladrielApiKey, H>;
104
105impl<T> ClientBuilder<T> {
106    pub fn fine_tune_api_key<S>(mut self, fine_tune_api_key: S) -> Self
107    where
108        S: AsRef<str>,
109    {
110        *self.ext_mut() = GaladrielBuilder {
111            fine_tune_api_key: Some(fine_tune_api_key.as_ref().into()),
112        };
113
114        self
115    }
116}
117
118impl ProviderClient for Client {
119    type Input = (String, Option<String>);
120    type Error = crate::client::ProviderClientError;
121
122    /// Create a new Galadriel client from the `GALADRIEL_API_KEY` environment variable,
123    /// and optionally from the `GALADRIEL_FINE_TUNE_API_KEY` environment variable.
124    fn from_env() -> Result<Self, Self::Error> {
125        let api_key = crate::client::required_env_var("GALADRIEL_API_KEY")?;
126        let fine_tune_api_key = crate::client::optional_env_var("GALADRIEL_FINE_TUNE_API_KEY")?;
127
128        let mut builder = Self::builder().api_key(api_key);
129
130        if let Some(fine_tune_api_key) = fine_tune_api_key {
131            builder = builder.fine_tune_api_key(fine_tune_api_key);
132        }
133
134        builder.build().map_err(Into::into)
135    }
136
137    fn from_val((api_key, fine_tune_api_key): Self::Input) -> Result<Self, Self::Error> {
138        let mut builder = Self::builder().api_key(api_key);
139
140        if let Some(fine_tune_key) = fine_tune_api_key {
141            builder = builder.fine_tune_api_key(fine_tune_key)
142        }
143
144        builder.build().map_err(Into::into)
145    }
146}
147
148#[derive(Debug, Deserialize)]
149struct ApiErrorResponse {
150    message: String,
151}
152
153#[derive(Debug, Deserialize)]
154#[serde(untagged)]
155enum ApiResponse<T> {
156    Ok(T),
157    Err(ApiErrorResponse),
158}
159
160#[derive(Clone, Debug, Deserialize, Serialize)]
161pub struct Usage {
162    pub prompt_tokens: usize,
163    pub total_tokens: usize,
164}
165
166impl std::fmt::Display for Usage {
167    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
168        write!(
169            f,
170            "Prompt tokens: {} Total tokens: {}",
171            self.prompt_tokens, self.total_tokens
172        )
173    }
174}
175
176// ================================================================
177// Galadriel Completion API
178// ================================================================
179
180/// `o1-preview` completion model
181pub const O1_PREVIEW: &str = "o1-preview";
182/// `o1-preview-2024-09-12` completion model
183pub const O1_PREVIEW_2024_09_12: &str = "o1-preview-2024-09-12";
184/// `o1-mini completion model
185pub const O1_MINI: &str = "o1-mini";
186/// `o1-mini-2024-09-12` completion model
187pub const O1_MINI_2024_09_12: &str = "o1-mini-2024-09-12";
188/// `gpt-4o` completion model
189pub const GPT_4O: &str = "gpt-4o";
190/// `gpt-4o-2024-05-13` completion model
191pub const GPT_4O_2024_05_13: &str = "gpt-4o-2024-05-13";
192/// `gpt-4-turbo` completion model
193pub const GPT_4_TURBO: &str = "gpt-4-turbo";
194/// `gpt-4-turbo-2024-04-09` completion model
195pub const GPT_4_TURBO_2024_04_09: &str = "gpt-4-turbo-2024-04-09";
196/// `gpt-4-turbo-preview` completion model
197pub const GPT_4_TURBO_PREVIEW: &str = "gpt-4-turbo-preview";
198/// `gpt-4-0125-preview` completion model
199pub const GPT_4_0125_PREVIEW: &str = "gpt-4-0125-preview";
200/// `gpt-4-1106-preview` completion model
201pub const GPT_4_1106_PREVIEW: &str = "gpt-4-1106-preview";
202/// `gpt-4-vision-preview` completion model
203pub const GPT_4_VISION_PREVIEW: &str = "gpt-4-vision-preview";
204/// `gpt-4-1106-vision-preview` completion model
205pub const GPT_4_1106_VISION_PREVIEW: &str = "gpt-4-1106-vision-preview";
206/// `gpt-4` completion model
207pub const GPT_4: &str = "gpt-4";
208/// `gpt-4-0613` completion model
209pub const GPT_4_0613: &str = "gpt-4-0613";
210/// `gpt-4-32k` completion model
211pub const GPT_4_32K: &str = "gpt-4-32k";
212/// `gpt-4-32k-0613` completion model
213pub const GPT_4_32K_0613: &str = "gpt-4-32k-0613";
214/// `gpt-3.5-turbo` completion model
215pub const GPT_35_TURBO: &str = "gpt-3.5-turbo";
216/// `gpt-3.5-turbo-0125` completion model
217pub const GPT_35_TURBO_0125: &str = "gpt-3.5-turbo-0125";
218/// `gpt-3.5-turbo-1106` completion model
219pub const GPT_35_TURBO_1106: &str = "gpt-3.5-turbo-1106";
220/// `gpt-3.5-turbo-instruct` completion model
221pub const GPT_35_TURBO_INSTRUCT: &str = "gpt-3.5-turbo-instruct";
222
223#[derive(Debug, Deserialize, Serialize)]
224pub struct CompletionResponse {
225    pub id: String,
226    pub object: String,
227    pub created: u64,
228    pub model: String,
229    pub system_fingerprint: Option<String>,
230    pub choices: Vec<Choice>,
231    pub usage: Option<Usage>,
232}
233
234impl From<ApiErrorResponse> for CompletionError {
235    fn from(err: ApiErrorResponse) -> Self {
236        CompletionError::ProviderError(err.message)
237    }
238}
239
240impl TryFrom<CompletionResponse> for completion::CompletionResponse<CompletionResponse> {
241    type Error = CompletionError;
242
243    fn try_from(response: CompletionResponse) -> Result<Self, Self::Error> {
244        let Choice { message, .. } = response.choices.first().ok_or_else(|| {
245            CompletionError::ResponseError("Response contained no choices".to_owned())
246        })?;
247
248        let mut content = message
249            .content
250            .as_ref()
251            .map(|c| vec![completion::AssistantContent::text(c)])
252            .unwrap_or_default();
253
254        content.extend(message.tool_calls.iter().map(|call| {
255            completion::AssistantContent::tool_call(
256                &call.function.name,
257                &call.function.name,
258                call.function.arguments.clone(),
259            )
260        }));
261
262        let choice = OneOrMany::many(content).map_err(|_| {
263            CompletionError::ResponseError(
264                "Response contained no message or tool call (empty)".to_owned(),
265            )
266        })?;
267        let usage = response
268            .usage
269            .as_ref()
270            .map(|usage| completion::Usage {
271                input_tokens: usage.prompt_tokens as u64,
272                output_tokens: (usage.total_tokens - usage.prompt_tokens) as u64,
273                total_tokens: usage.total_tokens as u64,
274                cached_input_tokens: 0,
275                cache_creation_input_tokens: 0,
276                reasoning_tokens: 0,
277            })
278            .unwrap_or_default();
279
280        Ok(completion::CompletionResponse {
281            choice,
282            usage,
283            raw_response: response,
284            message_id: None,
285        })
286    }
287}
288
289#[derive(Debug, Deserialize, Serialize)]
290pub struct Choice {
291    pub index: usize,
292    pub message: Message,
293    pub logprobs: Option<serde_json::Value>,
294    pub finish_reason: String,
295}
296
297#[derive(Debug, Serialize, Deserialize)]
298pub struct Message {
299    pub role: String,
300    pub content: Option<String>,
301    #[serde(default, deserialize_with = "json_utils::null_or_vec")]
302    pub tool_calls: Vec<openai::ToolCall>,
303}
304
305impl Message {
306    fn system(preamble: &str) -> Self {
307        Self {
308            role: "system".to_string(),
309            content: Some(preamble.to_string()),
310            tool_calls: Vec::new(),
311        }
312    }
313}
314
315impl TryFrom<Message> for message::Message {
316    type Error = message::MessageError;
317
318    fn try_from(message: Message) -> Result<Self, Self::Error> {
319        let tool_calls: Vec<message::ToolCall> = message
320            .tool_calls
321            .into_iter()
322            .map(|tool_call| tool_call.into())
323            .collect();
324
325        match message.role.as_str() {
326            "user" => Ok(Self::User {
327                content: OneOrMany::one(
328                    message
329                        .content
330                        .map(|content| message::UserContent::text(&content))
331                        .ok_or_else(|| {
332                            message::MessageError::ConversionError("Empty user message".to_string())
333                        })?,
334                ),
335            }),
336            "assistant" => Ok(Self::Assistant {
337                id: None,
338                content: OneOrMany::many(
339                    tool_calls
340                        .into_iter()
341                        .map(message::AssistantContent::ToolCall)
342                        .chain(
343                            message
344                                .content
345                                .map(|content| message::AssistantContent::text(&content))
346                                .into_iter(),
347                        ),
348                )
349                .map_err(|_| {
350                    message::MessageError::ConversionError("Empty assistant message".to_string())
351                })?,
352            }),
353            _ => Err(message::MessageError::ConversionError(format!(
354                "Unknown role: {}",
355                message.role
356            ))),
357        }
358    }
359}
360
361impl TryFrom<message::Message> for Message {
362    type Error = message::MessageError;
363
364    fn try_from(message: message::Message) -> Result<Self, Self::Error> {
365        match message {
366            message::Message::System { content } => Ok(Self {
367                role: "system".to_string(),
368                content: Some(content),
369                tool_calls: vec![],
370            }),
371            message::Message::User { content } => Ok(Self {
372                role: "user".to_string(),
373                content: content.iter().find_map(|c| match c {
374                    message::UserContent::Text(text) => Some(text.text.clone()),
375                    _ => None,
376                }),
377                tool_calls: vec![],
378            }),
379            message::Message::Assistant { content, .. } => {
380                let mut text_content: Option<String> = None;
381                let mut tool_calls = vec![];
382
383                for c in content.iter() {
384                    match c {
385                        message::AssistantContent::Text(text) => {
386                            text_content = Some(
387                                text_content
388                                    .map(|mut existing| {
389                                        existing.push('\n');
390                                        existing.push_str(&text.text);
391                                        existing
392                                    })
393                                    .unwrap_or_else(|| text.text.clone()),
394                            );
395                        }
396                        message::AssistantContent::ToolCall(tool_call) => {
397                            tool_calls.push(tool_call.clone().into());
398                        }
399                        message::AssistantContent::Reasoning(_) => {
400                            return Err(MessageError::ConversionError(
401                                "Galadriel currently doesn't support reasoning.".into(),
402                            ));
403                        }
404                        message::AssistantContent::Image(_) => {
405                            return Err(MessageError::ConversionError(
406                                "Galadriel currently doesn't support images.".into(),
407                            ));
408                        }
409                    }
410                }
411
412                Ok(Self {
413                    role: "assistant".to_string(),
414                    content: text_content,
415                    tool_calls,
416                })
417            }
418        }
419    }
420}
421
422#[derive(Clone, Debug, Deserialize, Serialize)]
423pub struct ToolDefinition {
424    pub r#type: String,
425    pub function: completion::ToolDefinition,
426}
427
428impl From<completion::ToolDefinition> for ToolDefinition {
429    fn from(tool: completion::ToolDefinition) -> Self {
430        Self {
431            r#type: "function".into(),
432            function: tool,
433        }
434    }
435}
436
437#[derive(Debug, Deserialize)]
438pub struct Function {
439    pub name: String,
440    pub arguments: String,
441}
442
443#[derive(Debug, Serialize, Deserialize)]
444pub(super) struct GaladrielCompletionRequest {
445    model: String,
446    pub messages: Vec<Message>,
447    #[serde(skip_serializing_if = "Option::is_none")]
448    temperature: Option<f64>,
449    #[serde(skip_serializing_if = "Vec::is_empty")]
450    tools: Vec<ToolDefinition>,
451    #[serde(skip_serializing_if = "Option::is_none")]
452    tool_choice: Option<crate::providers::openai::completion::ToolChoice>,
453    #[serde(flatten, skip_serializing_if = "Option::is_none")]
454    pub additional_params: Option<serde_json::Value>,
455}
456
457impl TryFrom<(&str, CompletionRequest)> for GaladrielCompletionRequest {
458    type Error = CompletionError;
459
460    fn try_from((model, req): (&str, CompletionRequest)) -> Result<Self, Self::Error> {
461        if req.output_schema.is_some() {
462            tracing::warn!("Structured outputs currently not supported for Galadriel");
463        }
464        let model = req.model.clone().unwrap_or_else(|| model.to_string());
465        // Build up the order of messages (context, chat_history, prompt)
466        let mut partial_history = vec![];
467        if let Some(docs) = req.normalized_documents() {
468            partial_history.push(docs);
469        }
470        partial_history.extend(req.chat_history);
471
472        // Add preamble to chat history (if available)
473        let mut full_history: Vec<Message> = match &req.preamble {
474            Some(preamble) => vec![Message::system(preamble)],
475            None => vec![],
476        };
477
478        // Convert and extend the rest of the history
479        full_history.extend(
480            partial_history
481                .into_iter()
482                .map(message::Message::try_into)
483                .collect::<Result<Vec<Message>, _>>()?,
484        );
485
486        let tool_choice = req
487            .tool_choice
488            .clone()
489            .map(crate::providers::openai::completion::ToolChoice::try_from)
490            .transpose()?;
491
492        Ok(Self {
493            model: model.to_string(),
494            messages: full_history,
495            temperature: req.temperature,
496            tools: req
497                .tools
498                .clone()
499                .into_iter()
500                .map(ToolDefinition::from)
501                .collect::<Vec<_>>(),
502            tool_choice,
503            additional_params: req.additional_params,
504        })
505    }
506}
507
508#[derive(Clone)]
509pub struct CompletionModel<T = reqwest::Client> {
510    client: Client<T>,
511    /// Name of the model (e.g.: gpt-3.5-turbo-1106)
512    pub model: String,
513}
514
515impl<T> CompletionModel<T>
516where
517    T: HttpClientExt,
518{
519    pub fn new(client: Client<T>, model: impl Into<String>) -> Self {
520        Self {
521            client,
522            model: model.into(),
523        }
524    }
525
526    pub fn with_model(client: Client<T>, model: &str) -> Self {
527        Self {
528            client,
529            model: model.into(),
530        }
531    }
532}
533
534impl<T> completion::CompletionModel for CompletionModel<T>
535where
536    T: HttpClientExt + Clone + Default + std::fmt::Debug + Send + 'static,
537{
538    type Response = CompletionResponse;
539    type StreamingResponse = openai::StreamingCompletionResponse;
540
541    type Client = Client<T>;
542
543    fn make(client: &Self::Client, model: impl Into<String>) -> Self {
544        Self::new(client.clone(), model.into())
545    }
546
547    async fn completion(
548        &self,
549        completion_request: CompletionRequest,
550    ) -> Result<completion::CompletionResponse<CompletionResponse>, CompletionError> {
551        let span = if tracing::Span::current().is_disabled() {
552            info_span!(
553                target: "rig::completions",
554                "chat",
555                gen_ai.operation.name = "chat",
556                gen_ai.provider.name = "galadriel",
557                gen_ai.request.model = self.model,
558                gen_ai.system_instructions = tracing::field::Empty,
559                gen_ai.response.id = tracing::field::Empty,
560                gen_ai.response.model = tracing::field::Empty,
561                gen_ai.usage.output_tokens = tracing::field::Empty,
562                gen_ai.usage.input_tokens = tracing::field::Empty,
563                gen_ai.usage.cache_read.input_tokens = tracing::field::Empty,
564            )
565        } else {
566            tracing::Span::current()
567        };
568
569        span.record("gen_ai.system_instructions", &completion_request.preamble);
570
571        let request =
572            GaladrielCompletionRequest::try_from((self.model.as_ref(), completion_request))?;
573
574        if enabled!(tracing::Level::TRACE) {
575            tracing::trace!(target: "rig::completions",
576                "Galadriel completion request: {}",
577                serde_json::to_string_pretty(&request)?
578            );
579        }
580
581        let body = serde_json::to_vec(&request)?;
582
583        let req = self
584            .client
585            .post("/chat/completions")?
586            .body(body)
587            .map_err(http_client::Error::from)?;
588
589        async move {
590            let response = self.client.send(req).await?;
591
592            if response.status().is_success() {
593                let t = http_client::text(response).await?;
594
595                if enabled!(tracing::Level::TRACE) {
596                    tracing::trace!(target: "rig::completions",
597                        "Galadriel completion response: {}",
598                        serde_json::to_string_pretty(&t)?
599                    );
600                }
601
602                match serde_json::from_str::<ApiResponse<CompletionResponse>>(&t)? {
603                    ApiResponse::Ok(response) => {
604                        let span = tracing::Span::current();
605                        span.record("gen_ai.response.id", response.id.clone());
606                        span.record("gen_ai.response.model", response.model.clone());
607                        if let Some(ref usage) = response.usage {
608                            span.record("gen_ai.usage.input_tokens", usage.prompt_tokens);
609                            span.record(
610                                "gen_ai.usage.output_tokens",
611                                usage.total_tokens - usage.prompt_tokens,
612                            );
613                        }
614                        response.try_into()
615                    }
616                    ApiResponse::Err(err) => Err(CompletionError::ProviderError(err.message)),
617                }
618            } else {
619                let text = http_client::text(response).await?;
620
621                Err(CompletionError::ProviderError(text))
622            }
623        }
624        .instrument(span)
625        .await
626    }
627
628    async fn stream(
629        &self,
630        completion_request: CompletionRequest,
631    ) -> Result<StreamingCompletionResponse<Self::StreamingResponse>, CompletionError> {
632        let preamble = completion_request.preamble.clone();
633        let mut request =
634            GaladrielCompletionRequest::try_from((self.model.as_ref(), completion_request))?;
635
636        let params = json_utils::merge(
637            request.additional_params.unwrap_or(serde_json::json!({})),
638            serde_json::json!({"stream": true, "stream_options": {"include_usage": true} }),
639        );
640
641        request.additional_params = Some(params);
642
643        let body = serde_json::to_vec(&request)?;
644
645        let req = self
646            .client
647            .post("/chat/completions")?
648            .body(body)
649            .map_err(http_client::Error::from)?;
650
651        let span = if tracing::Span::current().is_disabled() {
652            info_span!(
653                target: "rig::completions",
654                "chat_streaming",
655                gen_ai.operation.name = "chat_streaming",
656                gen_ai.provider.name = "galadriel",
657                gen_ai.request.model = self.model,
658                gen_ai.system_instructions = preamble,
659                gen_ai.response.id = tracing::field::Empty,
660                gen_ai.response.model = tracing::field::Empty,
661                gen_ai.usage.output_tokens = tracing::field::Empty,
662                gen_ai.usage.input_tokens = tracing::field::Empty,
663                gen_ai.usage.cache_read.input_tokens = tracing::field::Empty,
664                gen_ai.input.messages = serde_json::to_string(&request.messages)?,
665                gen_ai.output.messages = tracing::field::Empty,
666            )
667        } else {
668            tracing::Span::current()
669        };
670
671        send_compatible_streaming_request(self.client.clone(), req)
672            .instrument(span)
673            .await
674    }
675}
676#[cfg(test)]
677mod tests {
678    #[test]
679    fn test_client_initialization() {
680        let _client =
681            crate::providers::galadriel::Client::new("dummy-key").expect("Client::new() failed");
682        let _client_from_builder = crate::providers::galadriel::Client::builder()
683            .api_key("dummy-key")
684            .build()
685            .expect("Client::builder() failed");
686    }
687}