praisonaiagents 0.0.27__py3-none-any.whl → 0.0.29__py3-none-any.whl

Sign up to get free protection for your applications and to get access to all the features.
@@ -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