Skip to main content

systemprompt_models/
attribution.rs

1//! Scope attribution of an AI request.
2//!
3//! The value the request is charged to in each subject dimension a tenant
4//! registers (`project`, `cost_centre`, ...), which signal decided each value,
5//! and the API key that authenticated it.
6//!
7//! Core owns no organisation model. The entry layer resolves one
8//! [`AttributionEntry`] per registered dimension (the
9//! `x-systemprompt-scope-<dimension>` header, then the API key's bound value,
10//! then the dimension provider's first value) and the gateway carries the set
11//! unchanged onto the audit row, the quota windows and the OTLP export.
12//!
13//! Copyright (c) systemprompt.io — Business Source License 1.1.
14//! See <https://systemprompt.io> for licensing details.
15
16use serde::{Deserialize, Serialize};
17use systemprompt_identifiers::{ApiKeyId, ScopeDimension};
18
19/// Which signal decided an attributed value. Persisted as
20/// `ai_request_attributions.source`, a CHECK-constrained column.
21#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, Serialize, Deserialize)]
22#[serde(rename_all = "snake_case")]
23pub enum AttributionSource {
24    Header,
25    ApiKey,
26    Default,
27}
28
29impl AttributionSource {
30    #[must_use]
31    pub const fn as_str(self) -> &'static str {
32        match self {
33            Self::Header => "header",
34            Self::ApiKey => "api_key",
35            Self::Default => "default",
36        }
37    }
38
39    #[must_use]
40    pub fn parse(value: &str) -> Option<Self> {
41        match value {
42            "header" => Some(Self::Header),
43            "api_key" => Some(Self::ApiKey),
44            "default" => Some(Self::Default),
45            _ => None,
46        }
47    }
48}
49
50impl std::fmt::Display for AttributionSource {
51    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
52        f.write_str(self.as_str())
53    }
54}
55
56/// The value one request is attributed to in one dimension.
57#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
58pub struct AttributionEntry {
59    pub dimension: ScopeDimension,
60    pub value: String,
61    pub source: AttributionSource,
62}
63
64/// The resolved scope attribution of one request.
65#[derive(Debug, Clone, PartialEq, Eq, Default, Serialize, Deserialize)]
66pub struct RequestAttribution {
67    pub entries: Vec<AttributionEntry>,
68    pub api_key_id: Option<ApiKeyId>,
69}
70
71impl RequestAttribution {
72    #[must_use]
73    pub fn none() -> Self {
74        Self::default()
75    }
76
77    #[must_use]
78    pub fn value_for(&self, dimension: &str) -> Option<&str> {
79        self.entries
80            .iter()
81            .find(|entry| entry.dimension.as_str() == dimension)
82            .map(|entry| entry.value.as_str())
83    }
84}
85
86/// A value an API key is bound to in one dimension; it attributes the key's
87/// requests when they carry no header for that dimension.
88#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
89pub struct ScopeBinding {
90    pub dimension: ScopeDimension,
91    pub value: String,
92}
93
94impl ScopeBinding {
95    #[must_use]
96    pub fn traffic(class: TrafficClass) -> Self {
97        Self {
98            dimension: ScopeDimension::new(TrafficClass::DIMENSION),
99            value: class.as_str().to_owned(),
100        }
101    }
102}
103
104/// The canonical traffic class of a request, persisted as
105/// `ai_requests.traffic`. `Evaluation` is reachable only through an API key
106/// bound to it, so production usage cannot be relabelled by a header.
107#[derive(Debug, Clone, Copy, Default, PartialEq, Eq, Hash, Serialize, Deserialize)]
108#[serde(rename_all = "snake_case")]
109pub enum TrafficClass {
110    #[default]
111    Production,
112    Evaluation,
113}
114
115impl TrafficClass {
116    pub const DIMENSION: &'static str = "traffic";
117
118    #[must_use]
119    pub const fn as_str(self) -> &'static str {
120        match self {
121            Self::Production => "production",
122            Self::Evaluation => "evaluation",
123        }
124    }
125
126    #[must_use]
127    pub fn parse(value: &str) -> Option<Self> {
128        match value {
129            "production" => Some(Self::Production),
130            "evaluation" => Some(Self::Evaluation),
131            _ => None,
132        }
133    }
134
135    #[must_use]
136    pub fn of(attribution: &RequestAttribution) -> Option<Self> {
137        attribution
138            .value_for(Self::DIMENSION)
139            .map_or(Some(Self::Production), Self::parse)
140    }
141}
142
143impl std::fmt::Display for TrafficClass {
144    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
145        f.write_str(self.as_str())
146    }
147}