docent-python 0.1.20a0__py3-none-any.whl → 0.1.22a0__py3-none-any.whl

docent/judges/impl.py

@@ -0,0 +1,232 @@
+import json
+from abc import ABC, abstractmethod
+from typing import Any
+
+from docent._llm_util.data_models.llm_output import LLMOutput
+from docent._llm_util.data_models.simple_svc import BaseLLMService, SimpleLLMService
+from docent._log_util import get_logger
+from docent.data_models.agent_run import AgentRun
+from docent.judges.types import JudgeResult, ResultType, Rubric
+from docent.judges.util.parse_output import parse_and_validate_llm_output
+from docent.judges.util.voting import (
+    compute_output_distribution,
+    find_modal_result,
+    get_agreement_keys,
+)
+
+logger = get_logger(__name__)
+
+
+class BaseJudge(ABC):
+    def __init__(self, cfg: Rubric, llm_svc: BaseLLMService):
+        self.cfg = cfg
+        self.llm_svc = llm_svc
+
+    @abstractmethod
+    async def __call__(self, agent_run: AgentRun, *args: Any, **kwargs: Any) -> JudgeResult | None:
+        """Returns None if all rollouts failed to produce a valid output."""
+
+
+class MajorityVotingJudge(BaseJudge):
+    """Rolls out the judge multiple times, then uses majority voting to determine the final result."""
+
+    def __init__(
+        self,
+        cfg: Rubric,
+        n_rollouts_per_input: int,
+        llm_svc: BaseLLMService = SimpleLLMService(),
+    ):
+        super().__init__(cfg, llm_svc)
+        self.n_rollouts_per_input = n_rollouts_per_input
+
+    async def __call__(
+        self,
+        agent_run: AgentRun,
+        max_concurrency: int = 10,
+    ) -> JudgeResult | None:
+        async def _validation_callback(batch_index: int, llm_output: LLMOutput):
+            parse_and_validate_llm_output(llm_output, self.cfg.output_schema, agent_run)
+
+        prompt = [{"role": "user", "content": self.cfg.materialize_system_prompt(agent_run)}]
+        outputs = await self.llm_svc.get_completions(
+            inputs=[prompt for _ in range(self.n_rollouts_per_input)],
+            model_options=[self.cfg.judge_model],
+            max_new_tokens=16384,
+            timeout=180.0,
+            use_cache=False,
+            validation_callback=_validation_callback,
+            max_concurrency=max_concurrency,
+        )
+
+        # Process each rollout independently
+        indep_results: list[dict[str, Any]] = []
+        for output in outputs:
+            if validated_output := parse_and_validate_llm_output(
+                output, self.cfg.output_schema, agent_run
+            ):
+                indep_results.append(validated_output)
+
+        if not indep_results:
+            return None
+
+        # Get a list of the keys that we want to measure agreement on
+        agreement_keys = get_agreement_keys(self.cfg.output_schema)
+
+        # Find the result that best matches modal values
+        final_max_idx, final_agt_key_modes_and_counts = find_modal_result(
+            indep_results, agreement_keys
+        )
+        final_output = indep_results[final_max_idx]
+
+        # Compute the distribution of the output across the agreement keys
+        final_output_distribution = compute_output_distribution(
+            indep_results, self.cfg.output_schema, agreement_keys
+        )
+
+        return JudgeResult(
+            agent_run_id=agent_run.id,
+            rubric_id=self.cfg.id,
+            rubric_version=self.cfg.version,
+            output=final_output,
+            result_metadata={
+                "agt_keys": agreement_keys,
+                # Final measurements
+                "final_results": indep_results,
+                "final_agt_key_modes_and_counts": final_agt_key_modes_and_counts,
+                "final_max_idx": final_max_idx,
+                "final_output_distribution": final_output_distribution,
+            },
+            result_type=ResultType.DIRECT_RESULT,
+        )
+
+
+class MultiReflectionJudge(BaseJudge):
+    """Rolls out the judge multiple times, then uses reflection to determine the final result."""
+
+    def __init__(
+        self,
+        cfg: Rubric,
+        n_rollouts_per_input: int,
+        llm_svc: BaseLLMService = SimpleLLMService(),
+    ):
+        super().__init__(cfg, llm_svc)
+        self.n_rollouts_per_input = n_rollouts_per_input
+
+    async def __call__(
+        self,
+        agent_run: AgentRun,
+        max_concurrency: int = 10,
+    ) -> JudgeResult | None:
+        rubric = self.cfg
+
+        async def _validation_callback(batch_index: int, llm_output: LLMOutput):
+            parse_and_validate_llm_output(llm_output, rubric.output_schema, agent_run)
+
+        # Run several independent rollouts
+        prompt = [{"role": "user", "content": self.cfg.materialize_system_prompt(agent_run)}]
+        outputs = await self.llm_svc.get_completions(
+            inputs=[prompt for _ in range(self.n_rollouts_per_input)],
+            model_options=[rubric.judge_model],
+            max_new_tokens=16384,
+            timeout=180.0,
+            use_cache=False,
+            validation_callback=_validation_callback,
+            max_concurrency=max_concurrency,
+        )
+
+        # Process each rollout
+        indep_results: list[dict[str, Any]] = []
+        for output in outputs:
+            if output.first_text is None:
+                continue
+            if v_output := parse_and_validate_llm_output(output, rubric.output_schema, agent_run):
+                indep_results.append(v_output)
+
+        if not indep_results:
+            return None
+
+        # Compute initial modes
+        agreement_keys = get_agreement_keys(rubric.output_schema)
+        indep_max_idx, indep_agt_key_modes_and_counts = find_modal_result(
+            indep_results, agreement_keys
+        )
+
+        def _get_reflection_prompt(cur_index: int):
+            # Current result
+            result = indep_results[cur_index]
+            # Get other results (excluding the current one)
+            other_results = [r for j, r in enumerate(indep_results) if j != cur_index]
+
+            # Create the reflection message
+            other_results_text = "\n\n".join(
+                [f"Answer {j+1}:\n{json.dumps(r, indent=2)}" for j, r in enumerate(other_results)]
+            )
+
+            reflection_instruction = (
+                f"Here are {len(other_results)} other independent answers to the same rubric evaluation:\n\n"
+                f"{other_results_text}\n\n"
+                f"Please reflect on these other answers and your own answer. "
+                f"Consider if any of them have identified important aspects you missed, or if there are disagreements that should be resolved. "
+                f"Then provide your final answer in the same JSON format as before."
+            )
+
+            # Construct the multi-message prompt
+            # 1. Original user message
+            # 2. Assistant message with the rollout's result
+            # 3. New user message asking for reflection
+            return [
+                *prompt,  # Original user message(s)
+                {"role": "assistant", "content": json.dumps(result, indent=2)},
+                {"role": "user", "content": reflection_instruction},
+            ]
+
+        final_results = indep_results.copy()  # Shallow copy
+        if len(indep_results) > 1:
+            # Ask the judge to reflect on the others' results
+            reflection_outputs = await self.llm_svc.get_completions(
+                inputs=[_get_reflection_prompt(i) for i in range(len(indep_results))],
+                model_options=[rubric.judge_model],
+                max_new_tokens=16384,
+                timeout=180.0,
+                use_cache=False,
+                validation_callback=_validation_callback,
+                max_concurrency=max_concurrency,
+            )
+
+            # Process reflection outputs in the same way as the initial rollouts
+            reflected_results: list[dict[str, Any]] = []
+            for output in reflection_outputs:
+                if output.first_text is None:
+                    continue
+                if v_output := parse_and_validate_llm_output(
+                    output, rubric.output_schema, agent_run
+                ):
+                    reflected_results.append(v_output)
+
+            # Use reflected results if we got any, otherwise fall back to original results
+            if reflected_results:
+                final_results = reflected_results
+            else:
+                logger.warning("No reflected results found, falling back to original results")
+
+        final_max_idx, final_agt_key_modes_and_counts = find_modal_result(
+            final_results, agreement_keys
+        )
+        return JudgeResult(
+            agent_run_id=agent_run.id,
+            rubric_id=rubric.id,
+            rubric_version=rubric.version,
+            output=final_results[final_max_idx],
+            result_metadata={
+                "agt_keys": agreement_keys,
+                # Final measurements
+                "final_results": final_results,
+                "final_agt_key_modes_and_counts": final_agt_key_modes_and_counts,
+                "final_max_idx": final_max_idx,
+                # Also include initial measurements
+                "indep_results": indep_results,
+                "indep_max_idx": indep_max_idx,
+                "indep_agt_key_modes_and_counts": indep_agt_key_modes_and_counts,
+            },
+            result_type=ResultType.DIRECT_RESULT,
+        )

docent/judges/types.py

@@ -0,0 +1,240 @@
+import enum
+import json
+from string import Formatter
+from typing import Any, Callable, Protocol
+from uuid import uuid4
+
+from pydantic import BaseModel, Field, field_serializer, field_validator
+
+from docent._llm_util.providers.preference_types import PUBLIC_PROVIDER_PREFERENCES, ModelOption
+from docent._log_util import get_logger
+from docent.data_models.agent_run import AgentRun
+from docent.data_models.citation import parse_citations
+from docent.data_models.transcript import TEXT_RANGE_CITE_INSTRUCTION
+from docent.judges.util.meta_schema import validate_judge_result_schema
+
+logger = get_logger(__name__)
+
+DEFAULT_JUDGE_SYSTEM_PROMPT_TEMPLATE = """
+Here is a rubric that we are using to judge transcripts of AI agent runs.
+
+Rubric:
+{rubric}
+
+Agent run:
+{agent_run}
+
+Your response should convey your judgment of the agent run according to the criteria given in the rubric provided above. Your entire response must be a valid JSON string which can be parsed with python `json.loads` without any additional processing. Double quotes (`"`) in the middle of a string in the JSON object must be escaped with a backslash.
+
+The JSON object you produce must adhere to the following schema:
+{output_schema}
+
+{citation_instructions}
+""".strip()
+
+DEFAULT_CITATION_INSTRUCTIONS = f"""
+For strings which require citations (according to the `citations: True` property), you must also follow these instructions:
+{TEXT_RANGE_CITE_INSTRUCTION}
+""".strip()
+
+DEFAULT_JUDGE_OUTPUT_SCHEMA = {
+    "type": "object",
+    "properties": {
+        "label": {"type": "string", "enum": ["match", "no match"]},
+        "explanation": {"type": "string", "citations": True},
+    },
+    # Require these properties to be present
+    "required": ["label", "explanation"],
+    # Allow additional properties though, as their presence is not breaking
+}
+
+DEFAULT_JUDGE_MODEL = PUBLIC_PROVIDER_PREFERENCES.default_judge_models[0]
+
+
+class Rubric(BaseModel):
+    """TODO(mengk): this should really be called JudgeConfig,
+    but temporarily keeping this for consistency with docent_core."""
+
+    class Config:
+        frozen = True
+
+    # Primary key
+    id: str = Field(default_factory=lambda: str(uuid4()))
+    version: int = 1
+
+    # What the judge actually does
+    rubric_text: str
+
+    # Default instructions for the judge
+    system_prompt_template: str = DEFAULT_JUDGE_SYSTEM_PROMPT_TEMPLATE
+    citation_instructions: str = DEFAULT_CITATION_INSTRUCTIONS
+    output_schema: dict[str, Any] = DEFAULT_JUDGE_OUTPUT_SCHEMA
+
+    # How to run the judge
+    judge_model: ModelOption = DEFAULT_JUDGE_MODEL
+
+    def materialize_system_prompt(self, agent_run: AgentRun) -> str:
+        """Construct the full prompt text for rubric evaluation.
+
+        This is the canonical implementation of prompt construction - use this function
+        anywhere you need to construct a rubric evaluation prompt (including cost estimation).
+        """
+
+        output_schema_text = json.dumps(self.output_schema, indent=2)
+
+        # We've already validated that the system prompt template has these keys
+        prompt = self.system_prompt_template.format(
+            rubric=self.rubric_text,
+            agent_run=agent_run.to_text_new(),
+            output_schema=output_schema_text,
+            # Only include citation instructions if the schema requests citations
+            citation_instructions=(
+                self.citation_instructions if _schema_requests_citations(self.output_schema) else ""
+            ),
+        ).strip()
+
+        return prompt
+
+    @field_validator("system_prompt_template")
+    @classmethod
+    def validate_system_prompt_template(cls, system_prompt_template: str):
+        # Extract all field names from the template
+        formatter = Formatter()
+        field_names = {
+            field_name
+            for _, field_name, _, _ in formatter.parse(system_prompt_template)
+            if field_name is not None
+        }
+
+        # Check for required fields
+        required_fields = {"agent_run", "output_schema", "rubric", "citation_instructions"}
+        missing_fields = required_fields - field_names
+
+        if missing_fields:
+            raise ValueError(
+                f"system_prompt_template must contain the following placeholders: {missing_fields}"
+            )
+
+        return system_prompt_template
+
+    @field_validator("output_schema")
+    @classmethod
+    def validate_output_schema(cls, output_schema: dict[str, Any]):
+        """
+        Raises:
+            jsonschema.ValidationError: If the schema is invalid
+            jsonschema.SchemaError: If the schema is not a valid 2020-12 schema
+        """
+        validate_judge_result_schema(output_schema)
+        return output_schema
+
+
+class ResultType(enum.Enum):
+    """Enum for the type of result that a judge result can have."""
+
+    DIRECT_RESULT = "direct_result"
+    NEAR_MISS = "near_miss"
+
+
+class JudgeResult(BaseModel):
+    class Config:
+        frozen = True
+
+    id: str = Field(default_factory=lambda: str(uuid4()))
+    agent_run_id: str
+    rubric_id: str
+    rubric_version: int
+
+    # Outputs
+    output: dict[str, Any]
+    result_metadata: dict[str, Any] | None = None
+    result_type: ResultType
+
+    # Deprecated
+    value: str | None = None
+
+    @field_serializer("result_type")
+    def serialize_result_type(self, result_type: ResultType) -> str:
+        return result_type.value
+
+
+class JudgeResultWithCitations(JudgeResult):
+    @classmethod
+    def from_judge_result(
+        cls, result: JudgeResult, schema: dict[str, Any]
+    ) -> "JudgeResultWithCitations":
+        """Judge result must be validated against the schema before calling this function!"""
+
+        def _parse_citation_string(output: str) -> dict[str, Any]:
+            text, citations = parse_citations(output)
+            return {"text": text, "citations": citations}
+
+        data = result.model_dump()
+        try:
+            data["output"] = traverse_schema_and_transform(
+                data["output"], schema, _parse_citation_string
+            )
+        except Exception as e:
+            logger.error(f"Failed to parse citations: {e}")
+            logger.error(f"Output: {data['output']}")
+            data["output"] = {"raw": data["output"]}
+        return cls(**data)
+
+
+class JudgeResultCompletionCallback(Protocol):
+    """Called when some batch of judge results is completed.
+    Supports batched calls for cases where many results are pre-computed.
+    This avoids invoking the callback separately for each datapoint.
+    """
+
+    async def __call__(
+        self,
+        batch_index: int,
+        judge_results: list[JudgeResult] | None,
+    ) -> None: ...
+
+
+def traverse_schema_and_transform(
+    output: Any,
+    schema: dict[str, Any],
+    citation_string_handler: Callable[[str], Any],
+) -> Any:
+    """Recursively traverse output based on schema, applying citation_string_handler to citation strings."""
+    if schema.get("type") == "string" and schema.get("citations"):  # type: ignore
+        return citation_string_handler(output)
+    elif schema.get("type") == "object":
+        properties: dict[str, Any] = schema.get("properties", {})
+        result: dict[str, Any] = {}
+        for key in properties:
+            if key in output:
+                result[key] = traverse_schema_and_transform(
+                    output[key], properties[key], citation_string_handler
+                )
+        return result
+    elif schema.get("type") == "array":
+        item_schema: dict[str, Any] = schema.get("items", {})
+        return [
+            traverse_schema_and_transform(item, item_schema, citation_string_handler)
+            for item in output
+        ]
+    else:
+        return output
+
+
+def _schema_requests_citations(schema: dict[str, Any]) -> bool:
+    """Check if any field in the schema requests citations by having 'citations': 'true'."""
+
+    def _check_field(field_schema: Any) -> bool:
+        if isinstance(field_schema, dict):
+            if field_schema.get("citations"):  # type: ignore
+                return True
+            for value in field_schema.values():  # type: ignore
+                if isinstance(value, dict) and _check_field(value):
+                    return True
+                elif isinstance(value, list):
+                    for item in value:  # type: ignore
+                        if isinstance(item, dict) and _check_field(item):
+                            return True
+        return False
+
+    return _check_field(schema)

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