Repository navigation
Refactor/agent loop pattern - #3
Conversation
|
🤖 Hi @derailed-dash, I've received your request, and I'm working on it now! You can track my progress in the logs for more details. |
Summary of ChangesHello @derailed-dash, I'm Gemini Code Assist1! I'm currently reviewing this pull request and will post my feedback shortly. In the meantime, here's a summary to help you and other reviewers quickly get up to speed! This pull request introduces a major architectural refactor to the LLMS-Generator's document summarization process. The primary goal is to enhance the system's ability to handle large repositories by implementing a batching and iterative processing pattern. This change allows the agent to intelligently break down large tasks into smaller, manageable chunks, effectively mitigating the constraints of LLM context windows and leading to a more scalable and reliable summarization workflow. The refactor involves a complete overhaul of the Highlights
Using Gemini Code AssistThe full guide for Gemini Code Assist can be found on our documentation page, here are some quick tips. Invoking Gemini You can request assistance from Gemini at any point by creating a comment using either
Customization To customize Gemini Code Assist for GitHub experience, repository maintainers can create a configuration file and/or provide a custom code review style guide (such as PEP-8 for Python) by creating and adding files to a Limitations & Feedback Gemini Code Assist may make mistakes. Please leave feedback on any instances where its feedback is incorrect or counter productive. You can react with 👍 and 👎 on @gemini-code-assist comments. If you're interested in giving your feedback about your experience with Gemini Code Assist for Github and other Google products, sign up here. You can also get AI-powered code generation, chat, as well as code reviews directly in the IDE at no cost with the Gemini Code Assist IDE Extension. Footnotes
|
![gemini-code-assist[bot]](https://avatars.githubusercontent.com/in/956858?s=60&v=4){kind=small}
There was a problem hiding this comment.
Code Review
This pull request refactors the document summarization process to use a loop pattern with batching, which is a great improvement for handling large repositories. The implementation is well-structured with new agents and tools to manage the workflow. I've identified a significant issue with how the README file content is read for the project summary, which could lead to it being skipped. I've also found a small piece of redundant code. My comments provide detailed explanations and suggestions for how to address these points.
| # Agent to create the final project summary after the loop | ||
| project_summariser_agent = Agent( | ||
| name="project_summariser_agent", | ||
| description="Creates the final project summary from all file summaries.", | ||
| model=Gemini( | ||
| model=config.model, | ||
| retry_options=retry_options | ||
| ), | ||
| instruction="""Read the content of the project's README.md file (if available in session state as 'readme_content'). | ||
| Then, review the 'all_summaries' from the session state. | ||
| Generate a two-paragraph summary of the entire project based on these inputs. | ||
| The output should be a JSON object with a single key 'project_summary' containing the generated summary.""", | ||
| tools=[read_files], # To read the README | ||
| output_schema=ProjectSummaryOutput, | ||
| output_key="project_summary_raw" | ||
| ) | ||
|
|
||
| # This agent will process one batch sequentially | ||
| single_batch_processor = SequentialAgent( | ||
| name="single_batch_processor", | ||
| description="Reads and summarizes one batch of files.", | ||
| sub_agents=[ | ||
| file_reader_agent, # Reads files from 'current_batch' | ||
| content_summariser_agent, # Summarises files from 'current_batch' | ||
| update_summaries_agent # Appends batch summaries to 'all_summaries' | ||
| ] | ||
| ) | ||
|
|
||
| # This LoopAgent iteratively processes each batch of files until all are summarized. | ||
| batch_processing_loop = LoopAgent( | ||
| name="batch_processing_loop", | ||
| description="Processes all file batches in a loop.", | ||
| sub_agents=[ | ||
| batch_selector_agent, # Gets next batch or exits | ||
| single_batch_processor | ||
| ], | ||
| max_iterations=200 # A safeguard against infinite loops | ||
| ) | ||
|
|
||
| # This agent combines all collected summaries and the project summary into the final output. | ||
| final_summary_agent = Agent( | ||
| name="final_summary_agent", | ||
| description="Finalizes the document summaries by combining all individual and project summaries.", | ||
| model=Gemini( | ||
| model=config.model, | ||
| retry_options=retry_options | ||
| ), | ||
| instruction="""Call the `finalize_summaries` tool to combine all collected summaries and the project summary into the final output format.""", | ||
| tools=[finalize_summaries] | ||
| ) | ||
|
|
||
| # This is the main document summarizer agent, orchestrating the entire process. | ||
| document_summariser_agent = SequentialAgent( | ||
| name="document_summariser_agent", | ||
| description="A sequential agent that first reads file contents and then summarizes them.", | ||
| description="Orchestrates the entire file summarization process including batching and looping.", | ||
| sub_agents=[ | ||
| file_reader_agent, | ||
| content_summariser_agent | ||
| batch_creation_agent, # Step 1: Create batches of files | ||
| batch_processing_loop, # Step 2: Process each batch in a loop | ||
| project_summariser_agent, # Step 3: Generate overall project summary | ||
| final_summary_agent # Step 4: Finalize and combine all summaries | ||
| ] | ||
| ) |
There was a problem hiding this comment.
The current implementation for generating the project summary is flawed. The project_summariser_agent expects readme_content in the session state, but no prior step sets it. Using the read_files tool within this agent is inefficient as it would re-process all files.
To create a more robust and efficient workflow, I suggest introducing a dedicated agent to read the README file and then simplifying the project_summariser_agent. This makes the data flow explicit and avoids redundant file operations.
First, you would add a new tool to src/llms_gen_agent/sub_agents/doc_summariser/tools.py:
import os
def read_readme_to_state(tool_context: ToolContext) -> dict:
"""Finds and reads the root README.md file into session state."""
logger.debug("Executing read_readme_to_state")
repo_path = tool_context.state.get("repo_path")
all_files = tool_context.state.get("files", [])
if not repo_path or not all_files:
logger.warning("read_readme_to_state: missing repo_path or files in state.")
tool_context.state["readme_content"] = ""
return {"status": "failure"}
readme_path = None
for file in all_files:
if os.path.basename(file).lower() == 'readme.md' and os.path.dirname(file) == repo_path:
readme_path = file
break
if readme_path:
try:
with open(readme_path, 'r', encoding='utf-8') as f:
tool_context.state['readme_content'] = f.read()
return {"status": "success"}
except Exception as e:
logger.warning(f"Could not read README file {readme_path}: {e}")
tool_context.state['readme_content'] = ''
return {"status": "not_found"}Then, you can update the agent definitions as suggested below, assuming you also import the new read_readme_to_state tool.
# Agent to read the README file into session state
readme_reader_agent = Agent(
name="readme_reader_agent",
description="Reads the project's root README.md file into session state.",
model=Gemini(
model=config.model,
retry_options=retry_options
),
instruction="""You MUST call the `read_readme_to_state` tool. This is your ONLY task.""",
tools=[read_readme_to_state] # Assumes this tool is created and imported
)
# Agent to create the final project summary after the loop
project_summariser_agent = Agent(
name="project_summariser_agent",
description="Creates the final project summary from all file summaries and the README.",
model=Gemini(
model=config.model,
retry_options=retry_options
),
instruction="""Read the 'readme_content' and 'all_summaries' from the session state.
Generate a two-paragraph summary of the entire project based on these inputs.
The output should be a JSON object with a single key 'project_summary' containing the generated summary.""",
output_schema=ProjectSummaryOutput,
output_key="project_summary_raw"
)
# This agent will process one batch sequentially
single_batch_processor = SequentialAgent(
name="single_batch_processor",
description="Reads and summarizes one batch of files.",
sub_agents=[
file_reader_agent, # Reads files from 'current_batch'
content_summariser_agent, # Summarises files from 'current_batch'
update_summaries_agent # Appends batch summaries to 'all_summaries'
]
)
# This LoopAgent iteratively processes each batch of files until all are summarized.
batch_processing_loop = LoopAgent(
name="batch_processing_loop",
description="Processes all file batches in a loop.",
sub_agents=[
batch_selector_agent, # Gets next batch or exits
single_batch_processor
],
max_iterations=200 # A safeguard against infinite loops
)
# This agent combines all collected summaries and the project summary into the final output.
final_summary_agent = Agent(
name="final_summary_agent",
description="Finalizes the document summaries by combining all individual and project summaries.",
model=Gemini(
model=config.model,
retry_options=retry_options
),
instruction="""Call the `finalize_summaries` tool to combine all collected summaries and the project summary into the final output format.""",
tools=[finalize_summaries]
)
# This is the main document summarizer agent, orchestrating the entire process.
document_summariser_agent = SequentialAgent(
name="document_summariser_agent",
description="Orchestrates the entire file summarization process including batching and looping.",
sub_agents=[
batch_creation_agent, # Step 1: Create batches of files
batch_processing_loop, # Step 2: Process each batch in a loop
readme_reader_agent, # Step 3: Read the README file
project_summariser_agent, # Step 4: Generate overall project summary
final_summary_agent # Step 5: Finalize and combine all summaries
]
)
No description provided.