Coverage for astrocyte/pipeline/question_annotator.py: 0%

1"""PR2 D.5: LLM-backed question annotator. 

2 

3Replaces the regex-based ``_extract_question_entities`` heuristic with 

4a single LLM call that returns both: 

5 

61. **Entities** (proper + common nouns the question hinges on) — fed to 

7 the entity strategy. Catches lowercase nouns the regex misses 

8 ("pendant", "obesity", "dog treats", "taekwondo", "France"), which 

9 the open-domain failure analysis identified as the pivot for 

10 specific-fact questions. 

11 

122. **Date range** (start + end, anchored against a ``reference_date``) — 

13 fed to the temporal strategy as a *narrow* window. PR2-D.1-4 LME 

14 temporal-reasoning was 0% because we passed the temporal strategy 

15 the full conversation date range; without question-side date parsing 

16 it had nothing to filter on. This module gives it a real window. 

17 

18One LLM call per question (~$0.0002 at gpt-4o-mini prices). Both 

19fields are optional — the orchestrator handles missing entities or 

20date_range gracefully. 

21 

22See: 

23- docs/_design/recall.md §6 (recall pipeline, mode classifier slot) 

24- PR2-D.1-4 LME gate analysis (temporal-reasoning 0% root cause) 

25""" 

26 

27from __future__ import annotations 

28 

29import json 

30import logging 

31from dataclasses import dataclass 

32from datetime import datetime, timezone 

33from typing import TYPE_CHECKING 

34 

35if TYPE_CHECKING: 

36 from astrocyte.provider import LLMProvider 

37 

38from astrocyte.pipeline._json_tolerant import tolerant_json_loads_or_raise 

39from astrocyte.types import Message 

40 

41logger = logging.getLogger("astrocyte.pipeline.question_annotator") 

42 

43 

44@dataclass 

45class QuestionAnnotation: 

46 """One question's parsed structure for the section recall driver. 

47 

48 Both fields are optional. ``entities=[]`` skips the entity 

49 strategy; ``date_range=None`` skips the temporal strategy.""" 

50 

51 entities: list[str] 

52 date_range: tuple[datetime, datetime] | None 

53 

54 

55_PROMPT = """You are an analyst extracting search keys from a user's question about a long conversation transcript. 

56 

57Reference date (treat as "today" when resolving relative phrases): {reference_date} 

58 

59Return ONLY a JSON object with these keys: 

60 

61- "entities": array of strings — names, places, things, concepts, activities the question hinges on. Include BOTH proper nouns (people, places, brands) AND concrete common nouns (objects, activities, conditions). Skip stopwords (the, a, did, what, when, etc.) and tense markers. Aim for 1-6 entries. 

62 

63- "date_range": object with ISO-8601 "start" and "end" date strings, OR {{"start": null, "end": null}} when the question has no temporal anchor. Use the reference date to resolve relative phrases like "last week" or "two months ago". 

64 

65Examples (reference_date "22 October, 2023"): 

66 

67Q: "What did Caroline research?" 

68→ {{"entities": ["Caroline", "research"], "date_range": {{"start": null, "end": null}}}} 

69 

70Q: "In what country did Jolene's mother buy her the pendant?" 

71→ {{"entities": ["Jolene", "mother", "pendant", "country"], "date_range": {{"start": null, "end": null}}}} 

72 

73Q: "What are John's suspected health problems?" 

74→ {{"entities": ["John", "health problems"], "date_range": {{"start": null, "end": null}}}} 

75 

76Q: "What did Caroline say in May 2023?" 

77→ {{"entities": ["Caroline"], "date_range": {{"start": "2023-05-01", "end": "2023-05-31"}}}} 

78 

79Q: "Who did Maria have dinner with on May 3, 2023?" 

80→ {{"entities": ["Maria", "dinner"], "date_range": {{"start": "2023-05-03", "end": "2023-05-03"}}}} 

81 

82Q: "What was Caroline doing two months ago?" 

83→ {{"entities": ["Caroline"], "date_range": {{"start": "2023-08-01", "end": "2023-08-31"}}}} 

84 

85Q: "What temporary job did Jon take to cover expenses?" 

86→ {{"entities": ["Jon", "temporary job", "expenses"], "date_range": {{"start": null, "end": null}}}} 

87 

88Question: {question} 

89Output (JSON only):""" 

90 

91 

92async def annotate_question( 

93 provider: "LLMProvider", 

94 question: str, 

95 *, 

96 reference_date: str | None = None, 

97 model: str | None = None, 

98) -> QuestionAnnotation: 

99 """Single LLM call that returns entities + date_range. 

100 

101 ``reference_date`` is the human-readable date string from the 

102 conv_tree (e.g. "22 October, 2023"). When None, the prompt uses a 

103 placeholder; date phrases referencing "today" can't resolve, but 

104 explicit dates still parse. 

105 

106 Returns ``QuestionAnnotation(entities=[], date_range=None)`` on 

107 LLM failure or parse error — the orchestrator degrades gracefully 

108 (just skips the entity / temporal strategies for this question). 

109 """ 

110 if not question.strip(): 

111 return QuestionAnnotation(entities=[], date_range=None) 

112 

113 prompt = _PROMPT.format( 

114 question=question, 

115 reference_date=reference_date or "(unknown)", 

116 ) 

117 

118 try: 

119 completion = await provider.complete( 

120 messages=[Message(role="user", content=prompt)], 

121 model=model, 

122 max_tokens=200, 

123 temperature=0.0, 

124 response_format={"type": "json_object"}, 

125 ) 

126 except Exception as exc: # noqa: BLE001 — annotator failure shouldn't tank a question 

127 logger.warning( 

128 "annotate_question: LLM call failed for q=%r: %s: %s", 

129 question[:80], 

130 type(exc).__name__, 

131 exc, 

132 ) 

133 return QuestionAnnotation(entities=[], date_range=None) 

134 

135 try: 

136 parsed = tolerant_json_loads_or_raise(completion.text) 

137 except json.JSONDecodeError: 

138 logger.warning( 

139 "annotate_question: JSON parse failed for q=%r; raw=%r", 

140 question[:80], 

141 completion.text[:120], 

142 ) 

143 return QuestionAnnotation(entities=[], date_range=None) 

144 

145 raw_entities = parsed.get("entities") or [] 

146 entities: list[str] = [] 

147 seen: set[str] = set() 

148 for e in raw_entities: 

149 if not isinstance(e, str): 

150 continue 

151 clean = e.strip() 

152 if not clean: 

153 continue 

154 key = clean.casefold() 

155 if key in seen: 

156 continue 

157 seen.add(key) 

158 entities.append(clean) 

159 

160 date_range = _parse_iso_range(parsed.get("date_range")) 

161 return QuestionAnnotation(entities=entities, date_range=date_range) 

162 

163 

164def _parse_iso_range(raw) -> tuple[datetime, datetime] | None: 

165 """Coerce ``{"start": "YYYY-MM-DD", "end": "YYYY-MM-DD"}`` → tz-aware 

166 UTC datetime tuple. Returns ``None`` on missing / malformed input.""" 

167 if not isinstance(raw, dict): 

168 return None 

169 start_s = raw.get("start") 

170 end_s = raw.get("end") 

171 if not start_s or not end_s: 

172 return None 

173 try: 

174 start = datetime.strptime(start_s, "%Y-%m-%d").replace(tzinfo=timezone.utc) 

175 end = datetime.strptime(end_s, "%Y-%m-%d").replace( 

176 hour=23, 

177 minute=59, 

178 second=59, 

179 tzinfo=timezone.utc, 

180 ) 

181 except (ValueError, TypeError): 

182 return None 

183 if end < start: 

184 # LLM occasionally swaps; tolerate. 

185 start, end = end, start 

186 return (start, end)