docent-python 0.1.3a0__py3-none-any.whl → 0.1.5a0__py3-none-any.whl

docent/data_models/transcript.py

@@ -11,7 +11,6 @@ from docent.data_models._tiktoken_util import (
     truncate_to_token_limit,
 )
 from docent.data_models.chat import AssistantMessage, ChatMessage, ContentReasoning
-from docent.data_models.metadata import BaseMetadata
 
 # Template for formatting individual transcript blocks
 TRANSCRIPT_BLOCK_TEMPLATE = """
@@ -63,6 +62,53 @@ def format_chat_message(
     )
 
 
+class TranscriptGroup(BaseModel):
+    """Represents a group of transcripts that are logically related.
+
+    A transcript group can contain multiple transcripts and can have a hierarchical
+    structure with parent groups. This is useful for organizing transcripts into
+    logical units like experiments, tasks, or sessions.
+
+    Attributes:
+        id: Unique identifier for the transcript group, auto-generated by default.
+        name: Optional human-readable name for the transcript group.
+        description: Optional description of the transcript group.
+        parent_transcript_group_id: Optional ID of the parent transcript group.
+        metadata: Additional structured metadata about the transcript group.
+    """
+
+    id: str = Field(default_factory=lambda: str(uuid4()))
+    name: str | None = None
+    description: str | None = None
+    parent_transcript_group_id: str | 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:
+        if v is not None and not isinstance(v, dict):
+            raise ValueError(f"metadata must be a dictionary, got {type(v).__name__}")
+        return v  # type: ignore
+
+
+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)
+    """
+
+    class _FakeModel(BaseModel):
+        data: dict[str, Any]
+
+    return _FakeModel(data=obj).model_dump()["data"]
+
+
 class Transcript(BaseModel):
     """Represents a transcript of messages in a conversation with an AI agent.
 
@@ -74,6 +120,7 @@ class Transcript(BaseModel):
         id: Unique identifier for the transcript, auto-generated by default.
         name: Optional human-readable name for the transcript.
         description: Optional description of the transcript.
+        transcript_group_id: Optional ID of the transcript group this transcript belongs to.
         messages: List of chat messages in the transcript.
         metadata: Additional structured metadata about the transcript.
     """
@@ -81,27 +128,25 @@ class Transcript(BaseModel):
     id: str = Field(default_factory=lambda: str(uuid4()))
     name: str | None = None
     description: str | None = None
+    transcript_group_id: str | None = None
 
     messages: list[ChatMessage]
-    metadata: BaseMetadata = Field(default_factory=BaseMetadata)
-
+    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: BaseMetadata, _info: Any) -> dict[str, Any]:
+    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 metadata.model_dump(strip_internal_fields=False)
+        return fake_model_dump(metadata)
 
     @field_validator("metadata", mode="before")
     @classmethod
     def _validate_metadata_type(cls, v: Any) -> Any:
-        if v is not None and not isinstance(v, BaseMetadata):
-            raise ValueError(
-                f"metadata must be an instance of BaseMetadata, got {type(v).__name__}"
-            )
-        return v
+        if v is not None and not isinstance(v, dict):
+            raise ValueError(f"metadata must be a dict, got {type(v).__name__}")
+        return v  # type: ignore
 
     @property
     def units_of_action(self) -> list[list[int]]:
@@ -297,12 +342,7 @@ class Transcript(BaseModel):
         blocks_str = "\n".join(au_blocks)
 
         # Gather metadata
-        metadata_obj = self.metadata.model_dump(strip_internal_fields=True)
-        # Add the field descriptions if they exist
-        metadata_obj = {
-            (f"{k} ({d})" if (d := self.metadata.get_field_description(k)) is not None else k): v
-            for k, v in metadata_obj.items()
-        }
+        metadata_obj = fake_model_dump(self.metadata)
 
         yaml_width = float("inf")
         block_str = f"<blocks>\n{blocks_str}\n</blocks>\n"

docent/loaders/load_inspect.py

@@ -1,11 +1,13 @@
+from typing import Any
+
 from inspect_ai.log import EvalLog
 from inspect_ai.scorer import CORRECT, INCORRECT, NOANSWER, PARTIAL, Score
 
-from docent.data_models import AgentRun, InspectAgentRunMetadata, Transcript
+from docent.data_models import AgentRun, Transcript
 from docent.data_models.chat import parse_chat_message
 
 
-def _normalize_inspect_score(score: Score) -> float | None:
+def _normalize_inspect_score(score: Score) -> Any:
     """
     Normalize an inspect score to a float. This implements the same logic as inspect_ai.scorer._metric.value_to_float, but fails more conspicuously.
 
@@ -16,30 +18,39 @@ def _normalize_inspect_score(score: Score) -> float | None:
         The normalized score as a float, or None if the score is not a valid value.
     """
 
-    if isinstance(score.value, int | float | bool):
-        return float(score.value)
-    elif score.value == CORRECT:
-        return 1.0
-    elif score.value == PARTIAL:
-        return 0.5
-    elif score.value == INCORRECT or score.value == NOANSWER:
-        return 0
-    elif isinstance(score.value, str):
-        value = score.value.lower()
+    def _leaf_normalize(value: int | float | bool | str | None) -> float | str | None:
+        if value is None:
+            return None
+        if isinstance(value, int | float | bool):
+            return float(value)
+        if value == CORRECT:
+            return 1.0
+        if value == PARTIAL:
+            return 0.5
+        if value in [INCORRECT, NOANSWER]:
+            return 0
+        value = str(value).lower()
         if value in ["yes", "true"]:
             return 1.0
-        elif value in ["no", "false"]:
+        if value in ["no", "false"]:
             return 0.0
-        elif value.replace(".", "").isnumeric():
+        if value.replace(".", "").isnumeric():
             return float(value)
+        return value
 
-    raise ValueError(f"Unknown score value: {score.value}")
+    if isinstance(score.value, int | float | bool | str):
+        return _leaf_normalize(score.value)
+    if isinstance(score.value, list):
+        return [_leaf_normalize(v) for v in score.value]
+    assert isinstance(score.value, dict), "Inspect score must be leaf value, list, or dict"
+    return {k: _leaf_normalize(v) for k, v in score.value.items()}
 
 
 def load_inspect_log(log: EvalLog) -> list[AgentRun]:
     if log.samples is None:
         return []
 
+    # TODO(vincent): fix this
     agent_runs: list[AgentRun] = []
 
     for s in log.samples:
@@ -51,22 +62,23 @@ def load_inspect_log(log: EvalLog) -> list[AgentRun]:
         else:
             sample_scores = {k: _normalize_inspect_score(v) for k, v in s.scores.items()}
 
-        metadata = InspectAgentRunMetadata(
-            task_id=log.eval.task,
-            sample_id=str(sample_id),
-            epoch_id=epoch_id,
-            model=log.eval.model,
-            additional_metadata=s.metadata,
-            scores=sample_scores,
+        metadata = {
+            "task_id": log.eval.task,
+            "sample_id": str(sample_id),
+            "epoch_id": epoch_id,
+            "model": log.eval.model,
+            "additional_metadata": s.metadata,
+            "scores": sample_scores,
             # Scores could have answers, explanations, and other metadata besides the values we extract
-            scoring_metadata=s.scores,
-        )
+            "scoring_metadata": s.scores,
+        }
 
         agent_runs.append(
             AgentRun(
                 transcripts={
                     "main": Transcript(
-                        messages=[parse_chat_message(m.model_dump()) for m in s.messages]
+                        messages=[parse_chat_message(m.model_dump()) for m in s.messages],
+                        metadata={},
                     )
                 },
                 metadata=metadata,

docent/sdk/client.py

@@ -197,75 +197,85 @@ class Docent:
         return response.json()
 
     def list_searches(self, collection_id: str) -> list[dict[str, Any]]:
-        """List all searches for a given collection.
+        """List all rubrics for a given collection.
 
         Args:
             collection_id: ID of the Collection.
 
         Returns:
-            list: List of dictionaries containing search query information.
+            list: List of dictionaries containing rubric information.
 
         Raises:
             requests.exceptions.HTTPError: If the API request fails.
         """
-        url = f"{self._server_url}/{collection_id}/list_search_queries"
+        url = f"{self._server_url}/rubric/{collection_id}/rubrics"
         response = self._session.get(url)
         response.raise_for_status()
         return response.json()
 
-    def get_search_results(self, collection_id: str, search_query: str) -> list[dict[str, Any]]:
-        """Get search results for a given collection and search query.
-        Pass in either search_query or query_id.
+    def get_search_results(
+        self, collection_id: str, rubric_id: str, rubric_version: int
+    ) -> list[dict[str, Any]]:
+        """Get rubric results for a given collection, rubric and version.
 
         Args:
             collection_id: ID of the Collection.
-            search_query: The search query to get results for.
+            rubric_id: The ID of the rubric to get results for.
+            rubric_version: The version of the rubric to get results for.
 
         Returns:
-            list: List of dictionaries containing search result information.
+            list: List of dictionaries containing rubric result information.
 
         Raises:
             requests.exceptions.HTTPError: If the API request fails.
         """
-        url = f"{self._server_url}/{collection_id}/get_search_results"
-        response = self._session.post(url, json={"search_query": search_query})
+        url = f"{self._server_url}/rubric/{collection_id}/{rubric_id}/results"
+        response = self._session.get(url, params={"rubric_version": rubric_version})
         response.raise_for_status()
         return response.json()
 
-    def list_search_clusters(self, collection_id: str, search_query: str) -> list[dict[str, Any]]:
-        """List all search clusters for a given collection.
-        Pass in either search_query or query_id.
+    def list_search_clusters(
+        self, collection_id: str, rubric_id: str, rubric_version: int | None = None
+    ) -> list[dict[str, Any]]:
+        """List all centroids for a given collection and rubric.
 
         Args:
             collection_id: ID of the Collection.
-            search_query: The search query to get clusters for.
+            rubric_id: The ID of the rubric to get centroids for.
+            rubric_version: Optional version of the rubric. If not provided, uses latest.
 
         Returns:
-            list: List of dictionaries containing search cluster information.
+            list: List of dictionaries containing centroid information.
 
         Raises:
             requests.exceptions.HTTPError: If the API request fails.
         """
-        url = f"{self._server_url}/{collection_id}/list_search_clusters"
-        response = self._session.post(url, json={"search_query": search_query})
+        url = f"{self._server_url}/rubric/{collection_id}/{rubric_id}/centroids"
+        params: dict[str, int] = {}
+        if rubric_version is not None:
+            params["rubric_version"] = rubric_version
+        response = self._session.get(url, params=params)
         response.raise_for_status()
         return response.json()
 
-    def get_cluster_matches(self, collection_id: str, centroid: str) -> list[dict[str, Any]]:
-        """Get the matches for a given cluster.
+    def get_cluster_matches(
+        self, collection_id: str, rubric_id: str, rubric_version: int
+    ) -> list[dict[str, Any]]:
+        """Get centroid assignments for a given rubric.
 
         Args:
             collection_id: ID of the Collection.
-            cluster_id: The ID of the cluster to get matches for.
+            rubric_id: The ID of the rubric to get assignments for.
+            rubric_version: The version of the rubric to get assignments for.
 
         Returns:
-            list: List of dictionaries containing the search results that match the cluster.
+            list: List of dictionaries containing centroid assignment information.
 
         Raises:
             requests.exceptions.HTTPError: If the API request fails.
         """
-        url = f"{self._server_url}/{collection_id}/get_cluster_matches"
-        response = self._session.post(url, json={"centroid": centroid})
+        url = f"{self._server_url}/rubric/{collection_id}/{rubric_id}/assignments"
+        response = self._session.get(url, params={"rubric_version": rubric_version})
         response.raise_for_status()
         return response.json()