use std::pin::Pin;
use anyhow::Result;
use futures::{Stream, StreamExt};
use tokio_util::bytes::Bytes;
use crate::v1::api::Modelfarm;
use crate::v1::error::APIError;
use crate::v1::resources::modelfarm::chat_completion::{
ModelfarmChatCompletionRequest, ModelfarmChatCompletionResponse,
};
impl Modelfarm {
pub async fn chat(
&self,
modelfarm_params: ModelfarmChatCompletionRequest,
) -> Result<ModelfarmChatCompletionResponse, APIError> {
let endpoint = "/v1beta/chat";
println!("Sending request to {}", endpoint);
let response = self.post(endpoint, &modelfarm_params).await?;
let chat_completion_response: ModelfarmChatCompletionResponse =
serde_json::from_str(&response).unwrap();
Ok(chat_completion_response)
}
pub async fn stream_chat(
&self,
modelfarm_params: ModelfarmChatCompletionRequest,
) -> Result<
Pin<Box<dyn Stream<Item = Result<ModelfarmChatCompletionResponse, anyhow::Error>> + Send>>,
> {
let endpoint = "/v1beta/chat_streaming";
let res: Pin<Box<dyn Stream<Item = Result<Bytes, anyhow::Error>> + Send>> =
self.post_stream(endpoint, &modelfarm_params).await;
let stream = res.map(|res| {
let res = res.map_err(anyhow::Error::from)?;
let chat_response: ModelfarmChatCompletionResponse = serde_json::from_slice(&res)?;
Ok(chat_response)
});
Ok(Box::pin(stream))
}
}