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.
- praisonaiagents/__init__.py +4 -0
- praisonaiagents/agent/agent.py +7 -3
- praisonaiagents/agents/__init__.py +2 -1
- praisonaiagents/agents/agents.py +169 -35
- praisonaiagents/agents/autoagents.py +335 -0
- praisonaiagents/memory/memory.py +928 -0
- praisonaiagents/task/task.py +157 -9
- {praisonaiagents-0.0.27.dist-info → praisonaiagents-0.0.29.dist-info}/METADATA +4 -2
- {praisonaiagents-0.0.27.dist-info → praisonaiagents-0.0.29.dist-info}/RECORD +11 -17
- {praisonaiagents-0.0.27.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.27.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
|