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

docent/sdk/client.py

@@ -8,7 +8,8 @@ from tqdm import tqdm
 
 from docent._log_util.logger import get_logger
 from docent.data_models.agent_run import AgentRun
-from docent.data_models.judge import JudgeRunLabel
+from docent.data_models.judge import Label
+from docent.judges.util.meta_schema import validate_judge_result_schema
 from docent.loaders import load_inspect
 
 logger = get_logger(__name__)
@@ -50,9 +51,15 @@ class Docent:
         self._login(api_key)
 
     def _handle_response_errors(self, response: requests.Response):
-        """Handle API response and raise informative errors.
-        TODO: make this more informative."""
-        response.raise_for_status()
+        """Handle API response and raise informative errors."""
+        if response.status_code >= 400:
+            try:
+                error_data = response.json()
+                detail = error_data.get("detail", response.text)
+            except Exception:
+                detail = response.text
+
+            raise requests.HTTPError(f"HTTP {response.status_code}: {detail}", response=response)
 
     def _login(self, api_key: str):
         """Login with email/password to establish session."""
@@ -182,21 +189,24 @@ class Docent:
         self._handle_response_errors(response)
         return response.json()
 
-    def get_rubric_run_state(self, collection_id: str, rubric_id: str) -> dict[str, Any]:
+    def get_rubric_run_state(
+        self, collection_id: str, rubric_id: str, version: int | None = None
+    ) -> dict[str, Any]:
         """Get rubric run state for a given collection and rubric.
 
         Args:
             collection_id: ID of the Collection.
             rubric_id: The ID of the rubric to get run state for.
+            version: The version of the rubric to get run state for. If None, the latest version is used.
 
         Returns:
-            dict: Dictionary containing rubric run state with results, job_id, and total_agent_runs.
+            dict: Dictionary containing rubric run state with results, job_id, and total_results_needed.
 
         Raises:
             requests.exceptions.HTTPError: If the API request fails.
         """
         url = f"{self._server_url}/rubric/{collection_id}/{rubric_id}/rubric_run_state"
-        response = self._session.get(url)
+        response = self._session.get(url, params={"version": version})
         self._handle_response_errors(response)
         return response.json()
 
@@ -250,30 +260,59 @@ class Docent:
         clustering_state = self.get_clustering_state(collection_id, rubric_id)
         return clustering_state.get("assignments", {})
 
+    def create_label_set(
+        self,
+        collection_id: str,
+        name: str,
+        label_schema: dict[str, Any],
+        description: str | None = None,
+    ) -> str:
+        """Create a new label set with a JSON schema.
+
+        Args:
+            collection_id: ID of the collection.
+            name: Name of the label set.
+            label_schema: JSON schema for validating labels in this set.
+            description: Optional description of the label set.
+
+        Returns:
+            str: The ID of the created label set.
+
+        Raises:
+            ValueError: If the response is missing the label_set_id.
+            jsonschema.ValidationError: If the label schema is invalid.
+            requests.exceptions.HTTPError: If the API request fails.
+        """
+        validate_judge_result_schema(label_schema)
+
+        url = f"{self._server_url}/label/{collection_id}/label_set"
+        payload = {
+            "name": name,
+            "label_schema": label_schema,
+            "description": description,
+        }
+        response = self._session.post(url, json=payload)
+        self._handle_response_errors(response)
+        return response.json()["label_set_id"]
+
     def add_label(
         self,
         collection_id: str,
-        rubric_id: str,
-        label: JudgeRunLabel,
-    ) -> dict[str, Any]:
-        """Attach a manual label to an agent run for a rubric.
+        label: Label,
+    ) -> dict[str, str]:
+        """Create a label in a label set.
 
         Args:
-            collection_id: ID of the Collection that owns the rubric.
-            rubric_id: ID of the rubric the label applies to.
-            label: A `JudgeRunLabel` that must comply with the rubric's output schema.
+            collection_id: ID of the Collection.
+            label: A `Label` object that must comply with the label set's schema.
 
         Returns:
-            dict: API response containing a status message.
+            dict: API response containing the label_id.
 
         Raises:
-            ValueError: If the label does not target the rubric specified in the path.
             requests.exceptions.HTTPError: If the API request fails or validation errors occur.
         """
-        if label.rubric_id != rubric_id:
-            raise ValueError("Label rubric_id must match the rubric_id argument")
-
-        url = f"{self._server_url}/rubric/{collection_id}/rubric/{rubric_id}/label"
+        url = f"{self._server_url}/label/{collection_id}/label"
         payload = {"label": label.model_dump(mode="json")}
         response = self._session.post(url, json=payload)
         self._handle_response_errors(response)
@@ -282,55 +321,50 @@ class Docent:
     def add_labels(
         self,
         collection_id: str,
-        rubric_id: str,
-        labels: list[JudgeRunLabel],
+        labels: list[Label],
     ) -> dict[str, Any]:
-        """Attach multiple manual labels to a rubric.
+        """Create multiple labels.
 
         Args:
-            collection_id: ID of the Collection that owns the rubric.
-            rubric_id: ID of the rubric the labels apply to.
-            labels: List of `JudgeRunLabel` objects.
+            collection_id: ID of the Collection.
+            labels: List of `Label` objects.
 
         Returns:
-            dict: API response containing status information.
+            dict: API response containing label_ids list and optional errors list.
 
         Raises:
             ValueError: If no labels are provided.
-            ValueError: If any label targets a different rubric.
             requests.exceptions.HTTPError: If the API request fails.
         """
         if not labels:
             raise ValueError("labels must contain at least one entry")
 
-        rubric_ids = {label.rubric_id for label in labels}
-        if rubric_ids != {rubric_id}:
-            raise ValueError(
-                "All labels must specify the same rubric_id that is provided to add_labels"
-            )
-
-        payload = {"labels": [l.model_dump(mode="json") for l in labels]}
-
-        url = f"{self._server_url}/rubric/{collection_id}/rubric/{rubric_id}/labels"
+        url = f"{self._server_url}/label/{collection_id}/labels"
+        payload = {"labels": [label.model_dump(mode="json") for label in labels]}
         response = self._session.post(url, json=payload)
         self._handle_response_errors(response)
         return response.json()
 
-    def get_labels(self, collection_id: str, rubric_id: str) -> list[dict[str, Any]]:
-        """Retrieve all manual labels for a rubric.
+    def get_labels(
+        self, collection_id: str, label_set_id: str, filter_valid_labels: bool = False
+    ) -> list[dict[str, Any]]:
+        """Retrieve all labels in a label set.
 
         Args:
-            collection_id: ID of the Collection that owns the rubric.
-            rubric_id: ID of the rubric to fetch labels for.
+            collection_id: ID of the Collection.
+            label_set_id: ID of the label set to fetch labels for.
+            filter_valid_labels: If True, only return labels that match the label set schema
+                INCLUDING requirements. Default is False (returns all labels).
 
         Returns:
-            list: List of label dictionaries. Each includes agent_run_id and label content.
+            list: List of label dictionaries.
 
         Raises:
             requests.exceptions.HTTPError: If the API request fails.
         """
-        url = f"{self._server_url}/rubric/{collection_id}/rubric/{rubric_id}/labels"
-        response = self._session.get(url)
+        url = f"{self._server_url}/label/{collection_id}/label_set/{label_set_id}/labels"
+        params = {"filter_valid_labels": filter_valid_labels}
+        response = self._session.get(url, params=params)
         self._handle_response_errors(response)
         return response.json()
 
@@ -357,6 +391,24 @@ class Docent:
             # TODO(mengk): kinda hacky
             return AgentRun.model_validate(response.json())
 
+    def get_chat_sessions(self, collection_id: str, agent_run_id: str) -> list[dict[str, Any]]:
+        """Get all chat sessions for an agent run, excluding judge result sessions.
+
+        Args:
+            collection_id: ID of the Collection.
+            agent_run_id: The ID of the agent run to retrieve chat sessions for.
+
+        Returns:
+            list: List of chat session dictionaries.
+
+        Raises:
+            requests.exceptions.HTTPError: If the API request fails.
+        """
+        url = f"{self._server_url}/chat/{collection_id}/{agent_run_id}/sessions"
+        response = self._session.get(url)
+        self._handle_response_errors(response)
+        return response.json()
+
     def make_collection_public(self, collection_id: str) -> dict[str, Any]:
         """Make a collection publicly accessible to anyone with the link.
 
@@ -398,6 +450,91 @@ class Docent:
         logger.info(f"Successfully shared Collection '{collection_id}' with {email}")
         return response.json()
 
+    def get_dql_schema(self, collection_id: str) -> dict[str, Any]:
+        """Retrieve the DQL schema for a collection.
+
+        Args:
+            collection_id: ID of the Collection.
+
+        Returns:
+            dict: Dictionary containing available tables, columns, and metadata for DQL queries.
+
+        Raises:
+            requests.exceptions.HTTPError: If the API request fails.
+        """
+        url = f"{self._server_url}/dql/{collection_id}/schema"
+        response = self._session.get(url)
+        self._handle_response_errors(response)
+        return response.json()
+
+    def execute_dql(self, collection_id: str, dql: str) -> dict[str, Any]:
+        """Execute a DQL query against a collection.
+
+        Args:
+            collection_id: ID of the Collection.
+            dql: The DQL query string to execute.
+
+        Returns:
+            dict: Query execution results including rows, columns, execution metadata, and selected columns.
+
+        Raises:
+            ValueError: If `dql` is empty.
+            requests.exceptions.HTTPError: If the API request fails or the query is invalid.
+        """
+        if not dql.strip():
+            raise ValueError("dql must be a non-empty string")
+
+        url = f"{self._server_url}/dql/{collection_id}/execute"
+        response = self._session.post(url, json={"dql": dql})
+        self._handle_response_errors(response)
+        return response.json()
+
+    def select_agent_run_ids(
+        self,
+        collection_id: str,
+        where_clause: str | None = None,
+        limit: int | None = None,
+    ) -> list[str]:
+        """Convenience helper to fetch agent run IDs via DQL.
+
+        Args:
+            collection_id: ID of the Collection to query.
+            where_clause: Optional DQL WHERE clause applied to the agent_runs table.
+            limit: Optional LIMIT applied to the underlying DQL query.
+
+        Returns:
+            list[str]: Agent run IDs matching the criteria.
+
+        Raises:
+            ValueError: If the inputs are invalid.
+            requests.exceptions.HTTPError: If the API request fails.
+        """
+        query = "SELECT agent_runs.id AS agent_run_id FROM agent_runs"
+
+        if where_clause:
+            where_clause = where_clause.strip()
+            if not where_clause:
+                raise ValueError("where_clause must be a non-empty string when provided")
+            query += f" WHERE {where_clause}"
+
+        if limit is not None:
+            if limit <= 0:
+                raise ValueError("limit must be a positive integer when provided")
+            query += f" LIMIT {limit}"
+
+        result = self.execute_dql(collection_id, query)
+        rows = result.get("rows", [])
+        agent_run_ids = [str(row[0]) for row in rows if row]
+
+        if result.get("truncated"):
+            logger.warning(
+                "DQL query truncated at applied limit %s; returning %s agent run IDs",
+                result.get("applied_limit"),
+                len(agent_run_ids),
+            )
+
+        return agent_run_ids
+
     def list_agent_run_ids(self, collection_id: str) -> list[str]:
         """Get all agent run IDs for a collection.