Source code for langchain.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.callbacks.tracers.base import BaseTracer
from langchain.callbacks.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: dict) -> dict:
    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"]:
        """Converts a LangChain Run into a W&B Trace Span.