viruPVS/langchain-docs huggingface dataset:
self._parent_container = parent_container
        self._history_parent = parent_container.container()
        self._history_container: Optional[MutableExpander] = None
        self._current_thought: Optional[LLMThought] = None
        self._completed_thoughts: List[LLMThought] = []
        self._max_thought_containers = max(max_thought_containers, 1)
        self._expand_new_thoughts = expand_new_thoughts
        self._collapse_completed_thoughts = collapse_completed_thoughts
        self._thought_labeler = thought_labeler or LLMThoughtLabeler()
    def _require_current_thought(self) -> LLMThought:
        """Return our current LLMThought. Raise an error if we have no current
        thought.
        """
        if self._current_thought is None:
            raise RuntimeError("Current LLMThought is unexpectedly None!")
        return self._current_thought
    def _get_last_completed_thought(self) -> Optional[LLMThought]:
        """Return our most recent completed LLMThought, or None if we don't have one."""
        if len(self._completed_thoughts) > 0:
            return self._completed_thoughts[len(self._completed_thoughts) - 1]
        return None
    @property
    def _num_thought_containers(self) -> int:
        """The number of 'thought containers' we're currently showing: the
        number of completed thought containers, the history container (if it exists),
        and the current thought container (if it exists).
        """
        count = len(self._completed_thoughts)
        if self._history_container is not None:
            count += 1
        if self._current_thought is not None:
            count += 1
if self._current_thought is not None:
            count += 1
        return count
    def _complete_current_thought(self, final_label: Optional[str] = None) -> None:
        """Complete the current thought, optionally assigning it a new label.
        Add it to our _completed_thoughts list.
        """
        thought = self._require_current_thought()
        thought.complete(final_label)
        self._completed_thoughts.append(thought)
        self._current_thought = None
    def _prune_old_thought_containers(self) -> None:
        """If we have too many thoughts onscreen, move older thoughts to the
        'history container.'
        """
        while (
            self._num_thought_containers > self._max_thought_containers
            and len(self._completed_thoughts) > 0
        ):
            # Create our history container if it doesn't exist, and if
            # max_thought_containers is > 1. (if max_thought_containers is 1, we don't
            # have room to show history.)
            if self._history_container is None and self._max_thought_containers > 1:
                self._history_container = MutableExpander(
                    self._history_parent,
                    label=self._thought_labeler.get_history_label(),
                    expanded=False,
                )
            oldest_thought = self._completed_thoughts.pop(0)
            if self._history_container is not None:
                self._history_container.markdown(oldest_thought.container.label)
                self._history_container.append_copy(oldest_thought.container)
            oldest_thought.clear()
[docs]    def on_llm_start(
        self, serialized: Dict[str, Any], prompts: List[str], **kwargs: Any
self, serialized: Dict[str, Any], prompts: List[str], **kwargs: Any
    ) -> None:
        if self._current_thought is None:
            self._current_thought = LLMThought(
                parent_container=self._parent_container,
                expanded=self._expand_new_thoughts,
                collapse_on_complete=self._collapse_completed_thoughts,
                labeler=self._thought_labeler,
            )
        self._current_thought.on_llm_start(serialized, prompts)
        # We don't prune_old_thought_containers here, because our container won't
        # be visible until it has a child.
[docs]    def on_llm_new_token(self, token: str, **kwargs: Any) -> None:
        self._require_current_thought().on_llm_new_token(token, **kwargs)
        self._prune_old_thought_containers()
[docs]    def on_llm_end(self, response: LLMResult, **kwargs: Any) -> None:
        self._require_current_thought().on_llm_end(response, **kwargs)
        self._prune_old_thought_containers()
[docs]    def on_llm_error(self, error: BaseException, **kwargs: Any) -> None:
        self._require_current_thought().on_llm_error(error, **kwargs)
        self._prune_old_thought_containers()
[docs]    def on_tool_start(
        self, serialized: Dict[str, Any], input_str: str, **kwargs: Any
    ) -> None:
        self._require_current_thought().on_tool_start(serialized, input_str, **kwargs)
        self._prune_old_thought_containers()
[docs]    def on_tool_end(
        self,
        output: str,
[docs]    def on_tool_end(
        self,
        output: str,
        color: Optional[str] = None,
        observation_prefix: Optional[str] = None,
        llm_prefix: Optional[str] = None,
        **kwargs: Any,
    ) -> None:
        self._require_current_thought().on_tool_end(
            output, color, observation_prefix, llm_prefix, **kwargs
        )
        self._complete_current_thought()
[docs]    def on_tool_error(self, error: BaseException, **kwargs: Any) -> None:
        self._require_current_thought().on_tool_error(error, **kwargs)
        self._prune_old_thought_containers()
[docs]    def on_text(
        self,
        text: str,
        color: Optional[str] = None,
        end: str = "",
        **kwargs: Any,
    ) -> None:
        pass
[docs]    def on_chain_start(
        self, serialized: Dict[str, Any], inputs: Dict[str, Any], **kwargs: Any
    ) -> None:
        pass
[docs]    def on_chain_end(self, outputs: Dict[str, Any], **kwargs: Any) -> None:
        pass
[docs]    def on_chain_error(self, error: BaseException, **kwargs: Any) -> None:
        pass
[docs]    def on_agent_action(
        self, action: AgentAction, color: Optional[str] = None, **kwargs: Any
    ) -> Any:
        self._require_current_thought().on_agent_action(action, color, **kwargs)
        self._prune_old_thought_containers()
[docs]    def on_agent_finish(
Source code for langchain_community.callbacks.tracers.comet
from types import ModuleType, SimpleNamespace
from typing import TYPE_CHECKING, Any, Callable, Dict
from langchain_core.tracers import BaseTracer
if TYPE_CHECKING:
    from uuid import UUID
    from comet_llm import Span
    from comet_llm.chains.chain import Chain
    from langchain_community.callbacks.tracers.schemas import Run
def _get_run_type(run: "Run") -> str:
    if isinstance(run.run_type, str):
        return run.run_type
    elif hasattr(run.run_type, "value"):
        return run.run_type.value
    else:
        return str(run.run_type)
[docs]def import_comet_llm_api() -> SimpleNamespace:
    """Import comet_llm api and raise an error if it is not installed."""
    try:
        from comet_llm import (
            experiment_info,  # noqa: F401
            flush,  # noqa: F401
        )
        from comet_llm.chains import api as chain_api  # noqa: F401
        from comet_llm.chains import (
            chain,  # noqa: F401
            span,  # noqa: F401
        )
    except ImportError:
        raise ImportError(
            "To use the CometTracer you need to have the "
            "`comet_llm>=2.0.0` python package installed. Please install it with"
            " `pip install -U comet_llm`"
        )
    return SimpleNamespace(
        chain=chain,
        span=span,
        chain_api=chain_api,
        experiment_info=experiment_info,
        flush=flush,
    )
[docs]class CometTracer(BaseTracer):
)
[docs]class CometTracer(BaseTracer):
[docs]    def __init__(self, **kwargs: Any) -> None:
        super().__init__(**kwargs)
        self._span_map: Dict["UUID", "Span"] = {}
        self._chains_map: Dict["UUID", "Chain"] = {}
        self._initialize_comet_modules()
    def _initialize_comet_modules(self) -> None:
        comet_llm_api = import_comet_llm_api()
        self._chain: ModuleType = comet_llm_api.chain
        self._span: ModuleType = comet_llm_api.span
        self._chain_api: ModuleType = comet_llm_api.chain_api
        self._experiment_info: ModuleType = comet_llm_api.experiment_info
        self._flush: Callable[[], None] = comet_llm_api.flush
    def _persist_run(self, run: "Run") -> None:
        chain_ = self._chains_map[run.id]
        chain_.set_outputs(outputs=run.outputs)
        self._chain_api.log_chain(chain_)
    def _process_start_trace(self, run: "Run") -> None:
        if not run.parent_run_id:
            # This is the first run, which maps to a chain
            chain_: "Chain" = self._chain.Chain(
                inputs=run.inputs,
                metadata=None,
                experiment_info=self._experiment_info.get(),
            )
            self._chains_map[run.id] = chain_
        else:
            span: "Span" = self._span.Span(
                inputs=run.inputs,
                category=_get_run_type(run),
                metadata=run.extra,
                name=run.name,
            )
            span.__api__start__(self._chains_map[run.parent_run_id])
span.__api__start__(self._chains_map[run.parent_run_id])
            self._chains_map[run.id] = self._chains_map[run.parent_run_id]
            self._span_map[run.id] = span
    def _process_end_trace(self, run: "Run") -> None:
        if not run.parent_run_id:
            pass
            # Langchain will call _persist_run for us
        else:
            span = self._span_map[run.id]
            span.set_outputs(outputs=run.outputs)
            span.__api__end__()
[docs]    def flush(self) -> None:
        self._flush()
    def _on_llm_start(self, run: "Run") -> None:
        """Process the LLM Run upon start."""
        self._process_start_trace(run)
    def _on_llm_end(self, run: "Run") -> None:
        """Process the LLM Run."""
        self._process_end_trace(run)
    def _on_llm_error(self, run: "Run") -> None:
        """Process the LLM Run upon error."""
        self._process_end_trace(run)
    def _on_chain_start(self, run: "Run") -> None:
        """Process the Chain Run upon start."""
        self._process_start_trace(run)
    def _on_chain_end(self, run: "Run") -> None:
        """Process the Chain Run."""
        self._process_end_trace(run)
    def _on_chain_error(self, run: "Run") -> None:
        """Process the Chain Run upon error."""
        self._process_end_trace(run)
    def _on_tool_start(self, run: "Run") -> None:
        """Process the Tool Run upon start."""
        self._process_start_trace(run)
"""Process the Tool Run upon start."""
        self._process_start_trace(run)
    def _on_tool_end(self, run: "Run") -> None:
        """Process the Tool Run."""
        self._process_end_trace(run)
    def _on_tool_error(self, run: "Run") -> None:
        """Process the Tool Run upon error."""
        self._process_end_trace(run)
Source code for langchain_community.callbacks.tracers.wandb
"""A Tracer Implementation that records activity to Weights & Biases."""
from __future__ import annotations
import json
from typing import (
    TYPE_CHECKING,
    Any,
    Dict,
    List,
    Optional,
    Sequence,
    Tuple,
    TypedDict,
    Union,
)
from langchain_core.tracers.base import BaseTracer
from langchain_core.tracers.schemas import Run
if TYPE_CHECKING:
    from wandb import Settings as WBSettings
    from wandb.sdk.data_types.trace_tree import Span
    from wandb.sdk.lib.paths import StrPath
    from wandb.wandb_run import Run as WBRun
PRINT_WARNINGS = True
def _serialize_io(run_inputs: Optional[dict]) -> dict:
    if not run_inputs:
        return {}
    from google.protobuf.json_format import MessageToJson
    from google.protobuf.message import Message
    serialized_inputs = {}
    for key, value in run_inputs.items():
        if isinstance(value, Message):
            serialized_inputs[key] = MessageToJson(value)
        elif key == "input_documents":
            serialized_inputs.update(
                {f"input_document_{i}": doc.json() for i, doc in enumerate(value)}
            )
        else:
            serialized_inputs[key] = value
    return serialized_inputs
[docs]class RunProcessor:
    """Handles the conversion of a LangChain Runs into a WBTraceTree."""
[docs]    def __init__(self, wandb_module: Any, trace_module: Any):
        self.wandb = wandb_module
        self.trace_tree = trace_module
[docs]    def process_span(self, run: Run) -> Optional["Span"]:
[docs]    def process_span(self, run: Run) -> Optional["Span"]:
        """Converts a LangChain Run into a W&B Trace Span.
        :param run: The LangChain Run to convert.
        :return: The converted W&B Trace Span.
        """
        try:
            span = self._convert_lc_run_to_wb_span(run)
            return span
        except Exception as e:
            if PRINT_WARNINGS:
                self.wandb.termwarn(
                    f"Skipping trace saving - unable to safely convert LangChain Run "
                    f"into W&B Trace due to: {e}"
                )
            return None
    def _convert_run_to_wb_span(self, run: Run) -> "Span":
        """Base utility to create a span from a run.
        :param run: The run to convert.
        :return: The converted Span.
        """
        attributes = {**run.extra} if run.extra else {}
        attributes["execution_order"] = run.execution_order
        return self.trace_tree.Span(
            span_id=str(run.id) if run.id is not None else None,
            name=run.name,
            start_time_ms=int(run.start_time.timestamp() * 1000),
            end_time_ms=int(run.end_time.timestamp() * 1000)
            if run.end_time is not None
            else None,
            status_code=self.trace_tree.StatusCode.SUCCESS
            if run.error is None
            else self.trace_tree.StatusCode.ERROR,
            status_message=run.error,
            attributes=attributes,
        )
    def _convert_llm_run_to_wb_span(self, run: Run) -> "Span":
        """Converts a LangChain LLM Run into a W&B Trace Span.
"""Converts a LangChain LLM Run into a W&B Trace Span.
        :param run: The LangChain LLM Run to convert.
        :return: The converted W&B Trace Span.
        """
        base_span = self._convert_run_to_wb_span(run)
        if base_span.attributes is None:
            base_span.attributes = {}
        base_span.attributes["llm_output"] = (run.outputs or {}).get("llm_output", {})
        base_span.results = [
            self.trace_tree.Result(
                inputs={"prompt": prompt},
                outputs={
                    f"gen_{g_i}": gen["text"]
                    for g_i, gen in enumerate(run.outputs["generations"][ndx])
                }
                if (
                    run.outputs is not None
                    and len(run.outputs["generations"]) > ndx
                    and len(run.outputs["generations"][ndx]) > 0
                )
                else None,
            )
            for ndx, prompt in enumerate(run.inputs["prompts"] or [])
        ]
        base_span.span_kind = self.trace_tree.SpanKind.LLM
        return base_span
    def _convert_chain_run_to_wb_span(self, run: Run) -> "Span":
        """Converts a LangChain Chain Run into a W&B Trace Span.
        :param run: The LangChain Chain Run to convert.
        :return: The converted W&B Trace Span.
        """
        base_span = self._convert_run_to_wb_span(run)
        base_span.results = [
            self.trace_tree.Result(
                inputs=_serialize_io(run.inputs), outputs=_serialize_io(run.outputs)
            )
        ]
        base_span.child_spans = [
            self._convert_lc_run_to_wb_span(child_run) for child_run in run.child_runs
]
        base_span.span_kind = (
            self.trace_tree.SpanKind.AGENT
            if "agent" in run.name.lower()
            else self.trace_tree.SpanKind.CHAIN
        )
        return base_span
    def _convert_tool_run_to_wb_span(self, run: Run) -> "Span":
        """Converts a LangChain Tool Run into a W&B Trace Span.
        :param run: The LangChain Tool Run to convert.
        :return: The converted W&B Trace Span.
        """
        base_span = self._convert_run_to_wb_span(run)
        base_span.results = [
            self.trace_tree.Result(
                inputs=_serialize_io(run.inputs), outputs=_serialize_io(run.outputs)
            )
        ]
        base_span.child_spans = [
            self._convert_lc_run_to_wb_span(child_run) for child_run in run.child_runs
        ]
        base_span.span_kind = self.trace_tree.SpanKind.TOOL
        return base_span
    def _convert_lc_run_to_wb_span(self, run: Run) -> "Span":
        """Utility to convert any generic LangChain Run into a W&B Trace Span.
        :param run: The LangChain Run to convert.
        :return: The converted W&B Trace Span.
        """
        if run.run_type == "llm":
            return self._convert_llm_run_to_wb_span(run)
        elif run.run_type == "chain":
            return self._convert_chain_run_to_wb_span(run)
        elif run.run_type == "tool":
            return self._convert_tool_run_to_wb_span(run)
        else:
            return self._convert_run_to_wb_span(run)
else:
            return self._convert_run_to_wb_span(run)
[docs]    def process_model(self, run: Run) -> Optional[Dict[str, Any]]:
        """Utility to process a run for wandb model_dict serialization.
        :param run: The run to process.
        :return: The convert model_dict to pass to WBTraceTree.
        """
        try:
            data = json.loads(run.json())
            processed = self.flatten_run(data)
            keep_keys = (
                "id",
                "name",
                "serialized",
                "inputs",
                "outputs",
                "parent_run_id",
                "execution_order",
            )
            processed = self.truncate_run_iterative(processed, keep_keys=keep_keys)
            exact_keys, partial_keys = ("lc", "type"), ("api_key",)
            processed = self.modify_serialized_iterative(
                processed, exact_keys=exact_keys, partial_keys=partial_keys
            )
            output = self.build_tree(processed)
            return output
        except Exception as e:
            if PRINT_WARNINGS:
                self.wandb.termwarn(f"WARNING: Failed to serialize model: {e}")
            return None
[docs]    def flatten_run(self, run: Dict[str, Any]) -> List[Dict[str, Any]]:
        """Utility to flatten a nest run object into a list of runs.
        :param run: The base run to flatten.
        :return: The flattened list of runs.
        """
        def flatten(child_runs: List[Dict[str, Any]]) -> List[Dict[str, Any]]:
            """Utility to recursively flatten a list of child runs in a run.
            :param child_runs: The list of child runs to flatten.
:param child_runs: The list of child runs to flatten.
            :return: The flattened list of runs.
            """
            if child_runs is None:
                return []
            result = []
            for item in child_runs:
                child_runs = item.pop("child_runs", [])
                result.append(item)
                result.extend(flatten(child_runs))
            return result
        return flatten([run])
[docs]    def truncate_run_iterative(
        self, runs: List[Dict[str, Any]], keep_keys: Tuple[str, ...] = ()
    ) -> List[Dict[str, Any]]:
        """Utility to truncate a list of runs dictionaries to only keep the specified
            keys in each run.
        :param runs: The list of runs to truncate.
        :param keep_keys: The keys to keep in each run.
        :return: The truncated list of runs.
        """
        def truncate_single(run: Dict[str, Any]) -> Dict[str, Any]:
            """Utility to truncate a single run dictionary to only keep the specified
                keys.
            :param run: The run dictionary to truncate.
            :return: The truncated run dictionary
            """
            new_dict = {}
            for key in run:
                if key in keep_keys:
                    new_dict[key] = run.get(key)
            return new_dict
        return list(map(truncate_single, runs))
[docs]    def modify_serialized_iterative(
        self,
        runs: List[Dict[str, Any]],
        exact_keys: Tuple[str, ...] = (),
        partial_keys: Tuple[str, ...] = (),
    ) -> List[Dict[str, Any]]:
        """Utility to modify the serialized field of a list of runs dictionaries.
"""Utility to modify the serialized field of a list of runs dictionaries.
        removes any keys that match the exact_keys and any keys that contain any of the
        partial_keys.
        recursively moves the dictionaries under the kwargs key to the top level.
        changes the "id" field to a string "_kind" field that tells WBTraceTree how to
        visualize the run. promotes the "serialized" field to the top level.
        :param runs: The list of runs to modify.
        :param exact_keys: A tuple of keys to remove from the serialized field.
        :param partial_keys: A tuple of partial keys to remove from the serialized
            field.
        :return: The modified list of runs.
        """
        def remove_exact_and_partial_keys(obj: Dict[str, Any]) -> Dict[str, Any]:
            """Recursively removes exact and partial keys from a dictionary.
            :param obj: The dictionary to remove keys from.
            :return: The modified dictionary.
            """
            if isinstance(obj, dict):
                obj = {
                    k: v
                    for k, v in obj.items()
                    if k not in exact_keys
                    and not any(partial in k for partial in partial_keys)
                }
                for k, v in obj.items():
                    obj[k] = remove_exact_and_partial_keys(v)
            elif isinstance(obj, list):
                obj = [remove_exact_and_partial_keys(x) for x in obj]
            return obj
        def handle_id_and_kwargs(
            obj: Dict[str, Any], root: bool = False
        ) -> Dict[str, Any]:
            """Recursively handles the id and kwargs fields of a dictionary.
            changes the id field to a string "_kind" field that tells WBTraceTree how
changes the id field to a string "_kind" field that tells WBTraceTree how
            to visualize the run. recursively moves the dictionaries under the kwargs
            key to the top level.
            :param obj: a run dictionary with id and kwargs fields.
            :param root: whether this is the root dictionary or the serialized
                dictionary.
            :return: The modified dictionary.
            """
            if isinstance(obj, dict):
                if ("id" in obj or "name" in obj) and not root:
                    _kind = obj.get("id")
                    if not _kind:
                        _kind = [obj.get("name")]
                    obj["_kind"] = _kind[-1]
                    obj.pop("id", None)
                    obj.pop("name", None)
                    if "kwargs" in obj:
                        kwargs = obj.pop("kwargs")
                        for k, v in kwargs.items():
                            obj[k] = v
                for k, v in obj.items():
                    obj[k] = handle_id_and_kwargs(v)
            elif isinstance(obj, list):
                obj = [handle_id_and_kwargs(x) for x in obj]
            return obj
        def transform_serialized(serialized: Dict[str, Any]) -> Dict[str, Any]:
            """Transforms the serialized field of a run dictionary to be compatible
                with WBTraceTree.
            :param serialized: The serialized field of a run dictionary.
            :return: The transformed serialized field.
            """
            serialized = handle_id_and_kwargs(serialized, root=True)
            serialized = remove_exact_and_partial_keys(serialized)
            return serialized
        def transform_run(run: Dict[str, Any]) -> Dict[str, Any]:
            """Transforms a run dictionary to be compatible with WBTraceTree.