260 lines
8.1 KiB
Python
260 lines
8.1 KiB
Python
"""
|
|
Feedback database — SQLite-backed storage for sprite ratings.
|
|
|
|
Stores generated sprites with their prompts, params, PNG paths, and user ratings (1-5 stars).
|
|
Provides similarity search by prompt keywords for few-shot reference examples.
|
|
"""
|
|
|
|
import json
|
|
import os
|
|
import sqlite3
|
|
import time
|
|
import uuid
|
|
from dataclasses import dataclass, field
|
|
from typing import Optional
|
|
|
|
|
|
@dataclass
|
|
class FeedbackEntry:
|
|
id: str
|
|
prompt: str
|
|
params: dict
|
|
rating: int
|
|
feedback: Optional[str]
|
|
image_path: Optional[str]
|
|
created_at: str
|
|
|
|
|
|
@dataclass
|
|
class DBStats:
|
|
total: int = 0
|
|
rated: int = 0
|
|
unrated: int = 0
|
|
avg_rating: float = 0.0
|
|
|
|
|
|
@dataclass
|
|
class StyleReference:
|
|
id: str
|
|
image_path: str
|
|
name: str
|
|
role: str
|
|
notes: Optional[str]
|
|
priority: int
|
|
created_at: str
|
|
|
|
|
|
class FeedbackDB:
|
|
"""SQLite-backed feedback database for generated sprites."""
|
|
|
|
def __init__(self, conn: sqlite3.Connection):
|
|
self.conn = conn
|
|
|
|
@classmethod
|
|
def open(cls, path: str) -> "FeedbackDB":
|
|
conn = sqlite3.connect(path, check_same_thread=False)
|
|
conn.execute("""
|
|
CREATE TABLE IF NOT EXISTS feedback (
|
|
id TEXT PRIMARY KEY,
|
|
prompt TEXT NOT NULL,
|
|
params_json TEXT NOT NULL,
|
|
rating INTEGER DEFAULT 0,
|
|
feedback TEXT,
|
|
image_path TEXT,
|
|
created_at TEXT NOT NULL
|
|
)
|
|
""")
|
|
conn.execute("CREATE INDEX IF NOT EXISTS idx_prompt ON feedback(prompt)")
|
|
conn.execute("CREATE INDEX IF NOT EXISTS idx_rating ON feedback(rating)")
|
|
conn.execute("""
|
|
CREATE TABLE IF NOT EXISTS style_references (
|
|
id TEXT PRIMARY KEY,
|
|
image_path TEXT NOT NULL UNIQUE,
|
|
name TEXT NOT NULL,
|
|
role TEXT NOT NULL,
|
|
notes TEXT,
|
|
priority INTEGER NOT NULL DEFAULT 0,
|
|
created_at TEXT NOT NULL
|
|
)
|
|
""")
|
|
conn.execute("CREATE INDEX IF NOT EXISTS idx_style_priority ON style_references(priority DESC)")
|
|
conn.commit()
|
|
return cls(conn)
|
|
|
|
def add(
|
|
self,
|
|
prompt: str,
|
|
params: dict,
|
|
image_path: Optional[str] = None,
|
|
) -> str:
|
|
entry_id = str(uuid.uuid4())
|
|
params_json = json.dumps(params)
|
|
now = str(int(time.time()))
|
|
|
|
self.conn.execute(
|
|
"INSERT INTO feedback (id, prompt, params_json, rating, feedback, image_path, created_at) "
|
|
"VALUES (?, ?, ?, 0, NULL, ?, ?)",
|
|
(entry_id, prompt, params_json, image_path, now),
|
|
)
|
|
self.conn.commit()
|
|
return entry_id
|
|
|
|
def update_rating(
|
|
self,
|
|
entry_id: str,
|
|
rating: int,
|
|
feedback: Optional[str] = None,
|
|
) -> None:
|
|
rating = max(0, min(5, rating))
|
|
self.conn.execute(
|
|
"UPDATE feedback SET rating = ?, feedback = ? WHERE id = ?",
|
|
(rating, feedback, entry_id),
|
|
)
|
|
self.conn.commit()
|
|
|
|
def get_all(self) -> list[FeedbackEntry]:
|
|
cursor = self.conn.execute(
|
|
"SELECT id, prompt, params_json, rating, feedback, image_path, created_at "
|
|
"FROM feedback ORDER BY created_at DESC"
|
|
)
|
|
return [self._row_to_entry(row) for row in cursor]
|
|
|
|
def get_unrated(self) -> list[FeedbackEntry]:
|
|
cursor = self.conn.execute(
|
|
"SELECT id, prompt, params_json, rating, feedback, image_path, created_at "
|
|
"FROM feedback WHERE rating = 0 ORDER BY created_at DESC"
|
|
)
|
|
return [self._row_to_entry(row) for row in cursor]
|
|
|
|
def top_rated(self, limit: int = 10, min_rating: int = 1) -> list[FeedbackEntry]:
|
|
cursor = self.conn.execute(
|
|
"SELECT id, prompt, params_json, rating, feedback, image_path, created_at "
|
|
"FROM feedback WHERE rating >= ? ORDER BY rating DESC, created_at DESC LIMIT ?",
|
|
(min_rating, limit),
|
|
)
|
|
return [self._row_to_entry(row) for row in cursor]
|
|
|
|
def search_similar(self, query: str, limit: int = 5) -> list[FeedbackEntry]:
|
|
query_keywords = _tokenize(query)
|
|
|
|
if not query_keywords:
|
|
return self.top_rated(limit, 1)
|
|
|
|
all_entries = self.get_all()
|
|
|
|
scored = []
|
|
for entry in all_entries:
|
|
entry_keywords = _tokenize(entry.prompt)
|
|
match_count = sum(
|
|
1 for qk in query_keywords if any(ek == qk for ek in entry_keywords)
|
|
)
|
|
if match_count > 0:
|
|
scored.append((entry, match_count, entry.rating))
|
|
|
|
scored.sort(key=lambda x: (-x[1], -x[2]))
|
|
return [e for e, _, _ in scored[:limit]]
|
|
|
|
def stats(self) -> DBStats:
|
|
total = self.conn.execute("SELECT COUNT(*) FROM feedback").fetchone()[0]
|
|
rated = self.conn.execute(
|
|
"SELECT COUNT(*) FROM feedback WHERE rating > 0"
|
|
).fetchone()[0]
|
|
avg = (
|
|
self.conn.execute(
|
|
"SELECT AVG(rating) FROM feedback WHERE rating > 0"
|
|
).fetchone()[0]
|
|
or 0.0
|
|
)
|
|
|
|
return DBStats(
|
|
total=total,
|
|
rated=rated,
|
|
unrated=total - rated,
|
|
avg_rating=round(avg, 1),
|
|
)
|
|
|
|
def delete(self, entry_id: str) -> None:
|
|
self.conn.execute("DELETE FROM feedback WHERE id = ?", (entry_id,))
|
|
self.conn.commit()
|
|
|
|
def add_style_reference(
|
|
self,
|
|
image_path: str,
|
|
name: str,
|
|
role: str,
|
|
notes: Optional[str] = None,
|
|
priority: int = 0,
|
|
) -> str:
|
|
"""Register an existing image as a reusable visual style reference."""
|
|
existing = self.conn.execute(
|
|
"SELECT id FROM style_references WHERE image_path = ?", (image_path,)
|
|
).fetchone()
|
|
if existing:
|
|
self.conn.execute(
|
|
"UPDATE style_references SET name=?, role=?, notes=?, priority=? WHERE id=?",
|
|
(name, role, notes, priority, existing[0]),
|
|
)
|
|
self.conn.commit()
|
|
return existing[0]
|
|
|
|
reference_id = str(uuid.uuid4())
|
|
self.conn.execute(
|
|
"INSERT INTO style_references (id, image_path, name, role, notes, priority, created_at) "
|
|
"VALUES (?, ?, ?, ?, ?, ?, ?)",
|
|
(reference_id, image_path, name, role, notes, priority, str(int(time.time()))),
|
|
)
|
|
self.conn.commit()
|
|
return reference_id
|
|
|
|
def get_style_references(self, limit: int = 8) -> list[StyleReference]:
|
|
cursor = self.conn.execute(
|
|
"SELECT id, image_path, name, role, notes, priority, created_at "
|
|
"FROM style_references ORDER BY priority DESC, created_at DESC LIMIT ?",
|
|
(limit,),
|
|
)
|
|
return [StyleReference(*row) for row in cursor]
|
|
|
|
def delete_style_reference(self, reference_id: str) -> None:
|
|
self.conn.execute("DELETE FROM style_references WHERE id = ?", (reference_id,))
|
|
self.conn.commit()
|
|
|
|
def export_jsonl(self, path: str, min_rating: int = 4) -> int:
|
|
entries = self.top_rated(10000, min_rating)
|
|
lines = []
|
|
|
|
for entry in entries:
|
|
line = json.dumps(
|
|
{
|
|
"instruction": f"Generate a pixel-art sprite for: {entry.prompt}",
|
|
"response": entry.params,
|
|
"rating": entry.rating,
|
|
"feedback": entry.feedback,
|
|
"image_path": entry.image_path,
|
|
"prompt": entry.prompt,
|
|
}
|
|
)
|
|
lines.append(line)
|
|
|
|
with open(path, "w") as f:
|
|
f.write("\n".join(lines))
|
|
|
|
return len(entries)
|
|
|
|
def _row_to_entry(self, row: sqlite3.Row) -> FeedbackEntry:
|
|
return FeedbackEntry(
|
|
id=row[0],
|
|
prompt=row[1],
|
|
params=json.loads(row[2]),
|
|
rating=max(0, min(5, row[3])),
|
|
feedback=row[4],
|
|
image_path=row[5],
|
|
created_at=row[6],
|
|
)
|
|
|
|
|
|
def _tokenize(s: str) -> list[str]:
|
|
"""Tokenize a prompt into lowercase keywords."""
|
|
return [
|
|
w.lower() for w in s.replace("_", " ").replace("-", " ").split() if len(w) > 1
|
|
]
|