summaryrefslogtreecommitdiffstats
path: root/src/client/ollama.rs
diff options
context:
space:
mode:
authorsigoden <sigoden@gmail.com>2024-06-01 17:47:49 +0800
committerGitHub <noreply@github.com>2024-06-01 17:47:49 +0800
commit571d1022f628cb7d2a3125664bd3293bac4471b5 (patch)
tree67034aaf0a656ae9dce0cd0d3aae3605d2ea1767 /src/client/ollama.rs
parent259583f4f750e4ece7aed07858a9589110d7cf5d (diff)
downloadaichat-571d1022f628cb7d2a3125664bd3293bac4471b5.tar.gz
refactor: rename some client structs and methods (#555)
* rename `Completeion*` to `ChatCompletions*` * rename `send_message*` to `chat_completions*` * rename `request_builder` to `chat_completions_builder` * rename `build_body` to `build_chat_completions_body` * rename `extract_completion` to `extract_chat_completions` * format * remove unused config fields
Diffstat (limited to 'src/client/ollama.rs')
-rw-r--r--src/client/ollama.rs26
1 files changed, 15 insertions, 11 deletions
diff --git a/src/client/ollama.rs b/src/client/ollama.rs
index 014824b..beba8a1 100644
--- a/src/client/ollama.rs
+++ b/src/client/ollama.rs
@@ -1,6 +1,7 @@
use super::{
- catch_error, json_stream, message::*, Client, CompletionData, CompletionOutput, ExtraConfig,
- Model, ModelData, ModelPatches, OllamaClient, PromptAction, PromptKind, SseHandler,
+ catch_error, json_stream, message::*, ChatCompletionsData, ChatCompletionsOutput, Client,
+ ExtraConfig, Model, ModelData, ModelPatches, OllamaClient, PromptAction, PromptKind,
+ SseHandler,
};
use anyhow::{anyhow, bail, Result};
@@ -35,15 +36,15 @@ impl OllamaClient {
),
];
- fn request_builder(
+ fn chat_completions_builder(
&self,
client: &ReqwestClient,
- data: CompletionData,
+ data: ChatCompletionsData,
) -> Result<RequestBuilder> {
let api_base = self.get_api_base()?;
let api_auth = self.get_api_auth().ok();
- let mut body = build_body(data, &self.model)?;
+ let mut body = build_chat_completions_body(data, &self.model)?;
self.patch_request_body(&mut body);
let chat_endpoint = self.config.chat_endpoint.as_deref().unwrap_or("/api/chat");
@@ -61,9 +62,9 @@ impl OllamaClient {
}
}
-impl_client_trait!(OllamaClient, send_message, send_message_streaming);
+impl_client_trait!(OllamaClient, chat_completions, chat_completions_streaming);
-async fn send_message(builder: RequestBuilder) -> Result<CompletionOutput> {
+async fn chat_completions(builder: RequestBuilder) -> Result<ChatCompletionsOutput> {
let res = builder.send().await?;
let status = res.status();
let data = res.json().await?;
@@ -74,10 +75,13 @@ async fn send_message(builder: RequestBuilder) -> Result<CompletionOutput> {
let text = data["message"]["content"]
.as_str()
.ok_or_else(|| anyhow!("Invalid response data: {data}"))?;
- Ok(CompletionOutput::new(text))
+ Ok(ChatCompletionsOutput::new(text))
}
-async fn send_message_streaming(builder: RequestBuilder, handler: &mut SseHandler) -> Result<()> {
+async fn chat_completions_streaming(
+ builder: RequestBuilder,
+ handler: &mut SseHandler,
+) -> Result<()> {
let res = builder.send().await?;
let status = res.status();
if !status.is_success() {
@@ -105,8 +109,8 @@ async fn send_message_streaming(builder: RequestBuilder, handler: &mut SseHandle
Ok(())
}
-fn build_body(data: CompletionData, model: &Model) -> Result<Value> {
- let CompletionData {
+fn build_chat_completions_body(data: ChatCompletionsData, model: &Model) -> Result<Value> {
+ let ChatCompletionsData {
messages,
temperature,
top_p,