Richard Feldman/auto (#16696)91ffa022 years agoHistory
134 lines · 5.1 KB · Rust
RawDownload
use super::*;use sea_orm::{sea_query::OnConflict, QueryOrder};use std::str::FromStr;use strum::IntoEnumIterator as _; pub struct ModelParams {    pub provider: LanguageModelProvider,    pub name: String,    pub max_requests_per_minute: i64,    pub max_tokens_per_minute: i64,    pub max_tokens_per_day: i64,    pub price_per_million_input_tokens: i32,    pub price_per_million_output_tokens: i32,} impl LlmDatabase {    pub async fn initialize_providers(&mut self) -> Result<()> {        self.provider_ids = self            .transaction(|tx| async move {                let existing_providers = provider::Entity::find().all(&*tx).await?;                 let mut new_providers = LanguageModelProvider::iter()                    .filter(|provider| {                        !existing_providers                            .iter()                            .any(|p| p.name == provider.to_string())                    })                    .map(|provider| provider::ActiveModel {                        name: ActiveValue::set(provider.to_string()),                        ..Default::default()                    })                    .peekable();                 if new_providers.peek().is_some() {                    provider::Entity::insert_many(new_providers)                        .exec(&*tx)                        .await?;                }                 let all_providers: HashMap<_, _> = provider::Entity::find()                    .all(&*tx)                    .await?                    .iter()                    .filter_map(|provider| {                        LanguageModelProvider::from_str(&provider.name)                            .ok()                            .map(|p| (p, provider.id))                    })                    .collect();                 Ok(all_providers)            })            .await?;        Ok(())    }     pub async fn initialize_models(&mut self) -> Result<()> {        let all_provider_ids = &self.provider_ids;        self.models = self            .transaction(|tx| async move {                let all_models: HashMap<_, _> = model::Entity::find()                    .all(&*tx)                    .await?                    .into_iter()                    .filter_map(|model| {                        let provider = all_provider_ids.iter().find_map(|(provider, id)| {                            if *id == model.provider_id {                                Some(provider)                            } else {                                None                            }                        })?;                        Some(((*provider, model.name.clone()), model))                    })                    .collect();                Ok(all_models)            })            .await?;        Ok(())    }     pub async fn insert_models(&mut self, models: &[ModelParams]) -> Result<()> {        let all_provider_ids = &self.provider_ids;        self.transaction(|tx| async move {            model::Entity::insert_many(models.iter().map(|model_params| {                let provider_id = all_provider_ids[&model_params.provider];                model::ActiveModel {                    provider_id: ActiveValue::set(provider_id),                    name: ActiveValue::set(model_params.name.clone()),                    max_requests_per_minute: ActiveValue::set(model_params.max_requests_per_minute),                    max_tokens_per_minute: ActiveValue::set(model_params.max_tokens_per_minute),                    max_tokens_per_day: ActiveValue::set(model_params.max_tokens_per_day),                    price_per_million_input_tokens: ActiveValue::set(                        model_params.price_per_million_input_tokens,                    ),                    price_per_million_output_tokens: ActiveValue::set(                        model_params.price_per_million_output_tokens,                    ),                    ..Default::default()                }            }))            .on_conflict(                OnConflict::columns([model::Column::ProviderId, model::Column::Name])                    .update_columns([                        model::Column::MaxRequestsPerMinute,                        model::Column::MaxTokensPerMinute,                        model::Column::MaxTokensPerDay,                        model::Column::PricePerMillionInputTokens,                        model::Column::PricePerMillionOutputTokens,                    ])                    .to_owned(),            )            .exec_without_returning(&*tx)            .await?;            Ok(())        })        .await?;        self.initialize_models().await    }     /// Returns the list of LLM providers.    pub async fn list_providers(&self) -> Result<Vec<LanguageModelProvider>> {        self.transaction(|tx| async move {            Ok(provider::Entity::find()                .order_by_asc(provider::Column::Name)                .all(&*tx)                .await?                .into_iter()                .filter_map(|p| LanguageModelProvider::from_str(&p.name).ok())                .collect())        })        .await    }}