45 lines
1.2 KiB
Python
45 lines
1.2 KiB
Python
"""Load and render TOML-backed LLM prompt templates."""
|
|
|
|
from __future__ import annotations
|
|
|
|
import tomllib
|
|
from dataclasses import dataclass
|
|
from functools import cache
|
|
from pathlib import Path
|
|
|
|
|
|
@dataclass(frozen=True)
|
|
class Prompt:
|
|
"""A system and user prompt pair loaded from TOML."""
|
|
|
|
system: str
|
|
user: str
|
|
|
|
def messages(self, **values: str) -> list[dict[str, str]]:
|
|
"""Render this prompt as OpenAI-style chat messages."""
|
|
return [
|
|
{"role": "system", "content": self.system.format(**values)},
|
|
{"role": "user", "content": self.user.format(**values)},
|
|
]
|
|
|
|
|
|
@cache
|
|
def _get_prompt_dir() -> Path:
|
|
"""Return the directory containing prompt template files."""
|
|
return Path(__file__).resolve().parent
|
|
|
|
|
|
@cache
|
|
def load_prompt(name: str) -> Prompt:
|
|
"""Load and validate a named system and user prompt pair from TOML."""
|
|
path = _get_prompt_dir() / f"{name}.toml"
|
|
with path.open("rb") as file:
|
|
body = tomllib.load(file)
|
|
|
|
system = body.get("system")
|
|
user = body.get("user")
|
|
if not isinstance(system, str) or not isinstance(user, str):
|
|
msg = f"{path} must define string system and user prompts"
|
|
raise TypeError(msg)
|
|
return Prompt(system=system, user=user)
|