docent-python 0.1.41a0__py3-none-any.whl

docent/judges/util/forgiving_json.py

@@ -0,0 +1,108 @@
+import json
+from typing import Any
+
+
+def _repair_json(text: str) -> str:
+    """Strip leading/trailing text and fix unescaped quotes/newlines."""
+
+    json_start = None
+    for i, char in enumerate(text):
+        remaining = text[i:]
+        if (
+            char in '[{"'
+            or char.isdigit()
+            or char == "-"
+            or remaining.startswith("null")
+            or remaining.startswith("true")
+            or remaining.startswith("false")
+        ):
+            json_start = i
+            break
+    if json_start is None:
+        raise ValueError("No valid JSON start found")
+
+    result: list[str] = []
+    in_string = False
+    escape_next = False
+    depth = 0
+    started_with_container = text[json_start] in "[{"
+
+    for i in range(json_start, len(text)):
+        char = text[i]
+
+        if escape_next:
+            if in_string:
+                # Check if this is a valid escape sequence
+                is_valid_escape = char in '\\/bfnrt"' or (
+                    char == "u"
+                    and i + 4 < len(text)
+                    and all(c in "0123456789abcdefABCDEF" for c in text[i + 1 : i + 5])
+                )
+                if not is_valid_escape:
+                    # Invalid escape sequence - add another backslash to escape it
+                    result.append("\\")
+            result.append(char)
+            escape_next = False
+            continue
+
+        if char == "\\":
+            result.append(char)
+            escape_next = True
+            continue
+
+        if char == '"':
+            if in_string:
+                # Check if quote should be escaped by looking at what follows
+                remaining = text[i + 1 :].lstrip()
+                if remaining and remaining[0] not in ':,}]"':
+                    result.append('\\"')
+                    continue
+                in_string = False
+                result.append(char)
+                # If we're at depth 0 and closed a top-level string, we're done
+                if depth == 0 and not started_with_container:
+                    return "".join(result)
+            else:
+                in_string = True
+                result.append(char)
+        elif in_string and char == "\n":
+            result.append("\\n")
+        else:
+            result.append(char)
+
+        if not in_string:
+            if char in "[{":
+                depth += 1
+            elif char in "]}":
+                depth -= 1
+                if depth == 0:
+                    return "".join(result)
+            # For primitives at top level (depth 0), stop at whitespace if we've consumed content
+            elif depth == 0 and not started_with_container and result and char in " \t\n\r":
+                # Check if this is trailing whitespace after a complete primitive
+                current = "".join(result).strip()
+                if current:
+                    try:
+                        json.loads(current)
+                        return current
+                    except (json.JSONDecodeError, ValueError):
+                        pass
+
+    return "".join(result)
+
+
+def forgiving_json_loads(text: str) -> Any:
+    """
+    Parse JSON from text, applying heuristics to fix common LLM mistakes.
+
+    Repairs applied:
+    - Strip leading/trailing non-JSON text
+    - Escape unescaped quotes and newlines inside strings
+    - Fix invalid escape sequences inside strings
+    """
+    if not text or not text.strip():
+        raise ValueError("Empty or whitespace-only input")
+
+    text = _repair_json(text)
+
+    return json.loads(text)

docent/judges/util/meta_schema.json

@@ -0,0 +1,86 @@
+{
+  "$schema": "https://json-schema.org/draft/2020-12/schema",
+  "$id": "https://example.com/meta/mini-schema",
+  "title": "Meta-schema for Docent judge outputs. Makes some restrictions to 2020-12.",
+  "type": "object",
+  "additionalProperties": false,
+  "properties": {
+    "type": { "const": "object" },
+    "additionalProperties": { "const": false },
+    "required": {
+      "type": "array",
+      "items": { "type": "string" }
+    },
+
+    "properties": {
+      "type": "object",
+      "propertyNames": { "type": "string" },
+      "additionalProperties": {
+        "type": "object",
+        "additionalProperties": false,
+        "required": ["type"],
+
+        "properties": {
+          "type": {
+            "type": "string",
+            "enum": ["string", "integer", "number", "boolean"]
+          },
+          "description": {
+            "type": "string"
+          },
+          "citations": {
+            "type": "boolean"
+          },
+          "enum": {
+            "type": "array",
+            "items": {
+              "type": ["string", "integer", "boolean"]
+            }
+          },
+          "format": {
+            "type": "string",
+            "enum": [
+              "date-time",
+              "date",
+              "time",
+              "email",
+              "hostname",
+              "ipv4",
+              "ipv6",
+              "uri",
+              "uuid"
+            ]
+          },
+          "minLength": {
+            "type": "integer",
+            "minimum": 0
+          },
+          "maxLength": {
+            "type": "integer",
+            "minimum": 0
+          },
+          "pattern": {
+            "type": "string"
+          },
+          "minimum": {
+            "type": "number"
+          },
+          "maximum": {
+            "type": "number"
+          },
+          "exclusiveMinimum": {
+            "type": "number"
+          },
+          "exclusiveMaximum": {
+            "type": "number"
+          },
+          "multipleOf": {
+            "type": "number",
+            "exclusiveMinimum": 0
+          }
+        }
+      }
+    }
+  },
+  "required": ["type", "properties"]
+}

docent/judges/util/meta_schema.py

@@ -0,0 +1,29 @@
+import json
+from pathlib import Path
+from typing import Any
+
+import jsonschema
+
+
+def _load_meta_schema() -> dict[str, Any]:
+    """Load the rubric meta-schema from the adjacent JSON file."""
+    meta_schema_path = Path(__file__).with_suffix(".json")
+    with meta_schema_path.open("r", encoding="utf-8") as f:
+        return json.load(f)
+
+
+_META_VALIDATOR = jsonschema.Draft202012Validator(_load_meta_schema())
+
+
+def validate_judge_result_schema(schema: dict[str, Any]):
+    """Validate a proposed schema against the rubric meta-schema.
+
+    Raises:
+        jsonschema.ValidationError: If the schema is invalid
+        jsonschema.SchemaError: If the schema is not a valid 2020-12 schema
+    """
+    # First check that this is a valid 2020-12 schema
+    jsonschema.Draft202012Validator.check_schema(schema)
+
+    # Then check that it conforms to our subset of the 2020-12 schema
+    _META_VALIDATOR.validate(schema)  # type: ignore

docent/judges/util/parse_output.py

@@ -0,0 +1,68 @@
+from typing import Any, cast
+
+import jsonschema
+
+from docent._llm_util.data_models.exceptions import ValidationFailedException
+from docent._log_util import get_logger
+from docent.data_models.agent_run import AgentRun
+from docent.judges.util.forgiving_json import forgiving_json_loads
+
+logger = get_logger(__name__)
+
+
+def _validate_rubric_output(
+    output: dict[str, Any], output_schema: dict[str, Any], agent_run: AgentRun
+) -> dict[str, Any]:
+    """Validate that the output conforms to the output schema.
+
+    Args:
+        output: Raw results from LLM judge
+        agent_run: Agent run (unused, kept for backwards compatibility)
+        output_schema: Schema to validate against
+
+    Returns:
+        Validated result dict
+
+    Raises:
+        ValidationFailedException: If validation fails
+    """
+    try:
+        jsonschema.validate(output, output_schema)
+    except jsonschema.ValidationError as e:
+        raise ValidationFailedException(f"Schema validation failed: {e}", failed_output=str(output))
+
+    return output
+
+
+def parse_and_validate_output_str(
+    output_str: str, output_schema: dict[str, Any], agent_run: AgentRun
+) -> dict[str, Any]:
+    """Parse and validate LLM output for rubric evaluation.
+
+    Args:
+        llm_output: The LLM output to parse
+        output_schema: The schema to validate against
+        agent_run: Agent run for citation validation
+
+    Returns:
+        Validated output dict
+
+    Raises:
+        ValidationFailedException: If parsing or validation fails
+    """
+
+    try:
+        output = forgiving_json_loads(output_str)
+    except Exception as e:
+        raise ValidationFailedException(
+            f"Failed to parse JSON: {e}. Raw text: `{output_str}`",
+            failed_output=output_str,
+        )
+
+    if not isinstance(output, dict):
+        raise ValidationFailedException(
+            f"Expected dict output, got {type(output)}. Raw text: {output_str}",
+            failed_output=output_str,
+        )
+
+    return _validate_rubric_output(cast(dict[str, Any], output), output_schema, agent_run)

docent/judges/util/voting.py

@@ -0,0 +1,139 @@
+from collections import Counter
+from typing import Any, TypedDict, cast
+
+import numpy as np
+
+
+class EstimateWithCI(TypedDict):
+    mean: float
+    var: float
+    n: int
+    ci_95: float
+
+
+JudgeOutputDistribution = dict[str | bool | int | float, EstimateWithCI]
+
+
+def get_agreement_keys(schema: dict[str, Any]) -> list[str]:
+    """Get list of top-level keys in schema that we want to measure agreement on.
+
+    This includes enum and bool fields.
+
+    Args:
+        schema: JSON schema dict
+
+    Returns:
+        List of field names (keys) that should be used for measuring agreement
+    """
+    agreement_keys: list[str] = []
+
+    properties = schema.get("properties", {})
+    assert isinstance(properties, dict)
+    properties = cast(dict[str, Any], properties)
+
+    for key, field_schema in properties.items():
+        assert isinstance(field_schema, dict)
+        field_schema = cast(dict[str, Any], field_schema)
+
+        field_type = field_schema.get("type")
+        assert isinstance(field_type, str)
+
+        # Include boolean fields
+        if field_type == "boolean":
+            agreement_keys.append(key)
+        # Include enum fields (strings and numbers must be in this category)
+        elif "enum" in field_schema:
+            agreement_keys.append(key)
+
+    return agreement_keys
+
+
+def find_modal_result(indep_results: list[dict[str, Any]], agreement_keys: list[str]):
+    """Find the result that best matches modal values across agreement keys.
+
+    Args:
+        indep_results: List of independent results to analyze
+        agreement_keys: Keys to measure agreement on
+
+    Returns:
+        Tuple of (max_idx, agt_key_modes_and_counts) where:
+        - max_idx is the index of the result that best matches modal values
+        - agt_key_modes_and_counts maps each key to (modal_value, count) or None if no values exist for that key
+
+    Raises:
+        ValueError: If no results are provided
+    """
+    if not indep_results:
+        raise ValueError("No results to score")
+
+    # For each agreement key, compute the mode and count (or None, if no values exist for that key)
+    agt_key_modes_and_counts: dict[str, tuple[str | bool | int, int] | None] = {}
+    for key in agreement_keys:
+        key_modes = Counter(v for r in indep_results if (v := r.get(key)) is not None)
+        if most_common_one := key_modes.most_common(1):
+            agt_key_modes_and_counts[key] = most_common_one[0]
+        else:
+            agt_key_modes_and_counts[key] = None
+
+    # Score each rollout based on how many agreement keys they match
+    # If there is no mode for a key, or if a certain result doesn't have that key, it doesn't count.
+    # TODO(mengk): This may bias towards results that have more keys.
+    indep_result_scores: list[int] = []
+    for r in indep_results:
+        score = 0
+        for key in agreement_keys:
+            mode_and_count = agt_key_modes_and_counts[key]
+            if mode_and_count and r.get(key) == mode_and_count[0]:
+                score += 1
+        indep_result_scores.append(score)
+
+    # Argmax
+    max_idx = indep_result_scores.index(max(indep_result_scores))
+
+    return max_idx, agt_key_modes_and_counts
+
+
+def compute_output_distributions(
+    indep_results: list[dict[str, Any]], output_schema: dict[str, Any], agreement_keys: list[str]
+):
+    def _get_possible_values(key: str) -> list[str | bool | int | float]:
+        if "enum" in output_schema.get("properties", {}).get(key, {}):
+            return output_schema.get("properties", {}).get(key, {}).get("enum", [])
+        elif output_schema.get("properties", {}).get(key, {}).get("type") == "boolean":
+            return [True, False]
+        else:
+            return []
+
+    raw_counts: dict[str, dict[str | bool | int | float, int]] = {
+        key: {value: 0 for value in _get_possible_values(key)} for key in agreement_keys
+    }
+    # Collect counts for each possible value
+    for result in indep_results:
+        for key in agreement_keys:
+            if (value := result.get(key)) is not None:  # Could be none if the key is optional
+                assert (
+                    value in raw_counts[key]
+                ), "this should never happen; the value must be in possible values, since judge results have been validated against the schema"
+                raw_counts[key][value] += 1
+
+    distributions: dict[str, JudgeOutputDistribution] = {}
+    for agt_key in agreement_keys:
+        distributions[agt_key] = {}
+
+        # First normalize the counts to get probabilities
+        counts = raw_counts[agt_key]
+        total = sum(counts.values())
+        probs = {value: (count / total) if total > 0 else 0.0 for value, count in counts.items()}
+
+        for output_key, value in probs.items():
+            mean, estimate_var = value, (value * (1 - value))
+            # TODO(mengk): change to the wilson score interval
+            ci_95 = float(1.96 * np.sqrt(estimate_var / total)) if total > 0 else 0.0
+            distributions[agt_key][output_key] = {
+                "mean": mean,
+                "var": estimate_var,
+                "n": total,
+                "ci_95": ci_95,
+            }
+
+    return distributions

docent/loaders/load_inspect.py

@@ -0,0 +1,215 @@
+import json
+from pathlib import Path
+from typing import Any, BinaryIO, Generator, Tuple
+from zipfile import ZipFile
+
+from inspect_ai.log import EvalLog
+from inspect_ai.scorer import CORRECT, INCORRECT, NOANSWER, PARTIAL, Score
+
+from docent._log_util.logger import get_logger
+from docent.data_models import AgentRun, Transcript
+from docent.data_models.chat import parse_chat_message
+
+logger = get_logger(__name__)
+
+
+def _normalize_inspect_score(score: Score | dict[str, Any]) -> Any:
+    """
+    Normalize an inspect score to a float. Logic mirrors inspect_ai.scorer._metric.value_to_float.
+
+    Args:
+        score: The inspect score to normalize.
+
+    Returns:
+        The normalized score as a float, or None if the score is not a valid value.
+    """
+
+    def _leaf_normalize(value: Any) -> Any:
+        if value is None:
+            return None
+        if isinstance(value, int | float | bool):
+            return float(value)
+        if value == CORRECT:
+            return 1.0
+        if value == PARTIAL:
+            return 0.5
+        if value in [INCORRECT, NOANSWER]:
+            return 0
+        value = str(value).lower()
+        if value in ["yes", "true"]:
+            return 1.0
+        if value in ["no", "false"]:
+            return 0.0
+        if value.replace(".", "").isnumeric():
+            return float(value)
+        return value
+
+    if isinstance(score, dict):
+        value = score["value"]
+    else:
+        value = score.value
+
+    if isinstance(value, int | float | bool | str):
+        return _leaf_normalize(value)
+    if isinstance(value, list):
+        return [_leaf_normalize(v) for v in value]  # type: ignore
+    assert isinstance(value, dict), "Inspect score must be leaf value, list, or dict"
+    return {k: _leaf_normalize(v) for k, v in value.items()}  # type: ignore
+
+
+def load_inspect_log(log: EvalLog) -> list[AgentRun]:
+    if log.samples is None:
+        return []
+
+    # TODO(vincent): fix this
+    agent_runs: list[AgentRun] = []
+
+    for s in log.samples:
+        sample_id = s.id
+        epoch_id = s.epoch
+
+        if s.scores is None:
+            sample_scores = {}
+        else:
+            sample_scores = {k: _normalize_inspect_score(v) for k, v in s.scores.items()}
+
+        metadata = {
+            "task_id": log.eval.task,
+            "sample_id": str(sample_id),
+            "epoch_id": epoch_id,
+            "model": log.eval.model,
+            "additional_metadata": s.metadata,
+            "scores": sample_scores,
+            # Scores could have answers, explanations, and other metadata besides the values we extract
+            "scoring_metadata": s.scores,
+        }
+
+        agent_runs.append(
+            AgentRun(
+                transcripts=[
+                    Transcript(
+                        messages=[parse_chat_message(m.model_dump()) for m in s.messages],
+                        metadata={},
+                    )
+                ],
+                metadata=metadata,
+            )
+        )
+
+    return agent_runs
+
+
+def _read_sample_as_run(data: dict[str, Any], header_metadata: dict[str, Any] = {}) -> AgentRun:
+    if "scores" in data:
+        normalized_scores = {k: _normalize_inspect_score(v) for k, v in data["scores"].items()}
+    else:
+        normalized_scores = {}
+
+    if "metadata" in data:
+        sample_metadata = data["metadata"]
+    else:
+        sample_metadata = {}
+
+    run_metadata: dict[str, Any] = {
+        "sample_id": data.get("id"),
+        "epoch": data.get("epoch"),
+        "target": data.get("target"),
+        # Scores could have answers, explanations, and other metadata besides the values we extract
+        "scoring_metadata": data.get("scores"),
+        "scores": normalized_scores,
+        # If a key exists in header and sample, sample takes precedence
+        **header_metadata,
+        **sample_metadata,
+    }
+
+    run = AgentRun(
+        transcripts=[
+            Transcript(messages=[parse_chat_message(m) for m in data["messages"]], metadata={})
+        ],
+        metadata=run_metadata,
+    )
+    return run
+
+
+def _run_metadata_from_header(header: dict[str, Any]) -> dict[str, Any]:
+    """
+    Inspect logs often have a lot of metadata.
+    This function tries to get the most important stuff without adding clutter.
+    """
+    m: dict[str, Any] = {}
+    if e := header.get("eval"):
+        m["task"] = e["task"]
+        m["model"] = e["model"]
+    return m
+
+
+def get_total_samples(file_path: Path, format: str = "json") -> int:
+    """Return the total number of samples in the provided file."""
+    with open(file_path, "rb") as f:
+        if format == "json":
+            data = json.load(f)
+            return len(data.get("samples", []))
+        elif format == "eval":
+            z = ZipFile(f, mode="r")
+            try:
+                return sum(
+                    1
+                    for name in z.namelist()
+                    if name.startswith("samples/") and name.endswith(".json")
+                )
+            finally:
+                z.close()
+        else:
+            raise ValueError(f"Format must be 'json' or 'eval': {format}")
+
+
+def _runs_from_eval_file(
+    file: BinaryIO,
+) -> Tuple[dict[str, Any], Generator[AgentRun, None, None]]:
+    zip = ZipFile(file, mode="r")
+    try:
+        header: dict[str, Any] = json.load(zip.open("header.json", "r"))
+        header_metadata = _run_metadata_from_header(header)
+    except KeyError:
+        logger.warning(f"No header found in {file.name} file")
+        header_metadata = {}
+
+    def _iter_runs() -> Generator[AgentRun, None, None]:
+        try:
+            for sample_file in zip.namelist():
+                if not (sample_file.startswith("samples/") and sample_file.endswith(".json")):
+                    continue
+                with zip.open(sample_file, "r") as f:
+                    data = json.load(f)
+                run: AgentRun = _read_sample_as_run(data, header_metadata)
+                yield run
+        finally:
+            zip.close()
+
+    return header_metadata, _iter_runs()
+
+
+def _runs_from_json_file(
+    file: BinaryIO,
+) -> Tuple[dict[str, Any], Generator[AgentRun, None, None]]:
+    data = json.load(file)
+    header_metadata = _run_metadata_from_header(data)
+
+    def _iter_runs() -> Generator[AgentRun, None, None]:
+        for sample in data["samples"]:
+            run: AgentRun = _read_sample_as_run(sample, header_metadata)
+            yield run
+
+    return header_metadata, _iter_runs()
+
+
+def runs_from_file(
+    file: BinaryIO, format: str = "json"
+) -> Tuple[dict[str, Any], Generator[AgentRun, None, None]]:
+    if format == "json":
+        result = _runs_from_json_file(file)
+    elif format == "eval":
+        result = _runs_from_eval_file(file)
+    else:
+        raise ValueError(f"Format must be 'json' or 'eval': {format}")
+    return result

docent/samples/__init__.py

@@ -0,0 +1,3 @@
+from docent.samples.load import get_inspect_fpath, get_tau_bench_airline_fpath
+
+__all__ = ["get_inspect_fpath", "get_tau_bench_airline_fpath"]

docent/samples/load.py

@@ -0,0 +1,9 @@
+import os
+
+
+def get_inspect_fpath():
+    return os.path.join(os.path.dirname(__file__), "log.eval")
+
+
+def get_tau_bench_airline_fpath():
+    return os.path.join(os.path.dirname(__file__), "tb_airline.json")