core/memory.py

"""長期記憶 (JSON ファイル) の読み書き."""
from __future__ import annotations

import copy
import json
from datetime import datetime
from pathlib import Path
from typing import Any

import session_paths as sp

# パスはセッションごとに切り替わるため、定数ではなく関数で取得する。
ROOT = sp.ROOT

BATCH_SIZE = 20

DEFAULT_USER_MEMORY: dict[str, Any] = {
  "user_profile": {
    "user_name": "",
    "nickname": "",
    "goal": ""
  },
  "ai_persona": {
    "name": "",
    "persona": "",
    "traits": {
      "BigFive_Openness": 50,
      "BigFive_Conscientiousness": 50,
      "BigFive_Extraversion": 50,
      "BigFive_Agreeableness": 50,
      "BigFive_Neuroticism": 50,
      "KiSS18_BasicSkill": 50,
      "N": 0,
      "E": 0,
      "O": 0,
      "A": 0,
      "C": 0
    },
    "persona_style_description": "",
    "personality_text": ""
  },
  "relationship_prompt": "",
  "relationship_state": {
    "affection": 0,
    "tension": 0.0,
    "comfort_trust": 0,
    "insecurity": 0,
    "initiative_affection": 0
  },
  "history_summary": ""
}

def _new_default_user_memory() -> dict[str, Any]:
    return copy.deepcopy(DEFAULT_USER_MEMORY)

def _deep_merge(default: dict[str, Any], loaded: dict[str, Any]) -> dict[str, Any]:
    """default に loaded を重ねる。古い JSON に足りないキーがあっても落ちないようにする。"""
    merged = copy.deepcopy(default)
    for key, value in loaded.items():
        if isinstance(value, dict) and isinstance(merged.get(key), dict):
            merged[key] = _deep_merge(merged[key], value)
        else:
            merged[key] = value
    return merged


# user_memory.json を初期化する（無ければデフォルトで作成、あれば読み込み）
def init_user_memory(overwrite: bool = False) -> dict[str, Any]:
    user_memory_path = sp.user_memory_path()
    if overwrite or not user_memory_path.exists():
        data = _new_default_user_memory()
        save_json(user_memory_path, data)
        return data
    loaded = load_json(user_memory_path, {})
    return _deep_merge(_new_default_user_memory(), loaded)
    
def load_json(path: Path, default: Any) -> Any:
    if not path.exists():
        return default
    try:
        return json.loads(path.read_text(encoding="utf-8"))
    except json.JSONDecodeError:
        return default

def save_json(path: Path, data: Any) -> None:
    path.parent.mkdir(parents=True, exist_ok=True)
    path.write_text(
        json.dumps(data, ensure_ascii=False, indent=2),
        encoding="utf-8",
    )

def _summary_to_text(summary: Any) -> str:
    if isinstance(summary, dict):
        return str(summary.get("recent_summary") or "")
    return str(summary or "")

#load memory from json files, if not exist, return default values
def load_memory() -> dict[str, Any]:
    memory = init_user_memory()
    memory["conversation_history"] = load_json(
        sp.conversation_path(),
        memory.get("conversation_history", []),
    )
    summary_data = load_json(sp.summary_path(), {})
    memory["history_summary"] = _summary_to_text(summary_data)

    return memory

def update_history_summary(memory: dict[str, Any], summary: str) -> None:
    memory["history_summary"] = summary
    save_json(
        sp.summary_path(),
        {
            "recent_summary": summary,
            "updated_at": datetime.now().isoformat(timespec="seconds"),
        },
    )

def create_summary_from_txt(
    memory: dict[str, Any],
    txt_path: Path,
    self_name: str,
    batch_size: int = BATCH_SIZE,
    max_lines: int = BATCH_SIZE,
    *,
    tokenizer: Any | None = None,
    model: Any | None = None,
    device: str | None = None,
) -> None:
    from data.Models.Text2History.TinySwallow import summarize_history_from_txt

    summary = summarize_history_from_txt(
        txt_path,
        self_name=self_name,
        batch_size=batch_size,
        max_lines=max_lines,
        tokenizer=tokenizer,
        model=model,
        device=device,
    )
    memory["history_summary"] = _summary_to_text(summary)
    update_history_summary(memory, memory["history_summary"])


def save_memory(memory: dict[str, Any]) -> None:
    save_json(sp.conversation_path(), memory.get("conversation_history", []))

    update_history_summary(
        memory,
        _summary_to_text(memory.get("history_summary", "")),
    )

    user_memory = {
        key: value
        for key, value in memory.items()
        if key not in {"conversation_history", "episodic_memory"}
    }
    user_memory["history_summary"] = _summary_to_text(memory.get("history_summary", ""))
    save_json(sp.user_memory_path(), user_memory)

def reset_memory() -> dict[str, Any]:
    """全メモリを初期値に戻して保存する."""
    memory = _new_default_user_memory()
    save_memory(memory)
    return memory

def get_conversation(memory: dict[str, Any], n: int = 6) -> list[dict[str, str]]:
    return memory.get("conversation_history", [])[-n:]

def get_history(memory: dict[str, Any]):
    return _summary_to_text(memory.get("history_summary", ""))

def add_message(
    memory: dict[str, Any],
    role: str,
    text: str,
    *,
    keep_last: int | None = None,
) -> None:
    if role not in {"user", "assistant"}:
        raise ValueError("role must be 'user' or 'assistant'")

    history = memory.setdefault("conversation_history", [])
    history.append({
        "role": role,
        "text": text,
        "timestamp": datetime.now().isoformat(timespec="seconds"),
    })

    if keep_last is not None:
        memory["conversation_history"] = history[-keep_last:]

#初期設定のhistory_summary.json
def create_summary(memory: dict[str, Any], batch_size: int = BATCH_SIZE):
    from data.Models.Text2History.TinySwallow import summarize_history_from_txt
    summary = summarize_history_from_txt(
    sp.data_dir(),
    self_name="Rayta",
    batch_size=batch_size,
    max_lines=200,   
    )
    memory["history_summary"] = summary

#history_summary.jsonをUpdate
def update_summary(
    memory: dict[str, Any],
    *,
    tokenizer,
    model,
    device,
    batch_size: int = BATCH_SIZE,
) -> None:
    history = memory.get("conversation_history", [])
    if len(history) < batch_size:
        return

    from data.Models.Text2History.TinySwallow import summarize

    previous_summary = _summary_to_text(memory.get("history_summary", ""))
    new_summary = summarize(previous_summary, history, tokenizer, model, device)

    update_history_summary(memory, new_summary)
    memory["conversation_history"] = []