diffprompt 0.1.0__py3-none-any.whl

diffprompt/core/ontology.py

@@ -0,0 +1,160 @@
+"""
+Ontology inference and anchor-based tagging.
+
+Two steps:
+1. Infer dimensions + tags from the prompt (one LLM call)
+2. Generate anchor sentences per tag (one LLM call per tag)
+3. Tag test inputs using embedding similarity to anchors (no LLM)
+"""
+from __future__ import annotations
+import json
+import re
+from typing import Optional
+from sentence_transformers import SentenceTransformer
+from sklearn.metrics.pairwise import cosine_similarity
+import numpy as np
+
+from diffprompt.models.cascade import call_cascade
+
+
+INFER_PROMPT = """You are designing a test suite for an LLM prompt.
+
+Prompt: {prompt}
+
+Identify 3-4 dimensions that would reveal BEHAVIORAL differences in how this prompt responds.
+Focus on: user intent, emotional state, topic complexity, and request type.
+Avoid surface-level dimensions like length, format, or punctuation style.
+
+Return ONLY valid JSON. Example:
+{{
+  "tone": ["formal", "casual", "emotional", "urgent"],
+  "complexity": ["simple", "multi-part", "ambiguous"],
+  "intent": ["lookup", "reasoning", "emotional-support"]
+}}
+
+Each dimension must predict a meaningfully different response from this specific prompt."""
+
+
+ANCHOR_PROMPT = """For the tag "{tag}" in dimension "{dimension}", write 3 short example sentences
+that real users would send to this prompt: {prompt}
+
+Each sentence must clearly represent "{tag}" and be distinct from each other.
+Return ONLY a JSON array of 3 strings. No explanation."""
+
+
+def _extract_json_object(raw: str) -> str:
+    """Extract the first {...} block from raw LLM output."""
+    clean = re.sub(r"```(?:json)?|```", "", raw).strip()
+    match = re.search(r'\{.*\}', clean, re.DOTALL)
+    return match.group(0) if match else clean
+
+
+def _extract_json_array(raw: str) -> str:
+    """Extract the first [...] block from raw LLM output."""
+    clean = re.sub(r"```(?:json)?|```", "", raw).strip()
+    match = re.search(r'\[.*\]', clean, re.DOTALL)
+    return match.group(0) if match else clean
+
+
+class Ontology:
+    def __init__(self):
+        self.dimensions: dict[str, list[str]] = {}
+        self.anchors: dict[str, dict[str, list[str]]] = {}
+        self.anchor_embeddings: dict[str, dict[str, np.ndarray]] = {}
+        self._embedder: Optional[SentenceTransformer] = None
+
+    @property
+    def embedder(self) -> SentenceTransformer:
+        if self._embedder is None:
+            import io, contextlib
+            f = io.StringIO()
+            with contextlib.redirect_stdout(f), contextlib.redirect_stderr(f):
+                self._embedder = SentenceTransformer("all-MiniLM-L6-v2")
+        return self._embedder
+
+    async def infer(self, prompt: str, local_only: bool = False) -> None:
+        """Infer dimensions and tags from the prompt. One LLM call."""
+        raw, _ = await call_cascade(
+            INFER_PROMPT.format(prompt=prompt),
+            local_only=local_only,
+        )
+        try:
+            parsed = json.loads(_extract_json_object(raw))
+            if len(parsed) == 1 and isinstance(list(parsed.values())[0], dict):
+                parsed = list(parsed.values())[0]
+            self.dimensions = {k: v for k, v in parsed.items() if isinstance(v, list)}
+        except json.JSONDecodeError:
+            self.dimensions = {
+                "user_intent":    ["informational", "transactional", "complaint", "other"],
+                "emotional_state": ["neutral", "frustrated", "confused", "urgent"],
+                "request_type":   ["specific", "open_ended", "troubleshooting"],
+            }
+
+    async def build_anchors(self, prompt: str, local_only: bool = False) -> None:
+        """Generate 3 anchor sentences per tag. One LLM call per tag."""
+        for dimension, tags in self.dimensions.items():
+            self.anchors[dimension] = {}
+            for tag in tags:
+                raw, _ = await call_cascade(
+                    ANCHOR_PROMPT.format(tag=tag, dimension=dimension, prompt=prompt),
+                    local_only=local_only,
+                )
+                try:
+                    parsed = json.loads(_extract_json_array(raw))
+                    if isinstance(parsed, list) and parsed:
+                        self.anchors[dimension][tag] = [str(s) for s in parsed]
+                    else:
+                        self.anchors[dimension][tag] = [raw.strip()[:200]]
+                except Exception:
+                    self.anchors[dimension][tag] = [raw.strip()[:200]]
+
+        # Build embeddings for all anchors
+        for dimension, tag_anchors in self.anchors.items():
+            self.anchor_embeddings[dimension] = {}
+            for tag, sentences in tag_anchors.items():
+                self.anchor_embeddings[dimension][tag] = self.embedder.encode(sentences)
+
+    def tag(self, input_text: str) -> dict[str, str]:
+        """
+        Tag an input using max similarity across multiple anchors per tag.
+        No LLM call — pure embeddings. Deterministic and fast.
+        """
+        if not self.anchor_embeddings:
+            return {}
+
+        input_emb = self.embedder.encode([input_text])
+        result = {}
+
+        for dimension, tags in self.dimensions.items():
+            if dimension not in self.anchor_embeddings:
+                continue
+            best_tag = tags[0]
+            best_score = -1.0
+
+            for tag in tags:
+                if tag not in self.anchor_embeddings[dimension]:
+                    continue
+                anchor_embs = self.anchor_embeddings[dimension][tag]
+                sims = cosine_similarity(input_emb, anchor_embs)[0]
+                score = float(np.max(sims))
+                if score > best_score:
+                    best_score = score
+                    best_tag = tag
+
+            result[dimension] = best_tag
+
+        return result
+
+    def to_dict(self) -> dict:
+        return {"dimensions": self.dimensions, "anchors": self.anchors}
+
+    @classmethod
+    def from_dict(cls, data: dict) -> "Ontology":
+        o = cls()
+        o.dimensions = data["dimensions"]
+        o.anchors = data["anchors"]
+        for dimension, tag_anchors in o.anchors.items():
+            o.anchor_embeddings[dimension] = {}
+            for tag, sentences in tag_anchors.items():
+                o.anchor_embeddings[dimension][tag] = o.embedder.encode(sentences)
+        return o

diffprompt/core/runner.py

@@ -0,0 +1,65 @@
+"""
+Runs both prompts on all test cases.
+Returns RunResult for each (test_case, prompt_version) pair.
+"""
+from __future__ import annotations
+import asyncio
+import time
+from diffprompt.models import TestCase, RunResult
+from diffprompt.models.cascade import call_cascade
+
+
+async def run_single(
+    test_case: TestCase,
+    prompt: str,
+    version: str,
+    model: str,
+    local_only: bool = False,
+) -> RunResult:
+    start = time.monotonic()
+    output, model_used = await call_cascade(
+        test_case.input,
+        system=prompt,
+        local_only=local_only,
+    )
+    latency_ms = (time.monotonic() - start) * 1000
+
+    return RunResult(
+        test_id=test_case.id,
+        prompt_version=version,
+        output=output,
+        model_used=model_used,
+        latency_ms=latency_ms,
+    )
+
+
+async def run_both(
+    test_cases: list[TestCase],
+    prompt_v1: str,
+    prompt_v2: str,
+    model: str = "groq/llama-3.3-70b-versatile",
+    local_only: bool = False,
+    concurrency: int = 5,
+) -> tuple[dict[str, RunResult], dict[str, RunResult]]:
+    """
+    Run both prompts on all test cases concurrently.
+    Returns (v1_results, v2_results) as dicts keyed by test_id.
+    """
+    semaphore = asyncio.Semaphore(concurrency)
+
+    async def run_with_sem(tc, prompt, version):
+        async with semaphore:
+            return await run_single(tc, prompt, version, model, local_only)
+
+    # Fire all tasks concurrently (bounded by semaphore)
+    v1_tasks = [run_with_sem(tc, prompt_v1, "v1") for tc in test_cases]
+    v2_tasks = [run_with_sem(tc, prompt_v2, "v2") for tc in test_cases]
+
+    v1_raw, v2_raw = await asyncio.gather(
+        asyncio.gather(*v1_tasks),
+        asyncio.gather(*v2_tasks),
+    )
+
+    v1_results = {r.test_id: r for r in v1_raw}
+    v2_results = {r.test_id: r for r in v2_raw}
+    return v1_results, v2_results

diffprompt/core/scorer.py

@@ -0,0 +1,102 @@
+"""
+Regression scoring and key example selection.
+"""
+from __future__ import annotations
+import numpy as np
+from diffprompt.models import DiffResult, KeyExample, Verdict
+from diffprompt.models.cascade import call_cascade
+
+
+WHY_IT_MATTERS_PROMPT = """
+Input: {input}
+V1 output: {v1}
+V2 output: {v2}
+Verdict: {verdict}
+
+In ONE sentence, explain what this reveals about the prompt change.
+Focus on the mechanism — why did v2 behave differently here?
+"""
+
+
+def regression_score(diffs: list[DiffResult]) -> float:
+    if not diffs:
+        return 50.0
+    total_weight = sum(d.divergence for d in diffs)
+    if total_weight == 0:
+        return 100.0
+    weighted_sum = sum(
+        d.divergence * (1.0 if d.verdict == Verdict.IMPROVEMENT else -1.0 if d.verdict == Verdict.REGRESSION else 0.0)
+        for d in diffs
+    )
+    return float(round(((weighted_sum / total_weight) + 1) / 2 * 100, 1))
+
+
+def importance_score(diff: DiffResult) -> float:
+    divergence   = diff.divergence
+    centrality   = diff.cluster_centrality
+    input_length = len(diff.test_case.input.split())
+    simplicity   = max(0.0, 1.0 - input_length / 50)
+    surprise     = divergence * simplicity
+    return 0.4 * divergence + 0.3 * centrality + 0.3 * surprise
+
+
+async def select_key_examples(
+    diffs: list[DiffResult],
+    top_n: int = 3,
+    local_only: bool = False,
+) -> list[KeyExample]:
+    if not diffs:
+        return []
+
+    for d in diffs:
+        d.importance_score = importance_score(d)
+
+    selected: list[tuple[str, DiffResult]] = []
+    used_ids: set[str] = set()
+
+    def _pick(slot: str, candidate) -> None:
+        if candidate and candidate.test_case.id not in used_ids:
+            selected.append((slot, candidate))
+            used_ids.add(candidate.test_case.id)
+
+    _pick("most_important", max(diffs, key=lambda d: d.importance_score))
+
+    improvements = [d for d in diffs if d.verdict == Verdict.IMPROVEMENT and d.test_case.id not in used_ids]
+    if improvements:
+        _pick("best_improvement", max(improvements, key=lambda d: d.divergence))
+
+    remaining = [d for d in diffs if d.test_case.id not in used_ids]
+    if remaining:
+        _pick("most_surprising", max(
+            remaining,
+            key=lambda d: d.divergence * max(0, 1 - len(d.test_case.input.split()) / 50),
+        ))
+
+    if top_n > 3:
+        regressions = sorted(
+            [d for d in diffs if d.verdict == Verdict.REGRESSION and d.test_case.id not in used_ids],
+            key=lambda d: d.divergence, reverse=True,
+        )
+        for d in regressions[:top_n - 3]:
+            _pick(f"regression_{len(selected)}", d)
+
+    selected = selected[:top_n]
+
+    examples = []
+    for slot_name, diff in selected:
+        why, _ = await call_cascade(
+            WHY_IT_MATTERS_PROMPT.format(
+                input=diff.test_case.input,
+                v1=diff.v1_output[:600],
+                v2=diff.v2_output[:600],
+                verdict=diff.verdict.value,
+            ),
+            local_only=local_only,
+        )
+        examples.append(KeyExample(
+            slot=slot_name,
+            diff=diff,
+            why_it_matters=why.strip(),
+        ))
+
+    return examples

diffprompt/core/slicer.py

@@ -0,0 +1,152 @@
+"""
+Behavioral slicing.
+Groups diffs by input tags and computes per-slice performance.
+Recursive splitting for high-variance slices (max depth 3).
+"""
+from __future__ import annotations
+import numpy as np
+from diffprompt.models import DiffResult, SliceResult, Verdict, TestCategory
+
+
+MIN_SLICE_SIZE = 5
+MAX_DEPTH = 3
+VARIANCE_THRESHOLD = 0.1   # stop splitting if variance is already low
+HIGH_VARIANCE_THRESHOLD = 0.15  # trigger recursive split above this
+
+
+def compute_slices(diffs: list[DiffResult]) -> list[SliceResult]:
+    """Compute behavioral slices across all tag dimensions."""
+    if not diffs or not diffs[0].test_case.tags:
+        return []
+
+    dimensions = list(diffs[0].test_case.tags.keys())
+    slices = []
+
+    for dimension in dimensions:
+        # Group by tag value
+        groups: dict[str, list[DiffResult]] = {}
+        for d in diffs:
+            val = d.test_case.tags.get(dimension, "unknown")
+            groups.setdefault(val, []).append(d)
+
+        for value, group_diffs in groups.items():
+            slice_result = _compute_slice(
+                dimension=dimension,
+                value=value,
+                diffs=group_diffs,
+                depth=1,
+            )
+            if slice_result:
+                slices.append(slice_result)
+
+            # Recursive split if high variance + enough data + not too deep
+            if (
+                slice_result
+                and slice_result.variance > HIGH_VARIANCE_THRESHOLD
+                and len(group_diffs) >= MIN_SLICE_SIZE * 2
+                and slice_result.depth < MAX_DEPTH
+            ):
+                sub_slices = _recursive_split(group_diffs, dimension, depth=2)
+                slices.extend(sub_slices)
+
+    # Sort by mean_similarity ascending (worst slices first)
+    slices.sort(key=lambda s: s.mean_similarity)
+    return slices
+
+
+def _compute_slice(
+    dimension: str,
+    value: str,
+    diffs: list[DiffResult],
+    depth: int,
+) -> SliceResult | None:
+    if len(diffs) < 2:
+        return None
+
+    sims = [d.similarity for d in diffs]
+    mean_sim = float(np.mean(sims))
+    variance = float(np.var(sims))
+
+    typical_ratio = sum(
+        1 for d in diffs if d.test_case.category == TestCategory.TYPICAL
+    ) / len(diffs)
+
+    confidence = _compute_confidence(
+        n=len(diffs),
+        variance=variance,
+        typical_ratio=typical_ratio,
+    )
+
+    verdicts = [d.verdict for d in diffs]
+    verdict = _aggregate_verdict(verdicts)
+
+    return SliceResult(
+        dimension=dimension,
+        value=value,
+        label=f"{dimension}:{value}",
+        n=len(diffs),
+        mean_similarity=mean_sim,
+        variance=variance,
+        typical_ratio=typical_ratio,
+        confidence=confidence,
+        verdict=verdict,
+        depth=depth,
+    )
+
+
+def _recursive_split(
+    diffs: list[DiffResult],
+    parent_dimension: str,
+    depth: int,
+) -> list[SliceResult]:
+    """Try splitting a high-variance slice by other dimensions."""
+    if depth > MAX_DEPTH or not diffs:
+        return []
+
+    other_dimensions = [
+        k for k in diffs[0].test_case.tags.keys()
+        if k != parent_dimension
+    ]
+
+    results = []
+    for dimension in other_dimensions:
+        groups: dict[str, list[DiffResult]] = {}
+        for d in diffs:
+            val = d.test_case.tags.get(dimension, "unknown")
+            groups.setdefault(val, []).append(d)
+
+        for value, group in groups.items():
+            if len(group) < MIN_SLICE_SIZE:
+                continue
+            slice_result = _compute_slice(
+                dimension=f"{parent_dimension}+{dimension}",
+                value=value,
+                diffs=group,
+                depth=depth,
+            )
+            if slice_result:
+                results.append(slice_result)
+
+    return results
+
+
+def _compute_confidence(n: int, variance: float, typical_ratio: float) -> float:
+    """
+    Confidence = f(variance, typical_ratio, n).
+    Variance is checked first — high variance = impure slice = low confidence.
+    """
+    # Variance penalty (primary signal)
+    variance_score = max(0.0, 1.0 - variance * 5)
+
+    # Typical ratio bonus (real-world signal)
+    typical_score = typical_ratio
+
+    # N score (statistical power)
+    n_score = min(1.0, n / 20)
+
+    return float(0.5 * variance_score + 0.3 * typical_score + 0.2 * n_score)
+
+
+def _aggregate_verdict(verdicts: list[Verdict]) -> Verdict:
+    counts = {v: verdicts.count(v) for v in Verdict}
+    return max(counts, key=counts.get)

diffprompt/models/__init__.py

@@ -0,0 +1,105 @@
+"""
+Core data models for diffprompt.
+These types flow through the entire pipeline — generator → runner → diff → analysis → output.
+"""
+from __future__ import annotations
+from enum import Enum
+from typing import Optional
+from pydantic import BaseModel, Field
+
+
+class TestCategory(str, Enum):
+    TYPICAL     = "typical"
+    BOUNDARY    = "boundary"
+    ADVERSARIAL = "adversarial"
+    FORMAT      = "format"
+
+
+class Verdict(str, Enum):
+    IMPROVEMENT = "improvement"
+    REGRESSION  = "regression"
+    NEUTRAL     = "neutral"
+
+
+class OutputFormat(str, Enum):
+    TERMINAL = "terminal"
+    JSON     = "json"
+    HTML     = "html"
+    TXT      = "txt"
+
+
+class TestCase(BaseModel):
+    id: str
+    input: str
+    category: TestCategory
+    tags: dict[str, str] = Field(default_factory=dict)
+
+
+class RunResult(BaseModel):
+    test_id: str
+    prompt_version: str
+    output: str
+    model_used: str
+    latency_ms: Optional[float] = None
+
+
+class DiffResult(BaseModel):
+    test_case: TestCase
+    v1_output: str
+    v2_output: str
+    similarity: float
+    divergence: float
+    verdict: Verdict
+    reason: str
+    judge_confidence: float
+    importance_score: float = 0.0
+    cluster_label: int = -1
+    cluster_centrality: float = 0.0
+
+
+class SliceResult(BaseModel):
+    dimension: str
+    value: str
+    label: str
+    n: int
+    mean_similarity: float
+    variance: float
+    typical_ratio: float
+    confidence: float
+    verdict: Verdict
+    depth: int = 1
+
+
+class Cluster(BaseModel):
+    label: int
+    name: str
+    description: str
+    n: int
+    mean_similarity: float
+    test_ids: list[str]
+
+
+class KeyExample(BaseModel):
+    slot: str
+    diff: DiffResult
+    why_it_matters: str
+
+
+class DiffReport(BaseModel):
+    prompt_v1: str
+    prompt_v2: str
+    model: str
+    judge: str
+    test_cases: list[TestCase]
+    diversity_score: float
+    diffs: list[DiffResult]
+    slices: list[SliceResult]
+    clusters: list[Cluster]
+    unclustered: list[DiffResult]
+    key_examples: list[KeyExample]
+    regression_score: float
+    n_improved: int
+    n_regressed: int
+    n_neutral: int
+    verdict: Verdict
+    recommendation: str