Coverage for narrative_harm_classifier/classifier/tracking/store.py: 100%

1""" 

2classifier/tracking/store.py — Persistence for escalation-chain observations. 

3 

4Uses SQLAlchemy Core (not the ORM) against Settings.effective_tracking_db_url, 

5which defaults to a local SQLite file (zero-config) and works unchanged 

6against Postgres for real deployments. 

7""" 

8 

9from datetime import datetime 

10from functools import lru_cache 

11from typing import Optional 

12 

13from sqlalchemy import ( 

14 create_engine, 

15 MetaData, 

16 Table, 

17 Column, 

18 Integer, 

19 String, 

20 Float, 

21 Boolean, 

22 DateTime, 

23 select, 

24 insert, 

25) 

26 

27from narrative_harm_classifier.classifier.tracking.models import Observation, SeverityLevel 

28 

29metadata = MetaData() 

30 

31observations_table = Table( 

32 "observations", 

33 metadata, 

34 Column("id", Integer, primary_key=True, autoincrement=True), 

35 Column("source_id", String, nullable=False, index=True), 

36 Column("text_excerpt", String, nullable=False), 

37 Column("is_harmful", Boolean, nullable=False), 

38 Column("harm_category", String, nullable=False), 

39 Column("harm_mechanism", String, nullable=True), 

40 Column("confidence", Float, nullable=False), 

41 Column("severity", Integer, nullable=False), 

42 Column("observed_at", DateTime, nullable=False), 

43 Column("content_hash", String, nullable=False, server_default=""), 

44 Column("prev_hash", String, nullable=False, server_default=""), 

45 Column("record_hash", String, nullable=False, server_default=""), 

46) 

47 

48 

49class TrackingStore: 

50 def __init__(self, db_url: str): 

51 self.engine = create_engine(db_url, future=True) 

52 metadata.create_all(self.engine) 

53 

54 def add_observation(self, obs: Observation) -> Observation: 

55 with self.engine.begin() as conn: 

56 result = conn.execute( 

57 insert(observations_table).values( 

58 source_id=obs.source_id, 

59 text_excerpt=obs.text_excerpt, 

60 is_harmful=obs.is_harmful, 

61 harm_category=obs.harm_category, 

62 harm_mechanism=obs.harm_mechanism, 

63 confidence=obs.confidence, 

64 severity=int(obs.severity), 

65 observed_at=obs.observed_at, 

66 content_hash=obs.content_hash, 

67 prev_hash=obs.prev_hash, 

68 record_hash=obs.record_hash, 

69 ) 

70 ) 

71 obs.id = result.inserted_primary_key[0] 

72 return obs 

73 

74 def last_record_hash(self, source_id: str) -> str: 

75 """Most recent record_hash for a source, or GENESIS_HASH if it has no history yet.""" 

76 from narrative_harm_classifier.classifier.provenance import GENESIS_HASH 

77 

78 with self.engine.connect() as conn: 

79 query = ( 

80 select(observations_table.c.record_hash) 

81 .where(observations_table.c.source_id == source_id) 

82 .order_by(observations_table.c.observed_at.desc(), observations_table.c.id.desc()) 

83 .limit(1) 

84 ) 

85 row = conn.execute(query).first() 

86 return row[0] if row else GENESIS_HASH 

87 

88 def history(self, source_id: str, limit: Optional[int] = None) -> list[Observation]: 

89 with self.engine.connect() as conn: 

90 query = ( 

91 select(observations_table) 

92 .where(observations_table.c.source_id == source_id) 

93 .order_by(observations_table.c.observed_at.asc()) 

94 ) 

95 rows = conn.execute(query).mappings().all() 

96 

97 history = [ 

98 Observation( 

99 id=row["id"], 

100 source_id=row["source_id"], 

101 text_excerpt=row["text_excerpt"], 

102 is_harmful=row["is_harmful"], 

103 harm_category=row["harm_category"], 

104 harm_mechanism=row["harm_mechanism"], 

105 confidence=row["confidence"], 

106 severity=SeverityLevel(row["severity"]), 

107 observed_at=row["observed_at"], 

108 content_hash=row["content_hash"] or "", 

109 prev_hash=row["prev_hash"] or "", 

110 record_hash=row["record_hash"] or "", 

111 ) 

112 for row in rows 

113 ] 

114 if limit is not None: 

115 history = history[-limit:] 

116 return history 

117 

118 def list_source_ids(self) -> list[str]: 

119 with self.engine.connect() as conn: 

120 rows = conn.execute(select(observations_table.c.source_id).distinct()).all() 

121 return [r[0] for r in rows] 

122 

123 

124@lru_cache(maxsize=4) 

125def get_store(db_url: str) -> TrackingStore: 

126 return TrackingStore(db_url)