1"""VersionedPromptStore — version history and rollback for named templates."""
2
3from __future__ import annotations
4
5from dataclasses import dataclass, field
6from typing import TYPE_CHECKING, Any
7
8from lexigram.ai.prompt.exceptions import PromptNotFoundError, PromptVersionError
9
10if TYPE_CHECKING:
11 from lexigram.ai.prompt.template.base import AbstractPromptTemplate
12
13
14@dataclass
15class _VersionEntry:
16 version: int
17 template: AbstractPromptTemplate
18 metadata: dict[str, Any] = field(default_factory=dict)
19
20
21class VersionedPromptStore:
22 """A store that maintains a full version history for each named template.
23
24 Templates are added with :meth:`push`, which appends a new version and
25 makes it the *current* (latest) version. Previous versions remain
26 accessible via :meth:`get_version` and can be restored with
27 :meth:`rollback`.
28
29 Args:
30 max_versions: Maximum history depth per name. When the limit is
31 reached the oldest version is evicted. ``0`` means
32 unlimited. Defaults to ``0``.
33
34 Example::
35
36 store = VersionedPromptStore()
37 store.push("greeting", v1_template) # version 1
38 store.push("greeting", v2_template) # version 2
39 store.rollback("greeting") # restores version 1
40 """
41
42 def __init__(self, max_versions: int = 0) -> None:
43 self._history: dict[str, list[_VersionEntry]] = {}
44 self._current: dict[str, int] = {} # name → current version number
45 self._max = max_versions
46
47 def push(
48 self,
49 name: str,
50 template: AbstractPromptTemplate,
51 metadata: dict[str, Any] | None = None,
52 ) -> int:
53 """Add a new version of *template* under *name*.
54
55 Args:
56 name: Unique template name.
57 template: Template to store.
58 metadata: Arbitrary metadata attached to this version.
59
60 Returns:
61 The new version number (1-based, incrementing).
62 """
63 history = self._history.setdefault(name, [])
64 next_version = (history[-1].version + 1) if history else 1
65 entry = _VersionEntry(next_version, template, metadata or {})
66 history.append(entry)
67 self._current[name] = next_version
68
69 if self._max > 0 and len(history) > self._max:
70 history.pop(0)
71
72 return next_version
73
74 def get(self, name: str) -> AbstractPromptTemplate:
75 """Return the current (latest active) version of *name*.
76
77 Args:
78 name: Template name.
79
80 Raises:
81 :class:`~lexigram.ai.prompt.exceptions.PromptNotFoundError`:
82 *name* has never been pushed.
83 """
84 history = self._history.get(name)
85 if not history:
86 raise PromptNotFoundError(f"No template stored under '{name}'.")
87 current_version = self._current[name]
88 return self._find_entry(name, current_version).template
89
90 def get_version(self, name: str, version: int) -> AbstractPromptTemplate:
91 """Return a specific historical version.
92
93 Args:
94 name: Template name.
95 version: 1-based version number.
96
97 Raises:
98 :class:`~lexigram.ai.prompt.exceptions.PromptNotFoundError`:
99 *name* is unknown.
100 :class:`~lexigram.ai.prompt.exceptions.PromptVersionError`:
101 *version* does not exist (may have been evicted).
102 """
103 return self._find_entry(name, version).template
104
105 def rollback(self, name: str, steps: int = 1) -> AbstractPromptTemplate:
106 """Roll back *steps* version(s) and return the restored template.
107
108 Args:
109 name: Template name.
110 steps: Number of versions to roll back. Defaults to ``1``.
111
112 Returns:
113 The restored (now current) template.
114
115 Raises:
116 :class:`~lexigram.ai.prompt.exceptions.PromptVersionError`:
117 Cannot roll back further than the oldest stored version.
118 """
119 if name not in self._current:
120 raise PromptNotFoundError(f"No template stored under '{name}'.")
121
122 history = self._history[name]
123 current_index = next(
124 i for i, e in enumerate(history) if e.version == self._current[name]
125 )
126 target_index = current_index - steps
127 if target_index < 0:
128 raise PromptVersionError(
129 f"Cannot roll back {steps} step(s) for '{name}': "
130 f"only {current_index} version(s) are available in history."
131 )
132 target = history[target_index]
133 self._current[name] = target.version
134 return target.template
135
136 def list_versions(self, name: str) -> list[dict[str, Any]]:
137 """Return metadata for all stored versions of *name* (oldest first).
138
139 Args:
140 name: Template name.
141
142 Returns:
143 List of ``{"version": int, "current": bool, "metadata": dict}`` dicts.
144
145 Raises:
146 :class:`~lexigram.ai.prompt.exceptions.PromptNotFoundError`:
147 *name* is unknown.
148 """
149 if name not in self._history:
150 raise PromptNotFoundError(f"No template stored under '{name}'.")
151 current_version = self._current[name]
152 return [
153 {
154 "version": e.version,
155 "current": e.version == current_version,
156 "metadata": e.metadata,
157 }
158 for e in self._history[name]
159 ]
160
161 def _find_entry(self, name: str, version: int) -> _VersionEntry:
162 history = self._history.get(name)
163 if not history:
164 raise PromptNotFoundError(f"No template stored under '{name}'.")
165 for entry in history:
166 if entry.version == version:
167 return entry
168 raise PromptVersionError(
169 f"Version {version} of '{name}' is not available "
170 f"(may have been evicted by max_versions={self._max})."
171 )
172
173
174__all__ = ["VersionedPromptStore"]