1use std::collections::BTreeSet;
18
19use crate::api::{RedDBError, RedDBResult};
20
21use super::ask_rql_planner::{validate_candidate, CandidateDisposition, ValidatedCandidate};
22
23#[derive(Debug, Clone, Copy, PartialEq, Eq)]
25pub enum AskIntent {
26 Factual,
29 Synthesis,
31 HowTo,
33}
34
35impl AskIntent {
36 pub fn as_str(&self) -> &'static str {
38 match self {
39 AskIntent::Factual => "factual",
40 AskIntent::Synthesis => "synthesis",
41 AskIntent::HowTo => "how_to",
42 }
43 }
44
45 fn parse(raw: &str) -> Option<AskIntent> {
46 match raw.trim().to_ascii_lowercase().as_str() {
47 "factual" => Some(AskIntent::Factual),
48 "synthesis" => Some(AskIntent::Synthesis),
49 "how_to" | "howto" | "how-to" => Some(AskIntent::HowTo),
50 _ => None,
51 }
52 }
53}
54
55#[derive(Debug, Clone, PartialEq)]
57pub struct ScoredCollection {
58 pub collection: String,
59 pub score: f32,
60 pub columns: Vec<String>,
61}
62
63#[derive(Debug, Clone, Default, PartialEq)]
66pub struct NarrowedSlice {
67 pub collections: Vec<ScoredCollection>,
68}
69
70impl NarrowedSlice {
71 pub fn collection_names(&self) -> Vec<&str> {
73 self.collections
74 .iter()
75 .map(|c| c.collection.as_str())
76 .collect()
77 }
78
79 pub fn is_empty(&self) -> bool {
83 self.collections.is_empty()
84 }
85}
86
87#[derive(Debug, Clone, PartialEq, Eq)]
91pub struct RawSuggestion {
92 pub rql: String,
93 pub rationale: String,
94}
95
96#[derive(Debug, Clone, PartialEq, Eq)]
101pub struct SuggestedStatement {
102 pub rql: String,
104 pub mutating: bool,
107 pub statement_type: &'static str,
109 pub rationale: String,
111}
112
113#[derive(Debug, Clone, PartialEq, Eq)]
115pub struct AskPlan {
116 pub intent: AskIntent,
117 pub query: Option<String>,
119 pub answer: String,
122 pub suggestion: Vec<RawSuggestion>,
125 pub rationale: String,
127}
128
129impl AskPlan {
130 pub fn summary(&self) -> String {
132 let mut out = format!("intent={}", self.intent.as_str());
133 if let Some(query) = &self.query {
134 out.push_str("; query=");
135 out.push_str(query.trim());
136 }
137 out
138 }
139}
140
141pub fn validate_suggestions(raw: &[RawSuggestion]) -> Vec<SuggestedStatement> {
147 let mut out = Vec::new();
148 for item in raw {
149 if let Ok(candidate) = validate_candidate(&item.rql) {
152 out.push(SuggestedStatement {
153 mutating: !candidate.is_read_only(),
154 statement_type: candidate.statement_type,
155 rql: candidate.rql,
156 rationale: item.rationale.clone(),
157 });
158 }
159 }
160 out
161}
162
163pub trait PlannerModel {
170 fn plan(&self, prompt: &str) -> RedDBResult<String>;
171}
172
173impl<F> PlannerModel for F
174where
175 F: Fn(&str) -> RedDBResult<String>,
176{
177 fn plan(&self, prompt: &str) -> RedDBResult<String> {
178 self(prompt)
179 }
180}
181
182#[derive(Debug, Clone, PartialEq, Eq)]
184pub enum PlanRouting {
185 Execute { candidate: ValidatedCandidate },
188 RefuseMutating {
192 statement_type: &'static str,
193 rql: String,
194 },
195 Suggest {
201 answer: String,
202 suggestion: Vec<SuggestedStatement>,
203 },
204 Unsupported { intent: AskIntent },
207}
208
209#[derive(Debug, Clone, PartialEq, Eq)]
211pub struct PlannedRoute {
212 pub prompt: String,
213 pub plan: AskPlan,
214 pub routing: PlanRouting,
215}
216
217pub fn build_planner_prompt(question: &str, slice: &NarrowedSlice) -> String {
221 let mut prompt = String::new();
222 prompt.push_str(
223 "You are the RedDB ASK planner. Classify the user's question intent and, for a \
224 factual question, generate a single read-only RQL SELECT candidate over the \
225 collections below.\n\
226 Respond with ONLY a JSON object, no code fences or commentary:\n\
227 {\"intent\": \"factual\"|\"synthesis\"|\"how_to\", \"query\": \"<read-only RQL or null>\", \
228 \"rationale\": \"<one sentence>\"}\n\
229 Rules: use only the collections and columns listed; never invent a collection; \
230 a factual question MUST carry a read-only SELECT (joins and the global `any` source \
231 are allowed); never emit INSERT/UPDATE/DELETE/DDL.\n\n",
232 );
233
234 if slice.collections.is_empty() {
235 prompt.push_str("Candidate collections: (none)\n");
236 } else {
237 prompt.push_str("Candidate collections (name, score, columns):\n");
238 for c in &slice.collections {
239 prompt.push_str("- ");
240 prompt.push_str(&c.collection);
241 prompt.push_str(&format!(" (score {:.4})", c.score));
242 if !c.columns.is_empty() {
243 let mut cols: BTreeSet<&str> = BTreeSet::new();
244 for col in &c.columns {
245 cols.insert(col.as_str());
246 }
247 prompt.push_str(": ");
248 prompt.push_str(&cols.into_iter().collect::<Vec<_>>().join(", "));
249 }
250 prompt.push('\n');
251 }
252 }
253
254 prompt.push_str("\nQuestion: ");
255 prompt.push_str(question);
256 prompt
257}
258
259pub fn parse_plan(raw: &str) -> RedDBResult<AskPlan> {
262 let json = extract_json_object(raw)
263 .ok_or_else(|| RedDBError::Query("ASK planner returned no JSON plan object".to_string()))?;
264 let parsed: crate::json::Value = crate::json::from_str(json).map_err(|err| {
265 RedDBError::Query(format!("ASK planner returned invalid JSON plan: {err}"))
266 })?;
267 let obj = parsed
268 .as_object()
269 .ok_or_else(|| RedDBError::Query("ASK planner plan must be a JSON object".to_string()))?;
270
271 let intent_raw = obj
272 .get("intent")
273 .and_then(|v| v.as_str())
274 .ok_or_else(|| RedDBError::Query("ASK planner plan is missing `intent`".to_string()))?;
275 let intent = AskIntent::parse(intent_raw).ok_or_else(|| {
276 RedDBError::Query(format!(
277 "ASK planner returned an unknown intent `{intent_raw}`"
278 ))
279 })?;
280
281 let query = obj
282 .get("query")
283 .and_then(|v| v.as_str())
284 .map(str::trim)
285 .filter(|s| !s.is_empty())
286 .map(str::to_string);
287
288 let answer = obj
289 .get("answer")
290 .and_then(|v| v.as_str())
291 .unwrap_or("")
292 .trim()
293 .to_string();
294
295 let suggestion = obj
296 .get("suggestion")
297 .and_then(|v| v.as_array())
298 .map(parse_suggestions)
299 .unwrap_or_default();
300
301 let rationale = obj
302 .get("rationale")
303 .and_then(|v| v.as_str())
304 .unwrap_or("")
305 .trim()
306 .to_string();
307
308 Ok(AskPlan {
309 intent,
310 query,
311 answer,
312 suggestion,
313 rationale,
314 })
315}
316
317fn parse_suggestions(items: &[crate::json::Value]) -> Vec<RawSuggestion> {
322 let mut out = Vec::new();
323 for item in items {
324 if let Some(rql) = item.as_str() {
325 let rql = rql.trim();
326 if !rql.is_empty() {
327 out.push(RawSuggestion {
328 rql: rql.to_string(),
329 rationale: String::new(),
330 });
331 }
332 } else if let Some(obj) = item.as_object() {
333 let rql = obj
334 .get("rql")
335 .or_else(|| obj.get("statement"))
336 .and_then(|v| v.as_str())
337 .unwrap_or("")
338 .trim()
339 .to_string();
340 if rql.is_empty() {
341 continue;
342 }
343 let rationale = obj
344 .get("rationale")
345 .and_then(|v| v.as_str())
346 .unwrap_or("")
347 .trim()
348 .to_string();
349 out.push(RawSuggestion { rql, rationale });
350 }
351 }
352 out
353}
354
355pub fn route_plan(plan: &AskPlan) -> RedDBResult<PlanRouting> {
359 match plan.intent {
360 AskIntent::Factual => {
361 let rql = plan.query.as_deref().ok_or_else(|| {
362 RedDBError::Query(
363 "ASK planner classified the question as factual but produced no query \
364 candidate"
365 .to_string(),
366 )
367 })?;
368 let candidate = validate_candidate(rql)?;
371 match candidate.disposition {
372 CandidateDisposition::ReadOnly => Ok(PlanRouting::Execute { candidate }),
373 CandidateDisposition::Mutating => Ok(PlanRouting::RefuseMutating {
374 statement_type: candidate.statement_type,
375 rql: candidate.rql,
376 }),
377 }
378 }
379 AskIntent::HowTo => {
380 let suggestion = validate_suggestions(&plan.suggestion);
384 Ok(PlanRouting::Suggest {
385 answer: plan.answer.clone(),
386 suggestion,
387 })
388 }
389 other => Ok(PlanRouting::Unsupported { intent: other }),
390 }
391}
392
393pub fn plan_and_route<M: PlannerModel>(
396 question: &str,
397 slice: &NarrowedSlice,
398 model: &M,
399) -> RedDBResult<PlannedRoute> {
400 let prompt = build_planner_prompt(question, slice);
401 let raw = model.plan(&prompt)?;
402 let plan = parse_plan(&raw)?;
403 let routing = route_plan(&plan)?;
404 Ok(PlannedRoute {
405 prompt,
406 plan,
407 routing,
408 })
409}
410
411pub const DEFAULT_MAX_PLAN_STEPS: usize = 3;
417
418#[derive(Debug, Clone, Copy, PartialEq, Eq)]
421pub enum PlanStep {
422 RefineRetrieval,
425 Query,
427}
428
429impl PlanStep {
430 pub fn as_str(&self) -> &'static str {
431 match self {
432 PlanStep::RefineRetrieval => "refine_retrieval",
433 PlanStep::Query => "query",
434 }
435 }
436}
437
438pub fn clamp_plan_steps(requested: Option<usize>, cap: usize) -> usize {
442 let cap = cap.max(1);
443 match requested {
444 Some(n) => n.clamp(1, cap),
445 None => cap,
446 }
447}
448
449#[derive(Debug, Clone, PartialEq, Eq)]
453pub struct BudgetExhausted {
454 pub attempted: PlanStep,
455 pub executed_steps: usize,
456 pub max_steps: usize,
457}
458
459#[derive(Debug, Clone)]
462pub struct PlanBudget {
463 max_steps: usize,
464 executed: Vec<PlanStep>,
465}
466
467impl PlanBudget {
468 pub fn new(requested_steps: Option<usize>, cap: usize) -> Self {
471 Self {
472 max_steps: clamp_plan_steps(requested_steps, cap),
473 executed: Vec::new(),
474 }
475 }
476
477 pub fn max_steps(&self) -> usize {
479 self.max_steps
480 }
481
482 pub fn executed_count(&self) -> usize {
484 self.executed.len()
485 }
486
487 pub fn remaining(&self) -> usize {
489 self.max_steps.saturating_sub(self.executed.len())
490 }
491
492 pub fn executed_steps(&self) -> &[PlanStep] {
494 &self.executed
495 }
496
497 pub fn charge(&mut self, step: PlanStep) -> Result<(), BudgetExhausted> {
500 if self.executed.len() >= self.max_steps {
501 return Err(BudgetExhausted {
502 attempted: step,
503 executed_steps: self.executed.len(),
504 max_steps: self.max_steps,
505 });
506 }
507 self.executed.push(step);
508 Ok(())
509 }
510}
511
512pub trait RetrievalFunnel {
520 fn funnel(&self, expanded: bool) -> RedDBResult<NarrowedSlice>;
521}
522
523impl<F> RetrievalFunnel for F
524where
525 F: Fn(bool) -> RedDBResult<NarrowedSlice>,
526{
527 fn funnel(&self, expanded: bool) -> RedDBResult<NarrowedSlice> {
528 self(expanded)
529 }
530}
531
532#[derive(Debug, Clone, PartialEq)]
536pub enum GroundingOutcome {
537 Grounded { slice: NarrowedSlice, refined: bool },
540 NoMatchingSources,
543}
544
545pub fn ground_with_refine<F: RetrievalFunnel>(funnel: &F) -> RedDBResult<GroundingOutcome> {
550 let first = funnel.funnel(false)?;
551 if !first.is_empty() {
552 return Ok(GroundingOutcome::Grounded {
553 slice: first,
554 refined: false,
555 });
556 }
557 let refined = funnel.funnel(true)?;
559 if !refined.is_empty() {
560 return Ok(GroundingOutcome::Grounded {
561 slice: refined,
562 refined: true,
563 });
564 }
565 Ok(GroundingOutcome::NoMatchingSources)
566}
567
568fn extract_json_object(raw: &str) -> Option<&str> {
571 let start = raw.find('{')?;
572 let end = raw.rfind('}')?;
573 if end <= start {
574 return None;
575 }
576 Some(&raw[start..=end])
577}
578
579#[cfg(test)]
580mod tests {
581 use super::*;
582
583 fn slice() -> NarrowedSlice {
584 NarrowedSlice {
585 collections: vec![
586 ScoredCollection {
587 collection: "travelers".to_string(),
588 score: 0.91,
589 columns: vec!["passport".to_string(), "name".to_string()],
590 },
591 ScoredCollection {
592 collection: "trips".to_string(),
593 score: 0.44,
594 columns: vec!["passport".to_string(), "city".to_string()],
595 },
596 ],
597 }
598 }
599
600 fn mock_model(plan_json: &'static str) -> impl PlannerModel {
603 move |_prompt: &str| Ok(plan_json.to_string())
604 }
605
606 #[test]
607 fn prompt_contains_only_narrowed_slice() {
608 let prompt = build_planner_prompt("who owns passport FDD-1?", &slice());
609 assert!(prompt.contains("travelers"));
610 assert!(prompt.contains("trips"));
611 assert!(prompt.contains("passport"));
612 assert!(!prompt.contains("secret_ledger"));
614 }
615
616 #[test]
617 fn parse_plan_reads_intent_and_query() {
618 let plan = parse_plan(
619 "{\"intent\":\"factual\",\"query\":\"SELECT * FROM travelers WHERE passport = 'FDD-1'\",\"rationale\":\"lookup\"}",
620 )
621 .unwrap();
622 assert_eq!(plan.intent, AskIntent::Factual);
623 assert_eq!(
624 plan.query.as_deref(),
625 Some("SELECT * FROM travelers WHERE passport = 'FDD-1'")
626 );
627 assert_eq!(plan.rationale, "lookup");
628 }
629
630 #[test]
631 fn parse_plan_tolerates_code_fences() {
632 let plan = parse_plan(
633 "```json\n{\"intent\":\"synthesis\",\"query\":null,\"rationale\":\"summary\"}\n```",
634 )
635 .unwrap();
636 assert_eq!(plan.intent, AskIntent::Synthesis);
637 assert!(plan.query.is_none());
638 }
639
640 #[test]
641 fn parse_plan_rejects_non_json() {
642 let err = parse_plan("the answer is 42").unwrap_err();
643 assert!(
644 err.to_string().contains("no JSON plan object"),
645 "got: {err}"
646 );
647 }
648
649 #[test]
650 fn parse_plan_rejects_unknown_intent() {
651 let err = parse_plan("{\"intent\":\"chitchat\",\"query\":null}").unwrap_err();
652 assert!(err.to_string().contains("unknown intent"), "got: {err}");
653 }
654
655 #[test]
656 fn factual_plan_routes_to_executable_read_only_candidate() {
657 let route = plan_and_route(
658 "who owns passport FDD-1?",
659 &slice(),
660 &mock_model(
661 "{\"intent\":\"factual\",\"query\":\"SELECT * FROM travelers WHERE passport = 'FDD-1'\",\"rationale\":\"lookup\"}",
662 ),
663 )
664 .unwrap();
665 match route.routing {
666 PlanRouting::Execute { candidate } => {
667 assert!(candidate.is_read_only());
668 assert_eq!(candidate.statement_type, "select");
669 }
670 other => panic!("expected Execute, got {other:?}"),
671 }
672 }
673
674 #[test]
675 fn factual_join_candidate_is_executable() {
676 let route = plan_and_route(
677 "which cities did the owner of passport FDD-1 visit?",
678 &slice(),
679 &mock_model(
680 "{\"intent\":\"factual\",\"query\":\"SELECT * FROM travelers JOIN trips ON travelers.passport = trips.passport WHERE travelers.passport = 'FDD-1'\",\"rationale\":\"join\"}",
681 ),
682 )
683 .unwrap();
684 assert!(matches!(route.routing, PlanRouting::Execute { .. }));
685 }
686
687 #[test]
688 fn factual_plan_with_mutating_candidate_is_refused_not_executed() {
689 let route = plan_and_route(
690 "delete traveler FDD-1",
691 &slice(),
692 &mock_model(
693 "{\"intent\":\"factual\",\"query\":\"DELETE FROM travelers WHERE passport = 'FDD-1'\",\"rationale\":\"oops\"}",
694 ),
695 )
696 .unwrap();
697 match route.routing {
698 PlanRouting::RefuseMutating { statement_type, .. } => {
699 assert_eq!(statement_type, "delete")
700 }
701 other => panic!("expected RefuseMutating, got {other:?}"),
702 }
703 }
704
705 #[test]
706 fn factual_plan_with_malformed_candidate_is_rejected() {
707 let err = plan_and_route(
708 "who owns passport FDD-1?",
709 &slice(),
710 &mock_model(
711 "{\"intent\":\"factual\",\"query\":\"this is not rql\",\"rationale\":\"x\"}",
712 ),
713 )
714 .unwrap_err();
715 assert!(
716 err.to_string().contains("invalid RQL candidate"),
717 "got: {err}"
718 );
719 }
720
721 #[test]
722 fn factual_plan_without_query_is_rejected() {
723 let err = plan_and_route(
724 "who owns passport FDD-1?",
725 &slice(),
726 &mock_model("{\"intent\":\"factual\",\"query\":null,\"rationale\":\"x\"}"),
727 )
728 .unwrap_err();
729 assert!(err.to_string().contains("no query candidate"), "got: {err}");
730 }
731
732 #[test]
733 fn synthesis_intent_routes_unsupported_in_this_slice() {
734 let route = plan_and_route(
735 "summarise yesterday's incidents",
736 &slice(),
737 &mock_model("{\"intent\":\"synthesis\",\"query\":null,\"rationale\":\"rag\"}"),
738 )
739 .unwrap();
740 assert_eq!(
741 route.routing,
742 PlanRouting::Unsupported {
743 intent: AskIntent::Synthesis
744 }
745 );
746 }
747
748 #[test]
749 fn synthesis_intent_carries_the_routed_intent_for_the_rag_fallthrough() {
750 let route = plan_and_route(
754 "summarise the incidents from yesterday",
755 &slice(),
756 &mock_model("{\"intent\":\"synthesis\",\"query\":null,\"rationale\":\"rag\"}"),
757 )
758 .unwrap();
759 match route.routing {
760 PlanRouting::Unsupported { intent } => {
761 assert_eq!(intent, AskIntent::Synthesis);
762 assert_eq!(intent.as_str(), "synthesis");
763 }
764 other => panic!("expected Unsupported{{synthesis}}, got {other:?}"),
765 }
766 }
767
768 #[test]
769 fn calculation_question_routes_factual_not_synthesis() {
770 let route = plan_and_route(
776 "how many trips went to Lisbon?",
777 &slice(),
778 &mock_model(
779 "{\"intent\":\"factual\",\"query\":\"SELECT COUNT(*) FROM trips WHERE city = 'Lisbon'\",\"rationale\":\"aggregate count\"}",
780 ),
781 )
782 .unwrap();
783 assert_ne!(
784 route.plan.intent,
785 AskIntent::Synthesis,
786 "a calculation must never classify as synthesis (the LLM must not invent numbers)"
787 );
788 match route.routing {
789 PlanRouting::Execute { candidate } => {
790 assert!(candidate.is_read_only());
791 assert_eq!(candidate.statement_type, "select");
792 }
793 other => {
794 panic!("a calculation must route to an executable factual candidate, got {other:?}")
795 }
796 }
797 }
798
799 #[test]
800 fn how_to_intent_routes_to_a_validated_suggestion_envelope() {
801 let route = plan_and_route(
805 "how would I capture events from orders into a queue?",
806 &slice(),
807 &mock_model(
808 "{\"intent\":\"how_to\",\"answer\":\"Create a queue and backfill events into it.\",\
809 \"suggestion\":[\
810 {\"rql\":\"CREATE QUEUE events_q WORK\",\"rationale\":\"the sink queue\"},\
811 {\"rql\":\"EVENTS BACKFILL orders TO events_q\",\"rationale\":\"seed history\"},\
812 {\"rql\":\"SELECT * FROM travelers WHERE passport = 'FDD-1'\",\"rationale\":\"inspect\"}\
813 ],\"rationale\":\"guide\"}",
814 ),
815 )
816 .unwrap();
817 match route.routing {
818 PlanRouting::Suggest { answer, suggestion } => {
819 assert!(answer.contains("Create a queue"));
820 assert_eq!(suggestion.len(), 3);
821 assert!(suggestion[0].mutating, "CREATE QUEUE is mutating DDL");
824 assert!(suggestion[1].mutating, "EVENTS BACKFILL is mutating");
825 assert!(!suggestion[2].mutating, "the SELECT is read-only");
826 assert_eq!(suggestion[2].statement_type, "select");
827 assert_eq!(suggestion[0].rationale, "the sink queue");
828 }
829 other => panic!("expected Suggest, got {other:?}"),
830 }
831 }
832
833 #[test]
834 fn how_to_suggestion_drops_unparseable_statements_never_returns_them_raw() {
835 let route = plan_and_route(
838 "how do I list travelers?",
839 &slice(),
840 &mock_model(
841 "{\"intent\":\"how_to\",\"answer\":\"Select from the collection.\",\
842 \"suggestion\":[\
843 {\"rql\":\"SELECT * FROM travelers WHERE passport = 'FDD-1'\"},\
844 {\"rql\":\"this is not rql at all\"},\
845 \"DELETE FROM travelers WHERE passport = 'FDD-1'\"\
846 ]}",
847 ),
848 )
849 .unwrap();
850 match route.routing {
851 PlanRouting::Suggest { suggestion, .. } => {
852 assert_eq!(suggestion.len(), 2, "the unparseable statement is dropped");
853 assert_eq!(suggestion[0].statement_type, "select");
854 assert_eq!(suggestion[1].statement_type, "delete");
856 assert!(suggestion[1].mutating);
857 for s in &suggestion {
858 assert!(
859 !s.rql.contains("not rql"),
860 "raw unparseable text must never survive"
861 );
862 }
863 }
864 other => panic!("expected Suggest, got {other:?}"),
865 }
866 }
867
868 #[test]
869 fn routing_distinguishes_how_to_from_factual_at_the_closure_model_seam() {
870 let factual = plan_and_route(
873 "how would I capture events into a queue?",
874 &slice(),
875 &mock_model(
876 "{\"intent\":\"factual\",\"query\":\"SELECT * FROM travelers WHERE passport = 'FDD-1'\",\"rationale\":\"x\"}",
877 ),
878 )
879 .unwrap();
880 assert!(matches!(factual.routing, PlanRouting::Execute { .. }));
881
882 let how_to = plan_and_route(
883 "how would I capture events into a queue?",
884 &slice(),
885 &mock_model(
886 "{\"intent\":\"how_to\",\"answer\":\"Use a queue.\",\"suggestion\":[\"CREATE QUEUE events_q WORK\"],\"rationale\":\"x\"}",
887 ),
888 )
889 .unwrap();
890 assert!(matches!(how_to.routing, PlanRouting::Suggest { .. }));
891 }
892
893 #[test]
898 fn steps_clause_is_clamped_to_the_config_cap_never_exceeding_it() {
899 assert_eq!(clamp_plan_steps(Some(9), 3), 3);
901 assert_eq!(clamp_plan_steps(Some(2), 3), 2);
903 assert_eq!(clamp_plan_steps(Some(3), 3), 3);
904 assert_eq!(clamp_plan_steps(None, 3), 3);
906 assert_eq!(clamp_plan_steps(Some(0), 3), 1);
908 assert_eq!(clamp_plan_steps(None, 0), 1);
909 }
910
911 #[test]
912 fn plan_budget_charges_until_exhausted_then_refuses_more() {
913 let mut budget = PlanBudget::new(Some(2), 3);
914 assert_eq!(budget.max_steps(), 2);
915 assert_eq!(budget.remaining(), 2);
916
917 assert!(budget.charge(PlanStep::RefineRetrieval).is_ok());
918 assert_eq!(budget.remaining(), 1);
919 assert!(budget.charge(PlanStep::Query).is_ok());
920 assert_eq!(budget.remaining(), 0);
921 assert_eq!(
922 budget.executed_steps(),
923 &[PlanStep::RefineRetrieval, PlanStep::Query]
924 );
925
926 let err = budget.charge(PlanStep::Query).unwrap_err();
928 assert_eq!(err.max_steps, 2);
929 assert_eq!(err.executed_steps, 2);
930 assert_eq!(err.attempted, PlanStep::Query);
931 assert_eq!(budget.executed_count(), 2);
933 }
934
935 #[test]
936 fn plan_budget_request_above_cap_cannot_buy_extra_steps() {
937 let mut budget = PlanBudget::new(Some(5), 1);
939 assert_eq!(budget.max_steps(), 1);
940 assert!(budget.charge(PlanStep::Query).is_ok());
941 assert!(budget.charge(PlanStep::RefineRetrieval).is_err());
942 }
943
944 fn scripted_funnel(
951 script: Vec<bool>,
952 ) -> (impl RetrievalFunnel, std::rc::Rc<std::cell::Cell<usize>>) {
953 let calls = std::rc::Rc::new(std::cell::Cell::new(0usize));
954 let calls_inner = calls.clone();
955 let funnel = move |_expanded: bool| -> RedDBResult<NarrowedSlice> {
956 let n = calls_inner.get();
957 calls_inner.set(n + 1);
958 let non_empty = script.get(n).copied().unwrap_or(false);
959 Ok(if non_empty {
960 slice()
961 } else {
962 NarrowedSlice::default()
963 })
964 };
965 (funnel, calls)
966 }
967
968 #[test]
969 fn grounding_succeeds_on_first_pass_without_refining() {
970 let (funnel, calls) = scripted_funnel(vec![true]);
971 let outcome = ground_with_refine(&funnel).unwrap();
972 assert_eq!(calls.get(), 1, "no refine when the first pass grounds");
973 match outcome {
974 GroundingOutcome::Grounded { refined, .. } => assert!(!refined),
975 other => panic!("expected Grounded, got {other:?}"),
976 }
977 }
978
979 #[test]
980 fn empty_first_pass_triggers_exactly_one_refine_retrieval() {
981 let (funnel, calls) = scripted_funnel(vec![false, true]);
983 let outcome = ground_with_refine(&funnel).unwrap();
984 assert_eq!(calls.get(), 2, "exactly one refine re-funnel");
985 match outcome {
986 GroundingOutcome::Grounded { refined, .. } => assert!(refined),
987 other => panic!("expected Grounded after refine, got {other:?}"),
988 }
989 }
990
991 #[test]
992 fn second_grounding_failure_returns_no_matching_sources_not_an_invented_answer() {
993 let (funnel, calls) = scripted_funnel(vec![false, false]);
996 let outcome = ground_with_refine(&funnel).unwrap();
997 assert_eq!(calls.get(), 2, "exactly one refine retry, then give up");
998 assert_eq!(outcome, GroundingOutcome::NoMatchingSources);
999 }
1000
1001 #[test]
1002 fn plan_summary_includes_intent_and_query() {
1003 let plan = AskPlan {
1004 intent: AskIntent::Factual,
1005 query: Some("SELECT * FROM travelers WHERE passport = 'FDD-1'".to_string()),
1006 answer: String::new(),
1007 suggestion: Vec::new(),
1008 rationale: "lookup".to_string(),
1009 };
1010 let summary = plan.summary();
1011 assert!(summary.contains("intent=factual"));
1012 assert!(summary.contains("SELECT * FROM travelers"));
1013 }
1014}