49 lines
1.7 KiB
Python
49 lines
1.7 KiB
Python
from collections.abc import Iterable
|
|
from dataclasses import dataclass
|
|
from importlib.resources.abc import Traversable
|
|
from itertools import starmap
|
|
from typing import Self
|
|
|
|
from .document import MarkdownDocument
|
|
|
|
|
|
@dataclass(frozen=True, slots=True)
|
|
class PromptFilesBundle:
|
|
"""Represents a prompt and all of its associated markdown files."""
|
|
|
|
slug: str
|
|
prompt: MarkdownDocument
|
|
other: tuple[MarkdownDocument, ...]
|
|
|
|
@classmethod
|
|
def from_root(cls, root: Traversable) -> list[Self]:
|
|
# Should only be used for testing
|
|
return list(cls.from_docs(MarkdownDocument.from_root(root).values()))
|
|
|
|
@classmethod
|
|
def from_docs(cls, docs: Iterable[MarkdownDocument]) -> tuple[Self, ...]:
|
|
return tuple(starmap(cls.from_paths, group_prompt_paths(docs).items()))
|
|
|
|
@classmethod
|
|
def from_paths(cls, slug: str, paths: set[MarkdownDocument]) -> Self:
|
|
prompt = next(iter(p for p in paths if p.relpath.name == "PROMPT.md"))
|
|
sorted_paths = tuple(sorted(paths, key=lambda p: p.relpath.as_posix()))
|
|
other = tuple(p for p in sorted_paths if p != prompt)
|
|
return cls(
|
|
slug=slug,
|
|
prompt=prompt,
|
|
other=other,
|
|
)
|
|
|
|
|
|
def group_prompt_paths(docs: Iterable[MarkdownDocument]) -> dict[str, set[MarkdownDocument]]:
|
|
"""Group prompts from a list of markdown documents by their prompt slug."""
|
|
grouped: dict[str, set[MarkdownDocument]] = {}
|
|
for doc in sorted(
|
|
filter(lambda d: d.prompt_slug is not None, docs),
|
|
key=lambda d: d.relpath.as_posix(),
|
|
):
|
|
if doc.prompt_slug:
|
|
grouped.setdefault(doc.prompt_slug, set()).add(doc)
|
|
return grouped
|