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