docent-python 0.1.17a0__py3-none-any.whl → 0.1.19a0__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,7 +179,7 @@ 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]:
@@ -191,7 +197,7 @@ class Docent:
         """
         url = f"{self._server_url}/rubric/{collection_id}/{rubric_id}/rubric_run_state"
         response = self._session.get(url)
-        response.raise_for_status()
+        self._handle_response_errors(response)
         return response.json()
 
     def get_clustering_state(self, collection_id: str, rubric_id: str) -> dict[str, Any]:
@@ -209,7 +215,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 +250,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 +349,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 +371,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 +393,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 +412,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 +477,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 +490,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

{docent_python-0.1.17a0.dist-info → docent_python-0.1.19a0.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: docent-python
-Version: 0.1.17a0
+Version: 0.1.19a0
 Summary: Docent SDK
 Project-URL: Homepage, https://github.com/TransluceAI/docent
 Project-URL: Issues, https://github.com/TransluceAI/docent/issues

{docent_python-0.1.17a0.dist-info → docent_python-0.1.19a0.dist-info}/RECORD

@@ -1,20 +1,20 @@
 docent/__init__.py,sha256=fuhETwJPcesiB76Zxa64HBJxeaaTyRalIH-fs77TWsU,112
 docent/py.typed,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
-docent/trace.py,sha256=bCO66QFgQ9L_4XM4PXnZToMi1Crtc9n0723kNjwCdm0,66823
+docent/trace.py,sha256=CEDT7StBE6DaEffWEiW-Lntx5OxNmv9vyXFGI9UlW28,66357
 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/__init__.py,sha256=bE_Wy4Ql-9-0ZPcolMCPHhYvaE_Ug6h-jV7wOJ_DAi0,399
 docent/data_models/_tiktoken_util.py,sha256=hC0EDDWItv5-0cONBnHWgZtQOflDU7ZNEhXPFo4DvPc,3057
-docent/data_models/agent_run.py,sha256=F5oGsKVpWc-IUCX3QxChkgw6UR30MGTlP4YS2KeaUkQ,19315
-docent/data_models/citation.py,sha256=zpF9WuvVEfktltw1M9P3hwpg5yywizFUKF5zROBR2cY,5062
-docent/data_models/metadata.py,sha256=r0SYC4i2x096dXMLfw_rAMtcJQCsoV6EOMPZuEngbGA,9062
+docent/data_models/agent_run.py,sha256=7_37I9aS9rhDTkAvMPwoJGssQldvvKte8qVb93EnAiY,19329
+docent/data_models/citation.py,sha256=2_M1-_olVOJtjCGGFx1GIwGYWl0ILHxRsW8-EFDS9j0,7844
+docent/data_models/judge.py,sha256=zPbTqztn-yWu6tgD3R5JTyGnNiDhY6cWQ-gz3e_eM5k,340
+docent/data_models/metadata_util.py,sha256=E-EClAP5vVm9xbfTlPSz0tUyCalOfN9Jujd6JGoRnBg,487
 docent/data_models/regex.py,sha256=0ciIerkrNwb91bY5mTcyO5nDWH67xx2tZYObV52fmBo,1684
-docent/data_models/remove_invalid_citation_ranges.py,sha256=U-aIzRL-SuWFQZr1MqEGqXMNyIKQs7VQLxHDoFrMJwI,5658
+docent/data_models/remove_invalid_citation_ranges.py,sha256=3RSMsOzFO2cSjkxI549TAo12qdvD-AGHd05Jxu0amvs,6282
 docent/data_models/shared_types.py,sha256=jjm-Dh5S6v7UKInW7SEqoziOsx6Z7Uu4e3VzgCbTWvc,225
-docent/data_models/transcript.py,sha256=eT48m8rZ3STGdElH-B1aOlCuwdaCc673GNVdQTBgAt0,19429
-docent/data_models/yaml_util.py,sha256=6GrPWqbTZrryZh71cnSsiqbHkWVCd-8V3-6GeiEchUg,325
-docent/data_models/chat/__init__.py,sha256=GleyRzYqKRkwwSRm_tQJw5BudCbgu9WRSa71Fntz0L0,610
+docent/data_models/transcript.py,sha256=7cdj2KAO_e2k3rj7OPzJzmzrkxPHIW7fbHygKTr7EZg,19940
+docent/data_models/chat/__init__.py,sha256=ws77P3raDiOv6XesAMycUwu-uT75D5f9aNgjFeJbUH8,631
 docent/data_models/chat/content.py,sha256=Co-jO8frQa_DSP11wJuhPX0s-GpJk8yqtKqPeiAIZ_U,1672
 docent/data_models/chat/message.py,sha256=_72xeTdgv8ogQd4WLl1P3yXfIDkIEQrHlWgdvObeQxY,4291
 docent/data_models/chat/tool.py,sha256=MMglNHzkwHqUoK0xDWqs2FtelPsgHqwVpGpI1F8KZyw,3049
@@ -24,9 +24,9 @@ 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/agent_run_writer.py,sha256=0X5IP6wegoDsPiN_mblEvJa4sNWNxjPLbfNuYyA-qgY,9345
-docent/sdk/client.py,sha256=fuJrTF87OtUojULFY7acZuqg5xmE8F-4HgEeEV8_gq0,14781
-docent_python-0.1.17a0.dist-info/METADATA,sha256=D4flqV8wDdL0j8KMrVJbEezjQmdLhPGcCyZL7zlnMyA,1110
-docent_python-0.1.17a0.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
-docent_python-0.1.17a0.dist-info/licenses/LICENSE.md,sha256=QIMv2UiT6MppRasso4ymaA0w7ltkqmlL0HCt8CLD7Rc,580
-docent_python-0.1.17a0.dist-info/RECORD,,
+docent/sdk/agent_run_writer.py,sha256=0AWdxejoqZyuj9JSA39WlEwGcMSYTWNqnzIuluySY-M,11043
+docent/sdk/client.py,sha256=Y7vhb_auT4TJLy884QVD2SML7NplxqmUTaQfWpYzk-Y,18062
+docent_python-0.1.19a0.dist-info/METADATA,sha256=52oo-zJRM9keALjZ7eTF68xyVRF9RnzeKJb3xBtOPy0,1110
+docent_python-0.1.19a0.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
+docent_python-0.1.19a0.dist-info/licenses/LICENSE.md,sha256=QIMv2UiT6MppRasso4ymaA0w7ltkqmlL0HCt8CLD7Rc,580
+docent_python-0.1.19a0.dist-info/RECORD,,