Blame for engine/saveload.py

"""Validated, session-scoped save files. Imports never supply a storage path."""
import json
import re
import tempfile
from datetime import datetime, timezone
from engine.memory import load_conversation, save_conversation, validate_character_id
from engine.persona import PERSONA_FIELDS
from engine.memory_controls import upgrade_conversation
from engine.relationship import get_default_relationship, validate_relationship

VERSION = '3.0'
MAX_SAVE_BYTES = 4 * 1024 * 1024


def validate_text(value, name, limit):
    if not isinstance(value, str) or len(value) > limit:
        raise ValueError(f'{name} must be text of at most {limit} characters.')
    return value


def validate_memories(value):
    if not isinstance(value, list) or len(value) > 100:
        raise ValueError('Invalid long-term memory list.')
    result = []
    for item in value:
        if not isinstance(item, dict) or item.get('importance') not in ('high', 'medium', 'low'):
            raise ValueError('Invalid memory importance.')
        content = validate_text(item.get('content'), 'memory', 2000).strip()
        if not content:
            raise ValueError('Memory content cannot be empty.')
        result.append({'content': content, 'importance': item['importance']})
    return result


def validate_save(data):
    if not isinstance(data, dict) or data.get('version') not in ('1.0', '2.0', VERSION):
        raise ValueError('Unsupported or missing save version.')
    char_id = validate_character_id(data.get('character_id'))
    persona = data.get('persona', {})
    if not isinstance(persona, dict):
        raise ValueError('Persona must be an object.')
    persona = {key: validate_text(persona.get(key, ''), key, 4000) for key in PERSONA_FIELDS}
    conv = data.get('conversation')
    if not isinstance(conv, dict):
        raise ValueError('Conversation must be an object.')
    messages = conv.get('recent_messages')
    if not isinstance(messages, list) or len(messages) > 10000:
        raise ValueError('Invalid message list.')
    clean = []
    for item in messages:
        if not isinstance(item, dict) or item.get('role') not in ('user', 'assistant'):
            raise ValueError('Invalid message role.')
        clean.append({'role': item['role'], 'content': validate_text(item.get('content'), 'message', 32000)})
    count = conv.get('summarized_count', 0) if data['version'] != '1.0' else 0
    if type(count) is not int or not 0 <= count <= len(clean):
        raise ValueError('Invalid summary position.')
    summary = validate_text(conv.get('summary', ''), 'summary', 8000)
    if count and not summary.strip():
        raise ValueError('Summarized messages require a summary.')
    memories = validate_memories(conv.get('long_term_memory', []))
    result = {'recent_messages': clean, 'summary': summary, 'long_term_memory': memories,
              'summarized_count': count, 'revision': 0}
    upgrade_conversation(result)
    if data['version'] == VERSION:
        for item, raw in zip(memories, conv.get('long_term_memory', [])):
            identifier = raw.get('id')
            origin = raw.get('origin')
            start, end = raw.get('start'), raw.get('end')
            if not isinstance(identifier, str) or not re.fullmatch(r'[0-9a-f]{32}', identifier):
                raise ValueError('Invalid memory identifier.')
            if origin not in ('auto', 'manual'):
                raise ValueError('Invalid memory origin.')
            if type(start) is not int or type(end) is not int or not 0 <= start <= end <= len(clean):
                raise ValueError('Invalid memory source.')
            item.update(id=identifier, origin=origin, start=start, end=end)
        if len({m['id'] for m in memories}) != len(memories) or sum(m['origin'] == 'manual' for m in memories) > 60:
            raise ValueError('Duplicate memory IDs or too many personal memories.')
        excluded = conv.get('excluded_sources', [])
        if not isinstance(excluded, list) or len(excluded) > 10000:
            raise ValueError('Invalid source exclusions.')
        for span in excluded:
            if not isinstance(span, list) or len(span) != 2 or any(type(x) is not int for x in span) or not 0 <= span[0] < span[1] <= len(clean):
                raise ValueError('Invalid source exclusion.')
        forgotten = conv.get('forgotten_facts', [])
        if not isinstance(forgotten, list) or len(forgotten) > 1000 or any(not isinstance(x, str) or not re.fullmatch(r'[0-9a-f]{64}', x) for x in forgotten):
            raise ValueError('Invalid forgotten-fact list.')
        result.update(excluded_sources=excluded, forgotten_facts=forgotten,
                      relationship=validate_relationship(conv.get('relationship', get_default_relationship()), len(clean)))
    else:
        for item in memories:
            item['end'] = count or len(clean)
    return {'version': VERSION, 'character_id': char_id, 'persona': persona, 'conversation': result}


def export_save(character_id, persona, session_id):
    data = validate_save({'version': VERSION, 'character_id': character_id,
                          'persona': persona or {}, 'conversation': load_conversation(session_id, character_id)})
    data['saved_at'] = datetime.now(timezone.utc).isoformat()
    payload = json.dumps(data, ensure_ascii=False, indent=2).encode('utf-8')
    if len(payload) > MAX_SAVE_BYTES:
        raise ValueError('Conversation exceeds the 4 MB save limit.')
    with tempfile.NamedTemporaryFile(prefix=character_id + '_', suffix='_save.json', delete=False) as handle:
        handle.write(payload)
        return handle.name


def import_save(file_path):
    with open(file_path, 'rb') as handle:
        payload = handle.read(MAX_SAVE_BYTES + 1)
    if len(payload) > MAX_SAVE_BYTES:
        raise ValueError('Save file exceeds 4 MB.')
    try:
        return validate_save(json.loads(payload))
    except (UnicodeDecodeError, json.JSONDecodeError, RecursionError) as exc:
        raise ValueError('Save file is not valid JSON.') from exc


def restore_save(data, session_id):
    validated = validate_save(data)
    save_conversation(session_id, validated['character_id'], validated['conversation'])
    return validated