from __future__ import annotations import os import uuid import asyncio import json from typing import Optional from dotenv import load_dotenv import psycopg2 from psycopg2.extras import Json from sentence_transformers import SentenceTransformer from server.models import EpisodeMemory, AblationRecord load_dotenv() EMBED_MODEL = "sentence-transformers/all-MiniLM-L6-v2" def get_connection(): return psycopg2.connect( host = os.getenv("DB_HOST", "127.0.0.1"), port = int(os.getenv("DB_PORT", 5433)), dbname = os.getenv("DB_NAME", "sre_env"), user = os.getenv("DB_USER", "postgres"), password = os.getenv("DB_PASSWORD", ""), ) class AgenticMemory: """ Episode-level agentic memory using PostgreSQL + pgvector. Pattern: - Episode start: get_episodic() + search_semantic() → RAM cache - Steps 1..N: read from RAM cache (zero DB calls) - Episode end: asyncio.create_task(write()) → non-blocking """ def __init__(self, backend: str = "pgvector"): self.backend = backend self._model: Optional[SentenceTransformer] = None if backend == "pgvector": self._model = SentenceTransformer(EMBED_MODEL) # ------------------------------------------------------------------ # MCP tools # ------------------------------------------------------------------ def get_episodic(self, task_id: str, limit: int = 5) -> str: """ Retrieves last N episodes for this task_id. Called once at episode start — result cached in RAM. """ if self.backend == "disabled": return "" try: conn = get_connection() cur = conn.cursor() cur.execute(""" SELECT task_id, task_success, total_reward, steps_taken, summary FROM episodes WHERE task_id = %s ORDER BY created_at DESC LIMIT %s """, (task_id, limit)) rows = cur.fetchall() cur.close() conn.close() if not rows: return "" lines = ["Past episodes for this task:"] for row in rows: tid, success, reward, steps, summary = row lines.append( f"- task={tid} success={success} " f"reward={reward:.3f} steps={steps} | {summary[:100]}" ) return "\n".join(lines) except Exception as e: print(f"get_episodic error: {e}") return "" def search_semantic(self, state_summary: str, top_k: int = 3) -> str: """ Finds top-k semantically similar past episodes via vector search. Called once at episode start — result cached in RAM. """ if self.backend == "disabled" or not self._model: return "" try: emb = self._model.encode(state_summary, normalize_embeddings=True) emb_list = emb.tolist() conn = get_connection() cur = conn.cursor() cur.execute(""" SELECT task_id, task_success, total_reward, steps_taken, summary, 1 - (state_emb <=> %s::vector) AS similarity FROM episodes ORDER BY state_emb <=> %s::vector LIMIT %s """, (emb_list, emb_list, top_k)) rows = cur.fetchall() cur.close() conn.close() if not rows: return "" lines = ["Semantically similar past episodes:"] for row in rows: tid, success, reward, steps, summary, sim = row lines.append( f"- task={tid} success={success} " f"reward={reward:.3f} steps={steps} " f"similarity={sim:.3f} | {summary[:100]}" ) return "\n".join(lines) except Exception as e: print(f"search_semantic error: {e}") return "" def write(self, memory: EpisodeMemory): """ Writes episode to DB. Called non-blocking via asyncio.create_task(). Stores state embedding for future semantic search. """ if self.backend == "disabled": return try: emb = None if self._model and memory.summary: emb = self._model.encode( memory.summary, normalize_embeddings=True ).tolist() conn = get_connection() cur = conn.cursor() cur.execute(""" INSERT INTO episodes (episode_id, task_id, task_success, total_reward, steps_taken, actions, summary, state_emb) VALUES (%s, %s, %s, %s, %s, %s, %s, %s) """, ( memory.episode_id, memory.task_id, memory.task_success, memory.total_reward, memory.steps_taken, Json(memory.actions), memory.summary, emb, )) conn.commit() cur.close() conn.close() except Exception as e: print(f"memory.write error: {e}") async def write_async(self, memory: EpisodeMemory): """Async wrapper for non-blocking episode end write.""" loop = asyncio.get_event_loop() await loop.run_in_executor(None, self.write, memory) # ------------------------------------------------------------------ # Ablation tracking (point 7) # ------------------------------------------------------------------ def write_ablation_record(self, record: AblationRecord): """Saves ablation metrics after each epoch.""" if self.backend == "disabled": return try: conn = get_connection() cur = conn.cursor() cur.execute(""" INSERT INTO ablation_records (epoch, task_id, memory_backend, mean_reward, mean_steps_to_resolution, task_success_rate) VALUES (%s, %s, %s, %s, %s, %s) """, ( record.epoch, record.task_id, record.memory_backend, record.mean_reward, record.mean_steps_to_resolution, record.task_success_rate, )) conn.commit() cur.close() conn.close() except Exception as e: print(f"write_ablation_record error: {e}") def get_ablation_records(self, task_id: str) -> list: """Retrieves ablation records for Gradio dashboard plotting.""" if self.backend == "disabled": return [] try: conn = get_connection() cur = conn.cursor() cur.execute(""" SELECT epoch, memory_backend, mean_reward, mean_steps_to_resolution, task_success_rate FROM ablation_records WHERE task_id = %s ORDER BY epoch ASC """, (task_id,)) rows = cur.fetchall() cur.close() conn.close() return rows except Exception as e: print(f"get_ablation_records error: {e}") return [] # ------------------------------------------------------------------ # Cache builder — called once at episode start # ------------------------------------------------------------------ def build_system_context( self, task_id: str, state_summary: str, runbook: str = "", ) -> str: """ Merges episodic memory + semantic search + runbook into a single system_context string. Cached in RAM for the full episode. """ episodic = self.get_episodic(task_id) semantic = self.search_semantic(state_summary) parts = [] if runbook: parts.append(f"=== Runbook ===\n{runbook}") if episodic: parts.append(f"=== Episodic Memory ===\n{episodic}") if semantic: parts.append(f"=== Similar Episodes ===\n{semantic}") return "\n\n".join(parts)