docent-python 0.1.12a0__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_python-0.1.12a0.dist-info → docent_python-0.1.13a0.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: docent-python
-Version: 0.1.12a0
+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.12a0.dist-info → docent_python-0.1.13a0.dist-info}/RECORD

@@ -1,4 +1,5 @@
-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=C8oO5NuChSgyHIB5OI6qOfXWaIld7jdvlEqImK56a1E,66761
 docent/trace_temp.py,sha256=Z0lAPwVzXjFvxpiU-CuvfWIslq9Q4alNkZMoQ77Xudk,40711
@@ -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.12a0.dist-info/METADATA,sha256=OnxdikeOy69TZy_HGoWUxi7VDu20k8nwrikLcIPLXho,1038
-docent_python-0.1.12a0.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
-docent_python-0.1.12a0.dist-info/licenses/LICENSE.md,sha256=vOHzq3K4Ndu0UV9hPrtXvlD7pHOjyDQmGjHuLSIkRQY,1087
-docent_python-0.1.12a0.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,,