docs: refactor sequential agent example - clean imports and test code

examples/python/snippets/agents/workflow-agents/sequential_agent_code_development_agent.py

@@ -12,24 +12,12 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-import asyncio
-import os
+# --8<-- [start:init]
 from google.adk.agents.sequential_agent import SequentialAgent
 from google.adk.agents.llm_agent import LlmAgent
-from google.genai import types
-# Change 1: Import InMemoryRunner instead of Runner/InMemorySessionService
-from google.adk.runners import InMemoryRunner
-from typing import Optional # For type hints
 
 # --- Constants ---
-APP_NAME = "code_pipeline_app"
-USER_ID = "dev_user_01"
-SESSION_ID = "pipeline_session_02" # Use a unique session ID for each run if desired
-GEMINI_MODEL = "gemini-2.0-flash"
-
-
-# --8<-- [start:init]
-# Part of agent.py --> Follow https://google.github.io/adk-docs/get-started/quickstart/ to learn the setup
+GEMINI_MODEL = "gemini-2.5-flash"
 
 # --- 1. Define Sub-Agents for Each Pipeline Stage ---
 
@@ -38,44 +26,44 @@
 code_writer_agent = LlmAgent(
     name="CodeWriterAgent",
     model=GEMINI_MODEL,
-    # 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.
-""",
+    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']
+    output_key="generated_code"
 )
 
 # Code Reviewer Agent
 # Takes the code generated by the previous agent (read from state) and provides feedback.
 code_reviewer_agent = LlmAgent(
     name="CodeReviewerAgent",
     model=GEMINI_MODEL,
-    # Change 3: Improved instruction, correctly using state key injection
-    instruction="""You are an expert Python Code Reviewer. 
+    instruction="""
+    You are an expert Python Code Reviewer.
     Your task is to provide constructive feedback on the provided code.
 
     **Code to Review:**
     ```python
     {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?
-
-**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.
-""",
+    **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?
+
+    **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.",
-    output_key="review_comments", # Stores output in state['review_comments']
+    output_key="review_comments"
 )
 
 
@@ -84,29 +72,29 @@
 code_refactorer_agent = LlmAgent(
     name="CodeRefactorerAgent",
     model=GEMINI_MODEL,
-    # 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.
-
-  **Original Code:**
-  ```python
-  {generated_code}
-  ```
-
-  **Review Comments:**
-  {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.
-
-**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.
-""",
+    instruction="""
+    You are a Python Code Refactoring AI.
+    Your goal is to improve the given Python code based on the provided review comments.
+
+    **Original Code:**
+    ```python
+    {generated_code}
+    ```
+
+    **Review Comments:**
+    {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.
+
+    **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.",
-    output_key="refactored_code", # Stores output in state['refactored_code']
+    output_key="refactored_code"
 )
 
 
@@ -116,85 +104,7 @@
     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
 # --8<-- [end:init]
-
-# --- 3. Running the Agent (Using InMemoryRunner for local testing) ---
-
-# # Use InMemoryRunner
-# runner = InMemoryRunner(agent=root_agent, app_name=APP_NAME)
-# print(f"InMemoryRunner created for agent '{root_agent.name}'.")
-
-
-# # --- Interaction Function ---
-# async def call_code_pipeline(query: str, user_id: str, session_id: str) -> Optional[str]:
-#     """Runs the code pipeline for a given query and prints intermediate/final steps."""
-#     print(f"\n{'='*15} Running Code Pipeline for Query: '{query}' {'='*15}")
-#     print(f"Attempting to use Session ID: {session_id}")
-#     content = types.Content(role='user', parts=[types.Part(text=query)])
-#     final_code = None
-#     pipeline_step_outputs = {}
-
-#     # --- Explicit Session Creation/Check BEFORE run_async ---
-#     session_service = runner.session_service
-#     session = session_service.create_session(app_name=APP_NAME, user_id=user_id, session_id=session_id)
-
-#     # --- Run the Agent ---
-#     async for event in runner.run_async(
-#         user_id=user_id, session_id=session_id, new_message=content
-#     ):
-#         author_name = event.author or "System"
-#         is_final = event.is_final_response()
-
-#         if is_final and event.content and event.content.parts:
-#             output_text = event.content.parts[0].text.strip()
-#             pipeline_step_outputs[author_name] = output_text
-
-#             print(f"\n--- Output from: {author_name} ---")
-#             print(output_text)
-#             print("-" * (len(author_name) + 18))
-
-#             if author_name == code_refactorer_agent.name:
-#                 final_code = output_text
-
-#         elif event.error_message:
-#               # Log error but allow loop to continue if possible
-#               print(f"  -> Error from {author_name}: {event.error_message}")
-
-#     if final_code is None:
-#         print("\nPipeline did not produce final refactored code.")
-
-#     # --- Final State Retrieval ---
-#     print("\n--- Attempting to retrieve Final Session State ---")
-#     final_session_object = None
-#     try:
-#         # Use the correct method via the internal session_service
-#         final_session_object = runner.session_service.get_session(
-#             app_name=APP_NAME, user_id=user_id, session_id=session_id
-#         )
-#     except Exception as e:
-#          print(f"   -> Error retrieving final session object: {e}")
-
-#     if final_session_object:
-#         print(final_session_object.state) # Access the .state attribute
-#     else:
-#         print("State not found (Final session object could not be retrieved).")
-#     print("=" * 50)
-
-#     return final_code
-
-
-# query = "Write a Python function to calculate the factorial of a number using recursion."
-# # query = "Create a simple Python class for a 'Book' with attributes title and author."
-# # query = "Generate a Python function that takes a list of numbers and returns the sum of squares."
-# # query = "Write a Python script to read the first line of a text file named 'input.txt'."
-
-# # In Colab/Jupyter:
-# await call_code_pipeline(query, user_id=USER_ID, session_id=SESSION_ID)
-
-# # In a standalone Python script or if await is not supported/failing:
-# # asyncio.run(call_code_pipeline(query, user_id=USER_ID, session_id=SESSION_ID))