docent-python 0.1.0a7__py3-none-any.whl → 0.1.0a8__py3-none-any.whl

docent/data_models/metadata.py

@@ -34,7 +34,7 @@ class BaseMetadata(BaseModel):
     """
 
     model_config = ConfigDict(extra="allow")
-    allow_fields_without_descriptions: bool = False
+    allow_fields_without_descriptions: bool = True
 
     # Private attribute to store field descriptions
     _field_descriptions: dict[str, str | None] | None = PrivateAttr(default=None)
@@ -189,3 +189,41 @@ class BaseAgentRunMetadata(BaseMetadata):
     scores: dict[str, int | float | bool | None] = Field(
         description="A dict of score_key -> score_value. Use one key for each metric you're tracking."
     )
+
+
+class InspectAgentRunMetadata(BaseAgentRunMetadata):
+    """Extends BaseAgentRunMetadata with fields specific to Inspect runs.
+
+    Attributes:
+        task_id: The ID of the 'benchmark' or 'set of evals' that the transcript belongs to
+        sample_id: The specific task inside of the `task_id` benchmark that the transcript was run on
+        epoch_id: Each `sample_id` should be run multiple times due to stochasticity; `epoch_id` is the integer index of a specific run.
+        model: The model that was used to generate the transcript
+        scoring_metadata: Additional metadata about the scoring process
+        additional_metadata: Additional metadata about the transcript
+    """
+
+    task_id: str = Field(
+        description="The ID of the 'benchmark' or 'set of evals' that the transcript belongs to"
+    )
+
+    # Identification of this particular run
+    sample_id: str = Field(
+        description="The specific task inside of the `task_id` benchmark that the transcript was run on"
+    )
+    epoch_id: int = Field(
+        description="Each `sample_id` should be run multiple times due to stochasticity; `epoch_id` is the integer index of a specific run."
+    )
+
+    # Parameters for the run
+    model: str = Field(description="The model that was used to generate the transcript")
+
+    # Outcome
+    scoring_metadata: dict[str, Any] | None = Field(
+        description="Additional metadata about the scoring process"
+    )
+
+    # Inspect metadata
+    additional_metadata: dict[str, Any] | None = Field(
+        description="Additional metadata about the transcript"
+    )

docent/data_models/transcript.py

@@ -110,6 +110,8 @@ class Transcript(BaseModel):
         A unit of action represents a logical group of messages, such as a system message
         on its own or a user message followed by assistant responses and tool outputs.
 
+        For precise details on how action units are determined, refer to the _compute_units_of_action method implementation.
+
         Returns:
             list[list[int]]: List of units of action, where each unit is a list of message indices.
         """

docent/sdk/client.py

@@ -12,7 +12,7 @@ logger = get_logger(__name__)
 class Docent:
     """Client for interacting with the Docent API.
 
-    This client provides methods for creating and managing FrameGrids,
+    This client provides methods for creating and managing Collections,
     dimensions, agent runs, and filters in the Docent system.
 
     Args:
@@ -26,70 +26,61 @@ class Docent:
         self,
         server_url: str = "https://aws-docent-backend.transluce.org",
         web_url: str = "https://docent-alpha.transluce.org",
-        email: str | None = None,
-        password: str | None = None,
+        api_key: str | None = None,
     ):
         self._server_url = server_url.rstrip("/") + "/rest"
         self._web_url = web_url.rstrip("/")
 
-        self._email = email or os.getenv("DOCENT_EMAIL")
-        if self._email is None:
-            raise ValueError(
-                "Email address must be provided through keyword argument or DOCENT_EMAIL environment variable"
-            )
-
-        self._password = password or os.getenv("DOCENT_PASSWORD")
-        if self._password is None:
-            raise ValueError(
-                "Password must be provided through keyword argument or DOCENT_PASSWORD environment variable"
-            )
-
         # Use requests.Session for connection pooling and persistent headers
         self._session = requests.Session()
-        self._login()
 
-    def _login(self):
-        """Login with email/password to establish session."""
-        login_url = f"{self._server_url}/login"
-        response = self._session.post(
-            login_url, json={"email": self._email, "password": self._password}
-        )
+        api_key = api_key or os.getenv("DOCENT_API_KEY")
 
-        if response.status_code == 401:
+        if api_key is None:
             raise ValueError(
-                "Invalid username/password combination. "
-                f"If you don't already have an account, please sign up at {self._web_url}/signup"
+                "api_key is required. Please provide an "
+                "api_key or set the DOCENT_API_KEY environment variable."
             )
 
+        self._login(api_key)
+
+    def _login(self, api_key: str):
+        """Login with email/password to establish session."""
+        self._session.headers.update({"Authorization": f"Bearer {api_key}"})
+
+        url = f"{self._server_url}/api-keys/test"
+        response = self._session.get(url)
         response.raise_for_status()
-        logger.info(f"Successfully logged in as {self._email}")
 
-    def create_framegrid(
+        logger.info("Logged in with API key")
+        return
+
+    def create_collection(
         self,
-        fg_id: str | None = None,
+        collection_id: str | None = None,
         name: str | None = None,
         description: str | None = None,
     ) -> str:
-        """Creates a new FrameGrid.
+        """Creates a new Collection.
 
-        Creates a new FrameGrid and sets up a default MECE dimension
+        Creates a new Collection and sets up a default MECE dimension
         for grouping on the homepage.
 
         Args:
-            fg_id: Optional ID for the new FrameGrid. If not provided, one will be generated.
-            name: Optional name for the FrameGrid.
-            description: Optional description for the FrameGrid.
+            collection_id: Optional ID for the new Collection. If not provided, one will be generated.
+            name: Optional name for the Collection.
+            description: Optional description for the Collection.
 
         Returns:
-            str: The ID of the created FrameGrid.
+            str: The ID of the created Collection.
 
         Raises:
-            ValueError: If the response is missing the FrameGrid ID.
+            ValueError: If the response is missing the Collection ID.
             requests.exceptions.HTTPError: If the API request fails.
         """
         url = f"{self._server_url}/create"
         payload = {
-            "fg_id": fg_id,
+            "collection_id": collection_id,
             "name": name,
             "description": description,
         }
@@ -98,42 +89,46 @@ class Docent:
         response.raise_for_status()
 
         response_data = response.json()
-        fg_id = response_data.get("fg_id")
-        if fg_id is None:
-            raise ValueError("Failed to create frame grid: 'fg_id' missing in response.")
+        collection_id = response_data.get("collection_id")
+        if collection_id is None:
+            raise ValueError("Failed to create collection: 'collection_id' missing in response.")
 
-        logger.info(f"Successfully created FrameGrid with id='{fg_id}'")
+        logger.info(f"Successfully created Collection with id='{collection_id}'")
 
-        logger.info(f"FrameGrid creation complete. Frontend available at: {self._web_url}/{fg_id}")
-        return fg_id
+        logger.info(
+            f"Collection creation complete. Frontend available at: {self._web_url}/dashboard/{collection_id}"
+        )
+        return collection_id
 
-    def set_io_bin_keys(self, fg_id: str, inner_bin_key: str | None, outer_bin_key: str | None):
-        """Set inner and outer bin keys for a frame grid."""
+    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}/{fg_id}/set_io_bin_keys",
+            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, fg_id: str, dim: str):
-        """Set the inner bin key for a frame grid."""
-        current_io_bin_keys = self.get_io_bin_keys(fg_id)
+    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(fg_id, dim, current_io_bin_keys[1])  # Set inner, keep outer
+        self.set_io_bin_keys(collection_id, dim, current_io_bin_keys[1])  # Set inner, keep outer
 
-    def set_outer_bin_key(self, fg_id: str, dim: str):
-        """Set the outer bin key for a frame grid."""
-        current_io_bin_keys = self.get_io_bin_keys(fg_id)
+    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(fg_id, current_io_bin_keys[0], dim)  # Keep inner, set outer
+        self.set_io_bin_keys(collection_id, current_io_bin_keys[0], dim)  # Keep inner, set outer
 
-    def get_io_bin_keys(self, fg_id: str) -> tuple[str | None, str | None] | None:
-        """Gets the current inner and outer bin keys for a FrameGrid.
+    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:
-            fg_id: ID of the FrameGrid.
+            collection_id: ID of the Collection.
 
         Returns:
             tuple: (inner_bin_key | None, outer_bin_key | None)
@@ -141,20 +136,20 @@ class Docent:
         Raises:
             requests.exceptions.HTTPError: If the API request fails.
         """
-        url = f"{self._server_url}/{fg_id}/io_bin_keys"
+        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, fg_id: str, agent_runs: list[AgentRun]) -> dict[str, Any]:
-        """Adds agent runs to a FrameGrid.
+    def add_agent_runs(self, collection_id: str, agent_runs: list[AgentRun]) -> dict[str, Any]:
+        """Adds agent runs to a Collection.
 
         Agent runs represent execution traces that can be visualized and analyzed.
-        This method batches the insertion in groups of 5,000 for better performance.
+        This method batches the insertion in groups of 1,000 for better performance.
 
         Args:
-            fg_id: ID of the FrameGrid.
+            collection_id: ID of the Collection.
             agent_runs: List of AgentRun objects to add.
 
         Returns:
@@ -165,7 +160,7 @@ class Docent:
         """
         from tqdm import tqdm
 
-        url = f"{self._server_url}/{fg_id}/agent_runs"
+        url = f"{self._server_url}/{collection_id}/agent_runs"
         batch_size = 1000
         total_runs = len(agent_runs)
 
@@ -180,69 +175,96 @@ class Docent:
 
                 pbar.update(len(batch))
 
-        url = f"{self._server_url}/{fg_id}/compute_embeddings"
+        url = f"{self._server_url}/{collection_id}/compute_embeddings"
         response = self._session.post(url)
         response.raise_for_status()
 
-        logger.info(f"Successfully added {total_runs} agent runs to FrameGrid '{fg_id}'")
+        logger.info(f"Successfully added {total_runs} agent runs to Collection '{collection_id}'")
         return {"status": "success", "total_runs_added": total_runs}
 
-    def list_framegrids(self) -> list[dict[str, Any]]:
-        """Lists all available FrameGrids.
+    def list_collections(self) -> list[dict[str, Any]]:
+        """Lists all available Collections.
 
         Returns:
-            list: List of dictionaries containing FrameGrid information.
+            list: List of dictionaries containing Collection information.
 
         Raises:
             requests.exceptions.HTTPError: If the API request fails.
         """
-        url = f"{self._server_url}/framegrids"
+        url = f"{self._server_url}/collections"
         response = self._session.get(url)
         response.raise_for_status()
         return response.json()
 
-    def get_dimensions(self, fg_id: str, dim_ids: list[str] | None = None) -> list[dict[str, Any]]:
-        """Retrieves dimensions for a FrameGrid.
+    def list_searches(self, collection_id: str) -> list[dict[str, Any]]:
+        """List all searches for a given collection.
 
         Args:
-            fg_id: ID of the FrameGrid.
-            dim_ids: Optional list of dimension IDs to retrieve. If None, retrieves all dimensions.
+            collection_id: ID of the Collection.
 
         Returns:
-            list: List of dictionaries containing dimension information.
+            list: List of dictionaries containing search query information.
 
         Raises:
             requests.exceptions.HTTPError: If the API request fails.
         """
-        url = f"{self._server_url}/{fg_id}/get_dimensions"
-        payload = {
-            "dim_ids": dim_ids,
-        }
-        response = self._session.post(url, json=payload)
+        url = f"{self._server_url}/{collection_id}/list_search_queries"
+        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.
+
+        Args:
+            collection_id: ID of the Collection.
+            search_query: The search query to get results for.
+
+        Returns:
+            list: List of dictionaries containing search 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})
         response.raise_for_status()
         return response.json()
 
-    def list_attribute_searches(
-        self, fg_id: str, base_data_only: bool = True
-    ) -> list[dict[str, Any]]:
-        """Lists available attribute searches for a FrameGrid.
+    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.
 
-        Attribute searches allow finding frames with specific metadata attributes.
+        Args:
+            collection_id: ID of the Collection.
+            search_query: The search query to get clusters for.
+
+        Returns:
+            list: List of dictionaries containing search cluster 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})
+        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.
 
         Args:
-            fg_id: ID of the FrameGrid.
-            base_data_only: If True, returns only basic search information.
+            collection_id: ID of the Collection.
+            cluster_id: The ID of the cluster to get matches for.
 
         Returns:
-            list: List of dictionaries containing attribute search information.
+            list: List of dictionaries containing the search results that match the cluster.
 
         Raises:
             requests.exceptions.HTTPError: If the API request fails.
         """
-        url = f"{self._server_url}/{fg_id}/attribute_searches"
-        params = {
-            "base_data_only": base_data_only,
-        }
-        response = self._session.get(url, params=params)
+        url = f"{self._server_url}/{collection_id}/get_cluster_matches"
+        response = self._session.post(url, json={"centroid": centroid})
         response.raise_for_status()
         return response.json()

{docent_python-0.1.0a7.dist-info → docent_python-0.1.0a8.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: docent-python
-Version: 0.1.0a7
+Version: 0.1.0a8
 Summary: Docent SDK
 Project-URL: Homepage, https://github.com/TransluceAI/docent
 Project-URL: Issues, https://github.com/TransluceAI/docent/issues
@@ -12,3 +12,4 @@ Requires-Python: >=3.11
 Requires-Dist: pydantic>=2.11.7
 Requires-Dist: pyyaml>=6.0.2
 Requires-Dist: tiktoken>=0.7.0
+Requires-Dist: tqdm>=4.67.1

{docent_python-0.1.0a7.dist-info → docent_python-0.1.0a8.dist-info}/RECORD

@@ -6,10 +6,10 @@ docent/data_models/__init__.py,sha256=4NghEq_T9JqGTSo-hu_aZm35TDfahTzsbkrt8WB4Go
 docent/data_models/_tiktoken_util.py,sha256=hC0EDDWItv5-0cONBnHWgZtQOflDU7ZNEhXPFo4DvPc,3057
 docent/data_models/agent_run.py,sha256=sdvoUUpOhQAHqJHNR5KoHthCXrpJajdIREMacoR1ODk,9516
 docent/data_models/citation.py,sha256=WsVQZcBT2EJD24ysyeVOC5Xfo165RI7P5_cOnJBgHj0,10015
-docent/data_models/metadata.py,sha256=pb7fYQnqvjcmx7b7V6rRltJmuSRGJ58lmnPmNgab8rA,7112
+docent/data_models/metadata.py,sha256=XVPfUPwGcs6JlJg2u9Ry0A8NdHNZTvELwmrsmb7TagE,8702
 docent/data_models/regex.py,sha256=0ciIerkrNwb91bY5mTcyO5nDWH67xx2tZYObV52fmBo,1684
 docent/data_models/shared_types.py,sha256=jjm-Dh5S6v7UKInW7SEqoziOsx6Z7Uu4e3VzgCbTWvc,225
-docent/data_models/transcript.py,sha256=HK1NvtD06XqVto6O5qzQ_wCmxlx0O88wSBB2NRTPs44,13667
+docent/data_models/transcript.py,sha256=7xYCPi6gIUftX2tjOcteryQE9GWV7ThRZv4PmpteJhM,13793
 docent/data_models/chat/__init__.py,sha256=O04XQ2NmO8GTWqkkB_Iydj8j_CucZuLhoyMVTxJN_cs,570
 docent/data_models/chat/content.py,sha256=Co-jO8frQa_DSP11wJuhPX0s-GpJk8yqtKqPeiAIZ_U,1672
 docent/data_models/chat/message.py,sha256=iAo38kbV6wYbFh8S23cxLy6HY4C_i3PzQ6RpSQG5dxM,3861
@@ -19,8 +19,8 @@ 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=ptzgZq2vbxGrwPs6eK7TOs_FaymOkOtuM5kFVGklfuA,8832
-docent_python-0.1.0a7.dist-info/METADATA,sha256=Iqt4lNy7UxZjqRvy0Yvc7Ch5NuZuD2eXAB1O1cs6F8I,493
-docent_python-0.1.0a7.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
-docent_python-0.1.0a7.dist-info/licenses/LICENSE.md,sha256=vOHzq3K4Ndu0UV9hPrtXvlD7pHOjyDQmGjHuLSIkRQY,1087
-docent_python-0.1.0a7.dist-info/RECORD,,
+docent/sdk/client.py,sha256=2fS2bmO9wOvIdjoG_2mOGyeX2xcWXm-vzVAeqIf8BZQ,9784
+docent_python-0.1.0a8.dist-info/METADATA,sha256=KtUqXmRDX0QpnxiDox667QOkwHpmu7VTh8k3nYQIMHY,521
+docent_python-0.1.0a8.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
+docent_python-0.1.0a8.dist-info/licenses/LICENSE.md,sha256=vOHzq3K4Ndu0UV9hPrtXvlD7pHOjyDQmGjHuLSIkRQY,1087
+docent_python-0.1.0a8.dist-info/RECORD,,