summaryrefslogtreecommitdiffstats
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
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
-rw-r--r--src/client/azure_openai.rs12
-rw-r--r--src/client/bedrock.rs74
-rw-r--r--src/client/claude.rs34
-rw-r--r--src/client/cloudflare.rs31
-rw-r--r--src/client/cohere.rs31
-rw-r--r--src/client/common.rs54
-rw-r--r--src/client/ernie.rs47
-rw-r--r--src/client/gemini.rs16
-rw-r--r--src/client/ollama.rs26
-rw-r--r--src/client/openai.rs30
-rw-r--r--src/client/openai_compatible.rs12
-rw-r--r--src/client/qianwen.rs40
-rw-r--r--src/client/replicate.rs42
-rw-r--r--src/client/vertexai.rs49
-rw-r--r--src/client/vertexai_claude.rs29
-rw-r--r--src/config/input.rs10
-rw-r--r--src/main.rs10
-rw-r--r--src/repl/mod.rs2
-rw-r--r--src/serve.rs14
19 files changed, 298 insertions, 265 deletions
diff --git a/src/client/azure_openai.rs b/src/client/azure_openai.rs
index 75f8455..52c8a34 100644
--- a/src/client/azure_openai.rs
+++ b/src/client/azure_openai.rs
@@ -1,5 +1,5 @@
use super::{
- openai::*, AzureOpenAIClient, Client, CompletionData, ExtraConfig, Model, ModelData,
+ openai::*, AzureOpenAIClient, ChatCompletionsData, Client, ExtraConfig, Model, ModelData,
ModelPatches, PromptAction, PromptKind,
};
@@ -33,15 +33,15 @@ impl AzureOpenAIClient {
),
];
- 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_key = self.get_api_key()?;
- let mut body = openai_build_body(data, &self.model);
+ let mut body = openai_build_chat_completions_body(data, &self.model);
self.patch_request_body(&mut body);
let url = format!(
@@ -60,6 +60,6 @@ impl AzureOpenAIClient {
impl_client_trait!(
AzureOpenAIClient,
- crate::client::openai::openai_send_message,
- crate::client::openai::openai_send_message_streaming
+ crate::client::openai::openai_chat_completions,
+ crate::client::openai::openai_chat_completions_streaming
);
diff --git a/src/client/bedrock.rs b/src/client/bedrock.rs
index 355565f..3dfa977 100644
--- a/src/client/bedrock.rs
+++ b/src/client/bedrock.rs
@@ -1,7 +1,7 @@
use super::{
- catch_error, claude::*, prompt_format::*, BedrockClient, Client, CompletionData,
- CompletionOutput, ExtraConfig, Model, ModelData, ModelPatches, PromptAction, PromptKind,
- SseHandler,
+ catch_error, claude::*, prompt_format::*, BedrockClient, ChatCompletionsData,
+ ChatCompletionsOutput, Client, ExtraConfig, Model, ModelData, ModelPatches, PromptAction,
+ PromptKind, SseHandler,
};
use crate::utils::{base64_decode, encode_uri, hex_encode, hmac_sha256, sha256};
@@ -38,25 +38,25 @@ pub struct BedrockConfig {
impl Client for BedrockClient {
client_common_fns!();
- async fn send_message_inner(
+ async fn chat_completions_inner(
&self,
client: &ReqwestClient,
- data: CompletionData,
- ) -> Result<CompletionOutput> {
+ data: ChatCompletionsData,
+ ) -> Result<ChatCompletionsOutput> {
let model_category = ModelCategory::from_str(self.model.name())?;
- let builder = self.request_builder(client, data, &model_category)?;
- send_message(builder, &model_category).await
+ let builder = self.chat_completions_builder(client, data, &model_category)?;
+ chat_completions(builder, &model_category).await
}
- async fn send_message_streaming_inner(
+ async fn chat_completions_streaming_inner(
&self,
client: &ReqwestClient,
handler: &mut SseHandler,
- data: CompletionData,
+ data: ChatCompletionsData,
) -> Result<()> {
let model_category = ModelCategory::from_str(self.model.name())?;
- let builder = self.request_builder(client, data, &model_category)?;
- send_message_streaming(builder, handler, &model_category).await
+ let builder = self.chat_completions_builder(client, data, &model_category)?;
+ chat_completions_streaming(builder, handler, &model_category).await
}
}
@@ -81,10 +81,10 @@ impl BedrockClient {
("region", "AWS Region", true, PromptKind::String),
];
- fn request_builder(
+ fn chat_completions_builder(
&self,
client: &ReqwestClient,
- data: CompletionData,
+ data: ChatCompletionsData,
model_category: &ModelCategory,
) -> Result<RequestBuilder> {
let access_key_id = self.get_access_key_id()?;
@@ -101,7 +101,7 @@ impl BedrockClient {
let headers = IndexMap::new();
- let mut body = build_body(data, &self.model, model_category)?;
+ let mut body = build_chat_completions_body(data, &self.model, model_category)?;
self.patch_request_body(&mut body);
let builder = aws_fetch(
@@ -126,10 +126,10 @@ impl BedrockClient {
}
}
-async fn send_message(
+async fn chat_completions(
builder: RequestBuilder,
model_category: &ModelCategory,
-) -> Result<CompletionOutput> {
+) -> Result<ChatCompletionsOutput> {
let res = builder.send().await?;
let status = res.status();
let data: Value = res.json().await?;
@@ -140,13 +140,13 @@ async fn send_message(
debug!("non-stream-data: {data}");
match model_category {
- ModelCategory::Anthropic => claude_extract_completion(&data),
- ModelCategory::MetaLlama3 => llama_extract_completion(&data),
- ModelCategory::Mistral => mistral_extract_completion(&data),
+ ModelCategory::Anthropic => claude_extract_chat_completions(&data),
+ ModelCategory::MetaLlama3 => llama_extract_chat_completions(&data),
+ ModelCategory::Mistral => mistral_extract_chat_completions(&data),
}
}
-async fn send_message_streaming(
+async fn chat_completions_streaming(
builder: RequestBuilder,
handler: &mut SseHandler,
model_category: &ModelCategory,
@@ -211,14 +211,14 @@ async fn send_message_streaming(
Ok(())
}
-fn build_body(
- data: CompletionData,
+fn build_chat_completions_body(
+ data: ChatCompletionsData,
model: &Model,
model_category: &ModelCategory,
) -> Result<Value> {
match model_category {
ModelCategory::Anthropic => {
- let mut body = claude_build_body(data, model)?;
+ let mut body = claude_build_chat_completions_body(data, model)?;
if let Some(body_obj) = body.as_object_mut() {
body_obj.remove("model");
body_obj.remove("stream");
@@ -226,13 +226,19 @@ fn build_body(
body["anthropic_version"] = "bedrock-2023-05-31".into();
Ok(body)
}
- ModelCategory::MetaLlama3 => meta_llama_build_body(data, model, LLAMA3_PROMPT_FORMAT),
- ModelCategory::Mistral => mistral_build_body(data, model),
+ ModelCategory::MetaLlama3 => {
+ meta_llama_build_chat_completions_body(data, model, LLAMA3_PROMPT_FORMAT)
+ }
+ ModelCategory::Mistral => mistral_build_chat_completions_body(data, model),
}
}
-fn meta_llama_build_body(data: CompletionData, model: &Model, pt: PromptFormat) -> Result<Value> {
- let CompletionData {
+fn meta_llama_build_chat_completions_body(
+ data: ChatCompletionsData,
+ model: &Model,
+ pt: PromptFormat,
+) -> Result<Value> {
+ let ChatCompletionsData {
messages,
temperature,
top_p,
@@ -255,8 +261,8 @@ fn meta_llama_build_body(data: CompletionData, model: &Model, pt: PromptFormat)
Ok(body)
}
-fn mistral_build_body(data: CompletionData, model: &Model) -> Result<Value> {
- let CompletionData {
+fn mistral_build_chat_completions_body(data: ChatCompletionsData, model: &Model) -> Result<Value> {
+ let ChatCompletionsData {
messages,
temperature,
top_p,
@@ -279,11 +285,11 @@ fn mistral_build_body(data: CompletionData, model: &Model) -> Result<Value> {
Ok(body)
}
-fn llama_extract_completion(data: &Value) -> Result<CompletionOutput> {
+fn llama_extract_chat_completions(data: &Value) -> Result<ChatCompletionsOutput> {
let text = data["generation"]
.as_str()
.ok_or_else(|| anyhow!("Invalid response data: {data}"))?;
- let output = CompletionOutput {
+ let output = ChatCompletionsOutput {
text: text.to_string(),
tool_calls: vec![],
id: None,
@@ -293,11 +299,11 @@ fn llama_extract_completion(data: &Value) -> Result<CompletionOutput> {
Ok(output)
}
-fn mistral_extract_completion(data: &Value) -> Result<CompletionOutput> {
+fn mistral_extract_chat_completions(data: &Value) -> Result<ChatCompletionsOutput> {
let text = data["outputs"][0]["text"]
.as_str()
.ok_or_else(|| anyhow!("Invalid response data: {data}"))?;
- Ok(CompletionOutput::new(text))
+ Ok(ChatCompletionsOutput::new(text))
}
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
diff --git a/src/client/claude.rs b/src/client/claude.rs
index 3194fe2..6533a9a 100644
--- a/src/client/claude.rs
+++ b/src/client/claude.rs
@@ -1,7 +1,8 @@
use super::{
- catch_error, extract_system_message, message::*, sse_stream, ClaudeClient, Client,
- CompletionData, CompletionOutput, ExtraConfig, ImageUrl, MessageContent, MessageContentPart,
- Model, ModelData, ModelPatches, PromptAction, PromptKind, SseHandler, SseMmessage, ToolCall,
+ catch_error, extract_system_message, message::*, sse_stream, ChatCompletionsData,
+ ChatCompletionsOutput, ClaudeClient, Client, ExtraConfig, ImageUrl, MessageContent,
+ MessageContentPart, Model, ModelData, ModelPatches, PromptAction, PromptKind, SseHandler,
+ SseMmessage, ToolCall,
};
use anyhow::{bail, Context, Result};
@@ -27,14 +28,14 @@ impl ClaudeClient {
pub const PROMPTS: [PromptAction<'static>; 1] =
[("api_key", "API Key:", true, PromptKind::String)];
- fn request_builder(
+ fn chat_completions_builder(
&self,
client: &ReqwestClient,
- data: CompletionData,
+ data: ChatCompletionsData,
) -> Result<RequestBuilder> {
let api_key = self.get_api_key().ok();
- let mut body = claude_build_body(data, &self.model)?;
+ let mut body = claude_build_chat_completions_body(data, &self.model)?;
self.patch_request_body(&mut body);
let url = API_BASE;
@@ -55,11 +56,11 @@ impl ClaudeClient {
impl_client_trait!(
ClaudeClient,
- claude_send_message,
- claude_send_message_streaming
+ claude_chat_completions,
+ claude_chat_completions_streaming
);
-pub async fn claude_send_message(builder: RequestBuilder) -> Result<CompletionOutput> {
+pub async fn claude_chat_completions(builder: RequestBuilder) -> Result<ChatCompletionsOutput> {
let res = builder.send().await?;
let status = res.status();
let data: Value = res.json().await?;
@@ -67,10 +68,10 @@ pub async fn claude_send_message(builder: RequestBuilder) -> Result<CompletionOu
catch_error(&data, status.as_u16())?;
}
debug!("non-stream-data: {data}");
- claude_extract_completion(&data)
+ claude_extract_chat_completions(&data)
}
-pub async fn claude_send_message_streaming(
+pub async fn claude_chat_completions_streaming(
builder: RequestBuilder,
handler: &mut SseHandler,
) -> Result<()> {
@@ -135,8 +136,11 @@ pub async fn claude_send_message_streaming(
sse_stream(builder, handle).await
}
-pub fn claude_build_body(data: CompletionData, model: &Model) -> Result<Value> {
- let CompletionData {
+pub fn claude_build_chat_completions_body(
+ data: ChatCompletionsData,
+ model: &Model,
+) -> Result<Value> {
+ let ChatCompletionsData {
mut messages,
temperature,
top_p,
@@ -269,7 +273,7 @@ pub fn claude_build_body(data: CompletionData, model: &Model) -> Result<Value> {
Ok(body)
}
-pub fn claude_extract_completion(data: &Value) -> Result<CompletionOutput> {
+pub fn claude_extract_chat_completions(data: &Value) -> Result<ChatCompletionsOutput> {
let text = data["content"][0]["text"].as_str().unwrap_or_default();
let mut tool_calls = vec![];
@@ -303,7 +307,7 @@ pub fn claude_extract_completion(data: &Value) -> Result<CompletionOutput> {
bail!("Invalid response data: {data}");
}
- let output = CompletionOutput {
+ let output = ChatCompletionsOutput {
text: text.to_string(),
tool_calls,
id: data["id"].as_str().map(|v| v.to_string()),
diff --git a/src/client/cloudflare.rs b/src/client/cloudflare.rs
index 15dd5bb..966cee4 100644
--- a/src/client/cloudflare.rs
+++ b/src/client/cloudflare.rs
@@ -1,5 +1,5 @@
use super::{
- catch_error, sse_stream, Client, CloudflareClient, CompletionData, CompletionOutput,
+ catch_error, sse_stream, ChatCompletionsData, ChatCompletionsOutput, Client, CloudflareClient,
ExtraConfig, Model, ModelData, ModelPatches, PromptAction, PromptKind, SseHandler, SseMmessage,
};
@@ -30,15 +30,15 @@ impl CloudflareClient {
("api_key", "API Key:", true, PromptKind::String),
];
- fn request_builder(
+ fn chat_completions_builder(
&self,
client: &ReqwestClient,
- data: CompletionData,
+ data: ChatCompletionsData,
) -> Result<RequestBuilder> {
let account_id = self.get_account_id()?;
let api_key = self.get_api_key()?;
- 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 url = format!(
@@ -54,9 +54,13 @@ impl CloudflareClient {
}
}
-impl_client_trait!(CloudflareClient, send_message, send_message_streaming);
+impl_client_trait!(
+ CloudflareClient,
+ 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: Value = res.json().await?;
@@ -65,10 +69,13 @@ async fn send_message(builder: RequestBuilder) -> Result<CompletionOutput> {
}
debug!("non-stream-data: {data}");
- extract_completion(&data)
+ extract_chat_completions(&data)
}
-async fn send_message_streaming(builder: RequestBuilder, handler: &mut SseHandler) -> Result<()> {
+async fn chat_completions_streaming(
+ builder: RequestBuilder,
+ handler: &mut SseHandler,
+) -> Result<()> {
let handle = |message: SseMmessage| -> Result<bool> {
if message.data == "[DONE]" {
return Ok(true);
@@ -83,8 +90,8 @@ async fn send_message_streaming(builder: RequestBuilder, handler: &mut SseHandle
sse_stream(builder, handle).await
}
-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,
@@ -113,10 +120,10 @@ fn build_body(data: CompletionData, model: &Model) -> Result<Value> {
Ok(body)
}
-fn extract_completion(data: &Value) -> Result<CompletionOutput> {
+fn extract_chat_completions(data: &Value) -> Result<ChatCompletionsOutput> {
let text = data["result"]["response"]
.as_str()
.ok_or_else(|| anyhow!("Invalid response data: {data}"))?;
- Ok(CompletionOutput::new(text))
+ Ok(ChatCompletionsOutput::new(text))
}
diff --git a/src/client/cohere.rs b/src/client/cohere.rs
index e093778..e0a5eec 100644
--- a/src/client/cohere.rs
+++ b/src/client/cohere.rs
@@ -1,7 +1,7 @@
use super::{
- catch_error, extract_system_message, json_stream, message::*, Client, CohereClient,
- CompletionData, CompletionOutput, ExtraConfig, Model, ModelData, ModelPatches, PromptAction,
- PromptKind, SseHandler, ToolCall,
+ catch_error, extract_system_message, json_stream, message::*, ChatCompletionsData,
+ ChatCompletionsOutput, Client, CohereClient, ExtraConfig, Model, ModelData, ModelPatches,
+ PromptAction, PromptKind, SseHandler, ToolCall,
};
use anyhow::{bail, Result};
@@ -27,14 +27,14 @@ impl CohereClient {
pub const PROMPTS: [PromptAction<'static>; 1] =
[("api_key", "API Key:", true, PromptKind::String)];
- fn request_builder(
+ fn chat_completions_builder(
&self,
client: &ReqwestClient,
- data: CompletionData,
+ data: ChatCompletionsData,
) -> Result<RequestBuilder> {
let api_key = self.get_api_key()?;
- 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 url = API_URL;
@@ -47,9 +47,9 @@ impl CohereClient {
}
}
-impl_client_trait!(CohereClient, send_message, send_message_streaming);
+impl_client_trait!(CohereClient, 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: Value = res.json().await?;
@@ -58,10 +58,13 @@ async fn send_message(builder: RequestBuilder) -> Result<CompletionOutput> {
}
debug!("non-stream-data: {data}");
- extract_completion(&data)
+ extract_chat_completions(&data)
}
-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() {
@@ -97,8 +100,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 {
mut messages,
temperature,
top_p,
@@ -224,7 +227,7 @@ fn build_body(data: CompletionData, model: &Model) -> Result<Value> {
Ok(body)
}
-fn extract_completion(data: &Value) -> Result<CompletionOutput> {
+fn extract_chat_completions(data: &Value) -> Result<ChatCompletionsOutput> {
let text = data["text"].as_str().unwrap_or_default();
let mut tool_calls = vec![];
@@ -246,7 +249,7 @@ fn extract_completion(data: &Value) -> Result<CompletionOutput> {
if text.is_empty() && tool_calls.is_empty() {
bail!("Invalid response data: {data}");
}
- let output = CompletionOutput {
+ let output = ChatCompletionsOutput {
text: text.to_string(),
tool_calls,
id: data["generation_id"].as_str().map(|v| v.to_string()),
diff --git a/src/client/common.rs b/src/client/common.rs
index 336fa78..96ec90b 100644
--- a/src/client/common.rs
+++ b/src/client/common.rs
@@ -195,28 +195,28 @@ macro_rules! client_common_fns {
#[macro_export]
macro_rules! impl_client_trait {
- ($client:ident, $send_message:path, $send_message_streaming:path) => {
+ ($client:ident, $chat_completions:path, $chat_completions_streaming:path) => {
#[async_trait::async_trait]
impl $crate::client::Client for $crate::client::$client {
client_common_fns!();
- async fn send_message_inner(
+ async fn chat_completions_inner(
&self,
client: &reqwest::Client,
- data: $crate::client::CompletionData,
- ) -> anyhow::Result<$crate::client::CompletionOutput> {
- let builder = self.request_builder(client, data)?;
- $send_message(builder).await
+ data: $crate::client::ChatCompletionsData,
+ ) -> anyhow::Result<$crate::client::ChatCompletionsOutput> {
+ let builder = self.chat_completions_builder(client, data)?;
+ $chat_completions(builder).await
}
- async fn send_message_streaming_inner(
+ async fn chat_completions_streaming_inner(
&self,
client: &reqwest::Client,
handler: &mut $crate::client::SseHandler,
- data: $crate::client::CompletionData,
+ data: $crate::client::ChatCompletionsData,
) -> Result<()> {
- let builder = self.request_builder(client, data)?;
- $send_message_streaming(builder, handler).await
+ let builder = self.chat_completions_builder(client, data)?;
+ $chat_completions_streaming(builder, handler).await
}
}
};
@@ -282,20 +282,24 @@ pub trait Client: Sync + Send {
Ok(client)
}
- async fn send_message(&self, input: Input) -> Result<CompletionOutput> {
+ async fn chat_completions(&self, input: Input) -> Result<ChatCompletionsOutput> {
if self.global_config().read().dry_run {
let content = input.echo_messages();
- return Ok(CompletionOutput::new(&content));
+ return Ok(ChatCompletionsOutput::new(&content));
}
let client = self.build_client()?;
let data = input.prepare_completion_data(self.model(), false)?;
- self.send_message_inner(&client, data)
+ self.chat_completions_inner(&client, data)
.await
.with_context(|| "Failed to get answer")
}
- async fn send_message_streaming(&self, input: &Input, handler: &mut SseHandler) -> Result<()> {
+ async fn chat_completions_streaming(
+ &self,
+ input: &Input,
+ handler: &mut SseHandler,
+ ) -> Result<()> {
async fn watch_abort(abort: AbortSignal) {
loop {
if abort.aborted() {
@@ -319,7 +323,7 @@ pub trait Client: Sync + Send {
}
let client = self.build_client()?;
let data = input.prepare_completion_data(self.model(), true)?;
- self.send_message_streaming_inner(&client, handler, data).await
+ self.chat_completions_streaming_inner(&client, handler, data).await
} => {
handler.done()?;
ret.with_context(|| "Failed to get answer")
@@ -340,17 +344,17 @@ pub trait Client: Sync + Send {
}
}
- async fn send_message_inner(
+ async fn chat_completions_inner(
&self,
client: &ReqwestClient,
- data: CompletionData,
- ) -> Result<CompletionOutput>;
+ data: ChatCompletionsData,
+ ) -> Result<ChatCompletionsOutput>;
- async fn send_message_streaming_inner(
+ async fn chat_completions_streaming_inner(
&self,
client: &ReqwestClient,
handler: &mut SseHandler,
- data: CompletionData,
+ data: ChatCompletionsData,
) -> Result<()>;
}
@@ -391,7 +395,7 @@ pub fn select_model_patch<'a>(
}
#[derive(Debug)]
-pub struct CompletionData {
+pub struct ChatCompletionsData {
pub messages: Vec<Message>,
pub temperature: Option<f64>,
pub top_p: Option<f64>,
@@ -400,7 +404,7 @@ pub struct CompletionData {
}
#[derive(Debug, Clone, Default)]
-pub struct CompletionOutput {
+pub struct ChatCompletionsOutput {
pub text: String,
pub tool_calls: Vec<ToolCall>,
pub id: Option<String>,
@@ -408,7 +412,7 @@ pub struct CompletionOutput {
pub output_tokens: Option<u64>,
}
-impl CompletionOutput {
+impl ChatCompletionsOutput {
pub fn new(text: &str) -> Self {
Self {
text: text.to_string(),
@@ -473,7 +477,7 @@ pub async fn send_stream(
let mut handler = SseHandler::new(tx, abort.clone());
let (send_ret, rend_ret) = tokio::join!(
- client.send_message_streaming(input, &mut handler),
+ client.chat_completions_streaming(input, &mut handler),
render_stream(rx, config, abort.clone()),
);
if let Err(err) = rend_ret {
@@ -497,7 +501,7 @@ pub async fn send_stream(
}
#[allow(unused)]
-pub async fn send_message_as_streaming<F, Fut>(
+pub async fn chat_completions_as_streaming<F, Fut>(
builder: RequestBuilder,
handler: &mut SseHandler,
f: F,
diff --git a/src/client/ernie.rs b/src/client/ernie.rs
index 49e158b..097ee68 100644
--- a/src/client/ernie.rs
+++ b/src/client/ernie.rs
@@ -1,7 +1,7 @@
use super::{
- access_token::*, maybe_catch_error, patch_system_message, sse_stream, Client, CompletionData,
- CompletionOutput, ErnieClient, ExtraConfig, Model, ModelData, ModelPatches, PromptAction,
- PromptKind, SseHandler, SseMmessage,
+ access_token::*, maybe_catch_error, patch_system_message, sse_stream, ChatCompletionsData,
+ ChatCompletionsOutput, Client, ErnieClient, ExtraConfig, Model, ModelData, ModelPatches,
+ PromptAction, PromptKind, SseHandler, SseMmessage,
};
use anyhow::{anyhow, Context, Result};
@@ -31,12 +31,12 @@ impl ErnieClient {
("secret_key", "Secret Key:", true, PromptKind::String),
];
- fn request_builder(
+ fn chat_completions_builder(
&self,
client: &ReqwestClient,
- data: CompletionData,
+ data: ChatCompletionsData,
) -> Result<RequestBuilder> {
- 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 access_token = get_access_token(self.name())?;
@@ -81,36 +81,39 @@ impl ErnieClient {
impl Client for ErnieClient {
client_common_fns!();
- async fn send_message_inner(
+ async fn chat_completions_inner(
&self,
client: &ReqwestClient,
- data: CompletionData,
- ) -> Result<CompletionOutput> {
+ data: ChatCompletionsData,
+ ) -> Result<ChatCompletionsOutput> {
self.prepare_access_token().await?;
- let builder = self.request_builder(client, data)?;
- send_message(builder).await
+ let builder = self.chat_completions_builder(client, data)?;
+ chat_completions(builder).await
}
- async fn send_message_streaming_inner(
+ async fn chat_completions_streaming_inner(
&self,
client: &ReqwestClient,
handler: &mut SseHandler,
- data: CompletionData,
+ data: ChatCompletionsData,
) -> Result<()> {
self.prepare_access_token().await?;
- let builder = self.request_builder(client, data)?;
- send_message_streaming(builder, handler).await
+ let builder = self.chat_completions_builder(client, data)?;
+ chat_completions_streaming(builder, handler).await
}
}
-async fn send_message(builder: RequestBuilder) -> Result<CompletionOutput> {
+async fn chat_completions(builder: RequestBuilder) -> Result<ChatCompletionsOutput> {
let data: Value = builder.send().await?.json().await?;
maybe_catch_error(&data)?;
debug!("non-stream-data: {data}");
- extract_completion_text(&data)
+ extract_chat_completions_text(&data)
}
-async fn send_message_streaming(builder: RequestBuilder, handler: &mut SseHandler) -> Result<()> {
+async fn chat_completions_streaming(
+ builder: RequestBuilder,
+ handler: &mut SseHandler,
+) -> Result<()> {
let handle = |message: SseMmessage| -> Result<bool> {
let data: Value = serde_json::from_str(&message.data)?;
debug!("stream-data: {data}");
@@ -123,8 +126,8 @@ async fn send_message_streaming(builder: RequestBuilder, handler: &mut SseHandle
sse_stream(builder, handle).await
}
-fn build_body(data: CompletionData, model: &Model) -> Value {
- let CompletionData {
+fn build_chat_completions_body(data: ChatCompletionsData, model: &Model) -> Value {
+ let ChatCompletionsData {
mut messages,
temperature,
top_p,
@@ -155,11 +158,11 @@ fn build_body(data: CompletionData, model: &Model) -> Value {
body
}
-fn extract_completion_text(data: &Value) -> Result<CompletionOutput> {
+fn extract_chat_completions_text(data: &Value) -> Result<ChatCompletionsOutput> {
let text = data["result"]
.as_str()
.ok_or_else(|| anyhow!("Invalid response data: {data}"))?;
- let output = CompletionOutput {
+ let output = ChatCompletionsOutput {
text: text.to_string(),
tool_calls: vec![],
id: data["id"].as_str().map(|v| v.to_string()),
diff --git a/src/client/gemini.rs b/src/client/gemini.rs
index 05614ce..5cc45c5 100644
--- a/src/client/gemini.rs
+++ b/src/client/gemini.rs
@@ -1,6 +1,6 @@
use super::{
- vertexai::*, Client, CompletionData, ExtraConfig, GeminiClient, Model, ModelData, ModelPatches,
- PromptAction, PromptKind,
+ vertexai::*, ChatCompletionsData, Client, ExtraConfig, GeminiClient, Model, ModelData,
+ ModelPatches, PromptAction, PromptKind,
};
use anyhow::Result;
@@ -13,8 +13,6 @@ const API_BASE: &str = "https://generativelanguage.googleapis.com/v1beta/models/
pub struct GeminiConfig {
pub name: Option<String>,
pub api_key: Option<String>,
- #[serde(rename = "safetySettings")]
- pub safety_settings: Option<serde_json::Value>,
#[serde(default)]
pub models: Vec<ModelData>,
pub patches: Option<ModelPatches>,
@@ -27,10 +25,10 @@ impl GeminiClient {
pub const PROMPTS: [PromptAction<'static>; 1] =
[("api_key", "API Key:", true, PromptKind::String)];
- fn request_builder(
+ fn chat_completions_builder(
&self,
client: &ReqwestClient,
- data: CompletionData,
+ data: ChatCompletionsData,
) -> Result<RequestBuilder> {
let api_key = self.get_api_key()?;
@@ -39,7 +37,7 @@ impl GeminiClient {
false => "generateContent",
};
- let mut body = gemini_build_body(data, &self.model)?;
+ let mut body = gemini_build_chat_completions_body(data, &self.model)?;
self.patch_request_body(&mut body);
let model = &self.model.name();
@@ -56,6 +54,6 @@ impl GeminiClient {
impl_client_trait!(
GeminiClient,
- crate::client::vertexai::gemini_send_message,
- crate::client::vertexai::gemini_send_message_streaming
+ crate::client::vertexai::gemini_chat_completions,
+ crate::client::vertexai::gemini_chat_completions_streaming
);
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,
diff --git a/src/client/openai.rs b/src/client/openai.rs
index bf719ae..3cdea24 100644
--- a/src/client/openai.rs
+++ b/src/client/openai.rs
@@ -1,7 +1,7 @@
use super::{
- catch_error, message::*, sse_stream, Client, CompletionData, CompletionOutput, ExtraConfig,
- Model, ModelData, ModelPatches, OpenAIClient, PromptAction, PromptKind, SseHandler,
- SseMmessage, ToolCall,
+ catch_error, message::*, sse_stream, ChatCompletionsData, ChatCompletionsOutput, Client,
+ ExtraConfig, Model, ModelData, ModelPatches, OpenAIClient, PromptAction, PromptKind,
+ SseHandler, SseMmessage, ToolCall,
};
use anyhow::{bail, Result};
@@ -30,15 +30,15 @@ impl OpenAIClient {
pub const PROMPTS: [PromptAction<'static>; 1] =
[("api_key", "API Key:", true, PromptKind::String)];
- fn request_builder(
+ fn chat_completions_builder(
&self,
client: &ReqwestClient,
- data: CompletionData,
+ data: ChatCompletionsData,
) -> Result<RequestBuilder> {
let api_key = self.get_api_key()?;
let api_base = self.get_api_base().unwrap_or_else(|_| API_BASE.to_string());
- let mut body = openai_build_body(data, &self.model);
+ let mut body = openai_build_chat_completions_body(data, &self.model);
self.patch_request_body(&mut body);
let url = format!("{api_base}/chat/completions");
@@ -55,7 +55,7 @@ impl OpenAIClient {
}
}
-pub async fn openai_send_message(builder: RequestBuilder) -> Result<CompletionOutput> {
+pub async fn openai_chat_completions(builder: RequestBuilder) -> Result<ChatCompletionsOutput> {
let res = builder.send().await?;
let status = res.status();
let data: Value = res.json().await?;
@@ -64,10 +64,10 @@ pub async fn openai_send_message(builder: RequestBuilder) -> Result<CompletionOu
}
debug!("non-stream-data: {data}");
- openai_extract_completion(&data)
+ openai_extract_chat_completions(&data)
}
-pub async fn openai_send_message_streaming(
+pub async fn openai_chat_completions_streaming(
builder: RequestBuilder,
handler: &mut SseHandler,
) -> Result<()> {
@@ -125,8 +125,8 @@ pub async fn openai_send_message_streaming(
sse_stream(builder, handle).await
}
-pub fn openai_build_body(data: CompletionData, model: &Model) -> Value {
- let CompletionData {
+pub fn openai_build_chat_completions_body(data: ChatCompletionsData, model: &Model) -> Value {
+ let ChatCompletionsData {
messages,
temperature,
top_p,
@@ -201,7 +201,7 @@ pub fn openai_build_body(data: CompletionData, model: &Model) -> Value {
body
}
-pub fn openai_extract_completion(data: &Value) -> Result<CompletionOutput> {
+pub fn openai_extract_chat_completions(data: &Value) -> Result<ChatCompletionsOutput> {
let text = data["choices"][0]["message"]["content"]
.as_str()
.unwrap_or_default();
@@ -231,7 +231,7 @@ pub fn openai_extract_completion(data: &Value) -> Result<CompletionOutput> {
if text.is_empty() && tool_calls.is_empty() {
bail!("Invalid response data: {data}");
}
- let output = CompletionOutput {
+ let output = ChatCompletionsOutput {
text: text.to_string(),
tool_calls,
id: data["id"].as_str().map(|v| v.to_string()),
@@ -243,6 +243,6 @@ pub fn openai_extract_completion(data: &Value) -> Result<CompletionOutput> {
impl_client_trait!(
OpenAIClient,
- openai_send_message,
- openai_send_message_streaming
+ openai_chat_completions,
+ openai_chat_completions_streaming
);
diff --git a/src/client/openai_compatible.rs b/src/client/openai_compatible.rs
index f6c0a85..74cd954 100644
--- a/src/client/openai_compatible.rs
+++ b/src/client/openai_compatible.rs
@@ -1,5 +1,5 @@
use super::{
- openai::*, Client, CompletionData, ExtraConfig, Model, ModelData, ModelPatches,
+ openai::*, ChatCompletionsData, Client, ExtraConfig, Model, ModelData, ModelPatches,
OpenAICompatibleClient, PromptAction, PromptKind, OPENAI_COMPATIBLE_PLATFORMS,
};
@@ -36,10 +36,10 @@ impl OpenAICompatibleClient {
),
];
- fn request_builder(
+ fn chat_completions_builder(
&self,
client: &ReqwestClient,
- data: CompletionData,
+ data: ChatCompletionsData,
) -> Result<RequestBuilder> {
let api_base = match self.get_api_base() {
Ok(v) => v,
@@ -60,7 +60,7 @@ impl OpenAICompatibleClient {
};
let api_key = self.get_api_key().ok();
- let mut body = openai_build_body(data, &self.model);
+ let mut body = openai_build_chat_completions_body(data, &self.model);
self.patch_request_body(&mut body);
let chat_endpoint = self
@@ -84,6 +84,6 @@ impl OpenAICompatibleClient {
impl_client_trait!(
OpenAICompatibleClient,
- crate::client::openai::openai_send_message,
- crate::client::openai::openai_send_message_streaming
+ crate::client::openai::openai_chat_completions,
+ crate::client::openai::openai_chat_completions_streaming
);
diff --git a/src/client/qianwen.rs b/src/client/qianwen.rs
index 3f4b73a..0230e21 100644
--- a/src/client/qianwen.rs
+++ b/src/client/qianwen.rs
@@ -1,5 +1,5 @@
use super::{
- maybe_catch_error, message::*, sse_stream, Client, CompletionData, CompletionOutput,
+ maybe_catch_error, message::*, sse_stream, ChatCompletionsData, ChatCompletionsOutput, Client,
ExtraConfig, Model, ModelData, ModelPatches, PromptAction, PromptKind, QianwenClient,
SseHandler, SseMmessage,
};
@@ -38,10 +38,10 @@ impl QianwenClient {
pub const PROMPTS: [PromptAction<'static>; 1] =
[("api_key", "API Key:", true, PromptKind::String)];
- fn request_builder(
+ fn chat_completions_builder(
&self,
client: &ReqwestClient,
- data: CompletionData,
+ data: ChatCompletionsData,
) -> Result<RequestBuilder> {
let api_key = self.get_api_key()?;
@@ -51,7 +51,7 @@ impl QianwenClient {
true => API_URL_VL,
false => API_URL,
};
- let (mut body, has_upload) = build_body(data, &self.model)?;
+ let (mut body, has_upload) = build_chat_completions_body(data, &self.model)?;
self.patch_request_body(&mut body);
debug!("Qianwen Request: {url} {body}");
@@ -72,39 +72,39 @@ impl QianwenClient {
impl Client for QianwenClient {
client_common_fns!();
- async fn send_message_inner(
+ async fn chat_completions_inner(
&self,
client: &ReqwestClient,
- mut data: CompletionData,
- ) -> Result<CompletionOutput> {
+ mut data: ChatCompletionsData,
+ ) -> Result<ChatCompletionsOutput> {
let api_key = self.get_api_key()?;
patch_messages(self.model.name(), &api_key, &mut data.messages).await?;
- let builder = self.request_builder(client, data)?;
- send_message(builder, &self.model).await
+ let builder = self.chat_completions_builder(client, data)?;
+ chat_completions(builder, &self.model).await
}
- async fn send_message_streaming_inner(
+ async fn chat_completions_streaming_inner(
&self,
client: &ReqwestClient,
handler: &mut SseHandler,
- mut data: CompletionData,
+ mut data: ChatCompletionsData,
) -> Result<()> {
let api_key = self.get_api_key()?;
patch_messages(self.model.name(), &api_key, &mut data.messages).await?;
- let builder = self.request_builder(client, data)?;
- send_message_streaming(builder, handler, &self.model).await
+ let builder = self.chat_completions_builder(client, data)?;
+ chat_completions_streaming(builder, handler, &self.model).await
}
}
-async fn send_message(builder: RequestBuilder, model: &Model) -> Result<CompletionOutput> {
+async fn chat_completions(builder: RequestBuilder, model: &Model) -> Result<ChatCompletionsOutput> {
let data: Value = builder.send().await?.json().await?;
maybe_catch_error(&data)?;
debug!("non-stream-data: {data}");
- extract_completion_text(&data, model)
+ extract_chat_completions_text(&data, model)
}
-async fn send_message_streaming(
+async fn chat_completions_streaming(
builder: RequestBuilder,
handler: &mut SseHandler,
model: &Model,
@@ -133,8 +133,8 @@ async fn send_message_streaming(
sse_stream(builder, handle).await
}
-fn build_body(data: CompletionData, model: &Model) -> Result<(Value, bool)> {
- let CompletionData {
+fn build_chat_completions_body(data: ChatCompletionsData, model: &Model) -> Result<(Value, bool)> {
+ let ChatCompletionsData {
messages,
temperature,
top_p,
@@ -210,7 +210,7 @@ fn build_body(data: CompletionData, model: &Model) -> Result<(Value, bool)> {
Ok((body, has_upload))
}
-fn extract_completion_text(data: &Value, model: &Model) -> Result<CompletionOutput> {
+fn extract_chat_completions_text(data: &Value, model: &Model) -> Result<ChatCompletionsOutput> {
let err = || anyhow!("Invalid response data: {data}");
let text = if model.name() == "qwen-long" {
data["output"]["choices"][0]["message"]["content"]
@@ -223,7 +223,7 @@ fn extract_completion_text(data: &Value, model: &Model) -> Result<CompletionOutp
} else {
data["output"]["text"].as_str().ok_or_else(err)?
};
- let output = CompletionOutput {
+ let output = ChatCompletionsOutput {
text: text.to_string(),
tool_calls: vec![],
id: data["request_id"].as_str().map(|v| v.to_string()),
diff --git a/src/client/replicate.rs b/src/client/replicate.rs
index 3b0787a..92c7e18 100644
--- a/src/client/replicate.rs
+++ b/src/client/replicate.rs
@@ -1,5 +1,5 @@
use super::{
- catch_error, prompt_format::*, sse_stream, Client, CompletionData, CompletionOutput,
+ catch_error, prompt_format::*, sse_stream, ChatCompletionsData, ChatCompletionsOutput, Client,
ExtraConfig, Model, ModelData, ModelPatches, PromptAction, PromptKind, ReplicateClient,
SseHandler, SseMmessage,
};
@@ -29,13 +29,13 @@ impl ReplicateClient {
pub const PROMPTS: [PromptAction<'static>; 1] =
[("api_key", "API Key:", true, PromptKind::String)];
- fn request_builder(
+ fn chat_completions_builder(
&self,
client: &ReqwestClient,
- data: CompletionData,
+ data: ChatCompletionsData,
api_key: &str,
) -> Result<RequestBuilder> {
- 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 url = format!("{API_BASE}/models/{}/predictions", self.model.name());
@@ -52,33 +52,33 @@ impl ReplicateClient {
impl Client for ReplicateClient {
client_common_fns!();
- async fn send_message_inner(
+ async fn chat_completions_inner(
&self,
client: &ReqwestClient,
- data: CompletionData,
- ) -> Result<CompletionOutput> {
+ data: ChatCompletionsData,
+ ) -> Result<ChatCompletionsOutput> {
let api_key = self.get_api_key()?;
- let builder = self.request_builder(client, data, &api_key)?;
- send_message(client, builder, &api_key).await
+ let builder = self.chat_completions_builder(client, data, &api_key)?;
+ chat_completions(client, builder, &api_key).await
}
- async fn send_message_streaming_inner(
+ async fn chat_completions_streaming_inner(
&self,
client: &ReqwestClient,
handler: &mut SseHandler,
- data: CompletionData,
+ data: ChatCompletionsData,
) -> Result<()> {
let api_key = self.get_api_key()?;
- let builder = self.request_builder(client, data, &api_key)?;
- send_message_streaming(client, builder, handler).await
+ let builder = self.chat_completions_builder(client, data, &api_key)?;
+ chat_completions_streaming(client, builder, handler).await
}
}
-async fn send_message(
+async fn chat_completions(
client: &ReqwestClient,
builder: RequestBuilder,
api_key: &str,
-) -> Result<CompletionOutput> {
+) -> Result<ChatCompletionsOutput> {
let res = builder.send().await?;
let status = res.status();
let data: Value = res.json().await?;
@@ -101,14 +101,14 @@ async fn send_message(
let err = || anyhow!("Invalid response data: {prediction_data}");
let status = prediction_data["status"].as_str().ok_or_else(err)?;
if status == "succeeded" {
- return extract_completion(&prediction_data);
+ return extract_chat_completions(&prediction_data);
} else if status == "failed" || status == "canceled" {
return Err(err());
}
}
}
-async fn send_message_streaming(
+async fn chat_completions_streaming(
client: &ReqwestClient,
builder: RequestBuilder,
handler: &mut SseHandler,
@@ -135,8 +135,8 @@ async fn send_message_streaming(
sse_stream(sse_builder, handle).await
}
-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,
@@ -173,7 +173,7 @@ fn build_body(data: CompletionData, model: &Model) -> Result<Value> {
Ok(body)
}
-fn extract_completion(data: &Value) -> Result<CompletionOutput> {
+fn extract_chat_completions(data: &Value) -> Result<ChatCompletionsOutput> {
let text = data["output"]
.as_array()
.map(|parts| {
@@ -185,7 +185,7 @@ fn extract_completion(data: &Value) -> Result<CompletionOutput> {
})
.ok_or_else(|| anyhow!("Invalid response data: {data}"))?;
- let output = CompletionOutput {
+ let output = ChatCompletionsOutput {
text: text.to_string(),
tool_calls: vec![],
id: data["id"].as_str().map(|v| v.to_string()),
diff --git a/src/client/vertexai.rs b/src/client/vertexai.rs
index 102b5cb..b40247d 100644
--- a/src/client/vertexai.rs
+++ b/src/client/vertexai.rs
@@ -1,7 +1,7 @@
use super::{
- access_token::*, catch_error, json_stream, message::*, patch_system_message, Client,
- CompletionData, CompletionOutput, ExtraConfig, Model, ModelData, ModelPatches, PromptAction,
- PromptKind, SseHandler, ToolCall, VertexAIClient,
+ access_token::*, catch_error, json_stream, message::*, patch_system_message,
+ ChatCompletionsData, ChatCompletionsOutput, Client, ExtraConfig, Model, ModelData,
+ ModelPatches, PromptAction, PromptKind, SseHandler, ToolCall, VertexAIClient,
};
use anyhow::{anyhow, bail, Context, Result};
@@ -18,8 +18,6 @@ pub struct VertexAIConfig {
pub project_id: Option<String>,
pub location: Option<String>,
pub adc_file: Option<String>,
- #[serde(rename = "safetySettings")]
- pub safety_settings: Option<Value>,
#[serde(default)]
pub models: Vec<ModelData>,
pub patches: Option<ModelPatches>,
@@ -35,10 +33,10 @@ impl VertexAIClient {
("location", "Location", true, PromptKind::String),
];
- fn request_builder(
+ fn chat_completions_builder(
&self,
client: &ReqwestClient,
- data: CompletionData,
+ data: ChatCompletionsData,
) -> Result<RequestBuilder> {
let project_id = self.get_project_id()?;
let location = self.get_location()?;
@@ -52,7 +50,7 @@ impl VertexAIClient {
};
let url = format!("{base_url}/google/models/{}:{func}", self.model.name());
- let mut body = gemini_build_body(data, &self.model)?;
+ let mut body = gemini_build_chat_completions_body(data, &self.model)?;
self.patch_request_body(&mut body);
debug!("VertexAI Request: {url} {body}");
@@ -67,29 +65,29 @@ impl VertexAIClient {
impl Client for VertexAIClient {
client_common_fns!();
- async fn send_message_inner(
+ async fn chat_completions_inner(
&self,
client: &ReqwestClient,
- data: CompletionData,
- ) -> Result<CompletionOutput> {
+ data: ChatCompletionsData,
+ ) -> Result<ChatCompletionsOutput> {
prepare_gcloud_access_token(client, self.name(), &self.config.adc_file).await?;
- let builder = self.request_builder(client, data)?;
- gemini_send_message(builder).await
+ let builder = self.chat_completions_builder(client, data)?;
+ gemini_chat_completions(builder).await
}
- async fn send_message_streaming_inner(
+ async fn chat_completions_streaming_inner(
&self,
client: &ReqwestClient,
handler: &mut SseHandler,
- data: CompletionData,
+ data: ChatCompletionsData,
) -> Result<()> {
prepare_gcloud_access_token(client, self.name(), &self.config.adc_file).await?;
- let builder = self.request_builder(client, data)?;
- gemini_send_message_streaming(builder, handler).await
+ let builder = self.chat_completions_builder(client, data)?;
+ gemini_chat_completions_streaming(builder, handler).await
}
}
-pub async fn gemini_send_message(builder: RequestBuilder) -> Result<CompletionOutput> {
+pub async fn gemini_chat_completions(builder: RequestBuilder) -> Result<ChatCompletionsOutput> {
let res = builder.send().await?;
let status = res.status();
let data: Value = res.json().await?;
@@ -97,10 +95,10 @@ pub async fn gemini_send_message(builder: RequestBuilder) -> Result<CompletionOu
catch_error(&data, status.as_u16())?;
}
debug!("non-stream-data: {data}");
- gemini_extract_completion_text(&data)
+ gemini_extract_chat_completions_text(&data)
}
-pub async fn gemini_send_message_streaming(
+pub async fn gemini_chat_completions_streaming(
builder: RequestBuilder,
handler: &mut SseHandler,
) -> Result<()> {
@@ -140,7 +138,7 @@ pub async fn gemini_send_message_streaming(
Ok(())
}
-fn gemini_extract_completion_text(data: &Value) -> Result<CompletionOutput> {
+fn gemini_extract_chat_completions_text(data: &Value) -> Result<ChatCompletionsOutput> {
let text = data["candidates"][0]["content"]["parts"][0]["text"]
.as_str()
.unwrap_or_default();
@@ -171,7 +169,7 @@ fn gemini_extract_completion_text(data: &Value) -> Result<CompletionOutput> {
bail!("Invalid response data: {data}");
}
}
- let output = CompletionOutput {
+ let output = ChatCompletionsOutput {
text: text.to_string(),
tool_calls,
id: None,
@@ -181,8 +179,11 @@ fn gemini_extract_completion_text(data: &Value) -> Result<CompletionOutput> {
Ok(output)
}
-pub(crate) fn gemini_build_body(data: CompletionData, model: &Model) -> Result<Value> {
- let CompletionData {
+pub(crate) fn gemini_build_chat_completions_body(
+ data: ChatCompletionsData,
+ model: &Model,
+) -> Result<Value> {
+ let ChatCompletionsData {
mut messages,
temperature,
top_p,
diff --git a/src/client/vertexai_claude.rs b/src/client/vertexai_claude.rs
index fd06518..bdce7d8 100644
--- a/src/client/vertexai_claude.rs
+++ b/src/client/vertexai_claude.rs
@@ -1,6 +1,7 @@
use super::{
- access_token::*, claude::*, vertexai::*, Client, CompletionData, CompletionOutput, ExtraConfig,
- Model, ModelData, ModelPatches, PromptAction, PromptKind, SseHandler, VertexAIClaudeClient,
+ access_token::*, claude::*, vertexai::*, ChatCompletionsData, ChatCompletionsOutput, Client,
+ ExtraConfig, Model, ModelData, ModelPatches, PromptAction, PromptKind, SseHandler,
+ VertexAIClaudeClient,
};
use anyhow::Result;
@@ -29,10 +30,10 @@ impl VertexAIClaudeClient {
("location", "Location", true, PromptKind::String),
];
- fn request_builder(
+ fn chat_completions_builder(
&self,
client: &ReqwestClient,
- data: CompletionData,
+ data: ChatCompletionsData,
) -> Result<RequestBuilder> {
let project_id = self.get_project_id()?;
let location = self.get_location()?;
@@ -44,7 +45,7 @@ impl VertexAIClaudeClient {
self.model.name()
);
- let mut body = claude_build_body(data, &self.model)?;
+ let mut body = claude_build_chat_completions_body(data, &self.model)?;
self.patch_request_body(&mut body);
if let Some(body_obj) = body.as_object_mut() {
body_obj.remove("model");
@@ -63,24 +64,24 @@ impl VertexAIClaudeClient {
impl Client for VertexAIClaudeClient {
client_common_fns!();
- async fn send_message_inner(
+ async fn chat_completions_inner(
&self,
client: &ReqwestClient,
- data: CompletionData,
- ) -> Result<CompletionOutput> {
+ data: ChatCompletionsData,
+ ) -> Result<ChatCompletionsOutput> {
prepare_gcloud_access_token(client, self.name(), &self.config.adc_file).await?;
- let builder = self.request_builder(client, data)?;
- claude_send_message(builder).await
+ let builder = self.chat_completions_builder(client, data)?;
+ claude_chat_completions(builder).await
}
- async fn send_message_streaming_inner(
+ async fn chat_completions_streaming_inner(
&self,
client: &ReqwestClient,
handler: &mut SseHandler,
- data: CompletionData,
+ data: ChatCompletionsData,
) -> Result<()> {
prepare_gcloud_access_token(client, self.name(), &self.config.adc_file).await?;
- let builder = self.request_builder(client, data)?;
- claude_send_message_streaming(builder, handler).await
+ let builder = self.chat_completions_builder(client, data)?;
+ claude_chat_completions_streaming(builder, handler).await
}
}
diff --git a/src/config/input.rs b/src/config/input.rs
index 44bbb4b..ae94799 100644
--- a/src/config/input.rs
+++ b/src/config/input.rs
@@ -1,7 +1,7 @@
use super::{role::Role, session::Session, GlobalConfig};
use crate::client::{
- init_client, list_models, Client, CompletionData, ImageUrl, Message, MessageContent,
+ init_client, list_models, ChatCompletionsData, Client, ImageUrl, Message, MessageContent,
MessageContentPart, MessageRole, Model,
};
use crate::function::{ToolCallResult, ToolResults};
@@ -149,7 +149,11 @@ impl Input {
init_client(&self.config, Some(self.model()))
}
- pub fn prepare_completion_data(&self, model: &Model, stream: bool) -> Result<CompletionData> {
+ pub fn prepare_completion_data(
+ &self,
+ model: &Model,
+ stream: bool,
+ ) -> Result<ChatCompletionsData> {
if !self.medias.is_empty() && !model.supports_vision() {
bail!("The current model does not support vision.");
}
@@ -176,7 +180,7 @@ impl Input {
};
functions = config.function.select(function_matcher);
};
- Ok(CompletionData {
+ Ok(ChatCompletionsData {
messages,
temperature,
top_p,
diff --git a/src/main.rs b/src/main.rs
index 0bf5998..0a5e404 100644
--- a/src/main.rs
+++ b/src/main.rs
@@ -13,7 +13,7 @@ mod utils;
extern crate log;
use crate::cli::Cli;
-use crate::client::{list_models, send_stream, CompletionOutput};
+use crate::client::{list_models, send_stream, ChatCompletionsOutput};
use crate::config::{
Config, GlobalConfig, Input, InputContext, WorkingMode, CODE_ROLE, EXPLAIN_SHELL_ROLE,
SHELL_ROLE,
@@ -150,9 +150,9 @@ async fn start_directive(
let is_terminal_stdout = stdout().is_terminal();
let extract_code = !is_terminal_stdout && code_mode;
let (output, tool_call_results) = if no_stream || extract_code {
- let CompletionOutput {
+ let ChatCompletionsOutput {
text, tool_calls, ..
- } = client.send_message(input.clone()).await?;
+ } = client.chat_completions(input.clone()).await?;
if !tool_calls.is_empty() {
(String::new(), eval_tool_calls(config, tool_calls)?)
} else {
@@ -203,11 +203,11 @@ async fn shell_execute(config: &GlobalConfig, shell: &Shell, mut input: Input) -
let ret = if is_terminal_stdout {
let (spinner_tx, spinner_rx) = oneshot::channel();
tokio::spawn(run_spinner(" Generating", spinner_rx));
- let ret = client.send_message(input.clone()).await;
+ let ret = client.chat_completions(input.clone()).await;
let _ = spinner_tx.send(());
ret
} else {
- client.send_message(input.clone()).await
+ client.chat_completions(input.clone()).await
};
let mut eval_str = ret?.text;
if let Ok(true) = CODE_BLOCK_RE.is_match(&eval_str) {
diff --git a/src/repl/mod.rs b/src/repl/mod.rs
index f62f210..b78d39d 100644
--- a/src/repl/mod.rs
+++ b/src/repl/mod.rs
@@ -464,7 +464,7 @@ fn parse_command(line: &str) -> Option<(&str, Option<&str>)> {
async fn compress_session(config: &GlobalConfig) -> Result<()> {
let input = Input::from_str(config, config.read().summarize_prompt(), None);
let client = input.create_client()?;
- let summary = client.send_message(input).await?.text;
+ let summary = client.chat_completions(input).await?.text;
config.write().compress_session(&summary);
Ok(())
}
diff --git a/src/serve.rs b/src/serve.rs
index 5f43d88..2d375f8 100644
--- a/src/serve.rs
+++ b/src/serve.rs
@@ -1,7 +1,7 @@
use crate::{
client::{
- init_client, list_models, ClientConfig, CompletionData, CompletionOutput, Message, Model,
- ModelData, SseEvent, SseHandler,
+ init_client, list_models, ChatCompletionsData, ChatCompletionsOutput, ClientConfig,
+ Message, Model, ModelData, SseEvent, SseHandler,
},
config::{Config, GlobalConfig, Role},
utils::create_abort_signal,
@@ -270,7 +270,7 @@ impl Server {
let completion_id = generate_completion_id();
let created = Utc::now().timestamp();
- let completion_data: CompletionData = CompletionData {
+ let data: ChatCompletionsData = ChatCompletionsData {
messages,
temperature,
top_p,
@@ -306,7 +306,7 @@ impl Server {
}
tokio::select! {
_ = map_event(rx2, &tx, &mut is_first) => {}
- ret = client.send_message_streaming_inner(&http_client, &mut handler, completion_data) => {
+ ret = client.chat_completions_streaming_inner(&http_client, &mut handler, data) => {
if let Err(err) = ret {
send_first_event(&tx, Some(format!("{err:?}")), &mut is_first)
}
@@ -350,9 +350,7 @@ impl Server {
.body(BodyExt::boxed(StreamBody::new(stream)))?;
Ok(res)
} else {
- let output = client
- .send_message_inner(&http_client, completion_data)
- .await?;
+ let output = client.chat_completions_inner(&http_client, data).await?;
let res = Response::builder()
.header("Content-Type", "application/json")
.body(
@@ -452,7 +450,7 @@ fn create_frame(id: &str, model: &str, created: i64, content: &str, done: bool)
Frame::data(Bytes::from(output))
}
-fn ret_non_stream(id: &str, model: &str, created: i64, output: &CompletionOutput) -> Bytes {
+fn ret_non_stream(id: &str, model: &str, created: i64, output: &ChatCompletionsOutput) -> Bytes {
let id = output.id.as_deref().unwrap_or(id);
let input_tokens = output.input_tokens.unwrap_or_default();
let output_tokens = output.output_tokens.unwrap_or_default();