docent-python 0.1.4a0__tar.gz → 0.1.5a0__tar.gz

{docent_python-0.1.4a0 → docent_python-0.1.5a0}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: docent-python
-Version: 0.1.4a0
+Version: 0.1.5a0
 Summary: Docent SDK
 Project-URL: Homepage, https://github.com/TransluceAI/docent
 Project-URL: Issues, https://github.com/TransluceAI/docent/issues
@@ -22,4 +22,3 @@ Requires-Dist: pydantic>=2.11.7
 Requires-Dist: pyyaml>=6.0.2
 Requires-Dist: tiktoken>=0.7.0
 Requires-Dist: tqdm>=4.67.1
-Requires-Dist: traceloop-sdk>=0.44.1

{docent_python-0.1.4a0 → docent_python-0.1.5a0}/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()
 

{docent_python-0.1.4a0 → docent_python-0.1.5a0}/docent/trace.py

@@ -11,7 +11,8 @@ from collections import defaultdict
 from contextlib import asynccontextmanager, contextmanager
 from contextvars import ContextVar, Token
 from datetime import datetime, timezone
-from typing import Any, AsyncIterator, Callable, Dict, Iterator, List, Optional, Union
+from enum import Enum
+from typing import Any, AsyncIterator, Callable, Dict, Iterator, List, Optional, Set, Union
 
 import requests
 from opentelemetry import trace
@@ -42,6 +43,15 @@ DEFAULT_ENDPOINT = "https://api.docent.transluce.org/rest/telemetry"
 DEFAULT_COLLECTION_NAME = "default-collection-name"
 
 
+class Instruments(Enum):
+    """Enumeration of available instrument types."""
+
+    OPENAI = "openai"
+    ANTHROPIC = "anthropic"
+    BEDROCK = "bedrock"
+    LANGCHAIN = "langchain"
+
+
 def _is_notebook() -> bool:
     """Check if we're running in a Jupyter notebook."""
     try:
@@ -64,6 +74,8 @@ class DocentTracer:
         enable_console_export: bool = False,
         enable_otlp_export: bool = True,
         disable_batch: bool = False,
+        instruments: Optional[Set[Instruments]] = None,
+        block_instruments: Optional[Set[Instruments]] = None,
     ):
         """
         Initialize Docent tracing manager.
@@ -78,6 +90,8 @@ class DocentTracer:
             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.collection_name: str = collection_name
         self.collection_id: str = collection_id if collection_id else str(uuid.uuid4())
@@ -105,6 +119,9 @@ class DocentTracer:
         self.enable_console_export = enable_console_export
         self.enable_otlp_export = enable_otlp_export
         self.disable_batch = disable_batch
+        self.disabled_instruments: Set[Instruments] = {Instruments.LANGCHAIN}
+        self.instruments = instruments or (set(Instruments) - self.disabled_instruments)
+        self.block_instruments = block_instruments or set()
 
         # Use separate tracer provider to avoid interfering with existing OTEL setup
         self._tracer_provider: Optional[TracerProvider] = None
@@ -333,33 +350,39 @@ class DocentTracer:
             except Exception as e:
                 logger.warning(f"Failed to instrument threading: {e}")
 
+            enabled_instruments = self.instruments - self.block_instruments
+
             # Instrument OpenAI with our isolated tracer provider
-            try:
-                OpenAIInstrumentor().instrument(tracer_provider=self._tracer_provider)
-                logger.info("Instrumented OpenAI")
-            except Exception as e:
-                logger.warning(f"Failed to instrument OpenAI: {e}")
+            if Instruments.OPENAI in enabled_instruments:
+                try:
+                    OpenAIInstrumentor().instrument(tracer_provider=self._tracer_provider)
+                    logger.info("Instrumented OpenAI")
+                except Exception as e:
+                    logger.warning(f"Failed to instrument OpenAI: {e}")
 
             # Instrument Anthropic with our isolated tracer provider
-            try:
-                AnthropicInstrumentor().instrument(tracer_provider=self._tracer_provider)
-                logger.info("Instrumented Anthropic")
-            except Exception as e:
-                logger.warning(f"Failed to instrument Anthropic: {e}")
+            if Instruments.ANTHROPIC in enabled_instruments:
+                try:
+                    AnthropicInstrumentor().instrument(tracer_provider=self._tracer_provider)
+                    logger.info("Instrumented Anthropic")
+                except Exception as e:
+                    logger.warning(f"Failed to instrument Anthropic: {e}")
 
             # Instrument Bedrock with our isolated tracer provider
-            try:
-                BedrockInstrumentor().instrument(tracer_provider=self._tracer_provider)
-                logger.info("Instrumented Bedrock")
-            except Exception as e:
-                logger.warning(f"Failed to instrument Bedrock: {e}")
+            if Instruments.BEDROCK in enabled_instruments:
+                try:
+                    BedrockInstrumentor().instrument(tracer_provider=self._tracer_provider)
+                    logger.info("Instrumented Bedrock")
+                except Exception as e:
+                    logger.warning(f"Failed to instrument Bedrock: {e}")
 
             # Instrument LangChain with our isolated tracer provider
-            try:
-                LangchainInstrumentor().instrument(tracer_provider=self._tracer_provider)
-                logger.info("Instrumented LangChain")
-            except Exception as e:
-                logger.warning(f"Failed to instrument LangChain: {e}")
+            if Instruments.LANGCHAIN in enabled_instruments:
+                try:
+                    LangchainInstrumentor().instrument(tracer_provider=self._tracer_provider)
+                    logger.info("Instrumented LangChain")
+                except Exception as e:
+                    logger.warning(f"Failed to instrument LangChain: {e}")
 
             # Register cleanup handlers
             self._register_cleanup()
@@ -942,6 +965,8 @@ def initialize_tracing(
     enable_console_export: bool = False,
     enable_otlp_export: bool = True,
     disable_batch: bool = False,
+    instruments: Optional[Set[Instruments]] = None,
+    block_instruments: Optional[Set[Instruments]] = None,
 ) -> DocentTracer:
     """
     Initialize the global Docent tracer.
@@ -958,6 +983,8 @@ def initialize_tracing(
         enable_console_export: Whether to export spans to console
         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).
+        block_instruments: Set of instruments to explicitly disable.
 
     Returns:
         The initialized Docent tracer
@@ -966,6 +993,7 @@ def initialize_tracing(
         # Basic setup
         initialize_tracing("my-collection")
     """
+
     global _global_tracer
 
     # Check for API key in environment variable if not provided as parameter
@@ -983,6 +1011,8 @@ def initialize_tracing(
             enable_console_export=enable_console_export,
             enable_otlp_export=enable_otlp_export,
             disable_batch=disable_batch,
+            instruments=instruments,
+            block_instruments=block_instruments,
         )
         _global_tracer.initialize()
 

{docent_python-0.1.4a0 → docent_python-0.1.5a0}/pyproject.toml

@@ -1,7 +1,7 @@
 [project]
 name = "docent-python"
 description = "Docent SDK"
-version = "0.1.4-alpha"
+version = "0.1.5-alpha"
 authors = [
   { name="Transluce", email="info@transluce.org" },
 ]
@@ -25,7 +25,6 @@ dependencies = [
     "opentelemetry-instrumentation-langchain>=0.44.1",
     "opentelemetry-instrumentation-openai>=0.44.1",
     "opentelemetry-instrumentation-threading>=0.55b1",
-    "traceloop-sdk>=0.44.1",
 ]
 
 [build-system]