Coverage for astrocyte/pipeline/fact_extraction.py: 84%

1"""Structured fact extraction at retain time. 

2 

3Replaces the legacy two-pass approach (text chunking → separate 

4entity-extraction LLM call → separate fact-causal LLM call) with a 

5single LLM call that produces a list of structured facts. Each fact 

6carries five dimensions plus its embedded entities and intra-batch 

7causal relations: 

8 

9- ``what`` — the core factual content (1–2 sentences) 

10- ``when`` — temporal expression in natural language plus optional 

11 ISO ``occurred_start`` / ``occurred_end`` 

12- ``where`` — location or N/A 

13- ``who`` — people involved + relationships 

14- ``why`` — context or N/A 

15- ``fact_type`` — ``"world"`` (objective external fact) vs 

16 ``"experience"`` (first-person event) 

17- ``entities`` — named entities embedded in the fact 

18- ``causal_relations`` — directional ``caused_by`` references to other 

19 facts in the SAME batch by index 

20 

21Why this matters: every downstream signal (semantic embedding, cross- 

22encoder rerank, link expansion, reflect synthesis) gets cleaner inputs 

23when ``who`` and ``when`` are first-class fields rather than buried in 

24free text. Multi-hop and temporal questions especially benefit because 

25the structured fields can be filtered / matched directly. 

26 

27Cost: single LLM call per retain text replaces two (entity + causal). 

28Net cost approximately equal; output substantially richer. 

29 

30The module also exposes :func:`materialize_facts` which converts a 

31list of :class:`ExtractedFact` objects into the three artefacts the 

32orchestrator needs to persist: 

33 

34- :class:`VectorItem` per fact (replaces chunk-based VectorItems) 

35- :class:`Entity` per unique entity mentioned across the batch 

36- :class:`MemoryLink` per causal relation (memory-to-memory edges) 

37""" 

38 

39from __future__ import annotations 

40 

41import asyncio 

42import json 

43import logging 

44import re 

45import uuid 

46from dataclasses import dataclass, field 

47from datetime import datetime, timezone 

48 

49from astrocyte.types import ( 

50 Entity, 

51 MemoryLink, 

52 Message, 

53 VectorItem, 

54) 

55 

56_logger = logging.getLogger("astrocyte.fact_extraction") 

57 

58 

59# --------------------------------------------------------------------------- 

60# Data shapes 

61# --------------------------------------------------------------------------- 

62 

63 

64@dataclass 

65class FactEntity: 

66 """An entity mentioned within a fact.""" 

67 

68 name: str 

69 entity_type: str = "OTHER" # PERSON, ORG, LOCATION, CONCEPT, OTHER 

70 

71 

72@dataclass 

73class FactCausalRelation: 

74 """``caused_by`` reference from THIS fact to another fact in the 

75 same batch. Hindsight-style: target_index < source_index is 

76 enforced upstream when building MemoryLinks (we drop self-loops 

77 and out-of-range references).""" 

78 

79 target_fact_index: int 

80 strength: float = 1.0 # 0..1, defaults to strong 

81 

82 

83@dataclass 

84class ExtractedFact: 

85 """A single structured fact extracted from retain text. 

86 

87 The 5-dimension schema (what/when/where/who/why) materializes the 

88 semantic content the LLM extracted, plus embedded entities and 

89 causal links between facts in the same batch. 

90 """ 

91 

92 what: str # core content 

93 when: str = "N/A" # natural-language time expression 

94 where: str = "N/A" 

95 who: str = "N/A" 

96 why: str = "N/A" 

97 fact_type: str = "experience" # "world" | "experience" 

98 occurred_start: datetime | None = None 

99 occurred_end: datetime | None = None 

100 entities: list[FactEntity] = field(default_factory=list) 

101 causal_relations: list[FactCausalRelation] = field(default_factory=list) 

102 # Last, so positional construction stays compatible. 

103 #: False when ``fact_type`` is the default rather than the model's 

104 #: answer (its metadata entry was missing or invalid). Storage must not 

105 #: treat a default as a classification. 

106 fact_type_classified: bool = True 

107 

108 

109# --------------------------------------------------------------------------- 

110# Prompt 

111# --------------------------------------------------------------------------- 

112 

113 

114_VERBATIM_SYSTEM_PROMPT = """\ 

115You enrich PRE-CHUNKED text with structured metadata. Do NOT rewrite \ 

116or paraphrase the chunk text — the caller will use the original chunk \ 

117text verbatim. Your job is to produce per-chunk metadata. 

118 

119Output a JSON object: {"facts": [...]}. The facts list MUST be the \ 

120same length as the input chunks list, in the same order. For each \ 

121chunk, produce ONE entry with: 

122- "chunk_index" (integer): the [index] of the chunk this entry \ 

123describes, copied exactly from the input. 

124- "when" (string, default "N/A"): natural-language time expression \ 

125present in this chunk; "N/A" otherwise. 

126- "where" (string, default "N/A"): location. 

127- "who" (string, default "N/A"): people involved. 

128- "why" (string, default "N/A"): reason / motivation if explicit. 

129- "fact_type" ("world" | "experience"): classify the chunk. 

130- "occurred_start" (ISO 8601 string or null): resolve "when" to an \ 

131absolute date when possible; null otherwise. 

132- "occurred_end" (ISO 8601 string or null): instant or end of range. 

133- "entities" (list): each entity mentioned in this chunk, as \ 

134{"name": str, "entity_type": str}. Types: PERSON, ORG, LOCATION, \ 

135PRODUCT, CONCEPT, OTHER. 

136 

137Rules: 

1381. Output exactly one entry per input chunk, in the same order, each \ 

139tagged with its chunk_index. Never merge or skip chunks. 

1402. Don't invent. Use "N/A" / null / [] for absent metadata. 

1413. Output JSON only. 

142""" 

143 

144 

145# JSON Schema for OpenAI structured outputs (Phase 2 of cost-control 

146# port). When the provider supports ``response_format=json_schema`` (set 

147# at the call site), the model is decode-time constrained to this shape 

148# and malformed-JSON parse failures become impossible. Strict mode 

149# requires every property in ``required`` and ``additionalProperties: 

150# false`` at every level — keep it that way when extending. 

151# 

152# ``causal_relations`` is intentionally omitted from the schema: 

153# ``causal_links.enabled`` defaults to false in our research configs, 

154# and OpenAI strict-mode JSON schema enforces required-everywhere which 

155# makes optional cross-chunk index references awkward. If we re-enable 

156# causal extraction, add a separate schema variant rather than bolting 

157# it on here. 

158_VERBATIM_JSON_SCHEMA: dict = { 

159 "name": "verbatim_facts", 

160 "strict": True, 

161 "schema": { 

162 "type": "object", 

163 "additionalProperties": False, 

164 "required": ["facts"], 

165 "properties": { 

166 "facts": { 

167 "type": "array", 

168 "items": { 

169 "type": "object", 

170 "additionalProperties": False, 

171 "required": [ 

172 "chunk_index", 

173 "when", 

174 "where", 

175 "who", 

176 "why", 

177 "fact_type", 

178 "occurred_start", 

179 "occurred_end", 

180 "entities", 

181 ], 

182 "properties": { 

183 "chunk_index": {"type": "integer"}, 

184 "when": {"type": "string"}, 

185 "where": {"type": "string"}, 

186 "who": {"type": "string"}, 

187 "why": {"type": "string"}, 

188 "fact_type": { 

189 "type": "string", 

190 "enum": ["world", "experience"], 

191 }, 

192 "occurred_start": {"type": ["string", "null"]}, 

193 "occurred_end": {"type": ["string", "null"]}, 

194 "entities": { 

195 "type": "array", 

196 "items": { 

197 "type": "object", 

198 "additionalProperties": False, 

199 "required": ["name", "entity_type"], 

200 "properties": { 

201 "name": {"type": "string"}, 

202 "entity_type": { 

203 "type": "string", 

204 "enum": [ 

205 "PERSON", 

206 "ORG", 

207 "LOCATION", 

208 "PRODUCT", 

209 "CONCEPT", 

210 "OTHER", 

211 ], 

212 }, 

213 }, 

214 }, 

215 }, 

216 }, 

217 }, 

218 } 

219 }, 

220 }, 

221} 

222 

223 

224def _build_verbatim_user_prompt( 

225 chunk_texts: list[str], 

226 event_date: datetime | None = None, 

227) -> str: 

228 lines = [] 

229 if event_date is not None: 

230 lines.append(f"Reference date for relative time expressions: {event_date.isoformat()}") 

231 lines.append("") 

232 lines.append(f"Chunks ({len(chunk_texts)} total, indexed):") 

233 for i, text in enumerate(chunk_texts): 

234 snippet = text.strip() 

235 if len(snippet) > 800: 

236 snippet = snippet[:797] + "..." 

237 lines.append(f"[{i}] {snippet}") 

238 lines.append("") 

239 lines.append("Per-chunk metadata (JSON, one entry per chunk, each with its chunk_index):") 

240 return "\n".join(lines) 

241 

242 

243# --------------------------------------------------------------------------- 

244# Per-chunk parallel verbatim extraction (Phase 3 of cost-control port) 

245# --------------------------------------------------------------------------- 

246# 

247# Splits the verbatim metadata extraction into one LLM call per chunk, 

248# dispatched in parallel via ``asyncio.gather`` with a ``Semaphore`` 

249# bounding in-flight calls per session. Trade-offs vs the batched 

250# ``extract_facts_verbatim``: 

251# 

252# - Output tokens per call are tiny (~150 tokens for one entry vs ~150 * 

253# N for N-chunk batched), so per-call latency drops sharply on 

254# gpt-4o-mini (latency is roughly linear in output tokens). 

255# - More API calls overall, but they run in parallel, so wall time drops. 

256# - Cross-chunk ``causal_relations`` index references are dropped — they 

257# require co-located chunks in one prompt to reference. Acceptable 

258# default because ``causal_links.enabled`` is false in our research 

259# configs; if you need causal extraction, use the batched path. 

260 

261_VERBATIM_SINGLE_SYSTEM_PROMPT = """\ 

262You enrich a single PRE-CHUNKED text excerpt with structured metadata. \ 

263Do NOT rewrite or paraphrase the chunk — produce metadata only. 

264 

265Output a JSON object with these fields: 

266- "when" (string, default "N/A"): natural-language time expression in \ 

267the chunk; "N/A" otherwise. 

268- "where" (string, default "N/A"): location. 

269- "who" (string, default "N/A"): people involved. 

270- "why" (string, default "N/A"): reason / motivation if explicit. 

271- "fact_type" ("world" | "experience"): classify the chunk. 

272- "occurred_start" (ISO 8601 string or null): resolve "when" to an \ 

273absolute date when possible; null otherwise. 

274- "occurred_end" (ISO 8601 string or null): instant or end of range. 

275- "entities" (list): each entity as {"name": str, "entity_type": str}. \ 

276Types: PERSON, ORG, LOCATION, PRODUCT, CONCEPT, OTHER. 

277 

278Don't invent. Use "N/A" / null / [] for absent metadata. Output JSON \ 

279only. 

280""" 

281 

282 

283_VERBATIM_SINGLE_JSON_SCHEMA: dict = { 

284 "name": "verbatim_chunk_metadata", 

285 "strict": True, 

286 "schema": { 

287 "type": "object", 

288 "additionalProperties": False, 

289 "required": [ 

290 "when", 

291 "where", 

292 "who", 

293 "why", 

294 "fact_type", 

295 "occurred_start", 

296 "occurred_end", 

297 "entities", 

298 ], 

299 "properties": { 

300 "when": {"type": "string"}, 

301 "where": {"type": "string"}, 

302 "who": {"type": "string"}, 

303 "why": {"type": "string"}, 

304 "fact_type": { 

305 "type": "string", 

306 "enum": ["world", "experience"], 

307 }, 

308 "occurred_start": {"type": ["string", "null"]}, 

309 "occurred_end": {"type": ["string", "null"]}, 

310 "entities": { 

311 "type": "array", 

312 "items": { 

313 "type": "object", 

314 "additionalProperties": False, 

315 "required": ["name", "entity_type"], 

316 "properties": { 

317 "name": {"type": "string"}, 

318 "entity_type": { 

319 "type": "string", 

320 "enum": [ 

321 "PERSON", 

322 "ORG", 

323 "LOCATION", 

324 "PRODUCT", 

325 "CONCEPT", 

326 "OTHER", 

327 ], 

328 }, 

329 }, 

330 }, 

331 }, 

332 }, 

333 }, 

334} 

335 

336 

337def _build_verbatim_single_user_prompt( 

338 chunk: str, 

339 event_date: datetime | None = None, 

340) -> str: 

341 lines = [] 

342 if event_date is not None: 

343 lines.append(f"Reference date for relative time expressions: {event_date.isoformat()}") 

344 lines.append("") 

345 lines.append("Chunk:") 

346 snippet = chunk.strip() 

347 # Larger cap than batched (which truncates at 800 chars per chunk 

348 # to fit many in one prompt) — single-chunk path can afford the 

349 # full chunk text up to chunk_max_size. 

350 if len(snippet) > 8000: 

351 snippet = snippet[:7997] + "..." 

352 lines.append(snippet) 

353 lines.append("") 

354 lines.append("Metadata (JSON):") 

355 return "\n".join(lines) 

356 

357 

358@dataclass(slots=True, frozen=True) 

359class _ChunkResult: 

360 """Result of one parallel verbatim-extraction task. 

361 

362 Internal record paired with each input chunk. ``raw`` is the 

363 parsed metadata dict from the LLM (vendor-shaped JSON, no fixed 

364 schema beyond the constraints in ``_VERBATIM_SINGLE_JSON_SCHEMA``) 

365 or ``{}`` when extraction fell through after retries. 

366 """ 

367 

368 idx: int 

369 raw: dict 

370 

371 

372class _VerbatimChunkError(Exception): 

373 """Raised by ``_extract_one_chunk_verbatim_attempt`` on any failure 

374 that should be retried — LLM-side exception, malformed JSON, etc. 

375 Internal to the retry wrapper; never escapes the public surface.""" 

376 

377 

378async def _extract_one_chunk_verbatim_attempt( 

379 chunk: str, 

380 llm_provider, 

381 *, 

382 event_date: datetime | None = None, 

383) -> dict: 

384 """Single attempt at extracting verbatim metadata for ONE chunk. 

385 Raises :class:`_VerbatimChunkError` on any failure so the outer 

386 retry loop can back off and try again.""" 

387 messages_in = [ 

388 Message(role="system", content=_VERBATIM_SINGLE_SYSTEM_PROMPT), 

389 Message( 

390 role="user", 

391 content=_build_verbatim_single_user_prompt(chunk, event_date), 

392 ), 

393 ] 

394 try: 

395 completion = await llm_provider.complete( 

396 messages_in, 

397 # Per-chunk output is one metadata object; tiny in absolute 

398 # terms, but at chunk_max_size=2048 with dialogue-dense 

399 # chunks the LLM can emit dozens of entities and the JSON 

400 # blows past 512 tokens, getting truncated mid-emit. 

401 # Structured outputs can't recover from a hard max_tokens 

402 # cap, so we get parse failures even with json_schema mode. 

403 # 2048 gives ~10× headroom and is still tiny vs the 

404 # batched path's 4096 ceiling. 

405 max_tokens=2048, 

406 temperature=0.0, 

407 response_format={ 

408 "type": "json_schema", 

409 "json_schema": _VERBATIM_SINGLE_JSON_SCHEMA, 

410 }, 

411 ) 

412 except TypeError: 

413 # Provider's complete() pre-dates response_format kwarg — 

414 # retry once without (still JSON-parseable thanks to the 

415 # prompt). A real failure on this fallback raises through 

416 # the surrounding except. 

417 try: 

418 completion = await llm_provider.complete( 

419 messages_in, 

420 max_tokens=2048, 

421 temperature=0.0, 

422 ) 

423 except Exception as exc: 

424 raise _VerbatimChunkError(f"LLM call failed: {exc}") from exc 

425 except Exception as exc: 

426 raise _VerbatimChunkError(f"LLM call failed: {exc}") from exc 

427 parsed = _parse_json_object(completion.text) 

428 if parsed is None: 

429 raise _VerbatimChunkError("malformed JSON response") 

430 return parsed 

431 

432 

433async def _extract_one_chunk_verbatim( 

434 chunk: str, 

435 llm_provider, 

436 *, 

437 event_date: datetime | None = None, 

438 max_retries: int = 3, 

439 base_retry_delay: float = 2.0, 

440) -> dict: 

441 """Extract verbatim metadata for ONE chunk with retries. 

442 

443 Phase 4 of cost-control port — wraps 

444 :func:`_extract_one_chunk_verbatim_attempt` with the same retry 

445 policy Hindsight uses for ``_extract_chunk_with_retry``: up to 

446 ``max_retries`` attempts with exponential backoff 

447 (``base_retry_delay`` × 2^attempt). On final exhaustion returns 

448 ``{}`` so the caller falls through to a metadata-less 

449 ExtractedFact (chunk text still preserved). 

450 

451 Tests that need to exercise the bare single-attempt failure path 

452 can pass ``max_retries=1``. 

453 """ 

454 if not chunk.strip(): 

455 return {} 

456 last_exc: BaseException | None = None 

457 for attempt in range(max(1, max_retries)): 

458 try: 

459 return await _extract_one_chunk_verbatim_attempt( 

460 chunk, 

461 llm_provider, 

462 event_date=event_date, 

463 ) 

464 except _VerbatimChunkError as exc: 

465 last_exc = exc 

466 if attempt < max_retries - 1: 

467 delay = base_retry_delay * (2**attempt) 

468 _logger.warning( 

469 "fact_extraction (verbatim/parallel) chunk attempt %d/%d failed (%s); retrying in %.0fs", 

470 attempt + 1, 

471 max_retries, 

472 exc, 

473 delay, 

474 ) 

475 await asyncio.sleep(delay) 

476 _logger.warning( 

477 "fact_extraction (verbatim/parallel) chunk exhausted %d retries (%s); " 

478 "falling through to metadata-less ExtractedFact", 

479 max_retries, 

480 last_exc, 

481 ) 

482 return {} 

483 

484 

485def _build_extracted_fact_from_raw( 

486 chunk: str, 

487 raw: dict, 

488) -> ExtractedFact: 

489 """Convert one parsed metadata dict + the original chunk text into 

490 an :class:`ExtractedFact`. Shared by the batched and per-chunk 

491 parallel paths.""" 

492 entities: list[FactEntity] = [] 

493 for ent in raw.get("entities") or []: 

494 if not isinstance(ent, dict): 

495 continue 

496 name = str(ent.get("name") or "").strip() 

497 if not name: 

498 continue 

499 etype = str(ent.get("entity_type") or "OTHER").strip().upper() or "OTHER" 

500 entities.append(FactEntity(name=name, entity_type=etype)) 

501 ftype = str(raw.get("fact_type") or "").strip().lower() 

502 classified = ftype in {"world", "experience"} 

503 if not classified: 

504 ftype = "experience" 

505 return ExtractedFact( 

506 what=chunk, 

507 when=str(raw.get("when") or "N/A").strip() or "N/A", 

508 where=str(raw.get("where") or "N/A").strip() or "N/A", 

509 who=str(raw.get("who") or "N/A").strip() or "N/A", 

510 why=str(raw.get("why") or "N/A").strip() or "N/A", 

511 fact_type=ftype, 

512 fact_type_classified=classified, 

513 occurred_start=_parse_iso_datetime(raw.get("occurred_start")), 

514 occurred_end=_parse_iso_datetime(raw.get("occurred_end")), 

515 entities=entities, 

516 # Per-chunk parallel path drops cross-chunk causal_relations. 

517 # Caller using this helper from the batched path can override. 

518 causal_relations=[], 

519 ) 

520 

521 

522async def extract_facts_verbatim_parallel( 

523 chunk_texts: list[str], 

524 llm_provider, 

525 *, 

526 event_date: datetime | None = None, 

527 max_concurrency: int = 6, 

528 max_retries: int = 3, 

529 base_retry_delay: float = 2.0, 

530) -> list[ExtractedFact]: 

531 """Per-chunk parallel verbatim extraction. 

532 

533 Phase 3 of the Hindsight cost-control port. Sends one LLM call per 

534 chunk in parallel, bounded by ``max_concurrency`` per session call. 

535 Returns one :class:`ExtractedFact` per input chunk in the same 

536 order; failed chunks get a metadata-less ExtractedFact preserving 

537 the chunk text. 

538 

539 When to use vs :func:`extract_facts_verbatim`: 

540 - Sessions with many small chunks (LME-shaped traffic, ~6 chunks 

541 after Phase 1 chunk_size=2048): per-chunk parallel typically 

542 ~2× faster wall time because per-call output is small and they 

543 run concurrently. 

544 - Sessions with very few chunks (1–2): batched is comparable; the 

545 extra round trips don't pay for themselves. 

546 

547 Drops cross-chunk ``causal_relations`` index references — they make 

548 no sense per-chunk. ``causal_links.enabled=false`` in our research 

549 configs, so this is a non-loss; if causal extraction is on, prefer 

550 the batched path. 

551 """ 

552 if not chunk_texts: 

553 return [] 

554 if not any(t.strip() for t in chunk_texts): 

555 return [] 

556 sem = asyncio.Semaphore(max(1, max_concurrency)) 

557 

558 async def _one(idx: int, text: str) -> _ChunkResult: 

559 async with sem: 

560 return _ChunkResult( 

561 idx=idx, 

562 raw=await _extract_one_chunk_verbatim( 

563 text, 

564 llm_provider, 

565 event_date=event_date, 

566 max_retries=max_retries, 

567 base_retry_delay=base_retry_delay, 

568 ), 

569 ) 

570 

571 gathered = await asyncio.gather( 

572 *[_one(i, t) for i, t in enumerate(chunk_texts)], 

573 return_exceptions=True, 

574 ) 

575 

576 # Reassemble in input order. asyncio.gather preserves order, but we 

577 # store explicit ``_ChunkResult`` records to defend against future 

578 # refactors that change the dispatch pattern. 

579 raw_by_idx: dict[int, dict] = {} 

580 for r in gathered: 

581 if isinstance(r, BaseException): 

582 # The semaphore-wrapped coroutine swallows exceptions inside 

583 # _extract_one_chunk_verbatim and returns {}, so reaching 

584 # this branch means a programming error or task 

585 # cancellation. Log and continue with an empty metadata. 

586 _logger.warning( 

587 "fact_extraction (verbatim/parallel) task raised: %r", 

588 r, 

589 ) 

590 continue 

591 raw_by_idx[r.idx] = r.raw if isinstance(r.raw, dict) else {} 

592 

593 return [_build_extracted_fact_from_raw(chunk, raw_by_idx.get(idx, {})) for idx, chunk in enumerate(chunk_texts)] 

594 

595 

596# --------------------------------------------------------------------------- 

597# JSON parsing helpers 

598# --------------------------------------------------------------------------- 

599 

600 

601def _parse_json_object(raw: str) -> dict | None: 

602 text = raw.strip() 

603 if text.startswith("```"): 

604 text = re.sub(r"^```(?:json)?\s*", "", text) 

605 text = re.sub(r"\s*```$", "", text) 

606 match = re.search(r"\{.*\}", text, re.DOTALL) 

607 if match is None: 

608 return None 

609 try: 

610 payload = json.loads(match.group(0)) 

611 except json.JSONDecodeError: 

612 return None 

613 return payload if isinstance(payload, dict) else None 

614 

615 

616def _parse_iso_datetime(value) -> datetime | None: 

617 if not value: 

618 return None 

619 if isinstance(value, datetime): 

620 return value if value.tzinfo else value.replace(tzinfo=timezone.utc) 

621 if not isinstance(value, str): 

622 return None 

623 s = value.strip() 

624 if not s: 

625 return None 

626 # Tolerate trailing 'Z' (Python 3.11+ handles it, but be defensive). 

627 if s.endswith("Z"): 

628 s = s[:-1] + "+00:00" 

629 try: 

630 dt = datetime.fromisoformat(s) 

631 except ValueError: 

632 return None 

633 return dt if dt.tzinfo else dt.replace(tzinfo=timezone.utc) 

634 

635 

636# --------------------------------------------------------------------------- 

637# Extraction 

638# --------------------------------------------------------------------------- 

639 

640 

641#: Running alignment counters for verbatim extraction, so the rate of 

642#: dropped or misplaced metadata is observable instead of silent. 

643VERBATIM_ALIGNMENT: dict[str, int] = { 

644 "calls": 0, "exact": 0, "keyed": 0, "positional": 0, 

645 "dropped_unlocatable": 0, "unmatched_chunks": 0, 

646 "duplicate_indices": 0, "invalid_indices": 0, 

647} 

648 

649 

650def _as_index(value) -> int | None: 

651 """An integer chunk index, or None. Rejects bools (an int subclass).""" 

652 if isinstance(value, bool): 

653 return None 

654 if isinstance(value, int): 

655 return value 

656 if isinstance(value, str) and value.strip().isdigit(): 

657 return int(value.strip()) 

658 return None 

659 

660 

661def align_verbatim_metadata(raw_metadata: list, n_chunks: int) -> list[dict]: 

662 """Map the model's metadata entries onto chunks, never onto the wrong one. 

663 

664 The old join paired entries with chunks by list position. That is only 

665 correct when the model returns exactly one entry per chunk in order. A 

666 paired thinking on/off experiment (18 real LongMemEval batches, 2026-10-03) 

667 found the length was wrong in 25% of calls with extended thinking on and 

668 69% with it off; entity grounding in the positional chunk fell from 87.6% 

669 to 74.2% (63.2% in the second half of each list) — the signature of one 

670 skipped or merged entry shifting every later chunk onto its neighbour's 

671 dates and entities, silently. 

672 

673 Join rules, in order: 

674 

675 1. Entries carry ``chunk_index`` -> join on it. An entry with a missing, 

676 out-of-range or duplicate index is discarded; its chunk gets no 

677 metadata rather than someone else's. 

678 2. No entry carries an index (an older prompt or a provider that ignored 

679 it) and the count matches exactly -> positional join, which is safe at 

680 an exact length. 

681 3. No indices and the count is wrong -> no metadata for any chunk. The 

682 point where positions diverged cannot be located, and absent metadata 

683 is recoverable while misattributed metadata is not. 

684 

685 Returns a list of length ``n_chunks``; unmatched chunks get ``{}``. 

686 """ 

687 VERBATIM_ALIGNMENT["calls"] += 1 

688 entries = [e for e in raw_metadata if isinstance(e, dict)] 

689 aligned: list[dict] = [{} for _ in range(n_chunks)] 

690 

691 mode = "dropped" 

692 if any("chunk_index" in e for e in entries): 

693 mode = "keyed" 

694 VERBATIM_ALIGNMENT["keyed"] += 1 

695 taken: set[int] = set() 

696 for e in entries: 

697 idx = _as_index(e.get("chunk_index")) 

698 if idx is None or not (0 <= idx < n_chunks): 

699 VERBATIM_ALIGNMENT["invalid_indices"] += 1 

700 continue 

701 if idx in taken: 

702 VERBATIM_ALIGNMENT["duplicate_indices"] += 1 

703 continue 

704 taken.add(idx) 

705 aligned[idx] = e 

706 elif len(entries) == n_chunks: 

707 mode = "positional" 

708 VERBATIM_ALIGNMENT["positional"] += 1 

709 aligned = list(entries) 

710 else: 

711 VERBATIM_ALIGNMENT["dropped_unlocatable"] += 1 

712 _logger.info( 

713 "fact_extraction (verbatim): %d entries for %d chunks and no chunk_index; " 

714 "dropping metadata rather than risk misattribution", 

715 len(entries), n_chunks, 

716 ) 

717 

718 unmatched = sum(1 for a in aligned if not a) 

719 VERBATIM_ALIGNMENT["unmatched_chunks"] += unmatched 

720 if unmatched == 0: 

721 VERBATIM_ALIGNMENT["exact"] += 1 

722 elif mode == "keyed": 

723 _logger.info( 

724 "fact_extraction (verbatim): %d of %d chunks have no matching metadata entry", 

725 unmatched, n_chunks, 

726 ) 

727 return aligned 

728 

729 

730async def extract_facts_verbatim( 

731 chunk_texts: list[str], 

732 llm_provider, 

733 *, 

734 event_date: datetime | None = None, 

735) -> list[ExtractedFact]: 

736 """Extract per-chunk metadata WITHOUT paraphrasing the chunk text. 

737 

738 The "what" field of each returned :class:`ExtractedFact` is set to 

739 the original chunk text — not an LLM-generated summary. The LLM's 

740 job here is only to produce structured metadata (entities, 

741 causal_relations, temporal range, fact_type, where/who/why 

742 annotations) per chunk. 

743 

744 Why this exists (the design lesson from 2026-05-02): 

745 The "concise" mode :func:`extract_facts` replaces conversation text 

746 with structured paraphrases like "Caroline went hiking | Involving: 

747 Caroline | When: yesterday". That paraphrase loses the surface 

748 vocabulary of the original conversation, which question embeddings 

749 typically share — causing severe recall_hit_rate degradation. The 

750 verbatim mode preserves the original vocabulary while still 

751 enriching each chunk with the structured metadata needed for 

752 causal/temporal/per-fact retrieval signals. 

753 

754 Returns one :class:`ExtractedFact` per input chunk, in the same 

755 order, with ``what`` = the chunk text. Returns ``[]`` on any 

756 failure so retain falls back to legacy chunking. 

757 

758 Args: 

759 chunk_texts: List of pre-chunked source texts; one per memory. 

760 llm_provider: Producer of the metadata extraction. 

761 event_date: Reference for resolving relative time expressions. 

762 """ 

763 if not chunk_texts: 

764 return [] 

765 # Prefilter: drop empty chunks but preserve indices for the LLM. 

766 if not any(t.strip() for t in chunk_texts): 

767 return [] 

768 

769 # Phase 2 of cost-control port: prefer structured outputs so the 

770 # decoder is constrained to ``_VERBATIM_JSON_SCHEMA`` and the 

771 # "malformed JSON" failure mode becomes impossible. Falls back 

772 # transparently to the legacy free-form path when the provider 

773 # doesn't accept the kwarg (legacy fakes / non-OpenAI providers 

774 # that haven't ported the SPI extension yet). 

775 messages_in = [ 

776 Message(role="system", content=_VERBATIM_SYSTEM_PROMPT), 

777 Message( 

778 role="user", 

779 content=_build_verbatim_user_prompt(chunk_texts, event_date), 

780 ), 

781 ] 

782 try: 

783 completion = await llm_provider.complete( 

784 messages_in, 

785 max_tokens=4096, 

786 temperature=0.0, 

787 response_format={"type": "json_schema", "json_schema": _VERBATIM_JSON_SCHEMA}, 

788 ) 

789 except TypeError: 

790 # Provider's complete() pre-dates the response_format kwarg — 

791 # retry without it. Free-form path is still resilient via the 

792 # _parse_json_object fallback below. 

793 try: 

794 completion = await llm_provider.complete( 

795 messages_in, 

796 max_tokens=4096, 

797 temperature=0.0, 

798 ) 

799 except Exception as exc: 

800 _logger.warning("fact_extraction (verbatim) LLM call failed (%s)", exc) 

801 return [] 

802 except Exception as exc: 

803 _logger.warning("fact_extraction (verbatim) LLM call failed (%s)", exc) 

804 return [] 

805 

806 parsed = _parse_json_object(completion.text) 

807 if parsed is None: 

808 _logger.warning("fact_extraction (verbatim): malformed JSON response") 

809 return [] 

810 raw_metadata = parsed.get("facts") 

811 if not isinstance(raw_metadata, list): 

812 _logger.warning("fact_extraction (verbatim): 'facts' is not a list") 

813 return [] 

814 

815 aligned = align_verbatim_metadata(raw_metadata, len(chunk_texts)) 

816 out: list[ExtractedFact] = [] 

817 for idx, chunk in enumerate(chunk_texts): 

818 # Metadata is joined by the echoed chunk_index, not list position 

819 # (see align_verbatim_metadata). Unmatched chunks still keep their 

820 # text; they just carry no metadata. 

821 raw = aligned[idx] 

822 

823 # Entities 

824 entities: list[FactEntity] = [] 

825 for ent in raw.get("entities") or []: 

826 if not isinstance(ent, dict): 

827 continue 

828 name = str(ent.get("name") or "").strip() 

829 if not name: 

830 continue 

831 etype = str(ent.get("entity_type") or "OTHER").strip().upper() or "OTHER" 

832 entities.append(FactEntity(name=name, entity_type=etype)) 

833 

834 # Causal relations — same semantics as concise mode but indices 

835 # reference the chunk position (which IS the memory position). 

836 causal: list[FactCausalRelation] = [] 

837 for rel in raw.get("causal_relations") or []: 

838 if not isinstance(rel, dict): 

839 continue 

840 try: 

841 target = int(rel.get("target_fact_index")) 

842 except (TypeError, ValueError): 

843 continue 

844 if target == idx or not (0 <= target < len(chunk_texts)): 

845 continue 

846 try: 

847 strength = float(rel.get("strength", 1.0)) 

848 except (TypeError, ValueError): 

849 strength = 1.0 

850 causal.append(FactCausalRelation(target_fact_index=target, strength=strength)) 

851 

852 ftype = str(raw.get("fact_type") or "").strip().lower() 

853 classified = ftype in {"world", "experience"} 

854 if not classified: 

855 ftype = "experience" 

856 

857 out.append( 

858 ExtractedFact( 

859 # KEY: "what" is the ORIGINAL chunk text, not a paraphrase. 

860 what=chunk, 

861 when=str(raw.get("when") or "N/A").strip() or "N/A", 

862 where=str(raw.get("where") or "N/A").strip() or "N/A", 

863 who=str(raw.get("who") or "N/A").strip() or "N/A", 

864 why=str(raw.get("why") or "N/A").strip() or "N/A", 

865 fact_type=ftype, 

866 fact_type_classified=classified, 

867 occurred_start=_parse_iso_datetime(raw.get("occurred_start")), 

868 occurred_end=_parse_iso_datetime(raw.get("occurred_end")), 

869 entities=entities, 

870 causal_relations=causal, 

871 ) 

872 ) 

873 return out 

874 

875 

876# --------------------------------------------------------------------------- 

877# Materialization 

878# --------------------------------------------------------------------------- 

879 

880 

881def _slug(value: str) -> str: 

882 return re.sub(r"[^a-z0-9]+", "-", value.lower()).strip("-") or "unknown" 

883 

884 

885@dataclass 

886class MaterializedFacts: 

887 """The retain-pipeline artefacts produced from a fact-extraction batch. 

888 

889 - ``vector_items`` — one VectorItem per fact, ready for store_vectors. 

890 The item's ``id`` is also the fact's identifier for causal/entity 

891 linkage. Order matches the input ``ExtractedFact`` list. 

892 - ``entities`` — deduplicated Entity list across the batch (one 

893 entity per unique (name, type) tuple). Each entity uses a 

894 deterministic ID ``{type}:{slug}`` so subsequent retain calls 

895 collapse repeated mentions to the same canonical row. 

896 - ``memory_entity_associations`` — list of (vector_item_id, 

897 entity_id) tuples for the orchestrator's 

898 link_memories_to_entities call. 

899 - ``memory_links`` — directional caused_by edges; source is the 

900 effect's memory ID, target is the cause's memory ID. 

901 """ 

902 

903 vector_items: list[VectorItem] 

904 entities: list[Entity] 

905 memory_entity_associations: list[tuple[str, str]] 

906 memory_links: list[MemoryLink] 

907 

908 

909def materialize_facts( 

910 facts: list[ExtractedFact], 

911 *, 

912 bank_id: str, 

913 tags: list[str] | None = None, 

914 metadata: dict | None = None, 

915 occurred_at: datetime | None = None, 

916 embeddings: list[list[float]] | None = None, 

917) -> MaterializedFacts: 

918 """Convert extracted facts to retain-pipeline artefacts. 

919 

920 The orchestrator wires this output into its existing storage path 

921 (store_vectors → store_entities → link_memories_to_entities → 

922 store_memory_links) without further translation. 

923 

924 The VectorItem's ``text`` is always the raw chunk content 

925 (``fact.what``) — verbatim mode is now the only supported flow 

926 after the M9 concise-path removal. Storing the original vocabulary 

927 is what makes question embeddings match against retained memories; 

928 decorated text (``what | Involving: ... | When: ...``) caused 

929 severe recall_hit_rate degradation (2026-05-02 finding). 

930 

931 Args: 

932 facts: Output of :func:`extract_facts_verbatim` / 

933 :func:`extract_facts_verbatim_parallel`. 

934 bank_id: Target bank. 

935 tags: Tags applied to every produced VectorItem. 

936 metadata: Base metadata merged onto every VectorItem; the