-
Notifications
You must be signed in to change notification settings - Fork 1.5k
Open
Labels
documentation[Component] This issue is related to documentation, it will be transferred to adk-docs[Component] This issue is related to documentation, it will be transferred to adk-docsservices[Component] This issue is related to runtime services, e.g. sessions, memory, artifacts, etc[Component] This issue is related to runtime services, e.g. sessions, memory, artifacts, etc
Description
The session state update behaviour is confusing.
In this example the session state value of topic should be updated from 'a brave kitten exploring a haunted house' to 'a boy in australia' when the final state is logged.
To Reproduce
Steps to reproduce the behaviour:
- Run the script (a minimal version of the custom agent example in the docs
# Full runnable code for the StoryFlowAgent example
import logging
import json
from typing import AsyncGenerator
from typing_extensions import override
from google.adk.agents import LlmAgent, BaseAgent
from google.adk.agents.invocation_context import InvocationContext
from google.genai import types
from google.adk.sessions import InMemorySessionService
from google.adk.runners import Runner
from google.adk.events import Event
from pydantic import BaseModel, Field
from dotenv import load_dotenv
# Load environment variables from .env file
load_dotenv()
# --- Constants ---
APP_NAME = "story_app"
USER_ID = "12345"
SESSION_ID = "123344"
GEMINI_2_FLASH = "gemini-2.0-flash-lite"
# --- Configure Logging ---
logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)
class DummyAgent1Output(BaseModel):
test_value: str = Field(
..., description="The value of the key 'test_key' in the state session"
)
class StoryOutput(BaseModel):
story: str = Field(..., description="The story generated by the agent")
# --- Custom Orchestrator Agent ---
class StoryFlowAgent(BaseAgent):
"""
Custom agent for a story generation and refinement workflow.
This agent orchestrates a sequence of LLM agents to generate a story,
critique it, revise it, check grammar and tone, and potentially
regenerate the story if the tone is negative.
"""
# --- Field Declarations for Pydantic ---
# Declare the agents passed during initialization as class attributes with type hints
story_generator: LlmAgent
# model_config allows setting Pydantic configurations if needed, e.g., arbitrary_types_allowed
model_config = {"arbitrary_types_allowed": True}
def __init__(
self,
name: str,
story_generator: LlmAgent,
):
"""
Initializes the StoryFlowAgent.
Args:
name: The name of the agent.
story_generator: An LlmAgent to generate the initial story.
critic: An LlmAgent to critique the story.
reviser: An LlmAgent to revise the story based on criticism.
grammar_check: An LlmAgent to check the grammar.
tone_check: An LlmAgent to analyze the tone.
"""
# Define the sub_agents list for the framework
sub_agents_list = [
story_generator,
]
# Pydantic will validate and assign them based on the class annotations.
super().__init__(
name=name,
story_generator=story_generator,
sub_agents=sub_agents_list, # Pass the sub_agents list directly
)
@override
async def _run_async_impl(
self, ctx: InvocationContext
) -> AsyncGenerator[Event, None]:
"""
Implements the custom orchestration logic for the story workflow.
Uses the instance attributes assigned by Pydantic (e.g., self.story_generator).
"""
logger.info(f"[{self.name}] Starting story generation workflow.")
# 1. Initial Story Generation
logger.info(f"[{self.name}] Running StoryGenerator...")
async for event in self.story_generator.run_async(ctx):
# logger.info(
# f"[{self.name}] Event from StoryGenerator: {event.model_dump_json(indent=2, exclude_none=True)}"
# )
yield event
# Check if story was generated before proceeding
if (
"current_story" not in ctx.session.state
or not ctx.session.state["current_story"]
):
logger.error(
f"[{self.name}] Failed to generate initial story. Aborting workflow."
)
return # Stop processing if initial story failed
# logger.info(
# f"[{self.name}] Story state after generator: {ctx.session.state.get('current_story')}"
# )
logger.info(f"[{self.name}] Workflow finished.")
# --- Define the individual LLM agents ---
story_generator = LlmAgent(
name="StoryGenerator",
model=GEMINI_2_FLASH,
instruction="""You are a story writer. Write a short story (around 100 words) about a cat,
based on the topic provided in session state with key 'topic'""",
input_schema=None,
# output_schema=StoryOutput,
output_key="current_story", # Key for storing output in session state
)
state_var_picker = LlmAgent(
name="StoryGenerator",
model=GEMINI_2_FLASH,
instruction="""You are a story writer. Write a short story (around 100 words) about a cat,
based on the topic provided in session state with key 'topic'""",
input_schema=None,
# output_schema=StoryOutput,
output_key="current_story", # Key for storing output in session state
)
# --- Create the custom agent instance ---
story_flow_agent = StoryFlowAgent(
name="StoryFlowAgent",
story_generator=story_generator,
)
# --- Setup Runner and Session ---
session_service = InMemorySessionService()
initial_state = {"topic": "a brave kitten exploring a haunted house"}
session = session_service.create_session(
app_name=APP_NAME,
user_id=USER_ID,
session_id=SESSION_ID,
state=initial_state, # Pass initial state here
)
runner = Runner(
agent=story_flow_agent, # Pass the custom orchestrator agent
app_name=APP_NAME,
session_service=session_service,
)
# --- Function to Interact with the Agent ---
def call_dummy_agent(user_input_topic: str):
"""
Sends a new topic to the agent (overwriting the initial one if needed)
and runs the workflow.
"""
logger.info(f"Initial session state: {session.state}")
current_session = session_service.get_session(
app_name=APP_NAME, user_id=USER_ID, session_id=SESSION_ID
)
if not current_session:
logger.error("Session not found!")
return
current_session.state.update({"topic": user_input_topic})
updated_topic = current_session.state.get("topic")
# Shows updated value
logger.info(f"Updated session state topic to: {updated_topic}")
content = types.Content(
role="user",
parts=[types.Part(text=f"Generate a story about: {user_input_topic}")],
)
events = runner.run(user_id=USER_ID, session_id=SESSION_ID, new_message=content)
final_response = "No final response captured."
for event in events:
if event.is_final_response() and event.content and event.content.parts:
# logger.info(
# f"Potential final response from [{event.author}]: {event.content.parts[0].text}"
# )
final_response = event.content.parts[0].text
print("\n--- Agent Interaction Result ---")
print("Agent Final Response: ", final_response)
final_session = session_service.get_session(
app_name=APP_NAME, user_id=USER_ID, session_id=SESSION_ID
)
print("Final Session State:")
print(json.dumps(final_session.state, indent=2))
print("-------------------------------\n")
# Shows initial value
logger.info("updated topic in session state: %s", final_session.state.get("topic"))
call_dummy_agent("a boy in australia")
The output log
INFO:__main__:Initial session state: {'topic': 'a brave kitten exploring a haunted house'}
INFO:__main__:Updated session state topic to: a boy in australia
INFO:__main__:[StoryFlowAgent] Starting story generation workflow.
INFO:__main__:[StoryFlowAgent] Running StoryGenerator...
INFO:google.adk.models.google_llm:Sending out request, model: gemini-2.0-flash-lite, backend: ml_dev, stream: False
INFO:google.adk.models.google_llm:
LLM Request:
-----------------------------------------------------------
System Instruction:
You are a story writer. Write a short story (around 100 words) about a cat,
based on the topic provided in session state with key 'topic'
You are an agent. Your internal name is "StoryGenerator".
-----------------------------------------------------------
Contents:
{"parts":[{"text":"Generate a story about: a boy in australia"}],"role":"user"}
-----------------------------------------------------------
Functions:
-----------------------------------------------------------
INFO:google_genai.models:AFC is enabled with max remote calls: 10.
INFO:httpx:HTTP Request: POST https://generativelanguage.googleapis.com/v1beta/models/gemini-2.0-flash-lite:generateContent "HTTP/1.1 200 OK"
INFO:google.adk.models.google_llm:
LLM Response:
-----------------------------------------------------------
Text:
StoryGenerator:
The Australian sun beat down on ten-year-old Finn as he raced across the red dirt. His bare feet kicked up a cloud of ochre dust, chasing after a rogue cricket ball. He lived on a sprawling cattle station, miles from anything, and his best friends were his kelpie, Bluey, and the endless outback. Today, however, he felt a pang of loneliness. He missed the city, the shops, the other kids. Suddenly, a flash of emerald green caught his eye. Nestled under a gum tree, a sleek, black cat with jade eyes stared back at him. A new friend, perhaps? Finn grinned. The outback, he realized, always had a surprise or two.
-----------------------------------------------------------
Function calls:
-----------------------------------------------------------
Raw response:
{"candidates":[{"content":{"parts":[{"text":"StoryGenerator:\nThe Australian sun beat down on ten-year-old Finn as he raced across the red dirt. His bare feet kicked up a cloud of ochre dust, chasing after a rogue cricket ball. He lived on a sprawling cattle station, miles from anything, and his best friends were his kelpie, Bluey, and the endless outback. Today, however, he felt a pang of loneliness. He missed the city, the shops, the other kids. Suddenly, a flash of emerald green caught his eye. Nestled under a gum tree, a sleek, black cat with jade eyes stared back at him. A new friend, perhaps? Finn grinned. The outback, he realized, always had a surprise or two.\n"}],"role":"model"},"finish_reason":"STOP","avg_logprobs":-0.6154808869233003}],"model_version":"gemini-2.0-flash-lite","usage_metadata":{"candidates_token_count":148,"candidates_tokens_details":[{"modality":"TEXT","token_count":148}],"prompt_token_count":60,"prompt_tokens_details":[{"modality":"TEXT","token_count":60}],"total_token_count":208},"automatic_function_calling_history":[]}
-----------------------------------------------------------
INFO:__main__:[StoryFlowAgent] Workflow finished.
--- Agent Interaction Result ---
Agent Final Response: StoryGenerator:
The Australian sun beat down on ten-year-old Finn as he raced across the red dirt. His bare feet kicked up a cloud of ochre dust, chasing after a rogue cricket ball. He lived on a sprawling cattle station, miles from anything, and his best friends were his kelpie, Bluey, and the endless outback. Today, however, he felt a pang of loneliness. He missed the city, the shops, the other kids. Suddenly, a flash of emerald green caught his eye. Nestled under a gum tree, a sleek, black cat with jade eyes stared back at him. A new friend, perhaps? Finn grinned. The outback, he realized, always had a surprise or two.
Final Session State:
{
"topic": "a brave kitten exploring a haunted house",
"current_story": "StoryGenerator:\nThe Australian sun beat down on ten-year-old Finn as he raced across the red dirt. His bare feet kicked up a cloud of ochre dust, chasing after a rogue cricket ball. He lived on a sprawling cattle station, miles from anything, and his best friends were his kelpie, Bluey, and the endless outback. Today, however, he felt a pang of loneliness. He missed the city, the shops, the other kids. Suddenly, a flash of emerald green caught his eye. Nestled under a gum tree, a sleek, black cat with jade eyes stared back at him. A new friend, perhaps? Finn grinned. The outback, he realized, always had a surprise or two.\n"
}
-------------------------------
INFO:__main__:updated topic in session state: a brave kitten exploring a haunted house
Expected behavior
The final state log should reflect the updated value of the story topic but is not doing so. But the updated topic is being used by the Agent to produce the story
Desktop
- OS: debian 12
- Python version(python -V): 3.13.2
- ADK version(pip show google-adk): 0.2.0
Additional context
Use the following env vars to run
GOOGLE_GENAI_USE_VERTEXAI=FALSE
GOOGLE_API_KEY=<gemini_api_key>
aspen-yryr
Metadata
Metadata
Assignees
Labels
documentation[Component] This issue is related to documentation, it will be transferred to adk-docs[Component] This issue is related to documentation, it will be transferred to adk-docsservices[Component] This issue is related to runtime services, e.g. sessions, memory, artifacts, etc[Component] This issue is related to runtime services, e.g. sessions, memory, artifacts, etc