docent-python 0.1.11a0__py3-none-any.whl → 0.1.13a0__py3-none-any.whl

docent/__init__.py

@@ -1,3 +1,4 @@
-__all__ = ["Docent"]
+__all__ = ["Docent", "init"]
 
+from docent.agent_run_writer import init
 from docent.sdk.client import Docent

docent/agent_run_writer.py

@@ -0,0 +1,266 @@
+import atexit
+import os
+import queue
+import signal
+import threading
+import time
+from typing import Any, Callable, Coroutine, Optional
+
+import anyio
+import backoff
+import httpx
+from backoff.types import Details
+
+from docent._log_util.logger import get_logger
+from docent.data_models.agent_run import AgentRun
+from docent.sdk.client import Docent
+
+logger = get_logger(__name__)
+
+
+def _giveup(exc: BaseException) -> bool:
+    """Give up on client errors."""
+
+    if isinstance(exc, httpx.HTTPStatusError):
+        status = exc.response.status_code
+        return status < 500 and status != 429
+    return False
+
+
+def _print_backoff_message(e: Details):
+    logger.warning(
+        f"AgentRunWriter backing off for {e['wait']:.2f}s due to {e['exception'].__class__.__name__}"  # type: ignore
+    )
+
+
+class AgentRunWriter:
+    """Background thread for logging agent runs.
+
+    Args:
+        api_key (str): API key for the Docent API.
+        collection_id (str): ID of the collection to log agent runs to.
+        server_url (str): URL of the Docent server.
+        num_workers (int): Max number of concurrent tasks to run,
+            managed by anyio.CapacityLimiter.
+        queue_maxsize (int): Maximum size of the queue.
+            If maxsize is <= 0, the queue size is infinite.
+        request_timeout (float): Timeout for the HTTP request.
+        flush_interval (float): Interval to flush the queue.
+        batch_size (int): Number of agent runs to batch together.
+        max_retries (int): Maximum number of retries for the HTTP request.
+        shutdown_timeout (int): Timeout to wait for the background thread to finish
+            after the main thread has requested shutdown.
+    """
+
+    _instance: Optional["AgentRunWriter"] = None
+    _instance_lock = threading.Lock()
+
+    def __init__(
+        self,
+        api_key: str,
+        collection_id: str,
+        server_url: str = "https://api.docent.transluce.org",
+        num_workers: int = 2,
+        queue_maxsize: int = 20_000,
+        request_timeout: float = 30.0,
+        flush_interval: float = 1.0,
+        batch_size: int = 1_000,
+        max_retries: int = 5,
+        shutdown_timeout: int = 60,
+    ) -> None:
+        with self._instance_lock:
+            if AgentRunWriter._instance is not None:
+                return
+            AgentRunWriter._instance = self
+
+        # Request parameters
+        self._headers = {"Authorization": f"Bearer {api_key}"}
+        self._base_url = server_url.rstrip("/") + "/rest"
+        self._endpoint = f"{collection_id}/agent_runs"
+
+        self._num_workers = num_workers
+        self._request_timeout = request_timeout
+        self._flush_interval = flush_interval
+        self._batch_size = batch_size
+        self._max_retries = max_retries
+        self._shutdown_timeout = shutdown_timeout
+
+        self._queue: queue.Queue[AgentRun] = queue.Queue(maxsize=queue_maxsize)
+        self._cancel_event = threading.Event()
+
+        # Start background thread
+        self._thread = threading.Thread(
+            target=lambda: anyio.run(self._async_main),
+            name="AgentRunWriterThread",
+            daemon=True,
+        )
+        self._thread.start()
+        logger.info("AgentRunWriter thread started")
+
+        self._register_shutdown_hooks()
+
+    def _register_shutdown_hooks(self) -> None:
+        """Register shutdown hooks for atexit and signals."""
+
+        # Register shutdown hooks
+        atexit.register(self.finish)
+
+        # Register signal handlers for graceful shutdown
+        signal.signal(signal.SIGINT, lambda s, f: self._shutdown())  # Ctrl+C
+        signal.signal(signal.SIGTERM, lambda s, f: self._shutdown())  # Kill signal
+
+    def log_agent_runs(self, agent_runs: list[AgentRun]) -> None:
+        """Put a list of AgentRun objects into the queue.
+
+        If the queue is full, the method will block until the queue has space.
+
+        Args:
+            agent_runs (list[AgentRun]): List of AgentRun objects to put into the queue.
+        """
+
+        p_full = (
+            (self._queue.qsize() + len(agent_runs)) / self._queue.maxsize
+            if self._queue.maxsize > 0
+            else 0
+        )
+        if p_full >= 0.9:
+            logger.warning("AgentRunWriter queue is almost full (>=90%).")
+
+        for run in agent_runs:
+            try:
+                self._queue.put_nowait(run)
+            except queue.Full:
+                logger.warning("AgentRunWriter queue is full, blocking...")
+                self._queue.put(run, block=True)
+
+    def finish(self, force: bool = False) -> None:
+        """Request shutdown and wait up to timeout for pending tasks to complete.
+
+        Args:
+            force (bool): If True, shut down immediately. If False, wait for pending tasks to complete.
+        """
+        if not force:
+            # Wait for background thread to finish up to timeout
+            logger.info("Waiting for pending tasks to complete")
+
+            for i in range(0, self._shutdown_timeout, 5):
+                if not self._thread.is_alive():
+                    break
+
+                if self._queue.empty():
+                    break
+
+                logger.info(
+                    f"Waiting for pending tasks to complete " f"({i}/{self._shutdown_timeout})s"
+                )
+                time.sleep(5)
+
+        self._shutdown()
+
+    def _shutdown(self) -> None:
+        """Shutdown the AgentRunWriter thread."""
+        if self._thread.is_alive():
+            logger.info("Cancelling pending tasks...")
+            self._cancel_event.set()
+            n_pending = self._queue.qsize()
+            logger.info(f"Cancelled ~{n_pending} pending tasks")
+
+            # Give a brief moment to exit
+            logger.info("Waiting for thread to exit...")
+            self._thread.join(timeout=1.0)
+
+    def get_post_batch_fcn(
+        self, client: httpx.AsyncClient
+    ) -> Callable[[list[AgentRun], anyio.CapacityLimiter], Coroutine[Any, Any, None]]:
+        """Return a function that will post a batch of agent runs to the API."""
+
+        @backoff.on_exception(
+            backoff.expo,
+            exception=httpx.HTTPError,
+            giveup=_giveup,
+            max_tries=self._max_retries,
+            on_backoff=_print_backoff_message,
+        )
+        async def _post_batch(batch: list[AgentRun], limiter: anyio.CapacityLimiter) -> None:
+            async with limiter:
+                payload = {"agent_runs": [ar.model_dump(mode="json") for ar in batch]}
+                resp = await client.post(
+                    self._endpoint, json=payload, timeout=self._request_timeout
+                )
+                resp.raise_for_status()
+
+        return _post_batch
+
+    async def _async_main(self) -> None:
+        """Main async function for the AgentRunWriter thread."""
+
+        limiter = anyio.CapacityLimiter(self._num_workers)
+
+        async with httpx.AsyncClient(base_url=self._base_url, headers=self._headers) as client:
+            async with anyio.create_task_group() as tg:
+                _post_batch = self.get_post_batch_fcn(client)
+
+                async def batch_loop() -> None:
+                    while not self._cancel_event.is_set():
+                        batch = await self._gather_next_batch_from_queue()
+                        if not batch:
+                            continue
+
+                        tg.start_soon(_post_batch, batch, limiter)
+
+                tg.start_soon(batch_loop)
+
+    async def _gather_next_batch_from_queue(self) -> list[AgentRun]:
+        """Gather a batch of agent runs from the queue.
+
+        Fetches items from the queue until the batch is full or the timeout expires.
+        """
+        batch: list[AgentRun] = []
+        with anyio.move_on_after(self._flush_interval):
+            while len(batch) < self._batch_size:
+                try:
+                    item = self._queue.get_nowait()
+                    batch.append(item)
+                except queue.Empty:
+                    await anyio.sleep(0.1)
+
+        return batch
+
+
+def init(
+    collection_name: str = "Agent Run Collection",
+    collection_id: str | None = None,
+    server_url: str = "https://api.docent.transluce.org",
+    web_url: str = "https://docent.transluce.org",
+    api_key: str | None = None,
+):
+    """Initialize the AgentRunWriter thread.
+
+    Args:
+        collection_name (str): Name of the agent run collection.
+        collection_id (str): ID of the agent run collection.
+        server_url (str): URL of the Docent server.
+        web_url (str): URL of the Docent web UI.
+        api_key (str): API key for the Docent API.
+    """
+    api_key = api_key or os.getenv("DOCENT_API_KEY")
+
+    if api_key is None:
+        raise ValueError(
+            "api_key is required. Please provide an "
+            "api_key or set the DOCENT_API_KEY environment variable."
+        )
+
+    sdk = Docent(
+        server_url=server_url,
+        web_url=web_url,
+        api_key=api_key,
+    )
+
+    collection_id = collection_id or sdk.create_collection(name=collection_name)
+
+    return AgentRunWriter(
+        api_key=api_key,
+        collection_id=collection_id,
+        server_url=server_url,
+    )

docent/data_models/chat/tool.py

@@ -20,9 +20,9 @@ class ToolCall:
     """
 
     id: str
-    type: Literal["function"] | None
     function: str
     arguments: dict[str, Any]
+    type: Literal["function"] | None = None
     parse_error: str | None = None
     view: ToolCallContent | None = None
 

docent/trace.py

@@ -3,7 +3,6 @@ import contextvars
 import itertools
 import logging
 import os
-import signal
 import sys
 import threading
 import uuid
@@ -158,6 +157,7 @@ class DocentTracer:
             lambda: itertools.count(0)
         )
         self._transcript_counter_lock = threading.Lock()
+        self._flush_lock = threading.Lock()
 
     def get_current_agent_run_id(self) -> Optional[str]:
         """
@@ -179,14 +179,6 @@ class DocentTracer:
         # Register atexit handler
         atexit.register(self.cleanup)
 
-        # Register signal handlers for graceful shutdown
-        try:
-            signal.signal(signal.SIGINT, self._signal_handler)
-            signal.signal(signal.SIGTERM, self._signal_handler)
-        except (ValueError, OSError):
-            # Signal handlers might not work in all environments
-            pass
-
         self._cleanup_registered = True
 
     def _next_span_order(self, transcript_id: str) -> int:
@@ -197,10 +189,6 @@ class DocentTracer:
         with self._transcript_counter_lock:
             return next(self._transcript_counters[transcript_id])
 
-    def _signal_handler(self, signum: int, frame: Optional[object]):
-        """Handle shutdown signals."""
-        self.cleanup()
-
     def _init_spans_exporter(self, endpoint: str) -> Optional[Union[HTTPExporter, GRPCExporter]]:
         """Initialize the appropriate span exporter based on endpoint."""
         if not self.enable_otlp_export:
@@ -211,9 +199,11 @@ class DocentTracer:
                 http_exporter: HTTPExporter = HTTPExporter(
                     endpoint=f"{endpoint}/v1/traces", headers=self.headers
                 )
+                logger.debug(f"Initialized HTTP exporter for endpoint: {endpoint}/v1/traces")
                 return http_exporter
             else:
                 grpc_exporter: GRPCExporter = GRPCExporter(endpoint=endpoint, headers=self.headers)
+                logger.debug(f"Initialized gRPC exporter for endpoint: {endpoint}")
                 return grpc_exporter
         except Exception as e:
             logger.error(f"Failed to initialize span exporter for {endpoint}: {e}")
@@ -239,9 +229,11 @@ class DocentTracer:
         """Create appropriate span processor based on configuration."""
         if self.disable_batch or _is_notebook():
             simple_processor: SimpleSpanProcessor = SimpleSpanProcessor(exporter)
+            logger.debug("Created SimpleSpanProcessor for immediate export")
             return simple_processor
         else:
             batch_processor: BatchSpanProcessor = BatchSpanProcessor(exporter)
+            logger.debug("Created BatchSpanProcessor for batched export")
             return batch_processor
 
     def initialize(self):
@@ -310,8 +302,19 @@ class DocentTracer:
                         # attributes not available, skip them
                         pass
 
+                    # Debug logging for span creation
+                    span_name = getattr(span, "name", "unknown")
+                    span_attrs = getattr(span, "attributes", {})
+                    logger.debug(
+                        f"Created span: name='{span_name}', collection_id={self.manager.collection_id}, agent_run_id={span_attrs.get('agent_run_id')}, transcript_id={span_attrs.get('transcript_id')}"
+                    )
+
                 def on_end(self, span: ReadableSpan) -> None:
-                    pass
+                    # 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'}"
+                    )
 
                 def shutdown(self) -> None:
                     pass
@@ -422,15 +425,8 @@ class DocentTracer:
             return
 
         try:
-            # Notify backend that trace is done (no span creation)
-            try:
-                self._send_trace_done()
-            except Exception as e:
-                logger.warning(f"Failed to notify trace done: {e}")
-
-            self._root_context = None  # type: ignore
+            self.flush()
 
-            # Shutdown our isolated tracer provider
             if self._tracer_provider:
                 self._tracer_provider.shutdown()
                 self._tracer_provider = None
@@ -456,9 +452,12 @@ class DocentTracer:
             return
 
         try:
-            for processor in self._spans_processors:
+            logger.debug(f"Flushing {len(self._spans_processors)} span processors")
+            for i, processor in enumerate(self._spans_processors):
                 if hasattr(processor, "force_flush"):
-                    processor.force_flush()
+                    logger.debug(f"Flushing span processor {i}")
+                    processor.force_flush(timeout_millis=50)
+            logger.debug("Span flush completed")
         except Exception as e:
             logger.error(f"Error during flush: {e}")
 
@@ -476,29 +475,6 @@ class DocentTracer:
         """Verify if the manager is properly initialized."""
         return self._initialized
 
-    def __enter__(self) -> "DocentTracer":
-        """Context manager entry."""
-        self.initialize()
-        return self
-
-    def __exit__(self, exc_type: type[BaseException], exc_val: Any, exc_tb: Any) -> None:
-        """Context manager exit."""
-        self.close()
-
-    @property
-    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[Context]:
-        """Get the root context."""
-        if not self._initialized:
-            self.initialize()
-        return self._root_context
-
     @contextmanager
     def agent_run_context(
         self,
@@ -617,13 +593,15 @@ class DocentTracer:
         Get the API headers for HTTP requests.
 
         Returns:
-            Dictionary of headers including Authorization
+            Dictionary of headers including Authorization if set
         """
+        headers = {"Content-Type": "application/json"}
 
-        return {
-            "Content-Type": "application/json",
-            "Authorization": self.headers.get("Authorization", ""),
-        }
+        authorization = self.headers.get("Authorization")
+        if authorization:
+            headers["Authorization"] = authorization
+
+        return headers
 
     def _post_json(self, path: str, data: Dict[str, Any]) -> None:
         if not self._api_endpoint_base:
@@ -1157,7 +1135,10 @@ def close_tracing() -> None:
 def flush_tracing() -> None:
     """Force flush all spans to exporters."""
     if _global_tracer:
+        logger.debug("Flushing global tracer")
         _global_tracer.flush()
+    else:
+        logger.debug("No global tracer available to flush")
 
 
 def verify_initialized() -> bool:

{docent_python-0.1.11a0.dist-info → docent_python-0.1.13a0.dist-info}/METADATA

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

{docent_python-0.1.11a0.dist-info → docent_python-0.1.13a0.dist-info}/RECORD

@@ -1,6 +1,7 @@
-docent/__init__.py,sha256=J2BbO6rzilfw9WXRUeolr439EGFezqbMU_kCpCCryRA,59
+docent/__init__.py,sha256=KY_gsq7iKn5tZODvwbsDFMxAZukniec5nBoYawdhglo,108
+docent/agent_run_writer.py,sha256=QNCV4m36c9BuhzWCyuzs0wH9ql8uubzcQUXMhc3XVug,9135
 docent/py.typed,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
-docent/trace.py,sha256=-3QbmAjZJfKZp9pmhLne0EIB-akbz2XbEDNqpsROEco,66513
+docent/trace.py,sha256=C8oO5NuChSgyHIB5OI6qOfXWaIld7jdvlEqImK56a1E,66761
 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
@@ -16,7 +17,7 @@ docent/data_models/transcript.py,sha256=Gmy4lYdlvC5SXzpnerFJ83lIMPPiYUPgjOUbwg6a
 docent/data_models/chat/__init__.py,sha256=GleyRzYqKRkwwSRm_tQJw5BudCbgu9WRSa71Fntz0L0,610
 docent/data_models/chat/content.py,sha256=Co-jO8frQa_DSP11wJuhPX0s-GpJk8yqtKqPeiAIZ_U,1672
 docent/data_models/chat/message.py,sha256=xGt09keA6HRxw40xB_toNzEqA9ip7k53dnhXrEbKGO8,4157
-docent/data_models/chat/tool.py,sha256=x7NKINswPe0Kqvcx4ubjHzB-n0-i4DbFodvaBb2vitk,3042
+docent/data_models/chat/tool.py,sha256=MMglNHzkwHqUoK0xDWqs2FtelPsgHqwVpGpI1F8KZyw,3049
 docent/loaders/load_inspect.py,sha256=_cK2Qd6gyLQuJVzOlsvEZz7TrqzNmH6ZsLTkSCWAPqQ,6628
 docent/samples/__init__.py,sha256=roDFnU6515l9Q8v17Es_SpWyY9jbm5d6X9lV01V0MZo,143
 docent/samples/load.py,sha256=ZGE07r83GBNO4A0QBh5aQ18WAu3mTWA1vxUoHd90nrM,207
@@ -24,7 +25,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=rvOFXvyAr9QxCijN0_CWENbm8y3YQvR1msfFSBDZvOw,13309
-docent_python-0.1.11a0.dist-info/METADATA,sha256=6VpTCCXzOgvSPC3ox6eeIZepRxdcY9gP4SOh5QF5hQ4,1038
-docent_python-0.1.11a0.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
-docent_python-0.1.11a0.dist-info/licenses/LICENSE.md,sha256=vOHzq3K4Ndu0UV9hPrtXvlD7pHOjyDQmGjHuLSIkRQY,1087
-docent_python-0.1.11a0.dist-info/RECORD,,
+docent_python-0.1.13a0.dist-info/METADATA,sha256=NOSOTU8xJ5BHt_lA1Hq749RNknOBHMC2v7cQbF5ssEs,1038
+docent_python-0.1.13a0.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
+docent_python-0.1.13a0.dist-info/licenses/LICENSE.md,sha256=vOHzq3K4Ndu0UV9hPrtXvlD7pHOjyDQmGjHuLSIkRQY,1087
+docent_python-0.1.13a0.dist-info/RECORD,,