diff --git a/database/operations.py b/database/operations.py index f43e2ab..1c02af5 100644 --- a/database/operations.py +++ b/database/operations.py @@ -635,7 +635,9 @@ class PromptDatabase(PairOperations): """ with self.model.get_connection() as conn: cursor = conn.execute( - "SELECT DISTINCT TRIM(category) as category FROM prompts WHERE " + "SELECT TRIM(category) AS category FROM prompts WHERE " + "category IS NOT NULL AND TRIM(category) != '' UNION " + "SELECT TRIM(category) AS category FROM prompt_pairs WHERE " "category IS NOT NULL AND TRIM(category) != '' ORDER BY category" ) return [row["category"] for row in cursor.fetchall()] @@ -702,16 +704,16 @@ class PromptDatabase(PairOperations): def get_all_tags(self) -> List[str]: """ - Get all unique tags that are in use (linked to at least one prompt). + Get all unique tags linked to a standalone prompt or reusable pair. Returns: Sorted list of tag names """ with self.model.get_connection() as conn: cursor = conn.execute( - "SELECT DISTINCT t.name AS tag " - "FROM prompt_tags pt " - "JOIN tags t ON pt.tag_id = t.id " + "SELECT t.name AS tag FROM tags t WHERE " + "EXISTS(SELECT 1 FROM prompt_tags pt WHERE pt.tag_id=t.id) OR " + "EXISTS(SELECT 1 FROM prompt_pair_tags pp WHERE pp.tag_id=t.id) " "ORDER BY t.name" ) return [row["tag"] for row in cursor.fetchall()] @@ -787,7 +789,8 @@ class PromptDatabase(PairOperations): Returns tags whose name starts with ``prefix`` (case-insensitive), counted over the prompts that carry every tag in ``with_tags``; the - context tags themselves are left out. Ordered by count, then name. + context tags themselves are left out. Counts include separate standalone + and pair references, with owner types kept distinct. Ordered by count/name. """ seen = set() context = [] @@ -802,11 +805,11 @@ class PromptDatabase(PairOperations): if context: placeholders = ",".join("?" * len(context)) where.append( - "pt.prompt_id IN (" - " SELECT pt2.prompt_id FROM prompt_tags pt2" + "(pt.kind,pt.owner_id) IN (" + " SELECT pt2.kind,pt2.owner_id FROM tag_refs pt2" " JOIN tags t2 ON t2.id = pt2.tag_id" f" WHERE t2.name COLLATE NOCASE IN ({placeholders})" - " GROUP BY pt2.prompt_id" + " GROUP BY pt2.kind,pt2.owner_id" " HAVING COUNT(DISTINCT lower(t2.name)) = ?)" ) params.extend(context) @@ -816,8 +819,11 @@ class PromptDatabase(PairOperations): params.append(max(1, int(limit))) sql = ( + "WITH tag_refs(kind,owner_id,tag_id) AS (" + "SELECT 'prompt',prompt_id,tag_id FROM prompt_tags UNION ALL " + "SELECT 'pair',pair_id,tag_id FROM prompt_pair_tags) " "SELECT t.name AS name, COUNT(*) AS count" - " FROM prompt_tags pt JOIN tags t ON t.id = pt.tag_id" + " FROM tag_refs pt JOIN tags t ON t.id = pt.tag_id" f" WHERE {' AND '.join(where)}" " GROUP BY t.id" " ORDER BY count DESC, t.name COLLATE NOCASE ASC" diff --git a/tests/test_pair_vocabulary.py b/tests/test_pair_vocabulary.py new file mode 100644 index 0000000..89305ac --- /dev/null +++ b/tests/test_pair_vocabulary.py @@ -0,0 +1,106 @@ +"""Autocomplete shares vocabulary without mixing pair and prompt identities.""" + +import os +import tempfile +import unittest + +from database.operations import PromptDatabase +from utils.hashing import generate_prompt_hash + + +class TestPairVocabulary(unittest.TestCase): + def setUp(self): + directory = tempfile.TemporaryDirectory() + self.addCleanup(directory.cleanup) + self.db = PromptDatabase(os.path.join(directory.name, "vocabulary.db")) + self.addCleanup(self.db.close_all) + + def prompt(self, text: str, **kwargs) -> int: + return self.db.save_prompt( + text, prompt_hash=generate_prompt_hash(text), **kwargs + ) + + def suggestions(self, prefix="", context=None, limit=20): + return self.db.suggest_tags(prefix, context, limit) + + def test_pair_only_categories_and_tags_are_available_without_component_metadata( + self, + ): + pair = self.db.get_prompt_pair( + self.db.save_prompt_pair( + "portrait", "blurry", category="Pairs", tags=["studio"] + ) + ) + self.assertEqual(self.db.get_all_categories(), ["Pairs"]) + self.assertEqual(self.db.get_all_tags(), ["studio"]) + self.assertEqual(self.suggestions("stu"), [{"name": "studio", "count": 1}]) + for name in ("positive_prompt_id", "negative_prompt_id"): + component = self.db.get_prompt_by_id(pair[name]) + self.assertIsNone(component["category"]) + self.assertEqual(component["tags"], []) + + def test_categories_are_trimmed_deduplicated_and_include_archived_pairs(self): + self.prompt("standalone", category=" Portrait ") + pair = self.db.save_prompt_pair( + "portrait", "", category="Portrait", tags=["archived"] + ) + self.db.archive_prompt_pair(pair) + self.db.save_prompt_pair("landscape", "", category=" Landscape ") + self.db.save_prompt_pair("empty", "", category=" ") + self.assertEqual(self.db.get_all_categories(), ["Landscape", "Portrait"]) + self.assertEqual(self.db.get_all_tags(), ["archived"]) + self.assertEqual(self.suggestions("arc"), [{"name": "archived", "count": 1}]) + + def test_counts_include_distinct_pair_and_prompt_references(self): + self.prompt("portrait", tags=["shared", "legacy"]) + self.db.save_prompt_pair( + "portrait", "first negative", tags=["shared", "paired"] + ) + self.db.save_prompt_pair("portrait", "second negative", tags=["shared"]) + self.assertEqual(self.suggestions("sha"), [{"name": "shared", "count": 3}]) + self.assertEqual( + self.suggestions("sha", ["paired"]), [{"name": "shared", "count": 1}] + ) + self.assertEqual( + self.suggestions("sha", ["legacy"]), [{"name": "shared", "count": 1}] + ) + + def test_context_does_not_join_colliding_pair_and_prompt_ids(self): + prompt = self.prompt("portrait", tags=["legacy-context", "legacy-only"]) + pair = self.db.save_prompt_pair( + "portrait", "blurry", tags=["pair-context", "pair-only"] + ) + self.assertEqual(prompt, pair) + self.assertEqual( + self.suggestions("", ["legacy-context"]), + [{"name": "legacy-only", "count": 1}], + ) + self.assertEqual( + self.suggestions("", ["pair-context"]), [{"name": "pair-only", "count": 1}] + ) + self.assertEqual(self.suggestions("", ["legacy-context", "pair-context"]), []) + + def test_pair_context_matches_all_tags_case_insensitively(self): + self.db.save_prompt_pair("portrait", "blurry", tags=["One", "two", "three"]) + self.db.save_prompt_pair("other", "blurry", tags=["One", "extra"]) + self.assertEqual( + self.suggestions("TH", ["one", "TWO", "one"]), + [{"name": "three", "count": 1}], + ) + self.assertEqual(self.suggestions("", ["missing"]), []) + + def test_pair_prefix_wildcards_are_literal_and_limit_is_preserved(self): + self.db.save_prompt_pair( + "portrait", "", tags=["tag_one", "tag_two", "tag%literal"] + ) + self.assertEqual( + self.suggestions("tag_", limit=1), [{"name": "tag_one", "count": 1}] + ) + self.assertEqual( + self.suggestions("tag%"), [{"name": "tag%literal", "count": 1}] + ) + self.assertEqual(self.suggestions("%"), []) + + +if __name__ == "__main__": + unittest.main()