docent-python 0.1.19a0__py3-none-any.whl → 0.1.21a0__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,84 @@
+{
+  "$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" }
+          },
+          "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,95 @@
+import json
+from typing import Any, cast
+
+import jsonschema
+
+from docent._llm_util.data_models.exceptions import ValidationFailedException
+from docent._llm_util.data_models.llm_output import LLMOutput
+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_llm_output(
+    llm_output: LLMOutput,
+    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
+    """
+    if llm_output.first_text is None:
+        raise ValidationFailedException("LLM output has no text", failed_output=None)
+
+    try:
+        output = forgiving_json_loads(llm_output.first_text)
+    except json.JSONDecodeError as e:
+        raise ValidationFailedException(
+            f"Failed to parse JSON: {e}. Raw text: `{llm_output.first_text}`",
+            failed_output=llm_output.first_text,
+        )
+
+    if not isinstance(output, dict):
+        logger.error(f"Expected dict output, got {type(output)}")
+        logger.error(f"LLM output: {llm_output.first_text}")
+        raise ValidationFailedException(
+            f"Expected dict output, got {type(output)}. Raw text: {llm_output.first_text}",
+            failed_output=llm_output.first_text,
+        )
+
+    return _validate_rubric_output(cast(dict[str, Any], output), output_schema, agent_run)

docent/judges/util/voting.py

@@ -0,0 +1,84 @@
+from collections import Counter
+from typing import Any, cast
+
+
+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, bool, and int fields. We skip float and strings.
+
+    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 integer fields
+        elif field_type == "integer":
+            agreement_keys.append(key)
+        # Include enum fields (even strings)
+        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

docent/sdk/client.py

@@ -182,12 +182,15 @@ class Docent:
         self._handle_response_errors(response)
         return response.json()
 
-    def get_rubric_run_state(self, collection_id: str, rubric_id: str) -> dict[str, Any]:
+    def get_rubric_run_state(
+        self, collection_id: str, rubric_id: str, version: int | None = None
+    ) -> dict[str, Any]:
         """Get rubric run state for a given collection and rubric.
 
         Args:
             collection_id: ID of the Collection.
             rubric_id: The ID of the rubric to get run state for.
+            version: The version of the rubric to get run state for. If None, the latest version is used.
 
         Returns:
             dict: Dictionary containing rubric run state with results, job_id, and total_agent_runs.
@@ -196,7 +199,7 @@ class Docent:
             requests.exceptions.HTTPError: If the API request fails.
         """
         url = f"{self._server_url}/rubric/{collection_id}/{rubric_id}/rubric_run_state"
-        response = self._session.get(url)
+        response = self._session.get(url, params={"version": version})
         self._handle_response_errors(response)
         return response.json()
 

docent/trace.py

@@ -226,7 +226,7 @@ class DocentTracer:
         try:
 
             # Check for OTEL_SPAN_ATTRIBUTE_COUNT_LIMIT environment variable
-            default_attribute_limit = 1024
+            default_attribute_limit = 1024 * 16
             env_value = os.environ.get("OTEL_SPAN_ATTRIBUTE_COUNT_LIMIT", "0")
             env_limit = int(env_value) if env_value.isdigit() else 0
             attribute_limit = max(env_limit, default_attribute_limit)