커스텀 에이전트¶
ADK에서 지원Python v0.1.0Go v0.1.0Java v0.2.0
커스텀 에이전트는 BaseAgent를 직접 상속하고 자신만의 제어 흐름을 구현하여 임의의 오케스트레이션 로직을 정의할 수 있게 함으로써 ADK에서 최고의 유연성을 제공합니다. 이는 SequentialAgent, LoopAgent, ParallelAgent의 사전 정의된 패턴을 넘어서, 매우 특수하고 복잡한 에이전트 워크플로우를 구축할 수 있게 해줍니다.
고급 개념
_run_async_impl(또는 다른 언어의 해당 메서드)을 직접 구현하여 커스텀 에이전트를 구축하는 것은 강력한 제어 기능을 제공하지만, 사전 정의된 LlmAgent나 표준 WorkflowAgent 타입을 사용하는 것보다 더 복잡합니다. 커스텀 오케스트레이션 로직을 다루기 전에 이러한 기본 에이전트 타입들을 먼저 이해하는 것을 권장합니다.
소개: 사전 정의된 워크플로우를 넘어서¶
커스텀 에이전트란?¶
커스텀 에이전트는 본질적으로 google.adk.agents.BaseAgent를 상속하고 _run_async_impl 비동기 메서드 내에 핵심 실행 로직을 구현하여 생성하는 모든 클래스입니다. 이 메서드가 다른 에이전트(하위 에이전트)를 호출하고, 상태를 관리하며, 이벤트를 처리하는 방식을 완벽하게 제어할 수 있습니다.
Note
에이전트의 핵심 비동기 로직을 구현하는 특정 메서드 이름은 SDK 언어에 따라 약간 다를 수 있습니다(예: Java의 runAsyncImpl, Python의 _run_async_impl). 자세한 내용은 언어별 API 문서를 참조하세요.
왜 사용해야 하는가?¶
표준 워크플로우 에이전트(SequentialAgent, LoopAgent, ParallelAgent)가 일반적인 오케스트레이션 패턴을 다루지만, 요구 사항에 다음이 포함될 경우 커스텀 에이전트가 필요합니다.
- 조건부 로직: 런타임 조건이나 이전 단계의 결과에 따라 다른 하위 에이전트를 실행하거나 다른 경로를 택해야 할 경우.
- 복잡한 상태 관리: 단순한 순차적 전달을 넘어 워크플로우 전반에 걸쳐 상태를 유지하고 업데이트하기 위한 복잡한 로직을 구현해야 할 경우.
- 외부 통합: 오케스트레이션 흐름 제어 내에서 외부 API, 데이터베이스 또는 커스텀 라이브러리 호출을 직접 통합해야 할 경우.
- 동적 에이전트 선택: 상황이나 입력을 동적으로 평가하여 다음에 실행할 하위 에이전트를 선택해야 할 경우.
- 독특한 워크플로우 패턴: 표준적인 순차, 병렬 또는 루프 구조에 맞지 않는 오케스트레이션 로직을 구현해야 할 경우.
커스텀 로직 구현:¶
모든 커스텀 에이전트의 핵심은 고유한 비동기 동작을 정의하는 메서드입니다. 이 메서드를 통해 하위 에이전트를 오케스트레이션하고 실행 흐름을 관리할 수 있습니다.
모든 커스텀 에이전트의 심장은 _run_async_impl 메서드입니다. 이곳에서 고유한 동작을 정의합니다.
- 시그니처:
async def _run_async_impl(self, ctx: InvocationContext) -> AsyncGenerator[Event, None]: - 비동기 제너레이터:
async def함수여야 하며AsyncGenerator를 반환해야 합니다. 이를 통해 하위 에이전트나 자체 로직에서 생성된 이벤트를yield하여 실행기(runner)에 다시 전달할 수 있습니다. ctx(InvocationContext): 핵심적인 런타임 정보, 특히ctx.session.state에 대한 접근을 제공합니다. 이는 커스텀 에이전트가 오케스트레이션하는 단계 간에 데이터를 공유하는 주요 방법입니다.
Go에서는 agent.Agent 인터페이스를 만족하는 구조체의 일부로 Run 메서드를 구현합니다. 실제 로직은 일반적으로 커스텀 에이전트 구조체의 메서드가 됩니다.
- 시그니처:
Run(ctx agent.InvocationContext) iter.Seq2[*session.Event, error] - 이터레이터:
Run메서드는 이벤트와 오류를 반환하는 이터레이터(iter.Seq2)를 반환합니다. 이는 에이전트 실행에서 스트리밍 결과를 처리하는 표준적인 방법입니다. ctx(InvocationContext):agent.InvocationContext는 상태를 포함한 세션 및 기타 핵심적인 런타임 정보에 대한 접근을 제공합니다.- 세션 상태:
ctx.Session().State()를 통해 세션 상태에 접근할 수 있습니다.
모든 커스텀 에이전트의 심장은 BaseAgent로부터 오버라이드하는 runAsyncImpl 메서드입니다.
- 시그니처:
protected Flowable<Event> runAsyncImpl(InvocationContext ctx) - 리액티브 스트림 (
Flowable):io.reactivex.rxjava3.core.Flowable<Event>를 반환해야 합니다. 이Flowable은 커스텀 에이전트의 로직에 의해 생성될 이벤트의 스트림을 나타내며, 종종 여러 하위 에이전트의Flowable을 결합하거나 변환하여 만들어집니다. ctx(InvocationContext): 핵심적인 런타임 정보, 특히java.util.concurrent.ConcurrentMap<String, Object>인ctx.session().state()에 대한 접근을 제공합니다. 이는 커스텀 에이전트가 오케스트레이션하는 단계 간에 데이터를 공유하는 주요 방법입니다.
핵심 비동기 메서드 내 주요 기능:
-
하위 에이전트 호출:
run_async메서드를 사용하여 (일반적으로self.my_llm_agent와 같은 인스턴스 속성으로 저장된) 하위 에이전트를 호출하고 그 이벤트를 yield합니다. -
상태 관리: 세션 상태 딕셔너리(
ctx.session.state)를 읽고 써서 하위 에이전트 호출 간에 데이터를 전달하거나 결정을 내립니다. -
제어 흐름 구현: 표준 Python 구문(
if/elif/else,for/while루프,try/except)을 사용하여 하위 에이전트를 포함하는 정교하고, 조건부이거나, 반복적인 워크플로우를 만듭니다.
-
하위 에이전트 호출:
Run메서드를 호출하여 하위 에이전트를 호출합니다. -
상태 관리: 세션 상태를 읽고 써서 하위 에이전트 호출 간에 데이터를 전달하거나 결정을 내립니다.
// `ctx` (`agent.InvocationContext`)는 에이전트의 `Run` 함수에 직접 전달됨. // 이전 에이전트가 설정한 데이터 읽기 previousResult, err := ctx.Session().State().Get("some_key") if err != nil { // 키가 아직 존재하지 않을 수 있는 경우 처리 } // 상태에 기반한 결정 if val, ok := previousResult.(string); ok && val == "some_value" { // ... 특정 하위 에이전트 호출 ... } else { // ... 다른 하위 에이전트 호출 ... } // 나중 단계를 위한 결과 저장 if err := ctx.Session().State().Set("my_custom_result", "calculated_value"); err != nil { // 오류 처리 } -
제어 흐름 구현: 표준 Go 구문(
if/else,for/switch루프, 고루틴, 채널)을 사용하여 하위 에이전트를 포함하는 정교하고, 조건부이거나, 반복적인 워크플로우를 만듭니다.
-
하위 에이전트 호출: 비동기 실행 메서드를 사용하여 (일반적으로 인스턴스 속성이나 객체로 저장된) 하위 에이전트를 호출하고 그 이벤트 스트림을 반환합니다.
일반적으로 RxJava 연산자인
concatWith,flatMapPublisher,concatArray등을 사용하여 하위 에이전트의Flowable들을 연결합니다.// 예: 하나의 하위 에이전트 실행 // return someSubAgent.runAsync(ctx); // 예: 하위 에이전트를 순차적으로 실행 Flowable<Event> firstAgentEvents = someSubAgent1.runAsync(ctx) .doOnNext(event -> System.out.println("에이전트 1의 이벤트: " + event.id())); Flowable<Event> secondAgentEvents = Flowable.defer(() -> someSubAgent2.runAsync(ctx) .doOnNext(event -> System.out.println("에이전트 2의 이벤트: " + event.id())) ); return firstAgentEvents.concatWith(secondAgentEvents);Flowable.defer()는 후속 단계의 실행이 이전 단계의 완료나 상태에 따라 달라지는 경우에 자주 사용됩니다. -
상태 관리: 세션 상태를 읽고 써서 하위 에이전트 호출 간에 데이터를 전달하거나 결정을 내립니다. 세션 상태는
ctx.session().state()를 통해 얻는java.util.concurrent.ConcurrentMap<String, Object>입니다.// 이전 에이전트가 설정한 데이터 읽기 Object previousResult = ctx.session().state().get("some_key"); // 상태에 기반한 결정 if ("some_value".equals(previousResult)) { // ... 특정 하위 에이전트의 Flowable을 포함하는 로직 ... } else { // ... 다른 하위 에이전트의 Flowable을 포함하는 로직 ... } // 나중 단계를 위한 결과 저장 (종종 하위 에이전트의 output_key를 통해 수행됨) // ctx.session().state().put("my_custom_result", "calculated_value"); -
제어 흐름 구현: 표준 언어 구문(
if/else, 루프,try/catch)과 리액티브 연산자(RxJava)를 결합하여 정교한 워크플로우를 만듭니다.- 조건부: 조건에 따라 구독할
Flowable을 선택하기 위해Flowable.defer()를 사용하거나, 스트림 내에서 이벤트를 필터링하는 경우filter()를 사용합니다. - 반복:
repeat(),retry()같은 연산자를 사용하거나, 조건에 따라Flowable체인의 일부를 재귀적으로 호출하도록 구조화합니다(종종flatMapPublisher나concatMap으로 관리).
- 조건부: 조건에 따라 구독할
하위 에이전트와 상태 관리¶
일반적으로 커스텀 에이전트는 다른 에이전트(LlmAgent, LoopAgent 등)를 오케스트레이션합니다.
- 초기화: 보통 이러한 하위 에이전트의 인스턴스를 커스텀 에이전트의 생성자에 전달하고 인스턴스 필드/속성으로 저장합니다(예:
this.story_generator = story_generator_instance또는self.story_generator = story_generator_instance). 이렇게 하면 커스텀 에이전트의 핵심 비동기 실행 로직(예:_run_async_impl메서드) 내에서 접근할 수 있습니다. - 하위 에이전트 목록:
super()생성자를 사용하여BaseAgent를 초기화할 때,sub agents목록을 전달해야 합니다. 이 목록은 ADK 프레임워크에게 이 커스텀 에이전트의 직계 계층에 속한 에이전트가 무엇인지 알려줍니다. 이는 핵심 실행 로직(_run_async_impl)이self.xxx_agent를 통해 에이전트를 직접 호출하더라도 생명주기 관리, 인트로스펙션, 그리고 잠재적인 미래 라우팅 기능과 같은 프레임워크 기능에 중요합니다. 커스텀 로직이 최상위 수준에서 직접 호출하는 에이전트들을 포함시키세요. - 상태: 언급했듯이,
ctx.session.state는 하위 에이전트(특히output key를 사용하는LlmAgent)가 결과를 오케스트레이터에게 다시 전달하고 오케스트레이터가 필요한 입력을 전달하는 표준적인 방법입니다.
디자인 패턴 예시: StoryFlowAgent¶
조건부 로직이 포함된 다단계 콘텐츠 생성 워크플로우 예시 패턴으로 커스텀 에이전트의 강력함을 설명해 보겠습니다.
목표: 이야기를 생성하고, 비평과 수정을 통해 반복적으로 개선하며, 최종 검사를 수행하고, 결정적으로 최종 톤 검사에 실패하면 이야기를 재생성하는 시스템을 만드는 것입니다.
왜 커스텀인가? 여기서 커스텀 에이전트가 필요한 핵심 요구 사항은 톤 검사에 기반한 조건부 재생성입니다. 표준 워크플로우 에이전트는 하위 에이전트 작업의 결과에 따른 조건부 분기 기능이 내장되어 있지 않습니다. 오케스트레이터 내에 커스텀 로직(if tone == "negative": ...)이 필요합니다.
1부: 단순화된 커스텀 에이전트 초기화¶
BaseAgent를 상속하는 StoryFlowAgent를 정의합니다. __init__에서 (전달받은) 필요한 하위 에이전트를 인스턴스 속성으로 저장하고, 이 커스텀 에이전트가 직접 오케스트레이션할 최상위 에이전트를 BaseAgent 프레임워크에 알립니다.
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
critic: LlmAgent
reviser: LlmAgent
grammar_check: LlmAgent
tone_check: LlmAgent
loop_agent: LoopAgent
sequential_agent: SequentialAgent
# 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,
critic: LlmAgent,
reviser: LlmAgent,
grammar_check: LlmAgent,
tone_check: 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.
"""
# Create internal agents *before* calling super().__init__
loop_agent = LoopAgent(
name="CriticReviserLoop", sub_agents=[critic, reviser], max_iterations=2
)
sequential_agent = SequentialAgent(
name="PostProcessing", sub_agents=[grammar_check, tone_check]
)
# Define the sub_agents list for the framework
sub_agents_list = [
story_generator,
loop_agent,
sequential_agent,
]
# Pydantic will validate and assign them based on the class annotations.
super().__init__(
name=name,
story_generator=story_generator,
critic=critic,
reviser=reviser,
grammar_check=grammar_check,
tone_check=tone_check,
loop_agent=loop_agent,
sequential_agent=sequential_agent,
sub_agents=sub_agents_list, # Pass the sub_agents list directly
)
StoryFlowAgent 구조체와 생성자를 정의합니다. 생성자에서 필요한 하위 에이전트를 저장하고, 이 커스텀 에이전트가 직접 오케스트레이션할 최상위 에이전트를 BaseAgent 프레임워크에 알립니다.
// StoryFlowAgent is a custom agent that orchestrates a story generation workflow.
// It encapsulates the logic of running sub-agents in a specific sequence.
type StoryFlowAgent struct {
storyGenerator agent.Agent
revisionLoopAgent agent.Agent
postProcessorAgent agent.Agent
}
// NewStoryFlowAgent creates and configures the entire custom agent workflow.
// It takes individual LLM agents as input and internally creates the necessary
// workflow agents (loop, sequential), returning the final orchestrator agent.
func NewStoryFlowAgent(
storyGenerator,
critic,
reviser,
grammarCheck,
toneCheck agent.Agent,
) (agent.Agent, error) {
loopAgent, err := loopagent.New(loopagent.Config{
MaxIterations: 2,
AgentConfig: agent.Config{
Name: "CriticReviserLoop",
SubAgents: []agent.Agent{critic, reviser},
},
})
if err != nil {
return nil, fmt.Errorf("failed to create loop agent: %w", err)
}
sequentialAgent, err := sequentialagent.New(sequentialagent.Config{
AgentConfig: agent.Config{
Name: "PostProcessing",
SubAgents: []agent.Agent{grammarCheck, toneCheck},
},
})
if err != nil {
return nil, fmt.Errorf("failed to create sequential agent: %w", err)
}
// The StoryFlowAgent struct holds the agents needed for the Run method.
orchestrator := &StoryFlowAgent{
storyGenerator: storyGenerator,
revisionLoopAgent: loopAgent,
postProcessorAgent: sequentialAgent,
}
// agent.New creates the final agent, wiring up the Run method.
return agent.New(agent.Config{
Name: "StoryFlowAgent",
Description: "Orchestrates story generation, critique, revision, and checks.",
SubAgents: []agent.Agent{storyGenerator, loopAgent, sequentialAgent},
Run: orchestrator.Run,
})
}
BaseAgent를 확장하여 StoryFlowAgentExample을 정의합니다. 생성자에서 (매개변수로 전달된) 필요한 하위 에이전트 인스턴스를 인스턴스 필드로 저장합니다. 이 커스텀 에이전트가 직접 오케스트레이션할 이 최상위 하위 에이전트들은 BaseAgent의 super 생성자에도 목록으로 전달됩니다.
private final LlmAgent storyGenerator;
private final LoopAgent loopAgent;
private final SequentialAgent sequentialAgent;
public StoryFlowAgentExample(
String name, LlmAgent storyGenerator, LoopAgent loopAgent, SequentialAgent sequentialAgent) {
super(
name,
"Orchestrates story generation, critique, revision, and checks.",
List.of(storyGenerator, loopAgent, sequentialAgent),
null,
null);
this.storyGenerator = storyGenerator;
this.loopAgent = loopAgent;
this.sequentialAgent = sequentialAgent;
}
2부: 커스텀 실행 로직 정의¶
이 메서드는 표준 Python async/await 및 제어 흐름을 사용하여 하위 에이전트를 오케스트레이션합니다.
@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')}")
# 2. Critic-Reviser Loop
logger.info(f"[{self.name}] Running CriticReviserLoop...")
# Use the loop_agent instance attribute assigned during init
async for event in self.loop_agent.run_async(ctx):
logger.info(f"[{self.name}] Event from CriticReviserLoop: {event.model_dump_json(indent=2, exclude_none=True)}")
yield event
logger.info(f"[{self.name}] Story state after loop: {ctx.session.state.get('current_story')}")
# 3. Sequential Post-Processing (Grammar and Tone Check)
logger.info(f"[{self.name}] Running PostProcessing...")
# Use the sequential_agent instance attribute assigned during init
async for event in self.sequential_agent.run_async(ctx):
logger.info(f"[{self.name}] Event from PostProcessing: {event.model_dump_json(indent=2, exclude_none=True)}")
yield event
# 4. Tone-Based Conditional Logic
tone_check_result = ctx.session.state.get("tone_check_result")
logger.info(f"[{self.name}] Tone check result: {tone_check_result}")
if tone_check_result == "negative":
logger.info(f"[{self.name}] Tone is negative. Regenerating story...")
async for event in self.story_generator.run_async(ctx):
logger.info(f"[{self.name}] Event from StoryGenerator (Regen): {event.model_dump_json(indent=2, exclude_none=True)}")
yield event
else:
logger.info(f"[{self.name}] Tone is not negative. Keeping current story.")
pass
logger.info(f"[{self.name}] Workflow finished.")
- 초기
story_generator가 실행됩니다. 그 출력은ctx.session.state["current_story"]에 있을 것으로 예상됩니다. loop_agent가 실행되며, 내부적으로critic과reviser를max_iterations횟수만큼 순차적으로 호출합니다. 이들은 상태에서current_story와criticism을 읽고 씁니다.sequential_agent가 실행되며,grammar_check와tone_check를 차례로 호출하여current_story를 읽고grammar_suggestions와tone_check_result를 상태에 씁니다.- 커스텀 부분:
if문은 상태에서tone_check_result를 확인합니다. 만약 "negative"라면,story_generator가 다시 호출되어 상태의current_story를 덮어씁니다. 그렇지 않으면 흐름이 종료됩니다.
Run 메서드는 각 하위 에이전트의 Run 메서드를 루프에서 호출하고 그 이벤트를 반환하여 하위 에이전트를 오케스트레이션합니다.
// Run defines the custom execution logic for the StoryFlowAgent.
func (s *StoryFlowAgent) Run(ctx agent.InvocationContext) iter.Seq2[*session.Event, error] {
return func(yield func(*session.Event, error) bool) {
// Stage 1: Initial Story Generation
for event, err := range s.storyGenerator.Run(ctx) {
if err != nil {
yield(nil, fmt.Errorf("story generator failed: %w", err))
return
}
if !yield(event, nil) {
return
}
}
// Check if story was generated before proceeding
currentStory, err := ctx.Session().State().Get("current_story")
if err != nil || currentStory == "" {
log.Println("Failed to generate initial story. Aborting workflow.")
return
}
// Stage 2: Critic-Reviser Loop
for event, err := range s.revisionLoopAgent.Run(ctx) {
if err != nil {
yield(nil, fmt.Errorf("loop agent failed: %w", err))
return
}
if !yield(event, nil) {
return
}
}
// Stage 3: Post-Processing
for event, err := range s.postProcessorAgent.Run(ctx) {
if err != nil {
yield(nil, fmt.Errorf("sequential agent failed: %w", err))
return
}
if !yield(event, nil) {
return
}
}
// Stage 4: Conditional Regeneration
toneResult, err := ctx.Session().State().Get("tone_check_result")
if err != nil {
log.Printf("Could not read tone_check_result from state: %v. Assuming tone is not negative.", err)
return
}
if tone, ok := toneResult.(string); ok && tone == "negative" {
log.Println("Tone is negative. Regenerating story...")
for event, err := range s.storyGenerator.Run(ctx) {
if err != nil {
yield(nil, fmt.Errorf("story regeneration failed: %w", err))
return
}
if !yield(event, nil) {
return
}
}
} else {
log.Println("Tone is not negative. Keeping current story.")
}
}
}
- 초기
storyGenerator가 실행됩니다. 그 출력은 세션 상태의"current_story"키 아래에 있을 것으로 예상됩니다. revisionLoopAgent가 실행되며, 내부적으로critic과reviser를max_iterations횟수만큼 순차적으로 호출합니다. 이들은 상태에서current_story와criticism을 읽고 씁니다.postProcessorAgent가 실행되며,grammar_check와tone_check를 차례로 호출하여current_story를 읽고grammar_suggestions와tone_check_result를 상태에 씁니다.- 커스텀 부분: 코드는 상태에서
tone_check_result를 확인합니다. 만약 "negative"라면,story_generator가 다시 호출되어 상태의current_story를 덮어씁니다. 그렇지 않으면 흐름이 종료됩니다.
runAsyncImpl 메서드는 RxJava의 Flowable 스트림과 연산자를 사용하여 비동기 제어 흐름으로 하위 에이전트를 오케스트레이션합니다.
@Override
protected Flowable<Event> runAsyncImpl(InvocationContext invocationContext) {
// Implements the custom orchestration logic for the story workflow.
// Uses the instance attributes assigned by Pydantic (e.g., self.story_generator).
logger.log(Level.INFO, () -> String.format("[%s] Starting story generation workflow.", name()));
// Stage 1. Initial Story Generation
Flowable<Event> storyGenFlow = runStage(storyGenerator, invocationContext, "StoryGenerator");
// Stage 2: Critic-Reviser Loop (runs after story generation completes)
Flowable<Event> criticReviserFlow = Flowable.defer(() -> {
if (!isStoryGenerated(invocationContext)) {
logger.log(Level.SEVERE,() ->
String.format("[%s] Failed to generate initial story. Aborting after StoryGenerator.",
name()));
return Flowable.empty(); // Stop further processing if no story
}
logger.log(Level.INFO, () ->
String.format("[%s] Story state after generator: %s",
name(), invocationContext.session().state().get("current_story")));
return runStage(loopAgent, invocationContext, "CriticReviserLoop");
});
// Stage 3: Post-Processing (runs after critic-reviser loop completes)
Flowable<Event> postProcessingFlow = Flowable.defer(() -> {
logger.log(Level.INFO, () ->
String.format("[%s] Story state after loop: %s",
name(), invocationContext.session().state().get("current_story")));
return runStage(sequentialAgent, invocationContext, "PostProcessing");
});
// Stage 4: Conditional Regeneration (runs after post-processing completes)
Flowable<Event> conditionalRegenFlow = Flowable.defer(() -> {
String toneCheckResult = (String) invocationContext.session().state().get("tone_check_result");
logger.log(Level.INFO, () -> String.format("[%s] Tone check result: %s", name(), toneCheckResult));
if ("negative".equalsIgnoreCase(toneCheckResult)) {
logger.log(Level.INFO, () ->
String.format("[%s] Tone is negative. Regenerating story...", name()));
return runStage(storyGenerator, invocationContext, "StoryGenerator (Regen)");
} else {
logger.log(Level.INFO, () ->
String.format("[%s] Tone is not negative. Keeping current story.", name()));
return Flowable.empty(); // No regeneration needed
}
});
return Flowable.concatArray(storyGenFlow, criticReviserFlow, postProcessingFlow, conditionalRegenFlow)
.doOnComplete(() -> logger.log(Level.INFO, () -> String.format("[%s] Workflow finished.", name())));
}
// Helper method for a single agent run stage with logging
private Flowable<Event> runStage(BaseAgent agentToRun, InvocationContext ctx, String stageName) {
logger.log(Level.INFO, () -> String.format("[%s] Running %s...", name(), stageName));
return agentToRun
.runAsync(ctx)
.doOnNext(event ->
logger.log(Level.INFO,() ->
String.format("[%s] Event from %s: %s", name(), stageName, event.toJson())))
.doOnError(err ->
logger.log(Level.SEVERE,
String.format("[%s] Error in %s", name(), stageName), err))
.doOnComplete(() ->
logger.log(Level.INFO, () ->
String.format("[%s] %s finished.", name(), stageName)));
}
- 초기
storyGenerator.runAsync(invocationContext)Flowable이 실행됩니다. 그 출력은invocationContext.session().state().get("current_story")에 있을 것으로 예상됩니다. - (
Flowable.concatArray와Flowable.defer덕분에) 다음으로loopAgent의 Flowable이 실행됩니다. LoopAgent는 내부적으로critic과reviser하위 에이전트를 최대maxIterations만큼 순차적으로 호출합니다. 이들은 상태에서current_story와criticism을 읽고 씁니다. - 그 다음,
sequentialAgent의 Flowable이 실행됩니다. 이는grammar_check와tone_check를 차례로 호출하여current_story를 읽고grammar_suggestions와tone_check_result를 상태에 씁니다. - 커스텀 부분: sequentialAgent가 완료된 후,
Flowable.defer내의 로직이invocationContext.session().state()에서 "tone_check_result"를 확인합니다. 만약 "negative"라면,storyGeneratorFlowable이 조건부로 연결되어 다시 실행되고 "current_story"를 덮어씁니다. 그렇지 않으면 빈 Flowable이 사용되고 전체 워크플로우는 완료로 진행됩니다.
3부: LLM 하위 에이전트 정의¶
이것들은 특정 작업을 책임지는 표준 LlmAgent 정의입니다. 그들의 output key 매개변수는 결과를 session.state에 배치하여 다른 에이전트나 커스텀 오케스트레이터가 접근할 수 있도록 하는 데 중요합니다.
지시문에 상태 직접 주입
story_generator의 지시문을 주목하세요. {var} 구문은 플레이스홀더입니다. 지시문이 LLM으로 전송되기 전에, ADK 프레임워크는 (예: {topic})을 session.state['topic']의 값으로 자동으로 대체합니다. 이것은 지시문에서 템플릿을 사용하여 에이전트에 컨텍스트를 제공하는 권장 방법입니다. 자세한 내용은 상태 문서를 참조하세요.
GEMINI_2_FLASH = "gemini-2.0-flash" # 모델 상수 정의
# --- 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), on the following topic: {topic}""",
input_schema=None,
output_key="current_story", # Key for storing output in session state
)
critic = LlmAgent(
name="Critic",
model=GEMINI_2_FLASH,
instruction="""You are a story critic. Review the story provided: {{current_story}}. Provide 1-2 sentences of constructive criticism
on how to improve it. Focus on plot or character.""",
input_schema=None,
output_key="criticism", # Key for storing criticism in session state
)
reviser = LlmAgent(
name="Reviser",
model=GEMINI_2_FLASH,
instruction="""You are a story reviser. Revise the story provided: {{current_story}}, based on the criticism in
{{criticism}}. Output only the revised story.""",
input_schema=None,
output_key="current_story", # Overwrites the original story
)
grammar_check = LlmAgent(
name="GrammarCheck",
model=GEMINI_2_FLASH,
instruction="""You are a grammar checker. Check the grammar of the story provided: {current_story}. Output only the suggested
corrections as a list, or output 'Grammar is good!' if there are no errors.""",
input_schema=None,
output_key="grammar_suggestions",
)
tone_check = LlmAgent(
name="ToneCheck",
model=GEMINI_2_FLASH,
instruction="""You are a tone analyzer. Analyze the tone of the story provided: {current_story}. Output only one word: 'positive' if
the tone is generally positive, 'negative' if the tone is generally negative, or 'neutral'
otherwise.""",
input_schema=None,
output_key="tone_check_result", # This agent's output determines the conditional flow
)
// --- Define the individual LLM agents ---
LlmAgent storyGenerator =
LlmAgent.builder()
.name("StoryGenerator")
.model(MODEL_NAME)
.description("Generates the initial story.")
.instruction(
"""
You are a story writer. Write a short story (around 100 words) about a cat,
based on the topic: {topic}
""")
.inputSchema(null)
.outputKey("current_story") // Key for storing output in session state
.build();
LlmAgent critic =
LlmAgent.builder()
.name("Critic")
.model(MODEL_NAME)
.description("Critiques the story.")
.instruction(
"""
You are a story critic. Review the story: {current_story}. Provide 1-2 sentences of constructive criticism
on how to improve it. Focus on plot or character.
""")
.inputSchema(null)
.outputKey("criticism") // Key for storing criticism in session state
.build();
LlmAgent reviser =
LlmAgent.builder()
.name("Reviser")
.model(MODEL_NAME)
.description("Revises the story based on criticism.")
.instruction(
"""
You are a story reviser. Revise the story: {current_story}, based on the criticism: {criticism}. Output only the revised story.
""")
.inputSchema(null)
.outputKey("current_story") // Overwrites the original story
.build();
LlmAgent grammarCheck =
LlmAgent.builder()
.name("GrammarCheck")
.model(MODEL_NAME)
.description("Checks grammar and suggests corrections.")
.instruction(
"""
You are a grammar checker. Check the grammar of the story: {current_story}. Output only the suggested
corrections as a list, or output 'Grammar is good!' if there are no errors.
""")
.outputKey("grammar_suggestions")
.build();
LlmAgent toneCheck =
LlmAgent.builder()
.name("ToneCheck")
.model(MODEL_NAME)
.description("Analyzes the tone of the story.")
.instruction(
"""
You are a tone analyzer. Analyze the tone of the story: {current_story}. Output only one word: 'positive' if
the tone is generally positive, 'negative' if the tone is generally negative, or 'neutral'
otherwise.
""")
.outputKey("tone_check_result") // This agent's output determines the conditional flow
.build();
LoopAgent loopAgent =
LoopAgent.builder()
.name("CriticReviserLoop")
.description("Iteratively critiques and revises the story.")
.subAgents(critic, reviser)
.maxIterations(2)
.build();
SequentialAgent sequentialAgent =
SequentialAgent.builder()
.name("PostProcessing")
.description("Performs grammar and tone checks sequentially.")
.subAgents(grammarCheck, toneCheck)
.build();
// --- Define the individual LLM agents ---
storyGenerator, err := llmagent.New(llmagent.Config{
Name: "StoryGenerator",
Model: model,
Description: "Generates the initial story.",
Instruction: "You are a story writer. Write a short story (around 100 words) about a cat, based on the topic: {topic}",
OutputKey: "current_story",
})
if err != nil {
log.Fatalf("Failed to create StoryGenerator agent: %v", err)
}
critic, err := llmagent.New(llmagent.Config{
Name: "Critic",
Model: model,
Description: "Critiques the story.",
Instruction: "You are a story critic. Review the story: {current_story}. Provide 1-2 sentences of constructive criticism on how to improve it. Focus on plot or character.",
OutputKey: "criticism",
})
if err != nil {
log.Fatalf("Failed to create Critic agent: %v", err)
}
reviser, err := llmagent.New(llmagent.Config{
Name: "Reviser",
Model: model,
Description: "Revises the story based on criticism.",
Instruction: "You are a story reviser. Revise the story: {current_story}, based on the criticism: {criticism}. Output only the revised story.",
OutputKey: "current_story",
})
if err != nil {
log.Fatalf("Failed to create Reviser agent: %v", err)
}
grammarCheck, err := llmagent.New(llmagent.Config{
Name: "GrammarCheck",
Model: model,
Description: "Checks grammar and suggests corrections.",
Instruction: "You are a grammar checker. Check the grammar of the story: {current_story}. Output only the suggested corrections as a list, or output 'Grammar is good!' if there are no errors.",
OutputKey: "grammar_suggestions",
})
if err != nil {
log.Fatalf("Failed to create GrammarCheck agent: %v", err)
}
toneCheck, err := llmagent.New(llmagent.Config{
Name: "ToneCheck",
Model: model,
Description: "Analyzes the tone of the story.",
Instruction: "You are a tone analyzer. Analyze the tone of the story: {current_story}. Output only one word: 'positive' if the tone is generally positive, 'negative' if the tone is generally negative, or 'neutral' otherwise.",
OutputKey: "tone_check_result",
})
if err != nil {
log.Fatalf("Failed to create ToneCheck agent: %v", err)
}
4부: 커스텀 에이전트 인스턴스화 및 실행¶
마지막으로, StoryFlowAgent를 인스턴스화하고 평소처럼 Runner를 사용합니다.
# --- Create the custom agent instance ---
story_flow_agent = StoryFlowAgent(
name="StoryFlowAgent",
story_generator=story_generator,
critic=critic,
reviser=reviser,
grammar_check=grammar_check,
tone_check=tone_check,
)
INITIAL_STATE = {"topic": "a brave kitten exploring a haunted house"}
# --- Setup Runner and Session ---
async def setup_session_and_runner():
session_service = InMemorySessionService()
session = await session_service.create_session(app_name=APP_NAME, user_id=USER_ID, session_id=SESSION_ID, state=INITIAL_STATE)
logger.info(f"Initial session state: {session.state}")
runner = Runner(
agent=story_flow_agent, # Pass the custom orchestrator agent
app_name=APP_NAME,
session_service=session_service
)
return session_service, runner
# --- Function to Interact with the Agent ---
async def call_agent_async(user_input_topic: str):
"""
Sends a new topic to the agent (overwriting the initial one if needed)
and runs the workflow.
"""
session_service, runner = await setup_session_and_runner()
current_session = await 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["topic"] = user_input_topic
logger.info(f"Updated session state topic to: {user_input_topic}")
content = types.Content(role='user', parts=[types.Part(text=f"Generate a story about: {user_input_topic}")])
events = runner.run_async(user_id=USER_ID, session_id=SESSION_ID, new_message=content)
final_response = "No final response captured."
async 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 = await session_service.get_session(app_name=APP_NAME,
user_id=USER_ID,
session_id=SESSION_ID)
print("Final Session State:")
import json
print(json.dumps(final_session.state, indent=2))
print("-------------------------------\n")
# --- Run the Agent ---
# Note: In Colab, you can directly use 'await' at the top level.
# If running this code as a standalone Python script, you'll need to use asyncio.run() or manage the event loop.
await call_agent_async("a lonely robot finding a friend in a junkyard")
// Instantiate the custom agent, which encapsulates the workflow agents.
storyFlowAgent, err := NewStoryFlowAgent(
storyGenerator,
critic,
reviser,
grammarCheck,
toneCheck,
)
if err != nil {
log.Fatalf("Failed to create story flow agent: %v", err)
}
// --- Run the Agent ---
sessionService := session.InMemoryService()
initialState := map[string]any{
"topic": "a brave kitten exploring a haunted house",
}
sessionInstance, err := sessionService.Create(ctx, &session.CreateRequest{
AppName: appName,
UserID: userID,
State: initialState,
})
if err != nil {
log.Fatalf("Failed to create session: %v", err)
}
userTopic := "a lonely robot finding a friend in a junkyard"
r, err := runner.New(runner.Config{
AppName: appName,
Agent: storyFlowAgent,
SessionService: sessionService,
})
if err != nil {
log.Fatalf("Failed to create runner: %v", err)
}
input := genai.NewContentFromText("Generate a story about: "+userTopic, genai.RoleUser)
events := r.Run(ctx, userID, sessionInstance.Session.ID(), input, agent.RunConfig{
StreamingMode: agent.StreamingModeSSE,
})
var finalResponse string
for event, err := range events {
if err != nil {
log.Fatalf("An error occurred during agent execution: %v", err)
}
for _, part := range event.Content.Parts {
// Accumulate text from all parts of the final response.
finalResponse += part.Text
}
}
fmt.Println("\n--- Agent Interaction Result ---")
fmt.Println("Agent Final Response: " + finalResponse)
finalSession, err := sessionService.Get(ctx, &session.GetRequest{
UserID: userID,
AppName: appName,
SessionID: sessionInstance.Session.ID(),
})
if err != nil {
log.Fatalf("Failed to retrieve final session: %v", err)
}
fmt.Println("Final Session State:", finalSession.Session.State())
}
// --- Function to Interact with the Agent ---
// Sends a new topic to the agent (overwriting the initial one if needed)
// and runs the workflow.
public static void runAgent(StoryFlowAgentExample agent, String userTopic) {
// --- Setup Runner and Session ---
InMemoryRunner runner = new InMemoryRunner(agent);
Map<String, Object> initialState = new HashMap<>();
initialState.put("topic", "a brave kitten exploring a haunted house");
Session session =
runner
.sessionService()
.createSession(APP_NAME, USER_ID, new ConcurrentHashMap<>(initialState), SESSION_ID)
.blockingGet();
logger.log(Level.INFO, () -> String.format("Initial session state: %s", session.state()));
session.state().put("topic", userTopic); // Update the state in the retrieved session
logger.log(Level.INFO, () -> String.format("Updated session state topic to: %s", userTopic));
Content userMessage = Content.fromParts(Part.fromText("Generate a story about: " + userTopic));
// Use the modified session object for the run
Flowable<Event> eventStream = runner.runAsync(USER_ID, session.id(), userMessage);
final String[] finalResponse = {"No final response captured."};
eventStream.blockingForEach(
event -> {
if (event.finalResponse() && event.content().isPresent()) {
String author = event.author() != null ? event.author() : "UNKNOWN_AUTHOR";
Optional<String> textOpt =
event
.content()
.flatMap(Content::parts)
.filter(parts -> !parts.isEmpty())
.map(parts -> parts.get(0).text().orElse(""));
logger.log(Level.INFO, () ->
String.format("Potential final response from [%s]: %s", author, textOpt.orElse("N/A")));
textOpt.ifPresent(text -> finalResponse[0] = text);
}
});
System.out.println("\n--- Agent Interaction Result ---");
System.out.println("Agent Final Response: " + finalResponse[0]);
// Retrieve session again to see the final state after the run
Session finalSession =
runner
.sessionService()
.getSession(APP_NAME, USER_ID, SESSION_ID, Optional.empty())
.blockingGet();
assert finalSession != null;
System.out.println("Final Session State:" + finalSession.state());
System.out.println("-------------------------------\n");
}
(참고: import 및 실행 로직을 포함한 전체 실행 가능 코드는 아래 링크에서 찾을 수 있습니다.)
전체 코드 예제¶
Storyflow 에이전트
# StoryFlowAgent 예제의 전체 실행 가능 코드
# Copyright 2025 Google LLC
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
import logging
from typing import AsyncGenerator
from typing_extensions import override
from google.adk.agents import LlmAgent, BaseAgent, LoopAgent, SequentialAgent
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
# --- Constants ---
APP_NAME = "story_app"
USER_ID = "12345"
SESSION_ID = "123344"
GEMINI_2_FLASH = "gemini-2.0-flash"
# --- Configure Logging ---
logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)
# --- 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
critic: LlmAgent
reviser: LlmAgent
grammar_check: LlmAgent
tone_check: LlmAgent
loop_agent: LoopAgent
sequential_agent: SequentialAgent
# 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,
critic: LlmAgent,
reviser: LlmAgent,
grammar_check: LlmAgent,
tone_check: 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.
"""
# Create internal agents *before* calling super().__init__
loop_agent = LoopAgent(
name="CriticReviserLoop", sub_agents=[critic, reviser], max_iterations=2
)
sequential_agent = SequentialAgent(
name="PostProcessing", sub_agents=[grammar_check, tone_check]
)
# Define the sub_agents list for the framework
sub_agents_list = [
story_generator,
loop_agent,
sequential_agent,
]
# Pydantic will validate and assign them based on the class annotations.
super().__init__(
name=name,
story_generator=story_generator,
critic=critic,
reviser=reviser,
grammar_check=grammar_check,
tone_check=tone_check,
loop_agent=loop_agent,
sequential_agent=sequential_agent,
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')}")
# 2. Critic-Reviser Loop
logger.info(f"[{self.name}] Running CriticReviserLoop...")
# Use the loop_agent instance attribute assigned during init
async for event in self.loop_agent.run_async(ctx):
logger.info(f"[{self.name}] Event from CriticReviserLoop: {event.model_dump_json(indent=2, exclude_none=True)}")
yield event
logger.info(f"[{self.name}] Story state after loop: {ctx.session.state.get('current_story')}")
# 3. Sequential Post-Processing (Grammar and Tone Check)
logger.info(f"[{self.name}] Running PostProcessing...")
# Use the sequential_agent instance attribute assigned during init
async for event in self.sequential_agent.run_async(ctx):
logger.info(f"[{self.name}] Event from PostProcessing: {event.model_dump_json(indent=2, exclude_none=True)}")
yield event
# 4. Tone-Based Conditional Logic
tone_check_result = ctx.session.state.get("tone_check_result")
logger.info(f"[{self.name}] Tone check result: {tone_check_result}")
if tone_check_result == "negative":
logger.info(f"[{self.name}] Tone is negative. Regenerating story...")
async for event in self.story_generator.run_async(ctx):
logger.info(f"[{self.name}] Event from StoryGenerator (Regen): {event.model_dump_json(indent=2, exclude_none=True)}")
yield event
else:
logger.info(f"[{self.name}] Tone is not negative. Keeping current story.")
pass
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), on the following topic: {topic}""",
input_schema=None,
output_key="current_story", # Key for storing output in session state
)
critic = LlmAgent(
name="Critic",
model=GEMINI_2_FLASH,
instruction="""You are a story critic. Review the story provided: {{current_story}}. Provide 1-2 sentences of constructive criticism
on how to improve it. Focus on plot or character.""",
input_schema=None,
output_key="criticism", # Key for storing criticism in session state
)
reviser = LlmAgent(
name="Reviser",
model=GEMINI_2_FLASH,
instruction="""You are a story reviser. Revise the story provided: {{current_story}}, based on the criticism in
{{criticism}}. Output only the revised story.""",
input_schema=None,
output_key="current_story", # Overwrites the original story
)
grammar_check = LlmAgent(
name="GrammarCheck",
model=GEMINI_2_FLASH,
instruction="""You are a grammar checker. Check the grammar of the story provided: {current_story}. Output only the suggested
corrections as a list, or output 'Grammar is good!' if there are no errors.""",
input_schema=None,
output_key="grammar_suggestions",
)
tone_check = LlmAgent(
name="ToneCheck",
model=GEMINI_2_FLASH,
instruction="""You are a tone analyzer. Analyze the tone of the story provided: {current_story}. Output only one word: 'positive' if
the tone is generally positive, 'negative' if the tone is generally negative, or 'neutral'
otherwise.""",
input_schema=None,
output_key="tone_check_result", # This agent's output determines the conditional flow
)
# --- Create the custom agent instance ---
story_flow_agent = StoryFlowAgent(
name="StoryFlowAgent",
story_generator=story_generator,
critic=critic,
reviser=reviser,
grammar_check=grammar_check,
tone_check=tone_check,
)
INITIAL_STATE = {"topic": "a brave kitten exploring a haunted house"}
# --- Setup Runner and Session ---
async def setup_session_and_runner():
session_service = InMemorySessionService()
session = await session_service.create_session(app_name=APP_NAME, user_id=USER_ID, session_id=SESSION_ID, state=INITIAL_STATE)
logger.info(f"Initial session state: {session.state}")
runner = Runner(
agent=story_flow_agent, # Pass the custom orchestrator agent
app_name=APP_NAME,
session_service=session_service
)
return session_service, runner
# --- Function to Interact with the Agent ---
async def call_agent_async(user_input_topic: str):
"""
Sends a new topic to the agent (overwriting the initial one if needed)
and runs the workflow.
"""
session_service, runner = await setup_session_and_runner()
current_session = await 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["topic"] = user_input_topic
logger.info(f"Updated session state topic to: {user_input_topic}")
content = types.Content(role='user', parts=[types.Part(text=f"Generate a story about: {user_input_topic}")])
events = runner.run_async(user_id=USER_ID, session_id=SESSION_ID, new_message=content)
final_response = "No final response captured."
async 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 = await session_service.get_session(app_name=APP_NAME,
user_id=USER_ID,
session_id=SESSION_ID)
print("Final Session State:")
import json
print(json.dumps(final_session.state, indent=2))
print("-------------------------------\n")
# --- Run the Agent ---
# Note: In Colab, you can directly use 'await' at the top level.
# If running this code as a standalone Python script, you'll need to use asyncio.run() or manage the event loop.
await call_agent_async("a lonely robot finding a friend in a junkyard")
# StoryFlowAgent 예제의 전체 실행 가능 코드
package main
import (
"context"
"fmt"
"iter"
"log"
"google.golang.org/adk/agent/workflowagents/loopagent"
"google.golang.org/adk/agent/workflowagents/sequentialagent"
"google.golang.org/adk/agent"
"google.golang.org/adk/agent/llmagent"
"google.golang.org/adk/model/gemini"
"google.golang.org/adk/runner"
"google.golang.org/adk/session"
"google.golang.org/genai"
)
// StoryFlowAgent is a custom agent that orchestrates a story generation workflow.
// It encapsulates the logic of running sub-agents in a specific sequence.
type StoryFlowAgent struct {
storyGenerator agent.Agent
revisionLoopAgent agent.Agent
postProcessorAgent agent.Agent
}
// NewStoryFlowAgent creates and configures the entire custom agent workflow.
// It takes individual LLM agents as input and internally creates the necessary
// workflow agents (loop, sequential), returning the final orchestrator agent.
func NewStoryFlowAgent(
storyGenerator,
critic,
reviser,
grammarCheck,
toneCheck agent.Agent,
) (agent.Agent, error) {
loopAgent, err := loopagent.New(loopagent.Config{
MaxIterations: 2,
AgentConfig: agent.Config{
Name: "CriticReviserLoop",
SubAgents: []agent.Agent{critic, reviser},
},
})
if err != nil {
return nil, fmt.Errorf("failed to create loop agent: %w", err)
}
sequentialAgent, err := sequentialagent.New(sequentialagent.Config{
AgentConfig: agent.Config{
Name: "PostProcessing",
SubAgents: []agent.Agent{grammarCheck, toneCheck},
},
})
if err != nil {
return nil, fmt.Errorf("failed to create sequential agent: %w", err)
}
// The StoryFlowAgent struct holds the agents needed for the Run method.
orchestrator := &StoryFlowAgent{
storyGenerator: storyGenerator,
revisionLoopAgent: loopAgent,
postProcessorAgent: sequentialAgent,
}
// agent.New creates the final agent, wiring up the Run method.
return agent.New(agent.Config{
Name: "StoryFlowAgent",
Description: "Orchestrates story generation, critique, revision, and checks.",
SubAgents: []agent.Agent{storyGenerator, loopAgent, sequentialAgent},
Run: orchestrator.Run,
})
}
// Run defines the custom execution logic for the StoryFlowAgent.
func (s *StoryFlowAgent) Run(ctx agent.InvocationContext) iter.Seq2[*session.Event, error] {
return func(yield func(*session.Event, error) bool) {
// Stage 1: Initial Story Generation
for event, err := range s.storyGenerator.Run(ctx) {
if err != nil {
yield(nil, fmt.Errorf("story generator failed: %w", err))
return
}
if !yield(event, nil) {
return
}
}
// Check if story was generated before proceeding
currentStory, err := ctx.Session().State().Get("current_story")
if err != nil || currentStory == "" {
log.Println("Failed to generate initial story. Aborting workflow.")
return
}
// Stage 2: Critic-Reviser Loop
for event, err := range s.revisionLoopAgent.Run(ctx) {
if err != nil {
yield(nil, fmt.Errorf("loop agent failed: %w", err))
return
}
if !yield(event, nil) {
return
}
}
// Stage 3: Post-Processing
for event, err := range s.postProcessorAgent.Run(ctx) {
if err != nil {
yield(nil, fmt.Errorf("sequential agent failed: %w", err))
return
}
if !yield(event, nil) {
return
}
}
// Stage 4: Conditional Regeneration
toneResult, err := ctx.Session().State().Get("tone_check_result")
if err != nil {
log.Printf("Could not read tone_check_result from state: %v. Assuming tone is not negative.", err)
return
}
if tone, ok := toneResult.(string); ok && tone == "negative" {
log.Println("Tone is negative. Regenerating story...")
for event, err := range s.storyGenerator.Run(ctx) {
if err != nil {
yield(nil, fmt.Errorf("story regeneration failed: %w", err))
return
}
if !yield(event, nil) {
return
}
}
} else {
log.Println("Tone is not negative. Keeping current story.")
}
}
}
const (
modelName = "gemini-2.0-flash"
appName = "story_app"
userID = "user_12345"
)
func main() {
ctx := context.Background()
model, err := gemini.NewModel(ctx, modelName, &genai.ClientConfig{})
if err != nil {
log.Fatalf("Failed to create model: %v", err)
}
// --- Define the individual LLM agents ---
storyGenerator, err := llmagent.New(llmagent.Config{
Name: "StoryGenerator",
Model: model,
Description: "Generates the initial story.",
Instruction: "You are a story writer. Write a short story (around 100 words) about a cat, based on the topic: {topic}",
OutputKey: "current_story",
})
if err != nil {
log.Fatalf("Failed to create StoryGenerator agent: %v", err)
}
critic, err := llmagent.New(llmagent.Config{
Name: "Critic",
Model: model,
Description: "Critiques the story.",
Instruction: "You are a story critic. Review the story: {current_story}. Provide 1-2 sentences of constructive criticism on how to improve it. Focus on plot or character.",
OutputKey: "criticism",
})
if err != nil {
log.Fatalf("Failed to create Critic agent: %v", err)
}
reviser, err := llmagent.New(llmagent.Config{
Name: "Reviser",
Model: model,
Description: "Revises the story based on criticism.",
Instruction: "You are a story reviser. Revise the story: {current_story}, based on the criticism: {criticism}. Output only the revised story.",
OutputKey: "current_story",
})
if err != nil {
log.Fatalf("Failed to create Reviser agent: %v", err)
}
grammarCheck, err := llmagent.New(llmagent.Config{
Name: "GrammarCheck",
Model: model,
Description: "Checks grammar and suggests corrections.",
Instruction: "You are a grammar checker. Check the grammar of the story: {current_story}. Output only the suggested corrections as a list, or output 'Grammar is good!' if there are no errors.",
OutputKey: "grammar_suggestions",
})
if err != nil {
log.Fatalf("Failed to create GrammarCheck agent: %v", err)
}
toneCheck, err := llmagent.New(llmagent.Config{
Name: "ToneCheck",
Model: model,
Description: "Analyzes the tone of the story.",
Instruction: "You are a tone analyzer. Analyze the tone of the story: {current_story}. Output only one word: 'positive' if the tone is generally positive, 'negative' if the tone is generally negative, or 'neutral' otherwise.",
OutputKey: "tone_check_result",
})
if err != nil {
log.Fatalf("Failed to create ToneCheck agent: %v", err)
}
// Instantiate the custom agent, which encapsulates the workflow agents.
storyFlowAgent, err := NewStoryFlowAgent(
storyGenerator,
critic,
reviser,
grammarCheck,
toneCheck,
)
if err != nil {
log.Fatalf("Failed to create story flow agent: %v", err)
}
// --- Run the Agent ---
sessionService := session.InMemoryService()
initialState := map[string]any{
"topic": "a brave kitten exploring a haunted house",
}
sessionInstance, err := sessionService.Create(ctx, &session.CreateRequest{
AppName: appName,
UserID: userID,
State: initialState,
})
if err != nil {
log.Fatalf("Failed to create session: %v", err)
}
userTopic := "a lonely robot finding a friend in a junkyard"
r, err := runner.New(runner.Config{
AppName: appName,
Agent: storyFlowAgent,
SessionService: sessionService,
})
if err != nil {
log.Fatalf("Failed to create runner: %v", err)
}
input := genai.NewContentFromText("Generate a story about: "+userTopic, genai.RoleUser)
events := r.Run(ctx, userID, sessionInstance.Session.ID(), input, agent.RunConfig{
StreamingMode: agent.StreamingModeSSE,
})
var finalResponse string
for event, err := range events {
if err != nil {
log.Fatalf("An error occurred during agent execution: %v", err)
}
for _, part := range event.Content.Parts {
// Accumulate text from all parts of the final response.
finalResponse += part.Text
}
}
fmt.Println("\n--- Agent Interaction Result ---")
fmt.Println("Agent Final Response: " + finalResponse)
finalSession, err := sessionService.Get(ctx, &session.GetRequest{
UserID: userID,
AppName: appName,
SessionID: sessionInstance.Session.ID(),
})
if err != nil {
log.Fatalf("Failed to retrieve final session: %v", err)
}
fmt.Println("Final Session State:", finalSession.Session.State())
}
# StoryFlowAgent 예제의 전체 실행 가능 코드
import com.google.adk.agents.LlmAgent;
import com.google.adk.agents.BaseAgent;
import com.google.adk.agents.InvocationContext;
import com.google.adk.agents.LoopAgent;
import com.google.adk.agents.SequentialAgent;
import com.google.adk.events.Event;
import com.google.adk.runner.InMemoryRunner;
import com.google.adk.sessions.Session;
import com.google.genai.types.Content;
import com.google.genai.types.Part;
import io.reactivex.rxjava3.core.Flowable;
import java.util.HashMap;
import java.util.List;
import java.util.Map;
import java.util.Optional;
import java.util.concurrent.ConcurrentHashMap;
import java.util.logging.Level;
import java.util.logging.Logger;
public class StoryFlowAgentExample extends BaseAgent {
// --- Constants ---
private static final String APP_NAME = "story_app";
private static final String USER_ID = "user_12345";
private static final String SESSION_ID = "session_123344";
private static final String MODEL_NAME = "gemini-2.0-flash"; // Ensure this model is available
private static final Logger logger = Logger.getLogger(StoryFlowAgentExample.class.getName());
private final LlmAgent storyGenerator;
private final LoopAgent loopAgent;
private final SequentialAgent sequentialAgent;
public StoryFlowAgentExample(
String name, LlmAgent storyGenerator, LoopAgent loopAgent, SequentialAgent sequentialAgent) {
super(
name,
"Orchestrates story generation, critique, revision, and checks.",
List.of(storyGenerator, loopAgent, sequentialAgent),
null,
null);
this.storyGenerator = storyGenerator;
this.loopAgent = loopAgent;
this.sequentialAgent = sequentialAgent;
}
public static void main(String[] args) {
// --- Define the individual LLM agents ---
LlmAgent storyGenerator =
LlmAgent.builder()
.name("StoryGenerator")
.model(MODEL_NAME)
.description("Generates the initial story.")
.instruction(
"""
You are a story writer. Write a short story (around 100 words) about a cat,
based on the topic: {topic}
""")
.inputSchema(null)
.outputKey("current_story") // Key for storing output in session state
.build();
LlmAgent critic =
LlmAgent.builder()
.name("Critic")
.model(MODEL_NAME)
.description("Critiques the story.")
.instruction(
"""
You are a story critic. Review the story: {current_story}. Provide 1-2 sentences of constructive criticism
on how to improve it. Focus on plot or character.
""")
.inputSchema(null)
.outputKey("criticism") // Key for storing criticism in session state
.build();
LlmAgent reviser =
LlmAgent.builder()
.name("Reviser")
.model(MODEL_NAME)
.description("Revises the story based on criticism.")
.instruction(
"""
You are a story reviser. Revise the story: {current_story}, based on the criticism: {criticism}. Output only the revised story.
""")
.inputSchema(null)
.outputKey("current_story") // Overwrites the original story
.build();
LlmAgent grammarCheck =
LlmAgent.builder()
.name("GrammarCheck")
.model(MODEL_NAME)
.description("Checks grammar and suggests corrections.")
.instruction(
"""
You are a grammar checker. Check the grammar of the story: {current_story}. Output only the suggested
corrections as a list, or output 'Grammar is good!' if there are no errors.
""")
.outputKey("grammar_suggestions")
.build();
LlmAgent toneCheck =
LlmAgent.builder()
.name("ToneCheck")
.model(MODEL_NAME)
.description("Analyzes the tone of the story.")
.instruction(
"""
You are a tone analyzer. Analyze the tone of the story: {current_story}. Output only one word: 'positive' if
the tone is generally positive, 'negative' if the tone is generally negative, or 'neutral'
otherwise.
""")
.outputKey("tone_check_result") // This agent's output determines the conditional flow
.build();
LoopAgent loopAgent =
LoopAgent.builder()
.name("CriticReviserLoop")
.description("Iteratively critiques and revises the story.")
.subAgents(critic, reviser)
.maxIterations(2)
.build();
SequentialAgent sequentialAgent =
SequentialAgent.builder()
.name("PostProcessing")
.description("Performs grammar and tone checks sequentially.")
.subAgents(grammarCheck, toneCheck)
.build();
StoryFlowAgentExample storyFlowAgentExample =
new StoryFlowAgentExample(APP_NAME, storyGenerator, loopAgent, sequentialAgent);
// --- Run the Agent ---
runAgent(storyFlowAgentExample, "a lonely robot finding a friend in a junkyard");
}
// --- Function to Interact with the Agent ---
// Sends a new topic to the agent (overwriting the initial one if needed)
// and runs the workflow.
public static void runAgent(StoryFlowAgentExample agent, String userTopic) {
// --- Setup Runner and Session ---
InMemoryRunner runner = new InMemoryRunner(agent);
Map<String, Object> initialState = new HashMap<>();
initialState.put("topic", "a brave kitten exploring a haunted house");
Session session =
runner
.sessionService()
.createSession(APP_NAME, USER_ID, new ConcurrentHashMap<>(initialState), SESSION_ID)
.blockingGet();
logger.log(Level.INFO, () -> String.format("Initial session state: %s", session.state()));
session.state().put("topic", userTopic); // Update the state in the retrieved session
logger.log(Level.INFO, () -> String.format("Updated session state topic to: %s", userTopic));
Content userMessage = Content.fromParts(Part.fromText("Generate a story about: " + userTopic));
// Use the modified session object for the run
Flowable<Event> eventStream = runner.runAsync(USER_ID, session.id(), userMessage);
final String[] finalResponse = {"No final response captured."};
eventStream.blockingForEach(
event -> {
if (event.finalResponse() && event.content().isPresent()) {
String author = event.author() != null ? event.author() : "UNKNOWN_AUTHOR";
Optional<String> textOpt =
event
.content()
.flatMap(Content::parts)
.filter(parts -> !parts.isEmpty())
.map(parts -> parts.get(0).text().orElse(""));
logger.log(Level.INFO, () ->
String.format("Potential final response from [%s]: %s", author, textOpt.orElse("N/A")));
textOpt.ifPresent(text -> finalResponse[0] = text);
}
});
System.out.println("\n--- Agent Interaction Result ---");
System.out.println("Agent Final Response: " + finalResponse[0]);
// Retrieve session again to see the final state after the run
Session finalSession =
runner
.sessionService()
.getSession(APP_NAME, USER_ID, SESSION_ID, Optional.empty())
.blockingGet();
assert finalSession != null;
System.out.println("Final Session State:" + finalSession.state());
System.out.println("-------------------------------\n");
}
private boolean isStoryGenerated(InvocationContext ctx) {
Object currentStoryObj = ctx.session().state().get("current_story");
return currentStoryObj != null && !String.valueOf(currentStoryObj).isEmpty();
}
@Override
protected Flowable<Event> runAsyncImpl(InvocationContext invocationContext) {
// Implements the custom orchestration logic for the story workflow.
// Uses the instance attributes assigned by Pydantic (e.g., self.story_generator).
logger.log(Level.INFO, () -> String.format("[%s] Starting story generation workflow.", name()));
// Stage 1. Initial Story Generation
Flowable<Event> storyGenFlow = runStage(storyGenerator, invocationContext, "StoryGenerator");
// Stage 2: Critic-Reviser Loop (runs after story generation completes)
Flowable<Event> criticReviserFlow = Flowable.defer(() -> {
if (!isStoryGenerated(invocationContext)) {
logger.log(Level.SEVERE,() ->
String.format("[%s] Failed to generate initial story. Aborting after StoryGenerator.",
name()));
return Flowable.empty(); // Stop further processing if no story
}
logger.log(Level.INFO, () ->
String.format("[%s] Story state after generator: %s",
name(), invocationContext.session().state().get("current_story")));
return runStage(loopAgent, invocationContext, "CriticReviserLoop");
});
// Stage 3: Post-Processing (runs after critic-reviser loop completes)
Flowable<Event> postProcessingFlow = Flowable.defer(() -> {
logger.log(Level.INFO, () ->
String.format("[%s] Story state after loop: %s",
name(), invocationContext.session().state().get("current_story")));
return runStage(sequentialAgent, invocationContext, "PostProcessing");
});
// Stage 4: Conditional Regeneration (runs after post-processing completes)
Flowable<Event> conditionalRegenFlow = Flowable.defer(() -> {
String toneCheckResult = (String) invocationContext.session().state().get("tone_check_result");
logger.log(Level.INFO, () -> String.format("[%s] Tone check result: %s", name(), toneCheckResult));
if ("negative".equalsIgnoreCase(toneCheckResult)) {
logger.log(Level.INFO, () ->
String.format("[%s] Tone is negative. Regenerating story...", name()));
return runStage(storyGenerator, invocationContext, "StoryGenerator (Regen)");
} else {
logger.log(Level.INFO, () ->
String.format("[%s] Tone is not negative. Keeping current story.", name()));
return Flowable.empty(); // No regeneration needed
}
});
return Flowable.concatArray(storyGenFlow, criticReviserFlow, postProcessingFlow, conditionalRegenFlow)
.doOnComplete(() -> logger.log(Level.INFO, () -> String.format("[%s] Workflow finished.", name())));
}
// Helper method for a single agent run stage with logging
private Flowable<Event> runStage(BaseAgent agentToRun, InvocationContext ctx, String stageName) {
logger.log(Level.INFO, () -> String.format("[%s] Running %s...", name(), stageName));
return agentToRun
.runAsync(ctx)
.doOnNext(event ->
logger.log(Level.INFO,() ->
String.format("[%s] Event from %s: %s", name(), stageName, event.toJson())))
.doOnError(err ->
logger.log(Level.SEVERE,
String.format("[%s] Error in %s", name(), stageName), err))
.doOnComplete(() ->
logger.log(Level.INFO, () ->
String.format("[%s] %s finished.", name(), stageName)));
}
@Override
protected Flowable<Event> runLiveImpl(InvocationContext invocationContext) {
return Flowable.error(new UnsupportedOperationException("runLive not implemented."));
}
}