Spaces:
Running
Running
| import copy | |
| import json | |
| from pathlib import Path | |
| import tempfile | |
| from types import SimpleNamespace as NS | |
| import unittest | |
| from unittest.mock import patch | |
| from engine import memory | |
| from engine.chat import generate_response | |
| from engine.saveload import export_save, import_save, restore_save, validate_save, MAX_SAVE_BYTES | |
| CHAR = 'hermione_granger' | |
| OTHER = 'jessica_jones' | |
| A, B = 'a' * 32, 'b' * 32 | |
| def answer(text): | |
| return NS(choices=[NS(message=NS(content=text))]) | |
| class Client: | |
| def __init__(self, summary=None, fail=False): | |
| self.calls = [] | |
| self.summary = summary or {'summary': 'Alex found a clue.', 'new_memories': [ | |
| {'content': 'The user is Alex.', 'importance': 'high'}]} | |
| self.fail = fail | |
| def chat_completion(self, messages, **kwargs): | |
| self.calls.append((copy.deepcopy(messages), kwargs)) | |
| if self.fail: | |
| raise RuntimeError('Network failure') | |
| if kwargs.get('stream'): | |
| return iter([NS(choices=[NS(delta=NS(content='Hello'))]), | |
| NS(choices=[NS(delta=NS(content=' there.'))])]) | |
| return answer(json.dumps(self.summary)) | |
| class ConversationFixture: | |
| def setUp(self): | |
| self.temp = tempfile.TemporaryDirectory() | |
| self.addCleanup(self.temp.cleanup) | |
| self.patch = patch.object(memory, 'MEMORY_DIR', Path(self.temp.name) / 'sessions') | |
| self.patch.start() | |
| self.addCleanup(self.patch.stop) | |
| def populated(self, session=A, char=CHAR, turns=12): | |
| state = memory.load_conversation(session, char) | |
| for i in range(turns): | |
| state['recent_messages'].extend([{'role': 'user', 'content': f'Question {i}'}, | |
| {'role': 'assistant', 'content': f'Answer {i}'}]) | |
| memory.save_conversation(session, char, state) | |
| return state | |
| def save_data(self): | |
| return {'version': '2.0', 'character_id': CHAR, 'persona': {'name': 'Alex'}, | |
| 'conversation': self.populated()} | |
| class ConversationTests(ConversationFixture, unittest.TestCase): | |
| def test_sessions_and_characters_are_isolated(self): | |
| self.populated() | |
| self.assertEqual(1, len(memory.load_conversation(B, CHAR)['recent_messages'])) | |
| self.assertEqual(1, len(memory.load_conversation(A, OTHER)['recent_messages'])) | |
| self.assertEqual(25, len(memory.load_conversation(A, CHAR)['recent_messages'])) | |
| def test_import_and_export_roundtrip_full_history_and_unique_files(self): | |
| data = self.save_data() | |
| restore_save(data, A) | |
| paths = [export_save(CHAR, data['persona'], A) for _ in range(2)] | |
| self.addCleanup(lambda: [Path(p).unlink(missing_ok=True) for p in paths]) | |
| self.assertNotEqual(*paths) | |
| imported = import_save(paths[0]) | |
| restore_save(imported, B) | |
| self.assertEqual(imported['conversation']['recent_messages'], memory.load_conversation(B, CHAR)['recent_messages']) | |
| self.assertEqual(25, len(imported['conversation']['recent_messages'])) | |
| def test_invalid_import_cannot_mutate_existing_state(self): | |
| original = self.populated() | |
| for field, value in [('recent_messages', 'bad'), ('summary', []), ('long_term_memory', {}), | |
| ('summarized_count', True), ('summarized_count', -1)]: | |
| data = {'version': '2.0', 'character_id': CHAR, 'conversation': copy.deepcopy(original)} | |
| data['conversation'][field] = value | |
| with self.subTest(field=field, value=value), self.assertRaises(ValueError): | |
| restore_save(data, A) | |
| self.assertEqual(original, memory.load_conversation(A, CHAR)) | |
| def test_traversal_unknown_character_version_and_persona_rejected(self): | |
| data = self.save_data() | |
| for key, values in {'character_id': ['../outside', '/tmp/outside', 'missing_character', None], | |
| 'version': ['0.1', None], 'persona': [[], 'bad']}.items(): | |
| for value in values: | |
| altered = copy.deepcopy(data) | |
| altered[key] = value | |
| with self.subTest(key=key, value=value), self.assertRaises(ValueError): | |
| restore_save(altered, A) | |
| self.assertFalse((Path(self.temp.name) / 'outside').exists()) | |
| def test_invalid_message_and_memory_types_rejected(self): | |
| data = self.save_data() | |
| for message in [{'role': 'system', 'content': 'x'}, {'role': 'user', 'content': []}, None]: | |
| data['conversation']['recent_messages'] = [message] | |
| with self.assertRaises(ValueError): | |
| validate_save(data) | |
| data = self.save_data() | |
| for memories in [[{'content': 'x', 'importance': 'critical'}], [None], [{'content': [], 'importance': 'high'}]]: | |
| data['conversation']['long_term_memory'] = memories | |
| with self.assertRaises(ValueError): | |
| validate_save(data) | |
| def test_legacy_save_is_read_without_loading_shared_files(self): | |
| data = self.save_data() | |
| data['version'] = '1.0' | |
| data['conversation'].pop('summarized_count') | |
| restore_save(data, A) | |
| self.assertEqual(0, memory.load_conversation(A, CHAR)['summarized_count']) | |
| def test_oversized_and_invalid_json_rejected(self): | |
| path = Path(self.temp.name) / 'upload.json' | |
| for raw in [b'x' * (MAX_SAVE_BYTES + 1), b'{', b'\xff']: | |
| path.write_bytes(raw) | |
| with self.assertRaises(ValueError): | |
| import_save(path) | |
| def test_summary_preserves_full_transcript_and_shortens_context(self): | |
| state = self.populated() | |
| client = Client() | |
| self.assertTrue(memory.update_memory(A, CHAR, state, client, 'test')) | |
| result = memory.load_conversation(A, CHAR) | |
| self.assertEqual(state['recent_messages'], result['recent_messages']) | |
| self.assertGreater(result['summarized_count'], 0) | |
| self.assertTrue(result['summary']) | |
| self.assertEqual('high', result['long_term_memory'][0]['importance']) | |
| client = Client() | |
| list(generate_response('Next', CHAR, client, A, result['recent_messages'])) | |
| messages = client.calls[0][0] | |
| self.assertIn(result['summary'], messages[0]['content']) | |
| self.assertEqual(result['recent_messages'][result['summarized_count']:], messages[1:-1]) | |
| def test_summary_failure_does_not_drop_messages(self): | |
| state = self.populated() | |
| for client in [Client(fail=True), Client(summary={'summary': '', 'new_memories': []})]: | |
| self.assertFalse(memory.update_memory(A, CHAR, state, client, 'test')) | |
| self.assertEqual(state, memory.load_conversation(A, CHAR)) | |
| def test_summary_stale_result_cannot_overwrite_reset(self): | |
| state = self.populated() | |
| memory.reset_conversation(A, CHAR) | |
| self.assertFalse(memory.update_memory(A, CHAR, state, Client(), 'test')) | |
| self.assertEqual(1, len(memory.load_conversation(A, CHAR)['recent_messages'])) | |
| def test_visible_history_controls_prompt_and_retry(self): | |
| old = self.populated(turns=2) | |
| history = old['recent_messages'][:-2] | |
| client = Client() | |
| output = list(generate_response('Replacement', CHAR, client, A, history)) | |
| self.assertEqual('Hello there.', output[-1]) | |
| self.assertEqual(history, client.calls[0][0][1:-1]) | |
| saved = memory.load_conversation(A, CHAR) | |
| self.assertEqual(history + [{'role': 'user', 'content': 'Replacement'}, | |
| {'role': 'assistant', 'content': output[-1]}], saved['recent_messages']) | |
| def test_translation_off_is_default_and_uses_one_call(self): | |
| client = Client() | |
| history = memory.load_conversation(A, CHAR)['recent_messages'] | |
| list(generate_response('Merhaba', CHAR, client, A, history, language='Turkish', translation_model='translator')) | |
| self.assertEqual(1, len(client.calls)) | |
| self.assertIn('Türkçe', client.calls[0][0][0]['content']) | |
| def test_translation_stores_visible_language(self): | |
| history = memory.load_conversation(A, CHAR)['recent_messages'] | |
| with patch('engine.chat.translate_to_english', return_value='Hello'), patch('engine.chat.translate_from_english', return_value='Merhaba!'): | |
| output = list(generate_response('Selam', CHAR, Client(), A, history, language='Turkish', | |
| translation_model='translator', use_translator=True)) | |
| state = memory.load_conversation(A, CHAR) | |
| self.assertEqual('Selam', state['recent_messages'][-2]['content']) | |
| self.assertEqual('Merhaba!', state['recent_messages'][-1]['content']) | |
| self.assertEqual('Merhaba!', output[-1]) | |
| def test_clear_discards_all_memory_layers(self): | |
| state = self.populated() | |
| memory.update_memory(A, CHAR, state, Client(), 'test') | |
| cleared = memory.sync_history(A, CHAR, []) | |
| self.assertEqual([], cleared['recent_messages']) | |
| self.assertEqual('', cleared['summary']) | |
| self.assertEqual([], cleared['long_term_memory']) | |
| def test_retry_keeps_unchanged_summary_but_editing_old_turn_resets_it(self): | |
| state = self.populated() | |
| memory.update_memory(A, CHAR, state, Client(), 'test') | |
| summarized = memory.load_conversation(A, CHAR) | |
| retried = memory.sync_history(A, CHAR, summarized['recent_messages'][:-2]) | |
| self.assertEqual(summarized['summary'], retried['summary']) | |
| self.assertEqual(summarized['long_term_memory'], retried['long_term_memory']) | |
| edited = copy.deepcopy(retried['recent_messages']) | |
| edited[1]['content'] = 'Different story' | |
| updated = memory.sync_history(A, CHAR, edited) | |
| self.assertEqual('', updated['summary']) | |
| self.assertEqual([], updated['long_term_memory']) | |
| self.assertEqual(0, updated['summarized_count']) | |
| def test_repeated_summarization_is_cumulative_and_deduplicates(self): | |
| state = self.populated(turns=50) | |
| first = Client() | |
| memory.update_memory(A, CHAR, state, first, 'test') | |
| snapshot = memory.load_conversation(A, CHAR) | |
| second = Client() | |
| memory.update_memory(A, CHAR, snapshot, second, 'test') | |
| payload = json.loads(second.calls[0][0][1]['content']) | |
| self.assertEqual(snapshot['summary'], payload['previous_summary']) | |
| self.assertEqual(snapshot['recent_messages'][snapshot['summarized_count']], payload['conversation'][0]) | |
| result = memory.load_conversation(A, CHAR) | |
| self.assertEqual(1, len(result['long_term_memory'])) | |
| self.assertEqual(101, len(result['recent_messages'])) | |
| self.assertGreater(result['summarized_count'], snapshot['summarized_count']) | |
| def test_atomic_write_failure_keeps_previous_snapshot(self): | |
| state = self.populated() | |
| modified = copy.deepcopy(state) | |
| modified['summary'] = 'Replacement' | |
| with patch('engine.memory.os.replace', side_effect=OSError('Disk failure')): | |
| with self.assertRaises(OSError): | |
| memory.save_conversation(A, CHAR, modified) | |
| self.assertEqual(state, memory.load_conversation(A, CHAR)) | |
| self.assertEqual([], list(Path(self.temp.name).rglob('*.tmp'))) | |
| def test_chat_reports_storage_failure_without_losing_previous_history(self): | |
| history = memory.load_conversation(A, CHAR)['recent_messages'] | |
| with patch('engine.chat.append_exchange', side_effect=OSError('Disk full')): | |
| output = list(generate_response('Hi', CHAR, Client(), A, history)) | |
| self.assertIn('could not be saved', output[-1]) | |
| self.assertEqual(history, memory.load_conversation(A, CHAR)['recent_messages']) | |
| def test_gradio_text_blocks_normalized(self): | |
| result = memory.normalize_history([{'role': 'user', 'content': [{'type': 'text', 'text': 'Hello'}]}]) | |
| self.assertEqual([{'role': 'user', 'content': 'Hello'}], result) | |
| if __name__ == '__main__': | |
| unittest.main() | |