docent-python 0.1.17a0__py3-none-any.whl → 0.1.27a0__py3-none-any.whl

docent/data_models/chat/__init__.py

@@ -7,7 +7,12 @@ from docent.data_models.chat.message import (
     UserMessage,
     parse_chat_message,
 )
-from docent.data_models.chat.tool import ToolCall, ToolCallContent, ToolInfo, ToolParams
+from docent.data_models.chat.tool import (
+    ToolCall,
+    ToolCallContent,
+    ToolInfo,
+    ToolParams,
+)
 
 __all__ = [
     "ChatMessage",

docent/data_models/citation.py

@@ -1,15 +1,27 @@
 import re
+from dataclasses import dataclass
 
 from pydantic import BaseModel
 
 
+@dataclass
+class ParsedCitation:
+    """Represents a parsed citation before conversion to full Citation object."""
+
+    transcript_idx: int | None
+    block_idx: int | None
+    metadata_key: str | None = None
+    start_pattern: str | None = None
+
+
 class Citation(BaseModel):
     start_idx: int
     end_idx: int
     agent_run_idx: int | None = None
     transcript_idx: int | None = None
-    block_idx: int
+    block_idx: int | None = None
     action_unit_idx: int | None = None
+    metadata_key: str | None = None
     start_pattern: str | None = None
 
 
@@ -17,6 +29,9 @@ RANGE_BEGIN = "<RANGE>"
 RANGE_END = "</RANGE>"
 
 _SINGLE_RE = re.compile(r"T(\d+)B(\d+)")
+_METADATA_RE = re.compile(r"^M\.([^:]+)$")  # [M.key]
+_TRANSCRIPT_METADATA_RE = re.compile(r"^T(\d+)M\.([^:]+)$")  # [T0M.key]
+_MESSAGE_METADATA_RE = re.compile(r"^T(\d+)B(\d+)M\.([^:]+)$")  # [T0B1M.key]
 _RANGE_CONTENT_RE = re.compile(r":\s*" + re.escape(RANGE_BEGIN) + r".*?" + re.escape(RANGE_END))
 
 
@@ -70,41 +85,93 @@ def scan_brackets(text: str) -> list[tuple[int, int, str]]:
     return matches
 
 
-def parse_single_citation(part: str) -> tuple[int, int, str | None] | None:
+def parse_single_citation(part: str) -> ParsedCitation | None:
     """
     Parse a single citation token inside a bracket and return its components.
 
-    Returns (transcript_idx, block_idx, start_pattern) or None if invalid.
+    Returns ParsedCitation or None if invalid.
+    For metadata citations, transcript_idx may be None (for agent run metadata).
+    Supports optional text range for all valid citation kinds.
     """
     token = part.strip()
     if not token:
         return None
 
+    # Extract optional range part
+    start_pattern: str | None = None
+    citation_part = token
     if ":" in token:
-        citation_part, range_part = token.split(":", 1)
-        single_match = _SINGLE_RE.match(citation_part.strip())
-        if not single_match:
+        left, right = token.split(":", 1)
+        citation_part = left.strip()
+        start_pattern = _extract_range_pattern(right)
+
+    # Try matches in order of specificity
+    # 1) Message metadata [T0B0M.key]
+    m = _MESSAGE_METADATA_RE.match(citation_part)
+    if m:
+        transcript_idx = int(m.group(1))
+        block_idx = int(m.group(2))
+        metadata_key = m.group(3)
+        # Disallow nested keys like status.code per instruction
+        if "." in metadata_key:
             return None
-        transcript_idx = int(single_match.group(1))
-        block_idx = int(single_match.group(2))
-        start_pattern = _extract_range_pattern(range_part)
-        return transcript_idx, block_idx, start_pattern
-    else:
-        single_match = _SINGLE_RE.match(token)
-        if not single_match:
+        return ParsedCitation(
+            transcript_idx=transcript_idx,
+            block_idx=block_idx,
+            metadata_key=metadata_key,
+            start_pattern=start_pattern,
+        )
+
+    # 2) Transcript metadata [T0M.key]
+    m = _TRANSCRIPT_METADATA_RE.match(citation_part)
+    if m:
+        transcript_idx = int(m.group(1))
+        metadata_key = m.group(2)
+        if "." in metadata_key:
             return None
-        transcript_idx = int(single_match.group(1))
-        block_idx = int(single_match.group(2))
-        return transcript_idx, block_idx, None
+        return ParsedCitation(
+            transcript_idx=transcript_idx,
+            block_idx=None,
+            metadata_key=metadata_key,
+            start_pattern=start_pattern,
+        )
+
+    # 3) Agent run metadata [M.key]
+    m = _METADATA_RE.match(citation_part)
+    if m:
+        metadata_key = m.group(1)
+        if "." in metadata_key:
+            return None
+        return ParsedCitation(
+            transcript_idx=None,
+            block_idx=None,
+            metadata_key=metadata_key,
+            start_pattern=start_pattern,
+        )
+
+    # 4) Regular transcript block [T0B0]
+    m = _SINGLE_RE.match(citation_part)
+    if m:
+        transcript_idx = int(m.group(1))
+        block_idx = int(m.group(2))
+        return ParsedCitation(
+            transcript_idx=transcript_idx, block_idx=block_idx, start_pattern=start_pattern
+        )
+
+    return None
 
 
 def parse_citations(text: str) -> tuple[str, list[Citation]]:
     """
-    Parse citations from text in the format described by BLOCK_RANGE_CITE_INSTRUCTION.
+    Parse citations from text in the format described by TEXT_RANGE_CITE_INSTRUCTION.
 
     Supported formats:
     - Single block: [T<key>B<idx>]
     - Text range with start pattern: [T<key>B<idx>:<RANGE>start_pattern</RANGE>]
+    - Agent run metadata: [M.key]
+    - Transcript metadata: [T<key>M.key]
+    - Message metadata: [T<key>B<idx>M.key]
+    - Message metadata with text range: [T<key>B<idx>M.key:<RANGE>start_pattern</RANGE>]
 
     Args:
         text: The text to parse citations from
@@ -127,8 +194,21 @@ def parse_citations(text: str) -> tuple[str, list[Citation]]:
         # Parse a single citation token inside the bracket
         parsed = parse_single_citation(bracket_content)
         if parsed:
-            transcript_idx, block_idx, start_pattern = parsed
-            replacement = f"T{transcript_idx}B{block_idx}"
+            # Create appropriate replacement text based on citation type
+            if parsed.metadata_key:
+                if parsed.transcript_idx is None:
+                    # Agent run metadata [M.key]
+                    replacement = "run metadata"
+                elif parsed.block_idx is None:
+                    # Transcript metadata [T0M.key]
+                    replacement = f"T{parsed.transcript_idx}"
+                else:
+                    # Message metadata [T0B1M.key]
+                    replacement = f"T{parsed.transcript_idx}B{parsed.block_idx}"
+            else:
+                # Regular transcript block [T0B1]
+                replacement = f"T{parsed.transcript_idx}B{parsed.block_idx}"
+
             # Current absolute start position for this replacement in the cleaned text
             start_idx = len(cleaned_text)
             end_idx = start_idx + len(replacement)
@@ -137,10 +217,11 @@ def parse_citations(text: str) -> tuple[str, list[Citation]]:
                     start_idx=start_idx,
                     end_idx=end_idx,
                     agent_run_idx=None,
-                    transcript_idx=transcript_idx,
-                    block_idx=block_idx,
+                    transcript_idx=parsed.transcript_idx,
+                    block_idx=parsed.block_idx,
                     action_unit_idx=None,
-                    start_pattern=start_pattern,
+                    metadata_key=parsed.metadata_key,
+                    start_pattern=parsed.start_pattern,
                 )
             )
             cleaned_text += replacement

docent/data_models/judge.py

@@ -0,0 +1,19 @@
+"""Judge-related data models shared across Docent components."""
+
+from typing import Any
+from uuid import uuid4
+
+from pydantic import BaseModel, Field
+
+
+class Label(BaseModel):
+    id: str = Field(default_factory=lambda: str(uuid4()))
+
+    label_set_id: str
+
+    label_value: dict[str, Any]
+
+    agent_run_id: str
+
+
+__all__ = ["Label"]

docent/data_models/metadata_util.py

@@ -0,0 +1,16 @@
+import json
+from typing import Any
+
+from pydantic_core import to_jsonable_python
+
+
+def dump_metadata(metadata: dict[str, Any]) -> str | None:
+    """
+    Dump metadata to a JSON string.
+    We used to use YAML to save tokens, but JSON makes it easier to find cited ranges on the frontend because the frontend uses JSON.
+    """
+    if not metadata:
+        return None
+    metadata_obj = to_jsonable_python(metadata)
+    text = json.dumps(metadata_obj, indent=2)
+    return text.strip()

docent/data_models/remove_invalid_citation_ranges.py

@@ -1,3 +1,4 @@
+import json
 import re
 
 from docent.data_models.agent_run import AgentRun
@@ -52,7 +53,7 @@ def find_citation_matches_in_text(text: str, start_pattern: str) -> list[tuple[i
 
 def get_transcript_text_for_citation(agent_run: AgentRun, citation: Citation) -> str | None:
     """
-    Get the text content of a specific transcript block from an AgentRun,
+    Get the text content of a specific transcript block (or transcript/run metadata) from an AgentRun,
     using the same formatting as shown to LLMs via format_chat_message.
 
     Args:
@@ -62,19 +63,28 @@ def get_transcript_text_for_citation(agent_run: AgentRun, citation: Citation) ->
     Returns:
         Text content of the specified block (including tool calls), or None if not found
     """
-    if citation.transcript_idx is None:
-        return None
-
     try:
-        if citation.transcript_idx >= len(agent_run.get_transcript_ids_ordered()):
+        if citation.transcript_idx is None:
+            # At the run level, can only cite metadata
+            if citation.metadata_key is not None:
+                return json.dumps(agent_run.metadata.get(citation.metadata_key))
             return None
+
         transcript_id = agent_run.get_transcript_ids_ordered()[citation.transcript_idx]
         transcript = agent_run.transcript_dict[transcript_id]
 
-        if citation.block_idx >= len(transcript.messages):
+        if citation.block_idx is None:
+            # At the transcript level, can only cite metadata
+            if citation.metadata_key is not None:
+                return json.dumps(transcript.metadata.get(citation.metadata_key))
             return None
+
         message = transcript.messages[citation.block_idx]
 
+        # At the message level, can cite metadata or content
+        if citation.metadata_key is not None:
+            return json.dumps(message.metadata.get(citation.metadata_key))
+
         # Use the same formatting function that generates content for LLMs
         # This ensures consistent formatting between citation validation and LLM serialization
         return format_chat_message(
@@ -99,6 +109,9 @@ def validate_citation_text_range(agent_run: AgentRun, citation: Citation) -> boo
     if not citation.start_pattern:
         # Nothing to validate
         return True
+    if citation.metadata_key is not None:
+        # We don't need to remove invalid metadata citation ranges
+        return True
 
     text = get_transcript_text_for_citation(agent_run, citation)
     if text is None:
@@ -130,16 +143,16 @@ def remove_invalid_citation_ranges(text: str, agent_run: AgentRun) -> str:
         # Parse this bracket content to get citation info
         parsed = parse_single_citation(bracket_content)
         if parsed:
-            transcript_idx, block_idx, start_pattern = parsed
             # The citation spans from start to end in the original text
             citation = Citation(
                 start_idx=start,
                 end_idx=end,
                 agent_run_idx=None,
-                transcript_idx=transcript_idx,
-                block_idx=block_idx,
+                transcript_idx=parsed.transcript_idx,
+                block_idx=parsed.block_idx,
                 action_unit_idx=None,
-                start_pattern=start_pattern,
+                metadata_key=parsed.metadata_key,
+                start_pattern=parsed.start_pattern,
             )
             citations.append(citation)
 

docent/data_models/transcript.py

@@ -15,7 +15,7 @@ from docent.data_models._tiktoken_util import (
 )
 from docent.data_models.chat import AssistantMessage, ChatMessage, ContentReasoning
 from docent.data_models.citation import RANGE_BEGIN, RANGE_END
-from docent.data_models.yaml_util import yaml_dump_metadata
+from docent.data_models.metadata_util import dump_metadata
 
 # Template for formatting individual transcript blocks
 TRANSCRIPT_BLOCK_TEMPLATE = """
@@ -29,6 +29,12 @@ TEXT_RANGE_CITE_INSTRUCTION = f"""Anytime you quote the transcript, or refer to
 
 A citation may include a specific range of text within a block. Use {RANGE_BEGIN} and {RANGE_END} to mark the specific range of text. Add it after the block ID separated by a colon. For example, to cite the part of transcript 0, block 1, where the agent says "I understand the task", write [T0B1:{RANGE_BEGIN}I understand the task{RANGE_END}]. Citations must follow this exact format. The markers {RANGE_BEGIN} and {RANGE_END} must be used ONLY inside the brackets of a citation.
 
+- You may cite a top-level key in the agent run metadata like this: [M.task_description].
+- You may cite a top-level key in transcript metadata. For example, for transcript 0: [T0M.start_time].
+- You may cite a top-level key in message metadata for a block. For example, for transcript 0, block 1: [T0B1M.status].
+- You may not cite nested keys. For example, [T0B1M.status.code] is invalid.
+- Within a top-level metadata key you may cite a range of text that appears in the value. For example, [T0B1M.status:{RANGE_BEGIN}"running":false{RANGE_END}].
+
 Important notes:
 - You must include the full content of the text range {RANGE_BEGIN} and {RANGE_END}, EXACTLY as it appears in the transcript, word-for-word, including any markers or punctuation that appear in the middle of the text.
 - Citations must be as specific as possible. This means you should usually cite a specific text range within a block.
@@ -36,6 +42,8 @@ Important notes:
 - Citations are self-contained. Do NOT label them as citation or evidence. Just insert the citation by itself at the appropriate place in the text.
 - Citations must come immediately after the part of a claim that they support. This may be in the middle of a sentence.
 - Each pair of brackets must contain only one citation. To cite multiple blocks, use multiple pairs of brackets, like [T0B0] [T0B1].
+- Outside of citations, do not refer to transcript numbers or block numbers.
+- Outside of citations, avoid quoting or paraphrasing the transcript.
 """
 
 BLOCK_CITE_INSTRUCTION = """Each transcript and each block has a unique index. Cite the relevant indices in brackets when relevant, like [T<idx>B<idx>]. Use multiple tags to cite multiple blocks, like [T<idx1>B<idx1>][T<idx2>B<idx2>]. Remember to cite specific blocks and NOT action units."""
@@ -73,9 +81,9 @@ def format_chat_message(
                 cur_content += f"\n<tool call>\n{tool_call.function}({args})\n</tool call>"
 
     if message.metadata:
-        metadata_yaml = yaml_dump_metadata(message.metadata)
-        if metadata_yaml is not None:
-            cur_content += f"\n<|message metadata|>\n{metadata_yaml}\n</|message metadata|>"
+        metadata_text = dump_metadata(message.metadata)
+        if metadata_text is not None:
+            cur_content += f"\n<|message metadata|>\n{metadata_text}\n</|message metadata|>"
 
     return TRANSCRIPT_BLOCK_TEMPLATE.format(
         index_label=index_label, role=message.role, content=cur_content
@@ -127,13 +135,11 @@ class TranscriptGroup(BaseModel):
             str: XML-like wrapped text including the group's metadata.
         """
         # Prepare YAML metadata
-        yaml_text = yaml_dump_metadata(self.metadata)
-        if yaml_text is not None:
+        metadata_text = dump_metadata(self.metadata)
+        if metadata_text is not None:
             if indent > 0:
-                yaml_text = textwrap.indent(yaml_text, " " * indent)
-            inner = (
-                f"{children_text}\n<|{self.name} metadata|>\n{yaml_text}\n</|{self.name} metadata|>"
-            )
+                metadata_text = textwrap.indent(metadata_text, " " * indent)
+            inner = f"{children_text}\n<|{self.name} metadata|>\n{metadata_text}\n</|{self.name} metadata|>"
         else:
             inner = children_text
 
@@ -447,13 +453,11 @@ class Transcript(BaseModel):
         content_str = f"<|T{transcript_idx} blocks|>\n{blocks_str}\n</|T{transcript_idx} blocks|>"
 
         # Gather metadata and add to content
-        yaml_text = yaml_dump_metadata(self.metadata)
-        if yaml_text is not None:
+        metadata_text = dump_metadata(self.metadata)
+        if metadata_text is not None:
             if indent > 0:
-                yaml_text = textwrap.indent(yaml_text, " " * indent)
-            content_str += (
-                f"\n<|T{transcript_idx} metadata|>\n{yaml_text}\n</|T{transcript_idx} metadata|>"
-            )
+                metadata_text = textwrap.indent(metadata_text, " " * indent)
+            content_str += f"\n<|T{transcript_idx} metadata|>\n{metadata_text}\n</|T{transcript_idx} metadata|>"
 
         # Format content and return
         if indent > 0:

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,23 @@
+from docent.judges.impl import BaseJudge, MajorityVotingJudge, MultiReflectionJudge
+from docent.judges.types import (
+    JudgeResult,
+    JudgeResultCompletionCallback,
+    JudgeResultWithCitations,
+    JudgeVariant,
+    ResultType,
+    Rubric,
+)
+
+__all__ = [
+    # Judges
+    "MajorityVotingJudge",
+    "MultiReflectionJudge",
+    "BaseJudge",
+    # Types
+    "Rubric",
+    "JudgeResult",
+    "JudgeResultWithCitations",
+    "JudgeResultCompletionCallback",
+    "ResultType",
+    "JudgeVariant",
+]

docent/judges/analysis.py

@@ -0,0 +1,77 @@
+import json
+from pathlib import Path
+from typing import Any
+
+import anyio
+from pydantic import BaseModel
+from pydantic_core import to_jsonable_python
+from tqdm.auto import tqdm
+
+from docent._log_util import get_logger
+from docent.data_models.agent_run import AgentRun
+from docent.judges.impl import BaseJudge
+from docent.judges.util.voting import JudgeOutputDistribution
+
+logger = get_logger(__name__)
+
+
+class MultiReflectRollouts(BaseModel):
+    """Object is associated with a single agent run"""
+
+    agent_run_id: str
+
+    first_step_rollouts: list[dict[str, Any]]
+    first_step_rollout_metadata: list[dict[str, Any] | None]
+    # Each index in second_step_rollouts corresponds to an index in first_step_combinations
+    # Step 2 rollouts are computed by passing each step 1 combo into the judge several times
+    first_step_combinations: list[list[dict[str, Any]]] | None = None
+    second_step_rollouts: list[list[dict[str, Any]]] | None = None
+    second_step_rollout_metadata: list[list[dict[str, Any] | None]] | None = None
+
+    distributions: dict[str, JudgeOutputDistribution]
+
+
+async def collect_judge_pvs(
+    judge: BaseJudge,
+    agent_runs: list[AgentRun],
+    *,
+    results_path: Path,
+    estimate_output_distrs_kwargs: dict[str, Any],
+):
+    if results_path.exists():
+        raise FileExistsError(f"Results path already exists: {results_path}")
+    results_path.parent.mkdir(parents=True, exist_ok=True)
+
+    results = dict[str, MultiReflectRollouts]()
+    persist_lock = anyio.Lock()
+    pbar = tqdm(total=len(agent_runs), desc="Processing agent runs")
+
+    async def _persist():
+        async with persist_lock:
+            with open(str(results_path), "w") as f:
+                json.dump(to_jsonable_python(results), f, indent=2)
+
+    async def _execute_for_agent_run(agent_run: AgentRun):
+        result = await judge.estimate_output_distrs(agent_run, **estimate_output_distrs_kwargs)
+        if result is None:
+            pbar.update(1)
+            return
+
+        distrs, metadata = result
+        results[agent_run.id] = MultiReflectRollouts.model_validate(
+            {
+                "agent_run_id": agent_run.id,
+                "distributions": distrs,
+                **metadata,
+            }
+        )
+        await _persist()
+        pbar.update(1)
+
+    async with anyio.create_task_group() as tg_outer:
+        for agent_run in agent_runs:
+            tg_outer.start_soon(_execute_for_agent_run, agent_run)
+
+    pbar.close()
+
+    return results