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

docent/data_models/__init__.py

@@ -1,11 +1,13 @@
 from docent.data_models.agent_run import AgentRun
 from docent.data_models.citation import Citation
+from docent.data_models.judge import JudgeRunLabel
 from docent.data_models.regex import RegexSnippet
 from docent.data_models.transcript import Transcript, TranscriptGroup
 
 __all__ = [
     "AgentRun",
     "Citation",
+    "JudgeRunLabel",
     "RegexSnippet",
     "Transcript",
     "TranscriptGroup",

docent/data_models/agent_run.py

@@ -17,8 +17,8 @@ from pydantic_core import to_jsonable_python
 
 from docent._log_util import get_logger
 from docent.data_models._tiktoken_util import get_token_count, group_messages_into_ranges
+from docent.data_models.metadata_util import dump_metadata
 from docent.data_models.transcript import Transcript, TranscriptGroup
-from docent.data_models.yaml_util import yaml_dump_metadata
 
 logger = get_logger(__name__)
 
@@ -446,10 +446,10 @@ class AgentRun(BaseModel):
         text = _recurse("__global_root")
 
         # Append agent run metadata below the full 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)
-            text += f"\n<|agent run metadata|>\n{yaml_text}\n</|agent run metadata|>"
+                metadata_text = textwrap.indent(metadata_text, " " * indent)
+            text += f"\n<|agent run metadata|>\n{metadata_text}\n</|agent run metadata|>"
 
         return text

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,16 @@
+"""Judge-related data models shared across Docent components."""
+
+from typing import Any
+from uuid import uuid4
+
+from pydantic import BaseModel, Field
+
+
+class JudgeRunLabel(BaseModel):
+    id: str = Field(default_factory=lambda: str(uuid4()))
+    agent_run_id: str
+    rubric_id: str
+    label: dict[str, Any]
+
+
+__all__ = ["JudgeRunLabel"]

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.
@@ -73,9 +79,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 +133,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 +451,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/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,
     )