aimo-client 0.1.7

AiMo Network REST API types and client
Documentation
use aimo_core::{
    provider::{Model, ProviderMetadata},
    receipt::RequestReceipt,
};
use serde::{Deserialize, Serialize};

// Common pagination and filtering structures
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct PaginationQuery {
    pub page: Option<String>,
    pub limit: Option<String>,
}

impl PaginationQuery {
    pub fn get_page(&self) -> u64 {
        self.page
            .as_ref()
            .and_then(|p| p.parse().ok())
            .unwrap_or(1)
            .max(1)
    }

    pub fn get_limit(&self) -> u64 {
        self.limit
            .as_ref()
            .and_then(|l| l.parse().ok())
            .unwrap_or(20)
            .min(100)
            .max(1)
    }
}

impl Default for PaginationQuery {
    fn default() -> Self {
        Self {
            page: Some("1".to_string()),
            limit: Some("20".to_string()),
        }
    }
}

#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct SortQuery {
    pub sort_by: Option<String>,
    pub sort_order: Option<SortOrder>,
}

#[derive(Debug, Clone, Serialize, Deserialize)]
#[serde(rename_all = "lowercase")]
pub enum SortOrder {
    Asc,
    Desc,
}

impl Default for SortOrder {
    fn default() -> Self {
        Self::Desc
    }
}

// User activity endpoint types
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct UserActivityQuery {
    #[serde(flatten)]
    pub pagination: PaginationQuery,
    #[serde(flatten)]
    pub sort: SortQuery,
    pub provider_id: Option<String>,
    pub model_name: Option<String>,
    pub from_timestamp: Option<String>,
    pub to_timestamp: Option<String>,
}

impl UserActivityQuery {
    pub fn get_from_timestamp(&self) -> Option<i64> {
        self.from_timestamp.as_ref().and_then(|t| t.parse().ok())
    }

    pub fn get_to_timestamp(&self) -> Option<i64> {
        self.to_timestamp.as_ref().and_then(|t| t.parse().ok())
    }
}

#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct UserActivityItem {
    pub timestamp: i64,
    pub provider_id: String,
    pub model_name: Option<String>,
    pub prompt_tokens: u64,
    pub completion_tokens: u64,
    pub spending_token: String,
    pub spending_amount: u64,
    pub finish_reason: u8,
    pub request_id: String,
}

#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct UserActivityResponse {
    pub items: Vec<UserActivityItem>,
    pub pagination: PaginationInfo,
}

#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct PaginationInfo {
    pub page: u64,
    pub limit: u64,
    pub total_items: u64,
    pub total_pages: u64,
}

// Provider models endpoint types
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct ProviderModelsQuery {
    #[serde(flatten)]
    pub pagination: PaginationQuery,
    #[serde(flatten)]
    pub sort: SortQuery,
    pub provider_id: Option<String>,
    pub provider_name: Option<String>,
    pub model_name: Option<String>,
    pub category: Option<String>,
}

#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct ProviderModelItem {
    pub provider_id: String,
    pub provider_name: String,
    pub category: String,
    pub model: Model,
}

#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct PriceInfo {
    pub token: String,
    pub input_price: u64,
    pub output_price: u64,
}

#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct ProviderModelsResponse {
    pub items: Vec<ProviderModelItem>,
    pub pagination: PaginationInfo,
}

// Provider metadata endpoint types
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct ProviderMetadataResponse {
    pub provider_id: String,
    pub provider_name: String,
    pub category: String,
    pub models: Vec<ModelInfo>,
}

#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct ModelInfo {
    pub name: String,
    pub pricing: Vec<PriceInfo>,
}

// User balance endpoint types
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct UserBalanceResponse {
    pub balance: Vec<UserBalanceItem>,
}

#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct UserBalanceItem {
    pub token: String,
    pub available: u64,
    pub locked: u64,
    pub total: u64,
}

// Helper function to convert RequestSettled to UserActivityItem
impl UserActivityItem {
    pub fn from_request_settled(
        receipt: RequestReceipt,
        _provider_name: String,
        model_name: Option<String>,
    ) -> Self {
        Self {
            timestamp: receipt.timestamp,
            provider_id: receipt.service_id.to_string(),
            model_name,
            prompt_tokens: receipt.prompt_tokens,
            completion_tokens: receipt.completion_tokens,
            spending_token: receipt.token_mint.to_string(),
            spending_amount: receipt.amount,
            finish_reason: receipt.finish_reason,
            request_id: receipt.request_id.to_string(),
        }
    }
}

// Helper function to convert ProviderMetadata to responses
impl ProviderModelItem {
    pub fn from_provider_metadata(provider_id: String, metadata: &ProviderMetadata) -> Vec<Self> {
        match metadata {
            ProviderMetadata::ModelProvider(model_provider) => {
                // Properly serialize the category to match the serde snake_case format
                let category = serde_json::to_value(&model_provider.category)
                    .ok()
                    .and_then(|v| v.as_str().map(|s| s.to_string()))
                    .unwrap_or_else(|| "unknown".to_string());

                model_provider
                    .models
                    .iter()
                    .map(|model| Self {
                        provider_id: provider_id.clone(),
                        provider_name: model_provider.name.clone(),
                        category: category.clone(),
                        model: model.clone(),
                    })
                    .collect()
            }
        }
    }
}

impl ProviderMetadataResponse {
    pub fn from_provider_metadata(provider_id: String, metadata: ProviderMetadata) -> Self {
        match metadata {
            ProviderMetadata::ModelProvider(model_provider) => {
                // Properly serialize the category to match the serde snake_case format
                let category = serde_json::to_value(&model_provider.category)
                    .ok()
                    .and_then(|v| v.as_str().map(|s| s.to_string()))
                    .unwrap_or_else(|| "unknown".to_string());

                Self {
                    provider_id,
                    provider_name: model_provider.name,
                    category,
                    models: model_provider
                        .models
                        .into_iter()
                        .map(|model| ModelInfo {
                            name: model.name,
                            pricing: model
                                .pricing
                                .into_iter()
                                .map(|p| PriceInfo {
                                    token: p.token,
                                    input_price: p.input_price,
                                    output_price: p.output_price,
                                })
                                .collect(),
                        })
                        .collect(),
                }
            }
        }
    }
}