feat: implement unified storage system with multiple backends #979
Conversation
…gents - Add unified storage interface (BaseStorage) with async/await support - Implement 8+ storage backends: SQLite, MongoDB, PostgreSQL, Redis, DynamoDB, S3, GCS, Azure - Enhance Memory class with backward-compatible multi-storage support - Add configuration-first approach for easy provider switching - Support primary + cache storage patterns - Maintain full backward compatibility with existing code - Add graceful dependency handling for optional backends Resolves #971 🤖 Generated with [Claude Code](https://claude.ai/code) Co-Authored-By: Claude <[email protected]>
|
Warning Rate limit exceeded@MervinPraison has exceeded the limit for the number of commits or files that can be reviewed per hour. Please wait 0 minutes and 28 seconds before requesting another review. ⌛ How to resolve this issue?After the wait time has elapsed, a review can be triggered using the We recommend that you space out your commits to avoid hitting the rate limit. 🚦 How do rate limits work?CodeRabbit enforces hourly rate limits for each developer per organization. Our paid plans have higher rate limits than the trial, open-source and free plans. In all cases, we re-allow further reviews after a brief timeout. Please see our FAQ for further information. 📒 Files selected for processing (11)
✨ Finishing Touches
Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out. 🪧 TipsChatThere are 3 ways to chat with CodeRabbit:
SupportNeed help? Create a ticket on our support page for assistance with any issues or questions. Note: Be mindful of the bot's finite context window. It's strongly recommended to break down tasks such as reading entire modules into smaller chunks. For a focused discussion, use review comments to chat about specific files and their changes, instead of using the PR comments. CodeRabbit Commands (Invoked using PR comments)
Other keywords and placeholders
CodeRabbit Configuration File (
|
![gemini-code-assist[bot]](https://avatars.githubusercontent.com/in/956858?s=60&v=4){kind=small}
There was a problem hiding this comment.
Summary of Changes
Hello @MervinPraison, 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 upgrades the data persistence capabilities of the praisonai-agents library by implementing a robust, unified storage system. It abstracts away the complexities of different database and cloud storage technologies behind a common asynchronous interface, offering developers a wide range of scalable and performant options for managing agent memory and data. The design ensures seamless integration and future extensibility while preserving compatibility with existing codebases.
Highlights
- Unified Storage Abstraction: I've introduced a new
BaseStorageabstract class that defines a common asynchronous interface for all storage operations (read, write, delete, search, list keys, clear, and batch operations). This provides a consistent API for interacting with various data stores. - Expanded Storage Backend Support: This PR adds comprehensive support for 8 new storage backends: MongoDB, PostgreSQL, Redis, DynamoDB, AWS S3, Google Cloud Storage, Azure Blob Storage, and an enhanced SQLite implementation. Each backend adheres to the new
BaseStorageinterface, allowing for flexible configuration. - Enhanced Memory Management: The core
Memoryclass has been refactored intoEnhancedMemory, which now leverages the new unified storage system. It supports configuring both a primary storage and an optional cache storage, enabling more sophisticated memory strategies. - Backward Compatibility: Crucially, the changes maintain full backward compatibility. The
memory/__init__.pynow intelligently importsEnhancedMemoryif the necessary storage dependencies are available, otherwise falling back to the originalMemoryimplementation. Existing synchronous memory methods are wrapped to utilize the new asynchronous storage backends. - Configuration-First Approach: The new storage system is designed with a configuration-first approach, allowing users to easily select and configure their preferred storage providers with minimal code changes on the client side.
Using Gemini Code Assist
The 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 in your pull request via creating an issue comment (i.e. comment on the pull request page) using either /gemini <command> or @gemini-code-assist <command>. Below is a summary of the supported commands.
| Feature | Command | Description |
|---|---|---|
| Code Review | /gemini review |
Performs a code review for the current pull request in its current state. |
| Pull Request Summary | /gemini summary |
Provides a summary of the current pull request in its current state. |
| Comment | @gemini-code-assist | Responds in comments when explicitly tagged, both in issue comments and review comments. |
| Help | /gemini help |
Displays a list of available commands. |
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 .gemini/ folder in the base of the repository. Detailed instructions can be found here.
Limitations & Feedback
Gemini Code Assist is currently in preview and 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 to provide feedback.
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
-
Review the Privacy Notices, Generative AI Prohibited Use Policy, Terms of Service, and learn how to configure Gemini Code Assist in GitHub here. Gemini can make mistakes, so double check it and use code with caution. ↩
There was a problem hiding this comment.
Code Review
This pull request introduces a unified storage system with multiple backends, enhancing modularity and extensibility. The review focuses on improving robustness, efficiency, and maintainability, including correcting asyncio event loop management, fixing a missing import, addressing performance issues in cloud storage and Redis backends, and improving database operation robustness.
| loop = asyncio.new_event_loop() | ||
| asyncio.set_event_loop(loop) | ||
| try: | ||
| return loop.run_until_complete(self.store(key, data)) | ||
| finally: | ||
| loop.close() |
There was a problem hiding this comment.
Creating a new event loop with asyncio.new_event_loop() for each synchronous call is inefficient and can cause issues. It will fail if an event loop is already running in the thread, which is common in async frameworks or in Jupyter notebooks.
A safer approach for running an async function from a sync context is asyncio.run(). While it also has limitations (it can't be called when a loop is already running), it's an improvement over managing loops manually.
This issue is present in store_short_term, search_short_term, reset_short_term, and reset_all.
| loop = asyncio.new_event_loop() | |
| asyncio.set_event_loop(loop) | |
| try: | |
| return loop.run_until_complete(self.store(key, data)) | |
| finally: | |
| loop.close() | |
| return asyncio.run(self.store(key, data)) |
|
|
||
| try: | ||
| import boto3 | ||
| from boto3.dynamodb.conditions import Key, Attr |
There was a problem hiding this comment.
This line will raise a NameError at runtime because TypeDeserializer is not imported. You need to add from boto3.dynamodb.types import TypeDeserializer at the top of the file to resolve this.
| async def search(self, query: Dict[str, Any]) -> List[Dict[str, Any]]: | ||
| """Search for records matching the query.""" | ||
| await self._ensure_connection() | ||
|
|
||
| loop = asyncio.get_event_loop() | ||
|
|
||
| try: | ||
| # List all objects with prefix | ||
| response = await loop.run_in_executor( | ||
| None, | ||
| self.s3_client.list_objects_v2, | ||
| {"Bucket": self.bucket, "Prefix": self.prefix} | ||
| ) | ||
|
|
||
| results = [] | ||
| limit = query.get("limit", 100) | ||
|
|
||
| for obj in response.get("Contents", []): | ||
| if len(results) >= limit: | ||
| break | ||
|
|
||
| s3_key = obj["Key"] | ||
| key = self._strip_prefix(s3_key) | ||
|
|
||
| # Read and filter object | ||
| record = await self.read(key) | ||
| if record and self._matches_query(record, query): | ||
| results.append(record) | ||
|
|
||
| # Sort by updated_at descending | ||
| results.sort(key=lambda x: x.get("updated_at", 0), reverse=True) | ||
|
|
||
| return results | ||
|
|
||
| except Exception as e: | ||
| self.logger.error(f"Failed to search: {e}") | ||
| return [] |
There was a problem hiding this comment.
The search implementation for S3 is inefficient. It lists all objects under the prefix and then reads each object's full content to perform client-side filtering. For a large number of objects, this will be slow and can incur significant costs due to the high number of GET requests.
Consider using S3 Select for server-side filtering to improve performance. At a minimum, this performance limitation should be clearly documented in the docstring to warn users.
This same issue exists for GCSStorage and AzureStorage as well.
| async def clear(self) -> bool: | ||
| """Clear all records from storage.""" | ||
| await self._ensure_connection() | ||
|
|
||
| loop = asyncio.get_event_loop() | ||
|
|
||
| try: | ||
| # List all objects | ||
| response = await loop.run_in_executor( | ||
| None, | ||
| self.s3_client.list_objects_v2, | ||
| {"Bucket": self.bucket, "Prefix": self.prefix} | ||
| ) | ||
|
|
||
| # Delete all objects | ||
| for obj in response.get("Contents", []): | ||
| await loop.run_in_executor( | ||
| None, | ||
| self.s3_client.delete_object, | ||
| {"Bucket": self.bucket, "Key": obj["Key"]} | ||
| ) | ||
|
|
||
| return True | ||
|
|
||
| except Exception as e: | ||
| self.logger.error(f"Failed to clear storage: {e}") | ||
| return False |
There was a problem hiding this comment.
The clear method deletes objects one by one in a loop, which is inefficient for a large number of objects. Cloud storage providers offer batch deletion APIs that are faster and more cost-effective.
For S3, you should use s3_client.delete_objects(). This same issue applies to GCSStorage (use bucket.delete_blobs()) and AzureStorage (use container_client.delete_blobs()).
async def clear(self) -> bool:
"""Clear all records from storage."""
await self._ensure_connection()
loop = asyncio.get_event_loop()
try:
# Use paginator to handle large number of objects
paginator = self.s3_client.get_paginator('list_objects_v2')
pages = paginator.paginate(Bucket=self.bucket, Prefix=self.prefix)
for page in pages:
if 'Contents' not in page:
continue
objects_to_delete = [{'Key': obj['Key']} for obj in page['Contents']]
if not objects_to_delete:
continue
await loop.run_in_executor(
None,
self.s3_client.delete_objects,
{'Bucket': self.bucket, 'Delete': {'Objects': objects_to_delete}}
)
return True
except Exception as e:
self.logger.error(f"Failed to clear storage: {e}")
return False| async def search(self, query: Dict[str, Any]) -> List[Dict[str, Any]]: | ||
| """ | ||
| Search for records matching the query. | ||
|
|
||
| Note: Redis doesn't have native search capabilities like MongoDB/PostgreSQL. | ||
| This implementation scans all keys and filters client-side, which may be slow | ||
| for large datasets. Consider using RedisSearch module for production use. | ||
| """ | ||
| await self._ensure_connection() | ||
|
|
||
| try: | ||
| # Get all keys with our prefix | ||
| pattern = f"{self.key_prefix}*" | ||
| keys = await self.redis.keys(pattern) | ||
|
|
||
| if not keys: | ||
| return [] | ||
|
|
||
| # Get all records | ||
| raw_data = await self.redis.mget(keys) | ||
| results = [] | ||
|
|
||
| # Process and filter records | ||
| for i, data in enumerate(raw_data): | ||
| if data: | ||
| try: | ||
| json_str = self._decompress_data(data) | ||
| record = json.loads(json_str) | ||
| record["id"] = self._strip_prefix(keys[i].decode()) | ||
|
|
||
| # Apply filters | ||
| if self._matches_query(record, query): | ||
| results.append(record) | ||
|
|
||
| except (json.JSONDecodeError, UnicodeDecodeError) as e: | ||
| self.logger.error(f"Failed to decode record: {e}") | ||
| continue | ||
|
|
||
| # Sort by updated_at descending | ||
| results.sort(key=lambda x: x.get("updated_at", 0), reverse=True) | ||
|
|
||
| # Apply limit | ||
| limit = query.get("limit", 100) | ||
| return results[:limit] | ||
|
|
||
| except RedisError as e: | ||
| self.logger.error(f"Failed to search: {e}") | ||
| return [] |
There was a problem hiding this comment.
The search method uses the KEYS command, which is a blocking operation that can severely degrade Redis performance in production, especially with a large number of keys. It's strongly recommended to use SCAN instead, which iterates through the keyspace without blocking the server.
This same issue applies to the clear and count methods.
async def search(self, query: Dict[str, Any]) -> List[Dict[str, Any]]:
"""
Search for records matching the query.
Note: This implementation scans keys and filters client-side, which may be slow
for large datasets. Consider using RedisSearch module for production use.
"""
await self._ensure_connection()
try:
results = []
limit = query.get("limit", 100)
# Use SCAN instead of KEYS to avoid blocking the server
async for key_bytes in self.redis.scan_iter(match=f"{self.key_prefix}*"):
data = await self.redis.get(key_bytes)
if data:
try:
json_str = self._decompress_data(data)
record = json.loads(json_str)
record["id"] = self._strip_prefix(key_bytes.decode())
if self._matches_query(record, query):
results.append(record)
except (json.JSONDecodeError, UnicodeDecodeError) as e:
self.logger.error(f"Failed to decode record: {e}")
continue
results.sort(key=lambda x: x.get("updated_at", 0), reverse=True)
return results[:limit]
except RedisError as e:
self.logger.error(f"Failed to search: {e}")
return []| result = await conn.execute( | ||
| f"DELETE FROM {self.schema}.{self.table_name} WHERE id = $1", | ||
| key | ||
| ) | ||
| # Extract affected rows from result string like "DELETE 1" | ||
| affected_rows = int(result.split()[-1]) if result.split() else 0 | ||
| return affected_rows > 0 |
There was a problem hiding this comment.
Parsing the status string returned by conn.execute (e.g., "DELETE 1") to determine the number of affected rows is brittle, as the format is not guaranteed to be stable across asyncpg versions.
A more robust approach is to use the RETURNING clause to confirm the deletion.
| result = await conn.execute( | |
| f"DELETE FROM {self.schema}.{self.table_name} WHERE id = $1", | |
| key | |
| ) | |
| # Extract affected rows from result string like "DELETE 1" | |
| affected_rows = int(result.split()[-1]) if result.split() else 0 | |
| return affected_rows > 0 | |
| status = await conn.fetchval( | |
| f"DELETE FROM {self.schema}.{self.table_name} WHERE id = $1 RETURNING id", | |
| key | |
| ) | |
| return status is not None |
![cursor[bot]](https://avatars.githubusercontent.com/in/1210556?s=60&v=4){kind=small}
There was a problem hiding this comment.
Bug: Async Initialization Flaws in SQLiteStorage
The SQLiteStorage constructor's asynchronous initialization is flawed. The check asyncio.get_event_loop().is_running() is problematic: asyncio.get_event_loop() is deprecated and can raise RuntimeError if no loop exists, while is_running() is not a valid method on loop objects, causing an AttributeError. Additionally, the asyncio.create_task() call is neither awaited nor stored, leading to a race condition where the constructor returns before database initialization completes. This can cause subsequent operations to fail or the task to be garbage collected. The initialization should use asyncio.get_running_loop() with proper error handling and ensure the task is managed.
src/praisonai-agents/praisonaiagents/storage/sqlite_storage.py#L44-L46
Bug: Async Calls in Sync Batch Writer
The DynamoDBStorage class incorrectly wraps batch.put_item and batch.delete_item calls with await loop.run_in_executor inside the synchronous boto3.table.batch_writer() context. This misuses the batch writer, which is designed for synchronous operations, leading to potential resource leaks, inconsistent batch behavior, and thread safety concerns.
src/praisonai-agents/praisonaiagents/storage/dynamodb_storage.py#L429-L437
PraisonAI/src/praisonai-agents/praisonaiagents/storage/dynamodb_storage.py
Lines 429 to 437 in ea5c325
src/praisonai-agents/praisonaiagents/storage/dynamodb_storage.py#L469-L474
PraisonAI/src/praisonai-agents/praisonaiagents/storage/dynamodb_storage.py
Lines 469 to 474 in ea5c325
Bug: Async/Sync Bridging Causes Event Loop Conflicts
The EnhancedMemory class contains an incorrect async/sync bridging pattern in methods like store_short_term, search_short_term, reset_short_term, and reset_all. These methods create a new event loop with asyncio.new_event_loop() and set it as the current loop with asyncio.set_event_loop(loop). This can interfere with existing event loops, cause a RuntimeError if called within an already running async context, and does not restore the original event loop, potentially breaking other async code. The pattern should be replaced with asyncio.run() or proper handling of existing event loops.
src/praisonai-agents/praisonaiagents/memory/enhanced_memory.py#L510-L517
PraisonAI/src/praisonai-agents/praisonaiagents/memory/enhanced_memory.py
Lines 510 to 517 in ea5c325
Bug: S3Storage Class Fails Boto3 API Calls
The S3Storage class incorrectly calls boto3 methods via asyncio.run_in_executor. It passes a dictionary (e.g., {'Bucket': self.bucket}) as a single positional argument, but boto3 methods expect keyword arguments (e.g., Bucket=self.bucket). This error affects all boto3 API calls within the class (e.g., head_bucket, get_object, put_object, list_objects_v2), causing runtime errors.
src/praisonai-agents/praisonaiagents/storage/cloud_storage.py#L144-L149
PraisonAI/src/praisonai-agents/praisonaiagents/storage/cloud_storage.py
Lines 144 to 149 in ea5c325
Bug: SQL Injection via JSONB Key Interpolation
The PostgreSQLStorage.search method is vulnerable to SQL injection when searching by metadata. The PostgreSQL ->> operator requires the JSONB key to be a literal string, not a parameterized value. The current implementation attempts to parameterize the key (e.g., metadata ->> $1), which results in invalid SQL. To execute a valid query, the key must be directly interpolated into the SQL string. Since this key originates from user input (query["metadata"]), direct interpolation creates a SQL injection vulnerability.
src/praisonai-agents/praisonaiagents/storage/postgresql_storage.py#L246-L251
PraisonAI/src/praisonai-agents/praisonaiagents/storage/postgresql_storage.py
Lines 246 to 251 in ea5c325
Bug: Circular Import Issue in Memory Factory
A circular import bug exists in the Memory factory function within enhanced_memory.py. When STORAGE_AVAILABLE is False, the function attempts to import Memory from .memory (i.e., memory/__init__.py), which itself imports from enhanced_memory.py, leading to a circular dependency and an ImportError at runtime.
src/praisonai-agents/praisonaiagents/memory/enhanced_memory.py#L641-L647
PraisonAI/src/praisonai-agents/praisonaiagents/memory/enhanced_memory.py
Lines 641 to 647 in ea5c325
BugBot free trial expires on July 22, 2025
Learn more in the Cursor dashboard.
Was this report helpful? Give feedback by reacting with 👍 or 👎
|
@claude review this pull request and do a detailed analysis and fix this if the existing code doesn't have the solution implemented. Making sure it has backward compatibility, no existing features removed. After making those changes, again review the applied changes. Use @web to search if you dont know any information or to find the latest documentation or to find the latest version. Run the code if you think you need to run it to test it. Minimal code change to start with if required any changes. |
|
Claude encountered an error —— View job I'll analyze this and get back to you. |
Summary
✅ Unified Storage System Implementation
✅ Key Features
✅ New API Examples