docent-python 0.1.18a0__py3-none-any.whl → 0.1.20a0__py3-none-any.whl

docent/sdk/client.py

@@ -8,6 +8,7 @@ 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.loaders import load_inspect
 
 logger = get_logger(__name__)
@@ -48,13 +49,18 @@ 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()
+
     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()
+        self._handle_response_errors(response)
 
         logger.info("Logged in with API key")
         return
@@ -90,7 +96,7 @@ class Docent:
         }
 
         response = self._session.post(url, json=payload)
-        response.raise_for_status()
+        self._handle_response_errors(response)
 
         response_data = response.json()
         collection_id = response_data.get("collection_id")
@@ -134,13 +140,13 @@ class Docent:
                 payload = {"agent_runs": [ar.model_dump(mode="json") for ar in batch]}
 
                 response = self._session.post(url, json=payload)
-                response.raise_for_status()
+                self._handle_response_errors(response)
 
                 pbar.update(len(batch))
 
         url = f"{self._server_url}/{collection_id}/compute_embeddings"
         response = self._session.post(url)
-        response.raise_for_status()
+        self._handle_response_errors(response)
 
         logger.info(f"Successfully added {total_runs} agent runs to Collection '{collection_id}'")
         return {"status": "success", "total_runs_added": total_runs}
@@ -156,7 +162,7 @@ class Docent:
         """
         url = f"{self._server_url}/collections"
         response = self._session.get(url)
-        response.raise_for_status()
+        self._handle_response_errors(response)
         return response.json()
 
     def list_rubrics(self, collection_id: str) -> list[dict[str, Any]]:
@@ -173,15 +179,18 @@ class Docent:
         """
         url = f"{self._server_url}/rubric/{collection_id}/rubrics"
         response = self._session.get(url)
-        response.raise_for_status()
+        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.
@@ -190,8 +199,8 @@ class Docent:
             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.raise_for_status()
+        response = self._session.get(url, params={"version": version})
+        self._handle_response_errors(response)
         return response.json()
 
     def get_clustering_state(self, collection_id: str, rubric_id: str) -> dict[str, Any]:
@@ -209,7 +218,7 @@ class Docent:
         """
         url = f"{self._server_url}/rubric/{collection_id}/{rubric_id}/clustering_job"
         response = self._session.get(url)
-        response.raise_for_status()
+        self._handle_response_errors(response)
         return response.json()
 
     def get_cluster_centroids(self, collection_id: str, rubric_id: str) -> list[dict[str, Any]]:
@@ -244,6 +253,90 @@ class Docent:
         clustering_state = self.get_clustering_state(collection_id, rubric_id)
         return clustering_state.get("assignments", {})
 
+    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.
+
+        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.
+
+        Returns:
+            dict: API response containing a status message.
+
+        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"
+        payload = {"label": label.model_dump(mode="json")}
+        response = self._session.post(url, json=payload)
+        self._handle_response_errors(response)
+        return response.json()
+
+    def add_labels(
+        self,
+        collection_id: str,
+        rubric_id: str,
+        labels: list[JudgeRunLabel],
+    ) -> dict[str, Any]:
+        """Attach multiple manual labels to a rubric.
+
+        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.
+
+        Returns:
+            dict: API response containing status information.
+
+        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"
+        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.
+
+        Args:
+            collection_id: ID of the Collection that owns the rubric.
+            rubric_id: ID of the rubric to fetch labels for.
+
+        Returns:
+            list: List of label dictionaries. Each includes agent_run_id and label content.
+
+        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)
+        self._handle_response_errors(response)
+        return response.json()
+
     def get_agent_run(self, collection_id: str, agent_run_id: str) -> AgentRun | None:
         """Get a specific agent run by its ID.
 
@@ -259,7 +352,7 @@ class Docent:
         """
         url = f"{self._server_url}/{collection_id}/agent_run"
         response = self._session.get(url, params={"agent_run_id": agent_run_id})
-        response.raise_for_status()
+        self._handle_response_errors(response)
         if response.json() is None:
             return None
         else:
@@ -281,7 +374,7 @@ class Docent:
         """
         url = f"{self._server_url}/{collection_id}/make_public"
         response = self._session.post(url)
-        response.raise_for_status()
+        self._handle_response_errors(response)
 
         logger.info(f"Successfully made Collection '{collection_id}' public")
         return response.json()
@@ -303,13 +396,7 @@ class Docent:
         payload = {"email": email}
         response = self._session.post(url, json=payload)
 
-        try:
-            response.raise_for_status()
-        except requests.exceptions.HTTPError:
-            if response.status_code == 404:
-                raise ValueError(f"The user you are trying to share with ({email}) does not exist.")
-            else:
-                raise  # Re-raise the original exception
+        self._handle_response_errors(response)
 
         logger.info(f"Successfully shared Collection '{collection_id}' with {email}")
         return response.json()
@@ -328,7 +415,7 @@ class Docent:
         """
         url = f"{self._server_url}/{collection_id}/agent_run_ids"
         response = self._session.get(url)
-        response.raise_for_status()
+        self._handle_response_errors(response)
         return response.json()
 
     def recursively_ingest_inspect_logs(self, collection_id: str, fpath: str):
@@ -393,7 +480,7 @@ class Docent:
                         payload = {"agent_runs": [ar.model_dump(mode="json") for ar in batch_list]}
 
                         response = self._session.post(url, json=payload)
-                        response.raise_for_status()
+                        self._handle_response_errors(response)
 
                         runs_from_file += len(batch_list)
                         file_pbar.update(len(batch_list))
@@ -406,7 +493,7 @@ class Docent:
             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()
+            self._handle_response_errors(response)
 
         logger.info(
             f"Successfully ingested {total_runs_added} total agent runs from {len(eval_files)} files"

docent/trace.py

@@ -21,7 +21,7 @@ from opentelemetry.exporter.otlp.proto.grpc.trace_exporter import OTLPSpanExport
 from opentelemetry.exporter.otlp.proto.http.trace_exporter import OTLPSpanExporter as HTTPExporter
 from opentelemetry.instrumentation.threading import ThreadingInstrumentor
 from opentelemetry.sdk.resources import Resource
-from opentelemetry.sdk.trace import ReadableSpan, SpanProcessor, TracerProvider
+from opentelemetry.sdk.trace import ReadableSpan, SpanLimits, SpanProcessor, TracerProvider
 from opentelemetry.sdk.trace.export import (
     BatchSpanProcessor,
     ConsoleSpanExporter,
@@ -29,20 +29,13 @@ from opentelemetry.sdk.trace.export import (
 )
 from opentelemetry.trace import Span
 
-# Configure logging
 logger = logging.getLogger(__name__)
-logger.setLevel(logging.ERROR)
 
 # Default configuration
 DEFAULT_ENDPOINT = "https://api.docent.transluce.org/rest/telemetry"
 DEFAULT_COLLECTION_NAME = "default-collection-name"
 
 
-def _is_tracing_disabled() -> bool:
-    """Check if tracing is disabled via environment variable."""
-    return os.environ.get("DOCENT_DISABLE_TRACING", "").lower() == "true"
-
-
 class Instruments(Enum):
     """Enumeration of available instrument types."""
 
@@ -52,16 +45,10 @@ class Instruments(Enum):
     LANGCHAIN = "langchain"
 
 
-def _is_notebook() -> bool:
-    """Check if we're running in a Jupyter notebook."""
-    try:
-        return "ipykernel" in sys.modules
-    except Exception:
-        return False
-
-
 class DocentTracer:
-    """Manages Docent tracing setup and provides tracing utilities."""
+    """
+    Manages Docent tracing setup and provides tracing utilities.
+    """
 
     def __init__(
         self,
@@ -77,22 +64,6 @@ class DocentTracer:
         instruments: Optional[Set[Instruments]] = None,
         block_instruments: Optional[Set[Instruments]] = None,
     ):
-        """
-        Initialize Docent tracing manager.
-
-        Args:
-            collection_name: Name of the collection for resource attributes
-            collection_id: Optional collection ID (auto-generated if not provided)
-            agent_run_id: Optional agent_run_id to use for code outside of an agent run context (auto-generated if not provided)
-            endpoint: OTLP endpoint URL(s) - can be a single string or list of strings for multiple endpoints
-            headers: Optional headers for authentication
-            api_key: Optional API key for bearer token authentication (takes precedence over env var)
-            enable_console_export: Whether to export to console
-            enable_otlp_export: Whether to export to OTLP endpoint
-            disable_batch: Whether to disable batch processing (use SimpleSpanProcessor)
-            instruments: Set of instruments to enable (None = all instruments)
-            block_instruments: Set of instruments to explicitly disable
-        """
         self._initialized: bool = False
         # Check if tracing is disabled via environment variable
         if _is_tracing_disabled():
@@ -163,8 +134,12 @@ class DocentTracer:
         """
         Get the current agent run ID from context.
 
+        Retrieves the agent run ID that was set in the current execution context.
+        If no agent run context is active, returns the default agent run ID.
+
         Returns:
-            The current agent run ID if available, None otherwise
+            The current agent run ID if available, or the default agent run ID
+            if no context is active.
         """
         try:
             return self._agent_run_id_var.get()
@@ -249,12 +224,23 @@ class DocentTracer:
             return
 
         try:
+
+            # Check for OTEL_SPAN_ATTRIBUTE_COUNT_LIMIT environment variable
+            default_attribute_limit = 1024
+            env_value = os.environ.get("OTEL_SPAN_ATTRIBUTE_COUNT_LIMIT", "0")
+            env_limit = int(env_value) if env_value.isdigit() else 0
+            attribute_limit = max(env_limit, default_attribute_limit)
+
+            span_limits = SpanLimits(
+                max_attributes=attribute_limit,
+            )
+
             # Create our own isolated tracer provider
             self._tracer_provider = TracerProvider(
-                resource=Resource.create({"service.name": self.collection_name})
+                resource=Resource.create({"service.name": self.collection_name}),
+                span_limits=span_limits,
             )
 
-            # Add custom span processor for agent_run_id and transcript_id
             class ContextSpanProcessor(SpanProcessor):
                 def __init__(self, manager: "DocentTracer"):
                     self.manager: "DocentTracer" = manager
@@ -312,11 +298,7 @@ class DocentTracer:
                     )
 
                 def on_end(self, span: ReadableSpan) -> None:
-                    # Debug logging for span completion
-                    span_attrs = span.attributes or {}
-                    logger.debug(
-                        f"Completed span: name='{span.name}', collection_id={span_attrs.get('collection_id')}, agent_run_id={span_attrs.get('agent_run_id')}, transcript_id={span_attrs.get('transcript_id')}, duration_ns={span.end_time - span.start_time if span.end_time and span.start_time else 'unknown'}"
-                    )
+                    pass
 
                 def shutdown(self) -> None:
                     pass
@@ -422,7 +404,17 @@ class DocentTracer:
             raise
 
     def cleanup(self):
-        """Clean up Docent tracing resources and signal trace completion to backend."""
+        """
+        Clean up Docent tracing resources.
+
+        Flushes all pending spans to exporters and shuts down the tracer provider.
+        This method is automatically called during application shutdown via atexit
+        handlers, but can also be called manually for explicit cleanup.
+
+        The cleanup process:
+        1. Flushes all span processors to ensure data is exported
+        2. Shuts down the tracer provider and releases resources
+        """
         if self._disabled:
             return
 
@@ -473,7 +465,7 @@ class DocentTracer:
         if disabled and self._initialized:
             self.cleanup()
 
-    def verify_initialized(self) -> bool:
+    def is_initialized(self) -> bool:
         """Verify if the manager is properly initialized."""
         return self._initialized
 
@@ -1063,8 +1055,9 @@ def initialize_tracing(
         collection_id: Optional collection ID (auto-generated if not provided)
         endpoint: OTLP endpoint URL(s) for span export - can be a single string or list of strings for multiple endpoints
         headers: Optional headers for authentication
-        api_key: Optional API key for bearer token authentication (takes precedence over env var)
-        enable_console_export: Whether to export spans to console
+        api_key: Optional API key for bearer token authentication (takes precedence
+                over DOCENT_API_KEY environment variable)
+        enable_console_export: Whether to export spans to console for debugging
         enable_otlp_export: Whether to export spans to OTLP endpoint
         disable_batch: Whether to disable batch processing (use SimpleSpanProcessor)
         instruments: Set of instruments to enable (None = all instruments).
@@ -1074,7 +1067,6 @@ def initialize_tracing(
         The initialized Docent tracer
 
     Example:
-        # Basic setup
         initialize_tracing("my-collection")
     """
 
@@ -1137,17 +1129,17 @@ def close_tracing() -> None:
 def flush_tracing() -> None:
     """Force flush all spans to exporters."""
     if _global_tracer:
-        logger.debug("Flushing global tracer")
+        logger.debug("Flushing Docent tracer")
         _global_tracer.flush()
     else:
         logger.debug("No global tracer available to flush")
 
 
-def verify_initialized() -> bool:
+def is_initialized() -> bool:
     """Verify if the global Docent tracer is properly initialized."""
     if _global_tracer is None:
         return False
-    return _global_tracer.verify_initialized()
+    return _global_tracer.is_initialized()
 
 
 def is_disabled() -> bool:
@@ -1764,3 +1756,16 @@ def transcript_group_context(
     return TranscriptGroupContext(
         name, transcript_group_id, description, metadata, parent_transcript_group_id
     )
+
+
+def _is_tracing_disabled() -> bool:
+    """Check if tracing is disabled via environment variable."""
+    return os.environ.get("DOCENT_DISABLE_TRACING", "").lower() == "true"
+
+
+def _is_notebook() -> bool:
+    """Check if we're running in a Jupyter notebook."""
+    try:
+        return "ipykernel" in sys.modules
+    except Exception:
+        return False