use super::{BaseProvider, ModelPricing, Provider, ProviderError, ProviderType};
use crate::config::ProviderConfig;
use crate::core::models::{RequestContext, openai::*};
use crate::utils::error::Result;
use async_trait::async_trait;
use serde_json::json;
use std::collections::HashMap;
use tracing::{debug, info};
#[derive(Debug, Clone)]
pub struct AnyscaleProvider {
base: BaseProvider,
pricing_cache: HashMap<String, ModelPricing>,
}
impl AnyscaleProvider {
pub async fn new(config: &ProviderConfig) -> Result<Self> {
let base = BaseProvider::new(config)?;
let base_url = config
.base_url
.clone()
.unwrap_or_else(|| "https://api.endpoints.anyscale.com".to_string());
let provider = Self {
base: BaseProvider { base_url, ..base },
pricing_cache: Self::initialize_pricing_cache(),
};
info!(
"Anyscale provider '{}' initialized successfully",
config.name
);
Ok(provider)
}
fn initialize_pricing_cache() -> HashMap<String, ModelPricing> {
let mut cache = HashMap::new();
cache.insert(
"meta-llama/Llama-2-7b-chat-hf".to_string(),
ModelPricing {
model: "meta-llama/Llama-2-7b-chat-hf".to_string(),
input_cost_per_1k: 0.00015,
output_cost_per_1k: 0.00015,
currency: "USD".to_string(),
updated_at: chrono::Utc::now(),
},
);
cache.insert(
"meta-llama/Llama-2-13b-chat-hf".to_string(),
ModelPricing {
model: "meta-llama/Llama-2-13b-chat-hf".to_string(),
input_cost_per_1k: 0.00025,
output_cost_per_1k: 0.00025,
currency: "USD".to_string(),
updated_at: chrono::Utc::now(),
},
);
cache.insert(
"meta-llama/Llama-2-70b-chat-hf".to_string(),
ModelPricing {
model: "meta-llama/Llama-2-70b-chat-hf".to_string(),
input_cost_per_1k: 0.001,
output_cost_per_1k: 0.001,
currency: "USD".to_string(),
updated_at: chrono::Utc::now(),
},
);
cache.insert(
"mistralai/Mixtral-8x7B-Instruct-v0.1".to_string(),
ModelPricing {
model: "mistralai/Mixtral-8x7B-Instruct-v0.1".to_string(),
input_cost_per_1k: 0.0005,
output_cost_per_1k: 0.0005,
currency: "USD".to_string(),
updated_at: chrono::Utc::now(),
},
);
cache.insert(
"codellama/CodeLlama-34b-Instruct-hf".to_string(),
ModelPricing {
model: "codellama/CodeLlama-34b-Instruct-hf".to_string(),
input_cost_per_1k: 0.001,
output_cost_per_1k: 0.001,
currency: "USD".to_string(),
updated_at: chrono::Utc::now(),
},
);
cache
}
}
#[async_trait]
impl Provider for AnyscaleProvider {
fn name(&self) -> &str {
&self.base.name
}
fn provider_type(&self) -> ProviderType {
ProviderType::Custom("anyscale".to_string())
}
async fn supports_model(&self, model: &str) -> bool {
self.base.is_model_supported(model)
|| model.contains("llama")
|| model.contains("mixtral")
|| model.contains("codellama")
}
async fn supports_images(&self) -> bool {
false
}
async fn supports_embeddings(&self) -> bool {
false
}
async fn supports_streaming(&self) -> bool {
true
}
async fn list_models(&self) -> Result<Vec<Model>> {
let url = format!("{}/v1/models", self.base.base_url);
let response = self
.base
.client
.get(&url)
.header("Authorization", format!("Bearer {}", self.base.api_key))
.send()
.await
.map_err(|e| ProviderError::Network(e.to_string()))?;
if !response.status().is_success() {
let known_models = vec![
"meta-llama/Llama-2-7b-chat-hf",
"meta-llama/Llama-2-13b-chat-hf",
"meta-llama/Llama-2-70b-chat-hf",
"mistralai/Mixtral-8x7B-Instruct-v0.1",
"codellama/CodeLlama-34b-Instruct-hf",
];
let models = known_models
.into_iter()
.map(|model| Model {
id: model.to_string(),
object: "model".to_string(),
created: chrono::Utc::now().timestamp() as u64,
owned_by: "anyscale".to_string(),
})
.collect();
return Ok(models);
}
let models_response: serde_json::Value = self.base.parse_json_response(response).await?;
let models = models_response
.get("data")
.and_then(|d| d.as_array())
.unwrap_or(&vec![])
.iter()
.filter_map(|model| {
Some(Model {
id: model.get("id")?.as_str()?.to_string(),
object: model.get("object")?.as_str()?.to_string(),
created: model.get("created")?.as_u64()?,
owned_by: model.get("owned_by")?.as_str()?.to_string(),
})
})
.collect();
Ok(models)
}
async fn health_check(&self) -> Result<()> {
debug!("Performing Anyscale health check");
let url = format!("{}/v1/models", self.base.base_url);
let response = self
.base
.client
.get(&url)
.header("Authorization", format!("Bearer {}", self.base.api_key))
.send()
.await
.map_err(|e| ProviderError::Network(e.to_string()))?;
if response.status().is_success() {
Ok(())
} else {
Err(
ProviderError::Unknown(format!("Health check failed: {}", response.status()))
.into(),
)
}
}
async fn chat_completion(
&self,
request: ChatCompletionRequest,
_context: RequestContext,
) -> Result<ChatCompletionResponse> {
debug!("Anyscale chat completion for model: {}", request.model);
let mut body = json!({
"model": request.model,
"messages": request.messages
});
if let Some(max_tokens) = request.max_tokens {
body["max_tokens"] = json!(max_tokens);
}
if let Some(temperature) = request.temperature {
body["temperature"] = json!(temperature);
}
if let Some(top_p) = request.top_p {
body["top_p"] = json!(top_p);
}
if let Some(stream) = request.stream {
body["stream"] = json!(stream);
}
let url = format!("{}/v1/chat/completions", self.base.base_url);
let response = self
.base
.client
.post(&url)
.header("Authorization", format!("Bearer {}", self.base.api_key))
.header("Content-Type", "application/json")
.json(&body)
.send()
.await
.map_err(|e| ProviderError::Network(e.to_string()))?;
if !response.status().is_success() {
let status = response.status();
let error_text = response.text().await.unwrap_or_default();
return Err(match status.as_u16() {
401 => ProviderError::Authentication(error_text),
429 => ProviderError::RateLimit(error_text),
404 => ProviderError::ModelNotFound(error_text),
400 => ProviderError::InvalidRequest(error_text),
_ => ProviderError::Unknown(format!("HTTP {}: {}", status, error_text)),
}
.into());
}
let response_json: ChatCompletionResponse = self.base.parse_json_response(response).await?;
Ok(response_json)
}
async fn completion(
&self,
request: CompletionRequest,
_context: RequestContext,
) -> Result<CompletionResponse> {
let chat_request = ChatCompletionRequest {
model: request.model.clone(),
messages: vec![ChatMessage {
role: MessageRole::User,
content: Some(MessageContent::Text(request.prompt)),
name: None,
function_call: None,
tool_calls: None,
tool_call_id: None,
audio: None,
}],
max_tokens: request.max_tokens,
max_completion_tokens: None,
temperature: request.temperature.map(|t| t as f32),
top_p: request.top_p.map(|t| t as f32),
n: request.n,
stream: request.stream,
stream_options: None,
stop: request.stop,
presence_penalty: request.presence_penalty.map(|p| p as f32),
frequency_penalty: request.frequency_penalty.map(|f| f as f32),
logit_bias: request
.logit_bias
.map(|lb| lb.into_iter().map(|(k, v)| (k, v as f32)).collect()),
user: request.user,
functions: None,
function_call: None,
tools: None,
tool_choice: None,
response_format: None,
seed: None,
logprobs: None,
top_logprobs: None,
modalities: None,
audio: None,
};
let chat_response = self.chat_completion(chat_request, _context).await?;
let text = match &chat_response.choices.first().unwrap().message.content {
Some(MessageContent::Text(text)) => text.clone(),
Some(MessageContent::Parts(parts)) => parts
.iter()
.filter_map(|part| match part {
ContentPart::Text { text } => Some(text.clone()),
_ => None,
})
.collect::<Vec<String>>()
.join(" "),
None => String::new(),
};
Ok(CompletionResponse {
id: chat_response.id.replace("chatcmpl", "cmpl"),
object: "text_completion".to_string(),
created: chat_response.created,
model: request.model,
choices: vec![CompletionChoice {
text,
index: 0,
logprobs: None,
finish_reason: chat_response.choices.first().unwrap().finish_reason.clone(),
}],
usage: chat_response.usage,
})
}
async fn embedding(
&self,
_request: EmbeddingRequest,
_context: RequestContext,
) -> Result<EmbeddingResponse> {
Err(
ProviderError::InvalidRequest("Embeddings not supported by Anyscale".to_string())
.into(),
)
}
async fn image_generation(
&self,
_request: ImageGenerationRequest,
_context: RequestContext,
) -> Result<ImageGenerationResponse> {
Err(
ProviderError::InvalidRequest("Image generation not supported by Anyscale".to_string())
.into(),
)
}
async fn get_model_pricing(&self, model: &str) -> Result<ModelPricing> {
if let Some(pricing) = self.pricing_cache.get(model) {
Ok(pricing.clone())
} else {
Ok(ModelPricing {
model: model.to_string(),
input_cost_per_1k: 0.0005,
output_cost_per_1k: 0.0005,
currency: "USD".to_string(),
updated_at: chrono::Utc::now(),
})
}
}
async fn calculate_cost(
&self,
model: &str,
input_tokens: u32,
output_tokens: u32,
) -> Result<f64> {
let pricing = self.get_model_pricing(model).await?;
let input_cost = (input_tokens as f64 / 1000.0) * pricing.input_cost_per_1k;
let output_cost = (output_tokens as f64 / 1000.0) * pricing.output_cost_per_1k;
Ok(input_cost + output_cost)
}
}