How to define a correct state for multi-agent system

OSS Product Help LangChain
1

I am building a multi-agent app, it has some subagents, for example the following code, and there will more agents like Bulk, Surface and so on.

class MoleculeState(AgentState):
    atoms:Optional[MyAtoms] = Field(description="The atoms object. ", default=None)

@tool 
def from_formula(formula: str, runtime:ToolRuntime) -> Command:
    '''
    Create a molecule by its formula. 
    Args:
        formula: The chemical formula of the molecule.
        runtime: Runtime information.
    '''
    myatoms = None if not formula in g2 else MyAtoms.from_ase_atoms(molecule(formula))
    if myatoms is None:
        return Command(
            update={
                "atoms": None,
                "messages": [
                    ToolMessage(content=f"Failed to create molecule {formula}, it does not exist in ase g2 collection.",
                    tool_call_id=runtime.tool_call_id)
                ]
            }
        )
    else:
        return Command(
            update={
                "atoms": myatoms,
                "messages": [
                    ToolMessage(content=f"Successfully create molecule {formula}.",
                    tool_call_id=runtime.tool_call_id)
                ]
            }
        )

class MoleculeAgent:
    @staticmethod
    def create(model:BaseChatModel):
        system_prompt = '''You are an agent for constructing simple molecules, equipped with two tools:
- from_formula: Create a molecule from its formlua, e.g., CO2, H2O. The formula should exist as a key in ase g2 collection.
- from_smiles: Create a molecule from its SMILES string. 

If the user provide the formlua, try `from_formula` first. If no formlua is provided or `from_formula` fails, try to generate the corresponding SMILES string and call `from_smiles`.
'''
        agent = create_agent(model=model, 
                             system_prompt=system_prompt,
                             tools=[from_formula, from_smiles ],
                             context_schema=AtomsState)
        
        
        return agent

The question is, I want each subagents do their individual tasks and manage histories (e.g. ToolMessage, use these to let subagent retry task or handle errors), but if I define the total state for main agents, there will be multiple messages field for `TotalState`. Is this OK or not ? How to define an appropriate state for the main agent ?

class TotalState:
    state_molecule:MoleculeState
    state_surface:SurfaceState
    state_bulk:BulkState
2

hi @MSJavaScript

Concern 1 - multiple messages fields conflict?

Nested TotalState design trap. Graph state = flat channels, not recursive tree. Nest MoleculeState inside another - inner messages loses add_messages reducer, create_agent reads messages at top level only. Nesting doesn’t dodge collision - wrong model.

Concern 2 - right structure?

Each subagent already own subgraph (create_agent returns compiled graph w/ own messages). Compose, don’t nest TypedDicts. Two patterns (docs):

  • different schema (want isolated history) - wrap subagent in node, transform in/out. Subagent messages stay private
  • shared keys (want one transcript) → add agents as nodes → supervisor style

Parent TotalState holds results (molecule_atoms…) + one coordinator messages. Not nested AgentStates.

Concern 3 - per-subagent retry history?

Free w/ subgraph isolation. Each messages has add_messages → ToolMessage errors append to that subagent’s history → LLM sees fail, falls to from_smiles. Auto-retry via langchain.agents.middleware retry middleware.

Bugs I found:

  1. context_schema=AtomsState WRONG - that’s static runtime ctx. Need state_schema=MoleculeState (where atoms channel lives)
  2. Name mismatch: defined MoleculeState, passed AtomsState
  3. atoms: Optional[MyAtoms] = Field(...) invalid in TypedDict → atoms: NotRequired[MyAtoms]
3

Thanks for your detailed explanation, I’ll reconsider the structure of the data model and agent. :handshake: