Tony20100/langchain_full huggingface dataset:
Source code for langchain_core.tracers.run_collector
"""A tracer that collects all nested runs in a list."""
from typing import Any, List, Optional, Union
from uuid import UUID
from langchain_core.tracers.base import BaseTracer
from langchain_core.tracers.schemas import Run
[docs]class RunCollectorCallbackHandler(BaseTracer):
    """
    Tracer that collects all nested runs in a list.
    This tracer is useful for inspection and evaluation purposes.
    Parameters
    ----------
    example_id : Optional[Union[UUID, str]], default=None
        The ID of the example being traced. It can be either a UUID or a string.
    """
    name: str = "run-collector_callback_handler"
[docs]    def __init__(
        self, example_id: Optional[Union[UUID, str]] = None, **kwargs: Any
    ) -> None:
        """
        Initialize the RunCollectorCallbackHandler.
        Parameters
        ----------
        example_id : Optional[Union[UUID, str]], default=None
            The ID of the example being traced. It can be either a UUID or a string.
        """
        super().__init__(**kwargs)
        self.example_id = (
            UUID(example_id) if isinstance(example_id, str) else example_id
        )
        self.traced_runs: List[Run] = []
    def _persist_run(self, run: Run) -> None:
        """
        Persist a run by adding it to the traced_runs list.
        Parameters
        ----------
        run : Run
            The run to be persisted.
        """
        run_ = run.copy()
        run_.reference_example_id = self.example_id
        self.traced_runs.append(run_)
Source code for langchain_core.tracers.context
from __future__ import annotations
from contextlib import contextmanager
from contextvars import ContextVar
from typing import (
    TYPE_CHECKING,
    Any,
    Generator,
    List,
    Optional,
    Tuple,
    Type,
    Union,
    cast,
)
from uuid import UUID
from langsmith import utils as ls_utils
from langsmith.run_helpers import get_run_tree_context
from langchain_core.tracers.langchain import LangChainTracer
from langchain_core.tracers.run_collector import RunCollectorCallbackHandler
from langchain_core.tracers.schemas import TracerSessionV1
if TYPE_CHECKING:
    from langsmith import Client as LangSmithClient
    from langchain_core.callbacks.base import BaseCallbackHandler, Callbacks
    from langchain_core.callbacks.manager import AsyncCallbackManager, CallbackManager
# for backwards partial compatibility if this is imported by users but unused
tracing_callback_var: Any = None
tracing_v2_callback_var: ContextVar[Optional[LangChainTracer]] = ContextVar(
    "tracing_callback_v2", default=None
)
run_collector_var: ContextVar[Optional[RunCollectorCallbackHandler]] = ContextVar(
    "run_collector", default=None
)
[docs]@contextmanager
def tracing_enabled(
    session_name: str = "default",
) -> Generator[TracerSessionV1, None, None]:
    """Throw an error because this has been replaced by tracing_v2_enabled."""
    raise RuntimeError(
        "tracing_enabled is no longer supported. Please use tracing_enabled_v2 instead."
    )
[docs]@contextmanager
def tracing_v2_enabled(
    project_name: Optional[str] = None,
    *,
project_name: Optional[str] = None,
    *,
    example_id: Optional[Union[str, UUID]] = None,
    tags: Optional[List[str]] = None,
    client: Optional[LangSmithClient] = None,
) -> Generator[LangChainTracer, None, None]:
    """Instruct LangChain to log all runs in context to LangSmith.
    Args:
        project_name (str, optional): The name of the project.
            Defaults to "default".
        example_id (str or UUID, optional): The ID of the example.
            Defaults to None.
        tags (List[str], optional): The tags to add to the run.
            Defaults to None.
        client (LangSmithClient, optional): The client of the langsmith.
            Defaults to None.
    Returns:
        None
    Example:
        >>> with tracing_v2_enabled():
        ...     # LangChain code will automatically be traced
        You can use this to fetch the LangSmith run URL:
        >>> with tracing_v2_enabled() as cb:
        ...     chain.invoke("foo")
        ...     run_url = cb.get_run_url()
    """
    if isinstance(example_id, str):
        example_id = UUID(example_id)
    cb = LangChainTracer(
        example_id=example_id,
        project_name=project_name,
        tags=tags,
        client=client,
    )
    try:
        tracing_v2_callback_var.set(cb)
        yield cb
    finally:
        tracing_v2_callback_var.set(None)
[docs]@contextmanager
def collect_runs() -> Generator[RunCollectorCallbackHandler, None, None]:
    """Collect all run traces in context.
    Returns:
"""Collect all run traces in context.
    Returns:
        run_collector.RunCollectorCallbackHandler: The run collector callback handler.
    Example:
        >>> with collect_runs() as runs_cb:
                chain.invoke("foo")
                run_id = runs_cb.traced_runs[0].id
    """
    cb = RunCollectorCallbackHandler()
    run_collector_var.set(cb)
    yield cb
    run_collector_var.set(None)
def _get_trace_callbacks(
    project_name: Optional[str] = None,
    example_id: Optional[Union[str, UUID]] = None,
    callback_manager: Optional[Union[CallbackManager, AsyncCallbackManager]] = None,
) -> Callbacks:
    if _tracing_v2_is_enabled():
        project_name_ = project_name or _get_tracer_project()
        tracer = tracing_v2_callback_var.get() or LangChainTracer(
            project_name=project_name_,
            example_id=example_id,
        )
        if callback_manager is None:
            from langchain_core.callbacks.base import Callbacks
            cb = cast(Callbacks, [tracer])
        else:
            if not any(
                isinstance(handler, LangChainTracer)
                for handler in callback_manager.handlers
            ):
                callback_manager.add_handler(tracer, True)
                # If it already has a LangChainTracer, we don't need to add another one.
                # this would likely mess up the trace hierarchy.
            cb = callback_manager
    else:
        cb = None
    return cb
def _tracing_v2_is_enabled() -> bool:
    if tracing_v2_callback_var.get() is not None:
        return True
    return ls_utils.tracing_is_enabled()
def _get_tracer_project() -> str:
def _get_tracer_project() -> str:
    run_tree = get_run_tree_context()
    return getattr(
        run_tree,
        "session_name",
        getattr(
            # Note, if people are trying to nest @traceable functions and the
            # tracing_v2_enabled context manager, this will likely mess up the
            # tree structure.
            tracing_v2_callback_var.get(),
            "project",
            # Have to set this to a string even though it always will return
            # a string because `get_tracer_project` technically can return
            # None, but only when a specific argument is supplied.
            # Therefore, this just tricks the mypy type checker
            str(ls_utils.get_tracer_project()),
        ),
    )
_configure_hooks: List[
    Tuple[
        ContextVar[Optional[BaseCallbackHandler]],
        bool,
        Optional[Type[BaseCallbackHandler]],
        Optional[str],
    ]
] = []
[docs]def register_configure_hook(
    context_var: ContextVar[Optional[Any]],
    inheritable: bool,
    handle_class: Optional[Type[BaseCallbackHandler]] = None,
    env_var: Optional[str] = None,
) -> None:
    """Register a configure hook.
    Args:
        context_var (ContextVar[Optional[Any]]): The context variable.
        inheritable (bool): Whether the context variable is inheritable.
        handle_class (Optional[Type[BaseCallbackHandler]], optional):
          The callback handler class. Defaults to None.
        env_var (Optional[str], optional): The environment variable. Defaults to None.
    Raises:
        ValueError: If env_var is set, handle_class must also be set
          to a non-None value.
    """
to a non-None value.
    """
    if env_var is not None and handle_class is None:
        raise ValueError(
            "If env_var is set, handle_class must also be set to a non-None value."
        )
    from langchain_core.callbacks.base import BaseCallbackHandler
    _configure_hooks.append(
        (
            # the typings of ContextVar do not have the generic arg set as covariant
            # so we have to cast it
            cast(ContextVar[Optional[BaseCallbackHandler]], context_var),
            inheritable,
            handle_class,
            env_var,
        )
    )
register_configure_hook(run_collector_var, False)
Source code for langchain_core.tracers.event_stream
"""Internal tracer to power the event stream API."""
from __future__ import annotations
import asyncio
import logging
from typing import (
    TYPE_CHECKING,
    Any,
    AsyncIterator,
    Dict,
    Iterator,
    List,
    Optional,
    Sequence,
    TypeVar,
    Union,
    cast,
)
from uuid import UUID, uuid4
from typing_extensions import NotRequired, TypedDict
from langchain_core.callbacks.base import AsyncCallbackHandler
from langchain_core.messages import AIMessageChunk, BaseMessage, BaseMessageChunk
from langchain_core.outputs import (
    ChatGenerationChunk,
    GenerationChunk,
    LLMResult,
)
from langchain_core.runnables.schema import EventData, StreamEvent
from langchain_core.runnables.utils import (
    Input,
    Output,
    _RootEventFilter,
)
from langchain_core.tracers._streaming import _StreamingCallbackHandler
from langchain_core.tracers.log_stream import LogEntry
from langchain_core.tracers.memory_stream import _MemoryStream
from langchain_core.utils.aiter import py_anext
if TYPE_CHECKING:
    from langchain_core.documents import Document
    from langchain_core.runnables import Runnable, RunnableConfig
logger = logging.getLogger(__name__)
[docs]class RunInfo(TypedDict):
    """Information about a run."""
    name: str
    tags: List[str]
    metadata: Dict[str, Any]
    run_type: str
    inputs: NotRequired[Any]
    parent_run_id: Optional[UUID]
def _assign_name(name: Optional[str], serialized: Dict[str, Any]) -> str:
    """Assign a name to a run."""
    if name is not None:
        return name
if name is not None:
        return name
    if "name" in serialized:
        return serialized["name"]
    elif "id" in serialized:
        return serialized["id"][-1]
    return "Unnamed"
T = TypeVar("T")
class _AstreamEventsCallbackHandler(AsyncCallbackHandler, _StreamingCallbackHandler):
    """An implementation of an async callback handler for astream events."""
    def __init__(
        self,
        *args: Any,
        include_names: Optional[Sequence[str]] = None,
        include_types: Optional[Sequence[str]] = None,
        include_tags: Optional[Sequence[str]] = None,
        exclude_names: Optional[Sequence[str]] = None,
        exclude_types: Optional[Sequence[str]] = None,
        exclude_tags: Optional[Sequence[str]] = None,
        **kwargs: Any,
    ) -> None:
        """Initialize the tracer."""
        super().__init__(*args, **kwargs)
        # Map of run ID to run info.
        # the entry corresponding to a given run id is cleaned
        # up when each corresponding run ends.
        self.run_map: Dict[UUID, RunInfo] = {}
        # The callback event that corresponds to the end of a parent run
        # may be invoked BEFORE the callback event that corresponds to the end
        # of a child run, which results in clean up of run_map.
        # So we keep track of the mapping between children and parent run IDs
        # in a separate container. This container is GCed when the tracer is GCed.
        self.parent_map: Dict[UUID, Optional[UUID]] = {}
        self.is_tapped: Dict[UUID, Any] = {}
        # Filter which events will be sent over the queue.
# Filter which events will be sent over the queue.
        self.root_event_filter = _RootEventFilter(
            include_names=include_names,
            include_types=include_types,
            include_tags=include_tags,
            exclude_names=exclude_names,
            exclude_types=exclude_types,
            exclude_tags=exclude_tags,
        )
        loop = asyncio.get_event_loop()
        memory_stream = _MemoryStream[StreamEvent](loop)
        self.send_stream = memory_stream.get_send_stream()
        self.receive_stream = memory_stream.get_receive_stream()
    def _get_parent_ids(self, run_id: UUID) -> List[str]:
        """Get the parent IDs of a run (non-recursively) cast to strings."""
        parent_ids = []
        while parent_id := self.parent_map.get(run_id):
            str_parent_id = str(parent_id)
            if str_parent_id in parent_ids:
                raise AssertionError(
                    f"Parent ID {parent_id} is already in the parent_ids list. "
                    f"This should never happen."
                )
            parent_ids.append(str_parent_id)
            run_id = parent_id
        # Return the parent IDs in reverse order, so that the first
        # parent ID is the root and the last ID is the immediate parent.
        return parent_ids[::-1]
    def _send(self, event: StreamEvent, event_type: str) -> None:
        """Send an event to the stream."""
        if self.root_event_filter.include_event(event, event_type):
            self.send_stream.send_nowait(event)
    def __aiter__(self) -> AsyncIterator[Any]:
        """Iterate over the receive stream."""
        return self.receive_stream.__aiter__()
    async def tap_output_aiter(
return self.receive_stream.__aiter__()
    async def tap_output_aiter(
        self, run_id: UUID, output: AsyncIterator[T]
    ) -> AsyncIterator[T]:
        """Tap the output aiter."""
        sentinel = object()
        # atomic check and set
        tap = self.is_tapped.setdefault(run_id, sentinel)
        # wait for first chunk
        first = await py_anext(output, default=sentinel)
        if first is sentinel:
            return
        # get run info
        run_info = self.run_map.get(run_id)
        if run_info is None:
            # run has finished, don't issue any stream events
            yield cast(T, first)
            return
        if tap is sentinel:
            # if we are the first to tap, issue stream events
            event: StreamEvent = {
                "event": f"on_{run_info['run_type']}_stream",
                "run_id": str(run_id),
                "name": run_info["name"],
                "tags": run_info["tags"],
                "metadata": run_info["metadata"],
                "data": {},
                "parent_ids": self._get_parent_ids(run_id),
            }
            self._send({**event, "data": {"chunk": first}}, run_info["run_type"])
            yield cast(T, first)
            # consume the rest of the output
            async for chunk in output:
                self._send(
                    {**event, "data": {"chunk": chunk}},
                    run_info["run_type"],
                )
                yield chunk
        else:
            # otherwise just pass through
            yield cast(T, first)
            # consume the rest of the output
            async for chunk in output:
                yield chunk
async for chunk in output:
                yield chunk
    def tap_output_iter(self, run_id: UUID, output: Iterator[T]) -> Iterator[T]:
        """Tap the output aiter."""
        sentinel = object()
        # atomic check and set
        tap = self.is_tapped.setdefault(run_id, sentinel)
        # wait for first chunk
        first = next(output, sentinel)
        if first is sentinel:
            return
        # get run info
        run_info = self.run_map.get(run_id)
        if run_info is None:
            # run has finished, don't issue any stream events
            yield cast(T, first)
            return
        if tap is sentinel:
            # if we are the first to tap, issue stream events
            event: StreamEvent = {
                "event": f"on_{run_info['run_type']}_stream",
                "run_id": str(run_id),
                "name": run_info["name"],
                "tags": run_info["tags"],
                "metadata": run_info["metadata"],
                "data": {},
                "parent_ids": self._get_parent_ids(run_id),
            }
            self._send({**event, "data": {"chunk": first}}, run_info["run_type"])
            yield cast(T, first)
            # consume the rest of the output
            for chunk in output:
                self._send(
                    {**event, "data": {"chunk": chunk}},
                    run_info["run_type"],
                )
                yield chunk
        else:
            # otherwise just pass through
            yield cast(T, first)
            # consume the rest of the output
            for chunk in output:
                yield chunk
    def _write_run_start_info(
        self,
        run_id: UUID,
        *,
self,
        run_id: UUID,
        *,
        tags: Optional[List[str]],
        metadata: Optional[Dict[str, Any]],
        parent_run_id: Optional[UUID],
        name_: str,
        run_type: str,
        **kwargs: Any,
    ) -> None:
        """Update the run info."""
        info: RunInfo = {
            "tags": tags or [],
            "metadata": metadata or {},
            "name": name_,
            "run_type": run_type,
            "parent_run_id": parent_run_id,
        }
        if "inputs" in kwargs:
            # Handle inputs in a special case to allow inputs to be an
            # optionally provided and distinguish between missing value
            # vs. None value.
            info["inputs"] = kwargs["inputs"]
        self.run_map[run_id] = info
        self.parent_map[run_id] = parent_run_id
    async def on_chat_model_start(
        self,
        serialized: Dict[str, Any],
        messages: List[List[BaseMessage]],
        *,
        run_id: UUID,
        tags: Optional[List[str]] = None,
        parent_run_id: Optional[UUID] = None,
        metadata: Optional[Dict[str, Any]] = None,
        name: Optional[str] = None,
        **kwargs: Any,
    ) -> None:
        """Start a trace for an LLM run."""
        name_ = _assign_name(name, serialized)
        run_type = "chat_model"
        self._write_run_start_info(
            run_id,
            tags=tags,
            metadata=metadata,
            parent_run_id=parent_run_id,
            name_=name_,
            run_type=run_type,
name_=name_,
            run_type=run_type,
            inputs={"messages": messages},
        )
        self._send(
            {
                "event": "on_chat_model_start",
                "data": {
                    "input": {"messages": messages},
                },
                "name": name_,
                "tags": tags or [],
                "run_id": str(run_id),
                "metadata": metadata or {},
                "parent_ids": self._get_parent_ids(run_id),
            },
            run_type,
        )
    async def on_llm_start(
        self,
        serialized: Dict[str, Any],
        prompts: List[str],
        *,
        run_id: UUID,
        tags: Optional[List[str]] = None,
        parent_run_id: Optional[UUID] = None,
        metadata: Optional[Dict[str, Any]] = None,
        name: Optional[str] = None,
        **kwargs: Any,
    ) -> None:
        """Start a trace for an LLM run."""
        name_ = _assign_name(name, serialized)
        run_type = "llm"
        self._write_run_start_info(
            run_id,
            tags=tags,
            metadata=metadata,
            parent_run_id=parent_run_id,
            name_=name_,
            run_type=run_type,
            inputs={"prompts": prompts},
        )
        self._send(
            {
                "event": "on_llm_start",
                "data": {
                    "input": {
                        "prompts": prompts,
                    }
                },
                "name": name_,
                "tags": tags or [],
                "run_id": str(run_id),
                "metadata": metadata or {},
"run_id": str(run_id),
                "metadata": metadata or {},
                "parent_ids": self._get_parent_ids(run_id),
            },
            run_type,
        )
    async def on_llm_new_token(
        self,
        token: str,
        *,
        chunk: Optional[Union[GenerationChunk, ChatGenerationChunk]] = None,
        run_id: UUID,
        parent_run_id: Optional[UUID] = None,
        **kwargs: Any,
    ) -> None:
        """Run on new LLM token. Only available when streaming is enabled."""
        run_info = self.run_map.get(run_id)
        chunk_: Union[GenerationChunk, BaseMessageChunk]
        if run_info is None:
            raise AssertionError(f"Run ID {run_id} not found in run map.")
        if self.is_tapped.get(run_id):
            return
        if run_info["run_type"] == "chat_model":
            event = "on_chat_model_stream"
            if chunk is None:
                chunk_ = AIMessageChunk(content=token)
            else:
                chunk_ = cast(ChatGenerationChunk, chunk).message
        elif run_info["run_type"] == "llm":
            event = "on_llm_stream"
            if chunk is None:
                chunk_ = GenerationChunk(text=token)
            else:
                chunk_ = cast(GenerationChunk, chunk)
        else:
            raise ValueError(f"Unexpected run type: {run_info['run_type']}")
        self._send(
            {
                "event": event,
                "data": {
                    "chunk": chunk_,
                },
                "run_id": str(run_id),
                "name": run_info["name"],
"run_id": str(run_id),
                "name": run_info["name"],
                "tags": run_info["tags"],
                "metadata": run_info["metadata"],
                "parent_ids": self._get_parent_ids(run_id),
            },
            run_info["run_type"],
        )
    async def on_llm_end(
        self, response: LLMResult, *, run_id: UUID, **kwargs: Any
    ) -> None:
        """End a trace for an LLM run."""
        run_info = self.run_map.pop(run_id)
        inputs_ = run_info["inputs"]
        generations: Union[List[List[GenerationChunk]], List[List[ChatGenerationChunk]]]
        output: Union[dict, BaseMessage] = {}
        if run_info["run_type"] == "chat_model":
            generations = cast(List[List[ChatGenerationChunk]], response.generations)
            for gen in generations:
                if output != {}:
                    break
                for chunk in gen:
                    output = chunk.message
                    break
            event = "on_chat_model_end"
        elif run_info["run_type"] == "llm":
            generations = cast(List[List[GenerationChunk]], response.generations)
            output = {
                "generations": [
                    [
                        {
                            "text": chunk.text,
                            "generation_info": chunk.generation_info,
                            "type": chunk.type,
                        }
                        for chunk in gen
                    ]
                    for gen in generations
                ],
                "llm_output": response.llm_output,
            }
            event = "on_llm_end"
        else:
            raise ValueError(f"Unexpected run type: {run_info['run_type']}")
        self._send(
            {
                "event": event,