Blame for engine/memory.py

"""Session-scoped conversation snapshots and bounded model context.

The complete visible transcript is retained; only summarized prefixes leave the
model context. Legacy shared character files are deliberately never read.
"""
import copy
import json
import logging
import os
from pathlib import Path
from engine.static_files import load_static_text
import re
import tempfile
import threading
from uuid import uuid4
from engine.memory_controls import upgrade_conversation, context_messages, fingerprint
from engine.relationship import get_default_relationship, apply_relationship_changes, relationship_values

logger = logging.getLogger(__name__)
MEMORY_DIR = Path(__file__).resolve().parent.parent / 'memory' / 'sessions'
MAX_RECENT = 20
SUMMARY_BATCH = 10
MAX_SUMMARY_BATCH = 40
MAX_CONTEXT_MESSAGES = 40
MAX_CONTEXT_CHARS = 24000
MAX_MEMORIES = 40
_LOCK = threading.RLock()


def _chat_completion(client, messages, **kwargs):
    """Adapt chat completion calls for HF InferenceClient or OpenAI HF Router."""
    model = kwargs.get('model')
    if isinstance(model, str) and model.endswith(':featherless-ai'):
        model = kwargs.pop('model')
        return client.chat.completions.create(model=model, messages=messages, **kwargs)
    return client.chat_completion(messages, **kwargs)


def validate_character_id(character_id):
    if not isinstance(character_id, str) or not re.fullmatch(r'[a-zA-Z0-9_-]{1,100}', character_id):
        raise ValueError('Invalid character ID.')
    from engine.prompt_builder import load_character
    try:
        load_character(character_id)
    except (FileNotFoundError, ValueError) as exc:
        raise ValueError('Unknown character ID.') from exc
    return character_id


def _path(session_id, character_id):
    if not isinstance(session_id, str) or not re.fullmatch(r'[0-9a-f]{32}', session_id):
        raise ValueError('Invalid session ID.')
    validate_character_id(character_id)
    return Path(MEMORY_DIR) / session_id / (character_id + '.json')


def empty_conversation(greeting=''):
    return {'recent_messages': ([{'role': 'assistant', 'content': greeting}] if greeting else []),
            'summary': '', 'long_term_memory': [], 'summarized_count': 0, 'revision': 0, 'memory_revision': 0,
            'relationship': get_default_relationship(), 'excluded_sources': [], 'forgotten_facts': []}


def load_conversation(session_id, character_id):
    with _LOCK:
        path = _path(session_id, character_id)
        if path.exists():
            return upgrade_conversation(json.loads(path.read_text(encoding='utf-8')))
        from engine.prompt_builder import load_character
        return empty_conversation(load_character(character_id).get('greeting', 'Hello.'))


def save_conversation(session_id, character_id, conversation, expected_revision=None, *, append_only=False):
    """Atomically replace all layers; reject stale generation/summary results."""
    with _LOCK:
        path = _path(session_id, character_id)
        current = load_conversation(session_id, character_id)
        if expected_revision is not None and current['revision'] != expected_revision:
            return False
        data = copy.deepcopy(conversation)
        data['revision'] = current['revision'] + 1
        data['memory_revision'] = current.get('memory_revision', 0) + (0 if append_only else 1)
        _write_snapshot(path, data)
        conversation['revision'] = data['revision']
        conversation['memory_revision'] = data['memory_revision']
        return True


def _write_snapshot(path, data):
    """Caller holds _LOCK; replacing the file makes every layer atomic."""
    path.parent.mkdir(parents=True, exist_ok=True)
    fd, temporary = tempfile.mkstemp(dir=path.parent, suffix='.tmp')
    try:
        with os.fdopen(fd, 'w', encoding='utf-8') as handle:
            json.dump(data, handle, ensure_ascii=False)
        os.replace(temporary, path)
    finally:
        if os.path.exists(temporary):
            os.unlink(temporary)


def append_exchange(session_id, character_id, snapshot, message, response):
    """Keep a summary completed during streaming, but reject edited/reset history."""
    with _LOCK:
        current = load_conversation(session_id, character_id)
        if current['revision'] != snapshot['revision']:
            return None
        current['recent_messages'].extend([
            {'role': 'user', 'content': message},
            {'role': 'assistant', 'content': response},
        ])
        save_conversation(session_id, character_id, current, snapshot['revision'], append_only=True)
        return current


def _commit_memory(session_id, character_id, snapshot, updated):
    """Derived memory does not change transcript revision or replace its messages.

    Allow appended turns while rejecting edits, resets, imports and memory changes.
    Only derived fields are merged; newer transcript messages are never replaced.
    """
    with _LOCK:
        current = load_conversation(session_id, character_id)
        if current.get('memory_revision', 0) != snapshot.get('memory_revision', 0):
            return False
        if current['recent_messages'][:len(snapshot['recent_messages'])] != snapshot['recent_messages']:
            return False
        fields = ('summary', 'summarized_count', 'long_term_memory', 'relationship')
        if any(current[key] != snapshot[key] for key in fields):
            return False
        current.update({key: updated[key] for key in fields})
        _write_snapshot(_path(session_id, character_id), current)
        return True


def normalize_history(history):
    """Accept Gradio 6 text blocks as well as simple message dictionaries."""
    result = []
    for message in history or []:
        role, content = message.get('role'), message.get('content')
        if role not in ('user', 'assistant'):
            raise ValueError('Unsupported message role.')
        if isinstance(content, list):
            if any(not isinstance(b, dict) or b.get('type') != 'text' for b in content):
                raise ValueError('Only text messages are supported.')
            content = ''.join(b.get('text', '') for b in content)
        if not isinstance(content, str):
            raise ValueError('Message content must be text.')
        result.append({'role': role, 'content': content})
    return result


def sync_history(session_id, character_id, history):
    with _LOCK:
        conversation = load_conversation(session_id, character_id)
        messages = normalize_history(history)
        if messages != conversation['recent_messages']:
            old = conversation['recent_messages']
            common = 0
            for before, after in zip(old, messages):
                if before != after:
                    break
                common += 1
            if not messages:
                revision = conversation['revision']
                conversation = empty_conversation()
                conversation['revision'] = revision
            else:
                if common < conversation['summarized_count']:
                    conversation.update(summary='', summarized_count=0)
                # Preserve manually saved facts, remove evidence from discarded turns.
                conversation['long_term_memory'] = [m for m in conversation['long_term_memory']
                    if m['origin'] == 'manual' or m['end'] <= common]
                conversation['relationship']['events'] = [e for e in conversation['relationship']['events']
                    if e['end'] <= common]
                conversation['excluded_sources'] = [[a, min(b, common)] for a, b in conversation['excluded_sources'] if a < common]
            conversation['recent_messages'] = messages
            if not save_conversation(session_id, character_id, conversation, conversation['revision']):
                raise ValueError('Conversation changed. Please try again.')
        return conversation


def bounded_recent_context(conversation):
    """Bound inference input even if extraction is slow or unavailable.

    The full transcript stays on disk and remains eligible for summarization.
    Never claim temporarily omitted messages have already been summarized.
    """
    start = conversation['summarized_count']
    total = len(conversation['recent_messages'])
    first = max(start, total - MAX_CONTEXT_MESSAGES)
    messages = context_messages(conversation, first)
    kept, remaining = [], MAX_CONTEXT_CHARS
    clipped = False
    for message in reversed(messages):
        content = message['content']
        if len(content) > remaining:
            if not kept:
                marker = '[Earlier part omitted for context size.]\n'
                kept.append(dict(message, content=marker + content[-(remaining - len(marker)):]))
                clipped = True
            break
        kept.append(message)
        remaining -= len(content)
    kept.reverse()
    omitted = len(kept) < total - start or clipped
    return kept, omitted


def reset_conversation(session_id, character_id):
    from engine.prompt_builder import load_character
    state = empty_conversation(load_character(character_id).get('greeting', 'Hello.'))
    save_conversation(session_id, character_id, state)
    return state


def update_memory(session_id, character_id, conversation, client, model, persona=None):
    """Summarize a batch after streaming; failure retains every unsummarized turn."""
    messages = conversation['recent_messages']
    start = conversation['summarized_count']
    if len(messages) - start <= MAX_RECENT:
        return False
    # Catch up in one bounded batch when the user chats faster than extraction.
    pending = len(messages) - start
    batch = SUMMARY_BATCH if pending <= MAX_RECENT + SUMMARY_BATCH else min(MAX_SUMMARY_BATCH, pending - 10)
    end = min(start + batch, len(messages) - 10)
    # Do not split a user/assistant pair between summary and recent context.
    if end < len(messages) and messages[end]['role'] == 'assistant':
        end += 1
    try:
        from engine.prompt_builder import load_character
        character = load_character(character_id)
        prompt = load_static_text(Path(__file__).resolve().parent.parent / 'prompts/memory_prompt.txt')
        messages = [{'role': 'system', 'content': prompt}, {'role': 'user', 'content': json.dumps({
                'previous_summary': conversation['summary'],
                'existing_memories': conversation['long_term_memory'],
                'character': {key: character.get(key) for key in ('name', 'personality', 'identity', 'world', 'system_prompt')},
                'user_persona': persona or {},
                'relationship': {key: value for key, value in relationship_values(conversation['relationship']).items() if value is not None},
                'message_offset': start,
                'conversation': context_messages(conversation, start, end)}, ensure_ascii=False)}]
        response = _chat_completion(
            client, messages, model=model, max_tokens=1600, temperature=0.1, stream=False)
        raw = response.choices[0].message.content.strip()
        if raw.startswith('```'):
            raw = raw.split('\n', 1)[1].rsplit('```', 1)[0].strip()
        data = json.loads(raw)
        from engine.saveload import validate_memories, validate_text
        summary = validate_text(data['summary'], 'summary', 8000)
        if not summary.strip():
            raise ValueError('Empty memory summary.')
        new_memories = validate_memories(data['new_memories'])
        for item, raw_item in zip(new_memories, data['new_memories']):
            source = raw_item.get('source_index')
            if source is not None and (type(source) is not int or not start <= source < end):
                raise ValueError('Invalid extracted memory source.')
            item.update(id=uuid4().hex, origin='auto', start=start if source is None else source,
                        end=end if source is None else source + 1)
        combined = {}
        # Existing manual entries win over automated extraction of the same fact.
        for item in new_memories + conversation['long_term_memory']:
            key = fingerprint(item['content'])
            excluded = item['origin'] == 'auto' and any(item['start'] < b and a < item['end']
                                                        for a, b in conversation['excluded_sources'])
            if key not in conversation['forgotten_facts'] and not excluded:
                combined[key] = item
        rank = {'high': 0, 'medium': 1, 'low': 2}
        updated = copy.deepcopy(conversation)
        manual = [m for m in combined.values() if m['origin'] == 'manual']
        automatic = sorted((m for m in combined.values() if m['origin'] != 'manual'), key=lambda m: rank[m['importance']])
        relationship = conversation['relationship']
        changes = data.get('relationship_changes', [])
        # A malformed optional observation must not stall all summary progress.
        try:
            relationship = apply_relationship_changes(relationship, changes, start, end)
        except ValueError:
            logger.warning('Invalid relationship observations; keeping existing relationship state.')
        updated.update(summary=summary, summarized_count=end, relationship=relationship,
                       long_term_memory=manual + automatic[:MAX_MEMORIES])
        return _commit_memory(session_id, character_id, conversation, updated)
    except Exception:
        logger.warning('Memory update failed; unsummarized messages retained.', exc_info=True)
        return False