chore: Creates a sample main.py to demonstrate how to use adk as a library for the workflow agent in docs.

For Issue #573.

PiperOrigin-RevId: 760380509
This commit is contained in:
Wei Sun (Jack) 2025-05-18 16:16:12 -07:00 committed by Copybara-Service
parent 2ad1f79422
commit 9e767b3fe1
4 changed files with 241 additions and 43 deletions

View File

@ -0,0 +1,12 @@
# Workflow Agent Sample - SequentialAgent
Sample query:
* Write a quicksort method in python.
* Write a python function to do bubble sort.
To run in cli (after installing `google-adk`):
* `uv run main.py` (or `python main.py`)
Check sample output in `sample.output` file in this folder.

View File

@ -15,80 +15,97 @@
from google.adk.agents.llm_agent import LlmAgent
from google.adk.agents.sequential_agent import SequentialAgent
# Part of agent.py --> Follow https://google.github.io/adk-docs/get-started/quickstart/ to learn the setup
# --- 1. Define Sub-Agents for Each Pipeline Stage ---
# Code Writer Agent
# Takes the initial specification (from user query) and writes code.
code_writer_agent = LlmAgent(
name="code_writer_agent",
model="gemini-1.5-flash-001",
instruction="""You are a Code Writer AI.
Based on the user's request, write the initial Python code.
Output *only* the raw code block.
""",
description="Writes initial code based on a specification.",
# Stores its output (the generated code) into the session state
# under the key 'generated_code'.
output_key="generated_code",
name="CodeWriterAgent",
model="gemini-1.5-flash",
# Change 3: Improved instruction
instruction="""You are a Python Code Generator.
Based *only* on the user's request, write Python code that fulfills the requirement.
Output *only* the complete Python code block, enclosed in triple backticks (```python ... ```).
Do not add any other text before or after the code block.
""",
description="Writes initial Python code based on a specification.",
output_key="generated_code", # Stores output in state['generated_code']
)
# Code Reviewer Agent
# Takes the code generated by the previous agent (read from state) and provides feedback.
code_reviewer_agent = LlmAgent(
name="code_reviewer_agent",
model="gemini-2.0-flash-001",
instruction="""You are a Code Reviewer AI.
name="CodeReviewerAgent",
model="gemini-2.0-flash",
# Change 3: Improved instruction, correctly using state key injection
instruction="""You are an expert Python Code Reviewer.
Your task is to provide constructive feedback on the provided code.
Review the below Python code.
**Code to Review:**
```python
{generated_code}
```
```
{generated_code}
```
**Review Criteria:**
1. **Correctness:** Does the code work as intended? Are there logic errors?
2. **Readability:** Is the code clear and easy to understand? Follows PEP 8 style guidelines?
3. **Efficiency:** Is the code reasonably efficient? Any obvious performance bottlenecks?
4. **Edge Cases:** Does the code handle potential edge cases or invalid inputs gracefully?
5. **Best Practices:** Does the code follow common Python best practices?
Provide constructive feedback on potential errors, style issues, or improvements.
Focus on clarity and correctness.
Output only the review comments.
""",
**Output:**
Provide your feedback as a concise, bulleted list. Focus on the most important points for improvement.
If the code is excellent and requires no changes, simply state: "No major issues found."
Output *only* the review comments or the "No major issues" statement.
""",
description="Reviews code and provides feedback.",
# Stores its output (the review comments) into the session state
# under the key 'review_comments'.
output_key="review_comments",
output_key="review_comments", # Stores output in state['review_comments']
)
# Code Refactorer Agent
# Takes the original code and the review comments (read from state) and refactors the code.
code_refactorer_agent = LlmAgent(
name="code_refactorer_agent",
model="gemini-2.0-flash-001",
instruction="""You are a Code Refactorer AI.
name="CodeRefactorerAgent",
model="gemini-2.0-flash",
# Change 3: Improved instruction, correctly using state key injection
instruction="""You are a Python Code Refactoring AI.
Your goal is to improve the given Python code based on the provided review comments.
Below is the original Python code:
**Original Code:**
```python
{generated_code}
```
```
{generated_code}
```
**Review Comments:**
{review_comments}
Below are the review comments:
**Task:**
Carefully apply the suggestions from the review comments to refactor the original code.
If the review comments state "No major issues found," return the original code unchanged.
Ensure the final code is complete, functional, and includes necessary imports and docstrings.
{review_comments}
Refactor the code based on the provided feedback.
Output *only* the final, refactored code block.
""",
**Output:**
Output *only* the final, refactored Python code block, enclosed in triple backticks (```python ... ```).
Do not add any other text before or after the code block.
""",
description="Refactors code based on review comments.",
# Stores its output (the refactored code) into the session state
# under the key 'refactored_code'.
output_key="refactored_code",
output_key="refactored_code", # Stores output in state['refactored_code']
)
# --- 2. Create the SequentialAgent ---
# This agent orchestrates the pipeline by running the sub_agents in order.
code_pipeline_agent = SequentialAgent(
name="code_pipeline_agent",
name="CodePipelineAgent",
sub_agents=[code_writer_agent, code_reviewer_agent, code_refactorer_agent],
description=(
"Executes a sequence of code writing, reviewing, and refactoring."
),
# The agents will run in the order provided: Writer -> Reviewer -> Refactorer
)
# For ADK tools compatibility, the root agent must be named `root_agent`
root_agent = code_pipeline_agent

View File

@ -0,0 +1,87 @@
# 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 asyncio
from typing import cast
import agent
from dotenv import load_dotenv
from google.adk.cli.utils import logs
from google.adk.runners import InMemoryRunner
from google.adk.sessions import Session
from google.genai import types
load_dotenv(override=True)
logs.log_to_tmp_folder()
async def main():
app_name = 'my_app'
user_id_1 = 'user1'
runner = InMemoryRunner(
app_name=app_name,
agent=agent.root_agent,
)
async def run_prompt(session: Session, new_message: str) -> Session:
content = types.Content(
role='user', parts=[types.Part.from_text(text=new_message)]
)
print('** User says:', content.model_dump(exclude_none=True))
async for event in runner.run_async(
user_id=user_id_1,
session_id=session.id,
new_message=content,
):
if not event.content or not event.content.parts:
continue
if event.content.parts[0].text:
print(f'** {event.author}: {event.content.parts[0].text}')
elif event.content.parts[0].function_call:
print(
f'** {event.author}: fc /'
f' {event.content.parts[0].function_call.name} /'
f' {event.content.parts[0].function_call.args}\n'
)
elif event.content.parts[0].function_response:
print(
f'** {event.author}: fr /'
f' {event.content.parts[0].function_response.name} /'
f' {event.content.parts[0].function_response.response}\n'
)
return cast(
Session,
await runner.session_service.get_session(
app_name=app_name, user_id=user_id_1, session_id=session.id
),
)
session_1 = await runner.session_service.create_session(
app_name=app_name, user_id=user_id_1
)
print(f'----Session to create memory: {session_1.id} ----------------------')
session_1 = await run_prompt(
session_1, 'Write a python function to do quicksort.'
)
session_1 = await run_prompt(
session_1, 'Write another python function to do bubble sort.'
)
print('-------------------------------------------------------------------')
if __name__ == '__main__':
asyncio.run(main())

View File

@ -0,0 +1,82 @@
----Session to create memory: a05683c9-4a19-4ba8-b810-e90fc48be566 ----------------------
** User says: {'parts': [{'text': 'Write a python function to do quicksort.'}], 'role': 'user'}
** CodeWriterAgent: ```python
def quicksort(arr):
if len(arr) < 2:
return arr
else:
pivot = arr[0]
less = [i for i in arr[1:] if i <= pivot]
greater = [i for i in arr[1:] if i > pivot]
return quicksort(less) + [pivot] + quicksort(greater)
```
** CodeReviewerAgent: * **Correctness:** The code is functionally correct.
* **Readability:** The code is fairly readable, utilizing list comprehensions effectively.
* **Efficiency:** The space complexity is `O(n)` due to the creation of new lists (`less` and `greater`) in each recursive call. Choosing the first element as the pivot can lead to `O(n^2)` time complexity in the worst-case scenario (e.g., a sorted array). A better pivot selection strategy (e.g., random pivot) would improve average performance.
* **Best Practices:** While the code works, it's not the most efficient implementation of quicksort due to the space complexity and potential for worst-case time complexity. An in-place quicksort algorithm would be more efficient in terms of space.
** CodeRefactorerAgent: ```python
import random
def quicksort(arr):
"""
Sorts a list using the quicksort algorithm. This implementation
uses a randomly selected pivot to improve average-case performance
and performs the sort in-place to reduce space complexity.
Args:
arr (list): The list to be sorted.
Returns:
list: The sorted list.
"""
def _quicksort(arr, low, high):
if low < high:
pi = partition(arr, low, high)
_quicksort(arr, low, pi-1)
_quicksort(arr, pi+1, high)
def partition(arr, low, high):
# Choose a random pivot
pivot_index = random.randint(low, high)
arr[pivot_index], arr[high] = arr[high], arr[pivot_index]
pivot = arr[high]
i = (low - 1)
for j in range(low, high):
if arr[j] <= pivot:
i += 1
arr[i], arr[j] = arr[j], arr[i]
arr[i + 1], arr[high] = arr[high], arr[i + 1]
return (i + 1)
_quicksort(arr, 0, len(arr)-1)
return arr
```
** User says: {'parts': [{'text': 'Write another python function to do bubble sort.'}], 'role': 'user'}
** CodeWriterAgent: ```python
def bubble_sort(arr):
n = len(arr)
for i in range(n):
for j in range(0, n-i-1):
if arr[j] > arr[j+1]:
arr[j], arr[j+1] = arr[j+1], arr[j]
return arr
```
** CodeReviewerAgent: No major issues found.
** CodeRefactorerAgent: ```python
def bubble_sort(arr):
n = len(arr)
for i in range(n):
for j in range(0, n-i-1):
if arr[j] > arr[j+1]:
arr[j], arr[j+1] = arr[j+1], arr[j]
return arr
```
-------------------------------------------------------------------