Skip to main content

tea_context/
module.rs

1use std::collections::BTreeSet;
2
3use serde::{Deserialize, Serialize};
4use thiserror::Error;
5
6use crate::{PromptAuthority, PromptModuleId, PromptSegment};
7
8/// Maximum segments contributed by one module.
9pub const MAX_MODULE_SEGMENTS: usize = 256;
10
11/// Bounded priority used only within one authority class.
12#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Hash, Serialize, Deserialize)]
13#[serde(transparent)]
14pub struct PromptPriority(i16);
15
16impl PromptPriority {
17    /// Creates a priority from its signed bounded representation.
18    #[must_use]
19    pub const fn new(value: i16) -> Self {
20        Self(value)
21    }
22    /// Returns numeric priority; greater values sort first within authority.
23    #[must_use]
24    pub const fn get(self) -> i16 {
25        self.0
26    }
27}
28
29/// Ordered prompt fragments contributed at one authority and priority.
30#[derive(Debug, Clone, PartialEq, Eq, Serialize)]
31#[serde(rename_all = "camelCase")]
32pub struct PromptModule {
33    id: PromptModuleId,
34    authority: PromptAuthority,
35    priority: PromptPriority,
36    segments: Vec<PromptSegment>,
37}
38
39impl PromptModule {
40    /// Creates a validated non-empty prompt module.
41    ///
42    /// # Errors
43    ///
44    /// Returns an error for an empty/oversized collection or duplicate segment IDs.
45    pub fn new(
46        id: PromptModuleId,
47        authority: PromptAuthority,
48        priority: PromptPriority,
49        segments: Vec<PromptSegment>,
50    ) -> Result<Self, ModuleError> {
51        if segments.is_empty() || segments.len() > MAX_MODULE_SEGMENTS {
52            return Err(ModuleError::InvalidSegmentCount);
53        }
54        let mut ids = BTreeSet::new();
55        if segments.iter().any(|segment| !ids.insert(segment.id())) {
56            return Err(ModuleError::DuplicateSegmentId);
57        }
58        Ok(Self {
59            id,
60            authority,
61            priority,
62            segments,
63        })
64    }
65
66    /// Returns module identity.
67    #[must_use]
68    pub const fn id(&self) -> &PromptModuleId {
69        &self.id
70    }
71    /// Returns fixed authority.
72    #[must_use]
73    pub const fn authority(&self) -> PromptAuthority {
74        self.authority
75    }
76    /// Returns within-authority priority.
77    #[must_use]
78    pub const fn priority(&self) -> PromptPriority {
79        self.priority
80    }
81    /// Returns source-ordered segments.
82    #[must_use]
83    pub fn segments(&self) -> &[PromptSegment] {
84        &self.segments
85    }
86}
87
88#[derive(Deserialize)]
89#[serde(rename_all = "camelCase")]
90struct RawPromptModule {
91    id: PromptModuleId,
92    authority: PromptAuthority,
93    priority: PromptPriority,
94    segments: Vec<PromptSegment>,
95}
96
97impl<'de> Deserialize<'de> for PromptModule {
98    fn deserialize<D>(deserializer: D) -> Result<Self, D::Error>
99    where
100        D: serde::Deserializer<'de>,
101    {
102        let raw = RawPromptModule::deserialize(deserializer)?;
103        Self::new(raw.id, raw.authority, raw.priority, raw.segments)
104            .map_err(serde::de::Error::custom)
105    }
106}
107
108/// Prompt module invariant failure.
109#[derive(Debug, Clone, Copy, PartialEq, Eq, Error)]
110pub enum ModuleError {
111    /// Module contains no segments or too many segments.
112    #[error("prompt module segment count is invalid")]
113    InvalidSegmentCount,
114    /// One segment identity appears more than once in a module.
115    #[error("prompt module contains a duplicate segment ID")]
116    DuplicateSegmentId,
117}