Vito
3295ff8b20
* feat(maintenance): repair old Pair prompts' empty category and tags (#550) Squashed from PR #579. Closes #550. * feat(tagging): saved instructions with a default per engine (#431) Squashed from PR #587. * feat(tagging): pick saved instructions in the tagging dialogs (#431) Squashed from PR #588. Closes #431.
182 lines
7.8 KiB
Python
182 lines
7.8 KiB
Python
"""Pair category and tags repair through the maintenance routes (#550).
|
|
|
|
``pair_metadata`` is reviewed and applied like the other maintenance plans:
|
|
``options`` must be ``{}``, the job is titled "Pair category and tags", the
|
|
safety backup follows the ``maintenance.safety_backup`` policy, and the job's
|
|
downloadable report lists every filled prompt, never cut at 200 rows.
|
|
"""
|
|
|
|
import asyncio
|
|
import os
|
|
import sys
|
|
from unittest.mock import MagicMock
|
|
|
|
_mock_server = MagicMock()
|
|
_mock_server.PromptServer.instance.routes = MagicMock()
|
|
sys.modules.setdefault("server", _mock_server)
|
|
|
|
sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
|
|
|
|
from py.api.jobs import MAX_OUTCOME_ROWS # noqa: E402
|
|
from py.api.maintenance import MAINTENANCE_TITLES # noqa: E402
|
|
from py.config import PromptManagerConfig # noqa: E402
|
|
|
|
try:
|
|
from tests.test_api_maintenance import WAIT, ApplyTestCase # noqa: E402
|
|
except ImportError: # discovered with tests/ as the top-level directory
|
|
from test_api_maintenance import WAIT, ApplyTestCase # noqa: E402
|
|
|
|
OPERATION = "pair_metadata"
|
|
TITLE = "Pair category and tags"
|
|
|
|
|
|
class PairMetadataAPITestCase(ApplyTestCase):
|
|
def pair(self, positive, negative="", **metadata):
|
|
pair_id = self.api.db.save_prompt_pair(positive, negative, **metadata)
|
|
with self.api.db.model.get_connection() as conn:
|
|
prompt = conn.execute(
|
|
"SELECT positive_prompt_id FROM prompt_pairs WHERE id=?", (pair_id,)
|
|
).fetchone()[0]
|
|
return pair_id, prompt
|
|
|
|
def category(self, prompt_id):
|
|
with self.api.db.model.get_connection() as conn:
|
|
return conn.execute(
|
|
"SELECT category FROM prompts WHERE id=?", (prompt_id,)
|
|
).fetchone()[0]
|
|
|
|
|
|
class TestPairMetadataRoutes(PairMetadataAPITestCase):
|
|
async def test_preview_accepts_the_operation_and_no_options(self):
|
|
self.assertEqual(MAINTENANCE_TITLES[OPERATION], TITLE)
|
|
self.pair("a cat", category="portrait", tags=["a"])
|
|
plan = await self.preview(OPERATION)
|
|
self.assertEqual(plan["operation"], OPERATION)
|
|
self.assertEqual(plan["scope"]["count"], 1)
|
|
self.assertEqual(plan["backup"]["policy"], "on")
|
|
refused = await self.preview(OPERATION, 400, options={"all": True})
|
|
self.assertEqual(refused["code"], "invalid-request")
|
|
|
|
async def test_apply_fills_with_a_backup_and_names_it(self):
|
|
_pair, prompt = self.pair("a cat", "neg", category="portrait", tags=["a", "b"])
|
|
plan, summary, job = await self.run_apply(OPERATION)
|
|
self.assertEqual(summary["title"], TITLE)
|
|
self.assertEqual(summary["admission"], "records")
|
|
self.assertEqual(job.state, "finished", job.log)
|
|
self.assertEqual(self.category(prompt), "portrait")
|
|
[backup] = self.backups()
|
|
self.assertEqual(
|
|
job.status_line,
|
|
"Filled 1 prompt from its saved pairs (1 category, 1 tag set) · "
|
|
f"Safety backup {backup}",
|
|
)
|
|
self.assertEqual(
|
|
job.outcomes["changed"][0]["detail"],
|
|
f"Prompt {prompt}: category “portrait”; tags a, b",
|
|
)
|
|
self.assertEqual(job.counters["changed"], 1)
|
|
|
|
async def test_protected_prompts_are_counted(self):
|
|
_a, prompt = self.pair("conflict", "n1", category="x")
|
|
self.pair("conflict", "n2", category="y")
|
|
self.pair("fill me", category="portrait")
|
|
plan, _summary, job = await self.run_apply(OPERATION)
|
|
self.assertEqual(plan["protected"]["reasons"], ["1 where pairs disagree"])
|
|
self.assertEqual(job.counters["protected"], 1)
|
|
self.assertIsNone(self.category(prompt))
|
|
# Live test (#550): the count comes with its rows, in the review's wording.
|
|
self.assertEqual(
|
|
job.outcomes["protected"],
|
|
[
|
|
{
|
|
"path": "",
|
|
"root": "",
|
|
"root_index": None,
|
|
"detail": (
|
|
f"Prompt {prompt}: kept, pairs disagree on category: x, y"
|
|
),
|
|
}
|
|
],
|
|
)
|
|
report = job.report()["outcomes"]["protected"]
|
|
self.assertEqual((report["count"], report["truncated"]), (1, False))
|
|
|
|
async def test_report_lists_every_kept_prompt(self):
|
|
PromptManagerConfig.SAFETY_BACKUP = False
|
|
count = MAX_OUTCOME_ROWS + 5
|
|
for n in range(count):
|
|
_pair, prompt = self.pair(f"kept {n:03d}", category="portrait")
|
|
with self.api.db.model.get_connection() as conn:
|
|
conn.execute(
|
|
"UPDATE prompts SET updated_at='2099-01-01T00:00:00+00:00' "
|
|
"WHERE id=?",
|
|
(prompt,),
|
|
)
|
|
_plan, _summary, job = await self.run_apply(OPERATION)
|
|
self.assertEqual(job.counters["protected"], count)
|
|
self.assertEqual(len(job.outcomes["protected"]), MAX_OUTCOME_ROWS)
|
|
protected = job.report()["outcomes"]["protected"]
|
|
self.assertEqual(
|
|
(len(protected["rows"]), protected["truncated"]), (count, False)
|
|
)
|
|
self.assertTrue(
|
|
protected["rows"][0]["detail"].endswith(
|
|
": kept, edited after it was saved; its empty category or tags "
|
|
"may be intentional"
|
|
)
|
|
)
|
|
|
|
async def test_policy_off_writes_without_a_backup(self):
|
|
PromptManagerConfig.SAFETY_BACKUP = False
|
|
_pair, prompt = self.pair("no backup", category="portrait")
|
|
plan, _summary, job = await self.run_apply(OPERATION)
|
|
self.assertEqual(plan["backup"]["policy"], "off")
|
|
self.assertEqual(job.state, "finished", job.log)
|
|
self.assertIn("No safety backup (the policy is off)", "\n".join(job.log))
|
|
self.assertTrue(job.status_line.endswith("· No safety backup"))
|
|
self.assertEqual(self.backups(), [])
|
|
self.assertEqual(self.category(prompt), "portrait")
|
|
|
|
async def test_reviewed_with_the_policy_on_keeps_the_backup(self):
|
|
self.pair("promise", category="portrait")
|
|
plan = await self.preview(OPERATION)
|
|
PromptManagerConfig.SAFETY_BACKUP = False
|
|
body = await self.apply(plan["token"])
|
|
job = self.api._jobs.get(body["job"]["id"])
|
|
await asyncio.wait_for(asyncio.shield(job.task), WAIT)
|
|
self.assertEqual(len(self.backups()), 1)
|
|
|
|
async def test_plan_changed_writes_nothing(self):
|
|
pair_id, prompt = self.pair("drift", category="portrait")
|
|
plan = await self.preview(OPERATION)
|
|
self.api.db.update_prompt_pair_metadata(pair_id, category="landscape")
|
|
body = await self.apply(plan["token"])
|
|
job = self.api._jobs.get(body["job"]["id"])
|
|
await asyncio.wait_for(asyncio.shield(job.task), WAIT)
|
|
self.assertEqual(job.state, "failed")
|
|
self.assertEqual(job.last_event["code"], "plan-changed")
|
|
self.assertIsNone(self.category(prompt))
|
|
|
|
async def test_report_lists_every_filled_prompt(self):
|
|
PromptManagerConfig.SAFETY_BACKUP = False
|
|
count = MAX_OUTCOME_ROWS + 5
|
|
for n in range(count):
|
|
self.pair(f"bulk {n:03d}", category="portrait")
|
|
_plan, _summary, job = await self.run_apply(OPERATION)
|
|
self.assertEqual(job.state, "finished", job.log)
|
|
self.assertEqual(len(job.outcomes["changed"]), MAX_OUTCOME_ROWS)
|
|
response = await self.client.request(
|
|
"GET", f"/prompt_manager/jobs/{job.id}/report"
|
|
)
|
|
self.assertEqual(response.status, 200)
|
|
report = await response.json()
|
|
changed = report["outcomes"]["changed"]
|
|
self.assertEqual(changed["count"], count)
|
|
self.assertEqual(len(changed["rows"]), count)
|
|
self.assertFalse(changed["truncated"])
|
|
self.assertTrue(
|
|
all(
|
|
row["detail"].endswith("category “portrait”") for row in changed["rows"]
|
|
)
|
|
)
|