Files
ComfyAssets-ComfyUI_PromptM…/tests/test_prompt_graph.py
T
Vito 6c7e4d5b35 feat: release 3.2.4 — recently used sorting, run counts, positive-prompt image linking (#152)
* chore: start 3.2.4

* feat: track prompt usage (run_count, last_used_at) with server-side sorting

- Migration adds last_used_at and run_count; existing prompts are estimated
  from linked images (count, newest image time, else created_at). The
  backfill only touches NULL rows, so it self-heals rows written after a
  downgrade or an interrupted upgrade.
- Every node run records a use, including re-runs of an existing prompt;
  a tracking failure is logged and never blocks generation.
- search_prompts/get_recent_prompts accept a whitelisted sort key
  (last_used, created, run_count, rating, text); default stays newest first.

Refs #90

* feat: sort param on /recent and /search; microsecond last_used_at

- /prompt_manager/recent and /search accept ?sort= (whitelisted in the DB
  layer; unknown values fall back to newest first).
- last_used_at is written with microseconds from Python so a re-run in the
  same millisecond as another save still sorts first; ordering stays correct
  against millisecond values from the backfill.

Refs #90

* feat: server-side sort in admin UI with Recently Used / Most Used and run badges

- New web/js/prompt-list-sort.js owns the sort options, validation and
  request URLs; admin.js fills the dropdown from it and re-fetches from the
  server on change, so every sort spans all pages (the old client-side sort
  only reordered the current page).
- Default sort is Recently Used; prompt cards show a run-count badge with
  the last-used time on hover.
- The node's Recent button now lists recently used prompts.
- Contract test keeps the UI sort keys and the node's key in sync with the
  server whitelist.

Refs #90

* fix: link images and count runs by the positive prompt, not node execution

With a negative-prompt PromptManager node in the workflow, every executed
PromptManager node pushed onto the image-linking queue, so images could be
linked to the negative prompt or to a previous run. ComfyUI also skips
unchanged nodes, so re-runs of the same prompt were never counted.

- utils/prompt_graph.py finds the PromptManager nodes feeding a sampler's
  positive input (KSampler, CFG/Basic guiders, ControlNet and combine
  pass-throughs, PromptManagerText into encoders), falling back to all
  PromptManager nodes when none can be traced.
- Runs are counted by a ComfyUI on-prompt hook when a workflow is queued;
  prompts not yet saved are counted when the node first saves them.
- Nodes receive the queued graph (hidden PROMPT/UNIQUE_ID). Only the run's
  positive prompt becomes the current prompt; only batch items with linked
  text are pushed onto the FIFO queue.
- The image monitor links from the image's own metadata first, using the
  positive node, and falls back to the queue for batch items.

Refs #90

* fix: address review — count cached runs of linked text, keep repeated pending runs

- Text from known pure string nodes (core PrimitiveString/StringConcatenate,
  WAS Text Concatenate/Multiline, pysssss ShowText) is resolved at queue
  time, so cached re-runs of such prompts are counted and their images link
  from metadata. Unknown nodes are never guessed; the node compares its real
  text with the queue-time resolution and falls back to counting and
  queueing itself on any mismatch.
- A new prompt queued several times before its first run keeps every
  pending run; the first save records all of them.
- Registration flag is a threading.Event; consistent unittest imports;
  logging instead of print in the init fallback.

* refactor: unquote PendingFirstUse annotation so Tuple usage is visible
2026-10-07 13:52:09 -07:00

245 lines
8.2 KiB
Python

"""Tests for utils.prompt_graph: which PromptManager nodes feed a positive input."""
import os
import sys
import unittest
sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
from utils.prompt_graph import (
positive_prompt_nodes,
resolve_text,
run_prompt_nodes,
)
def pm(text="a cat", **inputs):
return {"class_type": "PromptManager", "inputs": {"text": text, **inputs}}
def ksampler(positive, negative, model="20"):
return {
"class_type": "KSampler",
"inputs": {
"positive": [positive, 0],
"negative": [negative, 0],
"model": [model, 0],
},
}
LOADER = {
"class_type": "CheckpointLoaderSimple",
"inputs": {"ckpt_name": "m.safetensors"},
}
class TestPositivePromptNodes(unittest.TestCase):
def test_positive_and_negative_prompt_manager_nodes(self):
# Negative listed first, as in the #75 / pmtest workflow
graph = {
"176": pm("blurry"),
"134": pm("a cat"),
"5": ksampler("134", "176"),
"20": LOADER,
}
self.assertEqual(positive_prompt_nodes(graph), ["134"])
def test_negative_only_prompt_manager_is_never_positive(self):
graph = {
"176": pm("blurry"),
"6": {"class_type": "CLIPTextEncode", "inputs": {"text": "a cat"}},
"5": ksampler("6", "176"),
"20": LOADER,
}
self.assertEqual(positive_prompt_nodes(graph), [])
def test_through_controlnet_and_combine_pass_throughs(self):
graph = {
"1": pm("subject"),
"2": pm("style"),
"3": pm("blurry"),
"4": {
"class_type": "ConditioningCombine",
"inputs": {"conditioning_1": ["1", 0], "conditioning_2": ["2", 0]},
},
"5": {
"class_type": "ControlNetApplyAdvanced",
"inputs": {"positive": ["4", 0], "negative": ["3", 0], "strength": 1.0},
},
"6": {
"class_type": "KSampler",
"inputs": {
"positive": ["5", 0],
"negative": ["5", 1],
"model": ["20", 0],
},
},
"20": LOADER,
}
self.assertEqual(sorted(positive_prompt_nodes(graph)), ["1", "2"])
def test_cfg_guider_and_basic_guider(self):
graph = {
"1": pm("fox"),
"2": pm("watermark"),
"3": {
"class_type": "CFGGuider",
"inputs": {
"positive": ["1", 0],
"negative": ["2", 0],
"model": ["20", 0],
"cfg": 5,
},
},
"4": pm("owl"),
"5": {
"class_type": "BasicGuider",
"inputs": {"conditioning": ["4", 0], "model": ["20", 0]},
},
"20": LOADER,
}
self.assertEqual(sorted(positive_prompt_nodes(graph)), ["1", "4"])
def test_prompt_manager_text_wired_into_an_encoder(self):
graph = {
"1": {"class_type": "PromptManagerText", "inputs": {"text": "a castle"}},
"2": {"class_type": "CLIPTextEncode", "inputs": {"text": ["1", 0]}},
"3": {"class_type": "PromptManagerText", "inputs": {"text": "lowres"}},
"4": {"class_type": "CLIPTextEncode", "inputs": {"text": ["3", 0]}},
"5": ksampler("2", "4"),
"20": LOADER,
}
self.assertEqual(positive_prompt_nodes(graph), ["1"])
def test_node_used_as_both_positive_and_negative_counts_as_positive(self):
graph = {"1": pm("same"), "5": ksampler("1", "1"), "20": LOADER}
self.assertEqual(positive_prompt_nodes(graph), ["1"])
def test_two_samplers_sharing_a_prompt_list_it_once(self):
graph = {
"1": pm("hero"),
"2": pm("blurry"),
"5": ksampler("1", "2"),
"14": ksampler("1", "2"),
"20": LOADER,
}
self.assertEqual(positive_prompt_nodes(graph), ["1"])
def test_graph_without_sampler_or_malformed_input(self):
self.assertEqual(positive_prompt_nodes({"1": pm()}), [])
for bad in (None, [], "x", {"1": None}, {"1": {"inputs": None}}):
self.assertEqual(positive_prompt_nodes(bad), [], repr(bad))
def test_cycles_terminate(self):
loop = {
"class_type": "Weird",
"inputs": {
"conditioning": ["9", 0],
"conditioning_1": ["9", 0],
"positive": ["9", 0],
},
}
graph = {"9": loop, "5": ksampler("9", "9"), "20": LOADER}
self.assertEqual(positive_prompt_nodes(graph), [])
class TestRunPromptNodes(unittest.TestCase):
"""The prompt nodes that represent a run, with a safe fallback."""
def test_uses_positive_nodes_when_found(self):
graph = {
"176": pm("blurry"),
"134": pm("a cat"),
"5": ksampler("134", "176"),
"20": LOADER,
}
self.assertEqual(run_prompt_nodes(graph), ["134"])
def test_unrecognised_sampler_falls_back_to_every_prompt_manager_node(self):
graph = {
"1": pm("a cat"),
"9": {"class_type": "ExoticSampler", "inputs": {"cond": ["1", 0]}},
"2": {"class_type": "PromptManagerText", "inputs": {"text": "x"}},
}
self.assertEqual(sorted(run_prompt_nodes(graph)), ["1", "2"])
def test_malformed_graph(self):
self.assertEqual(run_prompt_nodes(None), [])
class TestResolveText(unittest.TestCase):
"""Text known at queue time: typed, or from known pure string nodes."""
def _linked(self, source):
return {"1": pm(["2", 0]), "2": source}
def test_typed_text(self):
self.assertEqual(resolve_text({"1": pm(" a cat ")}, "1"), "a cat")
def test_core_primitive_string(self):
for cls in ("PrimitiveString", "PrimitiveStringMultiline"):
graph = self._linked({"class_type": cls, "inputs": {"value": "a cat"}})
self.assertEqual(resolve_text(graph, "1"), "a cat", cls)
def test_core_string_concatenate(self):
graph = self._linked(
{
"class_type": "StringConcatenate",
"inputs": {
"string_a": "a cat",
"string_b": "at dusk",
"delimiter": ", ",
},
}
)
self.assertEqual(resolve_text(graph, "1"), "a cat, at dusk")
def test_was_text_concatenate_skips_empty_and_cleans_whitespace(self):
graph = self._linked(
{
"class_type": "Text Concatenate",
"inputs": {
"delimiter": " ",
"clean_whitespace": "true",
"text_a": " a cat ",
"text_b": "",
"text_c": ["3", 0],
},
}
)
graph["3"] = {"class_type": "PrimitiveString", "inputs": {"value": "at dusk"}}
self.assertEqual(resolve_text(graph, "1"), "a cat at dusk")
def test_pass_through_nodes(self):
for cls in ("ShowText|pysssss", "Text Multiline"):
graph = self._linked({"class_type": cls, "inputs": {"text": "a cat"}})
self.assertEqual(resolve_text(graph, "1"), "a cat", cls)
def test_unknown_and_batch_sources_are_never_guessed(self):
for source in (
{"class_type": "PromptSearchList", "inputs": {"search_text": "cat"}},
{
"class_type": "SomeTextReplace",
"inputs": {"text": "a cat", "find": "cat", "replace": "dog"},
},
):
self.assertIsNone(
resolve_text(self._linked(source), "1"), source["class_type"]
)
def test_cycles_and_missing_links(self):
graph = {
"1": pm(["2", 0]),
"2": {
"class_type": "StringConcatenate",
"inputs": {"string_a": ["2", 0], "string_b": "x", "delimiter": ""},
},
}
self.assertIsNone(resolve_text(graph, "1"))
self.assertIsNone(resolve_text({"1": pm(["9", 0])}, "1"))
if __name__ == "__main__":
unittest.main()