Coverage for astrocyte/pipeline/query_plan.py: 97%

1"""Lightweight query planning for memory synthesis. 

2 

3This module deliberately stays deterministic and cheap. It identifies the 

4question shapes that need broader evidence assembly than a single fact lookup: 

5aggregate/list questions, temporal questions, inference/open-domain questions, 

6and adversarial/unanswerable-looking prompts. 

7""" 

8 

9from __future__ import annotations 

10 

11import re 

12from dataclasses import dataclass 

13 

14from astrocyte.pipeline.query_intent import QueryIntent, classify_query_intent 

15from astrocyte.pipeline.temporal import temporal_guidance_for_query 

16 

17 

18@dataclass(frozen=True) 

19class QueryPlan: 

20 """Derived plan for recall and reflect behavior.""" 

21 

22 intent: QueryIntent 

23 needs_temporal_reasoning: bool = False 

24 needs_multi_hop_synthesis: bool = False 

25 needs_aggregate_answer: bool = False 

26 needs_inference: bool = False 

27 may_be_adversarial: bool = False 

28 prompt_variant: str | None = None 

29 recall_max_results: int = 30 

30 reflect_rank_limit: int = 12 

31 reflect_expand_limit: int = 18 

32 guidance: str | None = None 

33 

34 

35_AGGREGATE_RE = re.compile( 

36 r"\b(" 

37 r"what\s+(activities|events|books|items|types|ways|symbols|instruments|artists|bands)|" 

38 r"how\s+many\s+times|" 

39 r"in\s+what\s+ways|" 

40 r"what\s+has\s+\w+\s+(painted|bought|read|done)" 

41 r")\b", 

42 re.IGNORECASE, 

43) 

44_INFERENCE_RE = re.compile( 

45 r"\b(would|likely|probably|considered|interested\s+in|prefer|leaning|pursue)\b", 

46 re.IGNORECASE, 

47) 

48_ADVERSARIAL_RE = re.compile( 

49 r"\b(type\s+of|what\s+type|did\s+\w+\s+make|who\s+is\s+\w+\s+a\s+fan\s+of|what\s+happened\s+to)\b", 

50 re.IGNORECASE, 

51) 

52_MULTI_HOP_RE = re.compile( 

53 r"\b(across|combine|relationship|support|participat(?:e|ed|ing)|activities|events|ways|types|how\s+many)\b", 

54 re.IGNORECASE, 

55) 

56 

57 

58def build_query_plan(query: str) -> QueryPlan: 

59 """Classify query shape and choose retrieval/synthesis guidance.""" 

60 

61 query_text = query or "" 

62 intent = classify_query_intent(query_text).intent 

63 temporal_guidance = temporal_guidance_for_query(query_text) 

64 needs_temporal = intent == QueryIntent.TEMPORAL or temporal_guidance is not None 

65 needs_aggregate = bool(_AGGREGATE_RE.search(query_text)) 

66 needs_inference = bool(_INFERENCE_RE.search(query_text)) 

67 may_be_adversarial = bool(_ADVERSARIAL_RE.search(query_text)) 

68 needs_multi_hop = needs_aggregate or bool(_MULTI_HOP_RE.search(query_text)) 

69 

70 prompt_variant: str | None = None 

71 if needs_temporal: 

72 prompt_variant = "temporal_aware" 

73 elif needs_aggregate or needs_multi_hop: 

74 prompt_variant = "grounded_synthesis" 

75 elif needs_inference: 

76 prompt_variant = "evidence_inference" 

77 elif may_be_adversarial: 

78 prompt_variant = "evidence_strict" 

79 

80 recall_max_results = 30 

81 reflect_rank_limit = 12 

82 reflect_expand_limit = 18 

83 if needs_aggregate or needs_multi_hop: 

84 recall_max_results = 40 

85 reflect_rank_limit = 18 

86 reflect_expand_limit = 26 

87 if needs_inference: 

88 recall_max_results = max(recall_max_results, 36) 

89 reflect_rank_limit = max(reflect_rank_limit, 16) 

90 reflect_expand_limit = max(reflect_expand_limit, 24) 

91 

92 guidance_parts: list[str] = [] 

93 if needs_aggregate: 

94 guidance_parts.append( 

95 "Aggregate/list question: scan all provided memories for distinct matching facts; " 

96 "do not stop after the first plausible memory." 

97 ) 

98 if needs_multi_hop: 

99 guidance_parts.append( 

100 "Multi-hop question: combine directly related facts across memories when they share " 

101 "the same person, event, object, or timeframe." 

102 ) 

103 if needs_inference: 

104 guidance_parts.append( 

105 "Inference question: answer with calibrated language only when the memories directly support it." 

106 ) 

107 if may_be_adversarial: 

108 guidance_parts.append( 

109 "Adversarial check: verify the named person and premise match the memories before answering." 

110 ) 

111 if temporal_guidance: 

112 guidance_parts.append(temporal_guidance) 

113 

114 return QueryPlan( 

115 intent=intent, 

116 needs_temporal_reasoning=needs_temporal, 

117 needs_multi_hop_synthesis=needs_multi_hop, 

118 needs_aggregate_answer=needs_aggregate, 

119 needs_inference=needs_inference, 

120 may_be_adversarial=may_be_adversarial, 

121 prompt_variant=prompt_variant, 

122 recall_max_results=recall_max_results, 

123 reflect_rank_limit=reflect_rank_limit, 

124 reflect_expand_limit=reflect_expand_limit, 

125 guidance="\n".join(guidance_parts) if guidance_parts else None, 

126 )