Skip to content

Commit fb099b8

Browse files
committed
Add PostgreSQL support for long-term memory storage
This PR adds comprehensive PostgreSQL integration for CrewAI's long-term memory system: - Create LTMPostgresStorage class supporting PostgreSQL 16+ database backend - Implement connection pooling with psycopg for improved performance - Add storage factory pattern via LTMStorageFactory for backend selection - Support environment variable configuration for simplified deployment - Add context manager pattern for improved resource management - Implement robust validation and security measures: - Input validation for all database parameters - Protection against SQL injection - Secure connection string handling - Comprehensive error types for better debugging - Add CI-compatible tests with mocks - Create detailed documentation with examples and best practices - Add optional dependency for psycopg[pool] in pyproject.toml
1 parent 2d6deee commit fb099b8

File tree

10 files changed

+5431
-2602
lines changed

10 files changed

+5431
-2602
lines changed

docs/how-to/postgres-long-term-memory.mdx

Lines changed: 730 additions & 0 deletions
Large diffs are not rendered by default.

pyproject.toml

Lines changed: 5 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -67,6 +67,9 @@ docling = [
6767
aisuite = [
6868
"aisuite>=0.1.10",
6969
]
70+
postgres = [
71+
"psycopg[pool]>=3.1.12",
72+
]
7073

7174
[tool.uv]
7275
dev-dependencies = [
@@ -87,6 +90,7 @@ dev-dependencies = [
8790
"pytest-recording>=0.13.2",
8891
"pytest-randomly>=3.16.0",
8992
"pytest-timeout>=2.3.1",
93+
"psycopg[pool]>=3.1.12",
9094
]
9195

9296
[project.scripts]
@@ -102,4 +106,4 @@ exclude_dirs = ["src/crewai/cli/templates"]
102106

103107
[build-system]
104108
requires = ["hatchling"]
105-
build-backend = "hatchling.build"
109+
build-backend = "hatchling.build"
Lines changed: 115 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -1,8 +1,9 @@
1-
from typing import Any, Dict, List
1+
from typing import Any, Dict, List, Optional
22

33
from crewai.memory.long_term.long_term_memory_item import LongTermMemoryItem
44
from crewai.memory.memory import Memory
5-
from crewai.memory.storage.ltm_sqlite_storage import LTMSQLiteStorage
5+
# Storage factory is used to create appropriate storage backend
6+
from crewai.memory.storage.ltm_storage_factory import LTMStorageFactory
67

78

89
class LongTermMemory(Memory):
@@ -14,14 +15,67 @@ class LongTermMemory(Memory):
1415
LongTermMemoryItem instances.
1516
"""
1617

17-
def __init__(self, storage=None, path=None):
18+
def __init__(
19+
self,
20+
storage=None,
21+
storage_type: str = "sqlite",
22+
path: Optional[str] = None,
23+
postgres_connection_string: Optional[str] = None,
24+
postgres_schema: Optional[str] = None,
25+
postgres_table_name: Optional[str] = None,
26+
postgres_min_pool_size: Optional[int] = None,
27+
postgres_max_pool_size: Optional[int] = None,
28+
postgres_use_connection_pool: Optional[bool] = None,
29+
):
30+
"""
31+
Initialize LongTermMemory with the specified storage backend.
32+
33+
Args:
34+
storage: Optional pre-configured storage instance
35+
storage_type: Type of storage to use ('sqlite' or 'postgres') when creating a new storage
36+
path: Path to SQLite database file (only used with SQLite storage)
37+
postgres_connection_string: Postgres connection string (only used with Postgres storage)
38+
postgres_schema: Postgres schema name (only used with Postgres storage)
39+
postgres_table_name: Postgres table name (only used with Postgres storage)
40+
postgres_min_pool_size: Minimum connection pool size (only used with Postgres storage)
41+
postgres_max_pool_size: Maximum connection pool size (only used with Postgres storage)
42+
postgres_use_connection_pool: Whether to use connection pooling (only used with Postgres storage)
43+
"""
1844
if not storage:
19-
storage = LTMSQLiteStorage(db_path=path) if path else LTMSQLiteStorage()
45+
storage = LTMStorageFactory.create_storage(
46+
storage_type=storage_type,
47+
path=path,
48+
connection_string=postgres_connection_string,
49+
schema=postgres_schema,
50+
table_name=postgres_table_name,
51+
min_pool_size=postgres_min_pool_size,
52+
max_pool_size=postgres_max_pool_size,
53+
use_connection_pool=postgres_use_connection_pool,
54+
)
2055
super().__init__(storage=storage)
2156

2257
def save(self, item: LongTermMemoryItem) -> None: # type: ignore # BUG?: Signature of "save" incompatible with supertype "Memory"
23-
metadata = item.metadata
24-
metadata.update({"agent": item.agent, "expected_output": item.expected_output})
58+
"""
59+
Save a memory item to storage.
60+
61+
Args:
62+
item: The LongTermMemoryItem to save
63+
"""
64+
# Create metadata dictionary with required values
65+
metadata = item.metadata.copy() if item.metadata else {}
66+
metadata.update({
67+
"agent": item.agent,
68+
"expected_output": item.expected_output
69+
})
70+
71+
# Ensure quality is in metadata (from item.quality if available)
72+
if "quality" not in metadata and item.quality is not None:
73+
metadata["quality"] = item.quality
74+
75+
# Check if quality is available
76+
if "quality" not in metadata:
77+
raise ValueError("Memory quality must be provided either in item.quality or item.metadata['quality']")
78+
2579
self.storage.save( # type: ignore # BUG?: Unexpected keyword argument "task_description","score","datetime" for "save" of "Storage"
2680
task_description=item.task,
2781
score=metadata["quality"],
@@ -30,7 +84,61 @@ def save(self, item: LongTermMemoryItem) -> None: # type: ignore # BUG?: Signat
3084
)
3185

3286
def search(self, task: str, latest_n: int = 3) -> List[Dict[str, Any]]: # type: ignore # signature of "search" incompatible with supertype "Memory"
33-
return self.storage.load(task, latest_n) # type: ignore # BUG?: "Storage" has no attribute "load"
87+
"""
88+
Search for memory items by task.
89+
90+
Args:
91+
task: The task description to search for
92+
latest_n: Maximum number of results to return
93+
94+
Returns:
95+
List of memory items matching the search criteria
96+
"""
97+
return self.storage.load(task, latest_n) or [] # type: ignore # BUG?: "Storage" has no attribute "load"
3498

3599
def reset(self) -> None:
100+
"""Reset the storage by deleting all memory items."""
36101
self.storage.reset()
102+
103+
def cleanup(self) -> None:
104+
"""
105+
Clean up resources and connections.
106+
107+
This method safely handles any exceptions that might occur during cleanup,
108+
ensuring resources are properly released.
109+
"""
110+
if hasattr(self.storage, 'close'):
111+
try:
112+
self.storage.close()
113+
except Exception as e:
114+
# Log the error but don't raise it to ensure cleanup continues
115+
print(f"WARNING: Error while closing memory storage: {e}")
116+
117+
# Keep close() for backward compatibility
118+
def close(self) -> None:
119+
"""
120+
Close any resources held by the storage.
121+
122+
For PostgreSQL storage with connection pooling enabled, this will
123+
close the connection pool. For other storage types, this is a no-op.
124+
125+
This method safely handles any exceptions that might occur during closing.
126+
127+
Note: This method is an alias for cleanup() and is maintained for backward compatibility.
128+
"""
129+
self.cleanup()
130+
131+
def __enter__(self):
132+
"""Support for using LongTermMemory as a context manager."""
133+
return self
134+
135+
def __exit__(self, exc_type, exc_val, exc_tb):
136+
"""
137+
Clean up resources when exiting a context manager block.
138+
139+
Args:
140+
exc_type: Exception type if an exception was raised in the context
141+
exc_val: Exception value if an exception was raised in the context
142+
exc_tb: Exception traceback if an exception was raised in the context
143+
"""
144+
self.cleanup()

src/crewai/memory/storage/__init__.py

Lines changed: 6 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1 +1,7 @@
11
"""Memory storage implementations for crewAI."""
2+
3+
from crewai.memory.storage.ltm_postgres_storage import LTMPostgresStorage
4+
from crewai.memory.storage.ltm_sqlite_storage import LTMSQLiteStorage
5+
from crewai.memory.storage.ltm_storage_factory import LTMStorageFactory
6+
7+
__all__ = ["LTMSQLiteStorage", "LTMPostgresStorage", "LTMStorageFactory"]

0 commit comments

Comments
 (0)