Writing agent state

Define your agent's state schema. AWS Strands maintains state throughout execution.

<Tabs groupId="language_strands_agent" items={['Python', 'TypeScript']} persist>
    <Tab value="Python">
        ```python title="agent/main.py"
        from ag_ui_strands import StrandsAgent, StrandsAgentConfig, create_strands_app
        from strands import Agent
        from strands.models.openai import OpenAIModel

        model = OpenAIModel(
            client_args={"api_key": os.getenv("OPENAI_API_KEY", "")},
            model_id="gpt-4o",
        )

        # Create the Strands agent
        strands_agent = Agent(
            model=model,
            system_prompt="Always communicate in the preferred language of the user as defined in your state. Do not communicate in any other language.",
        )

        # Inject state into the prompt
        def language_prompt(input_data, user_message: str) -> str:
            state_dict = getattr(input_data, "state", None)
            if isinstance(state_dict, dict) and "language" in state_dict:
                return f"Current language: {state_dict['language']}\n\nUser request: {user_message}"
            return user_message

        config = StrandsAgentConfig(state_context_builder=language_prompt)

        # Wrap with AG-UI integration
        agui_agent = StrandsAgent(
            agent=strands_agent,
            name="languageAgent",
            description="Always communicate in the preferred language of the user",
            config=config,
        )

        app = create_strands_app(agui_agent)
        ```
    </Tab>
    <Tab value="TypeScript">
        ```typescript title="agent/main.ts"
        import { Agent } from "@strands-agents/sdk";
        import { OpenAIModel } from "@strands-agents/sdk/models/openai";
        import { StrandsAgent, type StrandsAgentConfig } from "@ag-ui/aws-strands";
        import { createStrandsApp } from "@ag-ui/aws-strands/server";

        const model = new OpenAIModel({
          apiKey: process.env.OPENAI_API_KEY ?? "",
          modelId: "gpt-4o",
        });

        // Create the Strands agent
        const strandsAgent = new Agent({
          model,
          systemPrompt:
            "Always communicate in the preferred language of the user as defined in your state. Do not communicate in any other language.",
        });

        await strandsAgent.initialize();

        // Inject state into the prompt
        const config: StrandsAgentConfig = {
          stateContextBuilder: (inputData, userMessage) => {
            const state = (inputData.state ?? {}) as { language?: string };
            if (state.language) {
              return `Current language: ${state.language}\n\nUser request: ${userMessage}`;
            }
            return userMessage;
          },
        };

        // Wrap with AG-UI integration
        const aguiAgent = new StrandsAgent({
          agent: strandsAgent,
          name: "languageAgent",
          description: "Always communicate in the preferred language of the user",
          config,
        });

        const app = await createStrandsApp(aguiAgent);
        app.listen(8000);
        ```
    </Tab>
</Tabs>

With your agent connected and running, call the `useAgent` hook, pass the agent's name, and
use `agent.setState` to update the agent state.

```tsx title="ui/app/page.tsx"

// Define the agent state type to match your Strands agent
type AgentState = {
  language: "english" | "spanish";
};

function YourMainContent() {
  // [!code highlight:5]
  const { agent } = useAgent({
    agentId: "languageAgent",
    // optionally provide a type-safe initial state
    initialState: { language: "spanish" }
  });

  const toggleLanguage = () => {
    agent.setState({ language: agent.state?.language === "english" ? "spanish" : "english" }); // [!code highlight]
  };

  return (
    // style excluded for brevity
    <div>
      <h1>Your main content</h1>
      <p>Language: {agent.state?.language}</p>
      <button onClick={toggleLanguage}>Toggle Language</button>
    </div>
  );
}
```
<Callout type="info">
  The `agent.setState` function in `useAgent` will update the agent state and trigger a rerender when the state changes.
</Callout>

You can now use the `setState` function to update the agent state and `state` to read it. Try toggling the language button
and talking to your agent. You'll see the language change to match the agent's state.

<Callout type="info" title="Important">
  Shared-state in AWS Strands is prompt-driven. This means that your Agent's awareness of the shared state will be augmented
  by the instructions you provide to your agent.
</Callout>