Skip to main content

aft/commands/semantic_search/
telemetry.rs

1use std::fmt;
2
3use serde::{Deserialize, Serialize};
4
5use super::blocks::{BlockEntry, StabilityUnit};
6use super::evidence_descriptor::{EvidenceDescriptor, EvidenceKind, EvidenceTier};
7use super::generation_token::GenerationToken;
8use super::paging::SearchPage;
9use super::plan_table::{SearchLaneKind, SearchShape};
10use super::provenance::{LanePositions, LanePositionsAccessor, ProvenanceError};
11use super::trailer::{ExactPassState, MissingBoundedExactPassDisclosure, SearchTrailer};
12
13#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
14#[serde(rename_all = "snake_case")]
15pub enum ConfidenceTelemetry {
16    High,
17    Low,
18}
19
20/// Stable execution facts that are not derivable from the block reply itself.
21#[derive(Debug, Clone, PartialEq, Eq)]
22pub struct TelemetryRun {
23    pub shape: SearchShape,
24    pub confidence: Option<ConfidenceTelemetry>,
25    pub variants: Vec<String>,
26    pub embedding_calls: usize,
27    pub snapshot_generation: GenerationToken,
28}
29
30#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
31pub struct PlanTelemetry {
32    pub shape: SearchShape,
33    pub lanes_run: Vec<SearchLaneKind>,
34    pub exact_tier: EvidenceKind,
35    #[serde(skip_serializing_if = "Option::is_none")]
36    pub confidence: Option<ConfidenceTelemetry>,
37    pub variants: Vec<String>,
38    pub embedding_calls: usize,
39    pub snapshot_generation: GenerationToken,
40    pub retrieval_depth: usize,
41    pub depth_tier: usize,
42    pub lanes_exhausted: bool,
43    pub stability_void: bool,
44}
45
46/// One structured result. The ranked tuple stays flat to match search result
47/// records; provenance and the frozen evidence copy are added as sibling fields.
48#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
49pub struct StructuredResult {
50    #[serde(flatten)]
51    pub ranked_tuple: super::comparator::RankedTuple,
52    pub lane_positions: LanePositions,
53    pub evidence_descriptor: EvidenceDescriptor,
54}
55
56#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
57pub struct StructuredContent {
58    pub plan: PlanTelemetry,
59    pub results: Vec<StructuredResult>,
60}
61
62#[derive(Debug)]
63pub enum TelemetryError {
64    SnapshotGenerationMismatch,
65    Provenance(ProvenanceError),
66}
67
68impl fmt::Display for TelemetryError {
69    fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
70        match self {
71            Self::SnapshotGenerationMismatch => formatter.write_str(
72                "telemetry generation does not equal the generation that defined the canonical list",
73            ),
74            Self::Provenance(error) => error.fmt(formatter),
75        }
76    }
77}
78
79impl std::error::Error for TelemetryError {}
80
81impl From<ProvenanceError> for TelemetryError {
82    fn from(error: ProvenanceError) -> Self {
83        Self::Provenance(error)
84    }
85}
86
87/// Access to all decision inputs, intentionally excluding lane positions.
88/// Holding this view is sufficient for ranking, confidence, paging and trailer
89/// assembly, including when the output-only provenance accessor is unavailable.
90pub struct DecisionView<'a> {
91    page: &'a SearchPage,
92}
93
94impl<'a> DecisionView<'a> {
95    pub fn page_entries(&self) -> &'a [BlockEntry] {
96        &self.page.reply.page
97    }
98
99    pub fn page_stability_units(&self) -> Vec<StabilityUnit> {
100        self.page.reply.page_stability_units()
101    }
102
103    pub fn derive_confidence<T>(&self, derive: impl FnOnce(&[BlockEntry]) -> T) -> T {
104        derive(self.page_entries())
105    }
106
107    pub fn assemble_page<T>(&self, assemble: impl FnOnce(&[BlockEntry]) -> T) -> T {
108        assemble(self.page_entries())
109    }
110
111    pub fn assemble_trailer(
112        &self,
113        exact_pass: ExactPassState<'_>,
114    ) -> Result<SearchTrailer, MissingBoundedExactPassDisclosure> {
115        SearchTrailer::from_page(self.page, exact_pass)
116    }
117}
118
119/// The provenance capability is retained only for final structured output.
120/// Decision code receives `DecisionView`, which cannot call the accessor.
121pub struct TelemetryAssembler<'a, A> {
122    page: &'a SearchPage,
123    lane_positions: A,
124}
125
126impl<'a, A> TelemetryAssembler<'a, A>
127where
128    A: LanePositionsAccessor,
129{
130    pub fn new(page: &'a SearchPage, lane_positions: A) -> Self {
131        Self {
132            page,
133            lane_positions,
134        }
135    }
136
137    pub fn decision_view(&self) -> DecisionView<'a> {
138        DecisionView { page: self.page }
139    }
140
141    pub fn assemble(self, run: TelemetryRun) -> Result<StructuredContent, TelemetryError> {
142        if self.page.reply.canonical_list.key.snapshot_generation
143            != run.snapshot_generation.as_str()
144        {
145            return Err(TelemetryError::SnapshotGenerationMismatch);
146        }
147
148        let results = self
149            .page
150            .reply
151            .page
152            .iter()
153            .map(|entry| {
154                let stability = entry.stability_unit();
155                Ok(StructuredResult {
156                    ranked_tuple: stability.ranked_tuple,
157                    lane_positions: self.lane_positions.lane_positions(entry)?,
158                    evidence_descriptor: stability.evidence_descriptor,
159                })
160            })
161            .collect::<Result<Vec<_>, ProvenanceError>>()?;
162
163        let plan = PlanTelemetry {
164            shape: run.shape,
165            lanes_run: lanes_run(self.page),
166            exact_tier: exact_tier(self.page),
167            confidence: run.confidence,
168            variants: run.variants,
169            embedding_calls: run.embedding_calls,
170            snapshot_generation: run.snapshot_generation,
171            retrieval_depth: self.page.reply.retrieval_depth,
172            depth_tier: self.page.reply.depth_tier,
173            lanes_exhausted: self.page.reply.lanes_exhausted,
174            stability_void: self.page.stability_void,
175        };
176        Ok(StructuredContent { plan, results })
177    }
178}
179
180fn lanes_run(page: &SearchPage) -> Vec<SearchLaneKind> {
181    let mut lanes = page
182        .reply
183        .lane_enumeration_counts
184        .keys()
185        .copied()
186        .collect::<Vec<_>>();
187    lanes.sort_by_key(SearchLaneKind::default_plan_order_index);
188    lanes
189}
190
191fn exact_tier(page: &SearchPage) -> EvidenceKind {
192    page.reply
193        .canonical_list
194        .entries()
195        .find(|entry| entry.result.evidence.tier == EvidenceTier::Exact)
196        .map_or(EvidenceKind::None, |entry| entry.result.evidence.kind)
197}