docent-python 0.1.13a0__py3-none-any.whl → 0.1.15a0__py3-none-any.whl

docent/__init__.py

@@ -1,4 +1,4 @@
 __all__ = ["Docent", "init"]
 
-from docent.agent_run_writer import init
+from docent.sdk.agent_run_writer import init
 from docent.sdk.client import Docent

docent/data_models/agent_run.py

@@ -1,5 +1,7 @@
-import json
 import sys
+import textwrap
+from datetime import datetime
+from queue import Queue
 from typing import Any, Literal, TypedDict, cast
 from uuid import uuid4
 
@@ -7,18 +9,18 @@ import yaml
 from pydantic import (
     BaseModel,
     Field,
-    field_serializer,
+    PrivateAttr,
     field_validator,
     model_validator,
 )
+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.transcript import (
-    Transcript,
-    TranscriptGroup,
-    TranscriptWithoutMetadataValidator,
-    fake_model_dump,
-)
+from docent.data_models.transcript import Transcript, TranscriptGroup
+from docent.data_models.yaml_util import yaml_dump_metadata
+
+logger = get_logger(__name__)
 
 
 class FilterableField(TypedDict):
@@ -36,8 +38,8 @@ class AgentRun(BaseModel):
         id: Unique identifier for the agent run, auto-generated by default.
         name: Optional human-readable name for the agent run.
         description: Optional description of the agent run.
-        transcripts: Dict mapping transcript IDs to Transcript objects.
-        transcript_groups: Dict mapping transcript group IDs to TranscriptGroup objects.
+        transcripts: List of Transcript objects.
+        transcript_groups: List of TranscriptGroup objects.
         metadata: Additional structured metadata about the agent run as a JSON-serializable dictionary.
     """
 
@@ -45,36 +47,31 @@ class AgentRun(BaseModel):
     name: str | None = None
     description: str | None = None
 
-    transcripts: dict[str, Transcript]
-    transcript_groups: dict[str, TranscriptGroup] = Field(default_factory=dict)
+    transcripts: list[Transcript]
+    transcript_groups: list[TranscriptGroup] = Field(default_factory=list)
     metadata: dict[str, Any] = Field(default_factory=dict)
 
-    @field_serializer("metadata")
-    def serialize_metadata(self, metadata: dict[str, Any], _info: Any) -> dict[str, Any]:
-        """
-        Custom serializer for the metadata field - returns the dict as-is since it's already serializable.
-        """
-        return fake_model_dump(metadata)
-
-    @field_validator("metadata", mode="before")
+    @field_validator("transcripts", mode="before")
     @classmethod
-    def _validate_metadata_json_serializable(cls, v: Any) -> dict[str, Any]:
-        """
-        Validates that metadata is a dictionary and is JSON-serializable.
-        """
-        if v is None:
-            return {}
-
-        if not isinstance(v, dict):
-            raise ValueError(f"metadata must be a dictionary, got {type(v).__name__}")
-
-        # Check that the metadata is JSON serializable
-        try:
-            json.dumps(fake_model_dump(cast(dict[str, Any], v)))
-        except (TypeError, ValueError) as e:
-            raise ValueError(f"metadata must be JSON-serializable: {e}")
+    def _validate_transcripts_type(cls, v: Any) -> Any:
+        if isinstance(v, dict):
+            logger.warning(
+                "dict[str, Transcript] for transcripts is deprecated. Use list[Transcript] instead."
+            )
+            v = cast(dict[str, Transcript], v)
+            return [Transcript.model_validate(t) for t in v.values()]
+        return v
 
-        return cast(dict[str, Any], v)
+    @field_validator("transcript_groups", mode="before")
+    @classmethod
+    def _validate_transcript_groups_type(cls, v: Any) -> Any:
+        if isinstance(v, dict):
+            logger.warning(
+                "dict[str, TranscriptGroup] for transcript_groups is deprecated. Use list[TranscriptGroup] instead."
+            )
+            v = cast(dict[str, TranscriptGroup], v)
+            return [TranscriptGroup.model_validate(tg) for tg in v.values()]
+        return v
 
     @model_validator(mode="after")
     def _validate_transcripts_not_empty(self):
@@ -90,6 +87,52 @@ class AgentRun(BaseModel):
             raise ValueError("AgentRun must have at least one transcript")
         return self
 
+    def get_filterable_fields(self, max_depth: int = 1) -> list[FilterableField]:
+        """Returns a list of all fields that can be used to filter the agent run,
+        by recursively exploring the model_dump() for singleton types in dictionaries.
+
+        Returns:
+            list[FilterableField]: A list of filterable fields, where each field is a
+                                   dictionary containing its 'name' (path) and 'type'.
+        """
+
+        result: list[FilterableField] = []
+
+        def _explore_dict(d: dict[str, Any], prefix: str, depth: int):
+            nonlocal result
+
+            if depth > max_depth:
+                return
+
+            for k, v in d.items():
+                if isinstance(v, (str, int, float, bool)):
+                    result.append(
+                        {
+                            "name": f"{prefix}.{k}",
+                            "type": cast(Literal["str", "bool", "int", "float"], type(v).__name__),
+                        }
+                    )
+                elif isinstance(v, dict):
+                    _explore_dict(cast(dict[str, Any], v), f"{prefix}.{k}", depth + 1)
+
+        # Look at the agent run metadata
+        _explore_dict(to_jsonable_python(self.metadata), "metadata", 0)
+        # Look at the transcript metadata
+        # TODO(mengk): restore this later when we have the ability to integrate with SQL.
+        # for t_id, t in self.transcripts.items():
+        #     _explore_dict(
+        #         t.metadata.model_dump(strip_internal_fields=True), f"transcript.{t_id}.metadata", 0
+        #     )
+
+        # Append the text field
+        result.append({"name": "text", "type": "str"})
+
+        return result
+
+    ######################
+    # Converting to text #
+    ######################
+
     def _to_text_impl(self, token_limit: int = sys.maxsize, use_blocks: bool = False) -> list[str]:
         """
         Core implementation for converting agent run to text representation.
@@ -103,7 +146,7 @@ class AgentRun(BaseModel):
         """
         # Generate transcript strings using appropriate method
         transcript_strs: list[str] = []
-        for i, (t_key, t) in enumerate(self.transcripts.items()):
+        for i, t in enumerate(self.transcripts):
             if use_blocks:
                 transcript_content = t.to_str_blocks_with_token_limit(
                     token_limit=sys.maxsize,
@@ -116,14 +159,12 @@ class AgentRun(BaseModel):
                     transcript_idx=i,
                     agent_run_idx=None,
                 )[0]
-            transcript_strs.append(
-                f"<transcript {t_key}>\n{transcript_content}\n</transcript {t_key}>"
-            )
+            transcript_strs.append(f"<transcript>\n{transcript_content}\n</transcript>")
 
         transcripts_str = "\n\n".join(transcript_strs)
 
         # Gather metadata
-        metadata_obj = fake_model_dump(self.metadata)
+        metadata_obj = to_jsonable_python(self.metadata)
         if self.name is not None:
             metadata_obj["name"] = self.name
         if self.description is not None:
@@ -164,12 +205,12 @@ class AgentRun(BaseModel):
                     assert (
                         msg_range.end == msg_range.start + 1
                     ), "Ranges without metadata should be a single message"
-                    t_id, t = list(self.transcripts.items())[msg_range.start]
+                    t = self.transcripts[msg_range.start]
                     if msg_range.num_tokens < token_limit - 50:
                         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}>"
+                            transcript = f"<transcript>\n{t.to_str_blocks_with_token_limit(token_limit=sys.maxsize)[0]}\n</transcript>"
                         else:
-                            transcript = f"<transcript {t_id}>\n{t.to_str_with_token_limit(token_limit=sys.maxsize)[0]}\n</transcript {t_id}>"
+                            transcript = f"<transcript>\n{t.to_str_with_token_limit(token_limit=sys.maxsize)[0]}\n</transcript>"
                         result = (
                             f"Here is a partial agent run for analysis purposes only:\n{transcript}"
                         )
@@ -184,7 +225,7 @@ class AgentRun(BaseModel):
                                 token_limit=token_limit - 50,
                             )
                         for fragment in transcript_fragments:
-                            result = f"<transcript {t_id}>\n{fragment}\n</transcript {t_id}>"
+                            result = f"<transcript>\n{fragment}\n</transcript>"
                             result = (
                                 f"Here is a partial agent run for analysis purposes only:\n{result}"
                             )
@@ -229,71 +270,199 @@ class AgentRun(BaseModel):
         """
         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.
+    ##############################
+    # New text rendering methods #
+    ##############################
+
+    # Transcript ID -> Transcript
+    _transcript_dict: dict[str, Transcript] | None = PrivateAttr(default=None)
+    # Transcript Group ID -> Transcript Group
+    _transcript_group_dict: dict[str, TranscriptGroup] | None = PrivateAttr(default=None)
+    # Canonical tree cache keyed by full_tree flag
+    _canonical_tree_cache: dict[bool, dict[str | None, list[tuple[Literal["t", "tg"], str]]]] = (
+        PrivateAttr(default_factory=dict)
+    )
+    # Transcript IDs (depth-first) cache keyed by full_tree flag
+    _transcript_ids_ordered_cache: dict[bool, list[str]] = PrivateAttr(default_factory=dict)
+
+    @property
+    def transcript_dict(self) -> dict[str, Transcript]:
+        """Lazily compute and cache a mapping from transcript ID to Transcript."""
+        if self._transcript_dict is None:
+            self._transcript_dict = {t.id: t for t in self.transcripts}
+        return self._transcript_dict
+
+    @property
+    def transcript_group_dict(self) -> dict[str, TranscriptGroup]:
+        """Lazily compute and cache a mapping from transcript group ID to TranscriptGroup."""
+        if self._transcript_group_dict is None:
+            self._transcript_group_dict = {tg.id: tg for tg in self.transcript_groups}
+        return self._transcript_group_dict
+
+    def get_canonical_tree(
+        self, full_tree: bool = False
+    ) -> dict[str | None, list[tuple[Literal["t", "tg"], str]]]:
+        """Compute and cache the canonical, sorted transcript group tree.
 
         Args:
-            *args: Variable length argument list passed to parent method.
-            **kwargs: Arbitrary keyword arguments passed to parent method.
+            full_tree: If True, include all transcript groups regardless of whether
+                they contain transcripts. If False, include only the minimal tree
+                that connects relevant groups and transcripts.
 
         Returns:
-            dict[str, Any]: Dictionary representation of the model including the text property.
+            Canonical tree mapping parent group id (or "__global_root") to a list of
+            children (type, id) tuples sorted by creation time.
         """
-        return super().model_dump(*args, **kwargs) | {"text": self.text}
+        if (
+            full_tree not in self._canonical_tree_cache
+            or full_tree not in self._transcript_ids_ordered_cache
+        ):
+            canonical_tree, transcript_idx_map = self._build_canonical_tree(full_tree=full_tree)
+            self._canonical_tree_cache[full_tree] = canonical_tree
+            self._transcript_ids_ordered_cache[full_tree] = list(transcript_idx_map.keys())
+        return self._canonical_tree_cache[full_tree]
+
+    def get_transcript_ids_ordered(self, full_tree: bool = False) -> list[str]:
+        """Compute and cache the depth-first transcript id ordering.
 
-    def get_filterable_fields(self, max_depth: int = 1) -> list[FilterableField]:
-        """Returns a list of all fields that can be used to filter the agent run,
-        by recursively exploring the model_dump() for singleton types in dictionaries.
+        Args:
+            full_tree: Whether to compute based on the full tree or the minimal tree.
 
         Returns:
-            list[FilterableField]: A list of filterable fields, where each field is a
-                                   dictionary containing its 'name' (path) and 'type'.
+            List of transcript ids in depth-first order.
         """
-
-        result: list[FilterableField] = []
-
-        def _explore_dict(d: dict[str, Any], prefix: str, depth: int):
-            nonlocal result
-
-            if depth > max_depth:
-                return
-
-            for k, v in d.items():
-                if isinstance(v, (str, int, float, bool)):
-                    result.append(
-                        {
-                            "name": f"{prefix}.{k}",
-                            "type": cast(Literal["str", "bool", "int", "float"], type(v).__name__),
-                        }
+        if (
+            full_tree not in self._transcript_ids_ordered_cache
+            or full_tree not in self._canonical_tree_cache
+        ):
+            canonical_tree, transcript_idx_map = self._build_canonical_tree(full_tree=full_tree)
+            self._canonical_tree_cache[full_tree] = canonical_tree
+            self._transcript_ids_ordered_cache[full_tree] = list(transcript_idx_map.keys())
+        return self._transcript_ids_ordered_cache[full_tree]
+
+    def _build_canonical_tree(self, full_tree: bool = False):
+        t_dict = self.transcript_dict
+        tg_dict = self.transcript_group_dict
+
+        # Find all transcript groups that have direct transcript children
+        # Also keep track of transcripts that are not in a group
+        tgs_to_transcripts: dict[str, set[str]] = {}
+        for transcript in t_dict.values():
+            if transcript.transcript_group_id is None:
+                tgs_to_transcripts.setdefault("__global_root", set()).add(transcript.id)
+            else:
+                tgs_to_transcripts.setdefault(transcript.transcript_group_id, set()).add(
+                    transcript.id
+                )
+
+        # tg_tree maps from parent -> children. A child can be a group or a transcript.
+        #   A parent must be a group (or None, for transcripts that are not in a group).
+        tg_tree: dict[str, set[tuple[Literal["t", "tg"], str]]] = {}
+
+        if full_tree:
+            for tg_id, tg in tg_dict.items():
+                tg_tree.setdefault(tg.parent_transcript_group_id or "__global_root", set()).add(
+                    ("tg", tg_id)
+                )
+                for t_id in tgs_to_transcripts.get(tg_id, []):
+                    tg_tree.setdefault(tg_id, set()).add(("t", t_id))
+            for t_id, t in t_dict.items():
+                tg_tree.setdefault(t.transcript_group_id or "__global_root", set()).add(("t", t_id))
+        else:
+            # Initialize q with "important" tgs
+            q, seen = Queue[str](), set[str]()
+            for tg_id in tgs_to_transcripts.keys():
+                q.put(tg_id)
+                seen.add(tg_id)
+
+            # Do an "upwards BFS" from leaves up to the root. Builds a tree of only relevant nodes.
+            while q.qsize() > 0:
+                u_id = q.get()
+                u = tg_dict.get(u_id)  # None if __global_root
+
+                # Add the transcripts under this tg
+                for t_id in tgs_to_transcripts.get(u_id, []):
+                    tg_tree.setdefault(u_id, set()).add(("t", t_id))
+
+                # Add an edge from the parent
+                if u is not None:
+                    par_id = u.parent_transcript_group_id or "__global_root"
+                    # Mark u as a child of par
+                    tg_tree.setdefault(par_id, set()).add(("tg", u_id))
+                    # If we haven't investigated the parent before, add to q
+                    if par_id not in seen:
+                        q.put(par_id)
+                        seen.add(par_id)
+
+        # For each node, sort by created_at timestamp
+
+        def _cmp(element: tuple[Literal["t", "tg"], str]) -> datetime:
+            obj_type, obj_id = element
+            if obj_type == "tg":
+                return tg_dict[obj_id].created_at or datetime.max
+            else:
+                return t_dict[obj_id].created_at or datetime.max
+
+        c_tree: dict[str | None, list[tuple[Literal["t", "tg"], str]]] = {}
+        for tg_id in tg_tree:
+            children_ids = list(set(tg_tree[tg_id]))
+            sorted_children_ids = sorted(children_ids, key=_cmp)
+            c_tree[tg_id] = sorted_children_ids
+
+        # Compute transcript indices as the depth-first traversal index
+        transcript_idx_map: dict[str, int] = {}
+
+        def _assign_transcript_indices(cur_tg_id: str, next_idx: int) -> int:
+            children = c_tree.get(cur_tg_id, [])
+            for child_type, child_id in children:
+                if child_type == "tg":
+                    next_idx = _assign_transcript_indices(child_id, next_idx)
+                else:
+                    transcript_idx_map[child_id] = next_idx
+                    next_idx += 1
+            return next_idx
+
+        _assign_transcript_indices("__global_root", 0)
+
+        return c_tree, transcript_idx_map
+
+    def to_text_new(self, indent: int = 0, full_tree: bool = False):
+        c_tree = self.get_canonical_tree(full_tree=full_tree)
+        t_ids_ordered = self.get_transcript_ids_ordered(full_tree=full_tree)
+        t_idx_map = {t_id: i for i, t_id in enumerate(t_ids_ordered)}
+        t_dict = self.transcript_dict
+        tg_dict = self.transcript_group_dict
+
+        # Traverse the tree and render the string
+        def _recurse(tg_id: str) -> str:
+            children_ids = c_tree.get(tg_id, [])
+            children_texts: list[str] = []
+            for child_type, child_id in children_ids:
+                if child_type == "tg":
+                    children_texts.append(_recurse(child_id))
+                else:
+                    cur_text = t_dict[child_id].to_text_new(
+                        transcript_idx=t_idx_map[child_id],
+                        indent=indent,
                     )
-                elif isinstance(v, dict):
-                    _explore_dict(cast(dict[str, Any], v), f"{prefix}.{k}", depth + 1)
-
-        # Look at the agent run metadata
-        _explore_dict(fake_model_dump(self.metadata), "metadata", 0)
-        # Look at the transcript metadata
-        # TODO(mengk): restore this later when we have the ability to integrate with SQL.
-        # for t_id, t in self.transcripts.items():
-        #     _explore_dict(
-        #         t.metadata.model_dump(strip_internal_fields=True), f"transcript.{t_id}.metadata", 0
-        #     )
-
-        # Append the text field
-        result.append({"name": "text", "type": "str"})
-
-        return result
+                    children_texts.append(cur_text)
+            children_text = "\n".join(children_texts)
 
+            # No wrapper for global root
+            if tg_id == "__global_root":
+                return children_text
+            # Delegate rendering to TranscriptGroup
+            else:
+                tg = tg_dict[tg_id]
+                return tg.to_text_new(children_text=children_text, indent=indent)
 
-class AgentRunWithoutMetadataValidator(AgentRun):
-    """
-    A version of AgentRun that doesn't have the model_validator on metadata.
-    Needed for sending/receiving agent runs via JSON, since they incorrectly trip the existing model_validator.
-    """
+        text = _recurse("__global_root")
 
-    transcripts: dict[str, TranscriptWithoutMetadataValidator]  # type: ignore
+        # Append agent run metadata below the full content
+        yaml_text = yaml_dump_metadata(self.metadata)
+        if yaml_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|>"
 
-    @field_validator("metadata", mode="before")
-    @classmethod
-    def _validate_metadata_type(cls, v: Any) -> Any:
-        # Bypass the model_validator
-        return v
+        return text

docent/data_models/remove_invalid_citation_ranges.py

@@ -66,16 +66,13 @@ def get_transcript_text_for_citation(agent_run: AgentRun, citation: Citation) ->
         return None
 
     try:
-        transcript_keys = list(agent_run.transcripts.keys())
-        if citation.transcript_idx >= len(transcript_keys):
+        if citation.transcript_idx >= len(agent_run.get_transcript_ids_ordered()):
             return None
+        transcript_id = agent_run.get_transcript_ids_ordered()[citation.transcript_idx]
+        transcript = agent_run.transcript_dict[transcript_id]
 
-        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

docent/data_models/transcript.py

@@ -1,10 +1,12 @@
 import sys
+import textwrap
 from datetime import datetime
 from typing import Any
 from uuid import uuid4
 
 import yaml
-from pydantic import BaseModel, Field, PrivateAttr, field_serializer, field_validator
+from pydantic import BaseModel, Field, PrivateAttr, field_validator
+from pydantic_core import to_jsonable_python
 
 from docent.data_models._tiktoken_util import (
     get_token_count,
@@ -13,12 +15,13 @@ 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
 
 # Template for formatting individual transcript blocks
 TRANSCRIPT_BLOCK_TEMPLATE = """
-<{index_label} | role: {role}>
+<|{index_label}; role: {role}|>
 {content}
-</{index_label}>
+</|{index_label}; role: {role}|>
 """.strip()
 
 # Instructions for citing single transcript blocks
@@ -35,7 +38,7 @@ Important notes:
 - 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."""
+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."""
 
 
 def format_chat_message(
@@ -94,19 +97,11 @@ class TranscriptGroup(BaseModel):
     id: str = Field(default_factory=lambda: str(uuid4()))
     name: str | None = None
     description: str | None = None
-    collection_id: str
     agent_run_id: str
     parent_transcript_group_id: str | None = None
     created_at: datetime | None = None
     metadata: dict[str, Any] = Field(default_factory=dict)
 
-    @field_serializer("metadata")
-    def serialize_metadata(self, metadata: dict[str, Any], _info: Any) -> dict[str, Any]:
-        """
-        Custom serializer for the metadata field so the internal fields are explicitly preserved.
-        """
-        return fake_model_dump(metadata)
-
     @field_validator("metadata", mode="before")
     @classmethod
     def _validate_metadata_type(cls, v: Any) -> Any:
@@ -114,16 +109,33 @@ class TranscriptGroup(BaseModel):
             raise ValueError(f"metadata must be a dictionary, got {type(v).__name__}")
         return v  # type: ignore
 
+    def to_text_new(self, children_text: str, indent: int = 0) -> str:
+        """Render this transcript group with its children and metadata.
 
-def fake_model_dump(obj: dict[str, Any]) -> dict[str, Any]:
-    """
-    Emulate the action of pydantic.model_dump() for non-pydantic objects (to handle nested values)
-    """
+        Metadata appears below the rendered children content.
 
-    class _FakeModel(BaseModel):
-        data: dict[str, Any]
+        Args:
+            children_text: Pre-rendered text of this group's children (groups/transcripts).
+            indent: Number of spaces to indent the rendered output.
+
+        Returns:
+            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:
+            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|>"
+            )
+        else:
+            inner = children_text
 
-    return _FakeModel(data=obj).model_dump()["data"]
+        # Compose final text: content first, then metadata, all inside the group wrapper
+        if indent > 0:
+            inner = textwrap.indent(inner, " " * indent)
+        return f"<|{self.name}|>\n{inner}\n</|{self.name}|>"
 
 
 class Transcript(BaseModel):
@@ -152,13 +164,6 @@ class Transcript(BaseModel):
     metadata: dict[str, Any] = Field(default_factory=dict)
     _units_of_action: list[list[int]] | None = PrivateAttr(default=None)
 
-    @field_serializer("metadata")
-    def serialize_metadata(self, metadata: dict[str, Any], _info: Any) -> dict[str, Any]:
-        """
-        Custom serializer for the metadata field so the internal fields are explicitly preserved.
-        """
-        return fake_model_dump(metadata)
-
     @field_validator("metadata", mode="before")
     @classmethod
     def _validate_metadata_type(cls, v: Any) -> Any:
@@ -400,7 +405,7 @@ class Transcript(BaseModel):
         blocks_str = "\n".join(blocks)
 
         # Gather metadata
-        metadata_obj = fake_model_dump(self.metadata)
+        metadata_obj = to_jsonable_python(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>"
@@ -484,15 +489,32 @@ class Transcript(BaseModel):
             use_action_units=False,
         )
 
+    ##############################
+    # New text rendering methods #
+    ##############################
 
-class TranscriptWithoutMetadataValidator(Transcript):
-    """
-    A version of Transcript that doesn't have the model_validator on metadata.
-    Needed for sending/receiving transcripts via JSON, since they incorrectly trip the existing model_validator.
-    """
+    def to_text_new(self, transcript_idx: int = 0, indent: int = 0) -> str:
+        # Format individual message blocks
+        blocks: list[str] = []
+        for msg_idx, message in enumerate(self.messages):
+            block_text = format_chat_message(message, msg_idx, transcript_idx)
+            blocks.append(block_text)
+        blocks_str = "\n".join(blocks)
+        if indent > 0:
+            blocks_str = textwrap.indent(blocks_str, " " * indent)
+
+        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:
+            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|>"
+            )
 
-    @field_validator("metadata", mode="before")
-    @classmethod
-    def _validate_metadata_type(cls, v: Any) -> Any:
-        # Bypass the model_validator
-        return v
+        # Format content and return
+        if indent > 0:
+            content_str = textwrap.indent(content_str, " " * indent)
+        return f"<|T{transcript_idx}|>\n{content_str}\n</|T{transcript_idx}|>\n"

docent/data_models/yaml_util.py

@@ -0,0 +1,12 @@
+from typing import Any
+
+import yaml
+from pydantic_core import to_jsonable_python
+
+
+def yaml_dump_metadata(metadata: dict[str, Any]) -> str | None:
+    if not metadata:
+        return None
+    metadata_obj = to_jsonable_python(metadata)
+    yaml_text = yaml.dump(metadata_obj, width=float("inf"))
+    return yaml_text.strip()

docent/loaders/load_inspect.py

@@ -6,9 +6,12 @@ from zipfile import ZipFile
 from inspect_ai.log import EvalLog
 from inspect_ai.scorer import CORRECT, INCORRECT, NOANSWER, PARTIAL, Score
 
+from docent._log_util.logger import get_logger
 from docent.data_models import AgentRun, Transcript
 from docent.data_models.chat import parse_chat_message
 
+logger = get_logger(__name__)
+
 
 def _normalize_inspect_score(score: Score | dict[str, Any]) -> Any:
     """
@@ -83,12 +86,12 @@ def load_inspect_log(log: EvalLog) -> list[AgentRun]:
 
         agent_runs.append(
             AgentRun(
-                transcripts={
-                    "main": Transcript(
+                transcripts=[
+                    Transcript(
                         messages=[parse_chat_message(m.model_dump()) for m in s.messages],
                         metadata={},
                     )
-                },
+                ],
                 metadata=metadata,
             )
         )
@@ -120,11 +123,9 @@ def _read_sample_as_run(data: dict[str, Any], header_metadata: dict[str, Any] =
     }
 
     run = AgentRun(
-        transcripts={
-            "main": Transcript(
-                messages=[parse_chat_message(m) for m in data["messages"]], metadata={}
-            ),
-        },
+        transcripts=[
+            Transcript(messages=[parse_chat_message(m) for m in data["messages"]], metadata={})
+        ],
         metadata=run_metadata,
     )
     return run
@@ -166,8 +167,12 @@ def _runs_from_eval_file(
     file: BinaryIO,
 ) -> Tuple[dict[str, Any], Generator[AgentRun, None, None]]:
     zip = ZipFile(file, mode="r")
-    header: dict[str, Any] = json.load(zip.open("header.json", "r"))
-    header_metadata = _run_metadata_from_header(header)
+    try:
+        header: dict[str, Any] = json.load(zip.open("header.json", "r"))
+        header_metadata = _run_metadata_from_header(header)
+    except KeyError:
+        logger.warning(f"No header found in {file.name} file")
+        header_metadata = {}
 
     def _iter_runs() -> Generator[AgentRun, None, None]:
         try:

docent/sdk/client.py

@@ -1,10 +1,14 @@
+import itertools
 import os
+from pathlib import Path
 from typing import Any
 
 import requests
+from tqdm import tqdm
 
 from docent._log_util.logger import get_logger
-from docent.data_models.agent_run import AgentRun, AgentRunWithoutMetadataValidator
+from docent.data_models.agent_run import AgentRun
+from docent.loaders import load_inspect
 
 logger = get_logger(__name__)
 
@@ -100,49 +104,9 @@ class Docent:
         )
         return collection_id
 
-    def set_io_bin_keys(
-        self, collection_id: str, inner_bin_key: str | None, outer_bin_key: str | None
-    ):
-        """Set inner and outer bin keys for a collection."""
-        response = self._session.post(
-            f"{self._server_url}/{collection_id}/set_io_bin_keys",
-            json={"inner_bin_key": inner_bin_key, "outer_bin_key": outer_bin_key},
-        )
-        response.raise_for_status()
-
-    def set_inner_bin_key(self, collection_id: str, dim: str):
-        """Set the inner bin key for a collection."""
-        current_io_bin_keys = self.get_io_bin_keys(collection_id)
-        if current_io_bin_keys is None:
-            current_io_bin_keys = (None, None)
-        self.set_io_bin_keys(collection_id, dim, current_io_bin_keys[1])  # Set inner, keep outer
-
-    def set_outer_bin_key(self, collection_id: str, dim: str):
-        """Set the outer bin key for a collection."""
-        current_io_bin_keys = self.get_io_bin_keys(collection_id)
-        if current_io_bin_keys is None:
-            current_io_bin_keys = (None, None)
-        self.set_io_bin_keys(collection_id, current_io_bin_keys[0], dim)  # Keep inner, set outer
-
-    def get_io_bin_keys(self, collection_id: str) -> tuple[str | None, str | None] | None:
-        """Gets the current inner and outer bin keys for a Collection.
-
-        Args:
-            collection_id: ID of the Collection.
-
-        Returns:
-            tuple: (inner_bin_key | None, outer_bin_key | None)
-
-        Raises:
-            requests.exceptions.HTTPError: If the API request fails.
-        """
-        url = f"{self._server_url}/{collection_id}/io_bin_keys"
-        response = self._session.get(url)
-        response.raise_for_status()
-        data = response.json()
-        return (data.get("inner_bin_key"), data.get("outer_bin_key"))
-
-    def add_agent_runs(self, collection_id: str, agent_runs: list[AgentRun]) -> dict[str, Any]:
+    def add_agent_runs(
+        self, collection_id: str, agent_runs: list[AgentRun], batch_size: int = 1000
+    ) -> dict[str, Any]:
         """Adds agent runs to a Collection.
 
         Agent runs represent execution traces that can be visualized and analyzed.
@@ -161,7 +125,6 @@ class Docent:
         from tqdm import tqdm
 
         url = f"{self._server_url}/{collection_id}/agent_runs"
-        batch_size = 1000
         total_runs = len(agent_runs)
 
         # Process agent runs in batches
@@ -302,7 +265,7 @@ class Docent:
         else:
             # We do this to avoid metadata validation failing
             # TODO(mengk): kinda hacky
-            return AgentRunWithoutMetadataValidator.model_validate(response.json())
+            return AgentRun.model_validate(response.json())
 
     def make_collection_public(self, collection_id: str) -> dict[str, Any]:
         """Make a collection publicly accessible to anyone with the link.
@@ -367,3 +330,84 @@ class Docent:
         response = self._session.get(url)
         response.raise_for_status()
         return response.json()
+
+    def recursively_ingest_inspect_logs(self, collection_id: str, fpath: str):
+        """Recursively search directory for .eval files and ingest them as agent runs.
+
+        Args:
+            collection_id: ID of the Collection to add agent runs to.
+            fpath: Path to directory to search recursively.
+
+        Raises:
+            ValueError: If the path doesn't exist or isn't a directory.
+            requests.exceptions.HTTPError: If any API requests fail.
+        """
+        root_path = Path(fpath)
+        if not root_path.exists():
+            raise ValueError(f"Path does not exist: {fpath}")
+        if not root_path.is_dir():
+            raise ValueError(f"Path is not a directory: {fpath}")
+
+        # Find all .eval files recursively
+        eval_files = list(root_path.rglob("*.eval"))
+
+        if not eval_files:
+            logger.info(f"No .eval files found in {fpath}")
+            return
+
+        logger.info(f"Found {len(eval_files)} .eval files in {fpath}")
+
+        total_runs_added = 0
+        batch_size = 100
+
+        # Process each .eval file
+        for eval_file in tqdm(eval_files, desc="Processing .eval files", unit="files"):
+            # Get total samples for progress tracking
+            total_samples = load_inspect.get_total_samples(eval_file, format="eval")
+
+            if total_samples == 0:
+                logger.info(f"No samples found in {eval_file}")
+                continue
+
+            # Load runs from file
+            with open(eval_file, "rb") as f:
+                _, runs_generator = load_inspect.runs_from_file(f, format="eval")
+
+                # Process runs in batches
+                runs_from_file = 0
+                batches = itertools.batched(runs_generator, batch_size)
+
+                with tqdm(
+                    total=total_samples,
+                    desc=f"Processing {eval_file.name}",
+                    unit="runs",
+                    leave=False,
+                ) as file_pbar:
+                    for batch in batches:
+                        batch_list = list(batch)  # Convert generator batch to list
+                        if not batch_list:
+                            break
+
+                        # Add batch to collection
+                        url = f"{self._server_url}/{collection_id}/agent_runs"
+                        payload = {"agent_runs": [ar.model_dump(mode="json") for ar in batch_list]}
+
+                        response = self._session.post(url, json=payload)
+                        response.raise_for_status()
+
+                        runs_from_file += len(batch_list)
+                        file_pbar.update(len(batch_list))
+
+            total_runs_added += runs_from_file
+            logger.info(f"Added {runs_from_file} runs from {eval_file}")
+
+        # Compute embeddings after all files are processed
+        if total_runs_added > 0:
+            logger.info("Computing embeddings for added runs...")
+            url = f"{self._server_url}/{collection_id}/compute_embeddings"
+            response = self._session.post(url)
+            response.raise_for_status()
+
+        logger.info(
+            f"Successfully ingested {total_runs_added} total agent runs from {len(eval_files)} files"
+        )

docent/trace.py

@@ -197,12 +197,14 @@ class DocentTracer:
         try:
             if "http" in endpoint.lower() or "https" in endpoint.lower():
                 http_exporter: HTTPExporter = HTTPExporter(
-                    endpoint=f"{endpoint}/v1/traces", headers=self.headers
+                    endpoint=f"{endpoint}/v1/traces", headers=self.headers, timeout=30
                 )
                 logger.debug(f"Initialized HTTP exporter for endpoint: {endpoint}/v1/traces")
                 return http_exporter
             else:
-                grpc_exporter: GRPCExporter = GRPCExporter(endpoint=endpoint, headers=self.headers)
+                grpc_exporter: GRPCExporter = GRPCExporter(
+                    endpoint=endpoint, headers=self.headers, timeout=30
+                )
                 logger.debug(f"Initialized gRPC exporter for endpoint: {endpoint}")
                 return grpc_exporter
         except Exception as e:

{docent_python-0.1.13a0.dist-info → docent_python-0.1.15a0.dist-info}/METADATA

@@ -1,14 +1,16 @@
 Metadata-Version: 2.4
 Name: docent-python
-Version: 0.1.13a0
+Version: 0.1.15a0
 Summary: Docent SDK
 Project-URL: Homepage, https://github.com/TransluceAI/docent
 Project-URL: Issues, https://github.com/TransluceAI/docent/issues
 Project-URL: Docs, https://transluce-docent.readthedocs-hosted.com/en/latest
 Author-email: Transluce <info@transluce.org>
-License-Expression: MIT
+License-Expression: Apache-2.0
 License-File: LICENSE.md
 Requires-Python: >=3.11
+Requires-Dist: backoff>=2.2.1
+Requires-Dist: inspect-ai>=0.3.132
 Requires-Dist: opentelemetry-api>=1.34.1
 Requires-Dist: opentelemetry-exporter-otlp-proto-grpc>=1.34.1
 Requires-Dist: opentelemetry-exporter-otlp-proto-http>=1.34.1

{docent_python-0.1.13a0.dist-info → docent_python-0.1.15a0.dist-info}/RECORD

@@ -1,31 +1,32 @@
-docent/__init__.py,sha256=KY_gsq7iKn5tZODvwbsDFMxAZukniec5nBoYawdhglo,108
-docent/agent_run_writer.py,sha256=QNCV4m36c9BuhzWCyuzs0wH9ql8uubzcQUXMhc3XVug,9135
+docent/__init__.py,sha256=fuhETwJPcesiB76Zxa64HBJxeaaTyRalIH-fs77TWsU,112
 docent/py.typed,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
-docent/trace.py,sha256=C8oO5NuChSgyHIB5OI6qOfXWaIld7jdvlEqImK56a1E,66761
+docent/trace.py,sha256=bCO66QFgQ9L_4XM4PXnZToMi1Crtc9n0723kNjwCdm0,66823
 docent/trace_temp.py,sha256=Z0lAPwVzXjFvxpiU-CuvfWIslq9Q4alNkZMoQ77Xudk,40711
 docent/_log_util/__init__.py,sha256=3HXXrxrSm8PxwG4llotrCnSnp7GuroK1FNHsdg6f7aE,73
 docent/_log_util/logger.py,sha256=kwM0yRW1IJd6-XTorjWn48B4l8qvD2ZM6VDjY5eskQI,4422
 docent/data_models/__init__.py,sha256=4JbTDVzRhS5VZgo8MALwd_YI17GaN7X9E3rOc4Xl7kw,327
 docent/data_models/_tiktoken_util.py,sha256=hC0EDDWItv5-0cONBnHWgZtQOflDU7ZNEhXPFo4DvPc,3057
-docent/data_models/agent_run.py,sha256=AhokdyEscrlrg0q5aKaOv26cYvkA6LvAoQsz_WBg_pM,12240
+docent/data_models/agent_run.py,sha256=bsZGL0D3HIO8oxfaeUzaUWRS82u-IiO7vs0-Lv9spks,19970
 docent/data_models/citation.py,sha256=zpF9WuvVEfktltw1M9P3hwpg5yywizFUKF5zROBR2cY,5062
 docent/data_models/metadata.py,sha256=r0SYC4i2x096dXMLfw_rAMtcJQCsoV6EOMPZuEngbGA,9062
 docent/data_models/regex.py,sha256=0ciIerkrNwb91bY5mTcyO5nDWH67xx2tZYObV52fmBo,1684
-docent/data_models/remove_invalid_citation_ranges.py,sha256=0cn4Xg_tgg45nZvc-sjtqLgr1rywBBrsLJ_WBKEF0pY,5673
+docent/data_models/remove_invalid_citation_ranges.py,sha256=U-aIzRL-SuWFQZr1MqEGqXMNyIKQs7VQLxHDoFrMJwI,5658
 docent/data_models/shared_types.py,sha256=jjm-Dh5S6v7UKInW7SEqoziOsx6Z7Uu4e3VzgCbTWvc,225
-docent/data_models/transcript.py,sha256=Gmy4lYdlvC5SXzpnerFJ83lIMPPiYUPgjOUbwg6aWJQ,20238
+docent/data_models/transcript.py,sha256=xA6fcGwYn8ewgqWdIgrXcq1Qbt7rByCKqDabffvCL0A,21387
+docent/data_models/yaml_util.py,sha256=6GrPWqbTZrryZh71cnSsiqbHkWVCd-8V3-6GeiEchUg,325
 docent/data_models/chat/__init__.py,sha256=GleyRzYqKRkwwSRm_tQJw5BudCbgu9WRSa71Fntz0L0,610
 docent/data_models/chat/content.py,sha256=Co-jO8frQa_DSP11wJuhPX0s-GpJk8yqtKqPeiAIZ_U,1672
 docent/data_models/chat/message.py,sha256=xGt09keA6HRxw40xB_toNzEqA9ip7k53dnhXrEbKGO8,4157
 docent/data_models/chat/tool.py,sha256=MMglNHzkwHqUoK0xDWqs2FtelPsgHqwVpGpI1F8KZyw,3049
-docent/loaders/load_inspect.py,sha256=_cK2Qd6gyLQuJVzOlsvEZz7TrqzNmH6ZsLTkSCWAPqQ,6628
+docent/loaders/load_inspect.py,sha256=VLrtpvcVZ44n2DIPMwUivXqbvOWjaooGw6moY8UQ0VE,6789
 docent/samples/__init__.py,sha256=roDFnU6515l9Q8v17Es_SpWyY9jbm5d6X9lV01V0MZo,143
 docent/samples/load.py,sha256=ZGE07r83GBNO4A0QBh5aQ18WAu3mTWA1vxUoHd90nrM,207
 docent/samples/log.eval,sha256=orrW__9WBfANq7NwKsPSq9oTsQRcG6KohG5tMr_X_XY,397708
 docent/samples/tb_airline.json,sha256=eR2jFFRtOw06xqbEglh6-dPewjifOk-cuxJq67Dtu5I,47028
 docent/sdk/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
-docent/sdk/client.py,sha256=rvOFXvyAr9QxCijN0_CWENbm8y3YQvR1msfFSBDZvOw,13309
-docent_python-0.1.13a0.dist-info/METADATA,sha256=NOSOTU8xJ5BHt_lA1Hq749RNknOBHMC2v7cQbF5ssEs,1038
-docent_python-0.1.13a0.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
-docent_python-0.1.13a0.dist-info/licenses/LICENSE.md,sha256=vOHzq3K4Ndu0UV9hPrtXvlD7pHOjyDQmGjHuLSIkRQY,1087
-docent_python-0.1.13a0.dist-info/RECORD,,
+docent/sdk/agent_run_writer.py,sha256=QNCV4m36c9BuhzWCyuzs0wH9ql8uubzcQUXMhc3XVug,9135
+docent/sdk/client.py,sha256=fuJrTF87OtUojULFY7acZuqg5xmE8F-4HgEeEV8_gq0,14781
+docent_python-0.1.15a0.dist-info/METADATA,sha256=UP7y-vp_VE2RWb33COgh-mh0lhcJZhedCkKcsgBpUYc,1110
+docent_python-0.1.15a0.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
+docent_python-0.1.15a0.dist-info/licenses/LICENSE.md,sha256=QIMv2UiT6MppRasso4ymaA0w7ltkqmlL0HCt8CLD7Rc,580
+docent_python-0.1.15a0.dist-info/RECORD,,

docent_python-0.1.15a0.dist-info/licenses/LICENSE.md

@@ -0,0 +1,13 @@
+Copyright 2025 Clarity AI Research Inc., dba Transluce
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.

docent_python-0.1.13a0.dist-info/licenses/LICENSE.md

@@ -1,7 +0,0 @@
-Copyright 2025 Clarity AI Research, Inc. dba Transluce
-
-Permission is hereby granted, free of charge, to any person obtaining a copy of this software and associated documentation files (the “Software”), to deal in the Software without restriction, including without limitation the rights to use, copy, modify, merge, publish, distribute, sublicense, and/or sell copies of the Software, and to permit persons to whom the Software is furnished to do so, subject to the following conditions:
-
-The above copyright notice and this permission notice shall be included in all copies or substantial portions of the Software.
-
-THE SOFTWARE IS PROVIDED “AS IS”, WITHOUT WARRANTY OF ANY KIND, EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE.