Skip to content

Commit 2c5a455

Browse files
feat: implement memory agents with session summaries, agentic management, and references
- Add session summary configuration with auto-summarization every N turns - Implement agentic memory management with auto-classification and confidence thresholds - Add memory references with inline/footnote/metadata formatting options - Create MemoryTools class with remember(), update_memory(), forget(), search_memories() - Maintain full backward compatibility with existing memory functionality - Add comprehensive test suites for all new features Resolves #969 🤖 Generated with [Claude Code](https://claude.ai/code) Co-authored-by: Mervin Praison <[email protected]>
1 parent 9ae29b0 commit 2c5a455

File tree

5 files changed

+1023
-1
lines changed

5 files changed

+1023
-1
lines changed

src/praisonai-agents/praisonaiagents/memory/__init__.py

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -11,5 +11,6 @@
1111
"""
1212

1313
from .memory import Memory
14+
from .tools import MemoryTools, get_memory_tools
1415

15-
__all__ = ["Memory"]
16+
__all__ = ["Memory", "MemoryTools", "get_memory_tools"]

src/praisonai-agents/praisonaiagents/memory/memory.py

Lines changed: 334 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -117,6 +117,32 @@ def __init__(self, config: Dict[str, Any], verbose: int = 0):
117117
self.use_rag = (self.provider.lower() == "rag") and CHROMADB_AVAILABLE and self.cfg.get("use_embedding", False)
118118
self.graph_enabled = False # Initialize graph support flag
119119

120+
# Initialize session summary configuration
121+
self.session_summary_config = self.cfg.get("session_summary_config", {})
122+
self.session_enabled = self.session_summary_config.get("enabled", False)
123+
self.update_after_n_turns = self.session_summary_config.get("update_after_n_turns", 5)
124+
self.summary_model = self.session_summary_config.get("model", "gpt-4o-mini")
125+
self.include_in_context = self.session_summary_config.get("include_in_context", True)
126+
127+
# Initialize agentic memory configuration
128+
self.agentic_config = self.cfg.get("agentic_config", {})
129+
self.agentic_enabled = self.agentic_config.get("enabled", False)
130+
self.auto_classify = self.agentic_config.get("auto_classify", True)
131+
self.confidence_threshold = self.agentic_config.get("confidence_threshold", 0.7)
132+
self.management_model = self.agentic_config.get("management_model", "gpt-4o")
133+
134+
# Initialize memory reference configuration
135+
self.reference_config = self.cfg.get("reference_config", {})
136+
self.include_references = self.reference_config.get("include_references", False)
137+
self.reference_format = self.reference_config.get("reference_format", "inline")
138+
self.max_references = self.reference_config.get("max_references", 5)
139+
self.show_confidence = self.reference_config.get("show_confidence", False)
140+
141+
# Session tracking for summaries
142+
self.turn_counter = 0
143+
self.session_history = []
144+
self.current_session_summary = None
145+
120146
# Extract embedding model from config
121147
self.embedder_config = self.cfg.get("embedder", {})
122148
if isinstance(self.embedder_config, dict):
@@ -1144,3 +1170,311 @@ def search_with_quality(
11441170
logger.info(f"After quality filter: {len(filtered)} results")
11451171

11461172
return filtered
1173+
1174+
# -------------------------------------------------------------------------
1175+
# Session Summary Methods
1176+
# -------------------------------------------------------------------------
1177+
def add_to_session(self, role: str, content: str) -> None:
1178+
"""Add a conversation turn to the session history"""
1179+
if not self.session_enabled:
1180+
return
1181+
1182+
self.session_history.append({
1183+
"role": role,
1184+
"content": content,
1185+
"timestamp": time.time()
1186+
})
1187+
self.turn_counter += 1
1188+
1189+
# Check if we need to update the session summary
1190+
if self.turn_counter % self.update_after_n_turns == 0:
1191+
self._update_session_summary()
1192+
1193+
def _update_session_summary(self) -> None:
1194+
"""Update the session summary using the configured model"""
1195+
if not self.session_history:
1196+
return
1197+
1198+
# Create conversation text for summarization
1199+
conversation_text = "\n".join([
1200+
f"{turn['role']}: {turn['content']}"
1201+
for turn in self.session_history[-self.update_after_n_turns:]
1202+
])
1203+
1204+
summary_prompt = f"""
1205+
Summarize the following conversation, focusing on:
1206+
1. Key topics discussed
1207+
2. Important decisions made
1208+
3. Relevant context for future conversations
1209+
4. User preferences and requirements mentioned
1210+
1211+
Conversation:
1212+
{conversation_text}
1213+
1214+
Provide a concise summary in JSON format with keys: "text", "topics", "key_points"
1215+
"""
1216+
1217+
try:
1218+
if LITELLM_AVAILABLE:
1219+
import litellm
1220+
response = litellm.completion(
1221+
model=self.summary_model,
1222+
messages=[{"role": "user", "content": summary_prompt}],
1223+
response_format={"type": "json_object"},
1224+
temperature=0.3
1225+
)
1226+
summary_data = json.loads(response.choices[0].message.content)
1227+
elif OPENAI_AVAILABLE:
1228+
from openai import OpenAI
1229+
client = OpenAI()
1230+
response = client.chat.completions.create(
1231+
model=self.summary_model,
1232+
messages=[{"role": "user", "content": summary_prompt}],
1233+
response_format={"type": "json_object"},
1234+
temperature=0.3
1235+
)
1236+
summary_data = json.loads(response.choices[0].message.content)
1237+
else:
1238+
self._log_verbose("No LLM available for session summary", logging.WARNING)
1239+
return
1240+
1241+
self.current_session_summary = summary_data
1242+
1243+
# Store summary in long-term memory if enabled
1244+
if self.include_in_context:
1245+
self.store_long_term(
1246+
text=summary_data.get("text", ""),
1247+
metadata={
1248+
"type": "session_summary",
1249+
"topics": summary_data.get("topics", []),
1250+
"key_points": summary_data.get("key_points", []),
1251+
"turn_count": self.turn_counter
1252+
}
1253+
)
1254+
1255+
except Exception as e:
1256+
self._log_verbose(f"Error updating session summary: {e}", logging.ERROR)
1257+
1258+
async def aget_session_summary(self) -> Optional[Dict[str, Any]]:
1259+
"""Get the current session summary (async version)"""
1260+
return self.current_session_summary
1261+
1262+
def get_session_summary(self) -> Optional[Dict[str, Any]]:
1263+
"""Get the current session summary"""
1264+
return self.current_session_summary
1265+
1266+
# -------------------------------------------------------------------------
1267+
# Agentic Memory Management Methods
1268+
# -------------------------------------------------------------------------
1269+
def remember(self, fact: str, metadata: Optional[Dict[str, Any]] = None) -> bool:
1270+
"""Store important information with agentic classification"""
1271+
if not self.agentic_enabled:
1272+
# Fallback to regular long-term storage
1273+
self.store_long_term(fact, metadata=metadata)
1274+
return True
1275+
1276+
# Auto-classify the importance if enabled
1277+
if self.auto_classify:
1278+
importance_score = self._classify_importance(fact)
1279+
if importance_score < self.confidence_threshold:
1280+
self._log_verbose(f"Fact importance {importance_score} below threshold {self.confidence_threshold}")
1281+
return False
1282+
1283+
# Store with agentic metadata
1284+
agentic_metadata = metadata or {}
1285+
agentic_metadata.update({
1286+
"stored_by": "agentic_memory",
1287+
"importance_score": importance_score if self.auto_classify else 1.0,
1288+
"auto_classified": self.auto_classify
1289+
})
1290+
1291+
self.store_long_term(fact, metadata=agentic_metadata)
1292+
return True
1293+
1294+
def update_memory(self, memory_id: str, new_fact: str) -> bool:
1295+
"""Update existing memory by ID"""
1296+
try:
1297+
# Update in SQLite
1298+
conn = sqlite3.connect(self.long_db)
1299+
c = conn.cursor()
1300+
c.execute(
1301+
"UPDATE long_mem SET content = ?, meta = ? WHERE id = ?",
1302+
(new_fact, json.dumps({"updated": True, "updated_at": time.time()}), memory_id)
1303+
)
1304+
updated = c.rowcount > 0
1305+
conn.commit()
1306+
conn.close()
1307+
1308+
# Update in vector store if available
1309+
if self.use_rag and hasattr(self, "chroma_col"):
1310+
try:
1311+
# ChromaDB doesn't support direct updates, so we delete and re-add
1312+
self.chroma_col.delete(ids=[memory_id])
1313+
if LITELLM_AVAILABLE:
1314+
import litellm
1315+
response = litellm.embedding(
1316+
model=self.embedding_model,
1317+
input=new_fact
1318+
)
1319+
embedding = response.data[0]["embedding"]
1320+
elif OPENAI_AVAILABLE:
1321+
from openai import OpenAI
1322+
client = OpenAI()
1323+
response = client.embeddings.create(
1324+
input=new_fact,
1325+
model=self.embedding_model
1326+
)
1327+
embedding = response.data[0].embedding
1328+
else:
1329+
return updated
1330+
1331+
self.chroma_col.add(
1332+
documents=[new_fact],
1333+
metadatas=[{"updated": True, "updated_at": time.time()}],
1334+
ids=[memory_id],
1335+
embeddings=[embedding]
1336+
)
1337+
except Exception as e:
1338+
self._log_verbose(f"Error updating in ChromaDB: {e}", logging.ERROR)
1339+
1340+
return updated
1341+
1342+
except Exception as e:
1343+
self._log_verbose(f"Error updating memory: {e}", logging.ERROR)
1344+
return False
1345+
1346+
def forget(self, memory_id: str) -> bool:
1347+
"""Remove a memory by ID"""
1348+
try:
1349+
# Delete from SQLite
1350+
conn = sqlite3.connect(self.long_db)
1351+
c = conn.cursor()
1352+
c.execute("DELETE FROM long_mem WHERE id = ?", (memory_id,))
1353+
deleted = c.rowcount > 0
1354+
conn.commit()
1355+
conn.close()
1356+
1357+
# Delete from vector store if available
1358+
if self.use_rag and hasattr(self, "chroma_col"):
1359+
try:
1360+
self.chroma_col.delete(ids=[memory_id])
1361+
except Exception as e:
1362+
self._log_verbose(f"Error deleting from ChromaDB: {e}", logging.ERROR)
1363+
1364+
return deleted
1365+
1366+
except Exception as e:
1367+
self._log_verbose(f"Error forgetting memory: {e}", logging.ERROR)
1368+
return False
1369+
1370+
def search_memories(self, query: str, limit: int = 5, **kwargs) -> List[Dict[str, Any]]:
1371+
"""Search memories with agentic filtering"""
1372+
# Use existing search method but add agentic filtering
1373+
results = self.search_long_term(query, limit=limit, **kwargs)
1374+
1375+
# Filter by agentic metadata if enabled
1376+
if self.agentic_enabled:
1377+
results = [
1378+
r for r in results
1379+
if r.get("metadata", {}).get("stored_by") == "agentic_memory"
1380+
]
1381+
1382+
return results
1383+
1384+
def _classify_importance(self, fact: str) -> float:
1385+
"""Classify the importance of a fact using LLM"""
1386+
classification_prompt = f"""
1387+
Rate the importance of storing this information in long-term memory on a scale of 0.0 to 1.0:
1388+
- 1.0: Critical information (user preferences, key decisions, important facts)
1389+
- 0.7: Important information (useful context, relevant details)
1390+
- 0.5: Moderate information (might be useful later)
1391+
- 0.3: Low importance (casual conversation, temporary info)
1392+
- 0.0: Not worth storing (greetings, small talk)
1393+
1394+
Information: {fact}
1395+
1396+
Return only a number between 0.0 and 1.0.
1397+
"""
1398+
1399+
try:
1400+
if LITELLM_AVAILABLE:
1401+
import litellm
1402+
response = litellm.completion(
1403+
model=self.management_model,
1404+
messages=[{"role": "user", "content": classification_prompt}],
1405+
temperature=0.1
1406+
)
1407+
score_text = response.choices[0].message.content.strip()
1408+
elif OPENAI_AVAILABLE:
1409+
from openai import OpenAI
1410+
client = OpenAI()
1411+
response = client.chat.completions.create(
1412+
model=self.management_model,
1413+
messages=[{"role": "user", "content": classification_prompt}],
1414+
temperature=0.1
1415+
)
1416+
score_text = response.choices[0].message.content.strip()
1417+
else:
1418+
return 0.5 # Default moderate importance
1419+
1420+
return float(score_text)
1421+
1422+
except Exception as e:
1423+
self._log_verbose(f"Error classifying importance: {e}", logging.ERROR)
1424+
return 0.5 # Default moderate importance
1425+
1426+
# -------------------------------------------------------------------------
1427+
# Memory Reference Methods
1428+
# -------------------------------------------------------------------------
1429+
def search_with_references(self, query: str, limit: int = 5, **kwargs) -> Dict[str, Any]:
1430+
"""Search with memory references included"""
1431+
results = self.search_long_term(query, limit=limit, **kwargs)
1432+
1433+
if not self.include_references or not results:
1434+
return {
1435+
"content": "",
1436+
"references": []
1437+
}
1438+
1439+
# Format results with references
1440+
content_parts = []
1441+
references = []
1442+
1443+
for i, result in enumerate(results[:self.max_references], 1):
1444+
text = result.get("text", "")
1445+
metadata = result.get("metadata", {})
1446+
confidence = result.get("score", 0.0)
1447+
1448+
if self.reference_format == "inline":
1449+
content_parts.append(f"{text} [{i}]")
1450+
elif self.reference_format == "footnote":
1451+
content_parts.append(f"{text}")
1452+
else: # metadata format
1453+
content_parts.append(text)
1454+
1455+
ref_entry = {
1456+
"id": i,
1457+
"text": text,
1458+
"metadata": metadata
1459+
}
1460+
1461+
if self.show_confidence:
1462+
ref_entry["confidence"] = confidence
1463+
1464+
references.append(ref_entry)
1465+
1466+
content = " ".join(content_parts)
1467+
1468+
# Add footnotes if using footnote format
1469+
if self.reference_format == "footnote":
1470+
footnotes = [
1471+
f"[{ref['id']}] {ref['text']}" +
1472+
(f" (confidence: {ref['confidence']:.2f})" if self.show_confidence else "")
1473+
for ref in references
1474+
]
1475+
content += "\n\nReferences:\n" + "\n".join(footnotes)
1476+
1477+
return {
1478+
"content": content,
1479+
"references": references
1480+
}

0 commit comments

Comments
 (0)