docent-python 0.1.19a0__py3-none-any.whl → 0.1.27a0__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,87 @@
+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.data_models.remove_invalid_citation_ranges import remove_invalid_citation_ranges
+from docent.judges.types import traverse_schema_and_transform
+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 and filter citation text ranges in rubric results.
+    Also check that the output conforms to the output schema.
+
+    Args:
+        output: Raw results from LLM judge
+        agent_run: Agent run containing transcript data for validation
+
+    Returns:
+        Validated result dict with invalid citations removed
+
+    Raises:
+        ValidationFailedException: If validation fails
+    """
+
+    def _validate_citation_string(text: str) -> str:
+        validated_text = remove_invalid_citation_ranges(text, agent_run)
+        if validated_text != text:
+            logger.warning(
+                f"Citation validation removed invalid text range from citation in judge result. "
+                f"Agent run ID: {agent_run.id}, "
+                f"Original text: {text}, "
+                f"Validated text: {validated_text}, "
+            )
+        return validated_text
+
+    try:
+        jsonschema.validate(output, output_schema)
+    except jsonschema.ValidationError as e:
+        raise ValidationFailedException(f"Schema validation failed: {e}", failed_output=str(output))
+
+    try:
+        return traverse_schema_and_transform(output, output_schema, _validate_citation_string)
+    except Exception as e:
+        raise ValidationFailedException(
+            f"Citation validation failed: {e}", failed_output=str(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