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

This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
@@ -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