cognee_database/entities/
session_model_usage.rs

1//! SeaORM entity for the `session_model_usage` table (LIB-03).
2//!
3//! Per-`(session_id, user_id, model)` token + cost aggregate, populated by
4//! `accumulate_usage` (LIB-05) when an LLM call fires inside a tracked
5//! session scope. Normalizing this out of `session_records` lets
6//! mixed-model sessions (e.g. embedding calls + completion calls on
7//! different models) attribute cost correctly via
8//! `GET /api/v1/sessions/cost-by-model`.
9//!
10//! Mirrors Python's `SessionModelUsage` SQLAlchemy model at
11//! `cognee/modules/session_lifecycle/models.py:89-126`.
12
13use sea_orm::entity::prelude::*;
14use serde::{Deserialize, Serialize};
15
16#[derive(Clone, Debug, PartialEq, DeriveEntityModel, Serialize, Deserialize)]
17#[sea_orm(table_name = "session_model_usage")]
18pub struct Model {
19    #[sea_orm(primary_key, auto_increment = false)]
20    pub session_id: String,
21    /// Owning user id, hex-encoded UUID (see `session_record.rs`).
22    #[sea_orm(primary_key, auto_increment = false)]
23    pub user_id: String,
24    #[sea_orm(primary_key, auto_increment = false)]
25    pub model: String,
26
27    pub tokens_in: i32,
28    pub tokens_out: i32,
29    pub cost_usd: f64,
30
31    pub updated_at: DateTimeUtc,
32}
33
34#[derive(Copy, Clone, Debug, EnumIter, DeriveRelation)]
35pub enum Relation {}
36
37impl ActiveModelBehavior for ActiveModel {}
38
39impl Model {
40    /// Serialize to a JSON object whose key ordering matches Python's
41    /// [`SessionModelUsage.to_dict()`](https://github.com/topoteretes/cognee/blob/main/cognee/modules/session_lifecycle/models.py#L116-L126).
42    pub fn to_dict(&self) -> serde_json::Value {
43        serde_json::json!({
44            "session_id": self.session_id,
45            "user_id": self.user_id,
46            "model": self.model,
47            "tokens_in": self.tokens_in,
48            "tokens_out": self.tokens_out,
49            "cost_usd": self.cost_usd,
50            "updated_at": self.updated_at.to_rfc3339(),
51        })
52    }
53}
cognee_database/entities/session_model_usage.rs

cognee_database/entities/
session_model_usage.rs
