feat(database): add prompt storage and retrieval system
This commit is contained in:
Vito Sansevero
@@ -0,0 +1,6 @@
|
||||
"""Database module for KikoTextEncode prompt storage and management."""
|
||||
|
||||
from .operations import PromptDatabase
|
||||
from .models import PromptModel
|
||||
|
||||
__all__ = ["PromptDatabase", "PromptModel"]
|
||||
@@ -0,0 +1,209 @@
|
||||
"""
|
||||
Database schema and models for KikoTextEncode prompt storage.
|
||||
"""
|
||||
|
||||
import sqlite3
|
||||
import os
|
||||
from typing import Optional
|
||||
|
||||
|
||||
class PromptModel:
|
||||
"""Database model for prompt storage and schema management."""
|
||||
|
||||
def __init__(self, db_path: str = "prompts.db"):
|
||||
"""
|
||||
Initialize the database model.
|
||||
|
||||
Args:
|
||||
db_path: Path to the SQLite database file
|
||||
"""
|
||||
self.db_path = db_path
|
||||
self._ensure_database_exists()
|
||||
|
||||
def _ensure_database_exists(self) -> None:
|
||||
"""Create database and tables if they don't exist."""
|
||||
try:
|
||||
with sqlite3.connect(self.db_path) as conn:
|
||||
conn.execute("PRAGMA foreign_keys = ON")
|
||||
self._create_tables(conn)
|
||||
self._create_indexes(conn)
|
||||
conn.commit()
|
||||
except Exception as e:
|
||||
print(f"Error creating database: {e}")
|
||||
raise
|
||||
|
||||
def _create_tables(self, conn: sqlite3.Connection) -> None:
|
||||
"""Create the prompts table with all required columns."""
|
||||
conn.execute("""
|
||||
CREATE TABLE IF NOT EXISTS prompts (
|
||||
id INTEGER PRIMARY KEY AUTOINCREMENT,
|
||||
text TEXT NOT NULL,
|
||||
created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP,
|
||||
updated_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP,
|
||||
category TEXT,
|
||||
tags TEXT,
|
||||
rating INTEGER CHECK(rating >= 1 AND rating <= 5),
|
||||
notes TEXT,
|
||||
hash TEXT UNIQUE
|
||||
)
|
||||
""")
|
||||
|
||||
# Create images table for gallery functionality
|
||||
conn.execute("""
|
||||
CREATE TABLE IF NOT EXISTS generated_images (
|
||||
id INTEGER PRIMARY KEY AUTOINCREMENT,
|
||||
prompt_id TEXT NOT NULL,
|
||||
image_path TEXT NOT NULL,
|
||||
filename TEXT NOT NULL,
|
||||
generation_time TIMESTAMP DEFAULT CURRENT_TIMESTAMP,
|
||||
file_size INTEGER,
|
||||
width INTEGER,
|
||||
height INTEGER,
|
||||
format TEXT,
|
||||
workflow_data TEXT,
|
||||
prompt_metadata TEXT,
|
||||
parameters TEXT,
|
||||
FOREIGN KEY (prompt_id) REFERENCES prompts(id)
|
||||
)
|
||||
""")
|
||||
|
||||
# Check if we need to migrate from old schema with workflow_name
|
||||
self._migrate_workflow_name_removal(conn)
|
||||
|
||||
def _create_indexes(self, conn: sqlite3.Connection) -> None:
|
||||
"""Create indexes for better query performance."""
|
||||
indexes = [
|
||||
"CREATE INDEX IF NOT EXISTS idx_prompts_text ON prompts(text)",
|
||||
"CREATE INDEX IF NOT EXISTS idx_prompts_category ON prompts(category)",
|
||||
"CREATE INDEX IF NOT EXISTS idx_prompts_created_at ON prompts(created_at)",
|
||||
"CREATE INDEX IF NOT EXISTS idx_prompts_hash ON prompts(hash)",
|
||||
"CREATE INDEX IF NOT EXISTS idx_prompts_rating ON prompts(rating)",
|
||||
"CREATE INDEX IF NOT EXISTS idx_prompt_images ON generated_images(prompt_id)",
|
||||
"CREATE INDEX IF NOT EXISTS idx_image_path ON generated_images(image_path)",
|
||||
"CREATE INDEX IF NOT EXISTS idx_generation_time ON generated_images(generation_time)",
|
||||
]
|
||||
|
||||
for index_sql in indexes:
|
||||
conn.execute(index_sql)
|
||||
|
||||
def get_connection(self) -> sqlite3.Connection:
|
||||
"""
|
||||
Get a database connection with proper configuration.
|
||||
|
||||
Returns:
|
||||
sqlite3.Connection: Configured database connection
|
||||
"""
|
||||
conn = sqlite3.connect(self.db_path)
|
||||
conn.row_factory = sqlite3.Row # Enable dict-like access to rows
|
||||
conn.execute("PRAGMA foreign_keys = ON")
|
||||
return conn
|
||||
|
||||
def _migrate_workflow_name_removal(self, conn: sqlite3.Connection) -> None:
|
||||
"""Remove workflow_name column if it exists in existing database."""
|
||||
try:
|
||||
# Check if workflow_name column exists
|
||||
cursor = conn.execute("PRAGMA table_info(prompts)")
|
||||
columns = [column[1] for column in cursor.fetchall()]
|
||||
|
||||
if 'workflow_name' in columns:
|
||||
print("[KikoTextEncode] Migrating database: removing workflow_name column")
|
||||
|
||||
# Create new table without workflow_name
|
||||
conn.execute("""
|
||||
CREATE TABLE prompts_new (
|
||||
id INTEGER PRIMARY KEY AUTOINCREMENT,
|
||||
text TEXT NOT NULL,
|
||||
created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP,
|
||||
updated_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP,
|
||||
category TEXT,
|
||||
tags TEXT,
|
||||
rating INTEGER CHECK(rating >= 1 AND rating <= 5),
|
||||
notes TEXT,
|
||||
hash TEXT UNIQUE
|
||||
)
|
||||
""")
|
||||
|
||||
# Copy data from old table to new table
|
||||
conn.execute("""
|
||||
INSERT INTO prompts_new (id, text, created_at, updated_at, category, tags, rating, notes, hash)
|
||||
SELECT id, text, created_at, updated_at, category, tags, rating, notes, hash
|
||||
FROM prompts
|
||||
""")
|
||||
|
||||
# Drop old table and rename new one
|
||||
conn.execute("DROP TABLE prompts")
|
||||
conn.execute("ALTER TABLE prompts_new RENAME TO prompts")
|
||||
|
||||
print("[KikoTextEncode] Database migration completed")
|
||||
|
||||
except Exception as e:
|
||||
print(f"[KikoTextEncode] Migration error: {e}")
|
||||
# If migration fails, the table creation will handle it
|
||||
|
||||
def migrate_database(self) -> None:
|
||||
"""Apply any pending database migrations."""
|
||||
# Future migrations can be added here
|
||||
pass
|
||||
|
||||
def vacuum_database(self) -> None:
|
||||
"""Optimize database by running VACUUM command."""
|
||||
try:
|
||||
with sqlite3.connect(self.db_path) as conn:
|
||||
conn.execute("VACUUM")
|
||||
conn.commit()
|
||||
except Exception as e:
|
||||
print(f"Error vacuuming database: {e}")
|
||||
|
||||
def get_database_info(self) -> dict:
|
||||
"""
|
||||
Get information about the database.
|
||||
|
||||
Returns:
|
||||
dict: Database statistics and information
|
||||
"""
|
||||
try:
|
||||
with self.get_connection() as conn:
|
||||
cursor = conn.execute("SELECT COUNT(*) as total_prompts FROM prompts")
|
||||
total_prompts = cursor.fetchone()['total_prompts']
|
||||
|
||||
cursor = conn.execute(
|
||||
"SELECT COUNT(DISTINCT category) as unique_categories FROM prompts WHERE category IS NOT NULL"
|
||||
)
|
||||
unique_categories = cursor.fetchone()['unique_categories']
|
||||
|
||||
cursor = conn.execute(
|
||||
"SELECT AVG(rating) as avg_rating FROM prompts WHERE rating IS NOT NULL"
|
||||
)
|
||||
avg_rating = cursor.fetchone()['avg_rating']
|
||||
|
||||
# Get database file size
|
||||
db_size = os.path.getsize(self.db_path) if os.path.exists(self.db_path) else 0
|
||||
|
||||
return {
|
||||
'total_prompts': total_prompts,
|
||||
'unique_categories': unique_categories,
|
||||
'average_rating': round(avg_rating, 2) if avg_rating else None,
|
||||
'database_size_bytes': db_size,
|
||||
'database_path': os.path.abspath(self.db_path)
|
||||
}
|
||||
except Exception as e:
|
||||
print(f"Error getting database info: {e}")
|
||||
return {}
|
||||
|
||||
def backup_database(self, backup_path: str) -> bool:
|
||||
"""
|
||||
Create a backup of the database.
|
||||
|
||||
Args:
|
||||
backup_path: Path where the backup should be saved
|
||||
|
||||
Returns:
|
||||
bool: True if backup was successful, False otherwise
|
||||
"""
|
||||
try:
|
||||
import shutil
|
||||
shutil.copy2(self.db_path, backup_path)
|
||||
return True
|
||||
except Exception as e:
|
||||
print(f"Error creating database backup: {e}")
|
||||
return False
|
||||
@@ -0,0 +1,666 @@
|
||||
"""
|
||||
Database operations for KikoTextEncode prompt storage and retrieval.
|
||||
"""
|
||||
|
||||
import sqlite3
|
||||
import json
|
||||
import datetime
|
||||
import os
|
||||
from typing import Optional, List, Dict, Any, Union
|
||||
|
||||
from .models import PromptModel
|
||||
|
||||
|
||||
class PromptDatabase:
|
||||
"""Database operations class for managing prompts."""
|
||||
|
||||
def __init__(self, db_path: str = "prompts.db"):
|
||||
"""
|
||||
Initialize the database operations.
|
||||
|
||||
Args:
|
||||
db_path: Path to the SQLite database file
|
||||
"""
|
||||
self.model = PromptModel(db_path)
|
||||
|
||||
def save_prompt(
|
||||
self,
|
||||
text: str,
|
||||
category: Optional[str] = None,
|
||||
tags: Optional[List[str]] = None,
|
||||
rating: Optional[int] = None,
|
||||
notes: Optional[str] = None,
|
||||
prompt_hash: Optional[str] = None
|
||||
) -> int:
|
||||
"""
|
||||
Save a new prompt to the database.
|
||||
|
||||
Args:
|
||||
text: The prompt text
|
||||
category: Optional category
|
||||
tags: List of tags
|
||||
rating: Rating 1-5
|
||||
notes: Optional notes
|
||||
prompt_hash: SHA256 hash of the prompt
|
||||
|
||||
Returns:
|
||||
int: The ID of the saved prompt
|
||||
|
||||
Raises:
|
||||
ValueError: If required parameters are invalid
|
||||
sqlite3.Error: If database operation fails
|
||||
"""
|
||||
if not text or not text.strip():
|
||||
raise ValueError("Prompt text cannot be empty")
|
||||
|
||||
if rating is not None and (rating < 1 or rating > 5):
|
||||
raise ValueError("Rating must be between 1 and 5")
|
||||
|
||||
tags_json = json.dumps(tags) if tags else None
|
||||
|
||||
with self.model.get_connection() as conn:
|
||||
cursor = conn.execute(
|
||||
"""
|
||||
INSERT INTO prompts (
|
||||
text, category, tags, rating, notes, hash, created_at, updated_at
|
||||
) VALUES (?, ?, ?, ?, ?, ?, ?, ?)
|
||||
""",
|
||||
(
|
||||
text.strip(),
|
||||
category,
|
||||
tags_json,
|
||||
rating,
|
||||
notes,
|
||||
prompt_hash,
|
||||
datetime.datetime.now(datetime.timezone.utc).isoformat(),
|
||||
datetime.datetime.now(datetime.timezone.utc).isoformat()
|
||||
)
|
||||
)
|
||||
conn.commit()
|
||||
return cursor.lastrowid
|
||||
|
||||
def get_prompt_by_id(self, prompt_id: int) -> Optional[Dict[str, Any]]:
|
||||
"""
|
||||
Get a prompt by its ID.
|
||||
|
||||
Args:
|
||||
prompt_id: The prompt ID
|
||||
|
||||
Returns:
|
||||
Dict containing prompt data or None if not found
|
||||
"""
|
||||
with self.model.get_connection() as conn:
|
||||
cursor = conn.execute(
|
||||
"SELECT * FROM prompts WHERE id = ?", (prompt_id,)
|
||||
)
|
||||
row = cursor.fetchone()
|
||||
return self._row_to_dict(row) if row else None
|
||||
|
||||
def get_prompt_by_hash(self, prompt_hash: str) -> Optional[Dict[str, Any]]:
|
||||
"""
|
||||
Get a prompt by its hash.
|
||||
|
||||
Args:
|
||||
prompt_hash: SHA256 hash of the prompt
|
||||
|
||||
Returns:
|
||||
Dict containing prompt data or None if not found
|
||||
"""
|
||||
with self.model.get_connection() as conn:
|
||||
cursor = conn.execute(
|
||||
"SELECT * FROM prompts WHERE hash = ?", (prompt_hash,)
|
||||
)
|
||||
row = cursor.fetchone()
|
||||
return self._row_to_dict(row) if row else None
|
||||
|
||||
def search_prompts(
|
||||
self,
|
||||
text: Optional[str] = None,
|
||||
category: Optional[str] = None,
|
||||
tags: Optional[List[str]] = None,
|
||||
rating_min: Optional[int] = None,
|
||||
rating_max: Optional[int] = None,
|
||||
date_from: Optional[str] = None,
|
||||
date_to: Optional[str] = None,
|
||||
limit: int = 100,
|
||||
offset: int = 0
|
||||
) -> List[Dict[str, Any]]:
|
||||
"""
|
||||
Search prompts with various filters.
|
||||
|
||||
Args:
|
||||
text: Text to search for in prompt content
|
||||
category: Filter by category
|
||||
tags: Filter by tags (must contain all specified tags)
|
||||
rating_min: Minimum rating filter
|
||||
rating_max: Maximum rating filter
|
||||
date_from: Start date filter (ISO format)
|
||||
date_to: End date filter (ISO format)
|
||||
limit: Maximum number of results
|
||||
offset: Number of results to skip
|
||||
|
||||
Returns:
|
||||
List of dictionaries containing prompt data
|
||||
"""
|
||||
query_parts = ["SELECT * FROM prompts WHERE 1=1"]
|
||||
params = []
|
||||
|
||||
if text:
|
||||
query_parts.append("AND text LIKE ?")
|
||||
params.append(f"%{text}%")
|
||||
|
||||
if category:
|
||||
query_parts.append("AND category = ?")
|
||||
params.append(category)
|
||||
|
||||
if tags:
|
||||
for tag in tags:
|
||||
query_parts.append("AND tags LIKE ?")
|
||||
params.append(f"%{tag}%")
|
||||
|
||||
if rating_min is not None:
|
||||
query_parts.append("AND rating >= ?")
|
||||
params.append(rating_min)
|
||||
|
||||
if rating_max is not None:
|
||||
query_parts.append("AND rating <= ?")
|
||||
params.append(rating_max)
|
||||
|
||||
if date_from:
|
||||
query_parts.append("AND created_at >= ?")
|
||||
params.append(date_from)
|
||||
|
||||
if date_to:
|
||||
query_parts.append("AND created_at <= ?")
|
||||
params.append(date_to)
|
||||
|
||||
|
||||
query_parts.append("ORDER BY created_at DESC LIMIT ? OFFSET ?")
|
||||
params.extend([limit, offset])
|
||||
|
||||
query = " ".join(query_parts)
|
||||
|
||||
with self.model.get_connection() as conn:
|
||||
cursor = conn.execute(query, params)
|
||||
rows = cursor.fetchall()
|
||||
return [self._row_to_dict(row) for row in rows]
|
||||
|
||||
def get_recent_prompts(self, limit: int = 10) -> List[Dict[str, Any]]:
|
||||
"""
|
||||
Get the most recent prompts.
|
||||
|
||||
Args:
|
||||
limit: Maximum number of prompts to return
|
||||
|
||||
Returns:
|
||||
List of dictionaries containing prompt data
|
||||
"""
|
||||
with self.model.get_connection() as conn:
|
||||
cursor = conn.execute(
|
||||
"SELECT * FROM prompts ORDER BY created_at DESC LIMIT ?", (limit,)
|
||||
)
|
||||
rows = cursor.fetchall()
|
||||
return [self._row_to_dict(row) for row in rows]
|
||||
|
||||
def get_prompts_by_category(self, category: str, limit: int = 100) -> List[Dict[str, Any]]:
|
||||
"""
|
||||
Get all prompts in a specific category.
|
||||
|
||||
Args:
|
||||
category: The category name
|
||||
limit: Maximum number of prompts to return
|
||||
|
||||
Returns:
|
||||
List of dictionaries containing prompt data
|
||||
"""
|
||||
with self.model.get_connection() as conn:
|
||||
cursor = conn.execute(
|
||||
"SELECT * FROM prompts WHERE category = ? ORDER BY created_at DESC LIMIT ?",
|
||||
(category, limit)
|
||||
)
|
||||
rows = cursor.fetchall()
|
||||
return [self._row_to_dict(row) for row in rows]
|
||||
|
||||
def get_top_rated_prompts(self, limit: int = 10) -> List[Dict[str, Any]]:
|
||||
"""
|
||||
Get the highest rated prompts.
|
||||
|
||||
Args:
|
||||
limit: Maximum number of prompts to return
|
||||
|
||||
Returns:
|
||||
List of dictionaries containing prompt data
|
||||
"""
|
||||
with self.model.get_connection() as conn:
|
||||
cursor = conn.execute(
|
||||
"""
|
||||
SELECT * FROM prompts
|
||||
WHERE rating IS NOT NULL
|
||||
ORDER BY rating DESC, created_at DESC
|
||||
LIMIT ?
|
||||
""",
|
||||
(limit,)
|
||||
)
|
||||
rows = cursor.fetchall()
|
||||
return [self._row_to_dict(row) for row in rows]
|
||||
|
||||
def update_prompt_metadata(
|
||||
self,
|
||||
prompt_id: int,
|
||||
category: Optional[str] = None,
|
||||
tags: Optional[List[str]] = None,
|
||||
rating: Optional[int] = None,
|
||||
notes: Optional[str] = None
|
||||
) -> bool:
|
||||
"""
|
||||
Update metadata for an existing prompt.
|
||||
|
||||
Args:
|
||||
prompt_id: The prompt ID
|
||||
category: New category
|
||||
tags: New tags list
|
||||
rating: New rating
|
||||
notes: New notes
|
||||
|
||||
Returns:
|
||||
bool: True if update was successful, False otherwise
|
||||
"""
|
||||
updates = []
|
||||
params = []
|
||||
|
||||
if category is not None:
|
||||
updates.append("category = ?")
|
||||
params.append(category)
|
||||
|
||||
if tags is not None:
|
||||
updates.append("tags = ?")
|
||||
params.append(json.dumps(tags))
|
||||
|
||||
if rating is not None:
|
||||
if rating < 1 or rating > 5:
|
||||
raise ValueError("Rating must be between 1 and 5")
|
||||
updates.append("rating = ?")
|
||||
params.append(rating)
|
||||
|
||||
if notes is not None:
|
||||
updates.append("notes = ?")
|
||||
params.append(notes)
|
||||
|
||||
|
||||
if not updates:
|
||||
return False
|
||||
|
||||
updates.append("updated_at = ?")
|
||||
params.append(datetime.datetime.now(datetime.timezone.utc).isoformat())
|
||||
params.append(prompt_id)
|
||||
|
||||
query = f"UPDATE prompts SET {', '.join(updates)} WHERE id = ?"
|
||||
|
||||
with self.model.get_connection() as conn:
|
||||
cursor = conn.execute(query, params)
|
||||
conn.commit()
|
||||
return cursor.rowcount > 0
|
||||
|
||||
def delete_prompt(self, prompt_id: int) -> bool:
|
||||
"""
|
||||
Delete a prompt by its ID.
|
||||
|
||||
Args:
|
||||
prompt_id: The prompt ID to delete
|
||||
|
||||
Returns:
|
||||
bool: True if deletion was successful, False otherwise
|
||||
"""
|
||||
with self.model.get_connection() as conn:
|
||||
cursor = conn.execute("DELETE FROM prompts WHERE id = ?", (prompt_id,))
|
||||
conn.commit()
|
||||
return cursor.rowcount > 0
|
||||
|
||||
def get_all_categories(self) -> List[str]:
|
||||
"""
|
||||
Get all unique categories from the database.
|
||||
|
||||
Returns:
|
||||
List of category names
|
||||
"""
|
||||
with self.model.get_connection() as conn:
|
||||
cursor = conn.execute(
|
||||
"SELECT DISTINCT category FROM prompts WHERE category IS NOT NULL ORDER BY category"
|
||||
)
|
||||
return [row['category'] for row in cursor.fetchall()]
|
||||
|
||||
def get_all_tags(self) -> List[str]:
|
||||
"""
|
||||
Get all unique tags from the database.
|
||||
|
||||
Returns:
|
||||
List of tag names
|
||||
"""
|
||||
all_tags = set()
|
||||
|
||||
with self.model.get_connection() as conn:
|
||||
cursor = conn.execute("SELECT tags FROM prompts WHERE tags IS NOT NULL")
|
||||
for row in cursor.fetchall():
|
||||
try:
|
||||
tags = json.loads(row['tags'])
|
||||
if isinstance(tags, list):
|
||||
all_tags.update(tags)
|
||||
except (json.JSONDecodeError, TypeError):
|
||||
continue
|
||||
|
||||
return sorted(list(all_tags))
|
||||
|
||||
def _row_to_dict(self, row: sqlite3.Row) -> Dict[str, Any]:
|
||||
"""
|
||||
Convert a database row to a dictionary with parsed JSON fields.
|
||||
|
||||
Args:
|
||||
row: SQLite row object
|
||||
|
||||
Returns:
|
||||
Dictionary representation of the row
|
||||
"""
|
||||
data = dict(row)
|
||||
|
||||
# Parse tags JSON
|
||||
if data.get('tags'):
|
||||
try:
|
||||
data['tags'] = json.loads(data['tags'])
|
||||
except (json.JSONDecodeError, TypeError):
|
||||
data['tags'] = []
|
||||
else:
|
||||
data['tags'] = []
|
||||
|
||||
return data
|
||||
|
||||
def export_prompts(self, file_path: str, format: str = "json") -> bool:
|
||||
"""
|
||||
Export all prompts to a file.
|
||||
|
||||
Args:
|
||||
file_path: Path to save the export file
|
||||
format: Export format ("json" or "csv")
|
||||
|
||||
Returns:
|
||||
bool: True if export was successful, False otherwise
|
||||
"""
|
||||
try:
|
||||
prompts = self.search_prompts(limit=10000) # Get all prompts
|
||||
|
||||
if format.lower() == "json":
|
||||
with open(file_path, 'w', encoding='utf-8') as f:
|
||||
json.dump(prompts, f, indent=2, ensure_ascii=False)
|
||||
elif format.lower() == "csv":
|
||||
import csv
|
||||
if prompts:
|
||||
with open(file_path, 'w', newline='', encoding='utf-8') as f:
|
||||
writer = csv.DictWriter(f, fieldnames=prompts[0].keys())
|
||||
writer.writeheader()
|
||||
for prompt in prompts:
|
||||
# Convert lists to strings for CSV
|
||||
row = prompt.copy()
|
||||
if isinstance(row.get('tags'), list):
|
||||
row['tags'] = ', '.join(row['tags'])
|
||||
writer.writerow(row)
|
||||
else:
|
||||
raise ValueError(f"Unsupported export format: {format}")
|
||||
|
||||
return True
|
||||
except Exception as e:
|
||||
print(f"Error exporting prompts: {e}")
|
||||
return False
|
||||
|
||||
def cleanup_duplicates(self) -> int:
|
||||
"""
|
||||
Remove duplicate prompts based on text content.
|
||||
|
||||
Returns:
|
||||
int: Number of duplicates removed
|
||||
"""
|
||||
try:
|
||||
with self.model.get_connection() as conn:
|
||||
# Find duplicates by text content (case-insensitive)
|
||||
cursor = conn.execute("""
|
||||
SELECT LOWER(TRIM(text)) as normalized_text, COUNT(*) as count,
|
||||
GROUP_CONCAT(id) as ids
|
||||
FROM prompts
|
||||
GROUP BY LOWER(TRIM(text))
|
||||
HAVING COUNT(*) > 1
|
||||
""")
|
||||
|
||||
duplicates = cursor.fetchall()
|
||||
total_removed = 0
|
||||
|
||||
for duplicate in duplicates:
|
||||
ids = duplicate['ids'].split(',')
|
||||
# Keep the first one (oldest), delete the rest
|
||||
ids_to_delete = ids[1:] # Skip the first ID
|
||||
|
||||
for id_to_delete in ids_to_delete:
|
||||
conn.execute("DELETE FROM prompts WHERE id = ?", (int(id_to_delete),))
|
||||
total_removed += 1
|
||||
|
||||
conn.commit()
|
||||
|
||||
if total_removed > 0:
|
||||
print(f"[KikoTextEncode] Removed {total_removed} duplicate prompts")
|
||||
|
||||
return total_removed
|
||||
|
||||
except Exception as e:
|
||||
print(f"Error cleaning up duplicates: {e}")
|
||||
return 0
|
||||
|
||||
# Gallery-related methods
|
||||
def link_image_to_prompt(
|
||||
self,
|
||||
prompt_id: str,
|
||||
image_path: str,
|
||||
metadata: Optional[Dict[str, Any]] = None
|
||||
) -> int:
|
||||
"""
|
||||
Link a generated image to a prompt.
|
||||
|
||||
Args:
|
||||
prompt_id: ID of the prompt that generated this image
|
||||
image_path: Full path to the image file
|
||||
metadata: Optional metadata about the image
|
||||
|
||||
Returns:
|
||||
int: The ID of the created image record
|
||||
"""
|
||||
filename = os.path.basename(image_path)
|
||||
file_info = metadata.get('file_info', {}) if metadata else {}
|
||||
|
||||
with self.model.get_connection() as conn:
|
||||
cursor = conn.execute(
|
||||
"""
|
||||
INSERT INTO generated_images
|
||||
(prompt_id, image_path, filename, file_size, width, height, format,
|
||||
workflow_data, prompt_metadata, parameters)
|
||||
VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?)
|
||||
""",
|
||||
(
|
||||
prompt_id,
|
||||
image_path,
|
||||
filename,
|
||||
file_info.get('size'),
|
||||
file_info.get('dimensions', [None, None])[0] if file_info.get('dimensions') else None,
|
||||
file_info.get('dimensions', [None, None])[1] if file_info.get('dimensions') else None,
|
||||
file_info.get('format'),
|
||||
json.dumps(metadata.get('workflow', {}) if metadata else {}),
|
||||
json.dumps(metadata.get('prompt', {}) if metadata else {}),
|
||||
json.dumps(metadata.get('parameters', {}) if metadata else {})
|
||||
)
|
||||
)
|
||||
conn.commit()
|
||||
return cursor.lastrowid
|
||||
|
||||
def get_prompt_images(self, prompt_id: str) -> List[Dict[str, Any]]:
|
||||
"""
|
||||
Get all images associated with a prompt.
|
||||
|
||||
Args:
|
||||
prompt_id: The prompt ID
|
||||
|
||||
Returns:
|
||||
List of image records
|
||||
"""
|
||||
with self.model.get_connection() as conn:
|
||||
cursor = conn.execute(
|
||||
"""
|
||||
SELECT * FROM generated_images
|
||||
WHERE prompt_id = ?
|
||||
ORDER BY generation_time DESC
|
||||
""",
|
||||
(prompt_id,)
|
||||
)
|
||||
return [self._image_row_to_dict(row) for row in cursor.fetchall()]
|
||||
|
||||
def get_recent_images(self, limit: int = 50) -> List[Dict[str, Any]]:
|
||||
"""
|
||||
Get recently generated images across all prompts.
|
||||
|
||||
Args:
|
||||
limit: Maximum number of images to return
|
||||
|
||||
Returns:
|
||||
List of image records with prompt text
|
||||
"""
|
||||
with self.model.get_connection() as conn:
|
||||
cursor = conn.execute(
|
||||
"""
|
||||
SELECT gi.*, p.text as prompt_text
|
||||
FROM generated_images gi
|
||||
LEFT JOIN prompts p ON gi.prompt_id = p.id
|
||||
ORDER BY gi.generation_time DESC
|
||||
LIMIT ?
|
||||
""",
|
||||
(limit,)
|
||||
)
|
||||
return [self._image_row_to_dict(row) for row in cursor.fetchall()]
|
||||
|
||||
def search_images_by_prompt(self, search_term: str) -> List[Dict[str, Any]]:
|
||||
"""
|
||||
Search images by prompt text.
|
||||
|
||||
Args:
|
||||
search_term: Text to search for in prompt content
|
||||
|
||||
Returns:
|
||||
List of image records with prompt text
|
||||
"""
|
||||
with self.model.get_connection() as conn:
|
||||
cursor = conn.execute(
|
||||
"""
|
||||
SELECT gi.*, p.text as prompt_text
|
||||
FROM generated_images gi
|
||||
JOIN prompts p ON gi.prompt_id = p.id
|
||||
WHERE p.text LIKE ?
|
||||
ORDER BY gi.generation_time DESC
|
||||
""",
|
||||
(f"%{search_term}%",)
|
||||
)
|
||||
return [self._image_row_to_dict(row) for row in cursor.fetchall()]
|
||||
|
||||
def get_image_by_id(self, image_id: int) -> Optional[Dict[str, Any]]:
|
||||
"""
|
||||
Get an image record by its ID.
|
||||
|
||||
Args:
|
||||
image_id: The image ID
|
||||
|
||||
Returns:
|
||||
Image record or None if not found
|
||||
"""
|
||||
with self.model.get_connection() as conn:
|
||||
cursor = conn.execute(
|
||||
"SELECT * FROM generated_images WHERE id = ?",
|
||||
(image_id,)
|
||||
)
|
||||
row = cursor.fetchone()
|
||||
return self._image_row_to_dict(row) if row else None
|
||||
|
||||
def delete_image(self, image_id: int) -> bool:
|
||||
"""
|
||||
Delete an image record by its ID.
|
||||
|
||||
Args:
|
||||
image_id: The image ID to delete
|
||||
|
||||
Returns:
|
||||
bool: True if deletion was successful
|
||||
"""
|
||||
with self.model.get_connection() as conn:
|
||||
cursor = conn.execute(
|
||||
"DELETE FROM generated_images WHERE id = ?",
|
||||
(image_id,)
|
||||
)
|
||||
conn.commit()
|
||||
return cursor.rowcount > 0
|
||||
|
||||
def cleanup_missing_images(self) -> int:
|
||||
"""
|
||||
Remove image records where the actual file no longer exists.
|
||||
|
||||
Returns:
|
||||
int: Number of orphaned records removed
|
||||
"""
|
||||
removed_count = 0
|
||||
|
||||
with self.model.get_connection() as conn:
|
||||
cursor = conn.execute("SELECT id, image_path FROM generated_images")
|
||||
images = cursor.fetchall()
|
||||
|
||||
for image in images:
|
||||
if not os.path.exists(image['image_path']):
|
||||
conn.execute("DELETE FROM generated_images WHERE id = ?", (image['id'],))
|
||||
removed_count += 1
|
||||
|
||||
conn.commit()
|
||||
|
||||
return removed_count
|
||||
|
||||
def _clean_nan_values(self, obj):
|
||||
"""
|
||||
Recursively clean NaN values from nested data structures.
|
||||
|
||||
Args:
|
||||
obj: The object to clean
|
||||
|
||||
Returns:
|
||||
Cleaned object with NaN values replaced by None
|
||||
"""
|
||||
if isinstance(obj, dict):
|
||||
return {key: self._clean_nan_values(value) for key, value in obj.items()}
|
||||
elif isinstance(obj, list):
|
||||
return [self._clean_nan_values(item) for item in obj]
|
||||
elif isinstance(obj, float) and str(obj) == 'nan':
|
||||
return None
|
||||
else:
|
||||
return obj
|
||||
|
||||
def _image_row_to_dict(self, row: sqlite3.Row) -> Dict[str, Any]:
|
||||
"""
|
||||
Convert an image database row to a dictionary with parsed JSON fields.
|
||||
|
||||
Args:
|
||||
row: SQLite row object
|
||||
|
||||
Returns:
|
||||
Dictionary representation of the row
|
||||
"""
|
||||
data = dict(row)
|
||||
|
||||
# Parse JSON fields
|
||||
for field in ['workflow_data', 'prompt_metadata', 'parameters']:
|
||||
if data.get(field):
|
||||
try:
|
||||
parsed_data = json.loads(data[field])
|
||||
data[field] = self._clean_nan_values(parsed_data)
|
||||
except (json.JSONDecodeError, TypeError):
|
||||
data[field] = {}
|
||||
else:
|
||||
data[field] = {}
|
||||
|
||||
return data
|
||||
Reference in New Issue
Block a user