memory.py

#!/usr/bin/env python3
"""NEO Memory v2 - SQLite + ChromaDB Fallback"""
import sqlite3
import hashlib
from datetime import datetime
from pathlib import Path
from collections import deque
from typing import List, Dict

try:
    import chromadb
    from sentence_transformers import SentenceTransformer
    CHROMA_AVAILABLE = True
except Exception:
    CHROMA_AVAILABLE = False

from config_v2 import config


class Experience:
    def __init__(self, task_id, inp, action, result, reward, lesson=""):
        self.task_id   = task_id
        self.input     = inp
        self.action    = action
        self.result    = result
        self.reward    = reward
        self.lesson    = lesson
        self.timestamp = datetime.now().isoformat()
        self.id        = hashlib.md5(inp.encode()).hexdigest()


class NEOMemory:
    def __init__(self):
        self.db_path = config.MEMORY_PATH
        self.db_path.parent.mkdir(parents=True, exist_ok=True)
        self.replay_buffer = deque(maxlen=2000)
        self._init_sqlite()
        if CHROMA_AVAILABLE:
            self._init_chroma()

    # ── Init ──────────────────────────────────────────────────
    def _init_sqlite(self):
        conn = sqlite3.connect(self.db_path)
        conn.execute("""
            CREATE TABLE IF NOT EXISTS exp (
                id TEXT PRIMARY KEY, task TEXT, input TEXT, action TEXT,
                result TEXT, reward REAL, lesson TEXT, ts TEXT
            )
        """)
        conn.commit()
        conn.close()

    def _init_chroma(self):
        try:
            self.chroma   = chromadb.PersistentClient(path=str(config.CHROMA_PATH))
            self.coll     = self.chroma.get_or_create_collection("neo_exp")
            self.embedder = SentenceTransformer(config.EMBEDDING_MODEL)
        except Exception as e:
            print(f"[memory] ChromaDB init failed: {e} — using SQLite only.")
            global CHROMA_AVAILABLE
            CHROMA_AVAILABLE = False

    # ── Store ─────────────────────────────────────────────────
    def store(self, exp: Experience):
        conn = sqlite3.connect(self.db_path)
        conn.execute(
            "INSERT OR REPLACE INTO exp VALUES(?,?,?,?,?,?,?,?)",
            (exp.id, exp.task_id, exp.input, exp.action,
             exp.result, exp.reward, exp.lesson, exp.timestamp),
        )
        conn.commit()
        conn.close()
        self.replay_buffer.append(exp)

        if CHROMA_AVAILABLE:
            try:
                self.coll.add(
                    embeddings=[self.embedder.encode(exp.input).tolist()],
                    documents=[exp.input],
                    metadatas=[{"reward": exp.reward, "lesson": exp.lesson}],
                    ids=[exp.id],
                )
            except Exception:
                pass

    # ── Search ────────────────────────────────────────────────
    def search(self, query: str, n: int = 3) -> List[Dict]:
        if CHROMA_AVAILABLE:
            try:
                res = self.coll.query(
                    query_embeddings=[self.embedder.encode(query).tolist()],
                    n_results=min(n, self.coll.count()),
                )
                return res["metadatas"][0] if res["metadatas"] else []
            except Exception:
                pass

        # SQLite keyword fallback
        conn = sqlite3.connect(self.db_path)
        cur  = conn.execute(
            "SELECT input,action,result,lesson,reward FROM exp "
            "WHERE input LIKE ? ORDER BY reward DESC LIMIT ?",
            (f"%{query[:40]}%", n),
        )
        rows = cur.fetchall()
        conn.close()
        return [{"input": r[0], "action": r[1], "result": r[2],
                 "lesson": r[3], "reward": r[4]} for r in rows]

    # ── Stats ─────────────────────────────────────────────────
    def stats(self) -> Dict:
        conn = sqlite3.connect(self.db_path)
        c, r = conn.execute("SELECT COUNT(*), AVG(reward) FROM exp").fetchone()
        conn.close()
        return {"count": c or 0, "avg_reward": round(r or 0, 3),
                "replay": len(self.replay_buffer)}