Feat/agent usage #819
Feat/agent usage #819
Conversation
Summary of ChangesHello @suluyana, 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 suite of improvements focused on enhancing agent operational transparency and efficiency. It integrates robust token usage tracking for large language models, refines the agent's conversational memory management for more precise state handling, and optimizes resource loading for external tools. These changes collectively aim to provide deeper insights into agent performance and streamline its interaction with various components. Highlights
🧠 New Feature in Public Preview: You can now enable Memory to help Gemini Code Assist learn from your team's feedback. This makes future code reviews more consistent and personalized to your project's style. Click here to enable Memory in your admin console. 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 introduces token usage tracking for agents and includes several refactorings and bug fixes in the memory management component. My review focuses on improving the design of the state management for token counting, simplifying the execution of background tasks, and enhancing code readability in the memory module.
Key feedback points:
- The use of global variables for token counting should be replaced with an encapsulated state to improve modularity and testability.
- The method for running background async tasks is overly complex and can be greatly simplified using
asyncio.create_task. - A method in the memory module can be refactored for better clarity.
| def _add_memory(): | ||
| asyncio.run( | ||
| self.add_memory( | ||
| messages, add_type='add_after_task', **kwargs)) | ||
|
|
||
| loop = asyncio.get_running_loop() | ||
| loop.run_in_executor(None, _add_memory) |
There was a problem hiding this comment.
The current implementation for running add_memory in the background is overly complex and potentially problematic. Using loop.run_in_executor to call a function that itself uses asyncio.run() creates a new event loop within a worker thread. This is an unconventional pattern that can be inefficient and lead to subtle issues.
For running a "fire-and-forget" async task from within an async function, asyncio.create_task() is the standard and much simpler approach. It schedules the coroutine to run on the current event loop without blocking.
| def _add_memory(): | |
| asyncio.run( | |
| self.add_memory( | |
| messages, add_type='add_after_task', **kwargs)) | |
| loop = asyncio.get_running_loop() | |
| loop.run_in_executor(None, _add_memory) | |
| # Schedule add_memory to run in the background without blocking. | |
| asyncio.create_task( | |
| self.add_memory(messages, add_type='add_after_task', **kwargs)) | |
| TOTAL_PROMPT_TOKENS = 0 | ||
| TOTAL_COMPLETION_TOKENS = 0 | ||
| TOKEN_LOCK = asyncio.Lock() |
There was a problem hiding this comment.
The use of module-level global variables TOTAL_PROMPT_TOKENS and TOTAL_COMPLETION_TOKENS for tracking token usage introduces tight coupling between all agent instances within the same process. This can lead to incorrect accounting if multiple independent agents are running, and makes the code harder to test and maintain.
A better approach would be to encapsulate this state within a dedicated usage tracking class or within the LLMAgent instance itself. This would provide better isolation and more flexible usage tracking (e.g., per-agent or per-task).
| def parse_messages(self, messages: List[Message]) -> List[Message]: | ||
| new_messages = [] | ||
| for msg in messages: | ||
| role = getattr(msg, 'role', None) | ||
| content = getattr(msg, 'content', None) | ||
|
|
||
| if 'system' not in self.ignore_roles and role == 'system': | ||
| new_messages.append(msg) | ||
| if role == 'user': | ||
| new_messages.append(msg) | ||
| if 'assistant' not in self.ignore_roles and role == 'assistant' and content is not None: | ||
| new_messages.append(msg) | ||
| if 'tool' not in self.ignore_roles and role == 'tool': | ||
| new_messages.append(msg) | ||
|
|
||
| return new_messages |
There was a problem hiding this comment.
The parse_messages method contains a series of if statements that can be simplified to improve readability and reduce repetition. Consolidating the logic for checking roles against self.ignore_roles would make the function's intent clearer and easier to maintain.
def parse_messages(self, messages: List[Message]) -> List[Message]:
new_messages = []
for msg in messages:
role = getattr(msg, 'role', None)
if role == 'user':
new_messages.append(msg)
continue
if role in self.ignore_roles:
continue
if role == 'assistant' and getattr(msg, 'content', None) is None:
continue
new_messages.append(msg)
return new_messages1 participant
Change Summary
Related issue number
Checklist
pre-commit installandpre-commit run --all-filesbefore git commit, and passed lint check.