praisonaiagents 0.0.28__py3-none-any.whl → 0.0.30__py3-none-any.whl

Sign up to get free protection for your applications and to get access to all the features.
@@ -0,0 +1,931 @@
1
+ import os
2
+ import sqlite3
3
+ import json
4
+ import time
5
+ import shutil
6
+ from typing import Any, Dict, List, Optional, Union, Literal
7
+ import logging
8
+
9
+ # Set up logger
10
+ logger = logging.getLogger(__name__)
11
+
12
+ try:
13
+ import chromadb
14
+ from chromadb.config import Settings as ChromaSettings
15
+ CHROMADB_AVAILABLE = True
16
+ except ImportError:
17
+ CHROMADB_AVAILABLE = False
18
+ logger.warning("To use memory features, please run: pip install \"praisonaiagents[memory]\"")
19
+
20
+ try:
21
+ import mem0
22
+ MEM0_AVAILABLE = True
23
+ except ImportError:
24
+ MEM0_AVAILABLE = False
25
+
26
+ try:
27
+ import openai
28
+ OPENAI_AVAILABLE = True
29
+ except ImportError:
30
+ OPENAI_AVAILABLE = False
31
+
32
+
33
+
34
+
35
+ class Memory:
36
+ """
37
+ A single-file memory manager covering:
38
+ - Short-term memory (STM) for ephemeral context
39
+ - Long-term memory (LTM) for persistent knowledge
40
+ - Entity memory (structured data about named entities)
41
+ - User memory (preferences/history for each user)
42
+ - Quality score logic for deciding which data to store in LTM
43
+ - Context building from multiple memory sources
44
+
45
+ Config example:
46
+ {
47
+ "provider": "rag" or "mem0" or "none",
48
+ "use_embedding": True,
49
+ "short_db": "short_term.db",
50
+ "long_db": "long_term.db",
51
+ "rag_db_path": "rag_db", # optional path for local embedding store
52
+ "config": {
53
+ "api_key": "...", # if mem0 usage
54
+ "org_id": "...",
55
+ "project_id": "...",
56
+ ...
57
+ }
58
+ }
59
+ """
60
+
61
+ def __init__(self, config: Dict[str, Any], verbose: int = 0):
62
+ self.cfg = config or {}
63
+ self.verbose = verbose
64
+
65
+ # Set logger level based on verbose
66
+ if verbose >= 5:
67
+ logger.setLevel(logging.INFO)
68
+ else:
69
+ logger.setLevel(logging.WARNING)
70
+
71
+ # Also set ChromaDB and OpenAI client loggers to WARNING
72
+ logging.getLogger('chromadb').setLevel(logging.WARNING)
73
+ logging.getLogger('openai').setLevel(logging.WARNING)
74
+ logging.getLogger('httpx').setLevel(logging.WARNING)
75
+ logging.getLogger('httpcore').setLevel(logging.WARNING)
76
+ logging.getLogger('chromadb.segment.impl.vector.local_persistent_hnsw').setLevel(logging.ERROR)
77
+
78
+ self.provider = self.cfg.get("provider", "rag")
79
+ self.use_mem0 = (self.provider.lower() == "mem0") and MEM0_AVAILABLE
80
+ self.use_rag = (self.provider.lower() == "rag") and CHROMADB_AVAILABLE and self.cfg.get("use_embedding", False)
81
+
82
+ # Create .praison directory if it doesn't exist
83
+ os.makedirs(".praison", exist_ok=True)
84
+
85
+ # Short-term DB
86
+ self.short_db = self.cfg.get("short_db", ".praison/short_term.db")
87
+ self._init_stm()
88
+
89
+ # Long-term DB
90
+ self.long_db = self.cfg.get("long_db", ".praison/long_term.db")
91
+ self._init_ltm()
92
+
93
+ # Conditionally init Mem0 or local RAG
94
+ if self.use_mem0:
95
+ self._init_mem0()
96
+ elif self.use_rag:
97
+ self._init_chroma()
98
+
99
+ def _log_verbose(self, msg: str, level: int = logging.INFO):
100
+ """Only log if verbose >= 5"""
101
+ if self.verbose >= 5:
102
+ logger.log(level, msg)
103
+
104
+ # -------------------------------------------------------------------------
105
+ # Initialization
106
+ # -------------------------------------------------------------------------
107
+ def _init_stm(self):
108
+ """Creates or verifies short-term memory table."""
109
+ os.makedirs(os.path.dirname(self.short_db) or ".", exist_ok=True)
110
+ conn = sqlite3.connect(self.short_db)
111
+ c = conn.cursor()
112
+ c.execute("""
113
+ CREATE TABLE IF NOT EXISTS short_mem (
114
+ id TEXT PRIMARY KEY,
115
+ content TEXT,
116
+ meta TEXT,
117
+ created_at REAL
118
+ )
119
+ """)
120
+ conn.commit()
121
+ conn.close()
122
+
123
+ def _init_ltm(self):
124
+ """Creates or verifies long-term memory table."""
125
+ os.makedirs(os.path.dirname(self.long_db) or ".", exist_ok=True)
126
+ conn = sqlite3.connect(self.long_db)
127
+ c = conn.cursor()
128
+ c.execute("""
129
+ CREATE TABLE IF NOT EXISTS long_mem (
130
+ id TEXT PRIMARY KEY,
131
+ content TEXT,
132
+ meta TEXT,
133
+ created_at REAL
134
+ )
135
+ """)
136
+ conn.commit()
137
+ conn.close()
138
+
139
+ def _init_mem0(self):
140
+ """Initialize Mem0 client for agent or user memory."""
141
+ from mem0 import MemoryClient
142
+ mem_cfg = self.cfg.get("config", {})
143
+ api_key = mem_cfg.get("api_key", os.getenv("MEM0_API_KEY"))
144
+ org_id = mem_cfg.get("org_id")
145
+ proj_id = mem_cfg.get("project_id")
146
+ if org_id and proj_id:
147
+ self.mem0_client = MemoryClient(api_key=api_key, org_id=org_id, project_id=proj_id)
148
+ else:
149
+ self.mem0_client = MemoryClient(api_key=api_key)
150
+
151
+ def _init_chroma(self):
152
+ """Initialize a local Chroma client for embedding-based search."""
153
+ try:
154
+ # Create directory if it doesn't exist
155
+ rag_path = self.cfg.get("rag_db_path", "chroma_db")
156
+ os.makedirs(rag_path, exist_ok=True)
157
+
158
+ # Initialize ChromaDB with persistent storage
159
+ self.chroma_client = chromadb.PersistentClient(
160
+ path=rag_path,
161
+ settings=ChromaSettings(
162
+ anonymized_telemetry=False,
163
+ allow_reset=True
164
+ )
165
+ )
166
+
167
+ collection_name = "memory_store"
168
+ try:
169
+ self.chroma_col = self.chroma_client.get_collection(name=collection_name)
170
+ self._log_verbose("Using existing ChromaDB collection")
171
+ except Exception as e:
172
+ self._log_verbose(f"Collection '{collection_name}' not found. Creating new collection. Error: {e}")
173
+ self.chroma_col = self.chroma_client.create_collection(
174
+ name=collection_name,
175
+ metadata={"hnsw:space": "cosine"}
176
+ )
177
+ self._log_verbose("Created new ChromaDB collection")
178
+
179
+ except Exception as e:
180
+ self._log_verbose(f"Failed to initialize ChromaDB: {e}", logging.ERROR)
181
+ self.use_rag = False
182
+
183
+ # -------------------------------------------------------------------------
184
+ # Basic Quality Score Computation
185
+ # -------------------------------------------------------------------------
186
+ def compute_quality_score(
187
+ self,
188
+ completeness: float,
189
+ relevance: float,
190
+ clarity: float,
191
+ accuracy: float,
192
+ weights: Dict[str, float] = None
193
+ ) -> float:
194
+ """
195
+ Combine multiple sub-metrics into one final score, as an example.
196
+
197
+ Args:
198
+ completeness (float): 0-1
199
+ relevance (float): 0-1
200
+ clarity (float): 0-1
201
+ accuracy (float): 0-1
202
+ weights (Dict[str, float]): optional weighting like {"completeness": 0.25, "relevance": 0.3, ...}
203
+
204
+ Returns:
205
+ float: Weighted average 0-1
206
+ """
207
+ if not weights:
208
+ weights = {
209
+ "completeness": 0.25,
210
+ "relevance": 0.25,
211
+ "clarity": 0.25,
212
+ "accuracy": 0.25
213
+ }
214
+ total = (completeness * weights["completeness"]
215
+ + relevance * weights["relevance"]
216
+ + clarity * weights["clarity"]
217
+ + accuracy * weights["accuracy"]
218
+ )
219
+ return round(total, 3) # e.g. round to 3 decimal places
220
+
221
+ # -------------------------------------------------------------------------
222
+ # Short-Term Methods
223
+ # -------------------------------------------------------------------------
224
+ def store_short_term(
225
+ self,
226
+ text: str,
227
+ metadata: Dict[str, Any] = None,
228
+ completeness: float = None,
229
+ relevance: float = None,
230
+ clarity: float = None,
231
+ accuracy: float = None,
232
+ weights: Dict[str, float] = None,
233
+ evaluator_quality: float = None
234
+ ):
235
+ """Store in short-term memory with optional quality metrics"""
236
+ logger.info(f"Storing in short-term memory: {text[:100]}...")
237
+ logger.info(f"Metadata: {metadata}")
238
+
239
+ metadata = self._process_quality_metrics(
240
+ metadata, completeness, relevance, clarity,
241
+ accuracy, weights, evaluator_quality
242
+ )
243
+ logger.info(f"Processed metadata: {metadata}")
244
+
245
+ # Existing store logic
246
+ try:
247
+ conn = sqlite3.connect(self.short_db)
248
+ ident = str(time.time_ns())
249
+ conn.execute(
250
+ "INSERT INTO short_mem (id, content, meta, created_at) VALUES (?,?,?,?)",
251
+ (ident, text, json.dumps(metadata), time.time())
252
+ )
253
+ conn.commit()
254
+ conn.close()
255
+ logger.info(f"Successfully stored in short-term memory with ID: {ident}")
256
+ except Exception as e:
257
+ logger.error(f"Failed to store in short-term memory: {e}")
258
+ raise
259
+
260
+ def search_short_term(
261
+ self,
262
+ query: str,
263
+ limit: int = 5,
264
+ min_quality: float = 0.0,
265
+ relevance_cutoff: float = 0.0
266
+ ) -> List[Dict[str, Any]]:
267
+ """Search short-term memory with optional quality filter"""
268
+ self._log_verbose(f"Searching short memory for: {query}")
269
+
270
+ if self.use_mem0 and hasattr(self, "mem0_client"):
271
+ results = self.mem0_client.search(query=query, limit=limit)
272
+ filtered = [r for r in results if r.get("score", 1.0) >= relevance_cutoff]
273
+ return filtered
274
+
275
+ elif self.use_rag and hasattr(self, "chroma_col"):
276
+ try:
277
+ from openai import OpenAI
278
+ client = OpenAI()
279
+
280
+ response = client.embeddings.create(
281
+ input=query,
282
+ model="text-embedding-3-small"
283
+ )
284
+ query_embedding = response.data[0].embedding
285
+
286
+ resp = self.chroma_col.query(
287
+ query_embeddings=[query_embedding],
288
+ n_results=limit
289
+ )
290
+
291
+ results = []
292
+ if resp["ids"]:
293
+ for i in range(len(resp["ids"][0])):
294
+ metadata = resp["metadatas"][0][i] if "metadatas" in resp else {}
295
+ quality = metadata.get("quality", 0.0)
296
+ score = 1.0 - (resp["distances"][0][i] if "distances" in resp else 0.0)
297
+ if quality >= min_quality and score >= relevance_cutoff:
298
+ results.append({
299
+ "id": resp["ids"][0][i],
300
+ "text": resp["documents"][0][i],
301
+ "metadata": metadata,
302
+ "score": score
303
+ })
304
+ return results
305
+ except Exception as e:
306
+ self._log_verbose(f"Error searching ChromaDB: {e}", logging.ERROR)
307
+ return []
308
+
309
+ else:
310
+ # Local fallback
311
+ conn = sqlite3.connect(self.short_db)
312
+ c = conn.cursor()
313
+ rows = c.execute(
314
+ "SELECT id, content, meta FROM short_mem WHERE content LIKE ? LIMIT ?",
315
+ (f"%{query}%", limit)
316
+ ).fetchall()
317
+ conn.close()
318
+
319
+ results = []
320
+ for row in rows:
321
+ meta = json.loads(row[2] or "{}")
322
+ quality = meta.get("quality", 0.0)
323
+ if quality >= min_quality:
324
+ results.append({
325
+ "id": row[0],
326
+ "text": row[1],
327
+ "metadata": meta
328
+ })
329
+ return results
330
+
331
+ def reset_short_term(self):
332
+ """Completely clears short-term memory."""
333
+ conn = sqlite3.connect(self.short_db)
334
+ conn.execute("DELETE FROM short_mem")
335
+ conn.commit()
336
+ conn.close()
337
+
338
+ # -------------------------------------------------------------------------
339
+ # Long-Term Methods
340
+ # -------------------------------------------------------------------------
341
+ def _sanitize_metadata(self, metadata: Dict) -> Dict:
342
+ """Sanitize metadata for ChromaDB - convert to acceptable types"""
343
+ sanitized = {}
344
+ for k, v in metadata.items():
345
+ if v is None:
346
+ continue
347
+ if isinstance(v, (str, int, float, bool)):
348
+ sanitized[k] = v
349
+ elif isinstance(v, dict):
350
+ # Convert dict to string representation
351
+ sanitized[k] = str(v)
352
+ else:
353
+ # Convert other types to string
354
+ sanitized[k] = str(v)
355
+ return sanitized
356
+
357
+ def store_long_term(
358
+ self,
359
+ text: str,
360
+ metadata: Dict[str, Any] = None,
361
+ completeness: float = None,
362
+ relevance: float = None,
363
+ clarity: float = None,
364
+ accuracy: float = None,
365
+ weights: Dict[str, float] = None,
366
+ evaluator_quality: float = None
367
+ ):
368
+ """Store in long-term memory with optional quality metrics"""
369
+ logger.info(f"Storing in long-term memory: {text[:100]}...")
370
+ logger.info(f"Initial metadata: {metadata}")
371
+
372
+ # Process metadata
373
+ metadata = metadata or {}
374
+ metadata = self._process_quality_metrics(
375
+ metadata, completeness, relevance, clarity,
376
+ accuracy, weights, evaluator_quality
377
+ )
378
+ logger.info(f"Processed metadata: {metadata}")
379
+
380
+ # Generate unique ID
381
+ ident = str(time.time_ns())
382
+ created = time.time()
383
+
384
+ # Store in SQLite
385
+ try:
386
+ conn = sqlite3.connect(self.long_db)
387
+ conn.execute(
388
+ "INSERT INTO long_mem (id, content, meta, created_at) VALUES (?,?,?,?)",
389
+ (ident, text, json.dumps(metadata), created)
390
+ )
391
+ conn.commit()
392
+ conn.close()
393
+ logger.info(f"Successfully stored in SQLite with ID: {ident}")
394
+ except Exception as e:
395
+ logger.error(f"Error storing in SQLite: {e}")
396
+ return
397
+
398
+ # Store in vector database if enabled
399
+ if self.use_rag and hasattr(self, "chroma_col"):
400
+ try:
401
+ from openai import OpenAI
402
+ client = OpenAI(api_key=os.getenv("OPENAI_API_KEY")) # Ensure API key is correctly set
403
+
404
+ logger.info("Getting embeddings from OpenAI...")
405
+ logger.debug(f"Embedding input text: {text}") # Log the input text
406
+
407
+ response = client.embeddings.create(
408
+ input=text,
409
+ model="text-embedding-3-small"
410
+ )
411
+ embedding = response.data[0].embedding
412
+ logger.info("Successfully got embeddings")
413
+ logger.debug(f"Received embedding of length: {len(embedding)}") # Log embedding details
414
+
415
+ # Sanitize metadata for ChromaDB
416
+ sanitized_metadata = self._sanitize_metadata(metadata)
417
+
418
+ # Store in ChromaDB with embedding
419
+ self.chroma_col.add(
420
+ documents=[text],
421
+ metadatas=[sanitized_metadata],
422
+ ids=[ident],
423
+ embeddings=[embedding]
424
+ )
425
+ logger.info(f"Successfully stored in ChromaDB with ID: {ident}")
426
+ except Exception as e:
427
+ logger.error(f"Error storing in ChromaDB: {e}")
428
+
429
+ elif self.use_mem0 and hasattr(self, "mem0_client"):
430
+ try:
431
+ self.mem0_client.add(text, metadata=metadata)
432
+ logger.info("Successfully stored in Mem0")
433
+ except Exception as e:
434
+ logger.error(f"Error storing in Mem0: {e}")
435
+
436
+
437
+ def search_long_term(
438
+ self,
439
+ query: str,
440
+ limit: int = 5,
441
+ relevance_cutoff: float = 0.0,
442
+ min_quality: float = 0.0
443
+ ) -> List[Dict[str, Any]]:
444
+ """Search long-term memory with optional quality filter"""
445
+ self._log_verbose(f"Searching long memory for: {query}")
446
+ self._log_verbose(f"Min quality: {min_quality}")
447
+
448
+ found = []
449
+
450
+ if self.use_mem0 and hasattr(self, "mem0_client"):
451
+ results = self.mem0_client.search(query=query, limit=limit)
452
+ # Filter by quality
453
+ filtered = [r for r in results if r.get("metadata", {}).get("quality", 0.0) >= min_quality]
454
+ logger.info(f"Found {len(filtered)} results in Mem0")
455
+ return filtered
456
+
457
+ elif self.use_rag and hasattr(self, "chroma_col"):
458
+ try:
459
+ from openai import OpenAI
460
+ client = OpenAI()
461
+
462
+ # Get query embedding
463
+ response = client.embeddings.create(
464
+ input=query,
465
+ model="text-embedding-3-small" # Using consistent model
466
+ )
467
+ query_embedding = response.data[0].embedding
468
+
469
+ # Search ChromaDB with embedding
470
+ resp = self.chroma_col.query(
471
+ query_embeddings=[query_embedding],
472
+ n_results=limit,
473
+ include=["documents", "metadatas", "distances"]
474
+ )
475
+
476
+ results = []
477
+ if resp["ids"]:
478
+ for i in range(len(resp["ids"][0])):
479
+ metadata = resp["metadatas"][0][i] if "metadatas" in resp else {}
480
+ text = resp["documents"][0][i]
481
+ # Add memory record citation
482
+ text = f"{text} (Memory record: {text})"
483
+ found.append({
484
+ "id": resp["ids"][0][i],
485
+ "text": text,
486
+ "metadata": metadata,
487
+ "score": 1.0 - (resp["distances"][0][i] if "distances" in resp else 0.0)
488
+ })
489
+ logger.info(f"Found {len(found)} results in ChromaDB")
490
+
491
+ except Exception as e:
492
+ self._log_verbose(f"Error searching ChromaDB: {e}", logging.ERROR)
493
+
494
+ # Always try SQLite as fallback or additional source
495
+ conn = sqlite3.connect(self.long_db)
496
+ c = conn.cursor()
497
+ rows = c.execute(
498
+ "SELECT id, content, meta, created_at FROM long_mem WHERE content LIKE ? LIMIT ?",
499
+ (f"%{query}%", limit)
500
+ ).fetchall()
501
+ conn.close()
502
+
503
+ for row in rows:
504
+ meta = json.loads(row[2] or "{}")
505
+ text = row[1]
506
+ # Add memory record citation if not already present
507
+ if "(Memory record:" not in text:
508
+ text = f"{text} (Memory record: {text})"
509
+ # Only add if not already found by ChromaDB/Mem0
510
+ if not any(f["id"] == row[0] for f in found):
511
+ found.append({
512
+ "id": row[0],
513
+ "text": text,
514
+ "metadata": meta,
515
+ "created_at": row[3]
516
+ })
517
+ logger.info(f"Found {len(found)} total results after SQLite")
518
+
519
+ results = found
520
+
521
+ # Filter by quality if needed
522
+ if min_quality > 0:
523
+ self._log_verbose(f"Found {len(results)} initial results")
524
+ results = [
525
+ r for r in results
526
+ if r.get("metadata", {}).get("quality", 0.0) >= min_quality
527
+ ]
528
+ self._log_verbose(f"After quality filter: {len(results)} results")
529
+
530
+ # Apply relevance cutoff if specified
531
+ if relevance_cutoff > 0:
532
+ results = [r for r in results if r.get("score", 1.0) >= relevance_cutoff]
533
+ logger.info(f"After relevance filter: {len(results)} results")
534
+
535
+ return results[:limit]
536
+
537
+ def reset_long_term(self):
538
+ """Clear local LTM DB, plus Chroma or mem0 if in use."""
539
+ conn = sqlite3.connect(self.long_db)
540
+ conn.execute("DELETE FROM long_mem")
541
+ conn.commit()
542
+ conn.close()
543
+
544
+ if self.use_mem0 and hasattr(self, "mem0_client"):
545
+ # Mem0 has no universal reset API. Could implement partial or no-op.
546
+ pass
547
+ if self.use_rag and hasattr(self, "chroma_client"):
548
+ self.chroma_client.reset() # entire DB
549
+ self._init_chroma() # re-init fresh
550
+
551
+ # -------------------------------------------------------------------------
552
+ # Entity Memory Methods
553
+ # -------------------------------------------------------------------------
554
+ def store_entity(self, name: str, type_: str, desc: str, relations: str):
555
+ """
556
+ Save entity info in LTM (or mem0/rag).
557
+ We'll label the metadata type = entity for easy filtering.
558
+ """
559
+ data = f"Entity {name}({type_}): {desc} | relationships: {relations}"
560
+ self.store_long_term(data, metadata={"category": "entity"})
561
+
562
+ def search_entity(self, query: str, limit: int = 5) -> List[Dict[str, Any]]:
563
+ """
564
+ Filter to items that have metadata 'category=entity'.
565
+ """
566
+ all_hits = self.search_long_term(query, limit=20) # gather more
567
+ ents = []
568
+ for h in all_hits:
569
+ meta = h.get("metadata") or {}
570
+ if meta.get("category") == "entity":
571
+ ents.append(h)
572
+ return ents[:limit]
573
+
574
+ def reset_entity_only(self):
575
+ """
576
+ If you only want to drop entity items from LTM, you'd do a custom
577
+ delete from local DB where meta LIKE '%category=entity%'.
578
+ For brevity, we do a full LTM reset here.
579
+ """
580
+ self.reset_long_term()
581
+
582
+ # -------------------------------------------------------------------------
583
+ # User Memory Methods
584
+ # -------------------------------------------------------------------------
585
+ def store_user_memory(self, user_id: str, text: str, extra: Dict[str, Any] = None):
586
+ """
587
+ If mem0 is used, do user-based addition. Otherwise store in LTM with user in metadata.
588
+ """
589
+ meta = {"user_id": user_id}
590
+ if extra:
591
+ meta.update(extra)
592
+
593
+ if self.use_mem0 and hasattr(self, "mem0_client"):
594
+ self.mem0_client.add(text, user_id=user_id, metadata=meta)
595
+ else:
596
+ self.store_long_term(text, metadata=meta)
597
+
598
+ def search_user_memory(self, user_id: str, query: str, limit: int = 5) -> List[Dict[str, Any]]:
599
+ """
600
+ If mem0 is used, pass user_id in. Otherwise fallback to local filter on user in metadata.
601
+ """
602
+ if self.use_mem0 and hasattr(self, "mem0_client"):
603
+ return self.mem0_client.search(query=query, limit=limit, user_id=user_id)
604
+ else:
605
+ hits = self.search_long_term(query, limit=20)
606
+ filtered = []
607
+ for h in hits:
608
+ meta = h.get("metadata", {})
609
+ if meta.get("user_id") == user_id:
610
+ filtered.append(h)
611
+ return filtered[:limit]
612
+
613
+ def reset_user_memory(self):
614
+ """
615
+ Clear all user-based info. For simplicity, we do a full LTM reset.
616
+ Real usage might filter only metadata "user_id".
617
+ """
618
+ self.reset_long_term()
619
+
620
+ # -------------------------------------------------------------------------
621
+ # Putting it all Together: Task Finalization
622
+ # -------------------------------------------------------------------------
623
+ def finalize_task_output(
624
+ self,
625
+ content: str,
626
+ agent_name: str,
627
+ quality_score: float,
628
+ threshold: float = 0.7,
629
+ metrics: Dict[str, Any] = None,
630
+ task_id: str = None
631
+ ):
632
+ """Store task output in memory with appropriate metadata"""
633
+ logger.info(f"Finalizing task output: {content[:100]}...")
634
+ logger.info(f"Agent: {agent_name}, Quality: {quality_score}, Threshold: {threshold}")
635
+
636
+ metadata = {
637
+ "task_id": task_id,
638
+ "agent": agent_name,
639
+ "quality": quality_score,
640
+ "metrics": metrics,
641
+ "task_type": "output",
642
+ "stored_at": time.time()
643
+ }
644
+ logger.info(f"Prepared metadata: {metadata}")
645
+
646
+ # Always store in short-term memory
647
+ try:
648
+ logger.info("Storing in short-term memory...")
649
+ self.store_short_term(
650
+ text=content,
651
+ metadata=metadata
652
+ )
653
+ logger.info("Successfully stored in short-term memory")
654
+ except Exception as e:
655
+ logger.error(f"Failed to store in short-term memory: {e}")
656
+
657
+ # Store in long-term memory if quality meets threshold
658
+ if quality_score >= threshold:
659
+ try:
660
+ logger.info(f"Quality score {quality_score} >= {threshold}, storing in long-term memory...")
661
+ self.store_long_term(
662
+ text=content,
663
+ metadata=metadata
664
+ )
665
+ logger.info("Successfully stored in long-term memory")
666
+ except Exception as e:
667
+ logger.error(f"Failed to store in long-term memory: {e}")
668
+ else:
669
+ logger.info(f"Quality score {quality_score} < {threshold}, skipping long-term storage")
670
+
671
+ # -------------------------------------------------------------------------
672
+ # Building Context (Short, Long, Entities, User)
673
+ # -------------------------------------------------------------------------
674
+ def build_context_for_task(
675
+ self,
676
+ task_descr: str,
677
+ user_id: Optional[str] = None,
678
+ additional: str = "",
679
+ max_items: int = 3
680
+ ) -> str:
681
+ """
682
+ Merges relevant short-term, long-term, entity, user memories
683
+ into a single text block with deduplication and clean formatting.
684
+ """
685
+ q = (task_descr + " " + additional).strip()
686
+ lines = []
687
+ seen_contents = set() # Track unique contents
688
+
689
+ def normalize_content(content: str) -> str:
690
+ """Normalize content for deduplication"""
691
+ # Extract just the main content without citations for comparison
692
+ normalized = content.split("(Memory record:")[0].strip()
693
+ # Keep more characters to reduce false duplicates
694
+ normalized = ''.join(c.lower() for c in normalized if not c.isspace())
695
+ return normalized
696
+
697
+ def format_content(content: str, max_len: int = 150) -> str:
698
+ """Format content with clean truncation at word boundaries"""
699
+ if not content:
700
+ return ""
701
+
702
+ # Clean up content by removing extra whitespace and newlines
703
+ content = ' '.join(content.split())
704
+
705
+ # If content contains a memory citation, preserve it
706
+ if "(Memory record:" in content:
707
+ return content # Keep original citation format
708
+
709
+ # Regular content truncation
710
+ if len(content) <= max_len:
711
+ return content
712
+
713
+ truncate_at = content.rfind(' ', 0, max_len - 3)
714
+ if truncate_at == -1:
715
+ truncate_at = max_len - 3
716
+ return content[:truncate_at] + "..."
717
+
718
+ def add_section(title: str, hits: List[Any]) -> None:
719
+ """Add a section of memory hits with deduplication"""
720
+ if not hits:
721
+ return
722
+
723
+ formatted_hits = []
724
+ for h in hits:
725
+ content = h.get('text', '') if isinstance(h, dict) else str(h)
726
+ if not content:
727
+ continue
728
+
729
+ # Keep original format if it has a citation
730
+ if "(Memory record:" in content:
731
+ formatted = content
732
+ else:
733
+ formatted = format_content(content)
734
+
735
+ # Only add if we haven't seen this normalized content before
736
+ normalized = normalize_content(formatted)
737
+ if normalized not in seen_contents:
738
+ seen_contents.add(normalized)
739
+ formatted_hits.append(formatted)
740
+
741
+ if formatted_hits:
742
+ # Add section header
743
+ if lines:
744
+ lines.append("") # Space before new section
745
+ lines.append(title)
746
+ lines.append("=" * len(title)) # Underline the title
747
+ lines.append("") # Space after title
748
+
749
+ # Add formatted content with bullet points
750
+ for content in formatted_hits:
751
+ lines.append(f" • {content}")
752
+
753
+ # Add each section
754
+ # First get all results
755
+ short_term = self.search_short_term(q, limit=max_items)
756
+ long_term = self.search_long_term(q, limit=max_items)
757
+ entities = self.search_entity(q, limit=max_items)
758
+ user_mem = self.search_user_memory(user_id, q, limit=max_items) if user_id else []
759
+
760
+ # Add sections in order of priority
761
+ add_section("Short-term Memory Context", short_term)
762
+ add_section("Long-term Memory Context", long_term)
763
+ add_section("Entity Context", entities)
764
+ if user_id:
765
+ add_section("User Context", user_mem)
766
+
767
+ return "\n".join(lines) if lines else ""
768
+
769
+ # -------------------------------------------------------------------------
770
+ # Master Reset (Everything)
771
+ # -------------------------------------------------------------------------
772
+ def reset_all(self):
773
+ """
774
+ Fully wipes short-term, long-term, and any memory in mem0 or rag.
775
+ """
776
+ self.reset_short_term()
777
+ self.reset_long_term()
778
+ # Entities & user memory are stored in LTM or mem0, so no separate step needed.
779
+
780
+ def _process_quality_metrics(
781
+ self,
782
+ metadata: Dict[str, Any],
783
+ completeness: float = None,
784
+ relevance: float = None,
785
+ clarity: float = None,
786
+ accuracy: float = None,
787
+ weights: Dict[str, float] = None,
788
+ evaluator_quality: float = None
789
+ ) -> Dict[str, Any]:
790
+ """Process and store quality metrics in metadata"""
791
+ metadata = metadata or {}
792
+
793
+ # Handle sub-metrics if provided
794
+ if None not in [completeness, relevance, clarity, accuracy]:
795
+ metadata.update({
796
+ "completeness": completeness,
797
+ "relevance": relevance,
798
+ "clarity": clarity,
799
+ "accuracy": accuracy,
800
+ "quality": self.compute_quality_score(
801
+ completeness, relevance, clarity, accuracy, weights
802
+ )
803
+ })
804
+ # Handle external evaluator quality if provided
805
+ elif evaluator_quality is not None:
806
+ metadata["quality"] = evaluator_quality
807
+
808
+ return metadata
809
+
810
+ def calculate_quality_metrics(
811
+ self,
812
+ output: str,
813
+ expected_output: str,
814
+ llm: Optional[str] = None,
815
+ custom_prompt: Optional[str] = None
816
+ ) -> Dict[str, float]:
817
+ """Calculate quality metrics using LLM"""
818
+ logger.info("Calculating quality metrics for output")
819
+ logger.info(f"Output: {output[:100]}...")
820
+ logger.info(f"Expected: {expected_output[:100]}...")
821
+
822
+ # Default evaluation prompt
823
+ default_prompt = f"""
824
+ Evaluate the following output against expected output.
825
+ Score each metric from 0.0 to 1.0:
826
+ - Completeness: Does it address all requirements?
827
+ - Relevance: Does it match expected output?
828
+ - Clarity: Is it clear and well-structured?
829
+ - Accuracy: Is it factually correct?
830
+
831
+ Expected: {expected_output}
832
+ Actual: {output}
833
+
834
+ Return ONLY a JSON with these keys: completeness, relevance, clarity, accuracy
835
+ Example: {{"completeness": 0.95, "relevance": 0.8, "clarity": 0.9, "accuracy": 0.85}}
836
+ """
837
+
838
+ try:
839
+ # Use OpenAI client from main.py
840
+ from ..main import client
841
+
842
+ response = client.chat.completions.create(
843
+ model=llm or "gpt-4o",
844
+ messages=[{
845
+ "role": "user",
846
+ "content": custom_prompt or default_prompt
847
+ }],
848
+ response_format={"type": "json_object"},
849
+ temperature=0.3
850
+ )
851
+
852
+ metrics = json.loads(response.choices[0].message.content)
853
+
854
+ # Validate metrics
855
+ required = ["completeness", "relevance", "clarity", "accuracy"]
856
+ if not all(k in metrics for k in required):
857
+ raise ValueError("Missing required metrics in LLM response")
858
+
859
+ logger.info(f"Calculated metrics: {metrics}")
860
+ return metrics
861
+
862
+ except Exception as e:
863
+ logger.error(f"Error calculating metrics: {e}")
864
+ return {
865
+ "completeness": 0.0,
866
+ "relevance": 0.0,
867
+ "clarity": 0.0,
868
+ "accuracy": 0.0
869
+ }
870
+
871
+ def store_quality(
872
+ self,
873
+ text: str,
874
+ quality_score: float,
875
+ task_id: Optional[str] = None,
876
+ iteration: Optional[int] = None,
877
+ metrics: Optional[Dict[str, float]] = None,
878
+ memory_type: Literal["short", "long"] = "long"
879
+ ) -> None:
880
+ """Store quality metrics in memory"""
881
+ logger.info(f"Attempting to store in {memory_type} memory: {text[:100]}...")
882
+
883
+ metadata = {
884
+ "quality": quality_score,
885
+ "task_id": task_id,
886
+ "iteration": iteration
887
+ }
888
+
889
+ if metrics:
890
+ metadata.update({
891
+ k: v for k, v in metrics.items() # Remove metric_ prefix
892
+ })
893
+
894
+ logger.info(f"With metadata: {metadata}")
895
+
896
+ try:
897
+ if memory_type == "short":
898
+ self.store_short_term(text, metadata=metadata)
899
+ logger.info("Successfully stored in short-term memory")
900
+ else:
901
+ self.store_long_term(text, metadata=metadata)
902
+ logger.info("Successfully stored in long-term memory")
903
+ except Exception as e:
904
+ logger.error(f"Failed to store in memory: {e}")
905
+
906
+ def search_with_quality(
907
+ self,
908
+ query: str,
909
+ min_quality: float = 0.0,
910
+ memory_type: Literal["short", "long"] = "long",
911
+ limit: int = 5
912
+ ) -> List[Dict[str, Any]]:
913
+ """Search with quality filter"""
914
+ logger.info(f"Searching {memory_type} memory for: {query}")
915
+ logger.info(f"Min quality: {min_quality}")
916
+
917
+ search_func = (
918
+ self.search_short_term if memory_type == "short"
919
+ else self.search_long_term
920
+ )
921
+
922
+ results = search_func(query, limit=limit)
923
+ logger.info(f"Found {len(results)} initial results")
924
+
925
+ filtered = [
926
+ r for r in results
927
+ if r.get("metadata", {}).get("quality", 0.0) >= min_quality
928
+ ]
929
+ logger.info(f"After quality filter: {len(filtered)} results")
930
+
931
+ return filtered