Persist State Variables through Chain of Runnables (LCEL)

[juandiegobalbi]

Checked

• I searched existing ideas and did not find a similar one
• I added a very descriptive title
• I've clearly described the feature request and motivation for it

Feature request

I'd like to propose a new feature that would allow the user to store each Runnable's output in a state object, that can later be accessed on any of the following steps in the chain, adding this sort of 'state-machine' behavior to the chain,
similar to how LangGraph works, but on LCEL Runnables.

Motivation

Having worked with plenty of real-world use cases, I have come to realize that most of the time,
the solution you need is neither a complex cyclic graph nor a simple linear chain, but rather a complex, well-structured LCEL chain on Runnables.

That being said, after starting to nest chains within chains, I find that having to wrap step outputs on RunnablePassthroughs to store them and be able to
access them in the next steps is a bit cumbersome and even kind of annoying.

On the other hand, adding more and more keys to the chain's state with RunnablePassthroughs also makes debugging and tracing in LangSmith a bit more difficult,
since each step's output will carry the state object with all the added keys, making it harder to track the actual output of each step.

Motivated on LangGraph, I would like to propose a new feature that would allow the user to store each Runnable's output in a state object, that can later be accessed
on any of the following steps in the chain, adding this sort of 'state-machine' behavior to the chain.

Proposal (If applicable)

I have myself coded a State model to persist state, and would like your feedback and/or some comments on whether you guys have been working on a possible solution for this or not.

The State Base Model:

import typing as t

from langchain.pydantic_v1 import BaseModel
from langchain_core.runnables import Runnable, RunnableLambda


class State(BaseModel):
    class Config:
        extra = "allow"

    def write_with_key(self, data: t.Dict[str, t.Any]) -> t.Dict[str, t.Any]:
        """
        Write the state value to the state object with a key, using the key as the attribute name.

        Args:
            data (Dict[str, Any]): The data to be written to the state object.

        Returns:
            Dict[str, Any]: The data that was written to the state object.
        """

        data_key = list(data.keys())[0]
        setattr(self, data_key.lower(), data[data_key])

        return data[data_key]
    
    def write_without_key(
        self, state_value: t.Union[t.Type[BaseModel], t.Dict[str, t.Any]]
    ) -> t.Union[t.Type[BaseModel], t.Dict[str, t.Any]]:
        """
        write the state value to the state object without a key, using the class name as the attribute name.

        Args:
            state_value (Union[Type[BaseModel], Dict[str, Any]]): The value to be written to the state object.

        Returns:
            Union[Type[BaseModel], Dict[str, Any]]: The value that was written to the state object.
        """

        class_instance = state_value.__class__.__name__.lower()
        if class_instance == "dict":
            for key, value in state_value.items():
                setattr(self, key.lower(), value)
                
        else:
            if not class_instance in self.__dict__.keys():
                raise ValueError(
                    f"'{class_instance}' - State keys that want to be written directly without key-value assignment must be initialized from State definition."
                )

            setattr(self, class_instance.lower(), state_value)

        return state_value
    
    def store(self, key: t.Optional[str] = None) -> Runnable:
        """
        Store the value in the state object as a runnable object.

        If a key is provided, the value will be stored in the state object with the key as the attribute name.
        Otherwise, the value will be stored in the state object with the class name as the
        attribute name (the class name must have been initialized from the State definition).

        Args:
            key (Optional[str], optional): The key to store the value in the state object. Defaults to None.

        Returns:
            Runnable: The runnable object that can be used to store the value in the state object.
        """

        if key:
            return (lambda x: {key: x}) | RunnableLambda(self.write_with_key)
        return RunnableLambda(self.write_without_key)

    def read_from_state(self, key: t.Dict[str, str]) -> t.Any:
        """
        Read the value from the state object.

        Args:
            key (Dict[str, str]): The key to read the value from the state object.

        Returns:
            Any: The value that was read from the state object.
        """

        value = getattr(self, key["key"], None)
        if not value:
            raise KeyError(
                f"Tried to retrieve missing variable from State - {{{key['key']}}}. Make sure it has being stored previously to access it."
            )
        
        return value
    
    def retrieve(self, key: str) -> Runnable:
        """
        Retrieve the value from the state object as a runnable object.

        Args:
            key (str): The key to retrieve the value from the state object.
            
        Returns:
            Runnable: The runnable object that can be used to retrieve the value from the state object.
        """

        return (lambda x: {"key": key}) | RunnableLambda(self.read_from_state)
    

Example Implementation

from dotenv import load_dotenv
from langchain_core.output_parsers import StrOutputParser
from langchain_core.prompts import ChatPromptTemplate
from langchain_core.pydantic_v1 import BaseModel, Field
from langchain_openai.chat_models import ChatOpenAI

from state import State

load_dotenv(override=True)


# Initialize the state object with a number attribute to be used later in the chain
chain_state = State(number=20)

# Initialize the OpenAI model
llm = ChatOpenAI()



# Define the Person Chain
# Given a person's description, it will standardize the data to extract the person's name and age in a 
# 'Person' pydantic model.

class Person(BaseModel):
    """A pydantic model to store a person's attributes"""
    
    name: str = Field(..., description="The name of the person")
    age: int = Field(..., description="The age of the person")


person_prompt = ChatPromptTemplate.from_template(
    "Given the following person description, standardize the data to extract its name and age: \n\n{person_description}"
)

person_chain = (
    person_prompt 
    | llm.with_structured_output(Person) 
    | chain_state.store("person")
)


# Define the Number Chain
# Given the person's age and a number, it will prompt the LLM to add the person's age and the number together.
# We fetch the person's age from the previous chain result (as a regular LCEL chain) and the number from the initial state object.

addition_prompt = ChatPromptTemplate.from_template(
    "What is the result of {num1} + {num2}? Your response should SOLELY consist of the resulting number. No further text is required."
)
addition_chain = (
    {
        "num1": (lambda x: x.age), # From the previous chain
        "num2": chain_state.retrieve("number") # Fetch Number from Initial State
    }
    | addition_prompt
    | llm 
    | StrOutputParser()
)



multiplication_prompt = ChatPromptTemplate.from_template(
    "What is the result of {mult1} * {mult2} * {mult3}? Your response should SOLELY consist of the resulting number. No further text is required."
)

muiltiplication_chain = (
    {
        "mult1": lambda x: x, # From the previous chain
        "mult2": chain_state.retrieve("person") | (lambda x: x.age), # Fetch Person's Age from State (Stored in Person's Chain)
        "mult3": chain_state.retrieve("number") # Fetch Number from Initial State
    }
    | multiplication_prompt
    | llm 
    | StrOutputParser() 
    | chain_state.store("result")
)


chain = person_chain | addition_chain | muiltiplication_chain

person_description = "A person named John who is 30 years old."

result = chain.invoke({"person_description": person_description})


print('\nresult: ', result)
print('\nstate: ', chain_state)


# result:  30000
# number=20 person=Person(name='John', age=30) result='30000'

[minki-j]

I was looking for the same feature, but then discovered StateGraph. Have you checked this?
https://python.langchain.com/docs/langgraph/#define-the-agent-state