summaryrefslogtreecommitdiffstats
path: root/src/client/openai_compatible.rs
diff options
context:
space:
mode:
authorsigoden <sigoden@gmail.com>2024-09-24 07:42:24 +0800
committerGitHub <noreply@github.com>2024-09-24 07:42:24 +0800
commit912773c25a113f49c5df63cd3a8086d38c75103e (patch)
tree7d653edf095f0949d1eeb2f83ec3e5616972c3b1 /src/client/openai_compatible.rs
parent00c4a6e421f01590ffc3fa5601e93d5ec755fca7 (diff)
downloadaichat-912773c25a113f49c5df63cd3a8086d38c75103e.tar.gz
refactor: embeddings/rerank fn accept ref data (#878)
Diffstat (limited to 'src/client/openai_compatible.rs')
-rw-r--r--src/client/openai_compatible.rs13
1 files changed, 8 insertions, 5 deletions
diff --git a/src/client/openai_compatible.rs b/src/client/openai_compatible.rs
index b2797e7..f2b7ae3 100644
--- a/src/client/openai_compatible.rs
+++ b/src/client/openai_compatible.rs
@@ -66,7 +66,10 @@ fn prepare_chat_completions(
Ok(request_data)
}
-fn prepare_embeddings(self_: &OpenAICompatibleClient, data: EmbeddingsData) -> Result<RequestData> {
+fn prepare_embeddings(
+ self_: &OpenAICompatibleClient,
+ data: &EmbeddingsData,
+) -> Result<RequestData> {
let api_key = self_.get_api_key().ok();
let api_base = get_api_base_ext(self_)?;
@@ -83,7 +86,7 @@ fn prepare_embeddings(self_: &OpenAICompatibleClient, data: EmbeddingsData) -> R
Ok(request_data)
}
-fn prepare_rerank(self_: &OpenAICompatibleClient, data: RerankData) -> Result<RequestData> {
+fn prepare_rerank(self_: &OpenAICompatibleClient, data: &RerankData) -> Result<RequestData> {
let api_key = self_.get_api_key().ok();
let api_base = get_api_base_ext(self_)?;
@@ -145,7 +148,7 @@ pub struct GenericRerankResBody {
pub results: RerankOutput,
}
-pub fn generic_build_rerank_body(data: RerankData, model: &Model) -> Value {
+pub fn generic_build_rerank_body(data: &RerankData, model: &Model) -> Value {
let RerankData {
query,
documents,
@@ -158,9 +161,9 @@ pub fn generic_build_rerank_body(data: RerankData, model: &Model) -> Value {
"documents": documents,
});
if model.client_name() == "voyageai" {
- body["top_k"] = top_n.into()
+ body["top_k"] = (*top_n).into()
} else {
- body["top_n"] = top_n.into()
+ body["top_n"] = (*top_n).into()
}
body
}