# AUTOGENERATED! DO NOT EDIT! File to edit: ../../nbs/langchain.obs.ipynb.

# %% auto 0
__all__ = ['patch_wandb_tracer_serialize_io']

# %% ../../nbs/langchain.obs.ipynb 3
from typing import Optional

# %% ../../nbs/langchain.obs.ipynb 4
def _serialize_io(run_inputs: Optional[dict]) -> dict:
    """Different to the original function, this also handles the serialization of Pydantic output from llms."""
    if not run_inputs:
        return {}
    from google.protobuf.json_format import MessageToJson
    from google.protobuf.message import Message
    from langchain_core.pydantic_v1 import BaseModel

    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)}
            )
        elif isinstance(value, BaseModel): # this is added
            serialized_inputs[key] = value.json()
        else:
            serialized_inputs[key] = value
    return serialized_inputs

def patch_wandb_tracer_serialize_io():
    from langchain.callbacks.tracers import wandb
    wandb._serialize_io = _serialize_io
