1use std::collections::HashSet;
2
3use serde::Serialize;
4
5use crate::graph_index::GraphIndex;
6use crate::retrieval::ground_subgraph;
7use crate::schema::{Edge, Graph};
8
9use super::{ContextEdgeExpectation, ContextOmissionExpectation, GoldenCase};
10
11#[derive(Serialize, Clone, Debug)]
13pub struct ContextCaseResult {
14 pub query: String,
15 #[serde(skip_serializing_if = "Option::is_none")]
16 pub expected_focus_route: Option<String>,
17 #[serde(default, skip_serializing_if = "Vec::is_empty")]
18 pub expected_focus_routes: Vec<String>,
19 pub focus_route: String,
20 #[serde(skip_serializing_if = "Option::is_none")]
21 pub expected_compound_primary: Option<String>,
22 #[serde(skip_serializing_if = "Option::is_none")]
23 pub compound_primary: Option<String>,
24 #[serde(skip_serializing_if = "Vec::is_empty")]
25 pub vocabulary_gap_terms: Vec<String>,
26 #[serde(skip_serializing_if = "Option::is_none")]
27 pub expected_no_vocabulary_gap: Option<bool>,
28 #[serde(skip_serializing_if = "Vec::is_empty")]
29 pub editable_vocab_targets: Vec<String>,
30 #[serde(skip_serializing_if = "Vec::is_empty")]
31 pub non_editable_vocab_targets: Vec<String>,
32 #[serde(skip_serializing_if = "Vec::is_empty")]
33 pub omitted_context: Vec<ContextOmissionExpectation>,
34 #[serde(skip_serializing_if = "Vec::is_empty")]
35 pub omitted_compound_anchors: Vec<ContextOmissionExpectation>,
36 #[serde(skip_serializing_if = "Option::is_none")]
37 pub vocabulary_gap_ok: Option<bool>,
38 #[serde(skip_serializing_if = "Option::is_none")]
39 pub no_vocabulary_gap_ok: Option<bool>,
40 #[serde(skip_serializing_if = "Option::is_none")]
41 pub editable_vocab_targets_ok: Option<bool>,
42 #[serde(skip_serializing_if = "Option::is_none")]
43 pub non_editable_vocab_targets_ok: Option<bool>,
44 #[serde(skip_serializing_if = "Option::is_none")]
45 pub omitted_context_ok: Option<bool>,
46 #[serde(skip_serializing_if = "Option::is_none")]
47 pub omitted_compound_anchors_ok: Option<bool>,
48 #[serde(skip_serializing_if = "Option::is_none")]
49 pub expected_partition: Option<String>,
50 #[serde(skip_serializing_if = "Option::is_none")]
51 pub focus_route_partition: Option<String>,
52 #[serde(skip_serializing_if = "Option::is_none")]
53 pub partition_ok: Option<bool>,
54 #[serde(skip_serializing_if = "Option::is_none")]
55 pub focus_route_ok: Option<bool>,
56 #[serde(skip_serializing_if = "Option::is_none")]
57 pub compound_primary_ok: Option<bool>,
58 pub recall: Option<f64>,
60 pub noise: Option<f64>,
62 pub receipt_coverage: f64,
64 pub pack_size: usize,
66 pub edge_recall: Option<f64>,
68 pub missing: Vec<String>,
69 pub leaked: Vec<String>,
70 pub missing_edges: Vec<ContextEdgeExpectation>,
71 pub missing_omitted_context: Vec<ContextOmissionExpectation>,
72 pub missing_omitted_compound_anchors: Vec<ContextOmissionExpectation>,
73}
74
75#[derive(Serialize, Clone, Debug)]
77pub struct ContextReport {
78 pub context_cases: usize,
79 pub context_recall: f64,
81 pub context_noise: f64,
83 pub receipt_coverage: f64,
85 pub mean_pack_size: f64,
87 pub edge_recall: f64,
89 #[serde(skip_serializing_if = "Option::is_none")]
91 pub focus_route_match_rate: Option<f64>,
92 #[serde(skip_serializing_if = "Option::is_none")]
94 pub compound_primary_match_rate: Option<f64>,
95 #[serde(skip_serializing_if = "Option::is_none")]
97 pub vocabulary_gap_match_rate: Option<f64>,
98 #[serde(skip_serializing_if = "Option::is_none")]
100 pub no_vocabulary_gap_match_rate: Option<f64>,
101 #[serde(skip_serializing_if = "Option::is_none")]
103 pub editable_vocab_target_match_rate: Option<f64>,
104 #[serde(skip_serializing_if = "Option::is_none")]
106 pub non_editable_vocab_target_match_rate: Option<f64>,
107 #[serde(skip_serializing_if = "Option::is_none")]
109 pub omitted_context_match_rate: Option<f64>,
110 #[serde(skip_serializing_if = "Option::is_none")]
112 pub omitted_compound_anchor_match_rate: Option<f64>,
113 #[serde(skip_serializing_if = "Option::is_none")]
115 pub partition_match_rate: Option<f64>,
116 pub cases: Vec<ContextCaseResult>,
117}
118
119pub fn evaluate_context(
122 graph: &Graph,
123 cases: &[GoldenCase],
124 limit: usize,
125 depth: usize,
126 width: usize,
127) -> ContextReport {
128 let mut results = Vec::new();
129 let (mut recall_sum, mut recall_n) = (0.0f64, 0usize);
130 let (mut noise_sum, mut noise_n) = (0.0f64, 0usize);
131 let (mut edge_recall_sum, mut edge_recall_n) = (0.0f64, 0usize);
132 let (mut partition_ok_sum, mut partition_ok_n) = (0usize, 0usize);
133 let mut receipt_sum = 0.0f64;
134 let mut pack_size_sum = 0usize;
135 let (mut compound_primary_ok_sum, mut compound_primary_ok_n) = (0usize, 0usize);
136 let (mut vocabulary_gap_ok_sum, mut vocabulary_gap_ok_n) = (0usize, 0usize);
137 let (mut no_vocabulary_gap_ok_sum, mut no_vocabulary_gap_ok_n) = (0usize, 0usize);
138 let (mut editable_vocab_target_ok_sum, mut editable_vocab_target_ok_n) = (0usize, 0usize);
139 let (mut non_editable_vocab_target_ok_sum, mut non_editable_vocab_target_ok_n) =
140 (0usize, 0usize);
141 let (mut omitted_context_ok_sum, mut omitted_context_ok_n) = (0usize, 0usize);
142 let (mut omitted_compound_anchor_ok_sum, mut omitted_compound_anchor_ok_n) = (0usize, 0usize);
143 let index = GraphIndex::build(graph);
144
145 for c in cases {
146 if c.context_must.is_empty()
147 && c.context_must_not.is_empty()
148 && c.context_edges_must.is_empty()
149 && c.expected_focus_route.is_none()
150 && c.expected_compound_primary.is_none()
151 && c.expected_vocabulary_gap_terms.is_empty()
152 && !c.expected_no_vocabulary_gap
153 && c.expected_vocab_editable_targets.is_empty()
154 && c.expected_vocab_non_editable_targets.is_empty()
155 && c.expected_omitted_context.is_empty()
156 && c.expected_omitted_compound_anchors.is_empty()
157 {
158 continue;
159 }
160 let sg = ground_subgraph(graph, &c.query, limit, depth, width);
161 let ctx: HashSet<&str> = sg.context_order.iter().map(String::as_str).collect();
162 let focus_route_partition = index
163 .node(&sg.route)
164 .and_then(|node| node.partition.clone());
165 let partition_ok = c
166 .expected_partition
167 .as_ref()
168 .map(|expected| focus_route_partition.as_deref() == Some(expected.as_str()));
169 if let Some(ok) = partition_ok {
170 partition_ok_n += 1;
171 partition_ok_sum += ok as usize;
172 }
173 let compound_primary = sg
174 .compound
175 .as_ref()
176 .map(|compound| compound.primary.clone());
177 let compound_primary_ok = c
178 .expected_compound_primary
179 .as_ref()
180 .map(|expected| compound_primary.as_deref() == Some(expected.as_str()));
181 if let Some(ok) = compound_primary_ok {
182 compound_primary_ok_n += 1;
183 compound_primary_ok_sum += ok as usize;
184 }
185 let vocabulary_gap_terms = sg
186 .compound
187 .as_ref()
188 .map(|compound| compound.missing_terms.clone())
189 .unwrap_or_default();
190 let (editable_vocab_targets, non_editable_vocab_targets) =
191 vocabulary_target_candidates(&index, sg.compound.as_ref());
192 let vocabulary_gap_ok = (!c.expected_vocabulary_gap_terms.is_empty())
193 .then(|| contains_all(&vocabulary_gap_terms, &c.expected_vocabulary_gap_terms));
194 if let Some(ok) = vocabulary_gap_ok {
195 vocabulary_gap_ok_n += 1;
196 vocabulary_gap_ok_sum += ok as usize;
197 }
198 let no_vocabulary_gap_ok = c
199 .expected_no_vocabulary_gap
200 .then_some(vocabulary_gap_terms.is_empty());
201 if let Some(ok) = no_vocabulary_gap_ok {
202 no_vocabulary_gap_ok_n += 1;
203 no_vocabulary_gap_ok_sum += ok as usize;
204 }
205 let editable_vocab_targets_ok = (!c.expected_vocab_editable_targets.is_empty())
206 .then(|| contains_all(&editable_vocab_targets, &c.expected_vocab_editable_targets));
207 if let Some(ok) = editable_vocab_targets_ok {
208 editable_vocab_target_ok_n += 1;
209 editable_vocab_target_ok_sum += ok as usize;
210 }
211 let non_editable_vocab_targets_ok = (!c.expected_vocab_non_editable_targets.is_empty())
212 .then(|| {
213 contains_all(
214 &non_editable_vocab_targets,
215 &c.expected_vocab_non_editable_targets,
216 )
217 });
218 if let Some(ok) = non_editable_vocab_targets_ok {
219 non_editable_vocab_target_ok_n += 1;
220 non_editable_vocab_target_ok_sum += ok as usize;
221 }
222 let omitted_context = sg
223 .omitted_context_candidates
224 .iter()
225 .map(|candidate| ContextOmissionExpectation {
226 node_id: candidate.node_id.clone(),
227 reason: candidate.reason.clone(),
228 })
229 .collect::<Vec<_>>();
230 let omitted_compound_anchors = sg
231 .compound
232 .as_ref()
233 .map(|compound| {
234 compound
235 .omitted_anchors
236 .iter()
237 .map(|anchor| ContextOmissionExpectation {
238 node_id: anchor.node_id.clone(),
239 reason: anchor.reason.clone(),
240 })
241 .collect::<Vec<_>>()
242 })
243 .unwrap_or_default();
244 let missing_omitted_context =
245 missing_omissions(&omitted_context, &c.expected_omitted_context);
246 let omitted_context_ok =
247 (!c.expected_omitted_context.is_empty()).then_some(missing_omitted_context.is_empty());
248 if let Some(ok) = omitted_context_ok {
249 omitted_context_ok_n += 1;
250 omitted_context_ok_sum += ok as usize;
251 }
252 let missing_omitted_compound_anchors = missing_omissions(
253 &omitted_compound_anchors,
254 &c.expected_omitted_compound_anchors,
255 );
256 let omitted_compound_anchors_ok = (!c.expected_omitted_compound_anchors.is_empty())
257 .then_some(missing_omitted_compound_anchors.is_empty());
258 if let Some(ok) = omitted_compound_anchors_ok {
259 omitted_compound_anchor_ok_n += 1;
260 omitted_compound_anchor_ok_sum += ok as usize;
261 }
262
263 let recall = if c.context_must.is_empty() {
264 None
265 } else {
266 let hit = c
267 .context_must
268 .iter()
269 .filter(|m| ctx.contains(m.as_str()))
270 .count();
271 recall_sum += hit as f64 / c.context_must.len() as f64;
272 recall_n += 1;
273 Some(hit as f64 / c.context_must.len() as f64)
274 };
275 let noise = if c.context_must_not.is_empty() {
276 None
277 } else {
278 let leaked = c
279 .context_must_not
280 .iter()
281 .filter(|m| ctx.contains(m.as_str()))
282 .count();
283 noise_sum += leaked as f64 / c.context_must_not.len() as f64;
284 noise_n += 1;
285 Some(leaked as f64 / c.context_must_not.len() as f64)
286 };
287 let missing = c
288 .context_must
289 .iter()
290 .filter(|m| !ctx.contains(m.as_str()))
291 .cloned()
292 .collect();
293 let leaked = c
294 .context_must_not
295 .iter()
296 .filter(|m| ctx.contains(m.as_str()))
297 .cloned()
298 .collect();
299 let missing_edges = c
300 .context_edges_must
301 .iter()
302 .filter(|expected| !edge_present(expected, &sg.edges))
303 .cloned()
304 .collect::<Vec<_>>();
305 let edge_recall = if c.context_edges_must.is_empty() {
306 None
307 } else {
308 let recall = (c.context_edges_must.len() - missing_edges.len()) as f64
309 / c.context_edges_must.len() as f64;
310 edge_recall_sum += recall;
311 edge_recall_n += 1;
312 Some(recall)
313 };
314 let pack_size = sg.context_order.len();
315 let receipt_coverage = super::receipt_coverage(&index, &sg.context_order);
316 receipt_sum += receipt_coverage;
317 pack_size_sum += pack_size;
318 results.push(ContextCaseResult {
319 query: c.query.clone(),
320 expected_focus_route: c.expected_focus_route.clone(),
321 expected_focus_routes: c.expected_focus_routes.clone(),
322 focus_route: sg.route.clone(),
323 expected_compound_primary: c.expected_compound_primary.clone(),
324 compound_primary,
325 vocabulary_gap_terms,
326 expected_no_vocabulary_gap: c.expected_no_vocabulary_gap.then_some(true),
327 editable_vocab_targets,
328 non_editable_vocab_targets,
329 omitted_context,
330 omitted_compound_anchors,
331 vocabulary_gap_ok,
332 no_vocabulary_gap_ok,
333 editable_vocab_targets_ok,
334 non_editable_vocab_targets_ok,
335 omitted_context_ok,
336 omitted_compound_anchors_ok,
337 expected_partition: c.expected_partition.clone(),
338 focus_route_partition,
339 partition_ok,
340 focus_route_ok: acceptable_focus_routes(c).map(|expected| expected.contains(&sg.route)),
341 compound_primary_ok,
342 recall,
343 noise,
344 receipt_coverage,
345 pack_size,
346 edge_recall,
347 missing,
348 leaked,
349 missing_edges,
350 missing_omitted_context,
351 missing_omitted_compound_anchors,
352 });
353 }
354
355 let context_cases = results.len();
356
357 ContextReport {
358 context_cases,
359 context_recall: if recall_n == 0 {
360 0.0
361 } else {
362 recall_sum / recall_n as f64
363 },
364 context_noise: if noise_n == 0 {
365 0.0
366 } else {
367 noise_sum / noise_n as f64
368 },
369 receipt_coverage: if context_cases == 0 {
370 0.0
371 } else {
372 receipt_sum / context_cases as f64
373 },
374 mean_pack_size: if context_cases == 0 {
375 0.0
376 } else {
377 pack_size_sum as f64 / context_cases as f64
378 },
379 edge_recall: if edge_recall_n == 0 {
380 0.0
381 } else {
382 edge_recall_sum / edge_recall_n as f64
383 },
384 focus_route_match_rate: focus_route_match_rate(&results),
385 compound_primary_match_rate: (compound_primary_ok_n > 0)
386 .then_some(compound_primary_ok_sum as f64 / compound_primary_ok_n as f64),
387 vocabulary_gap_match_rate: (vocabulary_gap_ok_n > 0)
388 .then_some(vocabulary_gap_ok_sum as f64 / vocabulary_gap_ok_n as f64),
389 no_vocabulary_gap_match_rate: (no_vocabulary_gap_ok_n > 0)
390 .then_some(no_vocabulary_gap_ok_sum as f64 / no_vocabulary_gap_ok_n as f64),
391 editable_vocab_target_match_rate: (editable_vocab_target_ok_n > 0)
392 .then_some(editable_vocab_target_ok_sum as f64 / editable_vocab_target_ok_n as f64),
393 non_editable_vocab_target_match_rate: (non_editable_vocab_target_ok_n > 0).then_some(
394 non_editable_vocab_target_ok_sum as f64 / non_editable_vocab_target_ok_n as f64,
395 ),
396 omitted_context_match_rate: (omitted_context_ok_n > 0)
397 .then_some(omitted_context_ok_sum as f64 / omitted_context_ok_n as f64),
398 omitted_compound_anchor_match_rate: (omitted_compound_anchor_ok_n > 0)
399 .then_some(omitted_compound_anchor_ok_sum as f64 / omitted_compound_anchor_ok_n as f64),
400 partition_match_rate: (partition_ok_n > 0)
401 .then_some(partition_ok_sum as f64 / partition_ok_n as f64),
402 cases: results,
403 }
404}
405
406fn acceptable_focus_routes(c: &GoldenCase) -> Option<Vec<String>> {
407 let mut expected = c.expected_focus_routes.clone();
408 if let Some(route) = &c.expected_focus_route {
409 expected.push(route.clone());
410 }
411 expected.sort();
412 expected.dedup();
413 (!expected.is_empty()).then_some(expected)
414}
415
416fn vocabulary_target_candidates(
417 index: &GraphIndex<'_>,
418 compound: Option<&crate::retrieval::CompoundSubgraph>,
419) -> (Vec<String>, Vec<String>) {
420 let Some(compound) = compound else {
421 return (Vec::new(), Vec::new());
422 };
423 let mut candidates = Vec::new();
424 candidates.push(compound.primary.clone());
425 candidates.extend(compound.anchors.iter().map(|anchor| anchor.node_id.clone()));
426 candidates.extend(
427 compound
428 .omitted_anchors
429 .iter()
430 .map(|anchor| anchor.node_id.clone()),
431 );
432 candidates.sort();
433 candidates.dedup();
434
435 let mut editable = Vec::new();
436 let mut non_editable = Vec::new();
437 for node_id in candidates {
438 let target = if index.node(&node_id).is_some_and(vocab_target_is_editable) {
439 &mut editable
440 } else {
441 &mut non_editable
442 };
443 target.push(node_id);
444 }
445 (editable, non_editable)
446}
447
448fn vocab_target_is_editable(node: &crate::schema::Node) -> bool {
449 node.source_files
450 .first()
451 .is_some_and(|source| source.ends_with(".md"))
452}
453
454fn contains_all(actual: &[String], expected: &[String]) -> bool {
455 expected.iter().all(|item| actual.contains(item))
456}
457
458fn missing_omissions(
459 actual: &[ContextOmissionExpectation],
460 expected: &[ContextOmissionExpectation],
461) -> Vec<ContextOmissionExpectation> {
462 expected
463 .iter()
464 .filter(|expected| !actual.contains(expected))
465 .cloned()
466 .collect()
467}
468
469fn focus_route_match_rate(results: &[ContextCaseResult]) -> Option<f64> {
470 let judged = results
471 .iter()
472 .filter_map(|case| case.focus_route_ok)
473 .collect::<Vec<_>>();
474 if judged.is_empty() {
475 None
476 } else {
477 Some(judged.iter().filter(|ok| **ok).count() as f64 / judged.len() as f64)
478 }
479}
480
481fn edge_present(expected: &ContextEdgeExpectation, edges: &[Edge]) -> bool {
482 edges.iter().any(|edge| {
483 edge.from == expected.from && edge.to == expected.to && edge.relation == expected.relation
484 })
485}