A system for generating prompts in an automatic and customizable way
import sys

from launch import run_pip

packages = {
    ("jsonschema", True),
    ("toml", sys.version_info.major == 3 and sys.version_info.minor <= 10)  # Not required from 3.11 onward
}


for package, required in packages:
    if required:
        try:
            run_pip(f"install -U {package}", f"sd-webui-prompt-forge requirement: {package}")
        except Exception as e:
            print(e)
            print(f"Failed to install {package}, script prompt-forge might not work")

install.py
LilianBoulard/prompt-forge
import sys
from argparse import ArgumentParser
from pathlib import Path
from typing import TextIO

from scripts.prompt_forge import Generator


if __name__ == "__main__":
    parser = ArgumentParser(
        "Prompt-forge: create Stable Diffusion prompts from a configuration file"
    )

    parser.add_argument("content_file", type=Path, help="File containing the prompt generation configuration.")
    parser.add_argument("-m", "--mode", choices=("random", "exhaustive"), default="random", help="Mode of generation. Warning: exhaustive generates ALL possibilities, thus the number of prompts scales exponentially.")  # TODO
    parser.add_argument("-n", "--number", type=int, required=False, help="Number of prompts to generate. Ignored if `mode=exhaustive`.")
    parser.add_argument("-o", "--output", type=Path, default=sys.stdout, help="File to save the prompts to. By default, outputs to stdout.")

    args = parser.parse_args()

    if args.mode == "random" and not args.number:
        raise ValueError("Expected a number of prompts to be generated with `mode=random`")

    generator = Generator.from_file(args.content_file)

    def write(prompts: list[str], destination: Path | TextIO):
        prompts_txt = "\n".join(prompts)
        if isinstance(destination, Path):
            destination.write_text(prompts_txt)
        else:
            destination.write(prompts_txt)

    if args.mode == "random":
        write(generator.generate_random_prompts(n=args.number), args.output)
    elif args.mode == "exhaustive":
        write(generator.generate_exhaustive_prompts(), args.output)

prompt-forge-standalone.py
LilianBoulard/prompt-forge
"""
prompt-forge
Core logic code

Author: Lilian Boulard <https://github.com/LilianBoulard>

Licensed under the GNU Affero General Public License.
"""

from __future__ import annotations

import json
import random
import re
import sys
from itertools import product
from pathlib import Path
from typing import Literal

import jsonschema

if sys.version_info.major == 3 and sys.version_info.minor <= 10:
    import toml as tomllib
else:
    import tomllib


def is_balanced(parens: str) -> bool:
    # From: https://stackoverflow.com/a/73341167/
    parens_map ={"(":")","{":"}","[":"]"}
    stack = []
    for paren in parens:
        if paren in parens_map:  # is open
            stack.append(paren)
        elif paren in parens_map.values():  # is closed
            if (not stack) or (paren != parens_map[stack.pop()]):
                return False
    return not stack


def blocks_in_element(element: Block | Group | ExclusionGroup) -> set[Block]:
    """
    Given an element, lists the blocks it is constituted of.

    Parameters
    ----------
    element: Block, Group or ExclusionGroup
        Element to flatten.

    Returns
    -------
    set of Block
        The element, flattened to only the blocks it contains.
    """
    blocks = []
    members_queue = [element]
    while members_queue:
        member = members_queue.pop()
        if isinstance(member, Block):
            blocks.append(member)
        elif isinstance(member, (Group, ExclusionGroup)):
            members_queue.extend(member.members)
    return set(blocks)


def groups_in_group(group: Group | ExclusionGroup) -> list[Group | ExclusionGroup]:
    groups = []
    members_queue = group.members.copy()
    while members_queue:
        member = members_queue.pop()
        if isinstance(member, Block):
            continue
        elif isinstance(member, (Group, ExclusionGroup)):
            groups.append(member)
            members_queue.extend(member.members)
    return set(groups)


class Candidate:

    """
    A candidate node (used in a tree) either has a value
    (it's a leaf) or a list of children (it's a dummy).
    """

    def __init__(self, *, keyword: str | None = None, children: list | None = None, operator: Literal["OR", "AND"] | None = None):
        self.value = keyword
        self.children = children
        self.operator = operator

    def __len__(self) -> int:
        if self.value is not None:
            return 1
        else:
            return sum(len(child) for child in self.children)

    @classmethod
    def parse(cls, candidate: str) -> Candidate:
        """
        Takes the string representation of a candidate,
        and converts it to a tree-style structure, which
        can be further processed.

        Parameters
        ----------
        candidate: str
            The candidate keyword.

        Returns
        -------
        Candidate
            The tree representation of the candidate.
        """
        if not is_balanced(candidate):
            raise ValueError(f"Candidate {candidate!r} has unbalanced brackets/parentheses")

        candidate = candidate.replace("(", "[").replace(")", " | ]")

        # We will explicitly add spaces around
        # the pipes and the brackets, otherwise the splitting method
        # used down below might not work (it wouldn't find the empty string).
        # E.g. `a| [simple| example]|` -> `a | [ simple | example] | `.
        candidate = re.sub(r"(\s{0,1})\|(\s{0,1})", lambda x: f" {x.group(0).strip()} ", candidate)
        candidate = re.sub(r"(\s{0,1})\[(\s{0,1})", lambda x: f" {x.group(0).strip()} ", candidate)
        candidate = re.sub(r"(\s{0,1})\](\s{0,1})", lambda x: f" {x.group(0).strip()} ", candidate)

        return cls._parse(candidate)

    @classmethod
    def _parse(cls, candidate: str) -> Candidate:
        # First pass: split based on the pipes
        # (only on the first nesting layer,
        # we'll delegate the nested ones to recursive calls)
        split_parts = []
        buffer = []
        nesting_level = 0
        for c in candidate:
            if c == "[":
                nesting_level += 1
            elif c == "]":
                nesting_level -= 1

            buffer.append(c)

            if c == "|" and nesting_level == 0:
                split_parts.append("".join(buffer[:-1]).strip())
                buffer.clear()

        if buffer:
            split_parts.append("".join(buffer).strip())
            buffer.clear()

        # Second pass: isolate nested values
        isolated_parts: list[list[str]] = []
        buffer = []
        nesting_level = 0
        for part in split_parts:
            nested_parts_buffer: list[str] = []
            if part == "":
                isolated_parts.append([part])
                continue
            for c in part:
                if c == "[":
                    nesting_level += 1
                    if nesting_level == 1:
                        if buffer:
                            nested_parts_buffer.append("".join(buffer).strip())
                            buffer.clear()
                        continue
                elif c == "]":
                    nesting_level -= 1
                    if nesting_level == 0:
                        if buffer:
                            nested_parts_buffer.append("".join(buffer).strip())
                            buffer.clear()
                        continue
                buffer.append(c)
            if buffer:
                nested_parts_buffer.append("".join(buffer).strip())
                buffer.clear()
            isolated_parts.append(nested_parts_buffer)

        # Third pass: call recursively and collect the results
        children: list[Candidate] = []
        for part in isolated_parts:
            inner_buffer: list[Candidate] = []
            for nested_part in part:
                if "|" in nested_part or "[" in nested_part:
                    inner_buffer.append(cls.parse(nested_part))
                else:
                    inner_buffer.append(cls(keyword=nested_part))
            if len(inner_buffer) == 1:
                children.append(inner_buffer[0])
            else:
                children.append(cls(children=inner_buffer, operator="AND"))

        return cls(children=children, operator="OR")

    def expand(self, *, weighting: Literal["candidate-shallow", "candidate-deep", "keyword"]) -> list[tuple[str, int]]:
        """
        Expands the tree to enumerate all keywords.

        Parameters
        ----------
        weighting: {"candidate-shallow", "candidate-deep", "keyword"}
            Weighting system to use for this candidate (inherited from the block).
            Refer to the README for more information.

        Returns
        -------
        list of 2-tuples of str and int
            All keywords that can be picked from this candidate,
            along with the individual's weight.
        """
        if (
            (self.value is not None and self.children and self.operator)
            or (self.value is None and not self.children and not self.operator)
        ):
            raise ValueError(
                f"Candidate must either be a leaf "
                f"(have a value) or a dummy (have children and an operator). " 
                f"Has an incorrect combination: {self.value=}, "
                f"{self.children=}, {self.operator=}"
            )

        # If it's a leaf, return its value with a single unit of weight
        if self.value is not None:
            return [(self.value, 1)]

        children_and_op: list[tuple[list[tuple[str, int]], Literal["OR", "AND"]]] = [
            # If it's a leaf, it won't have an operator,
            # so we default to AND
            (child.expand(weighting=weighting), child.operator or "AND")
            for child in self.children
        ]

        if self.operator == "AND":
            # Explode AND lists with multiple elements.
            # Keep OR lists nested.
            # TODO: take care of the weights
            product_candidates = []
            for keywords, operator in children_and_op:
                if operator == "AND":
                    for item, weight in keywords:
                        product_candidates.append([(item, weight)])
                else:
                    product_candidates.append(keywords)
            prod = list(product(*product_candidates))
            joined_keywords = []
            for elements in prod:
                # FIXME: weighting is incorrect
                keywords, weights = zip(*elements)
                joined_keywords.append((
                    " ".join(filter(lambda keyword: not not keyword, keywords)),
                    min(weights),
                ))
            return joined_keywords

        elif self.operator == "OR":
            # Return everything together
            if weighting == "candidate-shallow":
                # Adjust the cumulative weights
                raise NotImplementedError("`candidate-shallow` is not implemented yet")
            elif weighting == "candidate-deep" or weighting == "keyword":
                # Do not adjust the cumulative weights
                return [
                    # weight is always 1
                    (item, weight)
                    for items, _ in children_and_op
                    for item, weight in items
                ]


class Group:

    __namespace__ = "groups"

    def __init__(self, name: str, members: list[Block]):
        self.name = name
        self.members = members

    def __hash__(self):
        return hash(self.name)
    
    def __repr__(self):
        return f"Group {self.name!r} with {len(self.members)} members"


class ExclusionGroup:

    __namespace__ = "exclusions"

    def __init__(self, name: str, members: list[Block | Group], weights: list[int] | None):
        self.name = name
        self.members = members
        self.weights = weights

    def choose_member(self) -> Block | Group:
        return random.choices(self.members, weights=self.weights)[0]

    def __hash__(self):
        return hash(self.name)
    
    def __repr__(self):
        return f"ExclusionGroup {self.name!r} with {len(self.members)} members"


class Block:

    """
    A block of keywords.

    Parameters
    ----------
    name: str
        Name of the block.
    candidates: list of str
        Candidates (can contain square brackets and parentheses).
    parameters: mapping of str to str or bool
        Parameters passed in the config.

    Attributes
    ----------
    name: str
        Name of the block
    num: range
        Possible number of keywords to pick.
    separator: str
        When picking multiple keywords from this block,
        the separator to use when joining them.
    candidates: list of list of 2-tuples of str and int
        The keywords for each candidate.
    weighting: {"candidate-shallow", "candidate-deep", "keyword"}
        How to tune probabilities when picking a keyword.
        Refer to the guide, section "Weighting" for more information.
    """

    __namespace__ = "blocks"

    name: str
    num: range
    separator: str
    candidates: list[list[tuple[str, int]]]
    weighting: Literal["candidate-shallow", "candidate-deep", "keyword"]

    def __init__(self, name: str, candidates: list[str], parameters: dict[str, any]) -> Block:
        self.name = name

        self.weighting = parameters.get("weighting", "candidate-deep")
        assert self.weighting in {"candidate-shallow", "candidate-deep", "keyword"}

        self.candidates = [
            Candidate.parse(candidate).expand(weighting=self.weighting)
            for candidate in candidates
        ]

        self.num = range(1, 2)
        if parameters.get("force"):
            self.num = range(len(self.candidates), len(self.candidates) + 1)
        elif parameters.get("optional"):
            self.num = range(0, 2)
        elif num_param := parameters.get("num"):
            # Parses `2` as `(2, 2)` and `2-3` as `(2, 3)`
            min_num, max_num = map(int, num_param.split("-")) if "-" in num_param else (int(num_param), int(num_param))
            self.num = range(min_num, max_num + 1)

        self.separator = parameters.get("separator", ", ")

    def __hash__(self):
        return hash(self.name)

    def __repr__(self):
        return (
            f"Block {self.name!r} "
            f"<{len(self.candidates)} candidates, "
            f"{sum(len(keywords) for keywords in self.candidates)} keywords>"
        )

    def generate_keyword(self) -> str:
        k = min(random.choice(self.num), len(self.candidates))
        chosen_candidates: list[list[tuple[str, int]]]

        if self.weighting in {"candidate-shallow", "candidate-deep"}:
            chosen_candidates = random.choices(self.candidates, k=k)
            return self.separator.join([
                random.choices(*zip(*candidate))[0]
                for candidate in chosen_candidates
            ])

        elif self.weighting == "keyword":
            candidate_weights = [len(candidate) for candidate in self.candidates]
            chosen_candidates = random.choices(self.candidates, weights=candidate_weights, k=k)

    def generate_all_keywords(self) -> list[str]:
        """
        Used when mode="exhaustive", returns the list of
        all possible keywords this block can generate.
        """
        return [
            keyword
            for candidate_keywords in self.candidates
            for keyword, _ in candidate_keywords
        ]


class Generator:

    """
    A generator created from a config.
    Capable of creating new prompts on demand.
    """

    def __init__(self, elements: list[Block | Group | ExclusionGroup], blocks_order: dict[str, int]) -> Generator:
        self.elements = elements
        self.blocks_order = blocks_order

    @classmethod
    def from_file(cls, file: Path) -> Generator:
        with file.open() as f:
            return cls.from_string(f.read())

    @classmethod
    def from_string(cls, configuration: str) -> Generator:
        # Parse the config
        config = tomllib.loads(configuration)

        # Since toml returns the config as an unordered JSON document,
        # we read the configuration to find the order of the blocks
        pattern = re.compile(r"\[" + Block.__namespace__ + r"\.(.+?)\]")
        blocks_order: dict[str, int] = {
            match.group(0): match.start()
            for match in pattern.finditer(configuration)
        }

        # Validate the JSON schema
        schema_file = Path(__file__).parent / "config-schema.json"
        if schema_file.is_file():
            with schema_file.open("rb") as f:
                schema = json.load(f)
            jsonschema.validate(config, schema)
        else:
            print(
                f"Did not find schema at {schema_file!s} "
                f"to validate the configuration against. "
            )

        # Create the blocks
        blocks = {
            Block(name, block.get("candidates", list()), block)
            for name, block in config.get(Block.__namespace__, {}).items()
        }
        mappings = {f"{Block.__namespace__}.{block.name}": block for block in blocks}

        # Create the groups
        groups = [
            Group(
                name=name,
                members=[
                    mappings[group_name]
                    for group_name in group.get("members", list())
                ],
            )
            for name, group in config.get(Group.__namespace__, {}).items()
        ]
        mappings.update({f"{Group.__namespace__}.{group.name}": group for group in groups})

        # Create the exclusion groups
        exclusion_groups = [
            ExclusionGroup(
                name=name, 
                members=[
                    mappings[member_name]
                    for member_name in group.get("members", list())
                ],
                weights=group.get("weights"),
            )
            for name, group in config.get(ExclusionGroup.__namespace__, {}).items()
        ]

        # List blocks that are present in at least one group
        used_blocks = {
            block
            for group in {*groups, *exclusion_groups}
            for block in blocks_in_element(group)
        }
        # List groups that are present in exclusion groups
        groups_in_exclusion_groups = {
            group
            for exclusion_group in exclusion_groups
            for group in groups_in_group(exclusion_group)
        }

        # List the blocks that are not present in any groups
        elements = blocks.difference(used_blocks)
        # List groups present in exclusion groups
        leftover_groups = {
            group
            for group in groups
            if group not in groups_in_exclusion_groups
        }

        # Add the remaining groups
        elements.update(leftover_groups)
        # And the exclusion groups
        elements.update(exclusion_groups)

        return cls(elements, blocks_order)

    def _sort_elements(self, element: Block | Group | ExclusionGroup) -> int:
        return min(
            self.blocks_order[f"[{Block.__namespace__}.{block.name}]"]
            for block in blocks_in_element(element)
        )

    def generate_random_prompts(self, n: int) -> list[str]:
        prompts = []
        for _ in range(n):
            prompt = []
            activated_blocks = []
            stack = [*self.elements]
            while stack:
                element = stack.pop()
                if isinstance(element, Block):
                    activated_blocks.append(element)
                elif isinstance(element, Group):
                    stack.extend(element.members)
                elif isinstance(element, ExclusionGroup):
                    stack.append(element.choose_member())
            for block in sorted(activated_blocks, key=self._sort_elements):
                keyword = block.generate_keyword()
                if keyword:  # Ignore empty keywords
                    prompt.append(keyword)
            prompts.append(", ".join(prompt))
        return prompts

    def generate_exhaustive_prompts(self) -> list[str]:
        # First step: create all the lines, that is,
        # lists of blocks that will generate prompts
        lines: list[list[Block]] = [[]]
        for element in self.elements:
            if isinstance(element, Block):
                for line in lines:
                    line.append(element)
            elif isinstance(element, ExclusionGroup):
                new_lines = []
                for member in element.members:
                    alternative_lines = lines.copy()
                    for line in alternative_lines:
                        line.extend(blocks_in_element(member))
                    new_lines.extend(alternative_lines)
                lines = new_lines

        # Second step: create the prompts
        prompts: list[str] = []
        for blocks in lines:
            prompt_parts: list[list[str]] = []
            for block in sorted(blocks, key=self._sort_elements):
                prompt_parts.append(block.generate_all_keywords())
            line_prompts = product(*prompt_parts)
            prompts.extend(map(lambda prompt_parts: ", ".join(prompt_parts), line_prompts))

        return prompts

scripts/prompt_forge.py
LilianBoulard/prompt-forge
"""
prompt-forge
Interface for AUTOMATIC1111's Stable Diffusion webui.

Author: Lilian Boulard <https://github.com/LilianBoulard>

Licensed under the GNU Affero General Public License.
"""

import copy
import random

import gradio as gr
import modules.scripts as scripts
from modules import errors
from modules.processing import (Processed, StableDiffusionProcessing,
                                process_images)
from modules.shared import state
from scripts.prompt_forge import Generator
from scripts.prompts_from_file import cmdargs


class Script(scripts.Script):

    """
    Class for interfacing with AUTOMATIC1111's webui.
    """
    
    def title(self) -> str:
        return "Prompt forge"

    def show(self, is_img2img: bool) -> bool:
        return True

    def ui(self, is_img2img: bool) -> tuple[gr.Textbox, gr.Radio, gr.Number, gr.Checkbox, gr.Checkbox]:

        def load_config_file(file: gr.File):
            """
            From https://github.com/AUTOMATIC1111/stable-diffusion-webui/blob/5ef669de080814067961f28357256e8fe27544f4/scripts/prompts_from_file.py#L96
            FIXME: why 7?
            """
            if file is None:
                return None, gr.update(), gr.update(lines=7)
            else:
                lines = file.decode('utf8', errors='ignore').splitlines()
                return None, "\n".join(lines), gr.update(lines=7)

        GUIDE_LINK = r"https://github.com/LilianBoulard/prompt-forge/blob/main/README.md"
        gr.HTML(value=f"<br>Confused/new? Read <a style=\"border-bottom: 1px #00ffff dotted;\" href=\"{GUIDE_LINK}\" target=\"_blank\" rel=\"noopener noreferrer\">the guide</a> for usage instructions!<br><br>")
        configuration = gr.Textbox(
            label="Configuration",

            elem_id=self.elem_id("prompt_txt"),
        )
        configuration_file = gr.File(
            label="Configuration file",
            info="See project's README for syntax definition",
            file_types=[".toml"],
            type="binary",
            elem_id=self.elem_id("configuration_file"),
        )
        with gr.Row():
            generation_type = gr.Radio(
                label="Mode",
                info=(
                    "Random will create a number of prompts randomly picked from the configuration, "
                    "exhaustive will generate all possible combinations."
                ),
                choices=["Random", "Exhaustive"],
                value="Random",
            )
            number_of_entries = gr.Number(
                label="Number of random prompts to generate",
                info="Note: ignore if using exhaustive mode",
                precision=0,
                minimum=0,
            )
        with gr.Row():
            warn_if_duplicates = gr.Checkbox(
                label="Warn when there are duplicate prompts (Note: can only happen when using mode 'Random')",
                value=True,
            )
            dry_run = gr.Checkbox(
                label="Dry run to test the configuration",
                value=False,
            )

        # Push configuration file's content into the textbox
        configuration_file.change(fn=load_config_file, inputs=[configuration_file], outputs=[configuration_file, configuration, configuration], show_progress=False)
        # TODO: Listen to changes on the mode selector, and enable/disable the entries
        # selector depending on the value
        # TODO: Disable duplicates warning when using mode exhaustive

        return (
            configuration,
            generation_type,
            number_of_entries,
            warn_if_duplicates,
            dry_run,
        )

    def run(
            self,
            p: StableDiffusionProcessing,
            configuration: str,
            generation_type: str,
            number_of_entries: int,
            duplicates_warning: bool,
            dry_run: bool,
        ):
        """
        Parts of this function are from script
        [prompts_from_file](https://github.com/AUTOMATIC1111/stable-diffusion-webui/blob/master/scripts/prompts_from_file.py)
        (https://github.com/AUTOMATIC1111/stable-diffusion-webui/blob/5ef669de080814067961f28357256e8fe27544f4/scripts/prompts_from_file.py).
        """
        if not configuration:
            raise RuntimeError("Configuration is required, please load one in. Read the script's guide for guidance.")

        generator: Generator = Generator.from_string(configuration)

        if generation_type == "Random":
            lines = generator.generate_random_prompts(number_of_entries)
        elif generation_type == "Exhaustive":
            lines = generator.generate_exhaustive_prompts()
        else:
            raise ValueError(f"Invalid generation type {generation_type!r}")

        if duplicates_warning:
            n_uniques = len(set(lines))
            duplicates = len(lines) - n_uniques
            if duplicates > 0:
                errors.report(f"The generated prompts contain duplicates ({duplicates}/{len(lines)})")

        if dry_run:
            print(f"Prompt-forge dry run: generated {len(lines)} prompts:")
            if len(lines) <= 8:
                for line in lines:
                    print(line)
            else:
                for line in lines[:3]:
                    print(line)
                print(f"... ({len(lines) - 6} prompts)")
                for line in lines[-3:]:
                    print(line)
            raise RuntimeWarning("Dry run completed without errors. Check the logs for details.")

        p.do_not_save_grid = True

        job_count = 0
        jobs = []

        for line in lines:
            if "--" in line:
                try:
                    args = cmdargs(line)
                except Exception:
                    errors.report(f"Error parsing prompt {line} as commandline", exc_info=True)
                    args = {"prompt": line}
            else:
                args = {"prompt": line}

            job_count += args.get("n_iter", p.n_iter)

            jobs.append(args)

        print(f"Will process {len(lines)} prompt in {job_count} jobs.")
        if p.seed == -1:
            p.seed = int(random.randrange(4294967294))

        state.job_count = job_count

        images = []
        all_prompts = []
        infotexts = []
        for args in jobs:
            state.job = f"{state.job_no + 1} out of {state.job_count}"

            copy_p = copy.copy(p)
            for k, v in args.items():
                setattr(copy_p, k, v)

            proc = process_images(copy_p)
            images += proc.images

            all_prompts += proc.all_prompts
            infotexts += proc.infotexts

        return Processed(p, images, p.seed, "", all_prompts=all_prompts, infotexts=infotexts)
scripts/webui_script.py
LilianBoulard/prompt-forge
from pathlib import Path

import jsonschema
import pytest
from scripts.prompt_forge import Generator, Candidate


def test_config_parsing():
    # Test that provided example config is valid
    Generator.from_file(Path(__file__).parent / "example-config.toml").generate_exhaustive_prompts()

    with pytest.raises(jsonschema.exceptions.ValidationError, match="required property"):
        # Empty config
        Generator.from_string("")
        # Not a config
        Generator.from_string("Invalid configuration")
        Generator.from_string('{"props": [1, 2, 3]}')
        # Malformed config
        Generator.from_string('[bloks.invalid]\ncandidates=[]')  # No `blocks`, but something else
        Generator.from_string('[blocks.test]')  # No candidates in block


def test_parse_candidate():
    """
    Tests that parsing a candidate returns the expected keywords.
    """
    test_candidates: list[tuple[str, tuple]] = [
        ("[[A | B] | [C | D]] E | F", ("A E", "B E", "C E", "D E", "F")),
        ("[A | B] [[C | D | ] E F | (G) H]", ('A C E F', 'A D E F', 'A E F', 'A G H', 'A H', 'B C E F', 'B D E F', 'B E F', 'B G H', 'B H')),
    ]
    for candidate, expected_keywords in test_candidates:
        # Dump the weights, we'll test that in another unit
        keywords, _ = zip(*Candidate.parse(candidate).expand(weighting="keyword"))
        assert keywords == expected_keywords


def test_exhaustive():
    prompts = Generator.from_string(
        "[blocks.test]\n"
        "candidates = [\n"
        "  \"[[A | B] | [C | D]] E | F\",\n"
        "  \"G | H\",\n"
        "]\n"
    ).generate_exhaustive_prompts()
    assert prompts == ['A E', 'B E', 'C E', 'D E', 'F', 'G', 'H']

test_forge.py
LilianBoulard/prompt-forge