Coverage for src/semware/services/vectordb.py: 71%

1"""LanceDB integration for vector storage and retrieval.""" 

2 

3import json 

4from datetime import datetime 

5from pathlib import Path 

6from typing import Any 

7 

8import lancedb 

9import numpy as np 

10import pyarrow as pa 

11from loguru import logger 

12 

13from ..config import settings 

14from ..models.schemas import DataRecord, SearchResult, TableSchema 

15 

16 

17class VectorDBService: 

18 """Service for managing LanceDB operations.""" 

19 

20 def __init__(self, db_path: Path | None = None): 

21 """Initialize the vector database service. 

22 

23 Args: 

24 db_path: Path to the database directory 

25 """ 

26 self.db_path = db_path or settings.db_path 

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

28 

29 # Connect to LanceDB 

30 self.db = lancedb.connect(str(self.db_path)) 

31 self.table_schemas: dict[str, TableSchema] = {} 

32 

33 # Load existing table schemas 

34 self._load_table_schemas() 

35 

36 logger.info(f"VectorDB initialized at: {self.db_path}") 

37 

38 def _load_table_schemas(self) -> None: 

39 """Load table schemas from metadata.""" 

40 schema_file = self.db_path / "table_schemas.json" 

41 if schema_file.exists(): 

42 try: 

43 with open(schema_file) as f: 

44 schemas_data = json.load(f) 

45 for table_name, schema_data in schemas_data.items(): 

46 self.table_schemas[table_name] = TableSchema(**schema_data) 

47 logger.info(f"Loaded {len(self.table_schemas)} table schemas") 

48 except Exception as e: 

49 logger.error(f"Failed to load table schemas: {e}") 

50 

51 def _save_table_schemas(self) -> None: 

52 """Save table schemas to metadata.""" 

53 schema_file = self.db_path / "table_schemas.json" 

54 try: 

55 schemas_data = { 

56 name: schema.model_dump() for name, schema in self.table_schemas.items() 

57 } 

58 with open(schema_file, "w") as f: 

59 json.dump(schemas_data, f, indent=2) 

60 logger.debug("Table schemas saved") 

61 except Exception as e: 

62 logger.error(f"Failed to save table schemas: {e}") 

63 

64 def _create_arrow_schema(self, table_schema: TableSchema) -> pa.Schema: 

65 """Create PyArrow schema from table schema. 

66 

67 Args: 

68 table_schema: Table schema definition 

69 

70 Returns: 

71 PyArrow schema 

72 """ 

73 fields = [] 

74 

75 # Add all columns from schema 

76 for col_name, col_type in table_schema.columns.items(): 

77 if col_type == "string": 

78 arrow_type = pa.string() 

79 elif col_type == "int": 

80 arrow_type = pa.int64() 

81 elif col_type == "float": 

82 arrow_type = pa.float64() 

83 elif col_type == "bool": 

84 arrow_type = pa.bool_() 

85 else: 

86 logger.warning( 

87 f"Unknown column type '{col_type}', defaulting to string" 

88 ) 

89 arrow_type = pa.string() 

90 

91 fields.append(pa.field(col_name, arrow_type)) 

92 

93 # Add embedding vector column (fixed size for LanceDB vector indexing) 

94 fields.append(pa.field("_embedding", pa.list_(pa.float32(), settings.embedding_dimension))) 

95 

96 # Add metadata columns 

97 fields.append(pa.field("_created_at", pa.timestamp("ms"))) 

98 fields.append(pa.field("_updated_at", pa.timestamp("ms"))) 

99 

100 return pa.schema(fields) 

101 

102 def _create_dummy_record(self, table_schema: TableSchema, arrow_schema: pa.Schema) -> pa.Table: 

103 """Create a dummy record for table initialization. 

104  

105 Args: 

106 table_schema: Table schema definition 

107 arrow_schema: PyArrow schema 

108  

109 Returns: 

110 PyArrow table with one dummy record 

111 """ 

112 data = {} 

113 current_time = datetime.now() 

114 

115 # Add dummy values for all columns 

116 for col_name, col_type in table_schema.columns.items(): 

117 if col_type == "string": 

118 data[col_name] = ["dummy"] 

119 elif col_type == "int": 

120 data[col_name] = [0] 

121 elif col_type == "float": 

122 data[col_name] = [0.0] 

123 elif col_type == "bool": 

124 data[col_name] = [False] 

125 else: 

126 data[col_name] = ["dummy"] 

127 

128 # Add dummy embedding (384 dimensions for MiniLM) 

129 data["_embedding"] = [[0.0] * settings.embedding_dimension] 

130 

131 # Add metadata 

132 data["_created_at"] = [current_time] 

133 data["_updated_at"] = [current_time] 

134 

135 return pa.table(data, schema=arrow_schema) 

136 

137 def create_table(self, schema: TableSchema) -> None: 

138 """Create a new table with the given schema. 

139 

140 Args: 

141 schema: Table schema definition 

142 

143 Raises: 

144 ValueError: If table already exists 

145 """ 

146 if schema.name in self.table_schemas: 

147 raise ValueError(f"Table '{schema.name}' already exists") 

148 

149 if schema.name in self.db.table_names(): 

150 raise ValueError(f"Table '{schema.name}' already exists in database") 

151 

152 # Create arrow schema 

153 arrow_schema = self._create_arrow_schema(schema) 

154 

155 # Create table using schema only (LanceDB will handle empty table creation) 

156 try: 

157 table = self.db.create_table(schema.name, schema=arrow_schema, mode="create") 

158 except Exception as e: 

159 logger.error(f"Failed to create table with schema-only approach: {e}") 

160 # Fallback: create with one dummy record then delete it 

161 dummy_data = self._create_dummy_record(schema, arrow_schema) 

162 table = self.db.create_table(schema.name, dummy_data, mode="create") 

163 # Delete the dummy record 

164 table.delete("_rowid = 0") 

165 

166 # Store schema 

167 self.table_schemas[schema.name] = schema 

168 self._save_table_schemas() 

169 

170 logger.info(f"Created table '{schema.name}' with {len(schema.columns)} columns") 

171 

172 def delete_table(self, table_name: str) -> None: 

173 """Delete a table. 

174 

175 Args: 

176 table_name: Name of the table to delete 

177 

178 Raises: 

179 ValueError: If table doesn't exist 

180 """ 

181 if table_name not in self.table_schemas: 

182 raise ValueError(f"Table '{table_name}' does not exist") 

183 

184 # Drop the table 

185 self.db.drop_table(table_name) 

186 

187 # Remove from schemas 

188 del self.table_schemas[table_name] 

189 self._save_table_schemas() 

190 

191 logger.info(f"Deleted table '{table_name}'") 

192 

193 def get_table_names(self) -> list[str]: 

194 """Get list of all table names. 

195 

196 Returns: 

197 List of table names 

198 """ 

199 return list(self.table_schemas.keys()) 

200 

201 def get_table_schema(self, table_name: str) -> TableSchema: 

202 """Get schema for a specific table. 

203 

204 Args: 

205 table_name: Name of the table 

206 

207 Returns: 

208 Table schema 

209 

210 Raises: 

211 ValueError: If table doesn't exist 

212 """ 

213 if table_name not in self.table_schemas: 

214 raise ValueError(f"Table '{table_name}' does not exist") 

215 

216 return self.table_schemas[table_name] 

217 

218 def get_table_record_count(self, table_name: str) -> int: 

219 """Get the number of records in a table. 

220 

221 Args: 

222 table_name: Name of the table 

223 

224 Returns: 

225 Number of records 

226 """ 

227 if table_name not in self.table_schemas: 

228 raise ValueError(f"Table '{table_name}' does not exist") 

229 

230 try: 

231 table = self.db.open_table(table_name) 

232 return table.count_rows() 

233 except Exception as e: 

234 logger.error(f"Failed to count rows in table '{table_name}': {e}") 

235 return 0 

236 

237 def upsert_records( 

238 self, table_name: str, records: list[DataRecord], embeddings: list[np.ndarray] 

239 ) -> tuple[int, int]: 

240 """Insert or update records in a table. 

241 

242 Args: 

243 table_name: Name of the table 

244 records: List of data records 

245 embeddings: List of embedding vectors (one per record) 

246 

247 Returns: 

248 Tuple of (inserted_count, updated_count) 

249 

250 Raises: 

251 ValueError: If table doesn't exist or data validation fails 

252 """ 

253 if table_name not in self.table_schemas: 

254 raise ValueError(f"Table '{table_name}' does not exist") 

255 

256 if len(records) != len(embeddings): 

257 raise ValueError("Number of records and embeddings must match") 

258 

259 schema = self.table_schemas[table_name] 

260 table = self.db.open_table(table_name) 

261 

262 # Prepare data for insertion 

263 insert_data = [] 

264 now = datetime.now() 

265 

266 for record, embedding in zip(records, embeddings, strict=False): 

267 # Validate record against schema 

268 self._validate_record(record, schema) 

269 

270 # Prepare row data 

271 row = record.data.copy() 

272 row["_embedding"] = embedding.tolist() 

273 row["_created_at"] = now 

274 row["_updated_at"] = now 

275 

276 insert_data.append(row) 

277 

278 # Convert to PyArrow table with proper schema 

279 arrow_schema = self._create_arrow_schema(schema) 

280 

281 # Convert dict data to column format 

282 columns_data = {} 

283 for field in arrow_schema: 

284 column_name = field.name 

285 if column_name == "_embedding": 

286 columns_data[column_name] = [row[column_name] for row in insert_data] 

287 else: 

288 columns_data[column_name] = [row[column_name] for row in insert_data] 

289 

290 arrow_table = pa.table(columns_data, schema=arrow_schema) 

291 

292 # Upsert data (merge on ID column) 

293 table.merge_insert( 

294 f"{schema.id_column}" 

295 ).when_matched_update_all().when_not_matched_insert_all().execute(arrow_table) 

296 

297 # For now, assume all are inserts (Lance doesn't return update counts easily) 

298 inserted_count = len(records) 

299 updated_count = 0 

300 

301 logger.info(f"Upserted {len(records)} records in table '{table_name}'") 

302 return inserted_count, updated_count 

303 

304 def delete_record(self, table_name: str, record_id: str | int) -> None: 

305 """Delete a record from a table. 

306 

307 Args: 

308 table_name: Name of the table 

309 record_id: ID of the record to delete 

310 

311 Raises: 

312 ValueError: If table doesn't exist 

313 """ 

314 if table_name not in self.table_schemas: 

315 raise ValueError(f"Table '{table_name}' does not exist") 

316 

317 schema = self.table_schemas[table_name] 

318 table = self.db.open_table(table_name) 

319 

320 # Delete the record 

321 table.delete(f"{schema.id_column} = '{record_id}'") 

322 

323 logger.info(f"Deleted record with ID '{record_id}' from table '{table_name}'") 

324 

325 def get_record( 

326 self, table_name: str, record_id: str | int 

327 ) -> dict[str, Any] | None: 

328 """Get a specific record by ID. 

329 

330 Args: 

331 table_name: Name of the table 

332 record_id: ID of the record 

333 

334 Returns: 

335 Record data or None if not found 

336 """ 

337 if table_name not in self.table_schemas: 

338 raise ValueError(f"Table '{table_name}' does not exist") 

339 

340 schema = self.table_schemas[table_name] 

341 table = self.db.open_table(table_name) 

342 

343 try: 

344 result = ( 

345 table.search().where(f"{schema.id_column} = '{record_id}'").to_pandas() 

346 ) 

347 if len(result) == 0: 

348 return None 

349 

350 # Convert first row to dict and remove internal columns 

351 record = result.iloc[0].to_dict() 

352 return {k: v for k, v in record.items() if not k.startswith("_")} 

353 except Exception as e: 

354 logger.error( 

355 f"Failed to get record '{record_id}' from table '{table_name}': {e}" 

356 ) 

357 return None 

358 

359 def similarity_search( 

360 self, 

361 table_name: str, 

362 query_embedding: np.ndarray, 

363 threshold: float, 

364 limit: int | None = None, 

365 ) -> list[SearchResult]: 

366 """Search for records by similarity threshold. 

367 

368 Args: 

369 table_name: Name of the table 

370 query_embedding: Query embedding vector 

371 threshold: Minimum similarity threshold 

372 limit: Maximum number of results 

373 

374 Returns: 

375 List of search results 

376 """ 

377 if table_name not in self.table_schemas: 

378 raise ValueError(f"Table '{table_name}' does not exist") 

379 

380 schema = self.table_schemas[table_name] 

381 table = self.db.open_table(table_name) 

382 

383 try: 

384 # Perform vector search with explicit vector column 

385 search_query = table.search(query_embedding.tolist(), vector_column_name="_embedding") 

386 

387 if limit: 

388 search_query = search_query.limit(limit) 

389 

390 results = search_query.to_pandas() 

391 

392 # Filter by threshold and convert to SearchResult 

393 search_results = [] 

394 for _, row in results.iterrows(): 

395 # LanceDB returns distance, convert to similarity (1 - normalized distance) 

396 similarity = max(0.0, min(1.0, 1.0 - row["_distance"])) 

397 

398 if similarity >= threshold: 

399 # Extract original data (exclude internal columns) 

400 data = {k: v for k, v in row.items() if not k.startswith("_")} 

401 

402 search_results.append( 

403 SearchResult( 

404 id=row[schema.id_column], 

405 data=data, 

406 similarity_score=similarity, 

407 ) 

408 ) 

409 

410 logger.info(f"Similarity search returned {len(search_results)} results") 

411 return search_results 

412 

413 except Exception as e: 

414 logger.error(f"Similarity search failed for table '{table_name}': {e}") 

415 return [] 

416 

417 def top_k_search( 

418 self, table_name: str, query_embedding: np.ndarray, k: int 

419 ) -> list[SearchResult]: 

420 """Search for top-k most similar records. 

421 

422 Args: 

423 table_name: Name of the table 

424 query_embedding: Query embedding vector 

425 k: Number of top results to return 

426 

427 Returns: 

428 List of search results 

429 """ 

430 if table_name not in self.table_schemas: 

431 raise ValueError(f"Table '{table_name}' does not exist") 

432 

433 schema = self.table_schemas[table_name] 

434 table = self.db.open_table(table_name) 

435 

436 try: 

437 # Perform vector search with explicit vector column 

438 results = table.search(query_embedding.tolist(), vector_column_name="_embedding").limit(k).to_pandas() 

439 

440 # Convert to SearchResult 

441 search_results = [] 

442 for _, row in results.iterrows(): 

443 # LanceDB returns distance, convert to similarity 

444 similarity = max(0.0, min(1.0, 1.0 - row["_distance"])) 

445 

446 # Extract original data (exclude internal columns) 

447 data = {k: v for k, v in row.items() if not k.startswith("_")} 

448 

449 search_results.append( 

450 SearchResult( 

451 id=row[schema.id_column], data=data, similarity_score=similarity 

452 ) 

453 ) 

454 

455 logger.info(f"Top-k search returned {len(search_results)} results") 

456 return search_results 

457 

458 except Exception as e: 

459 logger.error(f"Top-k search failed for table '{table_name}': {e}") 

460 return [] 

461 

462 def _validate_record(self, record: DataRecord, schema: TableSchema) -> None: 

463 """Validate a record against the table schema. 

464 

465 Args: 

466 record: Data record to validate 

467 schema: Table schema 

468 

469 Raises: 

470 ValueError: If validation fails 

471 """ 

472 # Check if all required columns are present 

473 for col_name in schema.columns.keys(): 

474 if col_name not in record.data: 

475 raise ValueError(f"Missing required column: {col_name}") 

476 

477 # Check for extra columns 

478 for col_name in record.data.keys(): 

479 if col_name not in schema.columns: 

480 raise ValueError(f"Unknown column: {col_name}") 

481 

482 # Validate ID column 

483 try: 

484 record.get_id(schema.id_column) 

485 except ValueError as e: 

486 raise ValueError(f"Invalid ID: {e}") 

487 

488 # Validate embedding column 

489 try: 

490 text = record.get_text_for_embedding(schema.embedding_column) 

491 if not text.strip(): 

492 raise ValueError("Embedding column cannot be empty") 

493 except ValueError as e: 

494 raise ValueError(f"Invalid embedding column: {e}") 

495 

496 

497# Global vector database instance 

498vectordb = VectorDBService()