src/kbdebugger/extraction/utils.py

import os
from random import random
import re
from time import time
from typing import Any, Callable, List, Optional, TypeVar

from kbdebugger.novelty.types import NoveltyDecision
from kbdebugger.types import ExtractionResult, TripletSubjectObjectPredicate
from kbdebugger.utils.json import write_json
from kbdebugger.utils.time import now_utc_compact
from .types import Qualities
from typing import Any, Dict


def coerce_triplets(item: Dict[str, Any], fallback_sentence: str) -> ExtractionResult:
    """
    Coerce a single item dict to ExtractionResult:
    { "sentence": str, "triplets": list[TripletSubjectObjectPredicate] }
    """
    sentence = item.get("sentence", fallback_sentence)
    raw_triplets = item.get("triplets", [])
    triplets: list[TripletSubjectObjectPredicate] = []

    if isinstance(raw_triplets, list):
        for t in raw_triplets:
            if isinstance(t, (list, tuple)) and len(t) == 3:
                subj, obj, rel = t
                if all(isinstance(x, str) for x in (subj, obj, rel)):
                    triplets.append((subj.strip(), obj.strip(), rel.strip()))

    return {"sentence": str(sentence), "triplets": triplets}


def coerce_triplets_batch(obj: Dict[str, Any], sentences: List[str]) -> List[ExtractionResult]:
    """
    Coerce the LLM batch output of shape:
    {
      "triplets_batch": [
        {"id": 0, "sentence": "...", "triplets": [...]},
        ...
      ]
    }
    into a list[ExtractionResult], aligned by input index.
    """
    # Default: one empty result per input sentence
    empty: ExtractionResult = {"sentence": "", "triplets": []}
    results: List[ExtractionResult] = [empty for _ in sentences]

    batch = obj.get("triplets_batch", [])
    if not isinstance(batch, list):
        return results

    # Map by id, but also be robust
    for item in batch:
        if not isinstance(item, dict):
            continue

        idx = item.get("id")
        if isinstance(idx, int) and 0 <= idx < len(sentences):
            results[idx] = coerce_triplets(item, sentences[idx])

    # Fill any missing entries with fallback (no triplets)
    for i, res in enumerate(results):
        if res["sentence"] == "":
            results[i] = {"sentence": sentences[i], "triplets": []}

    return results


def coerce_qualities(obj: Dict) -> Qualities:
    if not isinstance(obj, dict):
        return []
    qualities = obj.get("qualities")
    if not isinstance(qualities, list):
        return []
    out: Qualities = []
    for q in qualities:
        if isinstance(q, str):
            s = q.strip()
            if s:
                out.append(s)
    return out


def save_results_json(results: List[ExtractionResult]) -> None:
    """
    Write extraction results to a JSON file.
    """
    created_at = now_utc_compact()
    data = {
        "results": results,
    }
    path = f"logs/05_triplet_extraction_results_{created_at}.json"
    write_json(path, data)
    print(f"\n[INFO] Wrote JSON results to {path}")


# ---------------------------------------------------------------------------
# Helpers for `build_chunk_batch_decomposer`
# ---------------------------------------------------------------------------
_WS_RE = re.compile(r"\s+") # this matches all whitespace sequences i.e. newlines, tabs, multiple spaces, etc.

def sanitize_chunk(text: str) -> str:
    """
    Normalize a chunk into a single-line string.

    We intentionally avoid aggressive cleaning here: the upstream PDF cleaning
    stage already handles boilerplate/DOI stripping etc. Our goal is only to
    prevent formatting artifacts from confusing the LLM.
    """
    # replace all whitespace sequences (newlines, tabs, multiple spaces) with single space " "
    return _WS_RE.sub(" ", text or "").strip()


def coerce_batch_qualities(
    obj: Any,
    *,
    expected_n: int,
) -> Dict[int, Qualities]:
    """
    Parse the JSON object returned by the batch prompt into an id->qualities map.

    Expected schema (strict, by prompt contract):
        {
          "results": [
            {"id": 0, "qualities": ["...", "..."]},
            {"id": 1, "qualities": []}
          ]
        }

    This parser is defensive:
    - Accepts "id" as int or numeric string.
    - Accepts "qualities" as list[str] or other coercible structures.
    - Ignores unknown items; only keeps ids within range.
    - Returns a possibly sparse mapping; caller fills missing ids with [].
    """
    if not isinstance(obj, dict):
        return {}

    results = obj.get("results")
    if not isinstance(results, list):
        return {}

    out: Dict[int, Qualities] = {}

    for item in results:
        if not isinstance(item, dict):
            continue

        raw_id = item.get("id")
        if raw_id is None:
            continue

        # Coerce id -> int if possible
        chunk_id: Optional[int] = None
        if isinstance(raw_id, int):
            chunk_id = raw_id
        elif isinstance(raw_id, str) and raw_id.strip().isdigit():
            chunk_id = int(raw_id.strip())

        if chunk_id is None:
            continue
        if chunk_id < 0 or chunk_id >= expected_n:
            continue

        raw_qualities = item.get("qualities", [])
        # Try to coerce qualities robustly.
        # - If it's already a list, keep string-like entries.
        # - If it's a dict (rare), attempt coerce_qualities on it.
        qualities: Qualities = []

        if isinstance(raw_qualities, list):
            qualities = [str(x).strip() for x in raw_qualities if str(x).strip()]
        else:
            # Some models might accidentally return {"qualities": [...]} per item.
            # coerce_qualities can often salvage this.
            try:
                qualities = coerce_qualities(raw_qualities)  # type: ignore[arg-type]
            except Exception:
                qualities = []

        out[chunk_id] = qualities

    return out

def load_triplet_qualifying_decisions() -> set[NoveltyDecision]:
    """
    Load which novelty decisions qualify a quality for triplet extraction.

    Environment variable:
        KB_TRIPLET_QUALIFY_DECISIONS=PARTIALLY_NEW,NEW

    Defaults to:
        {"PARTIALLY_NEW", "NEW"}
    """
    raw = os.getenv("KB_TRIPLET_QUALIFY_DECISIONS", "").strip()

    fallback = {
        NoveltyDecision.PARTIALLY_NEW,
        NoveltyDecision.NEW,
    }

    if not raw:
        return fallback
    
    decisions: set[NoveltyDecision] = set()
    for token in raw.split(","):
        token = token.strip().upper()
        if not token:
            continue
        try:
            decisions.add(NoveltyDecision(token))
        except ValueError:
            # Ignore unknown tokens silently
            continue

    # Safety fallback
    if not decisions:
        decisions = fallback

    return decisions

# ---------------------------------------------------------------------------
# Parallelism helpers
# ---------------------------------------------------------------------------
T = TypeVar("T")

_RETRY_AFTER_RE = re.compile(r"try again in\s+([0-9]*\.?[0-9]+)s", re.IGNORECASE)


def _extract_retry_after_seconds(error_text: str) -> Optional[float]:
    """
    Extract a retry delay (in seconds) from Groq-style 429 error messages.

    Example message fragment:
        "Please try again in 13.45s."

    Returns
    -------
    Optional[float]
        The parsed delay in seconds if present, otherwise None.
    """
    match = _RETRY_AFTER_RE.search(error_text or "")
    if not match:
        return None

    try:
        return float(match.group(1))
    except ValueError:
        return None
    

def _call_with_rate_limit_retries(
    fn: Callable[[], T],
    *,
    max_retries: int = 8,
    default_backoff_s: float = 2.0,
    max_sleep_s: float = 30.0,
) -> T:
    """
    Call `fn()` with rate-limit-aware retries.

    Strategy
    --------
    - If the exception message contains "try again in Xs", sleep for X seconds
      (+ small jitter) and retry.
    - Otherwise, use a conservative exponential backoff.

    Why this exists
    ---------------
    Groq on-demand has strict TPM (Token-per-Minute) limits. When we batch or parallelize,
    occasional 429s are expected. Dropping a batch silently corrupts results.

    Raises
    ------
    RuntimeError
        If all retry attempts fail.
    """
    last_err: Exception | None = None

    for attempt in range(1, max_retries + 1):
        try:
            return fn()
        except Exception as e:  # SDKs often raise generic exceptions
            last_err = e
            msg = str(e)

            retry_after = _extract_retry_after_seconds(msg)
            if retry_after is not None:
                # Add a tiny jitter to avoid synchronizing retries across threads.
                sleep_s = retry_after + random.uniform(0.1, 0.4)
            else:
                # Exponential backoff for unknown transient failures
                sleep_s = default_backoff_s * (2 ** (attempt - 1))

            sleep_s = min(sleep_s, max_sleep_s)
            time.sleep(sleep_s)

    raise RuntimeError(f"LLM call failed after {max_retries} retries") from last_err