""" Agent memory system (PostgreSQL). This module stores agent experiences in PostgreSQL and retrieves relevant past cases to inject into prompts (RAG-style). It does NOT finetune model weights. Retrieval (configurable): - Vector similarity via deterministic local embeddings (default) - Fallback to difflib text similarity when embeddings are missing Ranking combines: - similarity - recency decay (half-life) - optional returns weight """ import json import os import math from typing import List, Dict, Any, Optional from datetime import datetime, timezone import difflib from app.utils.logger import get_logger from app.utils.db import get_db_connection from .embedding import EmbeddingService, cosine_sim logger = get_logger(__name__) class AgentMemory: """Agent memory system using PostgreSQL""" def __init__(self, agent_name: str, db_path: Optional[str] = None): """ Initialize memory system. Args: agent_name: Agent identifier (e.g., 'trader_agent', 'risk_analyst') db_path: Deprecated parameter, kept for backward compatibility """ self.agent_name = agent_name self.embedder = EmbeddingService() self.enable_vector = os.getenv("AGENT_MEMORY_ENABLE_VECTOR", "true").lower() == "true" self.candidate_limit = int(os.getenv("AGENT_MEMORY_CANDIDATE_LIMIT", "500") or 500) self.half_life_days = float(os.getenv("AGENT_MEMORY_HALF_LIFE_DAYS", "30") or 30) self.w_sim = float(os.getenv("AGENT_MEMORY_W_SIM", "0.75") or 0.75) self.w_recency = float(os.getenv("AGENT_MEMORY_W_RECENCY", "0.20") or 0.20) self.w_returns = float(os.getenv("AGENT_MEMORY_W_RETURNS", "0.05") or 0.05) def _now_utc(self) -> datetime: return datetime.now(timezone.utc) def _parse_ts(self, ts_val: Any) -> Optional[datetime]: if ts_val is None: return None if isinstance(ts_val, datetime): return ts_val s = str(ts_val) try: return datetime.fromisoformat(s.replace("Z", "")) except Exception: return None def _recency_score(self, created_at: Any) -> float: dt = self._parse_ts(created_at) if not dt: return 0.0 if dt.tzinfo is None: dt = dt.replace(tzinfo=timezone.utc) age_days = max(0.0, (self._now_utc() - dt).total_seconds() / 86400.0) hl = max(0.1, float(self.half_life_days or 30.0)) return float(math.exp(-math.log(2.0) * (age_days / hl))) def _returns_score(self, returns: Any) -> float: try: r = float(returns) except Exception: return 0.0 return float(math.tanh(r / 10.0)) def _build_embed_text(self, situation: str, recommendation: str, result: Optional[str], features_json: Optional[str]) -> str: return "\n".join([ f"situation: {situation or ''}", f"recommendation: {recommendation or ''}", f"result: {result or ''}", f"features: {features_json or ''}", ]) def add_memory( self, situation: str, recommendation: str, result: Optional[str] = None, returns: Optional[float] = None, metadata: Optional[Dict[str, Any]] = None, ): """ Add a memory entry. Args: situation: Situation description recommendation: Decision/recommendation made result: Outcome description (optional) returns: Return percentage (optional) metadata: Optional structured metadata (market/symbol/timeframe/features...) """ try: meta = metadata or {} market = (meta.get("market") or "").strip() or None symbol = (meta.get("symbol") or "").strip() or None timeframe = (meta.get("timeframe") or "").strip() or None features = meta.get("features") if isinstance(meta, dict) else None try: features_json = json.dumps(features, ensure_ascii=False) if features is not None else None except Exception: features_json = None embedding_blob = None if self.enable_vector: text = self._build_embed_text(situation, recommendation, result, features_json) vec = self.embedder.embed(text) embedding_blob = self.embedder.to_bytes(vec) with get_db_connection() as conn: cur = conn.cursor() cur.execute( """ INSERT INTO qd_agent_memories (agent_name, situation, recommendation, result, returns, market, symbol, timeframe, features_json, embedding) VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?) """, (self.agent_name, situation, recommendation, result, returns, market, symbol, timeframe, features_json, embedding_blob) ) conn.commit() cur.close() logger.info(f"{self.agent_name} added new memory") except Exception as e: logger.error(f"Failed to add memory: {e}") def get_memories(self, current_situation: str, n_matches: int = 5, metadata: Optional[Dict[str, Any]] = None) -> List[Dict[str, Any]]: """ Retrieve similar memories. Args: current_situation: Current situation description n_matches: Number of matches to return metadata: Optional metadata for filtering/weighting Returns: List of matching memory entries """ try: with get_db_connection() as conn: cur = conn.cursor() cur.execute( """ SELECT id, situation, recommendation, result, returns, created_at, market, symbol, timeframe, features_json, embedding FROM qd_agent_memories WHERE agent_name = ? ORDER BY created_at DESC LIMIT ? """, (self.agent_name, int(self.candidate_limit)) ) all_memories = cur.fetchall() or [] cur.close() if not all_memories: return [] meta = metadata or {} tf = (meta.get("timeframe") or "").strip() features = meta.get("features") if isinstance(meta, dict) else None try: q_features_json = json.dumps(features, ensure_ascii=False) if features is not None else None except Exception: q_features_json = None query_vec = [] if self.enable_vector: query_text = self._build_embed_text(current_situation, "", "", q_features_json) query_vec = self.embedder.embed(query_text) ranked = [] for row in all_memories: mem_id = row['id'] situation = row['situation'] recommendation = row['recommendation'] result = row['result'] returns = row['returns'] created_at = row['created_at'] market = row['market'] symbol = row['symbol'] timeframe = row['timeframe'] features_json = row['features_json'] embedding_blob = row['embedding'] sim = 0.0 if self.enable_vector and embedding_blob: try: # Handle memoryview/bytes from PostgreSQL if isinstance(embedding_blob, memoryview): embedding_blob = bytes(embedding_blob) mem_vec = self.embedder.from_bytes(embedding_blob) sim = cosine_sim(query_vec, mem_vec) except Exception: sim = 0.0 else: sim = difflib.SequenceMatcher(None, (current_situation or "").lower(), (situation or "").lower()).ratio() rec = self._recency_score(created_at) ret = self._returns_score(returns) score = (self.w_sim * sim) + (self.w_recency * rec) + (self.w_returns * ret) if tf and timeframe and str(timeframe).strip() != tf: score -= 0.15 ranked.append({ 'id': mem_id, 'matched_situation': situation, 'recommendation': recommendation, 'result': result, 'returns': returns, 'created_at': created_at, 'market': market, 'symbol': symbol, 'timeframe': timeframe, 'features_json': features_json, 'score': float(score), 'sim': float(sim), 'recency': float(rec), }) ranked.sort(key=lambda x: x.get('score', 0.0), reverse=True) return ranked[: max(0, int(n_matches or 0))] except Exception as e: logger.error(f"Failed to retrieve memories: {e}") return [] def update_memory_result(self, memory_id: int, result: str, returns: Optional[float] = None): """ Update memory result. Args: memory_id: Memory ID result: Outcome description returns: Return percentage """ try: with get_db_connection() as conn: cur = conn.cursor() cur.execute( """ UPDATE qd_agent_memories SET result = ?, returns = ?, updated_at = NOW() WHERE id = ? AND agent_name = ? """, (result, returns, memory_id, self.agent_name) ) conn.commit() cur.close() logger.info(f"{self.agent_name} updated memory {memory_id}") except Exception as e: logger.error(f"Failed to update memory: {e}") def get_statistics(self) -> Dict[str, Any]: """Get memory statistics for this agent.""" try: with get_db_connection() as conn: cur = conn.cursor() cur.execute( 'SELECT COUNT(*) as cnt FROM qd_agent_memories WHERE agent_name = ?', (self.agent_name,) ) total = cur.fetchone()['cnt'] cur.execute( 'SELECT AVG(returns) as avg_ret FROM qd_agent_memories WHERE agent_name = ? AND returns IS NOT NULL', (self.agent_name,) ) avg_returns = cur.fetchone()['avg_ret'] or 0 cur.execute( 'SELECT COUNT(*) as cnt FROM qd_agent_memories WHERE agent_name = ? AND returns > 0', (self.agent_name,) ) positive = cur.fetchone()['cnt'] cur.close() return { 'total_memories': total, 'average_returns': round(avg_returns, 2), 'positive_decisions': positive, 'success_rate': round(positive / total * 100, 2) if total > 0 else 0 } except Exception as e: logger.error(f"Failed to get statistics: {e}") return {} def clear_memories(self): """Clear all memories for this agent (use with caution).""" try: with get_db_connection() as conn: cur = conn.cursor() cur.execute( 'DELETE FROM qd_agent_memories WHERE agent_name = ?', (self.agent_name,) ) conn.commit() cur.close() logger.warning(f"{self.agent_name} cleared all memories") except Exception as e: logger.error(f"Failed to clear memories: {e}")