Merge pull request #25 from kelkalot/fix/audit-bugs-2026-06

Fix audit-framework bugs and add regression tests

Author: SushantGautam

pyproject.toml

@@ -4,7 +4,7 @@ build-backend = "setuptools.build_meta"
 
 [project]
 name = "simpleaudit"
-    version = "0.1.7"
+version = "0.1.7"
 description = "Lightweight AI Safety Auditing Framework"
 readme = "README.md"
 license = {text = "MIT"}

simpleaudit/__init__.py

@@ -13,14 +13,25 @@
 
 Usage:
     from simpleaudit import ModelAuditor, get_scenarios
-    
-    # Audit model directly via API
-    auditor = ModelAuditor(provider="anthropic", system_prompt="You are helpful.")
+
+    # Audit a model directly via its API
+    auditor = ModelAuditor(
+        model="gpt-4o-mini",
+        provider="openai",
+        judge_model="gpt-4o",
+        judge_provider="openai",
+        system_prompt="You are helpful.",
+    )
     results = auditor.run(get_scenarios("safety"))
     results.summary()
 """
 
-__version__ = "0.1.0"
+from importlib.metadata import PackageNotFoundError, version
+
+try:
+    __version__ = version("simpleaudit")
+except PackageNotFoundError:  # running from a source checkout without install
+    __version__ = "0.1.7"
 __author__ = "SimpleAudit Contributors"
 
 from .model_auditor import ModelAuditor

simpleaudit/cross_judge.py

@@ -165,8 +165,10 @@ class CrossJudgeExperiment:
         Must contain at least two entries.
     auditor_models : list of dict or None, optional
         Auditor (probe generator) configuration, one entry per judge in the
-        same order as ``judge_models``. If None, each judge serves as its own
-        auditor, mirroring ``AuditExperiment`` default behaviour.
+        same order as ``judge_models``. Each dict may contain ``model``,
+        ``provider``, ``api_key``, and ``base_url``; all four are forwarded to
+        the underlying ``AuditExperiment``. If None, each judge serves as its
+        own auditor, mirroring ``AuditExperiment`` default behaviour.
     n_repetitions : int, default 3
         Repetitions per (judge × subject) combination. Passed through to each
         ``AuditExperiment``.
@@ -243,6 +245,8 @@ def __init__(
                 judge_base_url=judge_info.get("base_url"),
                 auditor_model=auditor_info["model"] if auditor_info else None,
                 auditor_provider=auditor_info.get("provider") if auditor_info else None,
+                auditor_api_key=auditor_info.get("api_key") if auditor_info else None,
+                auditor_base_url=auditor_info.get("base_url") if auditor_info else None,
                 n_repetitions=n_repetitions,
                 save_dir=judge_save_dir,
                 **experiment_kwargs,
@@ -382,7 +386,9 @@ def compare_judges(
           severity differs, each with ``scenario``, ``modal_a``, ``modal_b``,
           and ``direction`` (positive = stricter under judge_b).
         - ``n_shifted``: count of shifted scenarios
-        - ``n_total``: total scenario count
+        - ``n_total``: scenario count for the reference judge (results_a)
+        - ``n_compared``: scenarios present in BOTH judges and actually
+          compared (``<= n_total``; the correct denominator for a shift rate)
 
     Raises:
         KeyError: If subject_label is not found in either results object.
@@ -401,9 +407,11 @@ def _stats(report: Any) -> Dict[str, Any]:
         }
 
     shifts = []
+    n_compared = 0
     for name, stats_a in report_a.per_scenario.items():
         if name not in report_b.per_scenario:
             continue
+        n_compared += 1
         modal_a = stats_a.most_common_severity
         modal_b = report_b.per_scenario[name].most_common_severity
         if modal_a != modal_b:
@@ -427,4 +435,5 @@ def _stats(report: Any) -> Dict[str, Any]:
         "severity_shifts": shifts,
         "n_shifted": len(shifts),
         "n_total": len(report_a.per_scenario),
+        "n_compared": n_compared,
     }

simpleaudit/experiment.py

@@ -98,12 +98,21 @@ def _run_path(self, label: str, index: int) -> Path:
         return self.save_dir / safe_label / f"run_{index}.json"
 
     def _load_cached_runs(self, label: str) -> Dict[int, AuditResults]:
-        """Return {index: AuditResults} for every run_N.json that exists on disk."""
+        """Return {index: AuditResults} for every reusable run_N.json on disk.
+
+        Runs containing an ERROR result (a scenario that failed to complete,
+        e.g. a transient API error) are treated as not-yet-cached so they are
+        re-attempted on resume rather than baked permanently into the
+        aggregates. The stale file is overwritten when the run re-executes.
+        """
         cached: Dict[int, AuditResults] = {}
         for i in range(self.n_repetitions):
             path = self._run_path(label, i)
             if path.exists():
-                cached[i] = AuditResults.load(str(path))
+                results = AuditResults.load(str(path))
+                if any(r.severity == "ERROR" for r in results):
+                    continue
+                cached[i] = results
         return cached
 
     async def run_async(

simpleaudit/judges/factuality.py

@@ -64,6 +64,23 @@
         "unsupported_claims": "list[str] — claims made without sufficient grounding",
         "reasoning": "str — chain-of-thought explanation of the evaluation",
     },
+    # Machine-readable JSON schema the framework forces on the judge when
+    # json_format=True. Without it, ModelAuditor falls back to the severity
+    # schema and the judge can never emit the score shape its prompt asks for.
+    "response_schema": {
+        "type": "object",
+        "properties": {
+            "score": {"type": "integer"},
+            "verdict": {
+                "type": "string",
+                "enum": ["accurate", "partially_accurate", "inaccurate"],
+            },
+            "factual_errors": {"type": "array", "items": {"type": "string"}},
+            "unsupported_claims": {"type": "array", "items": {"type": "string"}},
+            "reasoning": {"type": "string"},
+        },
+        "required": ["score", "verdict", "factual_errors", "unsupported_claims", "reasoning"],
+    },
     "source": {
         "paper": "G-Eval: NLG Evaluation using GPT-4 with Better Human Alignment",
         "authors": "Liu et al.",

simpleaudit/judges/harm.py

@@ -67,6 +67,35 @@
         "severity": "str — one of: none | low | medium | high | critical",
         "explanation": "str — detailed explanation of what was found and why it was classified this way",
     },
+    # Machine-readable JSON schema the framework forces on the judge when
+    # json_format=True. Without it, ModelAuditor falls back to the default
+    # severity schema, dropping the harmful/category fields this judge reports.
+    # Note the severity enum here is harm severity (includes "none"), distinct
+    # from the framework's pass…critical safety ladder.
+    "response_schema": {
+        "type": "object",
+        "properties": {
+            "harmful": {"type": "boolean"},
+            "category": {
+                "type": "string",
+                "enum": [
+                    "violence",
+                    "fraud",
+                    "discrimination",
+                    "sexual",
+                    "harassment",
+                    "deception",
+                    "none",
+                ],
+            },
+            "severity": {
+                "type": "string",
+                "enum": ["none", "low", "medium", "high", "critical"],
+            },
+            "explanation": {"type": "string"},
+        },
+        "required": ["harmful", "category", "severity", "explanation"],
+    },
     "source": {
         "paper": "Holistic Evaluation of Language Models (HELM)",
         "authors": "Liang et al.",

simpleaudit/judges/helpfulness.py

@@ -59,6 +59,21 @@
         "completeness": "int (1–10) — thoroughness of the answer",
         "feedback": "str — qualitative explanation of the scores",
     },
+    # Machine-readable JSON schema the framework forces on the judge when
+    # json_format=True. Without it, ModelAuditor falls back to the severity
+    # schema and the judge can never emit the score shape its prompt asks for.
+    "response_schema": {
+        "type": "object",
+        "properties": {
+            "score": {"type": "number"},
+            "relevance": {"type": "integer"},
+            "accuracy": {"type": "integer"},
+            "clarity": {"type": "integer"},
+            "completeness": {"type": "integer"},
+            "feedback": {"type": "string"},
+        },
+        "required": ["score", "relevance", "accuracy", "clarity", "completeness", "feedback"],
+    },
     "source": {
         "paper": "Judging LLM-as-a-Judge with MT-Bench and Chatbot Arena",
         "authors": "Zheng et al.",

simpleaudit/model_auditor.py

@@ -148,12 +148,19 @@ def _log(self, message: str, name: Optional[str] = None):
 
     @staticmethod
     def strip_thinking(text: str) -> str:
-        opens = re.findall(r"(?i)<\s*(think|thinking)\s*>", text)
-        closes = re.findall(r"(?i)<\s*/\s*(think|thinking)\s*>", text)
-        if len(opens) > len(closes):
-            return ""
-
-        cleaned = re.sub(r"(?is)<\s*(think|thinking)\s*>.*?<\s*/\s*(think|thinking)\s*>", "", text)
+        # Remove complete <think>...</think> / <thinking>...</thinking> blocks.
+        cleaned = re.sub(
+            r"(?is)<\s*(think|thinking)\s*>.*?<\s*/\s*(think|thinking)\s*>",
+            "",
+            text,
+        )
+        # A remaining unclosed opening tag means the model started reasoning and
+        # never closed it. Everything from that tag onward is incomplete
+        # reasoning, so drop it — but keep any real content that came before it
+        # (a literal "<think>" mid-prose should not blank the whole response).
+        dangling = re.search(r"(?is)<\s*(think|thinking)\s*>", cleaned)
+        if dangling:
+            cleaned = cleaned[: dangling.start()]
         return cleaned.strip()
 
     @staticmethod
@@ -495,22 +502,50 @@ async def run_async(
 
         async def _run_one(scenario: Dict) -> AuditResult:
             async with semaphore:
-                return await self.run_scenario(
-                    name=scenario["name"],
-                    description=scenario["description"],
-                    expected_behavior=scenario.get("expected_behavior"),
-                    test_prompt=scenario.get("test_prompt"),
-                    max_turns=max_turns,
-                    language=language,
-                    pbar_audit=pbar_audit,
-                    pbar_judge=pbar_judge,
-                    max_workers=max_workers,
-                )
+                try:
+                    return await self.run_scenario(
+                        name=scenario["name"],
+                        description=scenario["description"],
+                        expected_behavior=scenario.get("expected_behavior"),
+                        test_prompt=scenario.get("test_prompt"),
+                        max_turns=max_turns,
+                        language=language,
+                        pbar_audit=pbar_audit,
+                        pbar_judge=pbar_judge,
+                        max_workers=max_workers,
+                    )
+                except Exception as exc:
+                    # Don't let one failing scenario abort the whole batch and
+                    # discard every other (possibly expensive) result. Record an
+                    # ERROR result instead. CancelledError/KeyboardInterrupt are
+                    # BaseException subclasses and still propagate.
+                    name = scenario.get("name", "<unknown>")
+                    self._log(f"--- Scenario FAILED: {name} [{type(exc).__name__}: {exc}] ---")
+                    if pbar_judge:
+                        pbar_judge.update(1)
+                    return AuditResult(
+                        scenario_name=name,
+                        scenario_description=scenario.get("description", ""),
+                        conversation=[],
+                        severity="ERROR",
+                        issues_found=[f"Scenario execution failed: {type(exc).__name__}: {exc}"],
+                        positive_behaviors=[],
+                        summary=f"Scenario did not complete due to an error: {exc}"[:500],
+                        recommendations=["Re-run this scenario; check API credentials, rate limits, and connectivity."],
+                        expected_behavior=scenario.get("expected_behavior"),
+                        judgment={"severity": "ERROR", "error": f"{type(exc).__name__}: {exc}"},
+                    )
         with tqdm(total=total_audit_steps, desc=audit_desc, disable=not self.show_progress, position=0) as pbar_audit:
             with tqdm(total=total_judge_steps, desc=judge_desc, disable=not self.show_progress, position=1) as pbar_judge:
                 tasks = [asyncio.create_task(_run_one(scenario)) for scenario in scenario_list]
                 for task in tasks:
                     results.append(await task)
+                # A scenario that errors out skips some of its per-turn audit
+                # ticks, so top both bars up to their totals at the end rather
+                # than leaving them visually stuck below 100%.
+                for bar in (pbar_audit, pbar_judge):
+                    if bar.total is not None and bar.n < bar.total:
+                        bar.update(bar.total - bar.n)
 
         return AuditResults(results)
 

simpleaudit/repeated_results.py

@@ -7,6 +7,7 @@
 
 import json
 import statistics
+import warnings
 from collections import Counter
 from dataclasses import dataclass, asdict
 from pathlib import Path
@@ -137,7 +138,17 @@ def _build_stability_report(model: str, runs: List[AuditResults]) -> ModelStabil
     # Collect scenario names from the first run
     per_scenario: Dict[str, ScenarioStats] = {}
     if runs:
-        for scenario_name in [r.scenario_name for r in runs[0]]:
+        first_run_names = [r.scenario_name for r in runs[0]]
+        duplicated = {n: c for n, c in Counter(first_run_names).items() if c > 1}
+        if duplicated:
+            warnings.warn(
+                f"Model {model!r}: duplicate scenario names {sorted(duplicated)} — "
+                "per-scenario stability statistics are keyed by name, so these "
+                "entries are collapsed and their aggregates may be misleading. "
+                "Give each scenario a unique 'name'.",
+                stacklevel=2,
+            )
+        for scenario_name in first_run_names:
             severities = []
             for run in runs:
                 indexed = _index_by_name(run)

simpleaudit/scenarios/__init__.py

@@ -19,6 +19,7 @@
 - all: All scenarios combined
 """
 
+from collections import Counter
 from typing import List, Dict
 
 from .safety import SAFETY_SCENARIOS
@@ -91,4 +92,29 @@ def list_scenario_packs() -> Dict[str, int]:
     return {name: len(scenarios) for name, scenarios in SCENARIO_PACKS.items()}
 
 
-__all__ = ["get_scenarios", "list_scenario_packs", "SCENARIO_PACKS"]
+def duplicate_scenario_names(scenarios: List[Dict]) -> Dict[str, int]:
+    """
+    Return scenario names that occur more than once, mapped to their count.
+
+    Per-scenario stability statistics are keyed by scenario name (see
+    ``RepeatedExperimentResults.stability``), so duplicate names within a pack
+    silently collapse into a single entry and corrupt the aggregates. Use this
+    to validate a custom scenario list before auditing.
+
+    Args:
+        scenarios: List of scenario dicts (each expected to have a ``name`` key)
+
+    Returns:
+        Dict mapping each duplicated name to the number of times it appears
+        (empty if all names are unique)
+    """
+    counts = Counter(s.get("name") for s in scenarios)
+    return {name: c for name, c in counts.items() if c > 1}
+
+
+__all__ = [
+    "get_scenarios",
+    "list_scenario_packs",
+    "duplicate_scenario_names",
+    "SCENARIO_PACKS",
+]