Spaces:
Running
Running
| """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 | |