1use rpc::LanguageModelProvider;
 2
 3use super::*;
 4
 5impl LlmDatabase {
 6    pub async fn find_or_create_usage(
 7        &self,
 8        user_id: i32,
 9        provider: LanguageModelProvider,
10        model_name: &str,
11    ) -> Result<usage::Model> {
12        self.transaction(|tx| async move {
13            let provider_name = match provider {
14                LanguageModelProvider::Anthropic => "anthropic",
15                LanguageModelProvider::OpenAi => "open_ai",
16                LanguageModelProvider::Google => "google",
17                LanguageModelProvider::Zed => "zed",
18            };
19
20            let model = model::Entity::find()
21                .inner_join(provider::Entity)
22                .filter(
23                    provider::Column::Name
24                        .eq(provider_name)
25                        .and(model::Column::Name.eq(model_name)),
26                )
27                .one(&*tx)
28                .await?
29                // TODO: Create the model, if one doesn't exist.
30                .ok_or_else(|| anyhow!("no model found for {provider_name}:{model_name}"))?;
31            let model_id = model.id;
32
33            let existing_usage = usage::Entity::find()
34                .filter(
35                    usage::Column::UserId
36                        .eq(user_id)
37                        .and(usage::Column::ModelId.eq(model_id)),
38                )
39                .one(&*tx)
40                .await?;
41            if let Some(usage) = existing_usage {
42                return Ok(usage);
43            }
44
45            let usage = usage::Entity::insert(usage::ActiveModel {
46                user_id: ActiveValue::set(user_id),
47                model_id: ActiveValue::set(model_id),
48                ..Default::default()
49            })
50            .exec_with_returning(&*tx)
51            .await?;
52
53            Ok(usage)
54        })
55        .await
56    }
57}