praisonaiagents 0.0.28__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.
- praisonaiagents/__init__.py +2 -0
- praisonaiagents/agent/agent.py +2 -2
- praisonaiagents/agents/agents.py +169 -35
- praisonaiagents/memory/memory.py +928 -0
- praisonaiagents/task/task.py +157 -9
- {praisonaiagents-0.0.28.dist-info → praisonaiagents-0.0.29.dist-info}/METADATA +4 -2
- {praisonaiagents-0.0.28.dist-info → praisonaiagents-0.0.29.dist-info}/RECORD +9 -16
- {praisonaiagents-0.0.28.dist-info → praisonaiagents-0.0.29.dist-info}/WHEEL +1 -1
- praisonaiagents/build/lib/praisonaiagents/__init__.py +0 -1
- praisonaiagents/build/lib/praisonaiagents/agent/__init__.py +0 -4
- praisonaiagents/build/lib/praisonaiagents/agent/agent.py +0 -350
- praisonaiagents/build/lib/praisonaiagents/agents/__init__.py +0 -4
- praisonaiagents/build/lib/praisonaiagents/agents/agents.py +0 -318
- praisonaiagents/build/lib/praisonaiagents/main.py +0 -112
- praisonaiagents/build/lib/praisonaiagents/task/__init__.py +0 -4
- praisonaiagents/build/lib/praisonaiagents/task/task.py +0 -48
- {praisonaiagents-0.0.28.dist-info → praisonaiagents-0.0.29.dist-info}/top_level.txt +0 -0
@@ -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
|