1
- from typing import Any , Dict , List
1
+ from typing import Any , Dict , List , Optional
2
2
3
3
from crewai .memory .long_term .long_term_memory_item import LongTermMemoryItem
4
4
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
6
7
7
8
8
9
class LongTermMemory (Memory ):
@@ -14,14 +15,67 @@ class LongTermMemory(Memory):
14
15
LongTermMemoryItem instances.
15
16
"""
16
17
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
+ """
18
44
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
+ )
20
55
super ().__init__ (storage = storage )
21
56
22
57
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
+
25
79
self .storage .save ( # type: ignore # BUG?: Unexpected keyword argument "task_description","score","datetime" for "save" of "Storage"
26
80
task_description = item .task ,
27
81
score = metadata ["quality" ],
@@ -30,7 +84,61 @@ def save(self, item: LongTermMemoryItem) -> None: # type: ignore # BUG?: Signat
30
84
)
31
85
32
86
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"
34
98
35
99
def reset (self ) -> None :
100
+ """Reset the storage by deleting all memory items."""
36
101
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 ()
0 commit comments