diff --git a/.gitignore b/.gitignore index c18dd8d..0509835 100644 --- a/.gitignore +++ b/.gitignore @@ -1 +1,3 @@ __pycache__/ +.test-runtime/ +tmp*/ diff --git a/README.md b/README.md index 6fbaf48..93202ab 100644 --- a/README.md +++ b/README.md @@ -54,6 +54,36 @@ Condition example: Conditions are evaluated with `simpleeval`. `$variable` tokens are resolved from the input `vars`; missing variables make the condition false. +### Complex Prompt Parse JSON + +Creates or extends a `ComplexPromptVars` object from a JSON string. + +- `json_text` must be a JSON object. +- Object keys become variable names. +- String and number values are copied into `vars`. +- Boolean, null, object, empty array, and unsupported values are ignored. +- If a value is a non-empty array of strings or numbers, one item is selected + randomly using `seed`. +- `vars` is optional. If it is not connected, the node creates a new object. + +Example: + +```json +{ + "person": "Ada", + "color": ["red", "green", "blue"], + "age": 37 +} +``` + +This can produce variables such as: + +```text +person = Ada +color = green +age = 37 +``` + ## Configuration The configuration file is stored at: diff --git a/nodes.py b/nodes.py index 8c6944c..1aeae8f 100644 --- a/nodes.py +++ b/nodes.py @@ -1,5 +1,6 @@ import json import os +import random import re from pathlib import Path from typing import Any @@ -161,6 +162,41 @@ def apply_vars(prompt: str, vars: dict[str, str] | None = None) -> str: return VAR_PATTERN.sub(replace, prompt) +def is_json_variable_value(value: Any) -> bool: + return isinstance(value, (int, float, str)) and not isinstance(value, bool) + + +def parse_json_vars( + json_text: str, + seed: int, + vars: dict[str, Any] | None = None, +) -> dict[str, Any]: + try: + parsed = json.loads(json_text) + except json.JSONDecodeError as error: + raise ValueError(f"Invalid JSON: {error.msg}") from error + + if not isinstance(parsed, dict): + raise ValueError("JSON must be an object with variable names as keys.") + + result = dict(vars or {}) + rng = random.Random(seed) + + for key, value in parsed.items(): + if is_json_variable_value(value): + result[key] = value + continue + + if ( + isinstance(value, list) + and value + and all(is_json_variable_value(item) for item in value) + ): + result[key] = rng.choice(value) + + return result + + def normalize_condition(condition: str) -> str: normalized = VAR_PATTERN.sub(r"\1", condition) normalized = normalized.replace("&&", " and ") @@ -304,6 +340,47 @@ class ArtemKo7vComplexPromptSetVariable: return (result, True) +class ArtemKo7vComplexPromptParseJSON: + CATEGORY = "ArtemKo7v" + RETURN_TYPES = (ARTEMKO7V_COMPLEX_PROMPT_VARS,) + RETURN_NAMES = ("vars",) + FUNCTION = "parse_json" + + @classmethod + def INPUT_TYPES(cls): + return { + "required": { + "json_text": ( + "STRING", + { + "multiline": True, + "dynamicPrompts": False, + }, + ), + "seed": ( + "INT", + { + "default": 0, + "min": 0, + "max": 18446744073709551615, + "control_after_generate": "randomize", + }, + ), + }, + "optional": { + "vars": (ARTEMKO7V_COMPLEX_PROMPT_VARS,), + }, + } + + def parse_json( + self, + json_text: str, + seed: int, + vars: dict[str, Any] | None = None, + ): + return (parse_json_vars(json_text, seed, vars),) + + class ArtemKo7vComplexPromptEmptyString: CATEGORY = "ArtemKo7v" RETURN_TYPES = ("STRING",) @@ -321,11 +398,13 @@ class ArtemKo7vComplexPromptEmptyString: NODE_CLASS_MAPPINGS = { "ArtemKo7vComplexPrompt": ArtemKo7vComplexPrompt, "ArtemKo7vComplexPromptSetVariable": ArtemKo7vComplexPromptSetVariable, + "ArtemKo7vComplexPromptParseJSON": ArtemKo7vComplexPromptParseJSON, "ArtemKo7vComplexPromptEmptyString": ArtemKo7vComplexPromptEmptyString, } NODE_DISPLAY_NAME_MAPPINGS = { "ArtemKo7vComplexPrompt": "Complex Prompt", "ArtemKo7vComplexPromptSetVariable": "Complex Prompt Set Variable", + "ArtemKo7vComplexPromptParseJSON": "Complex Prompt Parse JSON", "ArtemKo7vComplexPromptEmptyString": "Complex Prompt Empty String", } diff --git a/tests/test_parse_json.py b/tests/test_parse_json.py new file mode 100644 index 0000000..07ae54c --- /dev/null +++ b/tests/test_parse_json.py @@ -0,0 +1,87 @@ +import atexit +import importlib +import random +import shutil +import sys +import unittest +from pathlib import Path +from unittest.mock import patch + + +PROJECT_ROOT = Path(__file__).resolve().parents[1] +TEST_RUNTIME_DIR = PROJECT_ROOT / ".test-runtime" +atexit.register(lambda: shutil.rmtree(TEST_RUNTIME_DIR, ignore_errors=True)) +sys.path.insert(0, str(PROJECT_ROOT)) +with patch("os.getcwd", return_value=str(TEST_RUNTIME_DIR)): + nodes = importlib.import_module("nodes") + + +class ParseJSONVarsTests(unittest.TestCase): + def test_creates_vars_from_string_and_number_values(self): + result = nodes.parse_json_vars( + '{"name": "Ada", "age": 37, "score": 9.5}', + seed=1, + ) + + self.assertEqual(result, {"name": "Ada", "age": 37, "score": 9.5}) + + def test_extends_existing_vars_and_overwrites_matching_keys(self): + result = nodes.parse_json_vars( + '{"name": "Ada", "age": 37}', + seed=1, + vars={"name": "Grace", "city": "London"}, + ) + + self.assertEqual(result, {"name": "Ada", "city": "London", "age": 37}) + + def test_selects_random_item_from_string_or_number_arrays(self): + seed = 42 + expected_color = random.Random(seed).choice(["red", "green", "blue"]) + rng = random.Random(seed) + rng.choice(["red", "green", "blue"]) + expected_size = rng.choice([1, 2, 3]) + + result = nodes.parse_json_vars( + '{"color": ["red", "green", "blue"], "size": [1, 2, 3]}', + seed=seed, + ) + + self.assertEqual(result, {"color": expected_color, "size": expected_size}) + + def test_ignores_unsupported_values(self): + result = nodes.parse_json_vars( + """ + { + "enabled": true, + "empty": [], + "missing": null, + "object": {"value": "nested"}, + "mixed": ["ok", false], + "valid": "kept" + } + """, + seed=1, + ) + + self.assertEqual(result, {"valid": "kept"}) + + def test_rejects_invalid_json(self): + with self.assertRaisesRegex(ValueError, "Invalid JSON"): + nodes.parse_json_vars("{", seed=1) + + def test_rejects_non_object_json(self): + with self.assertRaisesRegex(ValueError, "JSON must be an object"): + nodes.parse_json_vars('["name", "Ada"]', seed=1) + + +class ParseJSONNodeTests(unittest.TestCase): + def test_node_returns_vars_tuple(self): + node = nodes.ArtemKo7vComplexPromptParseJSON() + + result = node.parse_json('{"name": "Ada"}', seed=1) + + self.assertEqual(result, ({"name": "Ada"},)) + + +if __name__ == "__main__": + unittest.main()