use polyc_state::{
command::CommandScope,
digest::ContentDigest,
id::{AttemptId, WorkId},
model_attempt::{ModelConversationId, TenantId},
};
use crate::execution::ExecutionIdentity;
#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Hash)]
pub enum ModelPurpose {
Answering,
Summarization,
Classification,
}
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub struct ModelCostCeilings {
tenant: u64,
conversation: u64,
turn: u64,
call: u64,
}
impl ModelCostCeilings {
pub fn new(tenant: u64, conversation: u64, turn: u64, call: u64) -> Result<Self, ModelError> {
if [tenant, conversation, turn, call].contains(&0) {
return Err(ModelError::MissingBound);
}
Ok(Self {
tenant,
conversation,
turn,
call,
})
}
#[must_use]
pub fn effective(self) -> u64 {
self.tenant
.min(self.conversation)
.min(self.turn)
.min(self.call)
}
#[must_use]
pub const fn tenant(self) -> u64 {
self.tenant
}
#[must_use]
pub const fn conversation(self) -> u64 {
self.conversation
}
#[must_use]
pub const fn turn(self) -> u64 {
self.turn
}
#[must_use]
pub const fn call(self) -> u64 {
self.call
}
}
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct ModelRequestBinding {
pub tenant: TenantId,
pub conversation: ModelConversationId,
pub execution: ExecutionIdentity,
pub claim_attempt: AttemptId,
pub claim_scope: CommandScope,
pub claim_work: WorkId,
pub model_attempt: AttemptId,
pub purpose: ModelPurpose,
pub provider_target: String,
pub model: String,
pub request_digest: ContentDigest,
pub max_request_bytes: u64,
pub max_response_bytes: u64,
pub max_output_tokens: u64,
pub deadline_nanos: u64,
pub audience: String,
}
impl ModelRequestBinding {
#[must_use]
pub const fn claim_fence(&self) -> polyc_state::command::FencingToken {
self.execution.fence()
}
}
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct ModelCapability {
binding: ModelRequestBinding,
ceilings: ModelCostCeilings,
}
impl ModelCapability {
pub fn new(
binding: ModelRequestBinding,
ceilings: ModelCostCeilings,
) -> Result<Self, ModelError> {
validate_binding(&binding)?;
Ok(Self { binding, ceilings })
}
#[must_use]
pub const fn binding(&self) -> &ModelRequestBinding {
&self.binding
}
#[must_use]
pub fn effective_cost_ceiling(&self) -> u64 {
self.ceilings.effective()
}
#[must_use]
pub const fn ceilings(&self) -> ModelCostCeilings {
self.ceilings
}
pub fn verify(&self, request: &ModelRequestBinding) -> Result<(), ModelError> {
validate_binding(request)?;
if self.binding != *request {
return Err(ModelError::BindingMismatch);
}
Ok(())
}
}
#[derive(Debug, Clone, PartialEq, Eq, thiserror::Error)]
pub enum ModelError {
#[error("the model capability omits a required identity or bound")]
MissingBound,
#[error("the model request does not match its bounded capability")]
BindingMismatch,
}
fn validate_binding(binding: &ModelRequestBinding) -> Result<(), ModelError> {
if binding.tenant.is_empty()
|| binding.conversation.is_empty()
|| binding.claim_attempt.is_empty()
|| binding.claim_work.as_str().is_empty()
|| binding.claim_scope.aggregate().as_str().is_empty()
|| binding.claim_scope.partition().as_str().is_empty()
|| binding.claim_scope.namespace().as_str().is_empty()
|| binding.model_attempt.is_empty()
|| binding.provider_target.is_empty()
|| binding.model.is_empty()
|| binding.audience.is_empty()
|| binding.max_request_bytes == 0
|| binding.max_response_bytes == 0
|| binding.max_output_tokens == 0
|| binding.deadline_nanos == 0
{
return Err(ModelError::MissingBound);
}
Ok(())
}
#[cfg(test)]
mod tests {
use polyc_state::{
command::{CommandScope, FencingToken},
digest::ContentDigest,
id::{AggregateId, AttemptId, NamespaceId, PartitionId, WorkId},
model_attempt::{ModelConversationId, TenantId},
};
use super::{
ModelCapability, ModelCostCeilings, ModelError, ModelPurpose, ModelRequestBinding,
};
use crate::execution::{ExecutionId, ExecutionIdentity};
fn binding() -> ModelRequestBinding {
ModelRequestBinding {
tenant: TenantId::new("tenant-a"),
conversation: ModelConversationId::new("conversation-a"),
execution: ExecutionIdentity::new(
ExecutionId::new("execution-a"),
AttemptId::new("transport-a"),
FencingToken::new(7),
)
.expect("valid execution identity"),
claim_attempt: AttemptId::new("claim-a"),
claim_scope: CommandScope::new(
AggregateId::new("work-a"),
PartitionId::new("turns:conversation-a"),
NamespaceId::new("polychrome"),
),
claim_work: WorkId::new("work-a"),
model_attempt: AttemptId::new("model-a"),
purpose: ModelPurpose::Answering,
provider_target: "approved-provider".to_owned(),
model: "approved-model".to_owned(),
request_digest: ContentDigest::from_bytes([7; ContentDigest::LEN]),
max_request_bytes: 1024,
max_response_bytes: 2048,
max_output_tokens: 256,
deadline_nanos: 9,
audience: "model-broker".to_owned(),
}
}
#[test]
fn capability_rejects_each_changed_bound_field() {
let original = binding();
let capability = ModelCapability::new(
original.clone(),
ModelCostCeilings::new(100, 90, 80, 70).expect("non-zero ceilings"),
)
.expect("complete binding");
let mut changed = original.clone();
changed.tenant = TenantId::new("tenant-b");
assert_eq!(
capability.verify(&changed),
Err(ModelError::BindingMismatch)
);
changed = original.clone();
changed.conversation = ModelConversationId::new("conversation-b");
assert_eq!(
capability.verify(&changed),
Err(ModelError::BindingMismatch)
);
changed = original.clone();
changed.execution = ExecutionIdentity::new(
ExecutionId::new("execution-b"),
AttemptId::new("transport-a"),
FencingToken::new(7),
)
.expect("valid execution identity");
assert_eq!(
capability.verify(&changed),
Err(ModelError::BindingMismatch)
);
changed = original.clone();
changed.claim_attempt = AttemptId::new("claim-b");
assert_eq!(
capability.verify(&changed),
Err(ModelError::BindingMismatch)
);
changed = original.clone();
changed.model_attempt = AttemptId::new("model-b");
assert_eq!(
capability.verify(&changed),
Err(ModelError::BindingMismatch)
);
changed = original.clone();
changed.purpose = ModelPurpose::Classification;
assert_eq!(
capability.verify(&changed),
Err(ModelError::BindingMismatch)
);
changed = original.clone();
changed.provider_target = "other-provider".to_owned();
assert_eq!(
capability.verify(&changed),
Err(ModelError::BindingMismatch)
);
changed = original.clone();
changed.model = "other-model".to_owned();
assert_eq!(
capability.verify(&changed),
Err(ModelError::BindingMismatch)
);
changed = original.clone();
changed.request_digest = ContentDigest::from_bytes([8; ContentDigest::LEN]);
assert_eq!(
capability.verify(&changed),
Err(ModelError::BindingMismatch)
);
changed = original.clone();
changed.max_request_bytes = 1025;
assert_eq!(
capability.verify(&changed),
Err(ModelError::BindingMismatch)
);
changed = original.clone();
changed.max_response_bytes = 2049;
assert_eq!(
capability.verify(&changed),
Err(ModelError::BindingMismatch)
);
changed = original.clone();
changed.max_output_tokens = 257;
assert_eq!(
capability.verify(&changed),
Err(ModelError::BindingMismatch)
);
changed = original.clone();
changed.deadline_nanos = 10;
assert_eq!(
capability.verify(&changed),
Err(ModelError::BindingMismatch)
);
changed = original;
changed.audience = "other-broker".to_owned();
assert_eq!(
capability.verify(&changed),
Err(ModelError::BindingMismatch)
);
}
#[test]
fn effective_ceiling_is_the_lowest_required_scope() {
let ceilings = ModelCostCeilings::new(100, 90, 80, 70).expect("non-zero ceilings");
assert_eq!(ceilings.effective(), 70);
assert_eq!(
ModelCostCeilings::new(0, 90, 80, 70),
Err(ModelError::MissingBound)
);
}
}