coco-framework 0.1.0__py3-none-any.whl

coco_core/tracers/base.py

@@ -0,0 +1,955 @@
+"""Base interfaces for tracing runs."""
+
+from __future__ import annotations
+
+import asyncio
+import logging
+from abc import ABC, abstractmethod
+from typing import (
+    TYPE_CHECKING,
+    Any,
+)
+
+from typing_extensions import override
+
+from coco_core.callbacks.base import AsyncCallbackHandler, BaseCallbackHandler
+from coco_core.exceptions import TracerException  # noqa: F401
+from coco_core.tracers.core import _TracerCore
+
+if TYPE_CHECKING:
+    from collections.abc import Sequence
+    from uuid import UUID
+
+    from tenacity import RetryCallState
+
+    from coco_core.documents import Document
+    from coco_core.messages import BaseMessage
+    from coco_core.outputs import ChatGenerationChunk, GenerationChunk, LLMResult
+    from coco_core.tracers.schemas import Run
+
+logger = logging.getLogger(__name__)
+
+
+class BaseTracer(_TracerCore, BaseCallbackHandler, ABC):
+    """Base interface for tracers."""
+
+    @abstractmethod
+    def _persist_run(self, run: Run) -> None:
+        """Persist a run."""
+
+    def _start_trace(self, run: Run) -> None:
+        """Start a trace for a run."""
+        super()._start_trace(run)
+        self._on_run_create(run)
+
+    def _end_trace(self, run: Run) -> None:
+        """End a trace for a run."""
+        if not run.parent_run_id:
+            self._persist_run(run)
+        self.run_map.pop(str(run.id))
+        # If this run's parent was injected from an external tracing context
+        # (e.g. a langsmith @traceable), decrement its child refcount and
+        # remove it from run_map once the last child is done.
+        parent_id = str(run.parent_run_id) if run.parent_run_id else None
+        if parent_id and parent_id in self._external_run_ids:
+            self._external_run_ids[parent_id] -= 1
+            if self._external_run_ids[parent_id] <= 0:
+                self.run_map.pop(parent_id, None)
+                del self._external_run_ids[parent_id]
+        self._on_run_update(run)
+
+    def on_chat_model_start(
+        self,
+        serialized: dict[str, Any],
+        messages: list[list[BaseMessage]],
+        *,
+        run_id: UUID,
+        tags: list[str] | None = None,
+        parent_run_id: UUID | None = None,
+        metadata: dict[str, Any] | None = None,
+        name: str | None = None,
+        **kwargs: Any,
+    ) -> Run:
+        """Start a trace for a chat model run.
+
+        Note:
+            Naming can be confusing here: there is `on_chat_model_start`, but no
+            corresponding `on_chat_model_end` callback. Chat model completion is
+            routed through `on_llm_end` / `_on_llm_end`, which are shared with
+            text LLM runs.
+
+        Args:
+            serialized: The serialized model.
+            messages: The messages to start the chat with.
+            run_id: The run ID.
+            tags: The tags for the run.
+            parent_run_id: The parent run ID.
+            metadata: The metadata for the run.
+            name: The name of the run.
+            **kwargs: Additional arguments.
+
+        Returns:
+            The run.
+        """
+        chat_model_run = self._create_chat_model_run(
+            serialized=serialized,
+            messages=messages,
+            run_id=run_id,
+            parent_run_id=parent_run_id,
+            tags=tags,
+            metadata=metadata,
+            name=name,
+            **kwargs,
+        )
+        self._start_trace(chat_model_run)
+        self._on_chat_model_start(chat_model_run)
+        return chat_model_run
+
+    def on_llm_start(
+        self,
+        serialized: dict[str, Any],
+        prompts: list[str],
+        *,
+        run_id: UUID,
+        tags: list[str] | None = None,
+        parent_run_id: UUID | None = None,
+        metadata: dict[str, Any] | None = None,
+        name: str | None = None,
+        **kwargs: Any,
+    ) -> Run:
+        """Start a trace for an LLM run.
+
+        Args:
+            serialized: The serialized model.
+            prompts: The prompts to start the LLM with.
+            run_id: The run ID.
+            tags: The tags for the run.
+            parent_run_id: The parent run ID.
+            metadata: The metadata for the run.
+            name: The name of the run.
+            **kwargs: Additional arguments.
+
+        Returns:
+            The run.
+        """
+        llm_run = self._create_llm_run(
+            serialized=serialized,
+            prompts=prompts,
+            run_id=run_id,
+            parent_run_id=parent_run_id,
+            tags=tags,
+            metadata=metadata,
+            name=name,
+            **kwargs,
+        )
+        self._start_trace(llm_run)
+        self._on_llm_start(llm_run)
+        return llm_run
+
+    @override
+    def on_llm_new_token(
+        self,
+        token: str | list[str | dict[str, Any]],
+        *,
+        chunk: GenerationChunk | ChatGenerationChunk | None = None,
+        run_id: UUID,
+        parent_run_id: UUID | None = None,
+        **kwargs: Any,
+    ) -> Run:
+        """Run on new LLM token.
+
+        Only available when streaming is enabled.
+
+        Args:
+            token: The token, or a list of content blocks for structured output.
+            chunk: The chunk.
+            run_id: The run ID.
+            parent_run_id: The parent run ID.
+            **kwargs: Additional arguments.
+
+        Returns:
+            The run.
+        """
+        # "chat_model" is only used for the experimental new streaming_events format.
+        # This change should not affect any existing tracers.
+        llm_run = self._llm_run_with_token_event(
+            token=token,
+            run_id=run_id,
+            chunk=chunk,
+            parent_run_id=parent_run_id,
+        )
+        self._on_llm_new_token(llm_run, token, chunk)
+        return llm_run
+
+    @override
+    def on_retry(
+        self,
+        retry_state: RetryCallState,
+        *,
+        run_id: UUID,
+        **kwargs: Any,
+    ) -> Run:
+        """Run on retry.
+
+        Args:
+            retry_state: The retry state.
+            run_id: The run ID.
+            **kwargs: Additional arguments.
+
+        Returns:
+            The run.
+        """
+        return self._llm_run_with_retry_event(
+            retry_state=retry_state,
+            run_id=run_id,
+        )
+
+    @override
+    def on_llm_end(self, response: LLMResult, *, run_id: UUID, **kwargs: Any) -> Run:
+        """End a trace for an LLM or chat model run.
+
+        Note:
+            This is the end callback for both run types. Chat models start with
+            `on_chat_model_start`, but there is no `on_chat_model_end`;
+            completion is routed here for callback API compatibility.
+
+        Args:
+            response: The response.
+            run_id: The run ID.
+            **kwargs: Additional arguments.
+
+        Returns:
+            The run.
+        """
+        # "chat_model" is only used for the experimental new streaming_events format.
+        # This change should not affect any existing tracers.
+        llm_run = self._complete_llm_run(
+            response=response,
+            run_id=run_id,
+        )
+        self._end_trace(llm_run)
+        self._on_llm_end(llm_run)
+        return llm_run
+
+    def on_llm_error(
+        self,
+        error: BaseException,
+        *,
+        run_id: UUID,
+        **kwargs: Any,
+    ) -> Run:
+        """Handle an error for an LLM run.
+
+        Args:
+            error: The error.
+            run_id: The run ID.
+            **kwargs: Additional arguments.
+
+        Returns:
+            The run.
+        """
+        # "chat_model" is only used for the experimental new streaming_events format.
+        # This change should not affect any existing tracers.
+        llm_run = self._errored_llm_run(
+            error=error, run_id=run_id, response=kwargs.pop("response", None)
+        )
+        self._end_trace(llm_run)
+        self._on_llm_error(llm_run)
+        return llm_run
+
+    @override
+    def on_chain_start(
+        self,
+        serialized: dict[str, Any],
+        inputs: dict[str, Any],
+        *,
+        run_id: UUID,
+        tags: list[str] | None = None,
+        parent_run_id: UUID | None = None,
+        metadata: dict[str, Any] | None = None,
+        run_type: str | None = None,
+        name: str | None = None,
+        **kwargs: Any,
+    ) -> Run:
+        """Start a trace for a chain run.
+
+        Args:
+            serialized: The serialized chain.
+            inputs: The inputs for the chain.
+            run_id: The run ID.
+            tags: The tags for the run.
+            parent_run_id: The parent run ID.
+            metadata: The metadata for the run.
+            run_type: The type of the run.
+            name: The name of the run.
+            **kwargs: Additional arguments.
+
+        Returns:
+            The run.
+        """
+        chain_run = self._create_chain_run(
+            serialized=serialized,
+            inputs=inputs,
+            run_id=run_id,
+            tags=tags,
+            parent_run_id=parent_run_id,
+            metadata=metadata,
+            run_type=run_type,
+            name=name,
+            **kwargs,
+        )
+        self._start_trace(chain_run)
+        self._on_chain_start(chain_run)
+        return chain_run
+
+    @override
+    def on_chain_end(
+        self,
+        outputs: dict[str, Any],
+        *,
+        run_id: UUID,
+        inputs: dict[str, Any] | None = None,
+        **kwargs: Any,
+    ) -> Run:
+        """End a trace for a chain run.
+
+        Args:
+            outputs: The outputs for the chain.
+            run_id: The run ID.
+            inputs: The inputs for the chain.
+            **kwargs: Additional arguments.
+
+        Returns:
+            The run.
+        """
+        chain_run = self._complete_chain_run(
+            outputs=outputs,
+            run_id=run_id,
+            inputs=inputs,
+        )
+        self._end_trace(chain_run)
+        self._on_chain_end(chain_run)
+        return chain_run
+
+    @override
+    def on_chain_error(
+        self,
+        error: BaseException,
+        *,
+        inputs: dict[str, Any] | None = None,
+        run_id: UUID,
+        **kwargs: Any,
+    ) -> Run:
+        """Handle an error for a chain run.
+
+        Args:
+            error: The error.
+            inputs: The inputs for the chain.
+            run_id: The run ID.
+            **kwargs: Additional arguments.
+
+        Returns:
+            The run.
+        """
+        chain_run = self._errored_chain_run(
+            error=error,
+            run_id=run_id,
+            inputs=inputs,
+        )
+        self._end_trace(chain_run)
+        self._on_chain_error(chain_run)
+        return chain_run
+
+    def on_tool_start(
+        self,
+        serialized: dict[str, Any],
+        input_str: str,
+        *,
+        run_id: UUID,
+        tags: list[str] | None = None,
+        parent_run_id: UUID | None = None,
+        metadata: dict[str, Any] | None = None,
+        name: str | None = None,
+        inputs: dict[str, Any] | None = None,
+        **kwargs: Any,
+    ) -> Run:
+        """Start a trace for a tool run.
+
+        Args:
+            serialized: The serialized tool.
+            input_str: The input string.
+            run_id: The run ID.
+            tags: The tags for the run.
+            parent_run_id: The parent run ID.
+            metadata: The metadata for the run.
+            name: The name of the run.
+            inputs: The inputs for the tool.
+            **kwargs: Additional arguments.
+
+        Returns:
+            The run.
+        """
+        tool_run = self._create_tool_run(
+            serialized=serialized,
+            input_str=input_str,
+            run_id=run_id,
+            tags=tags,
+            parent_run_id=parent_run_id,
+            metadata=metadata,
+            name=name,
+            inputs=inputs,
+            **kwargs,
+        )
+        self._start_trace(tool_run)
+        self._on_tool_start(tool_run)
+        return tool_run
+
+    @override
+    def on_tool_end(self, output: Any, *, run_id: UUID, **kwargs: Any) -> Run:
+        """End a trace for a tool run.
+
+        Args:
+            output: The output for the tool.
+            run_id: The run ID.
+            **kwargs: Additional arguments.
+
+        Returns:
+            The run.
+        """
+        tool_run = self._complete_tool_run(
+            output=output,
+            run_id=run_id,
+        )
+        self._end_trace(tool_run)
+        self._on_tool_end(tool_run)
+        return tool_run
+
+    @override
+    def on_tool_error(
+        self,
+        error: BaseException,
+        *,
+        run_id: UUID,
+        **kwargs: Any,
+    ) -> Run:
+        """Handle an error for a tool run.
+
+        Args:
+            error: The error.
+            run_id: The run ID.
+            **kwargs: Additional arguments.
+
+        Returns:
+            The run.
+        """
+        tool_run = self._errored_tool_run(
+            error=error,
+            run_id=run_id,
+        )
+        self._end_trace(tool_run)
+        self._on_tool_error(tool_run)
+        return tool_run
+
+    def on_retriever_start(
+        self,
+        serialized: dict[str, Any],
+        query: str,
+        *,
+        run_id: UUID,
+        parent_run_id: UUID | None = None,
+        tags: list[str] | None = None,
+        metadata: dict[str, Any] | None = None,
+        name: str | None = None,
+        **kwargs: Any,
+    ) -> Run:
+        """Run when the `Retriever` starts running.
+
+        Args:
+            serialized: The serialized retriever.
+            query: The query.
+            run_id: The run ID.
+            parent_run_id: The parent run ID.
+            tags: The tags for the run.
+            metadata: The metadata for the run.
+            name: The name of the run.
+            **kwargs: Additional arguments.
+
+        Returns:
+            The run.
+        """
+        retrieval_run = self._create_retrieval_run(
+            serialized=serialized,
+            query=query,
+            run_id=run_id,
+            parent_run_id=parent_run_id,
+            tags=tags,
+            metadata=metadata,
+            name=name,
+            **kwargs,
+        )
+        self._start_trace(retrieval_run)
+        self._on_retriever_start(retrieval_run)
+        return retrieval_run
+
+    @override
+    def on_retriever_error(
+        self,
+        error: BaseException,
+        *,
+        run_id: UUID,
+        **kwargs: Any,
+    ) -> Run:
+        """Run when `Retriever` errors.
+
+        Args:
+            error: The error.
+            run_id: The run ID.
+            **kwargs: Additional arguments.
+
+        Returns:
+            The run.
+        """
+        retrieval_run = self._errored_retrieval_run(
+            error=error,
+            run_id=run_id,
+        )
+        self._end_trace(retrieval_run)
+        self._on_retriever_error(retrieval_run)
+        return retrieval_run
+
+    @override
+    def on_retriever_end(
+        self, documents: Sequence[Document], *, run_id: UUID, **kwargs: Any
+    ) -> Run:
+        """Run when the `Retriever` ends running.
+
+        Args:
+            documents: The documents.
+            run_id: The run ID.
+            **kwargs: Additional arguments.
+
+        Returns:
+            The run.
+        """
+        retrieval_run = self._complete_retrieval_run(
+            documents=documents,
+            run_id=run_id,
+        )
+        self._end_trace(retrieval_run)
+        self._on_retriever_end(retrieval_run)
+        return retrieval_run
+
+    def __deepcopy__(self, memo: dict[int, Any] | None = None) -> BaseTracer:
+        """Return self."""
+        return self
+
+    def __copy__(self) -> BaseTracer:
+        """Return self."""
+        return self
+
+
+class AsyncBaseTracer(_TracerCore, AsyncCallbackHandler, ABC):
+    """Async base interface for tracers."""
+
+    @abstractmethod
+    @override
+    async def _persist_run(self, run: Run) -> None:
+        """Persist a run."""
+
+    @override
+    async def _start_trace(self, run: Run) -> None:
+        """Start a trace for a run.
+
+        Starting a trace will run concurrently with each `_on_[run_type]_start` method.
+        No `_on_[run_type]_start` callback should depend on operations in
+        `_start_trace`.
+        """
+        super()._start_trace(run)
+        await self._on_run_create(run)
+
+    @override
+    async def _end_trace(self, run: Run) -> None:
+        """End a trace for a run.
+
+        Ending a trace will run concurrently with each `_on_[run_type]_end` method.
+        No `_on_[run_type]_end` callback should depend on operations in `_end_trace`.
+        """
+        if not run.parent_run_id:
+            await self._persist_run(run)
+        self.run_map.pop(str(run.id))
+        # If this run's parent was injected from an external tracing context
+        # (e.g. a langsmith @traceable), decrement its child refcount and
+        # remove it from run_map once the last child is done.
+        parent_id = str(run.parent_run_id) if run.parent_run_id else None
+        if parent_id and parent_id in self._external_run_ids:
+            self._external_run_ids[parent_id] -= 1
+            if self._external_run_ids[parent_id] <= 0:
+                self.run_map.pop(parent_id, None)
+                del self._external_run_ids[parent_id]
+        await self._on_run_update(run)
+
+    @override
+    async def on_chat_model_start(
+        self,
+        serialized: dict[str, Any],
+        messages: list[list[BaseMessage]],
+        *,
+        run_id: UUID,
+        parent_run_id: UUID | None = None,
+        tags: list[str] | None = None,
+        metadata: dict[str, Any] | None = None,
+        name: str | None = None,
+        **kwargs: Any,
+    ) -> Any:
+        chat_model_run = self._create_chat_model_run(
+            serialized=serialized,
+            messages=messages,
+            run_id=run_id,
+            parent_run_id=parent_run_id,
+            tags=tags,
+            metadata=metadata,
+            name=name,
+            **kwargs,
+        )
+        tasks = [
+            self._start_trace(chat_model_run),
+            self._on_chat_model_start(chat_model_run),
+        ]
+        await asyncio.gather(*tasks)
+        return chat_model_run
+
+    @override
+    async def on_llm_start(
+        self,
+        serialized: dict[str, Any],
+        prompts: list[str],
+        *,
+        run_id: UUID,
+        parent_run_id: UUID | None = None,
+        tags: list[str] | None = None,
+        metadata: dict[str, Any] | None = None,
+        **kwargs: Any,
+    ) -> None:
+        llm_run = self._create_llm_run(
+            serialized=serialized,
+            prompts=prompts,
+            run_id=run_id,
+            parent_run_id=parent_run_id,
+            tags=tags,
+            metadata=metadata,
+            **kwargs,
+        )
+        tasks = [self._start_trace(llm_run), self._on_llm_start(llm_run)]
+        await asyncio.gather(*tasks)
+
+    @override
+    async def on_llm_new_token(
+        self,
+        token: str | list[str | dict[str, Any]],
+        *,
+        chunk: GenerationChunk | ChatGenerationChunk | None = None,
+        run_id: UUID,
+        parent_run_id: UUID | None = None,
+        **kwargs: Any,
+    ) -> None:
+        llm_run = self._llm_run_with_token_event(
+            token=token,
+            run_id=run_id,
+            chunk=chunk,
+            parent_run_id=parent_run_id,
+        )
+        await self._on_llm_new_token(llm_run, token, chunk)
+
+    @override
+    async def on_retry(
+        self,
+        retry_state: RetryCallState,
+        *,
+        run_id: UUID,
+        **kwargs: Any,
+    ) -> None:
+        self._llm_run_with_retry_event(
+            retry_state=retry_state,
+            run_id=run_id,
+        )
+
+    @override
+    async def on_llm_end(
+        self,
+        response: LLMResult,
+        *,
+        run_id: UUID,
+        parent_run_id: UUID | None = None,
+        tags: list[str] | None = None,
+        **kwargs: Any,
+    ) -> None:
+        """End a trace for an LLM or chat model run.
+
+        Note:
+            This async callback also handles both run types. Async chat models
+            start with `on_chat_model_start`, but there is no
+            `on_chat_model_end`; completion is routed here for callback API
+            compatibility.
+        """
+        llm_run = self._complete_llm_run(
+            response=response,
+            run_id=run_id,
+        )
+        tasks = [self._on_llm_end(llm_run), self._end_trace(llm_run)]
+        await asyncio.gather(*tasks)
+
+    @override
+    async def on_llm_error(
+        self,
+        error: BaseException,
+        *,
+        run_id: UUID,
+        parent_run_id: UUID | None = None,
+        tags: list[str] | None = None,
+        **kwargs: Any,
+    ) -> None:
+        llm_run = self._errored_llm_run(
+            error=error,
+            run_id=run_id,
+        )
+        tasks = [self._on_llm_error(llm_run), self._end_trace(llm_run)]
+        await asyncio.gather(*tasks)
+
+    @override
+    async def on_chain_start(
+        self,
+        serialized: dict[str, Any],
+        inputs: dict[str, Any],
+        *,
+        run_id: UUID,
+        tags: list[str] | None = None,
+        parent_run_id: UUID | None = None,
+        metadata: dict[str, Any] | None = None,
+        run_type: str | None = None,
+        name: str | None = None,
+        **kwargs: Any,
+    ) -> None:
+        chain_run = self._create_chain_run(
+            serialized=serialized,
+            inputs=inputs,
+            run_id=run_id,
+            tags=tags,
+            parent_run_id=parent_run_id,
+            metadata=metadata,
+            run_type=run_type,
+            name=name,
+            **kwargs,
+        )
+        tasks = [self._start_trace(chain_run), self._on_chain_start(chain_run)]
+        await asyncio.gather(*tasks)
+
+    @override
+    async def on_chain_end(
+        self,
+        outputs: dict[str, Any],
+        *,
+        run_id: UUID,
+        inputs: dict[str, Any] | None = None,
+        **kwargs: Any,
+    ) -> None:
+        chain_run = self._complete_chain_run(
+            outputs=outputs,
+            run_id=run_id,
+            inputs=inputs,
+        )
+        tasks = [self._end_trace(chain_run), self._on_chain_end(chain_run)]
+        await asyncio.gather(*tasks)
+
+    @override
+    async def on_chain_error(
+        self,
+        error: BaseException,
+        *,
+        inputs: dict[str, Any] | None = None,
+        run_id: UUID,
+        **kwargs: Any,
+    ) -> None:
+        chain_run = self._errored_chain_run(
+            error=error,
+            inputs=inputs,
+            run_id=run_id,
+        )
+        tasks = [self._end_trace(chain_run), self._on_chain_error(chain_run)]
+        await asyncio.gather(*tasks)
+
+    @override
+    async def on_tool_start(
+        self,
+        serialized: dict[str, Any],
+        input_str: str,
+        *,
+        run_id: UUID,
+        tags: list[str] | None = None,
+        parent_run_id: UUID | None = None,
+        metadata: dict[str, Any] | None = None,
+        name: str | None = None,
+        inputs: dict[str, Any] | None = None,
+        **kwargs: Any,
+    ) -> None:
+        tool_run = self._create_tool_run(
+            serialized=serialized,
+            input_str=input_str,
+            run_id=run_id,
+            tags=tags,
+            parent_run_id=parent_run_id,
+            metadata=metadata,
+            inputs=inputs,
+            **kwargs,
+        )
+        tasks = [self._start_trace(tool_run), self._on_tool_start(tool_run)]
+        await asyncio.gather(*tasks)
+
+    @override
+    async def on_tool_end(
+        self,
+        output: Any,
+        *,
+        run_id: UUID,
+        **kwargs: Any,
+    ) -> None:
+        tool_run = self._complete_tool_run(
+            output=output,
+            run_id=run_id,
+        )
+        tasks = [self._end_trace(tool_run), self._on_tool_end(tool_run)]
+        await asyncio.gather(*tasks)
+
+    @override
+    async def on_tool_error(
+        self,
+        error: BaseException,
+        *,
+        run_id: UUID,
+        parent_run_id: UUID | None = None,
+        tags: list[str] | None = None,
+        **kwargs: Any,
+    ) -> None:
+        tool_run = self._errored_tool_run(
+            error=error,
+            run_id=run_id,
+        )
+        tasks = [self._end_trace(tool_run), self._on_tool_error(tool_run)]
+        await asyncio.gather(*tasks)
+
+    @override
+    async def on_retriever_start(
+        self,
+        serialized: dict[str, Any],
+        query: str,
+        *,
+        run_id: UUID,
+        parent_run_id: UUID | None = None,
+        tags: list[str] | None = None,
+        metadata: dict[str, Any] | None = None,
+        name: str | None = None,
+        **kwargs: Any,
+    ) -> None:
+        retriever_run = self._create_retrieval_run(
+            serialized=serialized,
+            query=query,
+            run_id=run_id,
+            parent_run_id=parent_run_id,
+            tags=tags,
+            metadata=metadata,
+            name=name,
+        )
+        tasks = [
+            self._start_trace(retriever_run),
+            self._on_retriever_start(retriever_run),
+        ]
+        await asyncio.gather(*tasks)
+
+    @override
+    async def on_retriever_error(
+        self,
+        error: BaseException,
+        *,
+        run_id: UUID,
+        parent_run_id: UUID | None = None,
+        tags: list[str] | None = None,
+        **kwargs: Any,
+    ) -> None:
+        retrieval_run = self._errored_retrieval_run(
+            error=error,
+            run_id=run_id,
+        )
+        tasks = [
+            self._end_trace(retrieval_run),
+            self._on_retriever_error(retrieval_run),
+        ]
+        await asyncio.gather(*tasks)
+
+    @override
+    async def on_retriever_end(
+        self,
+        documents: Sequence[Document],
+        *,
+        run_id: UUID,
+        parent_run_id: UUID | None = None,
+        tags: list[str] | None = None,
+        **kwargs: Any,
+    ) -> None:
+        retrieval_run = self._complete_retrieval_run(
+            documents=documents,
+            run_id=run_id,
+        )
+        tasks = [self._end_trace(retrieval_run), self._on_retriever_end(retrieval_run)]
+        await asyncio.gather(*tasks)
+
+    async def _on_run_create(self, run: Run) -> None:
+        """Process a run upon creation."""
+
+    async def _on_run_update(self, run: Run) -> None:
+        """Process a run upon update."""
+
+    async def _on_llm_start(self, run: Run) -> None:
+        """Process the LLM Run upon start."""
+
+    async def _on_llm_end(self, run: Run) -> None:
+        """Process LLM/chat model run completion."""
+
+    async def _on_llm_error(self, run: Run) -> None:
+        """Process the LLM Run upon error."""
+
+    async def _on_llm_new_token(
+        self,
+        run: Run,
+        token: str | list[str | dict[str, Any]],
+        chunk: GenerationChunk | ChatGenerationChunk | None,
+    ) -> None:
+        """Process new LLM token."""
+
+    async def _on_chain_start(self, run: Run) -> None:
+        """Process the Chain Run upon start."""
+
+    async def _on_chain_end(self, run: Run) -> None:
+        """Process the Chain Run."""
+
+    async def _on_chain_error(self, run: Run) -> None:
+        """Process the Chain Run upon error."""
+
+    async def _on_tool_start(self, run: Run) -> None:
+        """Process the Tool Run upon start."""
+
+    async def _on_tool_end(self, run: Run) -> None:
+        """Process the Tool Run."""
+
+    async def _on_tool_error(self, run: Run) -> None:
+        """Process the Tool Run upon error."""
+
+    async def _on_chat_model_start(self, run: Run) -> None:
+        """Process the Chat Model Run upon start."""
+
+    async def _on_retriever_start(self, run: Run) -> None:
+        """Process the Retriever Run upon start."""
+
+    async def _on_retriever_end(self, run: Run) -> None:
+        """Process the Retriever Run."""
+
+    async def _on_retriever_error(self, run: Run) -> None:
+        """Process the Retriever Run upon error."""