use futures::Stream;
use serde_json::Value;
use std::collections::HashMap;
use std::pin::Pin;
use std::sync::Arc;
use crate::core::providers::base::{
GlobalPoolManager, HeaderPair, HttpMethod, apply_headers, header, header_owned,
};
use crate::core::providers::unified_provider::ProviderError;
use crate::core::traits::{
provider::ProviderConfig, provider::llm_provider::trait_definition::LLMProvider,
};
use crate::core::types::{
chat::ChatRequest,
context::RequestContext,
health::HealthStatus,
image::ImageGenerationRequest,
model::ModelInfo,
model::ProviderCapability,
responses::{ChatChunk, ChatResponse, ImageGenerationResponse},
};
use super::{
ReplicateClient, ReplicateConfig, ReplicateErrorMapper,
models::ReplicateModelType,
prediction::{CreatePredictionRequest, PredictionResponse, PredictionStatus},
};
use crate::core::traits::error_mapper::trait_def::ErrorMapper;
#[derive(Debug, Clone)]
pub struct ReplicateProvider {
config: ReplicateConfig,
pool_manager: Arc<GlobalPoolManager>,
supported_models: Vec<ModelInfo>,
}
impl ReplicateProvider {
pub fn new(config: ReplicateConfig) -> Result<Self, ProviderError> {
config
.validate()
.map_err(|e| ProviderError::configuration("replicate", e))?;
let pool_manager = Arc::new(
GlobalPoolManager::new()
.map_err(|e| ProviderError::configuration("replicate", e.to_string()))?,
);
let supported_models = ReplicateClient::supported_models();
Ok(Self {
config,
pool_manager,
supported_models,
})
}
pub async fn with_api_token(api_token: impl Into<String>) -> Result<Self, ProviderError> {
let config = ReplicateConfig::new(api_token);
Self::new(config)
}
pub fn from_env() -> Result<Self, ProviderError> {
let config = ReplicateConfig::from_env();
Self::new(config)
}
fn get_request_headers(&self) -> Vec<HeaderPair> {
let mut headers = Vec::with_capacity(self.config.base.headers.len() + 2);
for (key, value) in &self.config.base.headers {
if key.eq_ignore_ascii_case("authorization") || key.eq_ignore_ascii_case("content-type")
{
continue;
}
headers.push(header_owned(key.clone(), value.clone()));
}
if let Some(api_key) = &self.config.base.api_key {
headers.push(header("Authorization", format!("Token {}", api_key)));
}
headers.push(header("Content-Type", "application/json".to_string()));
headers
}
async fn create_prediction_and_wait(
&self,
model: &str,
input: Value,
stream: bool,
) -> Result<PredictionResponse, ProviderError> {
let version_hash = ReplicateConfig::prediction_version(model);
let prediction_request =
ReplicateClient::create_prediction_request(input, version_hash, stream);
let prediction_url = self.config.get_prediction_url(model);
let prediction = self
.submit_prediction(&prediction_url, &prediction_request)
.await?;
let polling_url = prediction
.get_prediction_url()
.ok_or_else(|| {
ProviderError::replicate_response_parsing("No polling URL in prediction response")
})?
.to_string();
self.poll_prediction(&polling_url).await
}
fn transport_attempts(&self) -> u32 {
self.config.base.max_retries.saturating_add(1)
}
fn build_request(
&self,
url: &str,
method: &HttpMethod,
headers: Vec<HeaderPair>,
body: Option<&Value>,
) -> reqwest::RequestBuilder {
let request_builder = match method {
HttpMethod::GET => self.pool_manager.client().get(url),
HttpMethod::POST => self.pool_manager.client().post(url),
HttpMethod::PUT => self.pool_manager.client().put(url),
HttpMethod::DELETE => self.pool_manager.client().delete(url),
}
.timeout(self.config.timeout());
let mut request_builder = apply_headers(request_builder, headers);
if let Some(body_data) = body {
request_builder = request_builder
.header("Content-Type", "application/json")
.json(body_data);
}
request_builder
}
async fn execute_request_with_config(
&self,
url: &str,
method: HttpMethod,
headers: Vec<HeaderPair>,
body: Option<Value>,
) -> Result<reqwest::Response, ProviderError> {
let attempts = self.transport_attempts();
let mut last_error = None;
for attempt in 0..attempts {
let request_builder = self.build_request(url, &method, headers.clone(), body.as_ref());
match request_builder.send().await {
Ok(response) => return Ok(response),
Err(err) => {
let provider_err = if err.is_timeout() {
ProviderError::timeout("replicate", err.to_string())
} else {
ProviderError::network("replicate", err.to_string())
};
if attempt + 1 == attempts {
return Err(provider_err);
}
last_error = Some(provider_err);
}
}
}
Err(last_error.unwrap_or_else(|| {
ProviderError::network("replicate", "request failed before execution")
}))
}
async fn submit_prediction(
&self,
url: &str,
request: &CreatePredictionRequest,
) -> Result<PredictionResponse, ProviderError> {
let headers = self.get_request_headers();
let body = serde_json::to_value(request)
.map_err(|e| ProviderError::serialization("replicate", e.to_string()))?;
let response = self
.execute_request_with_config(url, HttpMethod::POST, headers, Some(body))
.await?;
let status = response.status();
let response_bytes = response
.bytes()
.await
.map_err(|e| ProviderError::network("replicate", e.to_string()))?;
if !status.is_success() {
let error_text = String::from_utf8_lossy(&response_bytes);
return Err(ProviderError::replicate_api_error(
status.as_u16(),
error_text.to_string(),
));
}
serde_json::from_slice(&response_bytes)
.map_err(|e| ProviderError::replicate_response_parsing(e.to_string()))
}
async fn poll_prediction(&self, url: &str) -> Result<PredictionResponse, ProviderError> {
let headers = self.get_request_headers();
let polling_delay = std::time::Duration::from_secs(self.config.polling_delay_seconds);
for _ in 0..self.config.polling_retries {
tokio::time::sleep(polling_delay).await;
let response = self
.execute_request_with_config(url, HttpMethod::GET, headers.clone(), None)
.await?;
let status = response.status();
let response_bytes = response
.bytes()
.await
.map_err(|e| ProviderError::network("replicate", e.to_string()))?;
if !status.is_success() {
let error_text = String::from_utf8_lossy(&response_bytes);
return Err(ProviderError::replicate_api_error(
status.as_u16(),
error_text.to_string(),
));
}
let prediction: PredictionResponse = serde_json::from_slice(&response_bytes)
.map_err(|e| ProviderError::replicate_response_parsing(e.to_string()))?;
match prediction.status {
PredictionStatus::Succeeded => return Ok(prediction),
PredictionStatus::Failed => {
let error = prediction
.error
.clone()
.unwrap_or_else(|| "Prediction failed".to_string());
return Err(ProviderError::replicate_prediction_failed(error));
}
PredictionStatus::Canceled => {
return Err(ProviderError::replicate_prediction_canceled(
"Prediction was canceled",
));
}
_ => {
}
}
}
Err(ProviderError::replicate_prediction_timeout(
"Maximum retries exceeded waiting for prediction",
))
}
async fn execute_image_generation(
&self,
request: ImageGenerationRequest,
_context: RequestContext,
) -> Result<ImageGenerationResponse, ProviderError> {
let model = request.model.as_deref().unwrap_or("stability-ai/sdxl");
match ReplicateClient::get_model_type(model) {
ReplicateModelType::ImageGeneration => {}
ReplicateModelType::Other => {
return Err(ProviderError::model_not_found("replicate", model));
}
_ => {
return Err(ProviderError::invalid_request(
"replicate",
"Cannot use a text model for image generation",
));
}
}
let input = ReplicateClient::transform_image_request(&request, model)?;
let prediction = self.create_prediction_and_wait(model, input, false).await?;
ReplicateClient::transform_prediction_to_image_response(&prediction)
}
}
impl LLMProvider for ReplicateProvider {
fn name(&self) -> &'static str {
"replicate"
}
fn capabilities(&self) -> &'static [ProviderCapability] {
&[
ProviderCapability::ChatCompletion,
ProviderCapability::ChatCompletionStream,
ProviderCapability::ImageGeneration,
]
}
fn models(&self) -> &[ModelInfo] {
&self.supported_models
}
fn supports_model(&self, model: &str) -> bool {
super::models::get_replicate_registry().has_model(model)
}
fn get_supported_openai_params(&self, _model: &str) -> &'static [&'static str] {
ReplicateClient::supported_openai_params()
}
async fn map_openai_params(
&self,
params: HashMap<String, Value>,
_model: &str,
) -> Result<HashMap<String, Value>, ProviderError> {
let mut mapped = HashMap::new();
for (key, value) in params {
let mapped_key = match key.as_str() {
"max_tokens" => "max_new_tokens".to_string(),
"stop" => "stop_sequences".to_string(),
_ => key,
};
mapped.insert(mapped_key, value);
}
Ok(mapped)
}
async fn transform_request(
&self,
request: ChatRequest,
_context: RequestContext,
) -> Result<Value, ProviderError> {
Ok(ReplicateClient::transform_chat_request(&request))
}
async fn transform_response(
&self,
raw_response: &[u8],
model: &str,
_request_id: &str,
) -> Result<ChatResponse, ProviderError> {
let prediction: PredictionResponse = serde_json::from_slice(raw_response)
.map_err(|e| ProviderError::replicate_response_parsing(e.to_string()))?;
ReplicateClient::transform_prediction_to_chat_response(&prediction, model)
}
fn get_error_mapper(&self) -> Box<dyn ErrorMapper<ProviderError>> {
Box::new(ReplicateErrorMapper)
}
async fn chat_completion(
&self,
request: ChatRequest,
_context: RequestContext,
) -> Result<ChatResponse, ProviderError> {
let model = &request.model;
let model_type = ReplicateClient::get_model_type(model);
if model_type == ReplicateModelType::Other {
return Err(ProviderError::model_not_found("replicate", model));
}
if model_type != ReplicateModelType::TextGeneration {
return Err(ProviderError::invalid_request(
"replicate",
"Cannot use image model for chat completion",
));
}
let input = ReplicateClient::transform_chat_request(&request);
let prediction = self.create_prediction_and_wait(model, input, false).await?;
ReplicateClient::transform_prediction_to_chat_response(&prediction, model)
}
async fn chat_completion_stream(
&self,
request: ChatRequest,
_context: RequestContext,
) -> Result<Pin<Box<dyn Stream<Item = Result<ChatChunk, ProviderError>> + Send>>, ProviderError>
{
let model = &request.model;
let model_type = ReplicateClient::get_model_type(model);
if model_type == ReplicateModelType::Other {
return Err(ProviderError::model_not_found("replicate", model));
}
if model_type != ReplicateModelType::TextGeneration {
return Err(ProviderError::invalid_request(
"replicate",
"Cannot use image model for chat completion",
));
}
let input = ReplicateClient::transform_chat_request(&request);
let version_hash = ReplicateConfig::prediction_version(model);
let prediction_request = ReplicateClient::create_prediction_request(
input,
version_hash,
self.config.use_streaming,
);
let prediction_url = self.config.get_prediction_url(model);
let prediction = self
.submit_prediction(&prediction_url, &prediction_request)
.await?;
if self.config.use_streaming
&& let Some(stream_url) = prediction.get_stream_url()
{
let client = crate::core::http::outbound::streaming_outbound_client().clone();
let request_builder = apply_headers(client.get(stream_url), self.get_request_headers())
.header("Accept", "text/event-stream");
let response =
crate::core::providers::base::connection_pool::send_streaming_request_with_timeout(
request_builder,
self.config.timeout(),
)
.await
.map_err(|err| err.into_provider_error("replicate"))?;
if !response.status().is_success() {
let status_code = response.status().as_u16();
let error_text =
crate::core::providers::base::connection_pool::read_streaming_error_body(
response,
)
.await
.map_err(|err| err.into_provider_error("replicate"))?;
return Err(ProviderError::replicate_api_error(status_code, error_text));
}
let stream = response.bytes_stream();
Ok(Box::pin(super::streaming::create_replicate_stream(stream)))
} else {
let polling_url = prediction
.get_prediction_url()
.ok_or_else(|| {
ProviderError::replicate_response_parsing(
"No polling URL in prediction response",
)
})?
.to_string();
let final_prediction = self.poll_prediction(&polling_url).await?;
let response =
ReplicateClient::transform_prediction_to_chat_response(&final_prediction, model)?;
let content = response.choices.first().and_then(|c| {
c.message.content.as_ref().and_then(|mc| match mc {
crate::core::types::message::MessageContent::Text(s) => Some(s.clone()),
_ => None,
})
});
let chunk = ChatChunk {
id: response.id.clone(),
object: "chat.completion.chunk".to_string(),
created: response.created,
model: response.model.clone(),
system_fingerprint: None,
choices: vec![crate::core::types::responses::ChatStreamChoice {
index: 0,
delta: crate::core::types::responses::ChatDelta {
role: Some(crate::core::types::message::MessageRole::Assistant),
content,
thinking: None,
tool_calls: None,
function_call: None,
audio: None,
},
logprobs: None,
finish_reason: Some(crate::core::types::responses::FinishReason::Stop),
}],
usage: None,
};
Ok(Box::pin(futures::stream::once(async move { Ok(chunk) })))
}
}
async fn image_generation(
&self,
request: ImageGenerationRequest,
context: RequestContext,
) -> Result<ImageGenerationResponse, ProviderError> {
self.execute_image_generation(request, context).await
}
async fn health_check(&self) -> HealthStatus {
let url = format!("{}/models", self.config.get_api_base());
let headers = self.get_request_headers();
match self
.execute_request_with_config(&url, HttpMethod::GET, headers, None)
.await
{
Ok(response) if response.status().is_success() => HealthStatus::Healthy,
_ => HealthStatus::Unhealthy,
}
}
async fn calculate_cost(
&self,
model: &str,
input_tokens: u32,
output_tokens: u32,
) -> Result<f64, ProviderError> {
let spec = super::models::get_replicate_registry()
.get_model_spec(model)
.ok_or_else(|| ProviderError::model_not_found("replicate", model))?;
if spec.model_type != ReplicateModelType::TextGeneration {
return Err(ProviderError::not_supported(
"replicate",
"Image cost requires output count, dimensions or hardware runtime; token counts cannot determine it",
));
}
let (Some(input_price), Some(output_price)) = (
spec.model_info.input_cost_per_1k_tokens,
spec.model_info.output_cost_per_1k_tokens,
) else {
return Err(ProviderError::not_supported(
"replicate",
"Verified token pricing is unavailable",
));
};
Ok(input_price * f64::from(input_tokens) / 1000.0
+ output_price * f64::from(output_tokens) / 1000.0)
}
}
#[cfg(test)]
mod tests {
use super::*;
fn image_request(model: &str) -> ImageGenerationRequest {
ImageGenerationRequest {
model: Some(model.into()),
prompt: "a bird".into(),
n: None,
size: None,
quality: None,
response_format: None,
style: None,
user: None,
}
}
#[tokio::test]
async fn audited_models_preserve_modality_and_reject_unknown_prices() {
let provider = ReplicateProvider::new(ReplicateConfig::new("fixture-token")).unwrap();
let wrapped = crate::core::providers::Provider::Replicate(provider.clone());
for model in [
"meta/meta-llama-3.1-405b-instruct",
"mistralai/mistral-7b-instruct-v0.2",
"mistralai/mixtral-8x7b-instruct-v0.1",
"lucataco/playground-v2.5-1024px-aesthetic",
"unknown/model",
] {
assert!(
!wrapped.supports_capability_for_model(model, &ProviderCapability::ChatCompletion)
);
assert!(
!wrapped.supports_capability_for_model(model, &ProviderCapability::ImageGeneration)
);
assert!(matches!(
provider.calculate_cost(model, 1, 1).await,
Err(ProviderError::ModelNotFound { .. })
));
assert!(matches!(
provider
.chat_completion(
ChatRequest {
model: model.into(),
..Default::default()
},
RequestContext::default()
)
.await,
Err(ProviderError::ModelNotFound { .. })
));
}
for model in provider.models() {
if model
.capabilities
.contains(&ProviderCapability::ImageGeneration)
{
assert!(
!wrapped.supports_capability_for_model(
&model.id,
&ProviderCapability::ChatCompletion
)
);
assert!(matches!(
provider.calculate_cost(&model.id, 0, 0).await,
Err(ProviderError::NotSupported { .. })
));
} else {
assert!(wrapped.supports_capability_for_model(
&model.id,
&ProviderCapability::ChatCompletionStream
));
assert!(!wrapped.supports_capability_for_model(
&model.id,
&ProviderCapability::ImageGeneration
));
}
}
let versioned = format!("meta/llama-2-70b-chat:{}", "a".repeat(64));
assert!(provider.supports_model(&versioned));
assert!(
wrapped.supports_capability_for_model(&versioned, &ProviderCapability::ChatCompletion)
);
let cost = provider
.calculate_cost(&versioned, 1000, 1000)
.await
.unwrap();
assert!((cost - 0.0034).abs() < 1e-12);
assert!(
provider
.image_generation(
image_request("meta/llama-2-70b-chat"),
RequestContext::default()
)
.await
.is_err()
);
}
#[tokio::test]
async fn http_community_version_and_official_prediction_paths() {
use tokio::{
io::{AsyncReadExt, AsyncWriteExt},
net::TcpListener,
};
let listener = TcpListener::bind(("127.0.0.1", 0)).await.unwrap();
let address = listener.local_addr().unwrap();
let server = tokio::spawn(async move {
let mut posts = Vec::new();
for _ in 0..6 {
let (mut socket, _) = listener.accept().await.unwrap();
let mut bytes = Vec::new();
let mut chunk = [0; 4096];
let (header_end, length) = loop {
let read = socket.read(&mut chunk).await.unwrap();
assert!(read > 0);
bytes.extend_from_slice(&chunk[..read]);
if let Some(end) = bytes.windows(4).position(|w| w == b"\r\n\r\n") {
let headers = String::from_utf8_lossy(&bytes[..end]);
let length: usize = headers
.lines()
.find_map(|l| {
l.to_ascii_lowercase()
.strip_prefix("content-length:")
.map(|v| v.trim().parse().unwrap())
})
.unwrap_or(0);
break (end + 4, length);
}
};
while bytes.len() < header_end + length {
let read = socket.read(&mut chunk).await.unwrap();
assert!(read > 0);
bytes.extend_from_slice(&chunk[..read]);
}
let headers = String::from_utf8_lossy(&bytes[..header_end]);
let body = if headers.starts_with("POST ") {
posts.push((headers.lines().next().unwrap().to_owned(), serde_json::from_slice::<Value>(&bytes[header_end..header_end + length]).unwrap()));
serde_json::json!({"id":"pred-test", "status":"processing", "urls":{"get":format!("http://{address}/v1/predictions/pred-test")}})
} else {
serde_json::json!({"id":"pred-test", "status":"succeeded", "output":["https://example.com/image.png"]})
}.to_string();
socket.write_all(format!("HTTP/1.1 200 OK\r\nContent-Type: application/json\r\nContent-Length: {}\r\nConnection: close\r\n\r\n{body}",body.len()).as_bytes()).await.unwrap();
}
posts
});
let mut config = ReplicateConfig::new("fixture-token");
config.base.api_base = Some(format!("http://{address}/v1"));
config.polling_delay_seconds = 1;
let provider = ReplicateProvider::new(config).unwrap();
let explicit_version = format!("black-forest-labs/flux-2-pro:{}", "b".repeat(64));
for model in [
"stability-ai/sdxl",
"black-forest-labs/flux-2-pro",
explicit_version.as_str(),
] {
let mut request = image_request(model);
request.size = Some("1024x768".into());
let response = provider
.image_generation(request, RequestContext::default())
.await
.unwrap();
assert_eq!(response.data.len(), 1);
}
let posts = server.await.unwrap();
assert_eq!(posts[0].0, "POST /v1/predictions HTTP/1.1");
assert_eq!(
posts[0].1["version"],
"7762fd07cf82c948538e41f63f77d685e02b063e37e496e96eefd46c929f9bdc"
);
assert_eq!(
posts[1].0,
"POST /v1/models/black-forest-labs/flux-2-pro/predictions HTTP/1.1"
);
assert!(posts[1].1.get("version").is_none());
assert_eq!(posts[1].1["input"]["aspect_ratio"], "custom");
assert_eq!(posts[1].1["input"]["width"], 1024);
assert!(posts[1].1["input"].get("num_outputs").is_none());
assert_eq!(posts[2].0, "POST /v1/predictions HTTP/1.1");
assert_eq!(posts[2].1["version"], "b".repeat(64));
assert_eq!(posts[2].1["input"]["aspect_ratio"], "custom");
assert!(posts[2].1["input"].get("num_outputs").is_none());
}
#[test]
fn test_provider_creation_without_api_key() {
let config = ReplicateConfig::default();
let result = ReplicateProvider::new(config);
assert!(result.is_err());
}
#[test]
fn test_provider_creation_with_api_key() {
let config = ReplicateConfig::new("test-token");
let result = ReplicateProvider::new(config);
assert!(result.is_ok());
}
#[test]
fn test_provider_name() {
let config = ReplicateConfig::new("test-token");
let provider = ReplicateProvider::new(config).unwrap();
assert_eq!(provider.name(), "replicate");
}
#[test]
fn test_provider_capabilities() {
let config = ReplicateConfig::new("test-token");
let provider = ReplicateProvider::new(config).unwrap();
let capabilities = provider.capabilities();
assert!(capabilities.contains(&ProviderCapability::ChatCompletion));
assert!(capabilities.contains(&ProviderCapability::ChatCompletionStream));
assert!(capabilities.contains(&ProviderCapability::ImageGeneration));
}
#[test]
fn test_provider_models() {
let config = ReplicateConfig::new("test-token");
let provider = ReplicateProvider::new(config).unwrap();
let models = provider.models();
assert!(!models.is_empty());
assert!(models.iter().any(|m| m.id.contains("llama")));
assert!(models.iter().any(|m| m.id.contains("sdxl")));
}
#[test]
fn test_supported_openai_params() {
let config = ReplicateConfig::new("test-token");
let provider = ReplicateProvider::new(config).unwrap();
let params = provider.get_supported_openai_params("any-model");
assert!(params.contains(&"temperature"));
assert!(params.contains(&"max_tokens"));
}
#[tokio::test]
async fn test_map_openai_params() {
let config = ReplicateConfig::new("test-token");
let provider = ReplicateProvider::new(config).unwrap();
let mut params = HashMap::new();
params.insert("max_tokens".to_string(), serde_json::json!(100));
params.insert("temperature".to_string(), serde_json::json!(0.7));
let mapped = provider.map_openai_params(params, "model").await.unwrap();
assert!(mapped.contains_key("max_new_tokens"));
assert!(mapped.contains_key("temperature"));
assert!(!mapped.contains_key("max_tokens"));
}
#[test]
fn test_get_request_headers() {
let config = ReplicateConfig::new("test-token");
let provider = ReplicateProvider::new(config).unwrap();
let headers = provider.get_request_headers();
assert!(headers.iter().any(|h| h.0 == "Authorization"));
assert!(headers.iter().any(|h| h.0 == "Content-Type"));
}
#[test]
fn test_get_request_headers_includes_custom_headers() {
let mut config = ReplicateConfig::new("test-token");
config
.base
.headers
.insert("x-replicate-test".to_string(), "ok".to_string());
let provider = match ReplicateProvider::new(config) {
Ok(provider) => provider,
Err(err) => panic!("replicate provider should accept test config: {err}"),
};
let headers = provider.get_request_headers();
assert!(
headers
.iter()
.any(|h| h.0.as_ref() == "x-replicate-test" && h.1.as_ref() == "ok")
);
}
#[test]
fn test_transport_attempts_include_initial_request() {
let mut config = ReplicateConfig::new("test-token");
config.base.max_retries = 4;
let provider = match ReplicateProvider::new(config) {
Ok(provider) => provider,
Err(err) => panic!("replicate provider should accept test config: {err}"),
};
assert_eq!(provider.transport_attempts(), 5);
}
#[test]
fn test_streaming_prediction_flag_follows_config() {
let mut config = ReplicateConfig::new("test-token");
config.use_streaming = false;
let provider = match ReplicateProvider::new(config) {
Ok(provider) => provider,
Err(err) => panic!("replicate provider should accept test config: {err}"),
};
let input = serde_json::json!({"prompt": "hello"});
let version_hash = ReplicateConfig::extract_version_hash("meta/llama-2-70b-chat");
let request = ReplicateClient::create_prediction_request(
input.clone(),
version_hash.clone(),
provider.config.use_streaming,
);
assert_eq!(request.stream, None);
let mut config = ReplicateConfig::new("test-token");
config.use_streaming = true;
let provider = match ReplicateProvider::new(config) {
Ok(provider) => provider,
Err(err) => panic!("replicate provider should accept test config: {err}"),
};
let request = ReplicateClient::create_prediction_request(
input,
version_hash,
provider.config.use_streaming,
);
assert_eq!(request.stream, Some(true));
}
#[test]
fn test_from_env_missing_token() {
unsafe {
std::env::remove_var("REPLICATE_API_TOKEN");
}
let result = ReplicateProvider::from_env();
assert!(result.is_err());
}
#[test]
fn test_error_mapper() {
let config = ReplicateConfig::new("test-token");
let provider = ReplicateProvider::new(config).unwrap();
let _mapper = provider.get_error_mapper();
}
#[tokio::test]
async fn test_transform_request() {
use crate::core::types::{
chat::ChatMessage, message::MessageContent, message::MessageRole,
};
let config = ReplicateConfig::new("test-token");
let provider = ReplicateProvider::new(config).unwrap();
let request = ChatRequest {
model: "meta/llama-2-70b-chat".to_string(),
messages: vec![ChatMessage {
role: MessageRole::User,
content: Some(MessageContent::Text("Hello".to_string())),
thinking: None,
audio: None,
name: None,
tool_calls: None,
tool_call_id: None,
function_call: None,
}],
temperature: Some(0.7),
max_tokens: Some(100),
top_p: None,
frequency_penalty: None,
presence_penalty: None,
stream: false,
stream_options: None,
tools: None,
tool_choice: None,
user: None,
response_format: None,
seed: None,
max_completion_tokens: None,
stop: None,
parallel_tool_calls: None,
n: None,
logit_bias: None,
functions: None,
function_call: None,
logprobs: None,
top_logprobs: None,
reasoning_effort: None,
store: None,
metadata: None,
service_tier: None,
thinking: None,
extra_params: HashMap::new(),
};
let context = RequestContext::default();
let result = provider.transform_request(request, context).await;
assert!(result.is_ok());
let value = result.unwrap();
assert!(value.get("prompt").is_some());
}
#[tokio::test]
async fn test_calculate_cost() {
let config = ReplicateConfig::new("test-token");
let provider = ReplicateProvider::new(config).unwrap();
let cost = provider
.calculate_cost("meta/llama-2-70b-chat", 100, 50)
.await
.unwrap();
assert!(
provider
.calculate_cost("black-forest-labs/flux-2-pro", 0, 0)
.await
.is_err()
);
assert!(cost >= 0.0);
}
#[tokio::test]
async fn test_calculate_cost_unknown_model() {
let config = ReplicateConfig::new("test-token");
let provider = ReplicateProvider::new(config).unwrap();
let cost = provider.calculate_cost("unknown/model", 100, 50).await;
assert!(matches!(cost, Err(ProviderError::ModelNotFound { .. })));
}
}