diff options
Diffstat (limited to 'src/client/gemini.rs')
| -rw-r--r-- | src/client/gemini.rs | 83 |
1 files changed, 45 insertions, 38 deletions
diff --git a/src/client/gemini.rs b/src/client/gemini.rs index aa1a5b1..2616218 100644 --- a/src/client/gemini.rs +++ b/src/client/gemini.rs @@ -23,57 +23,64 @@ impl GeminiClient { pub const PROMPTS: [PromptAction<'static>; 1] = [("api_key", "API Key:", true, PromptKind::String)]; +} + +impl_client_trait!( + GeminiClient, + ( + prepare_chat_completions, + gemini_chat_completions, + gemini_chat_completions_streaming + ), + (prepare_embeddings, gemini_embeddings), + (noop_prepare_rerank, noop_rerank), +); - fn prepare_chat_completions(&self, data: ChatCompletionsData) -> Result<RequestData> { - let api_key = self.get_api_key()?; +fn prepare_chat_completions( + self_: &GeminiClient, + data: ChatCompletionsData, +) -> Result<RequestData> { + let api_key = self_.get_api_key()?; - let func = match data.stream { - true => "streamGenerateContent", - false => "generateContent", - }; + let func = match data.stream { + true => "streamGenerateContent", + false => "generateContent", + }; - let url = format!("{API_BASE}{}:{}?key={}", &self.model.name(), func, api_key); + let url = format!("{API_BASE}{}:{}?key={}", self_.model.name(), func, api_key); - let body = gemini_build_chat_completions_body(data, &self.model)?; + let body = gemini_build_chat_completions_body(data, &self_.model)?; - let request_data = RequestData::new(url, body); + let request_data = RequestData::new(url, body); - Ok(request_data) - } + Ok(request_data) +} - fn prepare_embeddings(&self, data: EmbeddingsData) -> Result<RequestData> { - let api_key = self.get_api_key()?; +fn prepare_embeddings(self_: &GeminiClient, data: EmbeddingsData) -> Result<RequestData> { + let api_key = self_.get_api_key()?; - let url = format!( - "{API_BASE}{}:embedContent?key={}", - &self.model.name(), - api_key - ); + let url = format!( + "{API_BASE}{}:embedContent?key={}", + self_.model.name(), + api_key + ); - let body = json!({ - "content": { - "parts": [ - { - "text": data.texts[0], - } - ] - } - }); + let body = json!({ + "content": { + "parts": [ + { + "text": data.texts[0], + } + ] + } + }); - let request_data = RequestData::new(url, body); + let request_data = RequestData::new(url, body); - Ok(request_data) - } + Ok(request_data) } -impl_client_trait!( - GeminiClient, - gemini_chat_completions, - gemini_chat_completions_streaming, - gemini_embeddings -); - -async fn gemini_embeddings(builder: RequestBuilder) -> Result<EmbeddingsOutput> { +async fn gemini_embeddings(builder: RequestBuilder, _model: &Model) -> Result<EmbeddingsOutput> { let res = builder.send().await?; let status = res.status(); let data: Value = res.json().await?; |
