use anyhow::{Context, Result};
use crabrace::{CrabraceClient, Provider};
use serde::{Deserialize, Serialize};
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct CrabraceConfig {
#[serde(default = "default_enabled")]
pub enabled: bool,
#[serde(default = "default_base_url")]
pub base_url: String,
#[serde(default = "default_auto_update")]
pub auto_update: bool,
#[serde(default = "default_update_interval")]
pub update_interval_seconds: u64,
}
fn default_enabled() -> bool {
true
}
fn default_base_url() -> String {
"http://localhost:8080".to_string()
}
fn default_auto_update() -> bool {
true
}
fn default_update_interval() -> u64 {
3600 }
impl Default for CrabraceConfig {
fn default() -> Self {
Self {
enabled: default_enabled(),
base_url: default_base_url(),
auto_update: default_auto_update(),
update_interval_seconds: default_update_interval(),
}
}
}
pub struct CrabraceIntegration {
client: CrabraceClient,
config: CrabraceConfig,
}
impl CrabraceIntegration {
pub fn new(config: CrabraceConfig) -> Result<Self> {
let client = CrabraceClient::new(&config.base_url);
Ok(Self { client, config })
}
pub async fn fetch_providers(&self) -> Result<Vec<Provider>> {
self.client
.get_providers()
.await
.context("Failed to fetch providers from Crabrace")
}
pub async fn health_check(&self) -> Result<bool> {
self.client
.health_check()
.await
.context("Failed to check Crabrace health")
}
pub async fn get_provider(&self, provider_id: &str) -> Result<Option<Provider>> {
let providers = self.fetch_providers().await?;
Ok(providers.into_iter().find(|p| p.id == provider_id))
}
pub async fn get_all_model_ids(&self) -> Result<Vec<String>> {
let providers = self.fetch_providers().await?;
let model_ids: Vec<String> = providers
.into_iter()
.flat_map(|p| p.models.into_iter().map(|m| m.id))
.collect();
Ok(model_ids)
}
pub async fn is_provider_available(&self, provider_id: &str) -> Result<bool> {
Ok(self.get_provider(provider_id).await?.is_some())
}
pub fn config(&self) -> &CrabraceConfig {
&self.config
}
}