Coverage for src \ truenex_memory \ store \ sqlite.py: 93%

1"""SQLite schema management.""" 

2 

3from __future__ import annotations 

4 

5from dataclasses import dataclass 

6import json 

7from pathlib import Path 

8import sqlite3 

9from typing import Any 

10 

11 

12SCHEMA_VERSION = "4" 

13 

14 

15@dataclass(frozen=True) 

16class MemoryRecord: 

17 """Simple record used by the compatibility SQLite memory store.""" 

18 

19 id: int 

20 text: str 

21 metadata: dict[str, Any] 

22 

23 

24def connect(db_path: Path) -> sqlite3.Connection: 

25 """Open a SQLite connection with local-first defaults.""" 

26 

27 db_path.parent.mkdir(parents=True, exist_ok=True) 

28 conn = sqlite3.connect(db_path) 

29 conn.row_factory = sqlite3.Row 

30 conn.execute("PRAGMA foreign_keys = ON") 

31 return conn 

32 

33 

34def initialize_schema(conn: sqlite3.Connection) -> None: 

35 """Create the v1 local schema if it does not already exist.""" 

36 

37 conn.executescript( 

38 """ 

39 CREATE TABLE IF NOT EXISTS documents ( 

40 id TEXT PRIMARY KEY, 

41 project_id TEXT NOT NULL, 

42 path TEXT NOT NULL, 

43 filename TEXT NOT NULL, 

44 content_hash TEXT NOT NULL, 

45 last_indexed_at TEXT, 

46 created_at TEXT NOT NULL, 

47 updated_at TEXT NOT NULL 

48 ); 

49 

50 CREATE TABLE IF NOT EXISTS chunks ( 

51 id TEXT PRIMARY KEY, 

52 document_id TEXT NOT NULL, 

53 chunk_index INTEGER NOT NULL, 

54 heading_path TEXT, 

55 content TEXT NOT NULL, 

56 content_hash TEXT NOT NULL, 

57 token_count INTEGER NOT NULL DEFAULT 0, 

58 qdrant_point_id TEXT, 

59 embedding_model TEXT, 

60 embedding_vector_json TEXT, 

61 source_type TEXT, 

62 created_at TEXT NOT NULL, 

63 updated_at TEXT NOT NULL, 

64 FOREIGN KEY(document_id) REFERENCES documents(id) ON DELETE CASCADE 

65 ); 

66 

67 CREATE TABLE IF NOT EXISTS memory_nodes ( 

68 id TEXT PRIMARY KEY, 

69 project_id TEXT NOT NULL, 

70 type TEXT NOT NULL, 

71 title TEXT NOT NULL, 

72 content TEXT NOT NULL, 

73 status TEXT NOT NULL, 

74 source_kind TEXT NOT NULL, 

75 source_document_id TEXT, 

76 source_chunk_id TEXT, 

77 source_path TEXT, 

78 content_hash TEXT, 

79 created_by TEXT NOT NULL, 

80 model_name TEXT, 

81 confidence REAL, 

82 created_at TEXT NOT NULL, 

83 updated_at TEXT NOT NULL 

84 ); 

85 

86 CREATE TABLE IF NOT EXISTS edges ( 

87 id TEXT PRIMARY KEY, 

88 project_id TEXT NOT NULL, 

89 source_node_id TEXT NOT NULL, 

90 target_node_id TEXT NOT NULL, 

91 relation_type TEXT NOT NULL, 

92 created_by TEXT NOT NULL, 

93 confidence REAL, 

94 created_at TEXT NOT NULL 

95 ); 

96 

97 CREATE TABLE IF NOT EXISTS retrieval_logs ( 

98 id TEXT PRIMARY KEY, 

99 project_id TEXT NOT NULL, 

100 query TEXT NOT NULL, 

101 top_k INTEGER NOT NULL, 

102 result_count INTEGER NOT NULL, 

103 results_json TEXT NOT NULL, 

104 created_at TEXT NOT NULL 

105 ); 

106 

107 CREATE TABLE IF NOT EXISTS schema_migrations ( 

108 version TEXT PRIMARY KEY, 

109 applied_at TEXT NOT NULL 

110 ); 

111 

112 CREATE TABLE IF NOT EXISTS memories ( 

113 id INTEGER PRIMARY KEY AUTOINCREMENT, 

114 text TEXT NOT NULL, 

115 metadata_json TEXT NOT NULL DEFAULT '{}', 

116 created_at TEXT NOT NULL DEFAULT CURRENT_TIMESTAMP 

117 ); 

118 

119 CREATE TABLE IF NOT EXISTS source_ledger ( 

120 source_id TEXT PRIMARY KEY, 

121 source_path_or_alias TEXT NOT NULL, 

122 project_name TEXT, 

123 source_type TEXT NOT NULL, 

124 parser_version TEXT NOT NULL DEFAULT '1', 

125 content_hash TEXT, 

126 last_modified_at TEXT, 

127 last_indexed_at TEXT, 

128 status TEXT NOT NULL DEFAULT 'pending' 

129 CHECK(status IN ('active','pending','error','missing','skipped')), 

130 error_message TEXT, 

131 chunk_count INTEGER NOT NULL DEFAULT 0, 

132 created_at TEXT NOT NULL, 

133 updated_at TEXT NOT NULL 

134 ); 

135 

136 CREATE TABLE IF NOT EXISTS tasks ( 

137 task_id TEXT PRIMARY KEY, 

138 title TEXT NOT NULL, 

139 type TEXT NOT NULL CHECK(type IN ('bugfix','feature','refactor','review','query')), 

140 project TEXT, 

141 agent_session_id TEXT, 

142 human_outcome INTEGER CHECK(human_outcome IN (1, 0, -1)), 

143 human_comment TEXT, 

144 total_tokens INTEGER, 

145 total_duration_s REAL, 

146 status TEXT NOT NULL DEFAULT 'open' CHECK(status IN ('open','closed','unrated')), 

147 created_at TEXT NOT NULL, 

148 closed_at TEXT 

149 ); 

150 

151 CREATE TABLE IF NOT EXISTS task_steps ( 

152 step_id TEXT PRIMARY KEY, 

153 task_id TEXT NOT NULL, 

154 step_index INTEGER NOT NULL, 

155 prompt_used TEXT, 

156 output TEXT, 

157 brain_judgment TEXT CHECK(brain_judgment IN ('ok','needs_revision','rejected')), 

158 tokens_used INTEGER, 

159 duration_s REAL, 

160 model_used TEXT, 

161 created_at TEXT NOT NULL, 

162 FOREIGN KEY(task_id) REFERENCES tasks(task_id) ON DELETE CASCADE 

163 ); 

164 

165 CREATE TABLE IF NOT EXISTS verifier_rounds ( 

166 round_id TEXT PRIMARY KEY, 

167 task_id TEXT NOT NULL, 

168 step_id TEXT, 

169 suggestion_type TEXT NOT NULL, 

170 brain_accepted INTEGER NOT NULL CHECK(brain_accepted IN (0, 1)), 

171 brain_rationale TEXT, 

172 created_at TEXT NOT NULL, 

173 FOREIGN KEY(task_id) REFERENCES tasks(task_id) ON DELETE CASCADE 

174 ); 

175 """ 

176 ) 

177 conn.execute( 

178 "INSERT OR IGNORE INTO schema_migrations(version, applied_at) VALUES (?, datetime('now'))", 

179 (SCHEMA_VERSION,), 

180 ) 

181 _ensure_column(conn, "chunks", "embedding_model", "TEXT") 

182 _ensure_column(conn, "chunks", "embedding_vector_json", "TEXT") 

183 conn.commit() 

184 

185 

186def _ensure_column(conn: sqlite3.Connection, table: str, column: str, column_type: str) -> None: 

187 columns = {row["name"] for row in conn.execute(f"PRAGMA table_info({table})").fetchall()} 

188 if column not in columns: 

189 conn.execute(f"ALTER TABLE {table} ADD COLUMN {column} {column_type}") 

190 

191 

192def apply_column_upgrades(conn: sqlite3.Connection) -> None: 

193 """Add columns introduced in schema upgrades to existing tables. 

194 

195 Uses try/except because SQLite has no IF NOT EXISTS for ALTER TABLE. 

196 Safe to call repeatedly — duplicate-column errors are ignored. 

197 """ 

198 upgrades = [ 

199 "ALTER TABLE chunks ADD COLUMN source_type TEXT", 

200 ] 

201 for sql in upgrades: 

202 try: 

203 conn.execute(sql) 

204 conn.commit() 

205 except sqlite3.OperationalError: 

206 pass # column already exists 

207 

208 

209class SQLiteMemoryStore: 

210 """Small local text store kept for focused Task 2 tests.""" 

211 

212 def __init__(self, database: str | Path) -> None: 

213 self.database = database 

214 self._connection: sqlite3.Connection | None = None 

215 

216 def __enter__(self) -> "SQLiteMemoryStore": 

217 self.open() 

218 return self 

219 

220 def __exit__(self, exc_type, exc, tb) -> None: 

221 self.close() 

222 

223 def open(self) -> None: 

224 if self._connection is not None: 

225 return 

226 if self.database == ":memory:": 

227 conn = sqlite3.connect(":memory:") 

228 else: 

229 conn = connect(Path(self.database)) 

230 conn.row_factory = sqlite3.Row 

231 initialize_schema(conn) 

232 self._connection = conn 

233 

234 def close(self) -> None: 

235 if self._connection is not None: 

236 self._connection.close() 

237 self._connection = None 

238 

239 def add(self, text: str, metadata: dict[str, Any] | None = None) -> MemoryRecord: 

240 clean_text = text.strip() 

241 if not clean_text: 

242 raise ValueError("text cannot be empty") 

243 conn = self._conn() 

244 metadata_json = json.dumps(metadata or {}, sort_keys=True) 

245 cursor = conn.execute( 

246 "INSERT INTO memories(text, metadata_json) VALUES (?, ?)", 

247 (clean_text, metadata_json), 

248 ) 

249 conn.commit() 

250 return MemoryRecord(id=int(cursor.lastrowid), text=clean_text, metadata=dict(metadata or {})) 

251 

252 def search(self, query: str, *, limit: int = 5) -> list[MemoryRecord]: 

253 if limit < 1: 

254 raise ValueError("limit must be greater than zero") 

255 tokens = [token.lower() for token in query.split() if token.strip()] 

256 if not tokens: 

257 return [] 

258 rows = self._conn().execute("SELECT id, text, metadata_json FROM memories ORDER BY id").fetchall() 

259 results: list[MemoryRecord] = [] 

260 for row in rows: 

261 text = str(row["text"]) 

262 lowered = text.lower() 

263 if all(token in lowered for token in tokens): 

264 results.append( 

265 MemoryRecord( 

266 id=int(row["id"]), 

267 text=text, 

268 metadata=json.loads(row["metadata_json"]), 

269 ) 

270 ) 

271 if len(results) >= limit: 

272 break 

273 return results 

274 

275 def schema_version(self) -> int: 

276 row = self._conn().execute( 

277 "SELECT version FROM schema_migrations ORDER BY CAST(version AS INTEGER) DESC LIMIT 1" 

278 ).fetchone() 

279 return int(row["version"]) if row else 0 

280 

281 def _conn(self) -> sqlite3.Connection: 

282 if self._connection is None: 

283 self.open() 

284 assert self._connection is not None 

285 return self._connection