docent-python 0.1.2a0__py3-none-any.whl → 0.1.3a0__py3-none-any.whl

docent/trace.py

@@ -15,6 +15,7 @@ from contextvars import ContextVar, Token
 from typing import Any, AsyncIterator, Callable, Dict, Iterator, List, Optional, Union
 
 from opentelemetry import trace
+from opentelemetry.context import Context
 from opentelemetry.exporter.otlp.proto.grpc.trace_exporter import OTLPSpanExporter as GRPCExporter
 from opentelemetry.exporter.otlp.proto.http.trace_exporter import OTLPSpanExporter as HTTPExporter
 from opentelemetry.instrumentation.anthropic import AnthropicInstrumentor
@@ -23,17 +24,18 @@ from opentelemetry.instrumentation.langchain import LangchainInstrumentor
 from opentelemetry.instrumentation.openai import OpenAIInstrumentor
 from opentelemetry.instrumentation.threading import ThreadingInstrumentor
 from opentelemetry.sdk.resources import Resource
-from opentelemetry.sdk.trace import ReadableSpan, TracerProvider
+from opentelemetry.sdk.trace import ReadableSpan, SpanProcessor, TracerProvider
 from opentelemetry.sdk.trace.export import (
     BatchSpanProcessor,
     ConsoleSpanExporter,
     SimpleSpanProcessor,
 )
+from opentelemetry.trace import Span
 
 # Configure logging
 logging.basicConfig(level=logging.INFO)
 logger = logging.getLogger(__name__)
-logging.disable()
+logger.disabled = True
 
 # Default configuration
 DEFAULT_ENDPOINT = "https://api.docent.transluce.org/rest/telemetry"
@@ -130,26 +132,41 @@ class DocentTracer:
         self.span_postprocess_callback = span_postprocess_callback
 
         # Use separate tracer provider to avoid interfering with existing OTEL setup
-        self._tracer_provider: Optional[Any] = None
-        self._root_span: Optional[Any] = None
-        self._root_context: Optional[Any] = None
-        self._tracer: Optional[Any] = None
+        self._tracer_provider: Optional[TracerProvider] = None
+        self._root_span: Optional[Span] = None
+        self._root_context: Context = Context()
+        self._tracer: Optional[trace.Tracer] = None
         self._initialized: bool = False
         self._cleanup_registered: bool = False
         self._disabled: bool = False
-        self._spans_processors: List[Any] = []
+        self._spans_processors: List[Union[BatchSpanProcessor, SimpleSpanProcessor]] = []
 
         # Context variables for agent_run_id and transcript_id (thread/async safe)
-        self._collection_id_var: ContextVar[str] = contextvars.ContextVar("collection_id")
-        self._agent_run_id_var: ContextVar[str] = contextvars.ContextVar("agent_run_id")
-        self._transcript_id_var: ContextVar[str] = contextvars.ContextVar("transcript_id")
-        self._attributes_var: ContextVar[dict[str, Any]] = contextvars.ContextVar("attributes")
+        self._collection_id_var: ContextVar[str] = contextvars.ContextVar("docent_collection_id")
+        self._agent_run_id_var: ContextVar[str] = contextvars.ContextVar("docent_agent_run_id")
+        self._transcript_id_var: ContextVar[str] = contextvars.ContextVar("docent_transcript_id")
+        self._attributes_var: ContextVar[dict[str, Any]] = contextvars.ContextVar(
+            "docent_attributes"
+        )
         # Store atomic span order counters per transcript_id to persist across context switches
         self._transcript_counters: defaultdict[str, itertools.count[int]] = defaultdict(
             lambda: itertools.count(0)
         )
         self._transcript_counter_lock = threading.Lock()
 
+    def get_current_docent_span(self) -> Optional[Span]:
+        """
+        Get the current span from our isolated context.
+        This never touches the global OpenTelemetry context.
+        """
+        if self._root_context is None:
+            return None
+
+        try:
+            return trace.get_current_span(context=self._root_context)
+        except Exception:
+            return None
+
     def _register_cleanup(self):
         """Register cleanup handlers."""
         if self._cleanup_registered:
@@ -176,7 +193,7 @@ class DocentTracer:
         with self._transcript_counter_lock:
             return next(self._transcript_counters[transcript_id])
 
-    def _signal_handler(self, signum: int, frame: Any):
+    def _signal_handler(self, signum: int, frame: Optional[object]):
         """Handle shutdown signals."""
         self.cleanup()
         sys.exit(0)
@@ -213,13 +230,15 @@ class DocentTracer:
 
         return exporters
 
-    def _create_span_processor(self, exporter: Any) -> Any:
+    def _create_span_processor(
+        self, exporter: Union[HTTPExporter, GRPCExporter, ConsoleSpanExporter]
+    ) -> Union[SimpleSpanProcessor, BatchSpanProcessor]:
         """Create appropriate span processor based on configuration."""
         if self.disable_batch or _is_notebook():
-            simple_processor: Any = SimpleSpanProcessor(exporter)
+            simple_processor: SimpleSpanProcessor = SimpleSpanProcessor(exporter)
             return simple_processor
         else:
-            batch_processor: Any = BatchSpanProcessor(exporter)
+            batch_processor: BatchSpanProcessor = BatchSpanProcessor(exporter)
             return batch_processor
 
     def initialize(self):
@@ -234,11 +253,11 @@ class DocentTracer:
             )
 
             # Add custom span processor for run_id and transcript_id
-            class ContextSpanProcessor:
+            class ContextSpanProcessor(SpanProcessor):
                 def __init__(self, manager: "DocentTracer"):
                     self.manager: "DocentTracer" = manager
 
-                def on_start(self, span: Any, parent_context: Any = None) -> None:
+                def on_start(self, span: Span, parent_context: Optional[Context] = None) -> None:
                     # Add collection_id, agent_run_id, transcript_id, and any other current attributes
                     # Always add collection_id as it's always available
                     span.set_attribute("collection_id", self.manager.collection_id)
@@ -276,14 +295,14 @@ class DocentTracer:
                         # attributes not available, skip them
                         pass
 
-                def on_end(self, span: Any) -> None:
+                def on_end(self, span: ReadableSpan) -> None:
                     pass
 
                 def shutdown(self) -> None:
                     pass
 
-                def force_flush(self) -> None:
-                    pass
+                def force_flush(self, timeout_millis: Optional[float] = None) -> bool:
+                    return True
 
             # Configure span exporters for our isolated provider
             if self.enable_otlp_export:
@@ -294,7 +313,9 @@ class DocentTracer:
                 if otlp_exporters:
                     # Create a processor for each exporter
                     for exporter in otlp_exporters:
-                        otlp_processor: Any = self._create_span_processor(exporter)
+                        otlp_processor: Union[SimpleSpanProcessor, BatchSpanProcessor] = (
+                            self._create_span_processor(exporter)
+                        )
                         self._tracer_provider.add_span_processor(otlp_processor)
                         self._spans_processors.append(otlp_processor)
 
@@ -305,8 +326,10 @@ class DocentTracer:
                     logger.warning("Failed to initialize OTLP exporter")
 
             if self.enable_console_export:
-                console_exporter: Any = ConsoleSpanExporter()
-                console_processor: Any = self._create_span_processor(console_exporter)
+                console_exporter: ConsoleSpanExporter = ConsoleSpanExporter()
+                console_processor: Union[SimpleSpanProcessor, BatchSpanProcessor] = (
+                    self._create_span_processor(console_exporter)
+                )
                 self._tracer_provider.add_span_processor(console_processor)
                 self._spans_processors.append(console_processor)
 
@@ -318,9 +341,6 @@ class DocentTracer:
             self._tracer = self._tracer_provider.get_tracer(__name__)
 
             # Start root span
-            if self._tracer is None:
-                raise RuntimeError("Failed to get tracer from provider")
-
             self._root_span = self._tracer.start_span(
                 "application_session",
                 attributes={
@@ -328,8 +348,9 @@ class DocentTracer:
                     "session.type": "application_root",
                 },
             )
-            if self._root_span is not None:
-                self._root_context = trace.set_span_in_context(self._root_span)
+            self._root_context = trace.set_span_in_context(
+                self._root_span, context=self._root_context
+            )
 
             # Instrument threading for better context propagation
             try:
@@ -401,7 +422,7 @@ class DocentTracer:
                 self._root_span.end()
 
             self._root_span = None
-            self._root_context = None
+            self._root_context = None  # type: ignore
 
             # Shutdown our isolated tracer provider
             if self._tracer_provider:
@@ -451,21 +472,21 @@ class DocentTracer:
         self.close()
 
     @property
-    def tracer(self) -> Optional[Any]:
+    def tracer(self) -> Optional[trace.Tracer]:
         """Get the tracer instance."""
         if not self._initialized:
             self.initialize()
         return self._tracer
 
     @property
-    def root_context(self) -> Optional[Any]:
+    def root_context(self) -> Optional[Context]:
         """Get the root context."""
         if not self._initialized:
             self.initialize()
         return self._root_context
 
     @contextmanager
-    def span(self, name: str, attributes: Optional[Dict[str, Any]] = None) -> Iterator[Any]:
+    def span(self, name: str, attributes: Optional[Dict[str, Any]] = None) -> Iterator[Span]:
         """
         Context manager for creating spans with attributes.
         """
@@ -485,7 +506,7 @@ class DocentTracer:
     @asynccontextmanager
     async def async_span(
         self, name: str, attributes: Optional[Dict[str, Any]] = None
-    ) -> AsyncIterator[Any]:
+    ) -> AsyncIterator[Span]:
         """
         Async context manager for creating spans with attributes.
 
@@ -513,7 +534,7 @@ class DocentTracer:
         transcript_id: Optional[str] = None,
         metadata: Optional[Dict[str, Any]] = None,
         **attributes: Any,
-    ) -> Iterator[Any]:
+    ) -> Iterator[tuple[str, str]]:
         """
         Context manager for setting up an agent run context.
 
@@ -524,7 +545,7 @@ class DocentTracer:
             **attributes: Additional attributes to add to the context
 
         Yields:
-            Tuple of (context, agent_run_id, transcript_id)
+            Tuple of (agent_run_id, transcript_id)
         """
         if not self._initialized:
             self.initialize()
@@ -556,8 +577,7 @@ class DocentTracer:
                 if metadata:
                     _add_metadata_event_to_span(_span, metadata)
 
-                context = trace.get_current_span().get_span_context()
-                yield context, agent_run_id, transcript_id
+                yield agent_run_id, transcript_id
         finally:
             self._agent_run_id_var.reset(agent_run_id_token)
             self._transcript_id_var.reset(transcript_id_token)
@@ -570,7 +590,7 @@ class DocentTracer:
         transcript_id: Optional[str] = None,
         metadata: Optional[Dict[str, Any]] = None,
         **attributes: Any,
-    ) -> AsyncIterator[Any]:
+    ) -> AsyncIterator[tuple[str, str]]:
         """
         Async context manager for setting up an agent run context.
         Modifies the OpenTelemetry context so all spans inherit agent_run_id and transcript_id.
@@ -582,7 +602,7 @@ class DocentTracer:
             **attributes: Additional attributes to add to the context
 
         Yields:
-            Tuple of (context, agent_run_id, transcript_id)
+            Tuple of (agent_run_id, transcript_id)
         """
         if not self._initialized:
             self.initialize()
@@ -596,9 +616,9 @@ class DocentTracer:
             transcript_id = str(uuid.uuid4())
 
         # Set context variables for this execution context
-        agent_run_id_token: Any = self._agent_run_id_var.set(agent_run_id)
-        transcript_id_token: Any = self._transcript_id_var.set(transcript_id)
-        attributes_token: Any = self._attributes_var.set(attributes)
+        agent_run_id_token: Token[str] = self._agent_run_id_var.set(agent_run_id)
+        transcript_id_token: Token[str] = self._transcript_id_var.set(transcript_id)
+        attributes_token: Token[dict[str, Any]] = self._attributes_var.set(attributes)
 
         try:
             # Create a span with the agent run attributes
@@ -614,8 +634,7 @@ class DocentTracer:
                 if metadata:
                     _add_metadata_event_to_span(_span, metadata)
 
-                context = trace.get_current_span().get_span_context()
-                yield context, agent_run_id, transcript_id
+                yield agent_run_id, transcript_id
         finally:
             self._agent_run_id_var.reset(agent_run_id_token)
             self._transcript_id_var.reset(transcript_id_token)
@@ -661,7 +680,7 @@ class DocentTracer:
 
         return span, agent_run_id, transcript_id
 
-    def stop_transcript(self, span: Any) -> None:
+    def stop_transcript(self, span: Span) -> None:
         """
         Manually stop a transcript span.
 
@@ -671,7 +690,7 @@ class DocentTracer:
         if span and hasattr(span, "end"):
             span.end()
 
-    def start_span(self, name: str, attributes: Optional[Dict[str, Any]] = None) -> Any:
+    def start_span(self, name: str, attributes: Optional[Dict[str, Any]] = None) -> Span:
         """
         Manually start a span.
 
@@ -690,13 +709,13 @@ class DocentTracer:
 
         span_attributes: dict[str, Any] = attributes or {}
 
-        span: Any = self._tracer.start_span(
+        span: Span = self._tracer.start_span(
             name, context=self._root_context, attributes=span_attributes
         )
 
         return span
 
-    def stop_span(self, span: Any) -> None:
+    def stop_span(self, span: Span) -> None:
         """
         Manually stop a span.
 
@@ -829,7 +848,8 @@ def agent_run_score(name: str, score: float, attributes: Optional[Dict[str, Any]
         attributes: Optional additional attributes for the score event
     """
     try:
-        current_span: Any = trace.get_current_span()
+        # Get current span from our isolated context instead of global context
+        current_span: Optional[Span] = get_tracer().get_current_docent_span()
         if current_span and hasattr(current_span, "add_event"):
             event_attributes: dict[str, Any] = {
                 "score.name": name,
@@ -858,7 +878,7 @@ def _flatten_dict(d: Dict[str, Any], prefix: str = "") -> Dict[str, Any]:
     return flattened
 
 
-def _add_metadata_event_to_span(span: Any, metadata: Dict[str, Any]) -> None:
+def _add_metadata_event_to_span(span: Span, metadata: Dict[str, Any]) -> None:
     """
     Add metadata as an event to a span.
 
@@ -892,7 +912,7 @@ def agent_run_metadata(metadata: Dict[str, Any]) -> None:
         agent_run_metadata({"user": {"id": "123", "name": "John"}, "config": {"model": "gpt-4"}})
     """
     try:
-        current_span: Any = trace.get_current_span()
+        current_span: Optional[Span] = get_tracer().get_current_docent_span()
         if current_span:
             _add_metadata_event_to_span(current_span, metadata)
         else:
@@ -903,7 +923,7 @@ def agent_run_metadata(metadata: Dict[str, Any]) -> None:
 
 # Unified functions that automatically detect context
 @asynccontextmanager
-async def span(name: str, attributes: Optional[Dict[str, Any]] = None) -> AsyncIterator[Any]:
+async def span(name: str, attributes: Optional[Dict[str, Any]] = None) -> AsyncIterator[Span]:
     """
     Automatically choose sync or async span based on context.
     Can be used with both 'with' and 'async with'.
@@ -933,7 +953,7 @@ class AgentRunContext:
         self._sync_context: Optional[Any] = None
         self._async_context: Optional[Any] = None
 
-    def __enter__(self) -> Any:
+    def __enter__(self) -> tuple[str, str]:
         """Sync context manager entry."""
         self._sync_context = get_tracer().agent_run_context(
             self.agent_run_id, self.transcript_id, metadata=self.metadata, **self.attributes
@@ -945,7 +965,7 @@ class AgentRunContext:
         if self._sync_context:
             self._sync_context.__exit__(exc_type, exc_val, exc_tb)
 
-    async def __aenter__(self) -> Any:
+    async def __aenter__(self) -> tuple[str, str]:
         """Async context manager entry."""
         self._async_context = get_tracer().async_agent_run_context(
             self.agent_run_id, self.transcript_id, metadata=self.metadata, **self.attributes
@@ -963,13 +983,13 @@ def agent_run(
 ):
     """
     Decorator to wrap a function in an agent_run_context (sync or async).
-    Injects context, agent_run_id, and transcript_id as function attributes.
+    Injects agent_run_id and transcript_id as function attributes.
     Optionally accepts metadata to attach to the agent run context.
 
     Example:
         @agent_run
         def my_func(x, y):
-            print(my_func.docent.context, my_func.docent.agent_run_id, my_func.docent.transcript_id)
+            print(my_func.docent.agent_run_id, my_func.docent.transcript_id)
 
         @agent_run(metadata={"user": "John", "model": "gpt-4"})
         def my_func_with_metadata(x, y):
@@ -987,11 +1007,7 @@ def agent_run(
 
             @functools.wraps(f)
             async def async_wrapper(*args: Any, **kwargs: Any) -> Any:
-                async with AgentRunContext(metadata=metadata) as (
-                    context,
-                    agent_run_id,
-                    transcript_id,
-                ):
+                async with AgentRunContext(metadata=metadata) as (agent_run_id, transcript_id):
                     # Store docent data as function attributes
                     setattr(
                         async_wrapper,
@@ -1000,7 +1016,6 @@ def agent_run(
                             "DocentData",
                             (),
                             {
-                                "context": context,
                                 "agent_run_id": agent_run_id,
                                 "transcript_id": transcript_id,
                             },
@@ -1013,7 +1028,7 @@ def agent_run(
 
             @functools.wraps(f)
             def sync_wrapper(*args: Any, **kwargs: Any) -> Any:
-                with AgentRunContext(metadata=metadata) as (context, agent_run_id, transcript_id):
+                with AgentRunContext(metadata=metadata) as (agent_run_id, transcript_id):
                     # Store docent data as function attributes
                     setattr(
                         sync_wrapper,
@@ -1022,7 +1037,6 @@ def agent_run(
                             "DocentData",
                             (),
                             {
-                                "context": context,
                                 "agent_run_id": agent_run_id,
                                 "transcript_id": transcript_id,
                             },
@@ -1058,15 +1072,15 @@ def agent_run_context(
 
     Example:
         # Sync usage
-        with agent_run_context() as (context, agent_run_id, transcript_id):
+        with agent_run_context() as (agent_run_id, transcript_id):
             pass
 
         # Async usage
-        async with agent_run_context() as (context, agent_run_id, transcript_id):
+        async with agent_run_context() as (agent_run_id, transcript_id):
             pass
 
         # With metadata
-        with agent_run_context(metadata={"user": "John", "model": "gpt-4"}) as (context, agent_run_id, transcript_id):
+        with agent_run_context(metadata={"user": "John", "model": "gpt-4"}) as (agent_run_id, transcript_id):
             pass
     """
     return AgentRunContext(agent_run_id, transcript_id, metadata=metadata, **attributes)

docent/trace_alt.py

@@ -0,0 +1,497 @@
+import asyncio
+import atexit
+import functools
+import logging
+import os
+import uuid
+from contextlib import asynccontextmanager, contextmanager
+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
+
+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
+_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:
+            # If no agent_run_id in context, mark this as a default span
+            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-service")
+
+        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, _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()
+    _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
+    traceloop_processor = get_default_span_processor(
+        disable_batch=_disable_batch,
+        api_endpoint=_endpoint,
+        headers={"Authorization": f"Bearer {_api_key}"},
+    )
+
+    # Combine both processors
+    processors = [docent_processor, traceloop_processor]
+
+    os.environ["TRACELOOP_METRICS_ENABLED"] = "false"
+    os.environ["TRACELOOP_TRACE_ENABLED"] = "true"
+    # Temporarily redirect stdout to suppress Traceloop's 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 telemetry to traceloop's backend
+        disable_batch=_disable_batch,
+        instruments=_instruments,
+        block_instruments=_block_instruments,
+        processor=processors,  # Add both our context processor and Traceloop's 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 _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}")

docent_python-0.1.3a0.dist-info/METADATA

@@ -0,0 +1,25 @@
+Metadata-Version: 2.4
+Name: docent-python
+Version: 0.1.3a0
+Summary: Docent SDK
+Project-URL: Homepage, https://github.com/TransluceAI/docent
+Project-URL: Issues, https://github.com/TransluceAI/docent/issues
+Project-URL: Docs, https://transluce-docent.readthedocs-hosted.com/en/latest
+Author-email: Transluce <info@transluce.org>
+License-Expression: MIT
+License-File: LICENSE.md
+Requires-Python: >=3.11
+Requires-Dist: opentelemetry-api>=1.34.1
+Requires-Dist: opentelemetry-exporter-otlp-proto-grpc>=1.34.1
+Requires-Dist: opentelemetry-exporter-otlp-proto-http>=1.34.1
+Requires-Dist: opentelemetry-instrumentation-anthropic>=0.44.1
+Requires-Dist: opentelemetry-instrumentation-bedrock>=0.44.1
+Requires-Dist: opentelemetry-instrumentation-langchain>=0.44.1
+Requires-Dist: opentelemetry-instrumentation-openai>=0.44.1
+Requires-Dist: opentelemetry-instrumentation-threading>=0.55b1
+Requires-Dist: opentelemetry-sdk>=1.34.1
+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.2a0.dist-info → docent_python-0.1.3a0.dist-info}/RECORD

@@ -1,6 +1,7 @@
 docent/__init__.py,sha256=J2BbO6rzilfw9WXRUeolr439EGFezqbMU_kCpCCryRA,59
 docent/py.typed,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
-docent/trace.py,sha256=dWQiSo93YfHNuyC3zi_NHz_uymygMzEVtoQE_1wEyM0,39960
+docent/trace.py,sha256=Z0lAPwVzXjFvxpiU-CuvfWIslq9Q4alNkZMoQ77Xudk,40711
+docent/trace_alt.py,sha256=JmHJyDXdEAJ-qVg_e4qYlQrslj9q103wKZvq_c40G0o,17213
 docent/_log_util/__init__.py,sha256=3HXXrxrSm8PxwG4llotrCnSnp7GuroK1FNHsdg6f7aE,73
 docent/_log_util/logger.py,sha256=kwM0yRW1IJd6-XTorjWn48B4l8qvD2ZM6VDjY5eskQI,4422
 docent/data_models/__init__.py,sha256=-F8Er1RXPX6SEKFLnoi6OMhHq57KJqyHx1McP2rCsGQ,483
@@ -22,7 +23,7 @@ docent/samples/log.eval,sha256=orrW__9WBfANq7NwKsPSq9oTsQRcG6KohG5tMr_X_XY,39770
 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.2a0.dist-info/METADATA,sha256=u8jTcnP8WZ_qABVDi5zABxI3d34oLMY8hLHb55oeO_s,1037
-docent_python-0.1.2a0.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
-docent_python-0.1.2a0.dist-info/licenses/LICENSE.md,sha256=vOHzq3K4Ndu0UV9hPrtXvlD7pHOjyDQmGjHuLSIkRQY,1087
-docent_python-0.1.2a0.dist-info/RECORD,,
+docent_python-0.1.3a0.dist-info/METADATA,sha256=1MaUbt-Jf03yM3185m2sslYO5oH9BjcdjwxbQnWKfF8,1074
+docent_python-0.1.3a0.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
+docent_python-0.1.3a0.dist-info/licenses/LICENSE.md,sha256=vOHzq3K4Ndu0UV9hPrtXvlD7pHOjyDQmGjHuLSIkRQY,1087
+docent_python-0.1.3a0.dist-info/RECORD,,

docent_python-0.1.2a0.dist-info/METADATA

@@ -1,24 +0,0 @@
-Metadata-Version: 2.4
-Name: docent-python
-Version: 0.1.2a0
-Summary: Docent SDK
-Project-URL: Homepage, https://github.com/TransluceAI/docent
-Project-URL: Issues, https://github.com/TransluceAI/docent/issues
-Project-URL: Docs, https://transluce-docent.readthedocs-hosted.com/en/latest
-Author-email: Transluce <info@transluce.org>
-License-Expression: MIT
-License-File: LICENSE.md
-Requires-Python: >=3.11
-Requires-Dist: opentelemetry-api>=1.35.0
-Requires-Dist: opentelemetry-exporter-otlp-proto-grpc>=1.35.0
-Requires-Dist: opentelemetry-exporter-otlp-proto-http>=1.35.0
-Requires-Dist: opentelemetry-instrumentation-anthropic>=0.42.0
-Requires-Dist: opentelemetry-instrumentation-bedrock>=0.42.0
-Requires-Dist: opentelemetry-instrumentation-langchain>=0.43.1
-Requires-Dist: opentelemetry-instrumentation-openai>=0.42.0
-Requires-Dist: opentelemetry-instrumentation-threading>=0.56b0
-Requires-Dist: opentelemetry-sdk>=1.35.0
-Requires-Dist: pydantic>=2.11.7
-Requires-Dist: pyyaml>=6.0.2
-Requires-Dist: tiktoken>=0.7.0
-Requires-Dist: tqdm>=4.67.1