docent-python 0.1.4a0__py3-none-any.whl → 0.1.5a0__py3-none-any.whl

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/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.dist-info → docent_python-0.1.5a0.dist-info}/METADATA

@@ -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.dist-info → docent_python-0.1.5a0.dist-info}/RECORD

@@ -1,7 +1,6 @@
 docent/__init__.py,sha256=J2BbO6rzilfw9WXRUeolr439EGFezqbMU_kCpCCryRA,59
 docent/py.typed,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
-docent/trace.py,sha256=T6F9jbj0JcnGBxlcR65iDPF0lmeFR6laTlhDBcw0Mh0,61167
-docent/trace_alt.py,sha256=k0EZ_jyvVDqQ3HExDPZzdZYj6tKFL_cgbNRnOq4oGyg,17747
+docent/trace.py,sha256=aPbV9fy_JpFq9y3yuqhlHKPm2QiNh9SLsKOL2azxghs,62611
 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
@@ -23,8 +22,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=uyhTisb9bHk7Hd2G4UKLdfvuiAmYOOqJiwEPbYWN9IE,12371
-docent_python-0.1.4a0.dist-info/METADATA,sha256=AUptJVGzZtABJ8V1Hpzhi2EOaMdBJB_nSaldpN6J8Bg,1074
-docent_python-0.1.4a0.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
-docent_python-0.1.4a0.dist-info/licenses/LICENSE.md,sha256=vOHzq3K4Ndu0UV9hPrtXvlD7pHOjyDQmGjHuLSIkRQY,1087
-docent_python-0.1.4a0.dist-info/RECORD,,
+docent/sdk/client.py,sha256=fLdniy8JzMLoZpaS9SP2pHban_ToavgtI8VeHZLMNZo,12773
+docent_python-0.1.5a0.dist-info/METADATA,sha256=XRzyvevWQ36Mk6UB6HIGIuD6vCQBX0cazhZylY25Q24,1037
+docent_python-0.1.5a0.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
+docent_python-0.1.5a0.dist-info/licenses/LICENSE.md,sha256=vOHzq3K4Ndu0UV9hPrtXvlD7pHOjyDQmGjHuLSIkRQY,1087
+docent_python-0.1.5a0.dist-info/RECORD,,

docent/trace_alt.py

@@ -1,513 +0,0 @@
-import asyncio
-import atexit
-import functools
-import io
-import logging
-import os
-import uuid
-from contextlib import asynccontextmanager, contextmanager, redirect_stdout
-from contextvars import ContextVar, Token
-from typing import Any, AsyncIterator, Callable, Dict, Iterator, Optional, Set
-
-import requests
-from opentelemetry.context import Context
-from opentelemetry.sdk.trace import ReadableSpan, SpanProcessor
-from opentelemetry.trace import Span
-from traceloop.sdk import Traceloop
-
-# Configure logging
-logger = logging.getLogger(__name__)
-logger.disabled = True
-
-DEFAULT_ENDPOINT = "https://api.docent.transluce.org/rest/telemetry"
-
-# Context variables for tracking current agent run and collection
-_current_agent_run_id: ContextVar[Optional[str]] = ContextVar("current_agent_run_id", default=None)
-_current_collection_id: ContextVar[Optional[str]] = ContextVar(
-    "current_collection_id", default=None
-)
-
-# Global configuration
-_tracing_initialized = False
-_collection_name: Optional[str] = None
-_collection_id: Optional[str] = None
-_default_agent_run_id: Optional[str] = None
-_endpoint: Optional[str] = None
-_api_key: Optional[str] = None
-_enable_console_export = False
-_disable_batch = False
-_instruments: Optional[Set[Any]] = None
-_block_instruments: Optional[Set[Any]] = None
-
-
-class DocentSpanProcessor(SpanProcessor):
-    """Custom span processor to add Docent metadata to spans.
-
-    This processor integrates cleanly with Traceloop's existing span processing
-    and adds Docent-specific attributes to all spans.
-    """
-
-    def __init__(self, collection_id: str, enable_console_export: bool = False):
-        self.collection_id = collection_id
-        self.enable_console_export = enable_console_export
-
-    def on_start(self, span: Span, parent_context: Optional[Context] = None) -> None:
-        """Add Docent metadata when a span starts."""
-        # Always add collection_id
-        span.set_attribute("collection_id", self.collection_id)
-
-        # Add agent_run_id if available
-        agent_run_id = _get_current_agent_run_id()
-        if agent_run_id:
-            span.set_attribute("agent_run_id", agent_run_id)
-        else:
-            span.set_attribute("agent_run_id", _get_default_agent_run_id())
-            span.set_attribute("agent_run_id_default", True)
-
-        # Add service name for better integration with existing OTEL setups
-        span.set_attribute("service.name", _collection_name or "docent-trace")
-
-        if self.enable_console_export:
-            logging.debug(
-                f"Span started - collection_id: {self.collection_id}, agent_run_id: {agent_run_id}"
-            )
-
-    def on_end(self, span: ReadableSpan) -> None:
-        pass
-
-    def shutdown(self) -> None:
-        """Called when the processor is shut down."""
-
-    def force_flush(self, timeout_millis: float = 30000) -> bool:
-        """Force flush any pending spans."""
-        return True
-
-
-def initialize_tracing(
-    collection_name: str,
-    collection_id: Optional[str] = None,
-    endpoint: Optional[str] = None,
-    api_key: Optional[str] = None,
-    enable_console_export: bool = False,
-    disable_batch: bool = False,
-    instruments: Optional[Set[Any]] = None,
-    block_instruments: Optional[Set[Any]] = None,
-) -> None:
-    """Initialize Docent tracing with the specified configuration.
-
-    This function provides a comprehensive initialization that integrates cleanly
-    with existing OpenTelemetry setups and provides extensive configuration options.
-
-    Args:
-        collection_name: Name for your application/collection
-        collection_id: Optional collection ID (auto-generated if not provided)
-        endpoint: Optional OTLP endpoint URL (defaults to Docent's hosted service)
-        api_key: Optional API key (uses DOCENT_API_KEY environment variable if not provided)
-        enable_console_export: Whether to also export traces to console for debugging
-        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
-    """
-    global _tracing_initialized, _collection_name, _collection_id, _default_agent_run_id, _endpoint, _api_key
-    global _enable_console_export, _disable_batch, _instruments, _block_instruments
-
-    if _tracing_initialized:
-        logging.warning("Docent tracing already initialized")
-        return
-
-    _collection_name = collection_name
-    _collection_id = collection_id or _generate_id()
-    _default_agent_run_id = _get_default_agent_run_id()  # Generate default ID if not set
-    _endpoint = endpoint or DEFAULT_ENDPOINT
-    _api_key = api_key or os.getenv("DOCENT_API_KEY")
-    _enable_console_export = enable_console_export
-    _disable_batch = disable_batch
-    _instruments = instruments
-    _block_instruments = block_instruments
-
-    _set_current_collection_id(_collection_id)
-
-    if not _api_key:
-        raise ValueError(
-            "API key is required. Set DOCENT_API_KEY environment variable or pass api_key parameter."
-        )
-
-    # Initialize Traceloop with comprehensive configuration
-
-    # Get Traceloop's default span processor
-    from traceloop.sdk.tracing.tracing import get_default_span_processor
-
-    # Create our custom context span processor (only adds metadata, doesn't export)
-    docent_processor = DocentSpanProcessor(_collection_id, enable_console_export)
-
-    # Get Traceloop's default span processor for export
-    export_processor = get_default_span_processor(
-        disable_batch=_disable_batch,
-        api_endpoint=_endpoint,
-        headers={"Authorization": f"Bearer {_api_key}"},
-    )
-
-    # Combine both processors
-    processors = [docent_processor, export_processor]
-
-    os.environ["TRACELOOP_METRICS_ENABLED"] = "false"
-    os.environ["TRACELOOP_TRACE_ENABLED"] = "true"
-
-    # Temporarily redirect stdout to suppress print statements
-    with redirect_stdout(io.StringIO()):
-        Traceloop.init(  # type: ignore
-            app_name=collection_name,
-            api_endpoint=_endpoint,
-            api_key=_api_key,
-            telemetry_enabled=False,  # don't send analytics to traceloop's backend
-            disable_batch=_disable_batch,
-            instruments=_instruments,
-            block_instruments=_block_instruments,
-            processor=processors,  # Add both our context processor and export processor
-        )
-
-    _tracing_initialized = True
-    logging.info(
-        f"Docent tracing initialized for collection: {collection_name} with collection_id: {_collection_id}"
-    )
-
-    # Register cleanup handlers
-    atexit.register(_cleanup_tracing)
-
-
-def _cleanup_tracing() -> None:
-    """Clean up tracing resources on shutdown."""
-    global _tracing_initialized
-    if _tracing_initialized:
-        try:
-            # Notify API that the trace is over
-            _notify_trace_done()
-
-            logging.info("Docent tracing cleanup completed")
-        except Exception as e:
-            logging.warning(f"Error during tracing cleanup: {e}")
-        finally:
-            _tracing_initialized = False
-
-
-def _ensure_tracing_initialized():
-    """Ensure tracing has been initialized before use."""
-    if not _tracing_initialized:
-        raise RuntimeError("Docent tracing not initialized. Call initialize_tracing() first.")
-
-
-def _generate_id() -> str:
-    """Generate a unique ID for agent runs or collections."""
-    return str(uuid.uuid4())
-
-
-def _get_current_agent_run_id() -> Optional[str]:
-    """Get the current agent run ID from context."""
-    return _current_agent_run_id.get()
-
-
-def _get_current_collection_id() -> Optional[str]:
-    """Get the current collection ID from context."""
-    return _current_collection_id.get()
-
-
-def _get_default_agent_run_id() -> str:
-    """Get the default agent run ID, generating it if not set."""
-    global _default_agent_run_id
-    if _default_agent_run_id is None:
-        _default_agent_run_id = _generate_id()
-    return _default_agent_run_id
-
-
-def _set_current_agent_run_id(agent_run_id: Optional[str]) -> Token[Optional[str]]:
-    """Set the current agent run ID in context."""
-    return _current_agent_run_id.set(agent_run_id)
-
-
-def _set_current_collection_id(collection_id: Optional[str]) -> Token[Optional[str]]:
-    """Set the current collection ID in context."""
-    return _current_collection_id.set(collection_id)
-
-
-def _send_to_api(endpoint: str, data: Dict[str, Any]) -> None:
-    """Send data to the Docent API endpoint.
-
-    Args:
-        endpoint: The API endpoint URL
-        data: The data to send
-    """
-    try:
-        headers = {"Content-Type": "application/json", "Authorization": f"Bearer {_api_key}"}
-
-        response = requests.post(endpoint, json=data, headers=headers, timeout=10)
-        response.raise_for_status()
-
-        logging.debug(f"Successfully sent data to {endpoint}")
-    except requests.exceptions.RequestException as e:
-        logging.error(f"Failed to send data to {endpoint}: {e}")
-    except Exception as e:
-        logging.error(f"Unexpected error sending data to {endpoint}: {e}")
-
-
-def _notify_trace_done() -> None:
-    """Notify the Docent API that the trace is done."""
-    collection_id = _get_current_collection_id()
-    if collection_id and _endpoint:
-        data = {"collection_id": collection_id, "status": "completed"}
-        _send_to_api(f"{_endpoint}/v1/trace-done", data)
-
-
-def agent_run_score(name: str, score: float, attributes: Optional[Dict[str, Any]] = None) -> None:
-    """
-    Record a score event on the current span.
-    Automatically works in both sync and async contexts.
-
-    Args:
-        name: Name of the score metric
-        score: Numeric score value
-        attributes: Optional additional attributes for the score event
-    """
-    _ensure_tracing_initialized()
-
-    agent_run_id = _get_current_agent_run_id()
-    if not agent_run_id:
-        logging.warning("No active agent run context. Score will not be sent.")
-        return
-
-    collection_id = _get_current_collection_id() or _collection_id
-    if not collection_id:
-        logging.warning("No collection ID available. Score will not be sent.")
-        return
-
-    # Send score directly to API
-    score_data = {
-        "collection_id": collection_id,
-        "agent_run_id": agent_run_id,
-        "score_name": name,
-        "score_value": score,
-    }
-
-    # Add additional attributes if provided
-    if attributes:
-        score_data.update(attributes)
-
-    _send_to_api(f"{_endpoint}/v1/scores", score_data)
-
-
-def agent_run_metadata(metadata: Dict[str, Any]) -> None:
-    """Attach metadata to the current agent run.
-
-    Args:
-        metadata: Dictionary of metadata to attach
-    """
-    _ensure_tracing_initialized()
-
-    agent_run_id = _get_current_agent_run_id()
-    if not agent_run_id:
-        logging.warning("No active agent run context. Metadata will not be sent.")
-        return
-
-    collection_id = _get_current_collection_id() or _collection_id
-    if not collection_id:
-        logging.warning("No collection ID available. Metadata will not be sent.")
-        return
-
-    # Send metadata directly to API
-    metadata_data = {
-        "collection_id": collection_id,
-        "agent_run_id": agent_run_id,
-        "metadata": metadata,
-    }
-
-    _send_to_api(f"{_endpoint}/v1/metadata", metadata_data)
-
-
-@contextmanager
-def _agent_run_context_sync(
-    agent_run_id: Optional[str] = None,
-    metadata: Optional[Dict[str, Any]] = None,
-) -> Iterator[tuple[str, Optional[str]]]:
-    """Synchronous context manager for creating and managing agent runs."""
-    _ensure_tracing_initialized()
-
-    # Generate IDs if not provided
-    current_agent_run_id = agent_run_id or _generate_id()
-
-    # Set up context
-    agent_run_token = _set_current_agent_run_id(current_agent_run_id)
-
-    try:
-        # Send metadata to API if provided
-        if metadata:
-            agent_run_metadata(metadata)
-
-        # Yield the agent run ID and None for transcript_id (handled by backend)
-        # Traceloop will automatically create spans for any instrumented operations
-        # and our DocentSpanProcessor will add the appropriate metadata
-        yield (current_agent_run_id, None)
-    finally:
-        # Restore context
-        _current_agent_run_id.reset(agent_run_token)
-
-
-@asynccontextmanager
-async def _agent_run_context_async(
-    agent_run_id: Optional[str] = None,
-    metadata: Optional[Dict[str, Any]] = None,
-) -> AsyncIterator[tuple[str, Optional[str]]]:
-    """Asynchronous context manager for creating and managing agent runs."""
-    _ensure_tracing_initialized()
-
-    # Generate IDs if not provided
-    current_agent_run_id = agent_run_id or _generate_id()
-
-    # Set up context
-    agent_run_token = _set_current_agent_run_id(current_agent_run_id)
-
-    try:
-        # Send metadata to API if provided
-        if metadata:
-            agent_run_metadata(metadata)
-
-        # Yield the agent run ID and None for transcript_id (handled by backend)
-        # Traceloop will automatically create spans for any instrumented operations
-        # and our DocentSpanProcessor will add the appropriate metadata
-        yield (current_agent_run_id, None)
-    finally:
-        # Restore context
-        _current_agent_run_id.reset(agent_run_token)
-
-
-def agent_run_context(
-    agent_run_id: Optional[str] = None,
-    metadata: Optional[Dict[str, Any]] = None,
-):
-    """Context manager for creating and managing agent runs.
-
-    This context manager can be used in both synchronous and asynchronous contexts.
-    In async contexts, use it with `async with agent_run_context()`.
-    In sync contexts, use it with `with agent_run_context()`.
-
-    Args:
-        agent_run_id: Optional agent run ID (auto-generated if not provided)
-        metadata: Optional metadata to attach to the agent run
-
-    Returns:
-        A context manager that yields a tuple of (agent_run_id, transcript_id)
-        where transcript_id is None for now as it's handled by backend
-    """
-    # Check if we're in an async context by looking at the current frame
-    import inspect
-
-    frame = inspect.currentframe()
-    try:
-        # Look for async context indicators in the call stack
-        while frame:
-            if frame.f_code.co_flags & 0x80:  # CO_COROUTINE flag
-                return _agent_run_context_async(agent_run_id=agent_run_id, metadata=metadata)
-            frame = frame.f_back
-    finally:
-        # Clean up the frame reference
-        del frame
-
-    # Default to sync context manager
-    return _agent_run_context_sync(agent_run_id=agent_run_id, metadata=metadata)
-
-
-def agent_run(
-    func: Optional[Callable[..., Any]] = None,
-    *,
-    agent_run_id: Optional[str] = None,
-    metadata: Optional[Dict[str, Any]] = None,
-) -> Callable[..., Any]:
-    """Decorator for creating agent runs around functions.
-
-    Args:
-        func: Function to decorate
-        agent_run_id: Optional agent run ID (auto-generated if not provided)
-        metadata: Optional metadata to attach to the agent run
-
-    Returns:
-        Decorated function
-    """
-
-    def decorator(func: Callable[..., Any]) -> Callable[..., Any]:
-        @functools.wraps(func)
-        def sync_wrapper(*args: Any, **kwargs: Any) -> Any:
-            with _agent_run_context_sync(agent_run_id=agent_run_id, metadata=metadata) as (
-                run_id,
-                _,
-            ):
-                result = func(*args, **kwargs)
-                # Store agent run ID as an attribute for access
-                setattr(sync_wrapper, "docent", type("DocentInfo", (), {"agent_run_id": run_id})())  # type: ignore
-                return result
-
-        @functools.wraps(func)
-        async def async_wrapper(*args: Any, **kwargs: Any) -> Any:
-            async with _agent_run_context_async(agent_run_id=agent_run_id, metadata=metadata) as (
-                run_id,
-                _,
-            ):
-                result = await func(*args, **kwargs)
-                # Store agent run ID as an attribute for access
-                setattr(async_wrapper, "docent", type("DocentInfo", (), {"agent_run_id": run_id})())  # type: ignore
-                return result
-
-        # Return appropriate wrapper based on function type
-        if asyncio.iscoroutinefunction(func):
-            return async_wrapper
-        else:
-            return sync_wrapper
-
-    # Handle both @agent_run and @agent_run(agent_run_id=..., metadata=...)
-    if func is None:
-        return decorator
-    else:
-        return decorator(func)
-
-
-# Additional utility functions for better integration
-
-
-def get_current_agent_run_id() -> Optional[str]:
-    """Get the current agent run ID from context.
-
-    Returns:
-        The current agent run ID if available, None otherwise
-    """
-    return _get_current_agent_run_id()
-
-
-def get_current_collection_id() -> Optional[str]:
-    """Get the current collection ID from context.
-
-    Returns:
-        The current collection ID if available, None otherwise
-    """
-    return _get_current_collection_id()
-
-
-def is_tracing_initialized() -> bool:
-    """Check if tracing has been initialized.
-
-    Returns:
-        True if tracing is initialized, False otherwise
-    """
-    return _tracing_initialized
-
-
-def flush_spans() -> None:
-    """Force flush any pending spans to the backend.
-
-    This is useful for ensuring all spans are sent before shutdown
-    or for debugging purposes.
-    """
-    if _tracing_initialized:
-        try:
-            traceloop_instance = Traceloop.get()
-            if hasattr(traceloop_instance, "flush"):
-                traceloop_instance.flush()  # type: ignore
-            logging.debug("Spans flushed successfully")
-        except Exception as e:
-            logging.warning(f"Error flushing spans: {e}")