alith_models/api_model/
perplexity.rs

1use super::ApiLLMModel;
2use crate::{LLMModelBase, tokenizer::Tokenizer};
3use std::sync::Arc;
4
5impl ApiLLMModel {
6    pub fn perplexity_model_from_model_id(model_id: &str) -> ApiLLMModel {
7        if model_id.starts_with("llama-3.1-sonar-small") {
8            Self::sonar_small()
9        } else if model_id.starts_with("llama-3.1-sonar-large") {
10            Self::sonar_large()
11        } else if model_id.starts_with("llama-3.1-sonar-huge") {
12            Self::sonar_huge()
13        } else if model_id.contains("sonar-small") {
14            Self::sonar_small()
15        } else if model_id.contains("sonar-large") {
16            Self::sonar_large()
17        } else if model_id.contains("sonar-huge") {
18            Self::sonar_huge()
19        } else {
20            panic!("Model ID ({model_id}) not found for ApiLLMModel")
21        }
22    }
23
24    pub fn sonar_small() -> ApiLLMModel {
25        let model_id = "llama-3.1-sonar-small-128k-online".to_string();
26        let tokenizer = model_tokenizer(&model_id);
27        ApiLLMModel {
28            model_base: LLMModelBase {
29                model_id,
30                model_ctx_size: 127072,
31                inference_ctx_size: 8192,
32                tokenizer,
33            },
34            cost_per_m_in_tokens: 0.1,
35            cost_per_m_out_tokens: 0.1,
36            tokens_per_message: 3,
37            tokens_per_name: None,
38        }
39    }
40
41    pub fn sonar_large() -> ApiLLMModel {
42        let model_id = "llama-3.1-sonar-large-128k-online".to_string();
43        let tokenizer = model_tokenizer(&model_id);
44        ApiLLMModel {
45            model_base: LLMModelBase {
46                model_id,
47                model_ctx_size: 127072,
48                inference_ctx_size: 8192,
49                tokenizer,
50            },
51            cost_per_m_in_tokens: 0.5,
52            cost_per_m_out_tokens: 0.5,
53            tokens_per_message: 3,
54            tokens_per_name: None,
55        }
56    }
57
58    pub fn sonar_huge() -> ApiLLMModel {
59        let model_id = "llama-3.1-sonar-huge-128k-online".to_string();
60        let tokenizer = model_tokenizer(&model_id);
61        ApiLLMModel {
62            model_base: LLMModelBase {
63                model_id,
64                model_ctx_size: 127072,
65                inference_ctx_size: 8192,
66                tokenizer,
67            },
68            cost_per_m_in_tokens: 2.5,
69            cost_per_m_out_tokens: 2.5,
70            tokens_per_message: 3,
71            tokens_per_name: None,
72        }
73    }
74}
75
76pub fn model_tokenizer(_model_id: &str) -> Arc<Tokenizer> {
77    Arc::new(
78        Tokenizer::new_tiktoken("gpt-4")
79            .unwrap_or_else(|_| panic!("Failed to load tokenizer for gpt-4")),
80    )
81}
82
83pub trait PerplexityModelTrait: Sized {
84    fn model(&mut self) -> &mut ApiLLMModel;
85
86    /// Set the model using the model_id string.
87    fn model_id_str(mut self, model_id: &str) -> Self
88    where
89        Self: Sized,
90    {
91        *self.model() = ApiLLMModel::perplexity_model_from_model_id(model_id);
92        self
93    }
94
95    fn sonar_small(mut self) -> Self
96    where
97        Self: Sized,
98    {
99        *self.model() = ApiLLMModel::sonar_small();
100        self
101    }
102
103    fn sonar_large(mut self) -> Self
104    where
105        Self: Sized,
106    {
107        *self.model() = ApiLLMModel::sonar_large();
108        self
109    }
110
111    fn sonar_huge(mut self) -> Self
112    where
113        Self: Sized,
114    {
115        *self.model() = ApiLLMModel::sonar_huge();
116        self
117    }
118}
alith_models/api_model/perplexity.rs

alith_models/api_model/
perplexity.rs
