Coverage for astrocyte/pipeline/sync_push.py: 95%

1"""Team-memory batch push: one stored row per client record (team-memory.md §8, G2). 

2 

3``Astrocyte.push_records`` runs the policy layer (validation, PII, metadata, 

4provenance, rate limits) and hands the surviving records here. Each becomes 

5exactly one row with the client's id: no chunking, no extraction, no LLM call; 

6the text is re-embedded with this server's model. 

7 

8Ids are global in the stores (both SQL stores key rows on ``id`` alone), so a 

9pushed id is checked against every bank before anything is written, and the 

10write itself is insert-only (``VectorStore.insert_vectors``): a row in another 

11bank, a forgotten row, or a row a concurrent push wrote first is never 

12overwritten. 

13""" 

14 

15from __future__ import annotations 

16 

17import logging 

18from datetime import datetime, timezone 

19from typing import Any 

20 

21from astrocyte.pipeline.embedding import generate_embeddings 

22from astrocyte.types import MemoryChange, SyncPushRecord, SyncPushResult, VectorItem 

23 

24_logger = logging.getLogger("astrocyte.sync") 

25 

26#: Same wording whether the id is held in another bank or in this one with 

27#: other text, so a push can't be used to learn that another bank holds an id. 

28REASON_IN_USE = "id is already in use with different text; memory text is immutable" 

29REASON_FORGOTTEN = "id was forgotten; a forgotten memory can't be pushed again" 

30REASON_METADATA_IGNORED = "text unchanged; metadata updates are not accepted by push" 

31 

32 

33def _utc(dt: datetime | None) -> datetime | None: 

34 if dt is None: 

35 return None 

36 return dt.replace(tzinfo=timezone.utc) if dt.tzinfo is None else dt 

37 

38 

39def classify_existing(bank_id: str, record: SyncPushRecord, found: MemoryChange) -> SyncPushResult: 

40 """The result for a pushed id the store already holds. 

41 

42 Text is immutable; other fields (metadata, and later claim status, trust, 

43 staleness) will change through their own update path, so a re-push with 

44 the same text but different metadata is ``unchanged`` and the stored 

45 metadata is left as it is. 

46 """ 

47 if found.bank_id != bank_id: 

48 return SyncPushResult(id=record.id, status="rejected", reason=REASON_IN_USE) 

49 if found.deleted: 

50 return SyncPushResult(id=record.id, status="rejected", reason=REASON_FORGOTTEN) 

51 if found.text != record.text: 

52 return SyncPushResult(id=record.id, status="rejected", reason=REASON_IN_USE) 

53 same_fields = ( 

54 (found.metadata or None) == (record.metadata or None) 

55 and (found.tags or None) == (record.tags or None) 

56 and found.fact_type == record.fact_type 

57 and _utc(found.occurred_at) == _utc(record.occurred_at) 

58 ) 

59 return SyncPushResult( 

60 id=record.id, 

61 status="unchanged", 

62 reason=None if same_fields else REASON_METADATA_IGNORED, 

63 ) 

64 

65 

66class SyncPushStageMixin: 

67 """``push_records`` for :class:`~astrocyte.pipeline.orchestrator.PipelineOrchestrator`.""" 

68 

69 async def push_records(self: Any, bank_id: str, records: list[SyncPushRecord]) -> list[SyncPushResult]: 

70 """Store each record as one row with its own id; see the module docstring. 

71 

72 ``records`` have already passed the policy layer (their text is the 

73 possibly-redacted text, their metadata sanitized and stamped). The 

74 vector store must have ``lookup_ids`` and ``insert_vectors``. Returns 

75 one result per record, in order. 

76 """ 

77 store = self.vector_store 

78 existing = {c.id: c for c in await store.lookup_ids([r.id for r in records])} 

79 results: list[SyncPushResult | None] = [None] * len(records) 

80 pending: list[int] = [] 

81 for i, record in enumerate(records): 

82 found = existing.get(record.id) 

83 if found is None: 

84 pending.append(i) 

85 else: 

86 results[i] = classify_existing(bank_id, record, found) 

87 

88 embeddings = await generate_embeddings([records[i].text for i in pending], self.llm_provider) if pending else [] 

89 stored: list[VectorItem] = [] 

90 warn_only = self._config_dedup_action(bank_id) == "warn" 

91 for i, embedding in zip(pending, embeddings): 

92 record = records[i] 

93 # One at a time, so a record near-duplicating one stored earlier in 

94 # this batch is caught (the dedup cache learns each insert). 

95 [duplicate_of] = await self._find_duplicate_ids(bank_id, [record.text], [embedding], None) 

96 if duplicate_of is not None and duplicate_of != record.id and not warn_only: 

97 results[i] = SyncPushResult(id=record.id, status="duplicate", duplicate_of=duplicate_of) 

98 continue 

99 item = VectorItem( 

100 id=record.id, 

101 bank_id=bank_id, 

102 vector=embedding, 

103 text=record.text, 

104 metadata=record.metadata, 

105 tags=record.tags, 

106 fact_type=record.fact_type, 

107 occurred_at=record.occurred_at, 

108 retained_at=datetime.now(timezone.utc), 

109 ) 

110 if await store.insert_vectors([item]): 

111 results[i] = SyncPushResult(id=record.id, status="stored") 

112 self._dedup.add(bank_id, record.id, embedding, text=record.text) 

113 stored.append(item) 

114 continue 

115 # Someone else holds the id now: the same id earlier in this 

116 # batch, or a concurrent push. Answer as if it had been there. 

117 [*now_found] = await store.lookup_ids([record.id]) 

118 results[i] = ( 

119 classify_existing(bank_id, record, now_found[0]) 

120 if now_found 

121 else SyncPushResult(id=record.id, status="rejected", reason="could not be stored") 

122 ) 

123 

124 await self._mirror_pushed(bank_id, stored) 

125 return [r for r in results if r is not None] 

126 

127 async def _mirror_pushed(self: Any, bank_id: str, items: list[VectorItem]) -> None: 

128 """What retain does after storing vectors, minus anything that calls 

129 an LLM: the keyword index and the semantic-kNN graph. Best-effort.""" 

130 if not items: 

131 return 

132 if self.document_store is not None: 

133 from astrocyte.types import Document 

134 

135 for item in items: 

136 try: 

137 await self.document_store.store_document( 

138 Document(id=item.id, text=item.text, metadata=item.metadata, tags=item.tags), bank_id 

139 ) 

140 except Exception as exc: 

141 _logger.warning("sync push: document_store.store_document failed for %s: %s", item.id, exc) 

142 await self._persist_semantic_links(bank_id, [i.id for i in items], [i.vector for i in items])