Fix multi-agent state formatting

[logan-markewich]

Fixes #18416

We would actually format the state multiple time 😮‍💨 Whoops! The main cause is that while the agent is running, its actually populating its own scratchpad, and see the llm_input ends up being the same each time, and we format it each time

Fixed + added a PR

[RakeshReddyKondeti]

The fix prevents state duplication by using the formatted_input_with_state flag, ensuring the state is appended only once per workflow run. However, it does not address the issue updating the existing state in the chat history, leaving the "Current State" unchanged across iterations, even when tools or agents are expected to modify it.

---

Issue

While duplicate states are eliminated, the "Current State" remains static throughout iterations, failing to reflect updates made by tools or agents. This results in a chat history that does not accurately represent the progress of the workflow. (I am not sure whether I am still something here.)

Current Implementation Chat History (for example given here)

First Iteration

System prompt
Current State: 
{"research_notes": {}, "report_content": "Not written yet.", "review": "Review required."}
Current Message: Write research notes on example topic

Second Iteration (state unchanged)

System prompt
Current State: 
{"research_notes": {}, "report_content": "Not written yet.", "review": "Review required."}
Current Message: Write research notes on example topic
assistant: none
tool: <tavily search_web output>

Third Iteration (still unchanged)

System prompt
Current State: 
{"research_notes": {}, "report_content": "Not written yet.", "review": "Review required."}
Current Message: Write research notes on example topic
assistant: none
tool: <tavily search_web output>
assistant: none
tool: Notes Recorded

Fourth Iteration (still unchanged)

System prompt
Current State: 
{"research_notes": {}, "report_content": "Not written yet.", "review": "Review required."}
Current Message: Write research notes on example topic
assistant: none
tool: <tavily search_web output>
assistant: none
tool: Notes Recorded
assistant: None
Agent WriteAgent is now handling the request due to the following reason: I have recorded the necessary notes on the history of the web and its developments. Please continue with the current request.

---

Proposed Solution (Atleast this worked for me)

To address the limitations, I used regex-based solution . This approach dynamically detects and updates the state in the chat history, ensuring that the state reflects the latest updates made by tools or agents.

Regex-Based Implementation

  @step
  async def setup_agent(self, ctx: Context, ev: AgentInput) -> AgentSetup:
      """Main agent handling logic."""
      current_agent_name = ev.current_agent_name
      agent = self.agents[current_agent_name]
      llm_input = ev.input

      if agent.system_prompt:
          llm_input = [
              ChatMessage(role="system", content=agent.system_prompt),
              *llm_input,
          ]

      state = await ctx.get("state", default=None)

      import re
      if state:
          # Extract the last message
          for message in llm_input[::-1]:
              if message.role == MessageRole.USER:
                  last_message = message
                  break

          # last_message = llm_input[-1]

          # Generate a regex pattern from self.state_prompt
          state_placeholder = r"(?P<state>.*?)"
          msg_placeholder = r"(?P<msg>.*?)"
          state_prompt_pattern = re.escape(self.state_prompt.template).replace(
              r"\{state\}", state_placeholder
          ).replace(r"\{msg\}", msg_placeholder)

          # Check if the last message matches the state prompt structure
          for block in last_message.blocks[::-1]:
              if isinstance(block, TextBlock):
                  match = re.match(state_prompt_pattern, block.text, re.DOTALL)
                  if match:
                      # Extract the current state and message
                      original_msg = match.group("msg")
                      # Replace the state while keeping the original message
                      block.text = self.state_prompt.format(state=state, msg=original_msg)
                  else:
                      # State is not induced, apply the state prompt
                      block.text = self.state_prompt.format(state=state, msg=block.text)
                  break
                  
      return AgentSetup(
          input=llm_input,
          current_agent_name=ev.current_agent_name,
      )

With the above fix, the chat history:

First Iteration

System prompt
Current State: 
{"research_notes": {}, "report_content": "Not written yet.", "review": "Review required."}
Current Message: Write research notes on example topic

**Second Iteration **

System prompt
Current State: 
{"research_notes": {}, "report_content": "Not written yet.", "review": "Review required."}
Current Message: Write research notes on example topic
assistant: none
tool: <tavily search_web output>

Third Iteration (state updated)

System prompt
Current State: 
{"research_notes": {"notes_title": notes}, "report_content": "Not written yet.", "review": "Review required."}
Current Message: Write research notes on example topic
assistant: none
tool: <tavily search_web output>
assistant: none
tool: Notes Recorded

Fourth Iteration (state updated again)

System prompt
Current State: 
{"research_notes": {"notes_title": notes}, "report_content": Report Content, "review": "Review required."}
Current Message: Write research notes on example topic
assistant: none
tool: <tavily search_web output>
assistant: none
tool: Notes Recorded
assistant: None
Agent WriteAgent is now handling the request due to the following reason: I have recorded the necessary notes on the history of the web and its developments. Please continue with the current request.

---

Conclusion

I believe dynamically updating the chat history to ensure that state updates across iterations will help agents work effectively. My approach dynamically detects and replaces the state in the chat history, making it more meaningful and reflective of the workflow's progress.

[logan-markewich]

@RakeshReddyKondeti this is intended -- so that the agent can see how the state has changed across time.

The state is only appended into each NEW user message -- old messages will not be touched

The agent will see the new state on the next user message

Its debatable if tool outputs should be augmented with the updated state, but for now, I leave that in the hands of the user