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

1"""M11.1 — structured event-time extraction per section. 

2 

3Mirrors Hindsight's per-fact ``occurred_start`` / ``occurred_end`` 

4columns on ``memory_units``. At retain time, for each section, ask 

5the LLM to identify the most-prominent event the section describes 

6and emit ISO-8601 start (and optional end) timestamps. 

7 

8Why this differs from ``session_date`` (already on every section): 

9 

10- ``session_date`` is when the conversation session HAPPENED 

11 (e.g. May 8, 2023 — when the user typed the message) 

12- ``occurred_start`` is when the discussed EVENT happened 

13 (e.g. May 7 — "yesterday I went to the doctor") 

14 

15LME temporal-reasoning failures all share the same shape: the picker 

16finds the right SESSION but the synth uses the SESSION date instead 

17of the EVENT date. With ``occurred_start`` populated, 

18:func:`~astrocyte.pipeline.temporal_arithmetic.find_event_date` can 

19return the canonical event date directly. 

20 

21Relative phrases ("yesterday", "last week", "3 days ago") are anchored 

22against the section's ``session_date`` at extraction time so the 

23output is always an absolute ISO timestamp. 

24 

25See: 

26- ``docs/_design/recall.md`` §13 (M10 close-out) + §14 (M11 plan) 

27- ``hindsight/hindsight-api-slim/hindsight_api/engine/retain/fact_extraction.py`` 

28 for the canonical Hindsight pattern at memory_unit grain. 

29""" 

30 

31from __future__ import annotations 

32 

33import json 

34import logging 

35from datetime import datetime 

36from typing import TYPE_CHECKING 

37 

38from astrocyte.pipeline._json_tolerant import tolerant_json_loads_or_raise 

39from astrocyte.types import Message 

40 

41if TYPE_CHECKING: 

42 from astrocyte.provider import LLMProvider 

43 from astrocyte.types import PageIndexSection 

44 

45_logger = logging.getLogger("astrocyte.pipeline.section_event_extraction") 

46 

47 

48_EXTRACT_PROMPT = """\ 

49You are extracting the structured EVENT TIME from one section of a \ 

50conversation transcript. Your output drives query-time temporal \ 

51arithmetic — questions like "how many weeks ago did I visit the \ 

52doctor?" need the canonical date of the doctor visit, NOT the date \ 

53the user mentioned it. 

54 

55The section's conversation date is ``{session_date}``. Anchor any \ 

56relative time phrase ("yesterday", "last week", "3 days ago", "last \ 

57month") against THIS date. Output absolute ISO-8601 timestamps. 

58 

59Output a JSON object with EXACTLY these fields: 

60- "occurred_start": ISO-8601 timestamp of when the most-prominent \ 

61discussed event began. ``null`` if the section is generic chit-chat \ 

62with no specific event. 

63- "occurred_end": ISO-8601 timestamp of when that event ended. \ 

64``null`` if it's a single-day or instantaneous event. 

65- "event_description": 3-6 word description of the event (for \ 

66provenance / debugging). ``null`` if no specific event. 

67 

68Rules: 

69- "Yesterday I went to the doctor" with session_date=2023-05-08 → \ 

70 ``{{"occurred_start": "2023-05-07", "occurred_end": null, \ 

71"event_description": "doctor visit"}}`` 

72- "We had a wedding two Saturdays ago" with session_date=2023-05-15 \ 

73 → ``{{"occurred_start": "2023-05-06", "occurred_end": null, \ 

74"event_description": "wedding"}}`` 

75- "Spent last weekend camping" with session_date=2023-05-22 \ 

76 (Monday) → ``{{"occurred_start": "2023-05-20", "occurred_end": \ 

77"2023-05-21", "event_description": "weekend camping"}}`` 

78- "Trip from May 3-15" → ``{{"occurred_start": "2023-05-03", \ 

79"occurred_end": "2023-05-15", "event_description": "trip"}}`` 

80- Generic chit-chat ("How are you today?", recipe discussion with \ 

81no specific past event) → ``{{"occurred_start": null, "occurred_end": \ 

82null, "event_description": null}}`` 

83 

84When the section discusses MULTIPLE events, pick the most-prominent \ 

85one (the one the user is asking about / spent the most time on). Do \ 

86NOT try to list multiple events — one per section. 

87 

88If a relative phrase is ambiguous ("recently", "a while back"), \ 

89return ``null`` rather than guessing. 

90 

91OUTPUT MUST BE VALID JSON. No prose around it. 

92 

93Section content: 

94{section_text} 

95""" 

96 

97 

98def _parse_iso_date(s: str | None) -> datetime | None: 

99 if not s or not isinstance(s, str): 

100 return None 

101 try: 

102 # Accept both "2023-05-07" and "2023-05-07T12:00:00" forms. 

103 if "T" in s or ":" in s: 

104 return datetime.fromisoformat(s.replace("Z", "+00:00")) 

105 return datetime.fromisoformat(s) 

106 except (ValueError, TypeError) as exc: 

107 _logger.debug("section_event_extraction: bad ISO date %r: %s", s, exc) 

108 return None 

109 

110 

111async def extract_event_date_for_section( 

112 provider: "LLMProvider", 

113 section: "PageIndexSection", 

114 section_text: str, 

115 *, 

116 model: str | None = None, 

117) -> tuple[datetime | None, datetime | None]: 

118 """One LLM call → ``(occurred_start, occurred_end)`` for this section. 

119 

120 Returns ``(None, None)`` when: 

121 - the section is generic chit-chat with no specific event 

122 - the LLM output fails to parse 

123 - the LLM declines to commit a date (ambiguous phrase) 

124 

125 Caller persists the returned dates via 

126 :meth:`PageIndexStore.save_section_event_dates`. 

127 """ 

128 text = section_text.strip() 

129 if not text: 

130 return None, None 

131 sess_iso = section.session_date.strftime("%Y-%m-%d") if section.session_date is not None else "unknown" 

132 msg = _EXTRACT_PROMPT.format( 

133 session_date=sess_iso, 

134 section_text=text[:6000], # ~1500 tokens cap 

135 ) 

136 try: 

137 completion = await provider.complete( 

138 messages=[Message(role="user", content=msg)], 

139 model=model, 

140 max_tokens=200, 

141 temperature=0.0, 

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

143 ) 

144 except Exception as exc: # noqa: BLE001 

145 _logger.warning( 

146 "section_event_extraction: LLM call failed doc=%s line=%d: %s", 

147 section.document_id, 

148 section.line_num, 

149 exc, 

150 ) 

151 return None, None 

152 try: 

153 data = tolerant_json_loads_or_raise(completion.text) 

154 except json.JSONDecodeError: 

155 _logger.warning( 

156 "section_event_extraction: JSON parse failed doc=%s line=%d text=%r", 

157 section.document_id, 

158 section.line_num, 

159 completion.text[:200], 

160 ) 

161 return None, None 

162 start = _parse_iso_date(data.get("occurred_start")) 

163 end = _parse_iso_date(data.get("occurred_end")) 

164 return start, end