Skip to main content

tea_context/
segment.rs

1use serde::{Deserialize, Serialize};
2use thiserror::Error;
3
4use crate::{CacheScope, ConflictKey, PromptProvenance, PromptSegmentId, TrustLevel};
5
6/// Maximum UTF-8 bytes in one prompt segment.
7pub const MAX_SEGMENT_BYTES: usize = 1024 * 1024;
8
9/// Behavior when a selected segment cannot fit the remaining prompt budget.
10#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, Serialize, Deserialize)]
11#[serde(rename_all = "snake_case")]
12pub enum BudgetBehavior {
13    /// Compilation fails rather than dropping or changing the segment.
14    Required,
15    /// Content may be deterministically shortened with an explicit marker.
16    Truncate,
17    /// Segment may be omitted with a diagnostic.
18    Omit,
19}
20
21/// Whether a selected conflict claim may be replaced by higher precedence.
22#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, Serialize, Deserialize)]
23#[serde(rename_all = "snake_case")]
24pub enum ConflictMode {
25    /// Lower-precedence contenders are explicitly rejected as protected.
26    Protected,
27    /// A higher-precedence contender may replace this claim with diagnostics.
28    Replaceable,
29}
30
31/// Typed conflict claim attached to one segment.
32#[derive(Debug, Clone, PartialEq, Eq, Hash, Serialize, Deserialize)]
33#[serde(rename_all = "camelCase")]
34pub struct ConflictClaim {
35    key: ConflictKey,
36    mode: ConflictMode,
37}
38
39impl ConflictClaim {
40    /// Creates one conflict claim.
41    #[must_use]
42    pub const fn new(key: ConflictKey, mode: ConflictMode) -> Self {
43        Self { key, mode }
44    }
45    /// Returns the conflict key.
46    #[must_use]
47    pub const fn key(&self) -> &ConflictKey {
48        &self.key
49    }
50    /// Returns replacement behavior.
51    #[must_use]
52    pub const fn mode(&self) -> ConflictMode {
53        self.mode
54    }
55}
56
57/// One bounded sourced prompt fragment.
58#[derive(Debug, Clone, PartialEq, Eq, Hash, Serialize)]
59#[serde(rename_all = "camelCase")]
60pub struct PromptSegment {
61    id: PromptSegmentId,
62    content: String,
63    provenance: PromptProvenance,
64    trust: TrustLevel,
65    cache_scope: CacheScope,
66    #[serde(skip_serializing_if = "Option::is_none")]
67    conflict: Option<ConflictClaim>,
68    budget_behavior: BudgetBehavior,
69}
70
71impl PromptSegment {
72    /// Creates a validated prompt segment.
73    ///
74    /// # Errors
75    ///
76    /// Returns an error for empty, oversized, or null-containing content.
77    pub fn new(
78        id: PromptSegmentId,
79        content: impl Into<String>,
80        provenance: PromptProvenance,
81        trust: TrustLevel,
82        cache_scope: CacheScope,
83        budget_behavior: BudgetBehavior,
84    ) -> Result<Self, SegmentError> {
85        let content = content.into();
86        validate_content(&content)?;
87        Ok(Self {
88            id,
89            content,
90            provenance,
91            trust,
92            cache_scope,
93            conflict: None,
94            budget_behavior,
95        })
96    }
97
98    /// Adds a typed conflict claim.
99    #[must_use]
100    pub fn with_conflict(mut self, conflict: ConflictClaim) -> Self {
101        self.conflict = Some(conflict);
102        self
103    }
104
105    /// Returns segment identity.
106    #[must_use]
107    pub const fn id(&self) -> &PromptSegmentId {
108        &self.id
109    }
110    /// Returns exact segment content.
111    #[must_use]
112    pub fn content(&self) -> &str {
113        &self.content
114    }
115    /// Returns source attribution.
116    #[must_use]
117    pub const fn provenance(&self) -> &PromptProvenance {
118        &self.provenance
119    }
120    /// Returns source trust label.
121    #[must_use]
122    pub const fn trust(&self) -> TrustLevel {
123        self.trust
124    }
125    /// Returns intended cache scope.
126    #[must_use]
127    pub const fn cache_scope(&self) -> CacheScope {
128        self.cache_scope
129    }
130    /// Returns optional conflict claim.
131    #[must_use]
132    pub const fn conflict(&self) -> Option<&ConflictClaim> {
133        self.conflict.as_ref()
134    }
135    /// Returns overflow behavior.
136    #[must_use]
137    pub const fn budget_behavior(&self) -> BudgetBehavior {
138        self.budget_behavior
139    }
140}
141
142#[derive(Deserialize)]
143#[serde(rename_all = "camelCase")]
144struct RawPromptSegment {
145    id: PromptSegmentId,
146    content: String,
147    provenance: PromptProvenance,
148    trust: TrustLevel,
149    cache_scope: CacheScope,
150    #[serde(default)]
151    conflict: Option<ConflictClaim>,
152    budget_behavior: BudgetBehavior,
153}
154
155impl<'de> Deserialize<'de> for PromptSegment {
156    fn deserialize<D>(deserializer: D) -> Result<Self, D::Error>
157    where
158        D: serde::Deserializer<'de>,
159    {
160        let raw = RawPromptSegment::deserialize(deserializer)?;
161        let mut segment = Self::new(
162            raw.id,
163            raw.content,
164            raw.provenance,
165            raw.trust,
166            raw.cache_scope,
167            raw.budget_behavior,
168        )
169        .map_err(serde::de::Error::custom)?;
170        segment.conflict = raw.conflict;
171        Ok(segment)
172    }
173}
174
175/// Invalid prompt segment content.
176#[derive(Debug, Clone, Copy, PartialEq, Eq, Error)]
177#[error("prompt segment content is invalid")]
178pub struct SegmentError;
179
180fn validate_content(content: &str) -> Result<(), SegmentError> {
181    if content.is_empty() || content.len() > MAX_SEGMENT_BYTES || content.contains('\0') {
182        Err(SegmentError)
183    } else {
184        Ok(())
185    }
186}