[Q] Log protocol buffers with wandb tracing for LangChain

[jlewi]

I'm using Wandb's tracing feature with LangChain. I would like to log all the values as they pass through the steps in the chain. My chain consists of custom pre/post processing steps before/after standard chains like LLMChain. My custom chains use protocol buffers to represent the data. Wandb logs these protocol buffers as part of the trace but doesn't know how to serialize them so they just get logged as non-serializable

Here's a screenshot; Here CodeGenInput is a protocol buffer

Is it possible to configure how Wandb serializes different types? Protocol buffers can be serialized to JSON using json_format.

[rsanandres-wandb]

Hello @jlewi !

Could you send a link to your workspace? I will be reproducing this and reporting this to the relevant team and would like to bring the relevant workspace for them to view.

[jlewi]

Cool; https://wandb.ai/jeremylewi/autobuilder?workspace=user-jeremylewi

[parambharat]

Hi @jlewi, I'm trying to add serialization support for protocol buffers. Can you provide me with a sample script so that I can build and test the implementation ?

[jlewi]

Here are some snippets. I unfortunately haven't had a chance to try these out. This isn't super high priority for me right now so if you need a full repo from me feel free to downgrade or close.

Here's a rudimentary chain that takes as input a protobuf and returns a protobuf.

import langchain
from langchain import chains, chat_models, prompts, schema
from langchain.callbacks import manager
from langchain.chains.base import Chain

import codegen_pb2

class LLMTextToCodeGenOutput(Chain):
    """This chain takes the text output of the LLM which should contain markdown and turns it into patches.

    This also requires the orginal codegen_input which should be passed along
    """

    input_key: str = "text"
    output_key: str = "codegen_output"  #: :meta private:

    # The key to look for the input snippets.
    input_snippets_key = "input_snippets"

    @property
    def input_keys(self) -> List[str]:
        """Will be whatever keys the prompt expects.

        :meta private:
        """
        return [self.input_key, "codegen_input"]

    @property
    def output_keys(self) -> List[str]:
        """Will always return text key.

        :meta private:
        """
        return [self.output_key]

    def _call(
        self,
        inputs: Dict[str, Any],
        run_manager: Optional[manager.CallbackManagerForChainRun] = None,
    ) -> Dict[str, Any]:   
        output = codegen_pb2.CodeGenOutput(text="hello")
        return {self.output_key: output}

proto buff file might look like

syntax = "proto3";

option go_package = "github.com/jlewi/roboweb/go/protos/v1alpha1";

message CodeGenOutput {  
    string text = 1;
}

You'd need to install the protobuf toolchain and generate the protocol buffers. I unfortunately won't have time to put together a full repo.

[ayulockin]

Hey @jlewi , W&B Trace is an old product offering which is not maintained anymore. We were focused on shipping a better tooling called W&B Weave. Team is working on the Langchain integration.

Weave announcement Weave Documentation Closing this issue for now. Do let me know if you need anything more.