use std::pin::Pin;
use futures::Stream;
use tokio_stream::StreamExt;
use crate::v1::api::Ollama;
use crate::v1::error::APIError;
use crate::v1::resources::ollama::completion::{OllamaCompletionRequest, OllamaCompletionResponse};
impl Ollama {
pub async fn create(
&self,
ollama_params: OllamaCompletionRequest,
) -> Result<OllamaCompletionResponse, APIError> {
let endpoint = "/api/generate";
let response = self.post(endpoint, &ollama_params).await?;
println!("{}", response);
let completion_response: Result<OllamaCompletionResponse, _> =
serde_json::from_str(&response);
match completion_response {
Ok(data) => Ok(data),
Err(_) => Err(APIError::ParseError("Invalid JSON response".to_string())),
}
}
pub async fn create_stream(
&self,
ollama_params: OllamaCompletionRequest,
) -> Result<
Pin<Box<dyn Stream<Item = Result<OllamaCompletionResponse, APIError>> + Send>>,
APIError,
> {
let endpoint = "/api/generate";
let res = self
.post_stream(endpoint, &ollama_params.to_stream_request())
.await;
let stream = res.map(|res| match res {
Ok(res) => match serde_json::from_slice::<OllamaCompletionResponse>(&res) {
Ok(completion_response) => Ok(completion_response),
Err(_) => Err(APIError::ParseError("Invalid JSON response".to_string())),
},
Err(e) => Err(APIError::StreamError(e.to_string())),
});
Ok(Box::pin(stream))
}
}