cosmo-voice-assistant/cosmo/memory.py

"""
Персистентная память ассистента.
Хранится в data/memory.json — LLM читает и пишет её через инструменты.

Структура:
{
  "facts": {
    "user.name": "Даниил",
    "app.webstorm": "C:/Program Files/JetBrains/WebStorm.../webstorm64.exe",
    "user.browser": "chrome",
    ...
  },
  "history": [
    {"role": "user", "content": "..."},
    {"role": "assistant", "content": "..."}
  ]
}
"""

import json
import os
import threading
from datetime import datetime
from loguru import logger

MEMORY_PATH = os.path.join(
    os.path.dirname(os.path.dirname(os.path.abspath(__file__))),
    "data", "memory.json"
)


class Memory:
    def __init__(self, path: str = MEMORY_PATH, history_limit: int = 20):
        self.path = path
        self.history_limit = history_limit
        self._lock = threading.Lock()
        self._data = {"facts": {}, "history": []}
        self._load()

    # ------------------------------------------------------------------
    # Персистентность
    # ------------------------------------------------------------------

    def _load(self):
        os.makedirs(os.path.dirname(self.path), exist_ok=True)
        if os.path.exists(self.path):
            try:
                with open(self.path, "r", encoding="utf-8") as f:
                    self._data = json.load(f)
                logger.info(
                    f"Память загружена: {len(self._data.get('facts', {}))} фактов, "
                    f"{len(self._data.get('history', []))} сообщений в истории"
                )
            except Exception as e:
                logger.warning(f"Не удалось загрузить память: {e}. Начинаю с чистой.")
                self._data = {"facts": {}, "history": []}
        else:
            logger.info("Файл памяти не найден — создаю новый")

    def _save(self):
        try:
            with open(self.path, "w", encoding="utf-8") as f:
                json.dump(self._data, f, ensure_ascii=False, indent=2)
        except Exception as e:
            logger.error(f"Не удалось сохранить память: {e}")

    # ------------------------------------------------------------------
    # Факты (key-value долгосрочная память)
    # ------------------------------------------------------------------

    def get(self, key: str) -> str | None:
        """Получить факт по ключу. Возвращает None если не найден."""
        with self._lock:
            return self._data["facts"].get(key)

    def set(self, key: str, value: str):
        """Сохранить факт. Перезаписывает если уже существует."""
        with self._lock:
            self._data["facts"][key] = value
            self._save()
        logger.debug(f"Память: сохранено [{key}] = {value!r}")

    def delete(self, key: str) -> bool:
        """Удалить факт. Возвращает True если был."""
        with self._lock:
            existed = key in self._data["facts"]
            if existed:
                del self._data["facts"][key]
                self._save()
        return existed

    def list_facts(self, prefix: str = "") -> dict:
        """Вернуть все факты, опционально отфильтрованные по префиксу."""
        with self._lock:
            facts = self._data["facts"]
            if prefix:
                return {k: v for k, v in facts.items() if k.startswith(prefix)}
            return dict(facts)

    def facts_as_text(self) -> str:
        """Все факты в виде читаемого текста для системного промпта."""
        facts = self.list_facts()
        if not facts:
            return "Память пуста."
        lines = [f"  {k}: {v}" for k, v in sorted(facts.items())]
        return "\n".join(lines)

    # ------------------------------------------------------------------
    # История разговора (краткосрочная, последние N сообщений)
    # ------------------------------------------------------------------

    def add_message(self, role: str, content: str):
        """Добавить сообщение в историю (role: user/assistant/tool)."""
        with self._lock:
            self._data["history"].append({
                "role": role,
                "content": content,
                "ts": datetime.now().isoformat(timespec="seconds"),
            })
            # Ограничиваем длину истории
            if len(self._data["history"]) > self.history_limit:
                self._data["history"] = self._data["history"][-self.history_limit:]
            self._save()

    def get_history(self) -> list[dict]:
        """Вернуть историю в формате для LLM API (без поля ts)."""
        with self._lock:
            return [
                {"role": m["role"], "content": m["content"]}
                for m in self._data["history"]
            ]

    def clear_history(self):
        with self._lock:
            self._data["history"] = []
            self._save()