Conversation
Summary of ChangesHello, 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 significantly enhances the memory management system by introducing optimizations for when and how conversational memory is added. It implements an asynchronous early memory addition mechanism for specific message types to improve efficiency and responsiveness. Furthermore, it refines the logic for identifying and managing shared memory instances to prevent conflicts across different LLM configurations and includes a crucial fix for the underlying Milvus vector store integration to ensure robust and accurate data updates. 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. 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
The pull request introduces changes to memory management, including early memory addition and a fix for the mem0 library's update mechanism. The refactoring of memory handling in llm_agent.py and memory_manager.py appears to improve how memory instances are managed and identified. However, a critical issue was identified in the asynchronous scheduling of memory additions, which could lead to runtime errors.
| def _schedule_add_memory_after_task(self, messages, timestamp=None): | ||
|
|
||
| def _add_memory(): | ||
| asyncio.run( | ||
| self.add_memory( | ||
| messages, add_type='add_after_task', timestamp=timestamp)) | ||
|
|
||
| loop = asyncio.get_running_loop() | ||
| loop.run_in_executor(None, _add_memory) |
There was a problem hiding this comment.
The _schedule_add_memory_after_task function is called from an async context (run_loop). Using asyncio.run() inside a function that is already part of a running event loop (which asyncio.get_running_loop() implies) is an anti-pattern and will raise a RuntimeError: Cannot run asyncio.run() from a running event loop.
To run self.add_memory concurrently without blocking the main event loop, _schedule_add_memory_after_task should be an async function, and asyncio.create_task should be used to schedule the add_memory coroutine. This will allow add_memory to run in the background on the existing event loop.
async def _schedule_add_memory_after_task(self, messages, timestamp=None):
asyncio.create_task(
self.add_memory(
messages, add_type='add_after_task', timestamp=timestamp))| self._schedule_add_memory_after_task( | ||
| messages, timestamp='early') |
There was a problem hiding this comment.
Following the fix for _schedule_add_memory_after_task to be an async function, this call site must be updated to await the function.
await self._schedule_add_memory_after_task(
messages, timestamp='early')| asyncio.run( | ||
| self.add_memory( | ||
| messages, add_type='add_after_task', **kwargs)) | ||
| self._schedule_add_memory_after_task(messages) |
There was a problem hiding this comment.
Following the fix for _schedule_add_memory_after_task to be an async function, this call site must be updated to await the function.
await self._schedule_add_memory_after_task(messages)| import mem0.vector_stores.milvus | ||
| capture_event_origin = mem0.memory.main.capture_event | ||
| update_origin = mem0.vector_stores.milvus.MilvusDB.update | ||
|
|
||
| @wraps(update_origin) | ||
| def update(self, vector_id=None, vector=None, payload=None): | ||
| """ | ||
| Update a vector and its payload. | ||
|
|
||
| Args: | ||
| vector_id (str): ID of the vector to update. | ||
| vector (List[float], optional): Updated vector. | ||
| payload (Dict, optional): Updated payload. | ||
| """ | ||
| if vector is None: | ||
| res = self.client.get( | ||
| collection_name=self.collection_name, ids=[vector_id]) | ||
| if res: | ||
| vector = res[0]['vectors'] | ||
|
|
||
| schema = {'id': vector_id, 'vectors': vector, 'metadata': payload} | ||
| self.client.upsert( | ||
| collection_name=self.collection_name, data=schema) |
There was a problem hiding this comment.
Monkey-patching a third-party library's method (mem0.vector_stores.milvus.MilvusDB.update) can be a fragile practice. While it might be necessary for an immediate fix or specific functionality not available upstream, it introduces a dependency on the internal implementation details of mem0.
Consider documenting the specific reason for this monkey-patching (e.g., a bug in mem0 or a missing feature) and explore more robust integration methods in the long term, such as contributing the change upstream to the mem0 library or using a custom wrapper if mem0 provides extension points. This will improve maintainability and reduce the risk of breakage with future mem0 updates.
2 participants
Change Summary
Related issue number
Checklist
pre-commit installandpre-commit run --all-filesbefore git commit, and passed lint check.