mcp-plesk-dev-docs 0.4.2__py3-none-any.whl

plesk_unified/turboquant/compressors.py

@@ -0,0 +1,190 @@
+"""TurboQuant KV cache helpers."""
+
+from __future__ import annotations
+
+import math
+
+import torch
+
+
+# ---------------------------------------------------------------------------
+# Closed-form Gaussian integration helpers (replaces scipy.integrate.quad)
+# ---------------------------------------------------------------------------
+
+
+def _gauss_pdf(x: float, sigma: float) -> float:
+    return math.exp(-0.5 * (x / sigma) ** 2) / (sigma * math.sqrt(2.0 * math.pi))
+
+
+def _gauss_cdf(x: float, sigma: float) -> float:
+    return 0.5 * (1.0 + math.erf(x / (sigma * math.sqrt(2.0))))
+
+
+def _int_pdf(a: float, b: float, sigma: float) -> float:
+    return _gauss_cdf(b, sigma) - _gauss_cdf(a, sigma)
+
+
+def _int_x_pdf(a: float, b: float, sigma: float) -> float:
+    return sigma * sigma * (_gauss_pdf(a, sigma) - _gauss_pdf(b, sigma))
+
+
+class TurboQuantCompressorV2:
+    """Compressed key store with direct inner-product scoring."""
+
+    def __init__(self, head_dim: int, bits: int, seed: int, device: str = "cpu"):
+        self.head_dim = head_dim
+        self.bits = bits
+        self.mse_bits = max(bits - 1, 1)
+        self.device = device
+
+        gen = torch.Generator(device="cpu")
+        gen.manual_seed(seed)
+        G = torch.randn(head_dim, head_dim, generator=gen)
+        Q, R = torch.linalg.qr(G)
+        diag_sign = torch.sign(torch.diag(R))
+        diag_sign[diag_sign == 0] = 1.0
+        self.Pi = (Q * diag_sign.unsqueeze(0)).to(device)
+
+        self.centroids = self._solve_codebook(head_dim, self.mse_bits).to(device)
+
+        gen2 = torch.Generator(device="cpu")
+        gen2.manual_seed(seed + 10000)
+        self.S = torch.randn(head_dim, head_dim, generator=gen2).to(device)
+
+        self.PiT = self.Pi.T.contiguous()
+
+    def _solve_codebook(self, d: int, bits: int) -> torch.Tensor:
+        n_levels = 2**bits
+        sigma = 1.0 / math.sqrt(d)
+
+        lo, hi = -3.5 * sigma, 3.5 * sigma
+        centroids = [lo + (hi - lo) * (i + 0.5) / n_levels for i in range(n_levels)]
+
+        for _ in range(200):
+            boundaries = [
+                (centroids[i] + centroids[i + 1]) / 2.0 for i in range(n_levels - 1)
+            ]
+            edges = [lo * 3] + boundaries + [hi * 3]
+            new_centroids = []
+            for i in range(n_levels):
+                a, b = edges[i], edges[i + 1]
+                num = _int_x_pdf(a, b, sigma)
+                den = _int_pdf(a, b, sigma)
+                new_centroids.append(num / den if den > 1e-15 else centroids[i])
+            if (
+                max(abs(new_centroids[i] - centroids[i]) for i in range(n_levels))
+                < 1e-10
+            ):
+                break
+            centroids = new_centroids
+
+        return torch.tensor(centroids, dtype=torch.float32)
+
+    @torch.no_grad()
+    def compress(self, states: torch.Tensor) -> dict:
+        B, H, S, D = states.shape
+        flat = states.reshape(-1, D).float()
+
+        vec_norms = torch.norm(flat, dim=-1, keepdim=True)
+        flat_norm = flat / (vec_norms + 1e-8)
+
+        rotated = flat_norm @ self.Pi.T
+        diffs = rotated.unsqueeze(-1) - self.centroids
+        indices = diffs.abs().argmin(dim=-1).to(torch.uint8)
+
+        reconstructed_rotated = self.centroids[indices.long()]
+        k_mse = (reconstructed_rotated @ self.Pi) * vec_norms
+
+        residual = flat - k_mse
+        residual_norm = torch.norm(residual, dim=-1)
+
+        projected = residual @ self.S.T
+        signs = (projected >= 0).to(torch.int8) * 2 - 1
+
+        return {
+            "k_mse": k_mse.to(torch.float16).reshape(B, H, S, D),
+            "qjl_signs": signs.reshape(B, H, S, D),
+            "residual_norm": residual_norm.to(torch.float16).reshape(B, H, S),
+            "shape": (B, H, S, D),
+        }
+
+    @torch.no_grad()
+    def asymmetric_attention_scores(
+        self, queries: torch.Tensor, compressed: dict
+    ) -> torch.Tensor:
+        k_mse = compressed["k_mse"].float()
+        signs = compressed["qjl_signs"].float()
+        r_norm = compressed["residual_norm"].float()
+
+        term1 = torch.matmul(queries.float(), k_mse.transpose(-2, -1))
+        q_projected = torch.matmul(queries.float(), self.S.T)
+        qjl_ip = torch.matmul(q_projected, signs.transpose(-2, -1))
+
+        m = self.S.shape[0]
+        correction_scale = math.sqrt(math.pi / 2) / m
+        term2 = correction_scale * qjl_ip * r_norm.unsqueeze(-2)
+
+        return term1 + term2
+
+
+class TurboQuantCompressorMSE:
+    """MSE-only compressor for values."""
+
+    def __init__(self, head_dim: int, bits: int, seed: int, device: str = "cpu"):
+        self.head_dim = head_dim
+        self.bits = bits
+        self.device = device
+
+        gen = torch.Generator(device="cpu")
+        gen.manual_seed(seed)
+        G = torch.randn(head_dim, head_dim, generator=gen)
+        Q, R = torch.linalg.qr(G)
+        diag_sign = torch.sign(torch.diag(R))
+        diag_sign[diag_sign == 0] = 1.0
+        self.Pi = (Q * diag_sign.unsqueeze(0)).to(device)
+        self.centroids = self._solve_codebook(head_dim, bits).to(device)
+
+    def _solve_codebook(self, d, bits):
+        n_levels = 2**bits
+        sigma = 1.0 / math.sqrt(d)
+
+        lo, hi = -3.5 * sigma, 3.5 * sigma
+        centroids = [lo + (hi - lo) * (i + 0.5) / n_levels for i in range(n_levels)]
+        for _ in range(200):
+            boundaries = [
+                (centroids[i] + centroids[i + 1]) / 2.0 for i in range(n_levels - 1)
+            ]
+            edges = [lo * 3] + boundaries + [hi * 3]
+            new_c = []
+            for i in range(n_levels):
+                a, b = edges[i], edges[i + 1]
+                num = _int_x_pdf(a, b, sigma)
+                den = _int_pdf(a, b, sigma)
+                new_c.append(num / den if den > 1e-15 else centroids[i])
+            if max(abs(new_c[i] - centroids[i]) for i in range(n_levels)) < 1e-10:
+                break
+            centroids = new_c
+        return torch.tensor(centroids, dtype=torch.float32)
+
+    @torch.no_grad()
+    def compress(self, states: torch.Tensor) -> dict:
+        B, H, S, D = states.shape
+        flat = states.reshape(-1, D).float()
+        vec_norms = torch.norm(flat, dim=-1, keepdim=True)
+        flat_norm = flat / (vec_norms + 1e-8)
+        rotated = flat_norm @ self.Pi.T
+        diffs = rotated.unsqueeze(-1) - self.centroids
+        indices = diffs.abs().argmin(dim=-1).to(torch.uint8)
+        return {
+            "indices": indices,
+            "vec_norms": vec_norms.squeeze(-1).to(torch.float16),
+            "shape": (B, H, S, D),
+        }
+
+    @torch.no_grad()
+    def decompress(self, compressed: dict) -> torch.Tensor:
+        B, H, S, D = compressed["shape"]
+        indices = compressed["indices"].long()
+        reconstructed = self.centroids[indices] @ self.Pi
+        vec_norms = compressed["vec_norms"].float().unsqueeze(-1)
+        return (reconstructed * vec_norms).reshape(B, H, S, D)

plesk_unified/turboquant/lloyd_max.py

@@ -0,0 +1,190 @@
+# ruff: noqa
+"""Lloyd-Max scalar quantizer for rotated unit vectors.
+
+The coordinate distribution is approximately Beta-shaped on [-1, 1] after
+random rotation. For d >= 64, a Gaussian N(0, 1/d) is a good approximation.
+"""
+
+from __future__ import annotations
+
+import math
+
+import torch
+
+
+# ---------------------------------------------------------------------------
+# Pure-Python Gaussian integration helpers (replaces scipy.integrate.quad)
+# ---------------------------------------------------------------------------
+
+
+def _gauss_pdf(x: float, sigma: float) -> float:
+    """N(0, σ²) probability density at x."""
+    return math.exp(-0.5 * (x / sigma) ** 2) / (sigma * math.sqrt(2.0 * math.pi))
+
+
+def _gauss_cdf(x: float, sigma: float) -> float:
+    """N(0, σ²) cumulative distribution at x."""
+    return 0.5 * (1.0 + math.erf(x / (sigma * math.sqrt(2.0))))
+
+
+def _int_pdf(a: float, b: float, sigma: float) -> float:
+    """∫[a,b] N(0,σ²)(x) dx — closed form via erf."""
+    return _gauss_cdf(b, sigma) - _gauss_cdf(a, sigma)
+
+
+def _int_x_pdf(a: float, b: float, sigma: float) -> float:
+    """∫[a,b] x·N(0,σ²)(x) dx = σ²·[f(a) − f(b)]."""
+    return sigma * sigma * (_gauss_pdf(a, sigma) - _gauss_pdf(b, sigma))
+
+
+def _int_sq_pdf(a: float, b: float, sigma: float, c: float) -> float:
+    """∫[a,b] (x−c)²·N(0,σ²)(x) dx — closed form."""
+    fa, fb = _gauss_pdf(a, sigma), _gauss_pdf(b, sigma)
+    cdf_diff = _gauss_cdf(b, sigma) - _gauss_cdf(a, sigma)
+    sig2 = sigma * sigma
+    return (
+        sig2 * (a * fa - b * fb)
+        - 2.0 * c * sig2 * (fa - fb)
+        + (sig2 + c * c) * cdf_diff
+    )
+
+
+def _quad(f, a: float, b: float, n: int = 200) -> float:
+    """Composite Simpson's rule numerical integration over [a, b].
+
+    Used only for the ``use_exact=True`` (Beta-PDF) path; the Gaussian path
+    uses closed-form helpers above.
+    """
+    if n % 2 != 0:
+        n += 1
+    h = (b - a) / n
+    s = f(a) + f(b)
+    for i in range(1, n):
+        s += (4 if i % 2 else 2) * f(a + i * h)
+    return h / 3.0 * s
+
+
+def beta_pdf(x: float, d: int) -> float:
+    """PDF of a single coordinate after random rotation of a d-dim unit vector."""
+    if abs(x) >= 1.0:
+        return 0.0
+    coeff = math.gamma(d / 2) / (math.sqrt(math.pi) * math.gamma((d - 1) / 2))
+    return coeff * (1 - x * x) ** ((d - 3) / 2)
+
+
+def gaussian_approx_pdf(x: float, d: int) -> float:
+    """Gaussian approximation N(0, 1/d) -- accurate for d >= 64."""
+    sigma2 = 1.0 / d
+    return (1.0 / math.sqrt(2 * math.pi * sigma2)) * math.exp(-x * x / (2 * sigma2))
+
+
+def solve_lloyd_max(
+    d: int, bits: int, use_exact: bool = False, max_iter: int = 200, tol: float = 1e-10
+):
+    """
+    Solve Lloyd-Max optimal quantizer for the coordinate distribution.
+
+    Args:
+        d: vector dimension
+        bits: number of quantization bits
+        use_exact: if True, use exact Beta PDF; if False, use Gaussian approx
+        max_iter: maximum Lloyd-Max iterations
+        tol: convergence tolerance
+
+    Returns:
+        centroids: sorted tensor of 2^bits optimal centroids
+        boundaries: sorted tensor of 2^bits - 1 boundaries between centroids
+    """
+    n_levels = 2**bits
+    sigma = 1.0 / math.sqrt(d)
+
+    lo, hi = -3.5 * sigma, 3.5 * sigma
+    centroids = [lo + (hi - lo) * (i + 0.5) / n_levels for i in range(n_levels)]
+
+    for _ in range(max_iter):
+        boundaries = [
+            (centroids[i] + centroids[i + 1]) / 2.0 for i in range(n_levels - 1)
+        ]
+        edges = [lo * 3] + boundaries + [hi * 3]
+        new_centroids = []
+        for i in range(n_levels):
+            a, b = edges[i], edges[i + 1]
+
+            if use_exact:
+                numerator = _quad(lambda x: x * beta_pdf(x, d), a, b)
+                denominator = _quad(lambda x: beta_pdf(x, d), a, b)
+            else:
+                numerator = _int_x_pdf(a, b, sigma)
+                denominator = _int_pdf(a, b, sigma)
+
+            if denominator > 1e-15:
+                new_centroids.append(numerator / denominator)
+            else:
+                new_centroids.append(centroids[i])
+
+        max_shift = max(abs(new_centroids[i] - centroids[i]) for i in range(n_levels))
+        centroids = new_centroids
+
+        if max_shift < tol:
+            break
+
+    boundaries = [(centroids[i] + centroids[i + 1]) / 2.0 for i in range(n_levels - 1)]
+
+    return (
+        torch.tensor(centroids, dtype=torch.float32),
+        torch.tensor(boundaries, dtype=torch.float32),
+    )
+
+
+def compute_expected_distortion(
+    d: int,
+    bits: int,
+    centroids: torch.Tensor,
+    boundaries: torch.Tensor,
+    use_exact: bool = False,
+) -> float:
+    """Compute the expected MSE distortion per coordinate for the given quantizer."""
+    sigma = 1.0 / math.sqrt(d)
+    n_levels = len(centroids)
+
+    edges = [-3.5 * sigma * 3] + boundaries.tolist() + [3.5 * sigma * 3]
+    total_distortion = 0.0
+
+    for i in range(n_levels):
+        a, b = edges[i], edges[i + 1]
+        c = centroids[i].item()
+        if use_exact:
+            dist = _quad(lambda x, _c=c: (x - _c) ** 2 * beta_pdf(x, d), a, b)
+        else:
+            dist = _int_sq_pdf(a, b, sigma, c)
+        total_distortion += dist
+
+    return total_distortion
+
+
+class LloydMaxCodebook:
+    """Precomputed Lloyd-Max codebook for a given dimension and bit-width."""
+
+    def __init__(self, d: int, bits: int, use_exact: bool = False):
+        self.d = d
+        self.bits = bits
+        self.n_levels = 2**bits
+        self.centroids, self.boundaries = solve_lloyd_max(d, bits, use_exact)
+        self.distortion = compute_expected_distortion(
+            d, bits, self.centroids, self.boundaries, use_exact
+        )
+
+    def quantize(self, x: torch.Tensor) -> torch.Tensor:
+        """Quantize values to nearest centroid indices."""
+        diffs = x.unsqueeze(-1) - self.centroids.to(x.device)
+        return diffs.abs().argmin(dim=-1)
+
+    def dequantize(self, indices: torch.Tensor) -> torch.Tensor:
+        """Map indices back to centroid values."""
+        return self.centroids.to(indices.device)[indices]
+
+    def __repr__(self):
+        return (
+            f"LloydMaxCodebook(d={self.d}, bits={self.bits}, "
+            f"levels={self.n_levels}, distortion_per_coord={self.distortion:.6f})"
+        )

plesk_unified/turboquant/turboquant.py

@@ -0,0 +1,249 @@
+"""TurboQuant: two-stage vector quantization."""
+
+from __future__ import annotations
+
+import math
+from typing import Optional, Tuple, cast
+
+import torch
+from torch import nn
+
+from .lloyd_max import LloydMaxCodebook
+
+
+def generate_rotation_matrix(
+    d: int, seed: Optional[int] = None, device: str = "cpu"
+) -> torch.Tensor:
+    """Generate a random orthogonal rotation matrix via QR decomposition."""
+    gen = torch.Generator(device="cpu")
+    if seed is not None:
+        gen.manual_seed(seed)
+    G = torch.randn(d, d, generator=gen)
+    Q, R = torch.linalg.qr(G)
+    diag_sign = torch.sign(torch.diag(R))
+    diag_sign[diag_sign == 0] = 1.0
+    Q = Q * diag_sign.unsqueeze(0)
+    return Q.to(device)
+
+
+def generate_qjl_matrix(
+    d: int, m: Optional[int] = None, seed: Optional[int] = None, device: str = "cpu"
+) -> torch.Tensor:
+    """
+    Generate the random projection matrix S for QJL.
+    S has i.i.d. N(0,1) entries, shape (m, d).
+    Default m = d (same dimensionality).
+    """
+    if m is None:
+        m = d
+    gen = torch.Generator(device="cpu")
+    if seed is not None:
+        gen.manual_seed(seed)
+    S = torch.randn(m, d, generator=gen)
+    return S.to(device)
+
+
+class TurboQuantMSE(nn.Module):
+    """Stage 1: MSE-optimal quantizer."""
+
+    def __init__(self, d: int, bits: int, seed: int = 42, device: str = "cpu"):
+        super().__init__()
+        self.d = d
+        self.bits = bits
+        self.device = device
+
+        self.register_buffer(
+            "Pi", generate_rotation_matrix(d, seed=seed, device=device)
+        )
+        self.codebook = LloydMaxCodebook(d, bits)
+        self.register_buffer("centroids", self.codebook.centroids.to(device))
+        self.register_buffer("boundaries", self.codebook.boundaries.to(device))
+
+    def rotate(self, x: torch.Tensor) -> torch.Tensor:
+        Pi = cast("torch.Tensor", self.Pi)
+        return x @ Pi.T
+
+    def unrotate(self, y: torch.Tensor) -> torch.Tensor:
+        Pi = cast("torch.Tensor", self.Pi)
+        return y @ Pi
+
+    def quantize(self, x: torch.Tensor) -> torch.Tensor:
+        centroids = cast("torch.Tensor", self.centroids)
+        y = self.rotate(x)
+        diffs = y.unsqueeze(-1) - centroids
+        indices = diffs.abs().argmin(dim=-1)
+        return indices
+
+    def dequantize(self, indices: torch.Tensor) -> torch.Tensor:
+        centroids = cast("torch.Tensor", self.centroids)
+        y_hat = centroids[indices]
+        return self.unrotate(y_hat)
+
+    def forward(self, x: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]:
+        indices = self.quantize(x)
+        x_hat = self.dequantize(indices)
+        return x_hat, indices
+
+
+class TurboQuantProd(nn.Module):
+    """Stage 1 + Stage 2: Unbiased inner product quantizer."""
+
+    def __init__(
+        self,
+        d: int,
+        bits: int,
+        qjl_dim: Optional[int] = None,
+        seed: int = 42,
+        device: str = "cpu",
+    ):
+        super().__init__()
+        self.d = d
+        self.bits = bits
+        self.mse_bits = max(bits - 1, 1)
+        self.qjl_dim = qjl_dim or d
+        self.device = device
+
+        self.mse = TurboQuantMSE(d, self.mse_bits, seed=seed, device=device)
+        self.register_buffer(
+            "S", generate_qjl_matrix(d, m=self.qjl_dim, seed=seed + 1, device=device)
+        )
+
+    def quantize(self, x: torch.Tensor) -> dict:
+        x_hat, mse_indices = self.mse(x)
+        residual = x - x_hat
+        residual_norm = torch.norm(residual, dim=-1, keepdim=True)
+        S = cast("torch.Tensor", self.S)
+        projected = residual @ S.T
+        qjl_signs = torch.sign(projected)
+        qjl_signs[qjl_signs == 0] = 1.0
+
+        return {
+            "mse_indices": mse_indices,
+            "qjl_signs": qjl_signs,
+            "residual_norm": residual_norm.squeeze(-1),
+        }
+
+    def dequantize(self, compressed: dict) -> torch.Tensor:
+        return self.mse.dequantize(compressed["mse_indices"])
+
+    def inner_product(self, y: torch.Tensor, compressed: dict) -> torch.Tensor:
+        x_mse = self.mse.dequantize(compressed["mse_indices"])
+        term1 = (y * x_mse).sum(dim=-1)
+
+        S = cast("torch.Tensor", self.S)
+        y_projected = y @ S.T
+        qjl_ip = (y_projected * compressed["qjl_signs"]).sum(dim=-1)
+
+        m = self.qjl_dim
+        correction_scale = math.sqrt(math.pi / 2) / m
+        term2 = compressed["residual_norm"] * correction_scale * qjl_ip
+
+        return term1 + term2
+
+    def forward(self, x: torch.Tensor) -> dict:
+        return self.quantize(x)
+
+
+class TurboQuantKVCache:
+    """KV cache wrapper that uses TurboQuant to compress keys and values."""
+
+    def __init__(
+        self,
+        d_key: int,
+        d_value: int,
+        bits: int = 3,
+        seed: int = 42,
+        device: str = "cpu",
+    ):
+        self.d_key = d_key
+        self.d_value = d_value
+        self.bits = bits
+        self.device = device
+
+        self.key_quantizer = TurboQuantProd(d_key, bits, seed=seed, device=device)
+        self.value_quantizer = TurboQuantMSE(
+            d_value, bits, seed=seed + 100, device=device
+        )
+
+        self.key_cache = []
+        self.value_cache = []
+
+    def append(self, keys: torch.Tensor, values: torch.Tensor):
+        orig_shape = keys.shape
+        flat_keys = keys.reshape(-1, self.d_key)
+        flat_values = values.reshape(-1, self.d_value)
+
+        compressed_keys = self.key_quantizer.quantize(flat_keys)
+        value_indices = self.value_quantizer.quantize(flat_values)
+
+        self.key_cache.append(
+            {
+                "mse_indices": compressed_keys["mse_indices"],
+                "qjl_signs": compressed_keys["qjl_signs"],
+                "residual_norm": compressed_keys["residual_norm"],
+                "shape": orig_shape,
+            }
+        )
+        self.value_cache.append(
+            {
+                "indices": value_indices,
+                "shape": values.shape,
+            }
+        )
+
+    def attention_scores(self, queries: torch.Tensor) -> torch.Tensor:
+        scores = []
+        for cached in self.key_cache:
+            s = self.key_quantizer.inner_product(queries, cached)
+            scores.append(s)
+        return torch.cat(scores, dim=-1) if scores else torch.tensor([])
+
+    def get_values(self) -> torch.Tensor:
+        values = []
+        for cached in self.value_cache:
+            v = self.value_quantizer.dequantize(cached["indices"])
+            values.append(v)
+        return torch.cat(values, dim=0) if values else torch.tensor([])
+
+    def memory_usage_bits(self) -> dict:
+        n_keys = (
+            sum(c["mse_indices"].numel() for c in self.key_cache)
+            if self.key_cache
+            else 0
+        )
+        n_qjl = (
+            sum(c["qjl_signs"].numel() for c in self.key_cache) if self.key_cache else 0
+        )
+        n_norms = (
+            sum(c["residual_norm"].numel() for c in self.key_cache)
+            if self.key_cache
+            else 0
+        )
+        n_values = (
+            sum(c["indices"].numel() for c in self.value_cache)
+            if self.value_cache
+            else 0
+        )
+
+        key_bits = n_keys * self.key_quantizer.mse_bits + n_qjl * 1 + n_norms * 16
+        value_bits = n_values * self.bits
+        fp16_equivalent = (n_keys + n_values) * 16
+
+        return {
+            "key_bits": key_bits,
+            "value_bits": value_bits,
+            "total_bits": key_bits + value_bits,
+            "fp16_bits": fp16_equivalent,
+            "compression_ratio": (
+                fp16_equivalent / (key_bits + value_bits)
+                if (key_bits + value_bits) > 0
+                else 0
+            ),
+        }
+
+    def __len__(self):
+        return (
+            sum(c["mse_indices"].shape[0] for c in self.key_cache)
+            if self.key_cache
+            else 0
+        )

plesk_unified/types.py

@@ -0,0 +1,27 @@
+from enum import Enum
+from typing import Literal, Union
+
+
+class CategoryEnum(str, Enum):
+    """Supported Plesk documentation categories."""
+
+    GUIDE = "guide"
+    CLI = "cli"
+    API = "api"
+    PHP_STUBS = "php-stubs"
+    JS_SDK = "js-sdk"
+
+
+VALID_CATEGORIES: frozenset[str] = frozenset(c.value for c in CategoryEnum)
+
+
+def validate_category(category: str, allow_all: bool = False) -> None:
+    """Validate that a category string is valid."""
+    if allow_all and category == "all":
+        return
+    if category not in VALID_CATEGORIES:
+        raise ValueError(f"Invalid category: '{category}'")
+
+
+# Type alias for refresh_knowledge which accepts a specific category or "all"
+CategoryOrAll = Union[CategoryEnum, Literal["all"]]