use futures_util::StreamExt;
use reqwest::Method;
use serde_json::Value;
use crate::{
client::Client,
error::Result,
http::encode_model_path,
request::set_stream,
streaming::{events_from_bytes, text_from_event, EventStream, StreamProtocol, TextStream},
ChatModelList,
};
pub type ChatRequest = Value;
pub type ChatResponse = Value;
#[derive(Clone, Debug)]
pub struct ChatService {
client: Client,
}
impl ChatService {
pub(crate) fn new(client: Client) -> Self {
Self { client }
}
pub fn models(&self) -> ChatModelsService {
ChatModelsService::new(self.client.clone())
}
pub fn completions(&self) -> ChatCompletionsService {
ChatCompletionsService::new(self.client.clone())
}
pub fn messages(&self) -> MessagesService {
MessagesService::new(self.client.clone())
}
pub fn gemini(&self) -> GeminiService {
GeminiService::new(self.client.clone())
}
}
#[derive(Clone, Debug)]
pub struct ChatModelsService {
client: Client,
}
impl ChatModelsService {
pub(crate) fn new(client: Client) -> Self {
Self { client }
}
pub async fn list(&self) -> Result<ChatModelList> {
self.client
.request_json::<ChatModelList, ()>(Method::GET, "/api/v1/models", None, None)
.await
}
}
#[derive(Clone, Debug)]
pub struct ChatCompletionsService {
client: Client,
}
impl ChatCompletionsService {
pub(crate) fn new(client: Client) -> Self {
Self { client }
}
pub async fn create(&self, mut body: ChatRequest) -> Result<ChatResponse> {
set_stream(&mut body, false);
self.client
.request_value(Method::POST, "/api/v1/chat/completions", None, Some(&body))
.await
}
pub async fn stream(&self, mut body: ChatRequest) -> Result<EventStream> {
set_stream(&mut body, true);
let bytes = self
.client
.stream_json(Method::POST, "/api/v1/chat/completions", Some(&body))
.await?;
Ok(events_from_bytes(bytes))
}
pub async fn stream_text(&self, body: ChatRequest) -> Result<TextStream> {
let stream = self.stream(body).await?;
Ok(Box::pin(stream.filter_map(|event| async move {
match event {
Ok(event) => {
let text = text_from_event(&event, StreamProtocol::OpenAi);
if text.is_empty() {
None
} else {
Some(Ok(text))
}
}
Err(error) => Some(Err(error)),
}
})))
}
}
#[derive(Clone, Debug)]
pub struct MessagesService {
client: Client,
}
impl MessagesService {
pub(crate) fn new(client: Client) -> Self {
Self { client }
}
pub async fn create(&self, mut body: ChatRequest) -> Result<ChatResponse> {
set_stream(&mut body, false);
self.client
.request_value(Method::POST, "/api/v1/messages", None, Some(&body))
.await
}
pub async fn stream(&self, mut body: ChatRequest) -> Result<EventStream> {
set_stream(&mut body, true);
let bytes = self
.client
.stream_json(Method::POST, "/api/v1/messages", Some(&body))
.await?;
Ok(events_from_bytes(bytes))
}
pub async fn stream_text(&self, body: ChatRequest) -> Result<TextStream> {
let stream = self.stream(body).await?;
Ok(Box::pin(stream.filter_map(|event| async move {
match event {
Ok(event) => {
let text = text_from_event(&event, StreamProtocol::Anthropic);
if text.is_empty() {
None
} else {
Some(Ok(text))
}
}
Err(error) => Some(Err(error)),
}
})))
}
}
#[derive(Clone, Debug)]
pub struct GeminiService {
client: Client,
}
impl GeminiService {
pub(crate) fn new(client: Client) -> Self {
Self { client }
}
pub async fn generate_content(&self, model: &str, body: ChatRequest) -> Result<ChatResponse> {
let path = format!(
"/api/v1beta/models/{}:generateContent",
encode_model_path(model)
);
self.client
.request_value(Method::POST, &path, None, Some(&body))
.await
}
pub async fn stream_generate_content(
&self,
model: &str,
body: ChatRequest,
) -> Result<EventStream> {
let path = format!(
"/api/v1beta/models/{}:streamGenerateContent",
encode_model_path(model)
);
let bytes = self
.client
.stream_json(Method::POST, &path, Some(&body))
.await?;
Ok(events_from_bytes(bytes))
}
pub async fn stream_text(&self, model: &str, body: ChatRequest) -> Result<TextStream> {
let stream = self.stream_generate_content(model, body).await?;
Ok(Box::pin(stream.filter_map(|event| async move {
match event {
Ok(event) => {
let text = text_from_event(&event, StreamProtocol::Gemini);
if text.is_empty() {
None
} else {
Some(Ok(text))
}
}
Err(error) => Some(Err(error)),
}
})))
}
}