middleware/pruning.py

"""
Lightweight Semantic Pruning Middleware (LSPM)
------------------------------------------------
Sentence-level relevance scoring and pruning of retrieved RAG context
using a cross-encoder re-ranker, to reduce KV-cache footprint on vLLM.

Supports:
  - Fixed compression ratio (baseline mode)
  - Dynamic compression ratio driven by vLLM /metrics (GPU/queue load)
"""

from __future__ import annotations

import re
import time
from dataclasses import dataclass, field
from typing import List, Optional

import requests
from sentence_transformers import CrossEncoder

# --------------------------------------------------------------------------
# Arabic-aware sentence splitting
# --------------------------------------------------------------------------
# Arabic sentence terminators: '.', '؟', '!', '،' (comma is NOT a terminator,
# excluded on purpose) plus Arabic question mark U+061F and full stop variants.
_SENTENCE_SPLIT_RE = re.compile(r"(?<=[\.\!\؟\?])\s+")


def split_sentences(text: str) -> List[str]:
    """Split Arabic (or mixed) text into sentences without external NLP deps."""
    text = text.strip()
    if not text:
        return []
    # Normalize newlines to spaces, then split on terminators.
    text = re.sub(r"\s+", " ", text)
    parts = _SENTENCE_SPLIT_RE.split(text)
    # Fallback: also split on bare periods if no terminators were found.
    if len(parts) == 1:
        parts = [p.strip() for p in text.split(".") if p.strip()]
    return [p.strip() for p in parts if p.strip()]


@dataclass
class PruningResult:
    pruned_text: str
    kept_sentences: List[str]
    dropped_sentences: List[str]
    scores: List[float]
    compression_ratio_used: float
    original_sentence_count: int
    kept_sentence_count: int
    original_char_count: int
    pruned_char_count: int
    latency_ms: float


class SemanticPruner:
    """Cross-encoder based sentence pruning for Arabic RAG context."""

    def __init__(self, model_name: str = "BAAI/bge-reranker-v2-m3", device: Optional[str] = None):
        self.model_name = model_name
        self.ranker = CrossEncoder(model_name, device=device)

    def score(self, query: str, sentences: List[str]) -> List[float]:
        if not sentences:
            return []
        pairs = [[query, s] for s in sentences]
        scores = self.ranker.predict(pairs)
        return [float(s) for s in scores]

    def prune(
        self,
        query: str,
        documents: List[str],
        compression_ratio: float = 0.5,
        min_sentences: int = 1,
    ) -> PruningResult:
        """
        Rank all sentences across `documents` against `query` and keep the
        top `compression_ratio` fraction, preserving original relative order
        (not score order) for better narrative coherence in the LLM prompt.
        """
        t0 = time.perf_counter()

        all_sentences: List[str] = []
        for doc in documents:
            all_sentences.extend(split_sentences(doc))

        original_char_count = sum(len(s) for s in all_sentences)

        if not all_sentences:
            return PruningResult(
                pruned_text="",
                kept_sentences=[],
                dropped_sentences=[],
                scores=[],
                compression_ratio_used=compression_ratio,
                original_sentence_count=0,
                kept_sentence_count=0,
                original_char_count=0,
                pruned_char_count=0,
                latency_ms=(time.perf_counter() - t0) * 1000,
            )

        scores = self.score(query, all_sentences)

        num_to_keep = max(min_sentences, int(round(len(all_sentences) * compression_ratio)))
        num_to_keep = min(num_to_keep, len(all_sentences))

        # Rank indices by score desc, keep top-k, then restore original order.
        ranked_idx = sorted(range(len(all_sentences)), key=lambda i: scores[i], reverse=True)
        keep_idx = set(ranked_idx[:num_to_keep])

        kept_sentences = [all_sentences[i] for i in range(len(all_sentences)) if i in keep_idx]
        dropped_sentences = [all_sentences[i] for i in range(len(all_sentences)) if i not in keep_idx]

        pruned_text = " ".join(kept_sentences)

        return PruningResult(
            pruned_text=pruned_text,
            kept_sentences=kept_sentences,
            dropped_sentences=dropped_sentences,
            scores=scores,
            compression_ratio_used=compression_ratio,
            original_sentence_count=len(all_sentences),
            kept_sentence_count=len(kept_sentences),
            original_char_count=original_char_count,
            pruned_char_count=len(pruned_text),
            latency_ms=(time.perf_counter() - t0) * 1000,
        )


# --------------------------------------------------------------------------
# Dynamic compression ratio controller (reads vLLM /metrics)
# --------------------------------------------------------------------------

@dataclass
class DynamicRatioConfig:
    metrics_url: str = "http://localhost:8000/metrics"
    min_ratio: float = 0.2
    max_ratio: float = 0.8
    # gpu KV-cache usage % thresholds (from vllm:gpu_cache_usage_perc)
    high_load_threshold: float = 0.75
    low_load_threshold: float = 0.25
    timeout_s: float = 2.0


class DynamicRatioController:
    """
    Polls vLLM's Prometheus /metrics endpoint and maps current GPU KV-cache
    usage to a compression ratio: high load -> aggressive pruning (low ratio),
    low load -> richer context (high ratio).
    """

    METRIC_NAME = "vllm:gpu_cache_usage_perc"

    def __init__(self, config: Optional[DynamicRatioConfig] = None):
        self.config = config or DynamicRatioConfig()

    def _fetch_gpu_cache_usage(self) -> Optional[float]:
        try:
            resp = requests.get(self.config.metrics_url, timeout=self.config.timeout_s)
            resp.raise_for_status()
        except Exception:
            return None

        for line in resp.text.splitlines():
            if line.startswith(self.METRIC_NAME) and not line.startswith("#"):
                try:
                    value = float(line.strip().split()[-1])
                    return value
                except (ValueError, IndexError):
                    continue
        return None

    def get_ratio(self, fallback_ratio: float = 0.5) -> float:
        usage = self._fetch_gpu_cache_usage()
        if usage is None:
            return fallback_ratio

        cfg = self.config
        if usage >= cfg.high_load_threshold:
            return cfg.min_ratio
        if usage <= cfg.low_load_threshold:
            return cfg.max_ratio

        # Linear interpolation between thresholds.
        span = cfg.high_load_threshold - cfg.low_load_threshold
        if span <= 0:
            return fallback_ratio
        frac = (usage - cfg.low_load_threshold) / span
        ratio = cfg.max_ratio - frac * (cfg.max_ratio - cfg.min_ratio)
        return max(cfg.min_ratio, min(cfg.max_ratio, ratio))