mirror of
https://github.com/allaunthefox/Research-Stack.git
synced 2026-07-31 03:05:21 +00:00
456 lines
14 KiB
Python
456 lines
14 KiB
Python
# SPDX-License-Identifier: Apache-2.0
|
|
# Copyright (c) 2025 Allaun
|
|
|
|
"""Embedding repository for database operations.
|
|
|
|
This module provides embedding repository functionality for the database layer,
|
|
handling file embedding storage and retrieval.
|
|
|
|
Key Features:
|
|
- Embedding CRUD operations
|
|
- Model version management
|
|
- Batch operations
|
|
- Error handling
|
|
|
|
Dependencies:
|
|
- sqlite3 (standard library only)
|
|
- typing (standard library only)
|
|
"""
|
|
|
|
import json
|
|
from typing import Optional, List, Dict, Any
|
|
from .connection import DatabaseConnection
|
|
|
|
|
|
def _serialize_embedding(embedding: Any) -> str:
|
|
"""Serialize embedding to JSON string.
|
|
|
|
Args:
|
|
embedding: Embedding data (list, dict, or numpy array-like)
|
|
|
|
Returns:
|
|
JSON string representation
|
|
"""
|
|
# Handle numpy arrays and other array-like objects
|
|
if hasattr(embedding, 'tolist'):
|
|
embedding = embedding.tolist()
|
|
elif hasattr(embedding, '__iter__') and not isinstance(embedding, (str, bytes)):
|
|
embedding = list(embedding)
|
|
|
|
return json.dumps(embedding)
|
|
|
|
|
|
def _deserialize_embedding(data: str) -> Any:
|
|
"""Deserialize embedding from JSON string.
|
|
|
|
Args:
|
|
data: JSON string representation
|
|
|
|
Returns:
|
|
Deserialized embedding data
|
|
"""
|
|
if isinstance(data, (bytes, bytearray)):
|
|
try:
|
|
data = data.decode('utf-8')
|
|
except Exception:
|
|
data = data.decode('latin-1')
|
|
return json.loads(data)
|
|
|
|
|
|
class EmbeddingRepository:
|
|
"""Embedding repository for database operations.
|
|
|
|
Responsibilities:
|
|
- Manage file embeddings in database
|
|
- Handle embedding CRUD operations
|
|
- Manage model versions
|
|
- Support batch operations
|
|
"""
|
|
|
|
def __init__(self, db_connection: DatabaseConnection):
|
|
"""Initialize embedding repository.
|
|
|
|
Args:
|
|
db_connection: Database connection instance
|
|
"""
|
|
self.db = db_connection
|
|
|
|
def _get_embedding_dimensions(self, embedding: Any) -> int:
|
|
"""Return the number of dimensions for an embedding. Non-sequences -> 0."""
|
|
try:
|
|
if hasattr(embedding, 'shape'):
|
|
return int(getattr(embedding, 'shape')[-1])
|
|
if hasattr(embedding, '__len__') and not isinstance(embedding, (str, bytes)):
|
|
return len(embedding)
|
|
except Exception:
|
|
return 0
|
|
return 0
|
|
|
|
def add_embedding(self, file_id: int, embedding: Any, model_version: str, created_time: int) -> Optional[int]:
|
|
"""Add embedding to database.
|
|
|
|
Args:
|
|
file_id: File ID
|
|
embedding: Embedding data
|
|
model_version: Model version
|
|
created_time: Creation timestamp
|
|
|
|
Returns:
|
|
Embedding ID
|
|
"""
|
|
try:
|
|
# Compute embedding dimensions before serialization
|
|
try:
|
|
if hasattr(embedding, 'shape'):
|
|
# numpy-like
|
|
dims = int(getattr(embedding, 'shape')[-1])
|
|
elif hasattr(embedding, '__len__') and not isinstance(embedding, (str, bytes)):
|
|
dims = len(embedding)
|
|
else:
|
|
dims = 0
|
|
except Exception:
|
|
dims = 0
|
|
|
|
# Serialize embedding to JSON string (safer than pickle)
|
|
embedding_str = _serialize_embedding(embedding)
|
|
|
|
cursor = self.db.execute(
|
|
'''
|
|
INSERT INTO embeddings (file_id, embedding, model_version, created_time, dimensions)
|
|
VALUES (?, ?, ?, ?, ?)
|
|
''',
|
|
(file_id, embedding_str.encode('utf-8'), model_version, created_time, dims)
|
|
)
|
|
return cursor.lastrowid
|
|
except Exception as e:
|
|
print(f"[ERROR] Failed to add embedding: {e}")
|
|
raise
|
|
|
|
def get_embedding(self, embedding_id: int) -> Optional[Dict[str, Any]]:
|
|
"""Get embedding by ID.
|
|
|
|
Args:
|
|
embedding_id: Embedding ID
|
|
|
|
Returns:
|
|
Embedding data or None if not found
|
|
"""
|
|
try:
|
|
cursor = self.db.execute(
|
|
'SELECT * FROM embeddings WHERE id = ?',
|
|
(embedding_id,)
|
|
)
|
|
row = cursor.fetchone()
|
|
if row:
|
|
return {
|
|
'id': row[0],
|
|
'file_id': row[1],
|
|
'embedding': _deserialize_embedding(row[2]),
|
|
'model_version': row[3],
|
|
'created_time': row[4],
|
|
'dimensions': row[5] if len(row) > 5 else None
|
|
}
|
|
return None
|
|
except Exception as e:
|
|
print(f"[ERROR] Failed to get embedding: {e}")
|
|
raise
|
|
|
|
def get_embedding_by_file(self, file_id: int, model_version: str) -> Optional[Dict[str, Any]]:
|
|
"""Get embedding by file ID and model version.
|
|
|
|
Args:
|
|
file_id: File ID
|
|
model_version: Model version
|
|
|
|
Returns:
|
|
Embedding data or None if not found
|
|
"""
|
|
try:
|
|
cursor = self.db.execute(
|
|
'SELECT * FROM embeddings WHERE file_id = ? AND model_version = ?',
|
|
(file_id, model_version)
|
|
)
|
|
row = cursor.fetchone()
|
|
if row:
|
|
return {
|
|
'id': row[0],
|
|
'file_id': row[1],
|
|
'embedding': _deserialize_embedding(row[2]),
|
|
'model_version': row[3],
|
|
'created_time': row[4],
|
|
'dimensions': row[5] if len(row) > 5 else None
|
|
}
|
|
return None
|
|
except Exception as e:
|
|
print(f"[ERROR] Failed to get embedding by file: {e}")
|
|
raise
|
|
|
|
def get_embeddings_by_file(self, file_id: int) -> List[Dict[str, Any]]:
|
|
"""Get all embeddings for a file.
|
|
|
|
Args:
|
|
file_id: File ID
|
|
|
|
Returns:
|
|
List of embeddings for the file
|
|
"""
|
|
try:
|
|
cursor = self.db.execute(
|
|
'SELECT * FROM embeddings WHERE file_id = ? ORDER BY model_version',
|
|
(file_id,)
|
|
)
|
|
return [
|
|
{
|
|
'id': row[0],
|
|
'file_id': row[1],
|
|
'embedding': _deserialize_embedding(row[2]),
|
|
'model_version': row[3],
|
|
'created_time': row[4],
|
|
'dimensions': row[5] if len(row) > 5 else None
|
|
}
|
|
for row in cursor.fetchall()
|
|
]
|
|
except Exception as e:
|
|
print(f"[ERROR] Failed to get embeddings by file: {e}")
|
|
raise
|
|
|
|
def get_embeddings_by_model(self, model_version: str) -> List[Dict[str, Any]]:
|
|
"""Get all embeddings for a model version.
|
|
|
|
Args:
|
|
model_version: Model version
|
|
|
|
Returns:
|
|
List of embeddings for the model
|
|
"""
|
|
try:
|
|
cursor = self.db.execute(
|
|
'SELECT * FROM embeddings WHERE model_version = ? ORDER BY file_id',
|
|
(model_version,)
|
|
)
|
|
return [
|
|
{
|
|
'id': row[0],
|
|
'file_id': row[1],
|
|
'embedding': _deserialize_embedding(row[2]),
|
|
'model_version': row[3],
|
|
'created_time': row[4],
|
|
'dimensions': row[5] if len(row) > 5 else None
|
|
}
|
|
for row in cursor.fetchall()
|
|
]
|
|
except Exception as e:
|
|
print(f"[ERROR] Failed to get embeddings by model: {e}")
|
|
raise
|
|
|
|
def update_embedding(self, embedding_id: int, embedding: Any) -> bool:
|
|
"""Update embedding data.
|
|
|
|
Args:
|
|
embedding_id: Embedding ID
|
|
embedding: New embedding data
|
|
|
|
Returns:
|
|
True if updated, False if not found
|
|
"""
|
|
try:
|
|
# Compute new dimensions and serialize
|
|
try:
|
|
if hasattr(embedding, 'shape'):
|
|
dims = int(getattr(embedding, 'shape')[-1])
|
|
elif hasattr(embedding, '__len__') and not isinstance(embedding, (str, bytes)):
|
|
dims = len(embedding)
|
|
else:
|
|
dims = 0
|
|
except Exception:
|
|
dims = 0
|
|
|
|
embedding_str = _serialize_embedding(embedding)
|
|
|
|
cursor = self.db.execute(
|
|
'UPDATE embeddings SET embedding = ?, dimensions = ? WHERE id = ?',
|
|
(embedding_str.encode('utf-8'), dims, embedding_id)
|
|
)
|
|
return cursor.rowcount > 0
|
|
except Exception as e:
|
|
print(f"[ERROR] Failed to update embedding: {e}")
|
|
raise
|
|
|
|
def delete_embedding(self, embedding_id: int) -> bool:
|
|
"""Delete embedding from database.
|
|
|
|
Args:
|
|
embedding_id: Embedding ID to delete
|
|
|
|
Returns:
|
|
True if deleted, False if not found
|
|
"""
|
|
try:
|
|
cursor = self.db.execute(
|
|
'DELETE FROM embeddings WHERE id = ?',
|
|
(embedding_id,)
|
|
)
|
|
return cursor.rowcount > 0
|
|
except Exception as e:
|
|
print(f"[ERROR] Failed to delete embedding: {e}")
|
|
raise
|
|
|
|
def delete_embeddings_by_file(self, file_id: int) -> int:
|
|
"""Delete all embeddings for a file.
|
|
|
|
Args:
|
|
file_id: File ID
|
|
|
|
Returns:
|
|
Number of embeddings deleted
|
|
"""
|
|
try:
|
|
cursor = self.db.execute(
|
|
'DELETE FROM embeddings WHERE file_id = ?',
|
|
(file_id,)
|
|
)
|
|
return cursor.rowcount
|
|
except Exception as e:
|
|
print(f"[ERROR] Failed to delete embeddings by file: {e}")
|
|
raise
|
|
|
|
def delete_embeddings_by_model(self, model_version: str) -> int:
|
|
"""Delete all embeddings for a model version.
|
|
|
|
Args:
|
|
model_version: Model version
|
|
|
|
Returns:
|
|
Number of embeddings deleted
|
|
"""
|
|
try:
|
|
cursor = self.db.execute(
|
|
'DELETE FROM embeddings WHERE model_version = ?',
|
|
(model_version,)
|
|
)
|
|
return cursor.rowcount
|
|
except Exception as e:
|
|
print(f"[ERROR] Failed to delete embeddings by model: {e}")
|
|
raise
|
|
|
|
def get_all_embeddings(self) -> List[Dict[str, Any]]:
|
|
"""Get all embeddings from database.
|
|
|
|
Returns:
|
|
List of all embeddings
|
|
"""
|
|
try:
|
|
cursor = self.db.execute('SELECT * FROM embeddings ORDER BY file_id, model_version')
|
|
return [
|
|
{
|
|
'id': row[0],
|
|
'file_id': row[1],
|
|
'embedding': _deserialize_embedding(row[2]),
|
|
'model_version': row[3],
|
|
'created_time': row[4],
|
|
'dimensions': row[5] if len(row) > 5 else None
|
|
}
|
|
for row in cursor.fetchall()
|
|
]
|
|
except Exception as e:
|
|
print(f"[ERROR] Failed to get all embeddings: {e}")
|
|
raise
|
|
|
|
def count_embeddings(self) -> int:
|
|
"""Count total embeddings in database.
|
|
|
|
Returns:
|
|
Total embedding count
|
|
"""
|
|
try:
|
|
cursor = self.db.execute('SELECT COUNT(*) FROM embeddings')
|
|
return cursor.fetchone()[0]
|
|
except Exception as e:
|
|
print(f"[ERROR] Failed to count embeddings: {e}")
|
|
raise
|
|
|
|
def count_embeddings_by_model(self, model_version: str) -> int:
|
|
"""Count embeddings for a model version.
|
|
|
|
Args:
|
|
model_version: Model version
|
|
|
|
Returns:
|
|
Embedding count for the model
|
|
"""
|
|
try:
|
|
cursor = self.db.execute(
|
|
'SELECT COUNT(*) FROM embeddings WHERE model_version = ?',
|
|
(model_version,)
|
|
)
|
|
return cursor.fetchone()[0]
|
|
except Exception as e:
|
|
print(f"[ERROR] Failed to count embeddings by model: {e}")
|
|
raise
|
|
|
|
def batch_add_embeddings(self, embeddings: List[Dict[str, Any]]) -> int:
|
|
"""Add multiple embeddings in batch.
|
|
|
|
Args:
|
|
embeddings: List of embedding data dictionaries
|
|
|
|
Returns:
|
|
Number of embeddings added
|
|
"""
|
|
if not embeddings:
|
|
return 0
|
|
|
|
try:
|
|
data = []
|
|
for emb_data in embeddings:
|
|
emb = emb_data['embedding']
|
|
try:
|
|
if hasattr(emb, 'shape'):
|
|
d = int(getattr(emb, 'shape')[-1])
|
|
elif hasattr(emb, '__len__') and not isinstance(emb, (str, bytes)):
|
|
d = len(emb)
|
|
else:
|
|
d = 0
|
|
except Exception:
|
|
d = 0
|
|
|
|
data.append((
|
|
emb_data['file_id'],
|
|
_serialize_embedding(emb).encode('utf-8'),
|
|
emb_data['model_version'],
|
|
emb_data['created_time'],
|
|
d
|
|
))
|
|
|
|
self.db.executemany(
|
|
'''INSERT INTO embeddings
|
|
(file_id, embedding, model_version, created_time, dimensions)
|
|
VALUES (?, ?, ?, ?, ?)''',
|
|
[tuple(item) for item in data]
|
|
)
|
|
return len(embeddings)
|
|
except Exception as e: # pylint: disable=broad-exception-caught
|
|
print(f"[ERROR] Failed to batch add embeddings: {e}")
|
|
raise
|
|
|
|
def clear_all_embeddings(self) -> None:
|
|
"""Clear all embeddings from database."""
|
|
try:
|
|
self.db.execute('DELETE FROM embeddings')
|
|
self.db.commit()
|
|
except Exception as e:
|
|
print(f"[ERROR] Failed to clear all embeddings: {e}")
|
|
raise
|
|
|
|
|
|
def get_embedding_repository(db_path: str = "output/index.db") -> EmbeddingRepository:
|
|
"""Get embedding repository instance.
|
|
|
|
Args:
|
|
db_path: Path to SQLite database file
|
|
|
|
Returns:
|
|
EmbeddingRepository instance
|
|
"""
|
|
db_connection = DatabaseConnection.get_instance(db_path)
|
|
return EmbeddingRepository(db_connection)
|