StateGraph support

[dorringel]

Problem statement

While the current use of MessageGraph as the state object for all agents allows for some flexibility, some advance use-cases, such as those involving multi-agent collaboration, or task-oriented dialogue, require a more general state object, i.e. StateGraph, as can be seen, for example, in agent supervisor multi-agent langgraph example.

Supporting building agents based on StateGraph rather than MessageGraph will introduce flexibility and greatly extend the functionality as it'll allow building more sophisticated agents, that have the ability to persist and make use of information outside of the conversation messages.

Additionally, as StateGraph is a superset of MessageGraph, it should be relatively straightforward to adapt all existing agents to use StateGraph instead, by defining a base agent state with messages

class AgentStateWithMessages(TypedDict):
    messages: Annotated[Sequence[AnyMessage], operator.add]

    # subclasses can add additional properties as desired

Then adapt all agents to use the StateGraph instead, as can be seen in langgraph examples, i.e:

# async def _get_messages(messages):
async def _get_messages(state):
    messages = state["messages"]
    msgs = []
    for m in messages:
        if isinstance(m, LiberalToolMessage):
            _dict = m.dict()
            _dict["content"] = str(_dict["content"])
            m_c = ToolMessage(**_dict)
            msgs.append(m_c)
        else:
            msgs.append(m)
         
    # return [SystemMessage(content=system_message)] + msgs
    return {"messages": [SystemMessage(content=system_message)] + msgs}

Questions

1. Should the support for StateGraph be instead of, or in addition to MessageGraph? having to support both will introduce forks in many places in the code, having to route the logic based on the input type, i.e.:

if isinstance(inp, Sequence[AnyMessage]):
    # agent is MessageGraph based, do this..
elif isinstance(inp, AgentStateWithMessages):
    # agent is StateGraph based, do that...
else:
    raise ValueError(...)

When alternatively, having canonical support for StateGraph will make that complexity unnecessary.

@nfcampos would appreciate your thoughts on the above. I noticed an initial PR for this a while back so I imagine this has crossed your mind more than once.

Thank you