docent-python 0.1.19a0__py3-none-any.whl → 0.1.21a0__py3-none-any.whl

docent/data_models/util.py

@@ -0,0 +1,170 @@
+from __future__ import annotations
+
+from typing import Dict, Iterable, List, TypeVar
+from uuid import uuid4
+
+from pydantic import BaseModel
+
+from docent.data_models.agent_run import AgentRun
+
+T = TypeVar("T", bound=BaseModel)
+
+
+def _deep_copy_model(model: T) -> T:
+    """Create a deep copy of a Pydantic v2 model.
+
+    Using `model_copy(deep=True)` ensures nested models are fully copied and
+    mutations do not affect the original instance.
+    """
+    return model.model_copy(deep=True)
+
+
+def clone_agent_run_with_random_ids(agent_run: AgentRun) -> AgentRun:
+    """Clone an `AgentRun`, randomizing all IDs and fixing internal references.
+
+    The following transformations are performed on the cloned instance:
+    - Assign a new `AgentRun.id`.
+    - Assign new `Transcript.id` values and update any references to them (none today).
+    - Assign new `TranscriptGroup.id` values.
+    - Update `Transcript.transcript_group_id` to the new group IDs where applicable.
+    - Update `TranscriptGroup.agent_run_id` to the new `AgentRun.id`.
+    - Update `TranscriptGroup.parent_transcript_group_id` to the new group IDs where applicable.
+
+    Notes:
+    - If a `parent_transcript_group_id` or `transcript_group_id` references a group id that
+      is not present in the cloned run, the reference is left unchanged (mirrors importer behavior).
+
+    Args:
+        agent_run: The source `AgentRun` to clone.
+
+    Returns:
+        A new, independent `AgentRun` instance with randomized identifiers and consistent references.
+    """
+    # Validate source integrity before cloning
+    #  - No duplicate transcript or group IDs
+    #  - All transcript.group references exist if set
+    #  - All group.parent references exist if set
+    #  - All group.agent_run_id match the source run id
+    src_transcript_ids = [str(t.id) for t in agent_run.transcripts]
+    if len(src_transcript_ids) != len(set(src_transcript_ids)):
+        raise ValueError("Duplicate transcript ids detected in source AgentRun")
+
+    src_group_ids = [str(g.id) for g in agent_run.transcript_groups]
+    if len(src_group_ids) != len(set(src_group_ids)):
+        raise ValueError("Duplicate transcript group ids detected in source AgentRun")
+
+    src_group_id_set = set(src_group_ids)
+    for t in agent_run.transcripts:
+        if t.transcript_group_id is not None and str(t.transcript_group_id) not in src_group_id_set:
+            raise ValueError(
+                f"Transcript {t.id} references missing transcript_group_id {t.transcript_group_id}"
+            )
+
+    for g in agent_run.transcript_groups:
+        if (
+            g.parent_transcript_group_id is not None
+            and str(g.parent_transcript_group_id) not in src_group_id_set
+        ):
+            raise ValueError(
+                f"TranscriptGroup {g.id} references missing parent_transcript_group_id {g.parent_transcript_group_id}"
+            )
+        if str(g.agent_run_id) != str(agent_run.id):
+            raise ValueError(
+                f"TranscriptGroup {g.id} has agent_run_id {g.agent_run_id} which does not match AgentRun.id {agent_run.id}"
+            )
+
+    # Deep copy first so we never mutate the caller's instance
+    new_run = _deep_copy_model(agent_run)
+
+    # 1) Randomize AgentRun ID
+    new_agent_run_id = str(uuid4())
+    old_to_new_transcript_id: Dict[str, str] = {}
+    old_to_new_group_id: Dict[str, str] = {}
+
+    # 2) Pre-compute new IDs for transcripts and transcript groups without mutating yet
+    for transcript in new_run.transcripts:
+        old_to_new_transcript_id[str(transcript.id)] = str(uuid4())
+
+    for group in new_run.transcript_groups:
+        old_to_new_group_id[str(group.id)] = str(uuid4())
+
+    # 3) Mutate transcript groups: set new id, set agent_run_id, remap parents
+    for group in new_run.transcript_groups:
+        old_group_id = str(group.id)
+
+        # Assign new group id
+        group.id = old_to_new_group_id.get(old_group_id, str(uuid4()))
+
+        # Ensure group points to the new agent run id
+        group.agent_run_id = new_agent_run_id
+
+        # Remap parent id; raise if unknown
+        if group.parent_transcript_group_id is not None:
+            old_parent_id = str(group.parent_transcript_group_id)
+            if old_parent_id not in old_to_new_group_id:
+                raise ValueError(
+                    f"TranscriptGroup {old_group_id} parent_transcript_group_id {old_parent_id} not found in this AgentRun"
+                )
+            group.parent_transcript_group_id = old_to_new_group_id[old_parent_id]
+
+    # 4) Mutate transcripts: set new id, remap transcript_group_id
+    for transcript in new_run.transcripts:
+        old_transcript_id = str(transcript.id)
+
+        # Assign new transcript id
+        transcript.id = old_to_new_transcript_id.get(old_transcript_id, str(uuid4()))
+
+        # Remap group reference; raise if unknown
+        if transcript.transcript_group_id is not None:
+            old_group_id_ref = str(transcript.transcript_group_id)
+            if old_group_id_ref not in old_to_new_group_id:
+                raise ValueError(
+                    f"Transcript {old_transcript_id} references transcript_group_id {old_group_id_ref} not found in this AgentRun"
+                )
+            transcript.transcript_group_id = old_to_new_group_id[old_group_id_ref]
+
+    # 5) Finally set the new run id
+    new_run.id = new_agent_run_id
+
+    # Post-validate integrity on the cloned run
+    new_group_ids = [str(g.id) for g in new_run.transcript_groups]
+    if len(new_group_ids) != len(set(new_group_ids)):
+        raise ValueError("Duplicate transcript group ids detected after cloning")
+    new_group_id_set = set(new_group_ids)
+
+    new_transcript_ids = [str(t.id) for t in new_run.transcripts]
+    if len(new_transcript_ids) != len(set(new_transcript_ids)):
+        raise ValueError("Duplicate transcript ids detected after cloning")
+
+    for t in new_run.transcripts:
+        if t.transcript_group_id is not None and str(t.transcript_group_id) not in new_group_id_set:
+            raise ValueError(
+                f"Transcript {t.id} references missing transcript_group_id {t.transcript_group_id} after cloning"
+            )
+
+    for g in new_run.transcript_groups:
+        if (
+            g.parent_transcript_group_id is not None
+            and str(g.parent_transcript_group_id) not in new_group_id_set
+        ):
+            raise ValueError(
+                f"TranscriptGroup {g.id} references missing parent_transcript_group_id {g.parent_transcript_group_id} after cloning"
+            )
+        if str(g.agent_run_id) != str(new_run.id):
+            raise ValueError(
+                f"TranscriptGroup {g.id} has agent_run_id {g.agent_run_id} which does not match cloned AgentRun.id {new_run.id}"
+            )
+
+    return new_run
+
+
+def clone_agent_runs_with_random_ids(agent_runs: Iterable[AgentRun]) -> List[AgentRun]:
+    """Clone a sequence of `AgentRun` objects with randomized IDs.
+
+    Args:
+        agent_runs: Iterable of `AgentRun` instances to clone.
+
+    Returns:
+        A list of cloned `AgentRun` instances with fresh IDs and consistent references.
+    """
+    return [clone_agent_run_with_random_ids(ar) for ar in agent_runs]

docent/judges/__init__.py

@@ -0,0 +1,21 @@
+from docent.judges.impl import BaseJudge, MajorityVotingJudge, MultiReflectionJudge
+from docent.judges.types import (
+    JudgeResult,
+    JudgeResultCompletionCallback,
+    JudgeResultWithCitations,
+    ResultType,
+    Rubric,
+)
+
+__all__ = [
+    # Judges
+    "MajorityVotingJudge",
+    "MultiReflectionJudge",
+    "BaseJudge",
+    # Types
+    "Rubric",
+    "JudgeResult",
+    "JudgeResultWithCitations",
+    "JudgeResultCompletionCallback",
+    "ResultType",
+]

docent/judges/impl.py

@@ -0,0 +1,222 @@
+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 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]
+
+        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,
+            },
+            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)