8000 chore: Creates a sample main.py to demonstrate how to use adk as a li… · wongjoey/adk-python@9e767b3 · GitHub
[go: up one dir, main page]
Skip to content

Commit 9e767b3

Browse files
Jacksunweicopybara-github
Jacksunwei
authored and
copybara-github
committed
chore: Creates a sample main.py to demonstrate how to use adk as a library for the workflow agent in docs.
For Issue google#573. PiperOrigin-RevId: 760380509
1 parent 2ad1f79 commit 9e767b3

File tree

4 files changed

+248
-50
lines changed

4 files changed

+248
-50
lines changed

‎contributing/samples/workflow_agent_seq/README.md

Lines changed: 12 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,12 @@
1+
# Workflow Agent Sample - SequentialAgent
2+
3+
Sample query:
4+
5+
* Write a quicksort method in python.
6+
* Write a python function to do bubble sort.
7+
8+
To run in cli (after installing `google-adk`):
9+
10+
* `uv run main.py` (or `python main.py`)
11+
12+
Check sample output in `sample.output` file in this folder.

‎contributing/samples/workflow_agent_seq/agent.py

Lines changed: 67 additions & 50 deletions
Original file line numberDiff line numberDiff line change
@@ -15,80 +15,97 @@
1515
from google.adk.agents.llm_agent import LlmAgent
1616
from google.adk.agents.sequential_agent import SequentialAgent
1717

18+
# Part of agent.py --> Follow https://google.github.io/adk-docs/get-started/quickstart/ to learn the setup
19+
1820
# --- 1. Define Sub-Agents for Each Pipeline Stage ---
1921

2022
# Code Writer Agent
2123
# Takes the initial specification (from user query) and writes code.
2224
code_writer_agent = LlmAgent(
23-
name="code_writer_agent",
24-
model="gemini-1.5-flash-001",
25-
instruction="""You are a Code Writer AI.
26-
Based on the user's request, write the initial Python code.
27-
Output *only* the raw code block.
28-
""",
29-
description="Writes initial code based on a specification.",
30-
# Stores its output (the generated code) into the session state
31-
# under the key 'generated_code'.
32-
output_key="generated_code",
25+
name="CodeWriterAgent",
26+
model="gemini-1.5-flash",
27+
# Change 3: Improved instruction
28+
instruction="""You are a Python Code Generator.
29+
Based *only* on the user's request, write Python code that fulfills the requirement.
30+
Output *only* the complete Python code block, enclosed in triple backticks (```python ... ```).
31+
Do not add any other text before or after the code block.
32+
""",
33+
description="Writes initial Python code based on a specification.",
34+
output_key="generated_code", # Stores output in state['generated_code']
3335
)
3436

3537
# Code Reviewer Agent
3638
# Takes the code generated by the previous agent (read from state) and provides feedback.
3739
code_reviewer_agent = LlmAgent(
38-
name="code_reviewer_agent",
39-
model="gemini-2.0-flash-001",
40-
instruction="""You are a Code Reviewer AI.
41-
42-
Review the below Python code.
43-
44-
```
45-
{generated_code}
46-
```
47-
48-
Provide constructive feedback on potential errors, style issues, or improvements.
49-
Focus on clarity and correctness.
50-
Output only the review comments.
51-
52-
""",
40+
name="CodeReviewerAgent",
41+
model="gemini-2.0-flash",
42+
# Change 3: Improved instruction, correctly using state key injection
43+
instruction="""You are an expert Python Code Reviewer.
44+
Your task is to provide constructive feedback on the provided code.
45+
46+
**Code to Review:**
47+
```python
48+
{generated_code}
49+
```
50+
51+
**Review Criteria:**
52+
1. **Correctness:** Does the code work as intended? Are there logic errors?
53+
2. **Readability:** Is the code clear and easy to understand? Follows PEP 8 style guidelines?
54+
3. **Efficiency:** Is the code reasonably efficient? Any obvious performance bottlenecks?
55+
4. **Edge Cases:** Does the code handle potential edge cases or invalid inputs gracefully?
56+
5. **Best Practices:** Does the code follow common Python best practices?
57+
58+
**Output:**
59+
Provide your feedback as a concise, bulleted list. Focus on the most important points for improvement.
60+
If the code is excellent and requires no changes, simply state: "No major issues found."
61+
Output *only* the review comments or the "No major issues" statement.
62+
""",
5363
description="Reviews code and provides feedback.",
54-
# Stores its output (the review comments) into the session state
55-
# under the key 'review_comments'.
56-
output_key="review_comments",
64+
output_key="review_comments", # Stores output in state['review_comments']
5765
)
5866

67+
5968
# Code Refactorer Agent
6069
# Takes the original code and the review comments (read from state) and refactors the code.
6170
code_refactorer_agent = LlmAgent(
62-
name="code_refactorer_agent",
63-
model="gemini-2.0-flash-001",
64-
instruction="""You are a Code Refactorer AI.
65-
66-
Below is the original Python code:
67-
68-
```
69-
{generated_code}
70-
```
71-
72-
Below are the review comments:
73-
74-
{review_comments}
75-
76-
Refactor the code based on the provided feedback.
77-
78-
Output *only* the final, refactored code block.
79-
""",
71+
name="CodeRefactorerAgent",
72+
model="gemini-2.0-flash",
73+
# Change 3: Improved instruction, correctly using state key injection
74+
instruction="""You are a Python Code Refactoring AI.
75+
Your goal is to improve the given Python code based on the provided review comments.
76+
77+
**Original Code:**
78+
```python
79+
{generated_code}
80+
```
81+
82+
**Review Comments:**
83+
{review_comments}
84+
85+
**Task:**
86+
Carefully apply the suggestions from the review comments to refactor the original code.
87+
If the review comments state "No major issues found," return the original code unchanged.
88+
Ensure the final code is complete, functional, and includes necessary imports and docstrings.
89+
90+
**Output:**
91+
Output *only* the final, refactored Python code block, enclosed in triple backticks (```python ... ```).
92+
Do not add any other text before or after the code block.
93+
""",
8094
description="Refactors code based on review comments.",
81-
# Stores its output (the refactored code) into the session state
82-
# under the key 'refactored_code'.
83-
output_key="refactored_code",
95+
output_key="refactored_code", # Stores output in state['refactored_code']
8496
)
8597

98+
8699
# --- 2. Create the SequentialAgent ---
87100
# This agent orchestrates the pipeline by running the sub_agents in order.
88101
code_pipeline_agent = SequentialAgent(
89-
name="code_pipeline_agent",
102+
name="CodePipelineAgent",
90103
sub_agents=[code_writer_agent, code_reviewer_agent, code_refactorer_agent],
104+
description=(
105+
"Executes a sequence of code writing, reviewing, and refactoring."
106+
),
91107
# The agents will run in the order provided: Writer -> Reviewer -> Refactorer
92108
)
93109

110+
# For ADK tools compatibility, the root agent must be named `root_agent`
94111
root_agent = code_pipeline_agent

‎contributing/samples/workflow_agent_seq/main.py

Lines changed: 87 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,87 @@
1+
# Copyright 2025 Google LLC
2+
#
3+
# Licensed under the Apache License, Version 2.0 (the "License");
4+
# you may not use this file except in compliance with the License.
5+
# You may obtain a copy of the License at
6+
#
7+
# http://www.apache.org/licenses/LICENSE-2.0
8+
#
9+
# Unless required by applicable law or agreed to in writing, software
10+
# distributed under the License is distributed on an "AS IS" BASIS,
11+
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12+
# See the License for the specific language governing permissions and
13+
# limitations under the License.
14+
15+
16+
import asyncio
17+
from typing import cast
18+
19+
import agent
20+
from dotenv import load_dotenv
21+
from google.adk.cli.utils import logs
22+
from google.adk.runners import InMemoryRunner
23+
from google.adk.sessions import Session
24+
from google.genai import types
25+
26+
load_dotenv(override=True)
27+
logs.log_to_tmp_folder()
28+
29+
30+
async def main():
31+
app_name = 'my_app'
32+
user_id_1 = 'user1'
33+
runner = InMemoryRunner(
34+
app_name=app_name,
35+
agent=agent.root_agent,
36+
)
37+
38+
async def run_prompt(session: Session, new_message: str) -> Session:
39+
content = types.Content(
40+
role='user', parts=[types.Part.from_text(text=new_message)]
41+
)
42+
print('** User says:', content.model_dump(exclude_none=True))
43+
async for event in runner.run_async(
44+
user_id=user_id_1,
45+
session_id=session.id,
46+
new_message=content,
47+
):
48+
if not event.content or not event.content.parts:
49+
continue
50+
if event.content.parts[0].text:
51+
print(f'** {event.author}: {event.content.parts[0].text}')
52+
elif event.content.parts[0].function_call:
53+
print(
54+
f'** {event.author}: fc /'
55+
f' {event.content.parts[0].function_call.name} /'
56+
f' {event.content.parts[0].function_call.args}\n'
57+
)
58+
elif event.content.parts[0].function_response:
59+
print(
60+
f'** {event.author}: fr /'
61+
f' {event.content.parts[0].function_response.name} /'
62+
f' {event.content.parts[0].function_response.response}\n'
63+
)
64+
65+
return cast(
66+
Session,
67+
await runner.session_service.get_session(
68+
app_name=app_name, user_id=user_id_1, session_id=session.id
69+
),
70+
)
71+
72+
session_1 = await runner.session_service.create_session(
73+
app_name=app_name, user_id=user_id_1
74+
)
75+
76+
print(f'----Session to create memory: {session_1.id} ----------------------')
77+
session_1 = await run_prompt(
78+
session_1, 'Write a python function to do quicksort.'
79+
)
80+
session_1 = await run_prompt(
81+
session_1, 'Write another python function to do bubble sort.'
82+
)
83+
print('-------------------------------------------------------------------')
84+
85+
86+
if __name__ == '__main__':
87+
asyncio.run(main())

‎contributing/samples/workflow_agent_seq/sample.output

Lines changed: 82 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,82 @@
1+
----Session to create memory: a05683c9-4a19-4ba8-b810-e90fc48be566 ----------------------
2+
** User says: {'parts': [{'text': 'Write a python function to do quicksort.'}], 'role': 'user'}
3+
** CodeWriterAgent: ```python
4+
def quicksort(arr):
5+
if len(arr) < 2:
6+
return arr
7+
else:
8+
pivot = arr[0]
9+
less = [i for i in arr[1:] if i <= pivot]
10+
greater = [i for i in arr[1:] if i > pivot]
11+
return quicksort(less) + [pivot] + quicksort(greater)
12+
13+
```
14+
15+
** CodeReviewerAgent: * **Correctness:** The code is functionally correct.
16+
* **Readability:** The code is fairly readable, utilizing list comprehensions effectively.
17+
* **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.
18+
* **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.
19+
20+
** CodeRefactorerAgent: ```python
21+
import random
22+
23+
def quicksort(arr):
24+
"""
25+
Sorts a list using the quicksort algorithm. This implementation
26+
uses a randomly selected pivot to improve average-case performance
27+
and performs the sort in-place to reduce space complexity.
28+
29+
Args:
30+
arr (list): The list to be sorted.
31+
32+
Returns:
33+
list: The sorted list.
34+
"""
35+
def _quicksort(arr, low, high):
36+
if low < high:
37+
pi = partition(arr, low, high)
38+
_quicksort(arr, low, pi-1)
39+
_quicksort(arr, pi+1, high)
40+
41+
def partition(arr, low, high):
42+
# Choose a random pivot
43+
pivot_index = random.randint(low, high)
44+
arr[pivot_index], arr[high] = arr[high], arr[pivot_index]
45+
pivot = arr[high]
46+
47+
i = (low - 1)
48+
for j in range(low, high):
49+
if arr[j] <= pivot:
50+
i += 1
51+
arr[i], arr[j] = arr[j], arr[i]
52+
53+
arr[i + 1], arr[high] = arr[high], arr[i + 1]
54+
return (i + 1)
55+
56+
_quicksort(arr, 0, len(arr)-1)
57+
return arr
58+
```
59+
** User says: {'parts': [{'text': 'Write another python function to do bubble sort.'}], 'role': 'user'}
60+
** CodeWriterAgent: ```python
61+
def bubble_sort(arr):
62+
n = len(arr)
63+
for i in range(n):
64+
for j in range(0, n-i-1):
65+
if arr[j] > arr[j+1]:
66+
arr[j], arr[j+1] = arr[j+1], arr[j]
67+
return arr
68+
69+
```
70+
71+
** CodeReviewerAgent: No major issues found.
72+
73+
** CodeRefactorerAgent: ```python
74+
def bubble_sort(arr):
75+
n = len(arr)
76+
for i in range(n):
77+
for j in range(0, n-i-1):
78+
if arr[j] > arr[j+1]:
79+
arr[j], arr[j+1] = arr[j+1], arr[j]
80+
return arr
81+
```
82+
-------------------------------------------------------------------

0 commit comments

Comments
 (0)
0