Langchain runnable configuration POC

runnable_config.py

import uuid
from typing import Generic, TypeVar

from langchain_core.callbacks import CallbackManagerForChainRun
from langchain_core.runnables import RunnableConfig
from langchain_core.runnables.base import RunnableSerializable
from langchain_core.runnables.utils import ConfigurableField

T = TypeVar("T")


class Wrapper(Generic[T]):
    def __init__(self, value: T):
        self.value = value

    def __repr__(self):
        return f"Wrapper({self.value})"


class Demo(RunnableSerializable[Wrapper[int], Wrapper[str]]):
    factor: int = None

    def invoke(self, input: Wrapper[int], config: RunnableConfig | None = None) -> Wrapper[str]:
        return self._call_with_config(
            self._invoke,
            input,
            config,
        )

    def _invoke(
        self,
        input: Wrapper[int],
        run_manager: CallbackManagerForChainRun,
        config: RunnableConfig,
    ) -> Wrapper[str]:
        return Wrapper(str(input.value * self.factor))

    class Config:
        arbitrary_types_allowed = True


def get_with_config(run_id: str):
    return Demo().configurable_fields(
        factor=ConfigurableField(
            id=f"{run_id}_factor",
            name="Demo",
            description="Demo factor",
        ),
    )


if __name__ == "__main__":
    run_id = uuid.uuid4().hex
    demo = get_with_config(run_id)
    print(f"config specs: {demo.config_specs}\n  json schema: {demo.config_schema().schema_json()}")
    print(f"types:\n  input: {demo.InputType}\n    json schema: {demo.get_input_schema().schema_json()}")
    print(f"types:\n  output: {demo.OutputType}\n    json schema: {demo.get_output_schema().schema_json()}")

    config_dict = {
        f"{run_id}_factor": 2,
    }
    output = demo.with_config(
        configurable=config_dict
    ).invoke(Wrapper(10))
    print("invoke output", output)