J.A.R.V.I.S-py/memory_store.py

"""Long-term memory store — port of `crates/jarvis-core/src/long_term_memory.rs`.

Stores user facts in JSON at `<script_dir>/long_term_memory.json`.
Atomic write-through (write-then-rename).

Public API:
    remember(key, value) -> None
    recall(key) -> str | None
    forget(key) -> bool
    search(query, limit=5) -> list[dict]
    all_facts() -> list[dict]
    build_llm_context(prompt, limit=5) -> str   # used by chat handler
"""

import json
import os
import time

# JSON file lives in the same folder as main.py so the user can move the
# install and the data follows.
_HERE = os.path.dirname(os.path.abspath(__file__))
_PATH = os.path.join(_HERE, 'long_term_memory.json')

_store: dict[str, dict] = {}
_loaded = False


def _load():
    global _store, _loaded
    if _loaded:
        return
    if os.path.isfile(_PATH):
        try:
            with open(_PATH, encoding='utf-8') as f:
                data = json.load(f)
            _store = data.get('entries', {}) if isinstance(data, dict) else {}
        except (OSError, json.JSONDecodeError) as exc:
            print(f"[memory] corrupt store ({exc}) — starting empty")
            _store = {}
    _loaded = True


def _save():
    tmp = _PATH + '.tmp'
    try:
        with open(tmp, 'w', encoding='utf-8') as f:
            json.dump({'entries': _store}, f, ensure_ascii=False, indent=2)
        os.replace(tmp, _PATH)
    except OSError as exc:
        print(f"[memory] save failed: {exc}")


def _normalize_key(k: str) -> str:
    return (k or '').strip().lower()


def remember(key: str, value: str) -> None:
    _load()
    nk = _normalize_key(key)
    if not nk:
        return
    now = int(time.time())
    if nk in _store:
        _store[nk]['value'] = value
        _store[nk]['last_used_at'] = now
    else:
        _store[nk] = {
            'key': nk,
            'value': value,
            'created_at': now,
            'last_used_at': now,
            'use_count': 0,
        }
    _save()


def recall(key: str) -> str | None:
    _load()
    nk = _normalize_key(key)
    entry = _store.get(nk)
    if not entry:
        return None
    entry['last_used_at'] = int(time.time())
    entry['use_count'] = entry.get('use_count', 0) + 1
    _save()
    return entry['value']


def forget(key: str) -> bool:
    _load()
    nk = _normalize_key(key)
    if nk in _store:
        del _store[nk]
        _save()
        return True
    return False


def clear_all() -> int:
    """Wipe every fact. Returns count of removed entries. Atomic via _save."""
    _load()
    n = len(_store)
    if n == 0:
        return 0
    _store.clear()
    _save()
    return n


def search(query: str, limit: int = 5) -> list[dict]:
    _load()
    nq = _normalize_key(query)
    if not nq:
        return []
    hits = [
        e for e in _store.values()
        if nq in e['key'] or nq in e['value'].lower()
    ]
    hits.sort(key=lambda e: e.get('last_used_at', 0), reverse=True)
    return hits[: max(1, limit)]


def all_facts() -> list[dict]:
    _load()
    return list(_store.values())


def build_llm_context(prompt: str, limit: int = 5) -> str:
    """Format relevant memory facts as a system-prompt addendum."""
    hits = search(prompt, limit)
    if not hits:
        return ''
    lines = ["Известные факты о пользователе (используй если уместно):"]
    for h in hits:
        lines.append(f"- {h['key']} = {h['value']}")
    return '\n'.join(lines) + '\n'