How to maintain state in a LlamaIndex agent

Agent sessions often need to carry facts, tool results, or workflow counters beyond one prompt. A LlamaIndex agent keeps that session data in a Context object so a later run can continue from values saved during an earlier run.

The Context belongs to the workflow that created it. JsonSerializer converts the object into JSON-compatible data, while Context.from_dict() rebuilds it for another workflow.run(…, ctx=restored_ctx) call using the same workflow definition.

A deterministic FunctionCallingLLM test double exercises both agent turns without hosted model credentials. Its first tool call saves the name, and its post-restoration tool call reads that value from state before producing the later agent response.

1. Install the LlamaIndex core package in the active Python environment.

   $ python3 -m pip install --upgrade llama-index-core

   A project virtual environment prevents this package from changing a shared system Python environment.
   Related: How to install LlamaIndex with pip

2. Create llamaindex-agent-state.py with the foundational state-tool section.

   llamaindex-agent-state.py

   import asyncio
   import json
   from typing import Any
    
   from llama_index.core.agent.workflow import AgentWorkflow
   from llama_index.core.base.llms.types import (
       ChatMessage,
       ChatResponse,
       CompletionResponse,
       LLMMetadata,
       MessageRole,
   )
   from llama_index.core.llms.function_calling import FunctionCallingLLM
   from llama_index.core.llms.llm import ToolSelection
   from llama_index.core.workflow import Context, JsonSerializer
    
   TOOL_CALL_LOG: list[str] = []
    
    
   async def set_name(ctx: Context, name: str) -> str:
       """Save the user's preferred name in agent state."""
       async with ctx.store.edit_state() as ctx_state:
           ctx_state["state"]["name"] = name
       TOOL_CALL_LOG.append(f"set_name(name={name})")
       return name
    
    
   async def read_name(ctx: Context) -> str:
       """Read the user's preferred name from agent state."""
       state = await ctx.store.get("state")
       TOOL_CALL_LOG.append("read_name()")
       return state["name"]

   The workflow stores initial_state under the Context store's state key. edit_state() keeps the name change inside the Context that will be serialized.

3. Add the deterministic model metadata and tool-call parser below read_name().

   class StateAwareMockLLM(FunctionCallingLLM):
       @property
       def metadata(self) -> LLMMetadata:
           return LLMMetadata(
               is_chat_model=True,
               is_function_calling_model=True,
               model_name="state-aware-mock-llm",
           )
    
       def _prepare_chat_with_tools(
           self,
           tools: list[Any],
           user_msg: str | ChatMessage | None = None,
           chat_history: list[ChatMessage] | None = None,
           **kwargs: Any,
       ) -> dict[str, Any]:
           messages = list(chat_history or [])
           if isinstance(user_msg, str):
               messages.append(ChatMessage(role=MessageRole.USER, content=user_msg))
           elif isinstance(user_msg, ChatMessage):
               messages.append(user_msg)
           return {"messages": messages}
    
       def get_tool_calls_from_response(
           self,
           response: ChatResponse,
           error_on_no_tool_call: bool = True,
           **kwargs: Any,
       ) -> list[ToolSelection]:
           calls = response.message.additional_kwargs.get("tool_calls", [])
           if not calls and error_on_no_tool_call:
               raise ValueError("No tool calls returned")
           return calls

4. Add the mock model's state-tool selection below get_tool_calls_from_response().

       def chat(self, messages: list[ChatMessage], **kwargs: Any) -> ChatResponse:
           if messages and messages[-1].role == MessageRole.TOOL:
               return ChatResponse(
                   message=ChatMessage(
                       role=MessageRole.ASSISTANT,
                       content=f"The saved name is {messages[-1].content}.",
                   )
               )
    
           latest_user = next(
               str(message.content)
               for message in reversed(messages)
               if message.role == MessageRole.USER
           )
           if "What name did I ask" not in latest_user:
               tool_call = ToolSelection(
                   tool_id="call_set_name",
                   tool_name="set_name",
                   tool_kwargs={"name": "Laurie"},
               )
           else:
               tool_call = ToolSelection(
                   tool_id="call_read_name",
                   tool_name="read_name",
                   tool_kwargs={},
               )
    
           return ChatResponse(
               message=ChatMessage(
                   role=MessageRole.ASSISTANT,
                   content="",
                   additional_kwargs={"tool_calls": [tool_call]},
               )
           )

   The second response uses the tool message content rather than a hard-coded name. A provider-backed application can replace StateAwareMockLLM after this local state path works.

5. Complete the model adapters below chat().

       async def achat(
           self, messages: list[ChatMessage], **kwargs: Any
       ) -> ChatResponse:
           return self.chat(messages, **kwargs)
    
       def stream_chat(self, messages: list[ChatMessage], **kwargs: Any):
           yield self.chat(messages, **kwargs)
    
       async def astream_chat(self, messages: list[ChatMessage], **kwargs: Any):
           async def response_stream():
               yield self.chat(messages, **kwargs)
    
           return response_stream()
    
       def complete(self, prompt: str, **kwargs: Any) -> CompletionResponse:
           return CompletionResponse(text="")
    
       async def acomplete(self, prompt: str, **kwargs: Any) -> CompletionResponse:
           return self.complete(prompt, **kwargs)
    
       def stream_complete(self, prompt: str, **kwargs: Any):
           yield CompletionResponse(text="")
    
       async def astream_complete(self, prompt: str, **kwargs: Any):
           yield CompletionResponse(text="")

6. Add the workflow factory below StateAwareMockLLM.

   def build_workflow() -> AgentWorkflow:
       return AgentWorkflow.from_tools_or_functions(
           [set_name, read_name],
           llm=StateAwareMockLLM(),
           system_prompt="Use the state tools to save or retrieve the user's name.",
           initial_state={"name": "unset"},
       )

7. Start main() with the first agent turn below build_workflow().

   async def main() -> None:
       workflow = build_workflow()
       ctx = Context(workflow)
    
       saved_response = await workflow.run(
           user_msg="Remember that my name is Laurie.", ctx=ctx
       )
       print(f"saved_agent={saved_response}")

   The first agent turn calls set_name() with the active Context instead of changing the state directly outside the agent loop.

8. Append the Context serialization block inside main().

       ctx_dict = ctx.to_dict(serializer=JsonSerializer())
       with open("agent-state.json", "w", encoding="utf-8") as state_file:
           json.dump(ctx_dict, state_file)

   JsonSerializer accepts only JSON-serializable values in state. API clients, model handles, open files, and indexes belong outside the saved state.

9. Append the Context restoration block after the JSON write.

       with open("agent-state.json", encoding="utf-8") as state_file:
           saved_ctx_dict = json.load(state_file)
    
       restored_ctx = Context.from_dict(
           workflow,
           saved_ctx_dict,
           serializer=JsonSerializer(),
       )

   The restored dictionary must use the workflow definition that created the original Context.

10. Append the post-restoration agent turn below Context.from_dict().

        restored_response = await workflow.run(
            user_msg="What name did I ask you to remember?", ctx=restored_ctx
        )
        print(f"restored_agent={restored_response}")
        print(f"tool_calls={TOOL_CALL_LOG}")

    The second workflow.run() receives restored_ctx. Its read_name() tool call must return the value saved by the earlier agent turn.

11. Finish the program with the asynchronous entry point.

    if __name__ == "__main__":
        asyncio.run(main())

12. Inspect the assembled llamaindex-agent-state.py file for every constructed section.

    llamaindex-agent-state.py

    import asyncio
    import json
    from typing import Any
     
    from llama_index.core.agent.workflow import AgentWorkflow
    from llama_index.core.base.llms.types import (
        ChatMessage,
        ChatResponse,
        CompletionResponse,
        LLMMetadata,
        MessageRole,
    )
    from llama_index.core.llms.function_calling import FunctionCallingLLM
    from llama_index.core.llms.llm import ToolSelection
    from llama_index.core.workflow import Context, JsonSerializer
     
    TOOL_CALL_LOG: list[str] = []
     
     
    async def set_name(ctx: Context, name: str) -> str:
        """Save the user's preferred name in agent state."""
        async with ctx.store.edit_state() as ctx_state:
            ctx_state["state"]["name"] = name
        TOOL_CALL_LOG.append(f"set_name(name={name})")
        return name
     
     
    async def read_name(ctx: Context) -> str:
        """Read the user's preferred name from agent state."""
        state = await ctx.store.get("state")
        TOOL_CALL_LOG.append("read_name()")
        return state["name"]
     
     
    class StateAwareMockLLM(FunctionCallingLLM):
        @property
        def metadata(self) -> LLMMetadata:
            return LLMMetadata(
                is_chat_model=True,
                is_function_calling_model=True,
                model_name="state-aware-mock-llm",
            )
     
        def _prepare_chat_with_tools(
            self,
            tools: list[Any],
            user_msg: str | ChatMessage | None = None,
            chat_history: list[ChatMessage] | None = None,
            **kwargs: Any,
        ) -> dict[str, Any]:
            messages = list(chat_history or [])
            if isinstance(user_msg, str):
                messages.append(ChatMessage(role=MessageRole.USER, content=user_msg))
            elif isinstance(user_msg, ChatMessage):
                messages.append(user_msg)
            return {"messages": messages}
     
        def get_tool_calls_from_response(
            self,
            response: ChatResponse,
            error_on_no_tool_call: bool = True,
            **kwargs: Any,
        ) -> list[ToolSelection]:
            calls = response.message.additional_kwargs.get("tool_calls", [])
            if not calls and error_on_no_tool_call:
                raise ValueError("No tool calls returned")
            return calls
     
        def chat(self, messages: list[ChatMessage], **kwargs: Any) -> ChatResponse:
            if messages and messages[-1].role == MessageRole.TOOL:
                return ChatResponse(
                    message=ChatMessage(
                        role=MessageRole.ASSISTANT,
                        content=f"The saved name is {messages[-1].content}.",
                    )
                )
     
            latest_user = next(
                str(message.content)
                for message in reversed(messages)
                if message.role == MessageRole.USER
            )
            if "What name did I ask" not in latest_user:
                tool_call = ToolSelection(
                    tool_id="call_set_name",
                    tool_name="set_name",
                    tool_kwargs={"name": "Laurie"},
                )
            else:
                tool_call = ToolSelection(
                    tool_id="call_read_name",
                    tool_name="read_name",
                    tool_kwargs={},
                )
     
            return ChatResponse(
                message=ChatMessage(
                    role=MessageRole.ASSISTANT,
                    content="",
                    additional_kwargs={"tool_calls": [tool_call]},
                )
            )
     
        async def achat(
            self, messages: list[ChatMessage], **kwargs: Any
        ) -> ChatResponse:
            return self.chat(messages, **kwargs)
     
        def stream_chat(self, messages: list[ChatMessage], **kwargs: Any):
            yield self.chat(messages, **kwargs)
     
        async def astream_chat(self, messages: list[ChatMessage], **kwargs: Any):
            async def response_stream():
                yield self.chat(messages, **kwargs)
     
            return response_stream()
     
        def complete(self, prompt: str, **kwargs: Any) -> CompletionResponse:
            return CompletionResponse(text="")
     
        async def acomplete(self, prompt: str, **kwargs: Any) -> CompletionResponse:
            return self.complete(prompt, **kwargs)
     
        def stream_complete(self, prompt: str, **kwargs: Any):
            yield CompletionResponse(text="")
     
        async def astream_complete(self, prompt: str, **kwargs: Any):
            yield CompletionResponse(text="")
     
     
    def build_workflow() -> AgentWorkflow:
        return AgentWorkflow.from_tools_or_functions(
            [set_name, read_name],
            llm=StateAwareMockLLM(),
            system_prompt="Use the state tools to save or retrieve the user's name.",
            initial_state={"name": "unset"},
        )
     
     
    async def main() -> None:
        workflow = build_workflow()
        ctx = Context(workflow)
     
        saved_response = await workflow.run(
            user_msg="Remember that my name is Laurie.", ctx=ctx
        )
        print(f"saved_agent={saved_response}")
     
        ctx_dict = ctx.to_dict(serializer=JsonSerializer())
        with open("agent-state.json", "w", encoding="utf-8") as state_file:
            json.dump(ctx_dict, state_file)
     
        with open("agent-state.json", encoding="utf-8") as state_file:
            saved_ctx_dict = json.load(state_file)
     
        restored_ctx = Context.from_dict(
            workflow,
            saved_ctx_dict,
            serializer=JsonSerializer(),
        )
        restored_response = await workflow.run(
            user_msg="What name did I ask you to remember?", ctx=restored_ctx
        )
        print(f"restored_agent={restored_response}")
        print(f"tool_calls={TOOL_CALL_LOG}")
     
     
    if __name__ == "__main__":
        asyncio.run(main())

13. Run the completed state round trip.

    $ python3 llamaindex-agent-state.py
    saved_agent=The saved name is Laurie.
    restored_agent=The saved name is Laurie.
    tool_calls=['set_name(name=Laurie)', 'read_name()']

    The restored response must come after read_name() in the tool log. A missing call or an unset response means the later agent turn did not consume the saved Context value.