use std::time::Duration;
use super::{openai_wire, CompletionRequest, CompletionResponse, Provider};
use crate::error::{Error, Result};
pub use crate::net::REQUEST_TIMEOUT;
const ENDPOINT: &str = "https://api.openai.com/v1/chat/completions";
pub struct OpenAi {
client: reqwest::Client,
api_key: String,
model: String,
endpoint: String,
}
impl OpenAi {
pub fn new(api_key: impl Into<String>, model: impl Into<String>) -> Self {
Self {
client: crate::net::http_client(),
api_key: api_key.into(),
model: model.into(),
endpoint: ENDPOINT.to_string(),
}
}
pub fn with_timeout(mut self, timeout: Duration) -> Self {
self.client = crate::net::http_client_with_timeout(timeout);
self
}
#[cfg(test)]
pub(crate) fn at(endpoint: impl Into<String>, timeout: std::time::Duration) -> Self {
Self {
client: crate::net::http_client_with_timeout(timeout),
api_key: "test-key".into(),
model: "test-model".into(),
endpoint: endpoint.into(),
}
}
pub fn from_env() -> Result<Self> {
let api_key = std::env::var("OPENAI_API_KEY")
.map_err(|_| Error::Config("OPENAI_API_KEY is not set".into()))?;
let model = std::env::var("OPENAI_MODEL")
.map_err(|_| Error::Config("OPENAI_MODEL is not set".into()))?;
Ok(Self::new(api_key, model))
}
}
impl Provider for OpenAi {
fn name(&self) -> &str {
"openai"
}
fn endpoint(&self) -> Option<&str> {
Some(&self.endpoint)
}
#[cfg(feature = "media")]
fn accepts_images(&self) -> bool {
true
}
async fn complete(&self, request: CompletionRequest) -> Result<CompletionResponse> {
#[cfg(feature = "media")]
super::ensure_media_accepted(self.name(), self.accepts_images(), &request)?;
let resp = self
.client
.post(&self.endpoint)
.bearer_auth(&self.api_key)
.json(&openai_wire::body(&self.model, &request))
.send()
.await?;
openai_wire::parse_stream(super::ensure_success(resp).await?).await
}
}