diff options
| author | sigoden <sigoden@gmail.com> | 2024-07-27 21:33:04 +0800 |
|---|---|---|
| committer | GitHub <noreply@github.com> | 2024-07-27 21:33:04 +0800 |
| commit | f5e8e872b08d2aa9b3c063b0a98c59f842f94ecc (patch) | |
| tree | 6f6a6ce299d7383d290a7015b52ce777566d1e8e /src/client/ollama.rs | |
| parent | adf6716c8436b53854eb55588d258f4326580a4d (diff) | |
| download | aichat-f5e8e872b08d2aa9b3c063b0a98c59f842f94ecc.tar.gz | |
feat: support patching request url, headers and body (#756)
Diffstat (limited to 'src/client/ollama.rs')
| -rw-r--r-- | src/client/ollama.rs | 39 |
1 files changed, 14 insertions, 25 deletions
diff --git a/src/client/ollama.rs b/src/client/ollama.rs index 299cae7..5c26b99 100644 --- a/src/client/ollama.rs +++ b/src/client/ollama.rs @@ -1,7 +1,7 @@ use super::*; use anyhow::{bail, Context, Result}; -use reqwest::{Client as ReqwestClient, RequestBuilder}; +use reqwest::RequestBuilder; use serde::Deserialize; use serde_json::{json, Value}; @@ -12,7 +12,7 @@ pub struct OllamaConfig { pub api_auth: Option<String>, #[serde(default)] pub models: Vec<ModelData>, - pub patch: Option<ModelPatch>, + pub patch: Option<RequestPatch>, pub extra: Option<ExtraConfig>, } @@ -32,52 +32,41 @@ impl OllamaClient { ), ]; - fn chat_completions_builder( - &self, - client: &ReqwestClient, - data: ChatCompletionsData, - ) -> Result<RequestBuilder> { + fn prepare_chat_completions(&self, data: ChatCompletionsData) -> Result<RequestData> { let api_base = self.get_api_base()?; let api_auth = self.get_api_auth().ok(); - let mut body = build_chat_completions_body(data, &self.model)?; - self.patch_chat_completions_body(&mut body); - let url = format!("{api_base}/api/chat"); - debug!("Ollama Chat Completions Request: {url} {body}"); + let body = build_chat_completions_body(data, &self.model)?; + + let mut request_data = RequestData::new(url, body); - let mut builder = client.post(url).json(&body); if let Some(api_auth) = api_auth { - builder = builder.header("Authorization", api_auth) + request_data.header("Authorization", api_auth) } - Ok(builder) + Ok(request_data) } - fn embeddings_builder( - &self, - client: &ReqwestClient, - data: EmbeddingsData, - ) -> Result<RequestBuilder> { + fn prepare_embeddings(&self, data: EmbeddingsData) -> Result<RequestData> { let api_base = self.get_api_base()?; let api_auth = self.get_api_auth().ok(); + let url = format!("{api_base}/api/embed"); + let body = json!({ "model": self.model.name(), "input": data.texts, }); - let url = format!("{api_base}/api/embed"); - - debug!("Ollama Embeddings Request: {url} {body}"); + let mut request_data = RequestData::new(url, body); - let mut builder = client.post(url).json(&body); if let Some(api_auth) = api_auth { - builder = builder.header("Authorization", api_auth) + request_data.header("Authorization", api_auth) } - Ok(builder) + Ok(request_data) } } |
