docent-python 0.1.10a0__tar.gz → 0.1.12a0__tar.gz

{docent_python-0.1.10a0 → docent_python-0.1.12a0}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: docent-python
-Version: 0.1.10a0
+Version: 0.1.12a0
 Summary: Docent SDK
 Project-URL: Homepage, https://github.com/TransluceAI/docent
 Project-URL: Issues, https://github.com/TransluceAI/docent/issues

{docent_python-0.1.10a0 → docent_python-0.1.12a0}/docent/data_models/agent_run.py

@@ -90,19 +90,36 @@ class AgentRun(BaseModel):
             raise ValueError("AgentRun must have at least one transcript")
         return self
 
-    def to_text(self, token_limit: int = sys.maxsize) -> list[str]:
+    def _to_text_impl(self, token_limit: int = sys.maxsize, use_blocks: bool = False) -> list[str]:
         """
-        Represents an agent run as a list of strings, each of which is at most token_limit tokens
-        under the GPT-4 tokenization scheme.
+        Core implementation for converting agent run to text representation.
 
-        We'll try to split up long AgentRuns along transcript boundaries and include metadata.
-        For very long transcripts, we'll have to split them up further and remove metadata.
+        Args:
+            token_limit: Maximum tokens per returned string under the GPT-4 tokenization scheme
+            use_blocks: If True, use individual message blocks. If False, use action units.
+
+        Returns:
+            List of strings, each at most token_limit tokens
         """
+        # Generate transcript strings using appropriate method
+        transcript_strs: list[str] = []
+        for i, (t_key, t) in enumerate(self.transcripts.items()):
+            if use_blocks:
+                transcript_content = t.to_str_blocks_with_token_limit(
+                    token_limit=sys.maxsize,
+                    transcript_idx=i,
+                    agent_run_idx=None,
+                )[0]
+            else:
+                transcript_content = t.to_str_with_token_limit(
+                    token_limit=sys.maxsize,
+                    transcript_idx=i,
+                    agent_run_idx=None,
+                )[0]
+            transcript_strs.append(
+                f"<transcript {t_key}>\n{transcript_content}\n</transcript {t_key}>"
+            )
 
-        transcript_strs: list[str] = [
-            f"<transcript {t_key}>\n{t.to_str(agent_run_idx=None, transcript_idx=i)}\n</transcript {t_key}>"
-            for i, (t_key, t) in enumerate(self.transcripts.items())
-        ]
         transcripts_str = "\n\n".join(transcript_strs)
 
         # Gather metadata
@@ -128,7 +145,6 @@ class AgentRun(BaseModel):
             return [f"{transcripts_str}" f"{metadata_str}"]
 
         # Otherwise, split up the transcript and metadata into chunks
-        # TODO(vincent, mengk): does this code account for multiple transcripts correctly? a little confused.
         else:
             results: list[str] = []
             transcript_token_counts = [get_token_count(t) for t in transcript_strs]
@@ -150,13 +166,23 @@ class AgentRun(BaseModel):
                     ), "Ranges without metadata should be a single message"
                     t_id, t = list(self.transcripts.items())[msg_range.start]
                     if msg_range.num_tokens < token_limit - 50:
-                        transcript = f"<transcript {t_id}>\n{t.to_str()}\n</transcript {t_id}>"
+                        if use_blocks:
+                            transcript = f"<transcript {t_id}>\n{t.to_str_blocks_with_token_limit(token_limit=sys.maxsize)[0]}\n</transcript {t_id}>"
+                        else:
+                            transcript = f"<transcript {t_id}>\n{t.to_str_with_token_limit(token_limit=sys.maxsize)[0]}\n</transcript {t_id}>"
                         result = (
                             f"Here is a partial agent run for analysis purposes only:\n{transcript}"
                         )
                         results.append(result)
                     else:
-                        transcript_fragments = t.to_str_with_token_limit(token_limit - 50)
+                        if use_blocks:
+                            transcript_fragments = t.to_str_blocks_with_token_limit(
+                                token_limit=token_limit - 50,
+                            )
+                        else:
+                            transcript_fragments = t.to_str_with_token_limit(
+                                token_limit=token_limit - 50,
+                            )
                         for fragment in transcript_fragments:
                             result = f"<transcript {t_id}>\n{fragment}\n</transcript {t_id}>"
                             result = (
@@ -165,6 +191,26 @@ class AgentRun(BaseModel):
                             results.append(result)
             return results
 
+    def to_text(self, token_limit: int = sys.maxsize) -> list[str]:
+        """
+        Represents an agent run as a list of strings, each of which is at most token_limit tokens
+        under the GPT-4 tokenization scheme.
+
+        We'll try to split up long AgentRuns along transcript boundaries and include metadata.
+        For very long transcripts, we'll have to split them up further and remove metadata.
+        """
+        return self._to_text_impl(token_limit=token_limit, use_blocks=False)
+
+    def to_text_blocks(self, token_limit: int = sys.maxsize) -> list[str]:
+        """
+        Represents an agent run as a list of strings using individual message blocks,
+        each of which is at most token_limit tokens under the GPT-4 tokenization scheme.
+
+        Unlike to_text() which uses action units, this method formats each message
+        as an individual block.
+        """
+        return self._to_text_impl(token_limit=token_limit, use_blocks=True)
+
     @property
     def text(self) -> str:
         """Concatenates all transcript texts with double newlines as separators.
@@ -172,7 +218,16 @@ class AgentRun(BaseModel):
         Returns:
             str: A string representation of all transcripts.
         """
-        return self.to_text()[0]
+        return self._to_text_impl(token_limit=sys.maxsize, use_blocks=False)[0]
+
+    @property
+    def text_blocks(self) -> str:
+        """Concatenates all transcript texts using individual blocks format.
+
+        Returns:
+            str: A string representation of all transcripts using individual message blocks.
+        """
+        return self._to_text_impl(token_limit=sys.maxsize, use_blocks=True)[0]
 
     def model_dump(self, *args: Any, **kwargs: Any) -> dict[str, Any]:
         """Extends the parent model_dump method to include the text property.

{docent_python-0.1.10a0 → docent_python-0.1.12a0}/docent/data_models/chat/message.py

@@ -5,6 +5,7 @@ from pydantic import BaseModel, Discriminator
 
 from docent.data_models.chat.content import Content
 from docent.data_models.chat.tool import ToolCall
+from docent.data_models.citation import Citation
 
 logger = getLogger(__name__)
 
@@ -66,11 +67,15 @@ class AssistantMessage(BaseChatMessage):
         role: Always set to "assistant".
         model: Optional identifier for the model that generated this message.
         tool_calls: Optional list of tool calls made by the assistant.
+        citations: Optional list of citations referenced in the message content.
+        suggested_messages: Optional list of suggested followup messages.
     """
 
     role: Literal["assistant"] = "assistant"  # type: ignore
     model: str | None = None
     tool_calls: list[ToolCall] | None = None
+    citations: list[Citation] | None = None
+    suggested_messages: list[str] | None = None
 
 
 class ToolMessage(BaseChatMessage):

docent_python-0.1.12a0/docent/data_models/citation.py

@@ -0,0 +1,152 @@
+import re
+
+from pydantic import BaseModel
+
+
+class Citation(BaseModel):
+    start_idx: int
+    end_idx: int
+    agent_run_idx: int | None = None
+    transcript_idx: int | None = None
+    block_idx: int
+    action_unit_idx: int | None = None
+    start_pattern: str | None = None
+
+
+RANGE_BEGIN = "<RANGE>"
+RANGE_END = "</RANGE>"
+
+_SINGLE_RE = re.compile(r"T(\d+)B(\d+)")
+_RANGE_CONTENT_RE = re.compile(r":\s*" + re.escape(RANGE_BEGIN) + r".*?" + re.escape(RANGE_END))
+
+
+def _extract_range_pattern(range_part: str) -> str | None:
+    start_pattern: str | None = None
+
+    if RANGE_BEGIN in range_part and RANGE_END in range_part:
+        range_begin_idx = range_part.find(RANGE_BEGIN)
+        range_end_idx = range_part.find(RANGE_END)
+        if range_begin_idx != -1 and range_end_idx != -1:
+            range_content = range_part[range_begin_idx + len(RANGE_BEGIN) : range_end_idx]
+            start_pattern = range_content if range_content else None
+
+    return start_pattern
+
+
+def scan_brackets(text: str) -> list[tuple[int, int, str]]:
+    """Scan text for bracketed segments, respecting RANGE markers and nested brackets.
+
+    Returns a list of (start_index, end_index_exclusive, inner_content).
+    """
+    matches: list[tuple[int, int, str]] = []
+    i = 0
+    while i < len(text):
+        if text[i] == "[":
+            start = i
+            bracket_count = 1
+            j = i + 1
+            in_range = False
+
+            while j < len(text) and bracket_count > 0:
+                if text[j : j + len(RANGE_BEGIN)] == RANGE_BEGIN:
+                    in_range = True
+                elif text[j : j + len(RANGE_END)] == RANGE_END:
+                    in_range = False
+                elif text[j] == "[" and not in_range:
+                    bracket_count += 1
+                elif text[j] == "]" and not in_range:
+                    bracket_count -= 1
+                j += 1
+
+            if bracket_count == 0:
+                end_exclusive = j
+                bracket_content = text[start + 1 : end_exclusive - 1]
+                matches.append((start, end_exclusive, bracket_content))
+                i = j
+            else:
+                i += 1
+        else:
+            i += 1
+    return matches
+
+
+def parse_single_citation(part: str) -> tuple[int, int, str | None] | None:
+    """
+    Parse a single citation token inside a bracket and return its components.
+
+    Returns (transcript_idx, block_idx, start_pattern) or None if invalid.
+    """
+    token = part.strip()
+    if not token:
+        return None
+
+    if ":" in token:
+        citation_part, range_part = token.split(":", 1)
+        single_match = _SINGLE_RE.match(citation_part.strip())
+        if not single_match:
+            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 None
+        transcript_idx = int(single_match.group(1))
+        block_idx = int(single_match.group(2))
+        return transcript_idx, block_idx, None
+
+
+def parse_citations(text: str) -> tuple[str, list[Citation]]:
+    """
+    Parse citations from text in the format described by BLOCK_RANGE_CITE_INSTRUCTION.
+
+    Supported formats:
+    - Single block: [T<key>B<idx>]
+    - Text range with start pattern: [T<key>B<idx>:<RANGE>start_pattern</RANGE>]
+
+    Args:
+        text: The text to parse citations from
+
+    Returns:
+        A tuple of (cleaned_text, citations) where cleaned_text has brackets and range markers removed
+        and citations have start_idx and end_idx representing character positions
+        in the cleaned text
+    """
+    citations: list[Citation] = []
+    cleaned_text = ""
+
+    bracket_matches = scan_brackets(text)
+
+    last_end = 0
+    for start, end, bracket_content in bracket_matches:
+        # Append non-bracket text segment as-is
+        cleaned_text += text[last_end:start]
+
+        # 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}"
+            # Current absolute start position for this replacement in the cleaned text
+            start_idx = len(cleaned_text)
+            end_idx = start_idx + len(replacement)
+            citations.append(
+                Citation(
+                    start_idx=start_idx,
+                    end_idx=end_idx,
+                    agent_run_idx=None,
+                    transcript_idx=transcript_idx,
+                    block_idx=block_idx,
+                    action_unit_idx=None,
+                    start_pattern=start_pattern,
+                )
+            )
+            cleaned_text += replacement
+        last_end = end
+
+    # Append any remaining tail after the last bracket
+    cleaned_text += text[last_end:]
+
+    return cleaned_text, citations

docent_python-0.1.12a0/docent/data_models/remove_invalid_citation_ranges.py

@@ -0,0 +1,166 @@
+import re
+
+from docent.data_models.agent_run import AgentRun
+from docent.data_models.citation import Citation, parse_single_citation, scan_brackets
+from docent.data_models.transcript import format_chat_message
+
+
+def build_whitespace_flexible_regex(pattern: str) -> re.Pattern[str]:
+    """Build regex that is flexible with whitespace matching."""
+    out = ""
+    i = 0
+    while i < len(pattern):
+        ch = pattern[i]
+        if ch.isspace():
+            # Skip all consecutive whitespace
+            while i < len(pattern) and pattern[i].isspace():
+                i += 1
+            out += r"\s+"
+            continue
+        out += re.escape(ch)
+        i += 1
+    return re.compile(out, re.DOTALL)
+
+
+def find_citation_matches_in_text(text: str, start_pattern: str) -> list[tuple[int, int]]:
+    """
+    Find all matches of a citation pattern in text.
+
+    Args:
+        text: The text to search in
+        start_pattern: The pattern to search for
+
+    Returns:
+        List of (start_index, end_index) tuples for matches
+    """
+    if not start_pattern:
+        return []
+
+    try:
+        regex = build_whitespace_flexible_regex(start_pattern)
+        matches: list[tuple[int, int]] = []
+
+        for match in regex.finditer(text):
+            if match.group().strip():  # Only count non-empty matches
+                matches.append((match.start(), match.end()))
+
+        return matches
+
+    except re.error:
+        return []
+
+
+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,
+    using the same formatting as shown to LLMs via format_chat_message.
+
+    Args:
+        agent_run: The agent run containing transcript data
+        citation: Citation with transcript_idx and block_idx
+
+    Returns:
+        Text content of the specified block (including tool calls), or None if not found
+    """
+    if citation.transcript_idx is None:
+        return None
+
+    try:
+        transcript_keys = list(agent_run.transcripts.keys())
+        if citation.transcript_idx >= len(transcript_keys):
+            return None
+
+        transcript_key = transcript_keys[citation.transcript_idx]
+
+        transcript = agent_run.transcripts[transcript_key]
+        if citation.block_idx >= len(transcript.messages):
+            return None
+
+        message = transcript.messages[citation.block_idx]
+
+        # Use the same formatting function that generates content for LLMs
+        # This ensures consistent formatting between citation validation and LLM serialization
+        return format_chat_message(
+            message, citation.block_idx, citation.transcript_idx, citation.agent_run_idx
+        )
+
+    except (KeyError, IndexError, AttributeError):
+        return None
+
+
+def validate_citation_text_range(agent_run: AgentRun, citation: Citation) -> bool:
+    """
+    Validate that a citation's text range exists in the referenced transcript.
+
+    Args:
+        agent_run: The agent run containing transcript data
+        citation: Citation to validate
+
+    Returns:
+        True if the citation's text range exists in the transcript, False otherwise
+    """
+    if not citation.start_pattern:
+        # Nothing to validate
+        return True
+
+    text = get_transcript_text_for_citation(agent_run, citation)
+    if text is None:
+        return False
+
+    matches = find_citation_matches_in_text(text, citation.start_pattern)
+
+    return len(matches) > 0
+
+
+def remove_invalid_citation_ranges(text: str, agent_run: AgentRun) -> str:
+    """
+    Remove invalid citation ranges from chat message/judge result. We do this as a separate step before normal citation parsing.
+    Normal citation parsing happens every time we load chat/results from db,
+    but invalid ranges should never make it to the db.
+
+    Args:
+        text: Original text containing citations
+        agent_run: Agent run with transcript data
+
+    Returns:
+        Tuple of (cleaned_text, valid_citations)
+    """
+    # Find all bracket positions in the original text
+    bracket_matches = scan_brackets(text)
+    citations: list[Citation] = []
+
+    for start, end, bracket_content in bracket_matches:
+        # 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,
+                action_unit_idx=None,
+                start_pattern=start_pattern,
+            )
+            citations.append(citation)
+
+    # Filter to only citations with text ranges that need validation
+    citations_to_validate = [c for c in citations if c.start_pattern]
+
+    # Sort citations by start_idx in reverse order to avoid index shifting issues
+    sorted_citations = sorted(citations_to_validate, key=lambda c: c.start_idx, reverse=True)
+
+    invalid_citations: list[Citation] = [
+        c for c in sorted_citations if not validate_citation_text_range(agent_run, c)
+    ]
+
+    # Remove invalid text ranges from citations in the original text
+    modified_text = text
+    for citation in invalid_citations:
+        citation_without_range = f"[T{citation.transcript_idx}B{citation.block_idx}]"
+        before = modified_text[: citation.start_idx]
+        after = modified_text[citation.end_idx :]
+        modified_text = before + citation_without_range + after
+    return modified_text

{docent_python-0.1.10a0 → docent_python-0.1.12a0}/docent/data_models/transcript.py

@@ -12,6 +12,7 @@ from docent.data_models._tiktoken_util import (
     truncate_to_token_limit,
 )
 from docent.data_models.chat import AssistantMessage, ChatMessage, ContentReasoning
+from docent.data_models.citation import RANGE_BEGIN, RANGE_END
 
 # Template for formatting individual transcript blocks
 TRANSCRIPT_BLOCK_TEMPLATE = """
@@ -21,10 +22,20 @@ TRANSCRIPT_BLOCK_TEMPLATE = """
 """.strip()
 
 # Instructions for citing single transcript blocks
-SINGLE_RUN_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>]. Use an inner dash to cite a range of blocks, like [T<idx1>B<idx1>-T<idx2>B<idx2>]. Remember to cite specific blocks and NOT action units."
+TEXT_RANGE_CITE_INSTRUCTION = f"""Anytime you quote the transcript, or refer to something that happened in the transcript, or make any claim about the transcript, add an inline citation. Each transcript and each block has a unique index. Cite the relevant indices in brackets. For example, to cite the entirety of transcript 0, block 1, write [T0B1].
 
-# Instructions for citing multiple transcript blocks
-MULTI_RUN_CITE_INSTRUCTION = "Each run, each transcript, and each block has a unique index. Cite the relevant indices in brackets when relevant, like [R<idx>T<idx>B<idx>]. Use multiple tags to cite multiple blocks, like [R<idx1>T<idx1>B<idx1>][R<idx2>T<idx2>B<idx2>]. Use an inner dash to cite a range of blocks, like [R<idx1>T<idx1>B<idx1>-R<idx2>T<idx2>B<idx2>]. Remember to cite specific blocks and NOT action units."
+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.
+
+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.
+- A citation is not a quote. For brevity, text ranges will not be rendered inline. The user will have to click on the citation to see the full text range.
+- 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].
+"""
+
+BLOCK_CITE_INSTRUCTION = f"""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."""
 
 
 def format_chat_message(
@@ -291,66 +302,105 @@ class Transcript(BaseModel):
         agent_run_idx: int | None = None,
         highlight_action_unit: int | None = None,
     ) -> str:
-        return self.to_str_with_token_limit(
+        return self._to_str_with_token_limit_impl(
             token_limit=sys.maxsize,
-            agent_run_idx=agent_run_idx,
             transcript_idx=transcript_idx,
+            agent_run_idx=agent_run_idx,
+            use_action_units=True,
             highlight_action_unit=highlight_action_unit,
         )[0]
 
-    def to_str_with_token_limit(
+    def _generate_formatted_blocks(
         self,
-        token_limit: int,
         transcript_idx: int = 0,
         agent_run_idx: int | None = None,
+        use_action_units: bool = True,
         highlight_action_unit: int | None = None,
     ) -> list[str]:
-        """Represents the transcript as a list of strings, each of which is at most token_limit tokens
-        under the GPT-4 tokenization scheme.
+        """Generate formatted blocks for transcript representation.
 
-        We'll try to split up long transcripts along message boundaries and include metadata.
-        For very long messages, we'll have to truncate them and remove metadata.
+        Args:
+            transcript_idx: Index of the transcript
+            agent_run_idx: Optional agent run index
+            use_action_units: If True, group messages into action units. If False, use individual blocks.
+            highlight_action_unit: Optional action unit to highlight (only used with action units)
 
         Returns:
-            list[str]: A list of strings, each of which is at most token_limit tokens
-            under the GPT-4 tokenization scheme.
+            list[str]: List of formatted blocks
         """
-        if highlight_action_unit is not None and not (
-            0 <= highlight_action_unit < len(self._units_of_action or [])
-        ):
-            raise ValueError(f"Invalid action unit index: {highlight_action_unit}")
-
-        # Format blocks by units of action
-        au_blocks: list[str] = []
-        for unit_idx, unit in enumerate(self._units_of_action or []):
-            unit_blocks: list[str] = []
-            for msg_idx in unit:
-                unit_blocks.append(
+        if use_action_units:
+            if highlight_action_unit is not None and not (
+                0 <= highlight_action_unit < len(self._units_of_action or [])
+            ):
+                raise ValueError(f"Invalid action unit index: {highlight_action_unit}")
+
+            blocks: list[str] = []
+            for unit_idx, unit in enumerate(self._units_of_action or []):
+                unit_blocks: list[str] = []
+                for msg_idx in unit:
+                    unit_blocks.append(
+                        format_chat_message(
+                            self.messages[msg_idx],
+                            msg_idx,
+                            transcript_idx,
+                            agent_run_idx,
+                        )
+                    )
+
+                unit_content = "\n".join(unit_blocks)
+
+                # Add highlighting if requested
+                if highlight_action_unit and unit_idx == highlight_action_unit:
+                    blocks_str_template = "<HIGHLIGHTED>\n{}\n</HIGHLIGHTED>"
+                else:
+                    blocks_str_template = "{}"
+                blocks.append(
+                    blocks_str_template.format(
+                        f"<action unit {unit_idx}>\n{unit_content}\n</action unit {unit_idx}>"
+                    )
+                )
+        else:
+            # Individual message blocks
+            blocks = []
+            for msg_idx, message in enumerate(self.messages):
+                blocks.append(
                     format_chat_message(
-                        self.messages[msg_idx],
+                        message,
                         msg_idx,
                         transcript_idx,
                         agent_run_idx,
                     )
                 )
 
-            unit_content = "\n".join(unit_blocks)
+        return blocks
 
-            # Add highlighting if requested
-            if highlight_action_unit and unit_idx == highlight_action_unit:
-                blocks_str_template = "<HIGHLIGHTED>\n{}\n</HIGHLIGHTED>"
-            else:
-                blocks_str_template = "{}"
-            au_blocks.append(
-                blocks_str_template.format(
-                    f"<action unit {unit_idx}>\n{unit_content}\n</action unit {unit_idx}>"
-                )
-            )
-        blocks_str = "\n".join(au_blocks)
+    def _to_str_with_token_limit_impl(
+        self,
+        token_limit: int,
+        transcript_idx: int = 0,
+        agent_run_idx: int | None = None,
+        use_action_units: bool = True,
+        highlight_action_unit: int | None = None,
+    ) -> list[str]:
+        """Core implementation for string representation with token limits.
+
+        Args:
+            token_limit: Maximum tokens per returned string
+            transcript_idx: Index of the transcript
+            agent_run_idx: Optional agent run index
+            use_action_units: If True, group messages into action units. If False, use individual blocks.
+            highlight_action_unit: Optional action unit to highlight (only used with action units)
+
+        Returns:
+            list[str]: List of strings, each within token limit
+        """
+        blocks = self._generate_formatted_blocks(
+            transcript_idx, agent_run_idx, use_action_units, highlight_action_unit
+        )
+        blocks_str = "\n".join(blocks)
 
         # Gather metadata
         metadata_obj = fake_model_dump(self.metadata)
-
         yaml_width = float("inf")
         block_str = f"<blocks>\n{blocks_str}\n</blocks>\n"
         metadata_str = f"<metadata>\n{yaml.dump(metadata_obj, width=yaml_width)}\n</metadata>"
@@ -365,25 +415,75 @@ class Transcript(BaseModel):
             return [f"{block_str}" f"{metadata_str}"]
         else:
             results: list[str] = []
-            block_token_counts = [get_token_count(block) for block in au_blocks]
+            block_token_counts = [get_token_count(block) for block in blocks]
             ranges = group_messages_into_ranges(
                 block_token_counts, metadata_token_count, token_limit
             )
             for msg_range in ranges:
                 if msg_range.include_metadata:
-                    cur_au_blocks = "\n".join(au_blocks[msg_range.start : msg_range.end])
-                    results.append(f"<blocks>\n{cur_au_blocks}\n</blocks>\n" f"{metadata_str}")
+                    cur_blocks = "\n".join(blocks[msg_range.start : msg_range.end])
+                    results.append(f"<blocks>\n{cur_blocks}\n</blocks>\n" f"{metadata_str}")
                 else:
                     assert (
                         msg_range.end == msg_range.start + 1
                     ), "Ranges without metadata should be a single message"
-                    result = str(au_blocks[msg_range.start])
+                    result = str(blocks[msg_range.start])
                     if msg_range.num_tokens > token_limit - 10:
                         result = truncate_to_token_limit(result, token_limit - 10)
                     results.append(f"<blocks>\n{result}\n</blocks>\n")
 
             return results
 
+    def to_str_blocks(
+        self,
+        transcript_idx: int = 0,
+        agent_run_idx: int | None = None,
+    ) -> str:
+        """Represents the transcript as a string using individual message blocks.
+
+        Unlike to_str() which groups messages into action units, this method
+        formats each message as an individual block.
+
+        Returns:
+            str: A string representation with individual message blocks.
+        """
+        return self._to_str_with_token_limit_impl(
+            token_limit=sys.maxsize,
+            transcript_idx=transcript_idx,
+            agent_run_idx=agent_run_idx,
+            use_action_units=False,
+        )[0]
+
+    def to_str_with_token_limit(
+        self,
+        token_limit: int,
+        transcript_idx: int = 0,
+        agent_run_idx: int | None = None,
+        highlight_action_unit: int | None = None,
+    ) -> list[str]:
+        """Represents the transcript as a list of strings using action units with token limit handling."""
+        return self._to_str_with_token_limit_impl(
+            token_limit=token_limit,
+            transcript_idx=transcript_idx,
+            agent_run_idx=agent_run_idx,
+            use_action_units=True,
+            highlight_action_unit=highlight_action_unit,
+        )
+
+    def to_str_blocks_with_token_limit(
+        self,
+        token_limit: int,
+        transcript_idx: int = 0,
+        agent_run_idx: int | None = None,
+    ) -> list[str]:
+        """Represents the transcript as individual blocks with token limit handling."""
+        return self._to_str_with_token_limit_impl(
+            token_limit=token_limit,
+            transcript_idx=transcript_idx,
+            agent_run_idx=agent_run_idx,
+            use_action_units=False,
+        )
+
 
 class TranscriptWithoutMetadataValidator(Transcript):
     """

{docent_python-0.1.10a0 → docent_python-0.1.12a0}/docent/sdk/client.py

@@ -350,3 +350,20 @@ class Docent:
 
         logger.info(f"Successfully shared Collection '{collection_id}' with {email}")
         return response.json()
+
+    def list_agent_run_ids(self, collection_id: str) -> list[str]:
+        """Get all agent run IDs for a collection.
+
+        Args:
+            collection_id: ID of the Collection.
+
+        Returns:
+            str: JSON string containing the list of agent run IDs.
+
+        Raises:
+            requests.exceptions.HTTPError: If the API request fails.
+        """
+        url = f"{self._server_url}/{collection_id}/agent_run_ids"
+        response = self._session.get(url)
+        response.raise_for_status()
+        return response.json()

{docent_python-0.1.10a0 → docent_python-0.1.12a0}/docent/trace.py

@@ -3,7 +3,6 @@ import contextvars
 import itertools
 import logging
 import os
-import signal
 import sys
 import threading
 import uuid
@@ -158,6 +157,7 @@ class DocentTracer:
             lambda: itertools.count(0)
         )
         self._transcript_counter_lock = threading.Lock()
+        self._flush_lock = threading.Lock()
 
     def get_current_agent_run_id(self) -> Optional[str]:
         """
@@ -179,14 +179,6 @@ class DocentTracer:
         # Register atexit handler
         atexit.register(self.cleanup)
 
-        # Register signal handlers for graceful shutdown
-        try:
-            signal.signal(signal.SIGINT, self._signal_handler)
-            signal.signal(signal.SIGTERM, self._signal_handler)
-        except (ValueError, OSError):
-            # Signal handlers might not work in all environments
-            pass
-
         self._cleanup_registered = True
 
     def _next_span_order(self, transcript_id: str) -> int:
@@ -197,10 +189,6 @@ class DocentTracer:
         with self._transcript_counter_lock:
             return next(self._transcript_counters[transcript_id])
 
-    def _signal_handler(self, signum: int, frame: Optional[object]):
-        """Handle shutdown signals."""
-        self.cleanup()
-
     def _init_spans_exporter(self, endpoint: str) -> Optional[Union[HTTPExporter, GRPCExporter]]:
         """Initialize the appropriate span exporter based on endpoint."""
         if not self.enable_otlp_export:
@@ -211,9 +199,11 @@ class DocentTracer:
                 http_exporter: HTTPExporter = HTTPExporter(
                     endpoint=f"{endpoint}/v1/traces", headers=self.headers
                 )
+                logger.debug(f"Initialized HTTP exporter for endpoint: {endpoint}/v1/traces")
                 return http_exporter
             else:
                 grpc_exporter: GRPCExporter = GRPCExporter(endpoint=endpoint, headers=self.headers)
+                logger.debug(f"Initialized gRPC exporter for endpoint: {endpoint}")
                 return grpc_exporter
         except Exception as e:
             logger.error(f"Failed to initialize span exporter for {endpoint}: {e}")
@@ -239,9 +229,11 @@ class DocentTracer:
         """Create appropriate span processor based on configuration."""
         if self.disable_batch or _is_notebook():
             simple_processor: SimpleSpanProcessor = SimpleSpanProcessor(exporter)
+            logger.debug("Created SimpleSpanProcessor for immediate export")
             return simple_processor
         else:
             batch_processor: BatchSpanProcessor = BatchSpanProcessor(exporter)
+            logger.debug("Created BatchSpanProcessor for batched export")
             return batch_processor
 
     def initialize(self):
@@ -310,8 +302,19 @@ class DocentTracer:
                         # attributes not available, skip them
                         pass
 
+                    # Debug logging for span creation
+                    span_name = getattr(span, "name", "unknown")
+                    span_attrs = getattr(span, "attributes", {})
+                    logger.debug(
+                        f"Created span: name='{span_name}', collection_id={self.manager.collection_id}, agent_run_id={span_attrs.get('agent_run_id')}, transcript_id={span_attrs.get('transcript_id')}"
+                    )
+
                 def on_end(self, span: ReadableSpan) -> None:
-                    pass
+                    # Debug logging for span completion
+                    span_attrs = span.attributes or {}
+                    logger.debug(
+                        f"Completed span: name='{span.name}', collection_id={span_attrs.get('collection_id')}, agent_run_id={span_attrs.get('agent_run_id')}, transcript_id={span_attrs.get('transcript_id')}, duration_ns={span.end_time - span.start_time if span.end_time and span.start_time else 'unknown'}"
+                    )
 
                 def shutdown(self) -> None:
                     pass
@@ -422,15 +425,8 @@ class DocentTracer:
             return
 
         try:
-            # Notify backend that trace is done (no span creation)
-            try:
-                self._send_trace_done()
-            except Exception as e:
-                logger.warning(f"Failed to notify trace done: {e}")
-
-            self._root_context = None  # type: ignore
+            self.flush()
 
-            # Shutdown our isolated tracer provider
             if self._tracer_provider:
                 self._tracer_provider.shutdown()
                 self._tracer_provider = None
@@ -456,9 +452,12 @@ class DocentTracer:
             return
 
         try:
-            for processor in self._spans_processors:
+            logger.debug(f"Flushing {len(self._spans_processors)} span processors")
+            for i, processor in enumerate(self._spans_processors):
                 if hasattr(processor, "force_flush"):
-                    processor.force_flush()
+                    logger.debug(f"Flushing span processor {i}")
+                    processor.force_flush(timeout_millis=50)
+            logger.debug("Span flush completed")
         except Exception as e:
             logger.error(f"Error during flush: {e}")
 
@@ -476,29 +475,6 @@ class DocentTracer:
         """Verify if the manager is properly initialized."""
         return self._initialized
 
-    def __enter__(self) -> "DocentTracer":
-        """Context manager entry."""
-        self.initialize()
-        return self
-
-    def __exit__(self, exc_type: type[BaseException], exc_val: Any, exc_tb: Any) -> None:
-        """Context manager exit."""
-        self.close()
-
-    @property
-    def tracer(self) -> Optional[trace.Tracer]:
-        """Get the tracer instance."""
-        if not self._initialized:
-            self.initialize()
-        return self._tracer
-
-    @property
-    def root_context(self) -> Optional[Context]:
-        """Get the root context."""
-        if not self._initialized:
-            self.initialize()
-        return self._root_context
-
     @contextmanager
     def agent_run_context(
         self,
@@ -617,13 +593,15 @@ class DocentTracer:
         Get the API headers for HTTP requests.
 
         Returns:
-            Dictionary of headers including Authorization
+            Dictionary of headers including Authorization if set
         """
+        headers = {"Content-Type": "application/json"}
 
-        return {
-            "Content-Type": "application/json",
-            "Authorization": self.headers.get("Authorization", ""),
-        }
+        authorization = self.headers.get("Authorization")
+        if authorization:
+            headers["Authorization"] = authorization
+
+        return headers
 
     def _post_json(self, path: str, data: Dict[str, Any]) -> None:
         if not self._api_endpoint_base:
@@ -1157,7 +1135,10 @@ def close_tracing() -> None:
 def flush_tracing() -> None:
     """Force flush all spans to exporters."""
     if _global_tracer:
+        logger.debug("Flushing global tracer")
         _global_tracer.flush()
+    else:
+        logger.debug("No global tracer available to flush")
 
 
 def verify_initialized() -> bool:

{docent_python-0.1.10a0 → docent_python-0.1.12a0}/pyproject.toml

@@ -1,7 +1,7 @@
 [project]
 name = "docent-python"
 description = "Docent SDK"
-version = "0.1.10-alpha"
+version = "0.1.12-alpha"
 authors = [
   { name="Transluce", email="info@transluce.org" },
 ]

docent_python-0.1.10a0/docent/data_models/citation.py

@@ -1,223 +0,0 @@
-import re
-from typing import TypedDict
-
-
-class Citation(TypedDict):
-    start_idx: int
-    end_idx: int
-    agent_run_idx: int | None
-    transcript_idx: int | None
-    block_idx: int
-    action_unit_idx: int | None
-
-
-def parse_citations_single_run(text: str) -> list[Citation]:
-    """
-    Parse citations from text in the format described by SINGLE_BLOCK_CITE_INSTRUCTION.
-
-    Supported formats:
-    - Single block: [T<key>B<idx>]
-    - Multiple blocks: [T<key1>B<idx1>, T<key2>B<idx2>, ...]
-    - Dash-separated blocks: [T<key1>B<idx1>-T<key2>B<idx2>]
-
-    Args:
-        text: The text to parse citations from
-
-    Returns:
-        A list of Citation objects with start_idx and end_idx representing
-        the character positions in the text (excluding brackets)
-    """
-    citations: list[Citation] = []
-
-    # Find all bracketed content first
-    bracket_pattern = r"\[(.*?)\]"
-    bracket_matches = re.finditer(bracket_pattern, text)
-
-    for bracket_match in bracket_matches:
-        bracket_content = bracket_match.group(1)
-        # Starting position of the bracket content (excluding '[')
-        content_start_pos = bracket_match.start() + 1
-
-        # Split by commas if present
-        parts = [part.strip() for part in bracket_content.split(",")]
-
-        for part in parts:
-            # Check if this part contains a dash (range citation)
-            if "-" in part:
-                # Split by dash and process each sub-part
-                dash_parts = [dash_part.strip() for dash_part in part.split("-")]
-                for dash_part in dash_parts:
-                    # Check for single block citation: T<key>B<idx>
-                    single_match = re.match(r"T(\d+)B(\d+)", dash_part)
-                    if single_match:
-                        transcript_idx = int(single_match.group(1))
-                        block_idx = int(single_match.group(2))
-
-                        # Find position within the original text
-                        citation_text = f"T{transcript_idx}B{block_idx}"
-                        part_pos_in_content = bracket_content.find(dash_part)
-                        ref_pos = content_start_pos + part_pos_in_content
-                        ref_end = ref_pos + len(citation_text)
-
-                        # Check if this citation overlaps with any existing citation
-                        if not any(
-                            citation["start_idx"] <= ref_pos < citation["end_idx"]
-                            or citation["start_idx"] < ref_end <= citation["end_idx"]
-                            for citation in citations
-                        ):
-                            citations.append(
-                                Citation(
-                                    start_idx=ref_pos,
-                                    end_idx=ref_end,
-                                    agent_run_idx=None,
-                                    transcript_idx=transcript_idx,
-                                    block_idx=block_idx,
-                                    action_unit_idx=None,
-                                )
-                            )
-            else:
-                # Check for single block citation: T<key>B<idx>
-                single_match = re.match(r"T(\d+)B(\d+)", part)
-                if single_match:
-                    transcript_idx = int(single_match.group(1))
-                    block_idx = int(single_match.group(2))
-
-                    # Find position within the original text
-                    citation_text = f"T{transcript_idx}B{block_idx}"
-                    part_pos_in_content = bracket_content.find(part)
-                    ref_pos = content_start_pos + part_pos_in_content
-                    ref_end = ref_pos + len(citation_text)
-
-                    # Check if this citation overlaps with any existing citation
-                    if not any(
-                        citation["start_idx"] <= ref_pos < citation["end_idx"]
-                        or citation["start_idx"] < ref_end <= citation["end_idx"]
-                        for citation in citations
-                    ):
-                        citations.append(
-                            Citation(
-                                start_idx=ref_pos,
-                                end_idx=ref_end,
-                                agent_run_idx=None,
-                                transcript_idx=transcript_idx,
-                                block_idx=block_idx,
-                                action_unit_idx=None,
-                            )
-                        )
-
-    return citations
-
-
-def parse_citations_multi_run(text: str) -> list[Citation]:
-    """
-    Parse citations from text in the format described by MULTI_BLOCK_CITE_INSTRUCTION.
-
-    Supported formats:
-    - Single block in transcript: [R<idx>T<key>B<idx>] or ([R<idx>T<key>B<idx>])
-    - Multiple blocks: [R<idx1>T<key1>B<idx1>][R<idx2>T<key2>B<idx2>]
-    - Comma-separated blocks: [R<idx1>T<key1>B<idx1>, R<idx2>T<key2>B<idx2>, ...]
-    - Dash-separated blocks: [R<idx1>T<key1>B<idx1>-R<idx2>T<key2>B<idx2>]
-
-    Args:
-        text: The text to parse citations from
-
-    Returns:
-        A list of Citation objects with start_idx and end_idx representing
-        the character positions in the text (excluding brackets)
-    """
-    citations: list[Citation] = []
-
-    # Find all content within brackets - this handles nested brackets too
-    bracket_pattern = r"\[([^\[\]]*(?:\[[^\[\]]*\][^\[\]]*)*)\]"
-    # Also handle optional parentheses around the brackets
-    paren_bracket_pattern = r"\(\[([^\[\]]*(?:\[[^\[\]]*\][^\[\]]*)*)\]\)"
-
-    # Single citation pattern
-    single_pattern = r"R(\d+)T(\d+)B(\d+)"
-
-    # Find all bracket matches
-    for pattern in [bracket_pattern, paren_bracket_pattern]:
-        matches = re.finditer(pattern, text)
-        for match in matches:
-            # Get the content inside brackets
-            if pattern == bracket_pattern:
-                content = match.group(1)
-                start_pos = match.start() + 1  # +1 to skip the opening bracket
-            else:
-                content = match.group(1)
-                start_pos = match.start() + 2  # +2 to skip the opening parenthesis and bracket
-
-            # Split by comma if present
-            items = [item.strip() for item in content.split(",")]
-
-            for item in items:
-                # Check if this item contains a dash (range citation)
-                if "-" in item:
-                    # Split by dash and process each sub-item
-                    dash_items = [dash_item.strip() for dash_item in item.split("-")]
-                    for dash_item in dash_items:
-                        # Check for single citation
-                        single_match = re.match(single_pattern, dash_item)
-                        if single_match:
-                            agent_run_idx = int(single_match.group(1))
-                            transcript_idx = int(single_match.group(2))
-                            block_idx = int(single_match.group(3))
-
-                            # Calculate position in the original text
-                            citation_text = f"R{agent_run_idx}T{transcript_idx}B{block_idx}"
-                            citation_start = text.find(citation_text, start_pos)
-                            citation_end = citation_start + len(citation_text)
-
-                            # Move start_pos for the next item if there are more items
-                            start_pos = citation_end
-
-                            # Avoid duplicate citations
-                            if not any(
-                                citation["start_idx"] == citation_start
-                                and citation["end_idx"] == citation_end
-                                for citation in citations
-                            ):
-                                citations.append(
-                                    Citation(
-                                        start_idx=citation_start,
-                                        end_idx=citation_end,
-                                        agent_run_idx=agent_run_idx,
-                                        transcript_idx=transcript_idx,
-                                        block_idx=block_idx,
-                                        action_unit_idx=None,
-                                    )
-                                )
-                else:
-                    # Check for single citation
-                    single_match = re.match(single_pattern, item)
-                    if single_match:
-                        agent_run_idx = int(single_match.group(1))
-                        transcript_idx = int(single_match.group(2))
-                        block_idx = int(single_match.group(3))
-
-                        # Calculate position in the original text
-                        citation_text = f"R{agent_run_idx}T{transcript_idx}B{block_idx}"
-                        citation_start = text.find(citation_text, start_pos)
-                        citation_end = citation_start + len(citation_text)
-
-                        # Move start_pos for the next item if there are more items
-                        start_pos = citation_end
-
-                        # Avoid duplicate citations
-                        if not any(
-                            citation["start_idx"] == citation_start
-                            and citation["end_idx"] == citation_end
-                            for citation in citations
-                        ):
-                            citations.append(
-                                Citation(
-                                    start_idx=citation_start,
-                                    end_idx=citation_end,
-                                    agent_run_idx=agent_run_idx,
-                                    transcript_idx=transcript_idx,
-                                    block_idx=block_idx,
-                                    action_unit_idx=None,
-                                )
-                            )
-
-    return citations