1use crate::llm::db::{ModelId, UsageId, UsageMeasureId};
 2use sea_orm::entity::prelude::*;
 3
 4/// An LLM usage record.
 5#[derive(Clone, Debug, PartialEq, DeriveEntityModel)]
 6#[sea_orm(table_name = "usages")]
 7pub struct Model {
 8    #[sea_orm(primary_key)]
 9    pub id: UsageId,
10    /// The ID of the Zed user.
11    ///
12    /// Corresponds to the `users` table in the primary collab database.
13    pub user_id: i32,
14    pub model_id: ModelId,
15    pub measure_id: UsageMeasureId,
16    pub timestamp: DateTime,
17    pub buckets: Vec<i64>,
18    pub is_staff: bool,
19}
20
21#[derive(Copy, Clone, Debug, EnumIter, DeriveRelation)]
22pub enum Relation {
23    #[sea_orm(
24        belongs_to = "super::model::Entity",
25        from = "Column::ModelId",
26        to = "super::model::Column::Id"
27    )]
28    Model,
29    #[sea_orm(
30        belongs_to = "super::usage_measure::Entity",
31        from = "Column::MeasureId",
32        to = "super::usage_measure::Column::Id"
33    )]
34    UsageMeasure,
35}
36
37impl Related<super::model::Entity> for Entity {
38    fn to() -> RelationDef {
39        Relation::Model.def()
40    }
41}
42
43impl Related<super::usage_measure::Entity> for Entity {
44    fn to() -> RelationDef {
45        Relation::UsageMeasure.def()
46    }
47}
48
49impl ActiveModelBehavior for ActiveModel {}