|
| 1 | +"""Choosing what survives when a prompt does not fit its budget. |
| 2 | +
|
| 3 | +Prompt assembly is a packing problem: a user profile, a watch history and a |
| 4 | +catalogue of candidates rarely fit together in a context window, and something |
| 5 | +has to give. :func:`priority_select` decides what, keeping sections from the |
| 6 | +most important down until the budget is spent:: |
| 7 | +
|
| 8 | + selection = priority_select( |
| 9 | + [ |
| 10 | + ("critical", system_instructions), |
| 11 | + ("high", f"User profile: {profile}"), |
| 12 | + ("medium", f"Recent history: {history}"), |
| 13 | + ("low", f"Device: {device}"), |
| 14 | + ], |
| 15 | + TokenBudget.of(4096), |
| 16 | + ) |
| 17 | + prompt = selection.render() |
| 18 | +
|
| 19 | +The decision is deterministic in the sections and the budget alone, so an |
| 20 | +OnDemandFeatureView that calls it drops the same sections during training as it |
| 21 | +does at serving time. |
| 22 | +""" |
| 23 | + |
| 24 | +from dataclasses import dataclass |
| 25 | +from enum import Enum, IntEnum |
| 26 | +from typing import Iterable, Iterator, Optional, Union |
| 27 | + |
| 28 | +from feast.context.errors import RequiredSectionError |
| 29 | +from feast.context.token_budget import TokenBudget |
| 30 | + |
| 31 | +#: Placed between kept sections when a selection is rendered. |
| 32 | +DEFAULT_SEPARATOR = "\n\n" |
| 33 | + |
| 34 | +#: Characters of content quoted in messages about an unnamed section. |
| 35 | +_LABEL_PREVIEW_CHARS = 40 |
| 36 | + |
| 37 | + |
| 38 | +class Priority(IntEnum): |
| 39 | + """How readily a section may be dropped to make room for another. |
| 40 | +
|
| 41 | + Ordered, so ``Priority.HIGH > Priority.LOW``, and spaced by ten to leave |
| 42 | + room for levels a caller wants in between:: |
| 43 | +
|
| 44 | + Priority.HIGH - 1 # sits between medium and high |
| 45 | + """ |
| 46 | + |
| 47 | + #: Nice to have: dropped first. |
| 48 | + LOW = 10 |
| 49 | + #: Useful context, expendable under pressure. |
| 50 | + MEDIUM = 20 |
| 51 | + #: Kept unless the budget is genuinely tight. |
| 52 | + HIGH = 30 |
| 53 | + #: Never dropped; the selection fails instead. |
| 54 | + CRITICAL = 40 |
| 55 | + |
| 56 | + @classmethod |
| 57 | + def parse(cls, spec: "PrioritySpec") -> "Priority": |
| 58 | + """Resolve a priority or its case-insensitive name.""" |
| 59 | + if isinstance(spec, Priority): |
| 60 | + return spec |
| 61 | + if isinstance(spec, str): |
| 62 | + try: |
| 63 | + return cls[spec.strip().upper()] |
| 64 | + except KeyError: |
| 65 | + pass |
| 66 | + raise ValueError( |
| 67 | + f"Unknown priority {spec!r}. Use one of: " |
| 68 | + f"{', '.join(level.name.lower() for level in sorted(cls, reverse=True))}." |
| 69 | + ) |
| 70 | + |
| 71 | + def __str__(self) -> str: |
| 72 | + return self.name.lower() |
| 73 | + |
| 74 | + |
| 75 | +#: What callers may pass wherever a priority is expected. |
| 76 | +PrioritySpec = Union[str, Priority] |
| 77 | + |
| 78 | + |
| 79 | +@dataclass(frozen=True) |
| 80 | +class Section: |
| 81 | + """A candidate piece of a prompt, and how hard it is to lose. |
| 82 | +
|
| 83 | + Attributes: |
| 84 | + content: The text itself. Blank content is ignored by selection. |
| 85 | + priority: How readily the section is dropped. |
| 86 | + name: Optional identifier, quoted in error messages and useful for |
| 87 | + asserting on what a selection kept. |
| 88 | + """ |
| 89 | + |
| 90 | + content: str |
| 91 | + priority: Priority = Priority.MEDIUM |
| 92 | + name: Optional[str] = None |
| 93 | + |
| 94 | + def __post_init__(self) -> None: |
| 95 | + if not isinstance(self.content, str): |
| 96 | + raise TypeError(f"Section content must be a string, got {self.content!r}.") |
| 97 | + if not isinstance(self.priority, Priority): |
| 98 | + # A name from config lands here; resolve it once, so comparisons |
| 99 | + # and sorting see a Priority. |
| 100 | + object.__setattr__(self, "priority", Priority.parse(self.priority)) |
| 101 | + |
| 102 | + @classmethod |
| 103 | + def of(cls, spec: "SectionSpec") -> "Section": |
| 104 | + """Normalize a section or a ``(priority, content[, name])`` tuple.""" |
| 105 | + if isinstance(spec, Section): |
| 106 | + return spec |
| 107 | + if isinstance(spec, tuple) and len(spec) in (2, 3): |
| 108 | + priority, content, *rest = spec |
| 109 | + return cls( |
| 110 | + content=content, |
| 111 | + priority=Priority.parse(priority), |
| 112 | + name=rest[0] if rest else None, |
| 113 | + ) |
| 114 | + raise TypeError( |
| 115 | + f"Expected a Section or a (priority, content) tuple, got {spec!r}." |
| 116 | + ) |
| 117 | + |
| 118 | + @property |
| 119 | + def is_required(self) -> bool: |
| 120 | + """Whether the section must be kept, budget notwithstanding.""" |
| 121 | + return self.priority is Priority.CRITICAL |
| 122 | + |
| 123 | + @property |
| 124 | + def label(self) -> str: |
| 125 | + """Quoted identifier for messages: the name, else a content preview.""" |
| 126 | + if self.name: |
| 127 | + return repr(self.name) |
| 128 | + flattened = " ".join(self.content.split()) |
| 129 | + if len(flattened) > _LABEL_PREVIEW_CHARS: |
| 130 | + return repr(f"{flattened[: _LABEL_PREVIEW_CHARS - 3]}...") |
| 131 | + return repr(flattened) |
| 132 | + |
| 133 | + |
| 134 | +#: A section, or the tuple shorthand for one. |
| 135 | +SectionSpec = Union[ |
| 136 | + Section, |
| 137 | + tuple[PrioritySpec, str], |
| 138 | + tuple[PrioritySpec, str, Optional[str]], |
| 139 | +] |
| 140 | + |
| 141 | + |
| 142 | +class SectionOrder(str, Enum): |
| 143 | + """The order kept sections come back in.""" |
| 144 | + |
| 145 | + #: As the caller listed them, so the prompt reads as it was written. |
| 146 | + DECLARED = "declared" |
| 147 | + #: Most important first, ties broken by declaration order. |
| 148 | + PRIORITY = "priority" |
| 149 | + |
| 150 | + |
| 151 | +@dataclass(frozen=True) |
| 152 | +class Selection: |
| 153 | + """The outcome of a :func:`priority_select`: what fit and what did not. |
| 154 | +
|
| 155 | + Iterating yields the kept content, so a selection drops straight into a |
| 156 | + join or a template:: |
| 157 | +
|
| 158 | + for text in selection: ... |
| 159 | + prompt = selection.render() |
| 160 | +
|
| 161 | + Attributes: |
| 162 | + budget: The budget after charging every kept section. |
| 163 | + kept: Sections that made it in, ordered as requested. |
| 164 | + dropped: Sections left out, in declaration order. |
| 165 | + separator: What :meth:`render` joins kept sections with. |
| 166 | + """ |
| 167 | + |
| 168 | + budget: TokenBudget |
| 169 | + kept: tuple[Section, ...] = () |
| 170 | + dropped: tuple[Section, ...] = () |
| 171 | + separator: str = DEFAULT_SEPARATOR |
| 172 | + |
| 173 | + @property |
| 174 | + def contents(self) -> tuple[str, ...]: |
| 175 | + """The text of each kept section.""" |
| 176 | + return tuple(section.content for section in self.kept) |
| 177 | + |
| 178 | + @property |
| 179 | + def names(self) -> tuple[Optional[str], ...]: |
| 180 | + """The name of each kept section, None where unnamed.""" |
| 181 | + return tuple(section.name for section in self.kept) |
| 182 | + |
| 183 | + @property |
| 184 | + def is_complete(self) -> bool: |
| 185 | + """Whether everything offered was kept.""" |
| 186 | + return not self.dropped |
| 187 | + |
| 188 | + def render(self) -> str: |
| 189 | + """The kept sections joined by :attr:`separator`.""" |
| 190 | + return self.separator.join(self.contents) |
| 191 | + |
| 192 | + def __iter__(self) -> Iterator[str]: |
| 193 | + return iter(self.contents) |
| 194 | + |
| 195 | + def __len__(self) -> int: |
| 196 | + return len(self.kept) |
| 197 | + |
| 198 | + def __str__(self) -> str: |
| 199 | + return self.render() |
| 200 | + |
| 201 | + |
| 202 | +def priority_select( |
| 203 | + sections: Iterable[SectionSpec], |
| 204 | + budget: TokenBudget, |
| 205 | + *, |
| 206 | + separator: str = DEFAULT_SEPARATOR, |
| 207 | + order: SectionOrder = SectionOrder.DECLARED, |
| 208 | +) -> Selection: |
| 209 | + """Keep as much of ``sections`` as ``budget`` allows, least important out first. |
| 210 | +
|
| 211 | + Sections are considered from the highest priority down, ties in declaration |
| 212 | + order, and each is kept if what remains of the budget covers it. A section |
| 213 | + too large to fit is dropped and the scan continues, so a small low-priority |
| 214 | + section can still be kept after a large high-priority one was refused. |
| 215 | +
|
| 216 | + Sections whose content is blank are ignored: they would only double a |
| 217 | + separator. ``critical`` sections are never dropped. |
| 218 | +
|
| 219 | + The budget is charged for ``separator`` between each pair of kept sections, |
| 220 | + so :meth:`Selection.render` fits within the allowance. Tokenizers may merge |
| 221 | + text across a section boundary, which makes the charge an upper bound on |
| 222 | + what the joined string really costs. |
| 223 | +
|
| 224 | + Args: |
| 225 | + sections: Sections, or ``(priority, content[, name])`` tuples. |
| 226 | + budget: The allowance to fill. Anything already consumed on it stays |
| 227 | + consumed; only what remains is available here. |
| 228 | + separator: Placed between kept sections by |
| 229 | + :meth:`Selection.render`, and charged against the budget. |
| 230 | + order: Whether kept sections come back in declaration or priority |
| 231 | + order. |
| 232 | +
|
| 233 | + Returns: |
| 234 | + A :class:`Selection` holding the kept sections, the dropped ones, and |
| 235 | + the budget left over. |
| 236 | +
|
| 237 | + Raises: |
| 238 | + RequiredSectionError: A ``critical`` section does not fit. |
| 239 | + TypeError: A section is neither a :class:`Section` nor a tuple. |
| 240 | + ValueError: A priority or an order is not a known name. |
| 241 | + """ |
| 242 | + if not isinstance(separator, str): |
| 243 | + raise TypeError(f"separator must be a string, got {separator!r}.") |
| 244 | + order = SectionOrder(order) |
| 245 | + |
| 246 | + candidates = [ |
| 247 | + section |
| 248 | + for section in (Section.of(spec) for spec in sections) |
| 249 | + if section.content.strip() |
| 250 | + ] |
| 251 | + separator_tokens = budget.count_tokens(separator) |
| 252 | + |
| 253 | + remaining = budget |
| 254 | + kept: list[tuple[int, Section]] = [] |
| 255 | + dropped: list[tuple[int, Section]] = [] |
| 256 | + # Most important first, and among equals whoever was listed first. |
| 257 | + by_importance = sorted( |
| 258 | + enumerate(candidates), key=lambda pair: (-pair[1].priority, pair[0]) |
| 259 | + ) |
| 260 | + |
| 261 | + for position, section in by_importance: |
| 262 | + cost = remaining.count_tokens(section.content) |
| 263 | + if kept: |
| 264 | + cost += separator_tokens |
| 265 | + if remaining.fits_tokens(cost): |
| 266 | + remaining = remaining.consume_tokens(cost) |
| 267 | + kept.append((position, section)) |
| 268 | + elif section.is_required: |
| 269 | + raise RequiredSectionError( |
| 270 | + section.label, cost, remaining.remaining, remaining.max_tokens |
| 271 | + ) |
| 272 | + else: |
| 273 | + dropped.append((position, section)) |
| 274 | + |
| 275 | + if order is SectionOrder.DECLARED: |
| 276 | + kept.sort(key=lambda pair: pair[0]) |
| 277 | + dropped.sort(key=lambda pair: pair[0]) |
| 278 | + |
| 279 | + return Selection( |
| 280 | + budget=remaining, |
| 281 | + kept=tuple(section for _, section in kept), |
| 282 | + dropped=tuple(section for _, section in dropped), |
| 283 | + separator=separator, |
| 284 | + ) |
0 commit comments