docent-python 0.1.17a0__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

docent/sdk/agent_run_writer.py

@@ -4,11 +4,12 @@ import queue
 import signal
 import threading
 import time
-from typing import Any, Callable, Coroutine, Optional
+from typing import Any, AsyncGenerator, Callable, Coroutine, Optional
 
 import anyio
 import backoff
 import httpx
+import orjson
 from backoff.types import Details
 
 from docent._log_util.logger import get_logger
@@ -19,11 +20,16 @@ logger = get_logger(__name__)
 
 
 def _giveup(exc: BaseException) -> bool:
-    """Give up on client errors."""
+    """Give up on timeouts and client errors (4xx except 429). Retry others."""
+
+    # Give up immediately on any timeout (connect/read/write/pool)
+    if isinstance(exc, httpx.TimeoutException):
+        return True
 
     if isinstance(exc, httpx.HTTPStatusError):
         status = exc.response.status_code
         return status < 500 and status != 429
+
     return False
 
 
@@ -33,6 +39,15 @@ def _print_backoff_message(e: Details):
     )
 
 
+async def _generate_payload_chunks(runs: list[AgentRun]) -> AsyncGenerator[bytes, None]:
+    yield b'{"agent_runs": ['
+    for i, ar in enumerate(runs):
+        if i > 0:
+            yield b","
+        yield orjson.dumps(ar.model_dump(mode="json"))
+    yield b"]}"
+
+
 class AgentRunWriter:
     """Background thread for logging agent runs.
 
@@ -92,7 +107,6 @@ class AgentRunWriter:
         self._thread = threading.Thread(
             target=lambda: anyio.run(self._async_main),
             name="AgentRunWriterThread",
-            daemon=True,
         )
         self._thread.start()
         logger.info("AgentRunWriter thread started")
@@ -171,7 +185,7 @@ class AgentRunWriter:
             logger.info("Cancelling pending tasks...")
             self._cancel_event.set()
             n_pending = self._queue.qsize()
-            logger.info(f"Cancelled ~{n_pending} pending tasks")
+            logger.info(f"Cancelled ~{n_pending} pending runs")
 
             # Give a brief moment to exit
             logger.info("Waiting for thread to exit...")
@@ -179,7 +193,7 @@ class AgentRunWriter:
 
     def get_post_batch_fcn(
         self, client: httpx.AsyncClient
-    ) -> Callable[[list[AgentRun], anyio.CapacityLimiter], Coroutine[Any, Any, None]]:
+    ) -> Callable[[list[AgentRun]], Coroutine[Any, Any, None]]:
         """Return a function that will post a batch of agent runs to the API."""
 
         @backoff.on_exception(
@@ -189,34 +203,37 @@ class AgentRunWriter:
             max_tries=self._max_retries,
             on_backoff=_print_backoff_message,
         )
-        async def _post_batch(batch: list[AgentRun], limiter: anyio.CapacityLimiter) -> None:
-            async with limiter:
-                payload = {"agent_runs": [ar.model_dump(mode="json") for ar in batch]}
-                resp = await client.post(
-                    self._endpoint, json=payload, timeout=self._request_timeout
-                )
-                resp.raise_for_status()
+        async def _post_batch(batch: list[AgentRun]) -> None:
+            resp = await client.post(
+                self._endpoint,
+                content=_generate_payload_chunks(batch),
+                timeout=self._request_timeout,
+            )
+            resp.raise_for_status()
 
         return _post_batch
 
     async def _async_main(self) -> None:
         """Main async function for the AgentRunWriter thread."""
 
-        limiter = anyio.CapacityLimiter(self._num_workers)
-
         async with httpx.AsyncClient(base_url=self._base_url, headers=self._headers) as client:
+            _post_batch = self.get_post_batch_fcn(client)
             async with anyio.create_task_group() as tg:
-                _post_batch = self.get_post_batch_fcn(client)
 
-                async def batch_loop() -> None:
+                async def worker():
                     while not self._cancel_event.is_set():
                         batch = await self._gather_next_batch_from_queue()
                         if not batch:
                             continue
+                        try:
+                            await _post_batch(batch)
+                        except Exception as e:
+                            logger.error(
+                                f"Failed to post batch of {len(batch)} agent runs: {e.__class__.__name__}: {e}"
+                            )
 
-                        tg.start_soon(_post_batch, batch, limiter)
-
-                tg.start_soon(batch_loop)
+                for _ in range(self._num_workers):
+                    tg.start_soon(worker)
 
     async def _gather_next_batch_from_queue(self) -> list[AgentRun]:
         """Gather a batch of agent runs from the queue.
@@ -241,6 +258,14 @@ def init(
     server_url: str = "https://api.docent.transluce.org",
     web_url: str = "https://docent.transluce.org",
     api_key: str | None = None,
+    # Writer arguments
+    num_workers: int = 4,
+    queue_maxsize: int = 20_000,
+    request_timeout: float = 30.0,
+    flush_interval: float = 1.0,
+    batch_size: int = 1_000,
+    max_retries: int = 5,
+    shutdown_timeout: int = 60,
 ):
     """Initialize the AgentRunWriter thread.
 
@@ -250,6 +275,16 @@ def init(
         server_url (str): URL of the Docent server.
         web_url (str): URL of the Docent web UI.
         api_key (str): API key for the Docent API.
+        num_workers (int): Max number of concurrent tasks to run,
+            managed by anyio.CapacityLimiter.
+        queue_maxsize (int): Maximum size of the queue.
+            If maxsize is <= 0, the queue size is infinite.
+        request_timeout (float): Timeout for the HTTP request.
+        flush_interval (float): Interval to flush the queue.
+        batch_size (int): Number of agent runs to batch together.
+        max_retries (int): Maximum number of retries for the HTTP request.
+        shutdown_timeout (int): Timeout to wait for the background thread to finish
+            after the main thread has requested shutdown.
     """
     api_key = api_key or os.getenv("DOCENT_API_KEY")
 
@@ -271,4 +306,12 @@ def init(
         api_key=api_key,
         collection_id=collection_id,
         server_url=server_url,
+        # Writer arguments
+        num_workers=num_workers,
+        queue_maxsize=queue_maxsize,
+        request_timeout=request_timeout,
+        flush_interval=flush_interval,
+        batch_size=batch_size,
+        max_retries=max_retries,
+        shutdown_timeout=shutdown_timeout,
     )