|
|
@@ -0,0 +1,346 @@
|
|
|
+"""Synthetic Redis-only tests: no configuration or external connections."""
|
|
|
+
|
|
|
+import json
|
|
|
+import threading
|
|
|
+import unittest
|
|
|
+from concurrent.futures import ThreadPoolExecutor
|
|
|
+from unittest.mock import patch
|
|
|
+
|
|
|
+from langgraph.checkpoint.base import empty_checkpoint
|
|
|
+from langgraph.types import Command
|
|
|
+from redis.exceptions import ConnectionError
|
|
|
+
|
|
|
+from src.step3_question_answer.context import (
|
|
|
+ ContextConfig, ContextError, MIN_REUSE_BYTES, READ_CHECKPOINT, RedisContextSaver,
|
|
|
+)
|
|
|
+from src.step3_question_answer.graph import build_graph
|
|
|
+from src.step3_question_answer.tests.test_context import FakeRedis, fake_saver
|
|
|
+from src.step3_question_answer.timing import Timings, bind
|
|
|
+from src.step4_api.runtime import GraphRuntime
|
|
|
+
|
|
|
+
|
|
|
+def config(thread='synthetic-reuse', ns='', cid=None):
|
|
|
+ values = {'thread_id': thread, 'checkpoint_ns': ns}
|
|
|
+ if cid is not None:
|
|
|
+ values['checkpoint_id'] = cid
|
|
|
+ return {'configurable': values}
|
|
|
+
|
|
|
+
|
|
|
+def checkpoint(cid='001', text=None):
|
|
|
+ result = empty_checkpoint()
|
|
|
+ result['id'] = cid
|
|
|
+ result['channel_values'] = {
|
|
|
+ 'response': '仅限合成地区;有效期至合成日期;不含例外主体。',
|
|
|
+ 'recommendation': {
|
|
|
+ 'sources': [{'source_id': 'synthetic-source', 'path': '条件.例外',
|
|
|
+ 'id': 'synthetic-source',
|
|
|
+ 'content': ('合成原文条件、否定和尾部例外。' * 14000) if text is None else text}],
|
|
|
+ 'summary': {'text': '合成答复[依据](https://example.invalid/policy)'},
|
|
|
+ },
|
|
|
+ }
|
|
|
+ return result
|
|
|
+
|
|
|
+
|
|
|
+def measured_read(saver, cfg):
|
|
|
+ timing = Timings()
|
|
|
+ with bind(timing, 'session_result_load'):
|
|
|
+ result = saver.get_tuple(cfg)
|
|
|
+ return result, timing.snapshot()['measurements']
|
|
|
+
|
|
|
+
|
|
|
+def total(counts, name):
|
|
|
+ return counts[name + '.session_result_load']['total']
|
|
|
+
|
|
|
+
|
|
|
+class CheckpointReuseTests(unittest.TestCase):
|
|
|
+ def setUp(self):
|
|
|
+ self.saver = fake_saver()
|
|
|
+ self.cfg = config()
|
|
|
+ self.key = self.saver.key('synthetic-reuse')
|
|
|
+
|
|
|
+ def put(self, cid='001', cfg=None, text=None):
|
|
|
+ return self.saver.put(cfg or self.cfg, checkpoint(cid, text), {'step': 1}, {})
|
|
|
+
|
|
|
+ def test_large_body_wire_reduction_full_answer_and_fresh_pending(self):
|
|
|
+ with self.saver.reuse_written_checkpoint('synthetic-reuse'):
|
|
|
+ saved = self.put()
|
|
|
+ # New writes after checkpoint storage must never come from cache.
|
|
|
+ self.saver.put_writes(saved, [('answer', '完整合成回答'), ('__error__', '旧错误')], 'task')
|
|
|
+ self.saver.put_writes(saved, [('__error__', '新错误')], 'task')
|
|
|
+ cached, counts = measured_read(self.saver, self.cfg)
|
|
|
+ self.assertEqual(total(counts, 'redis_checkpoint_cache_hit'), 1)
|
|
|
+ self.assertIn(('task', '__error__', '新错误'), cached.pending_writes)
|
|
|
+ self.assertIn(('task', 'answer', '完整合成回答'), cached.pending_writes)
|
|
|
+ self.assertEqual(len(self.saver._reuse_scopes[self.key]), 1)
|
|
|
+ normal, full = measured_read(self.saver, self.cfg)
|
|
|
+ self.assertEqual(cached, normal)
|
|
|
+ self.assertEqual(self.saver._reuse_scopes, {})
|
|
|
+ self.assertLess(total(counts, 'redis_returned_bytes'), total(full, 'redis_returned_bytes') // 100)
|
|
|
+ self.assertEqual(total(full, 'redis_returned_bytes') - total(counts, 'redis_returned_bytes'),
|
|
|
+ total(counts, 'redis_checkpoint_cache_saved_bytes'))
|
|
|
+ self.assertEqual(total(counts, 'checkpoint_decoded_input_bytes'), total(full, 'checkpoint_decoded_input_bytes'))
|
|
|
+ self.assertNotIn('合成原文', json.dumps(counts, ensure_ascii=False))
|
|
|
+
|
|
|
+ def test_new_request_restart_and_old_storage_use_full_read(self):
|
|
|
+ with self.saver.reuse_written_checkpoint('synthetic-reuse'):
|
|
|
+ self.put()
|
|
|
+ original, counts = measured_read(self.saver, self.cfg)
|
|
|
+ self.assertEqual(total(counts, 'redis_checkpoint_cache_hit'), 1)
|
|
|
+ second = RedisContextSaver(ContextConfig('redis://localhost/0', 30), self.saver.client)
|
|
|
+ for saver in (self.saver, second):
|
|
|
+ with saver.reuse_written_checkpoint('synthetic-reuse'):
|
|
|
+ restored, counts = measured_read(saver, self.cfg)
|
|
|
+ self.assertEqual(restored, original)
|
|
|
+ self.assertEqual(total(counts, 'redis_checkpoint_cache_hit'), 0)
|
|
|
+ self.assertGreater(total(counts, 'redis_returned_bytes'), MIN_REUSE_BYTES)
|
|
|
+ self.assertTrue(all(json.loads(name)[0] == 'c' for name in self.saver.client.rows[self.key]))
|
|
|
+
|
|
|
+ def test_external_newer_checkpoint_and_same_id_rewrite_are_authoritative(self):
|
|
|
+ other = RedisContextSaver(ContextConfig('redis://localhost/0', 30), self.saver.client)
|
|
|
+ with self.saver.reuse_written_checkpoint('synthetic-reuse'):
|
|
|
+ self.put()
|
|
|
+ # Same field and same byte length, different content: length/ID alone
|
|
|
+ # would miss this update, the Redis fingerprint must reject reuse.
|
|
|
+ other.put(self.cfg, checkpoint('001', '合成新原文。' * 30000), {}, {})
|
|
|
+ state, counts = measured_read(self.saver, self.cfg)
|
|
|
+ self.assertEqual(total(counts, 'redis_checkpoint_cache_hit'), 0)
|
|
|
+ self.assertEqual(state.checkpoint['channel_values']['recommendation']['sources'][0]['content'], '合成新原文。' * 30000)
|
|
|
+ other.put(self.cfg, checkpoint('002'), {}, {})
|
|
|
+ state, counts = measured_read(self.saver, self.cfg)
|
|
|
+ self.assertEqual(state.checkpoint['id'], '002')
|
|
|
+ self.assertEqual(total(counts, 'redis_checkpoint_cache_hit'), 0)
|
|
|
+
|
|
|
+ def test_same_length_changed_body_corrupt_body_and_marker_string_fail_closed(self):
|
|
|
+ with self.saver.reuse_written_checkpoint('synthetic-reuse'):
|
|
|
+ self.put(text='甲' * 60000)
|
|
|
+ name = self.saver._field('c', '', '001')
|
|
|
+ before = self.saver.client.rows[self.key][name]
|
|
|
+ # Reuse exact metadata from the original to hold encoded length fixed.
|
|
|
+ value, metadata, parent = self.saver._decode(before)
|
|
|
+ value['channel_values']['recommendation']['sources'][0]['content'] = '乙' * 60000
|
|
|
+ replaced = self.saver._encode((value, metadata, parent))
|
|
|
+ self.assertEqual(len(before), len(replaced))
|
|
|
+ self.saver.client.rows[self.key][name] = replaced.encode()
|
|
|
+ changed, counts = measured_read(self.saver, self.cfg)
|
|
|
+ self.assertEqual(total(counts, 'redis_checkpoint_cache_hit'), 0)
|
|
|
+ self.assertEqual(changed.checkpoint['channel_values']['recommendation']['sources'][0]['content'], '乙' * 60000)
|
|
|
+ for invalid in ('1', '__reuse_written_checkpoint_v1__', 'invalid-json'):
|
|
|
+ self.saver.client.rows[self.key][name] = invalid
|
|
|
+ with self.assertRaisesRegex(ContextError, '^context_data_invalid$'):
|
|
|
+ self.saver.get_tuple(self.cfg)
|
|
|
+
|
|
|
+ def test_namespace_explicit_history_parent_and_latest_selection(self):
|
|
|
+ with self.saver.reuse_written_checkpoint('synthetic-reuse'):
|
|
|
+ first = self.put('001')
|
|
|
+ second = self.put('002', first)
|
|
|
+ current, counts = measured_read(self.saver, second)
|
|
|
+ self.assertEqual(current.parent_config, first)
|
|
|
+ self.assertEqual(total(counts, 'redis_checkpoint_cache_hit'), 1)
|
|
|
+ old, counts = measured_read(self.saver, first)
|
|
|
+ self.assertEqual(old.checkpoint['id'], '001')
|
|
|
+ self.assertEqual(total(counts, 'redis_checkpoint_cache_hit'), 0)
|
|
|
+ self.put('999', config(ns='other'))
|
|
|
+ main, counts = measured_read(self.saver, self.cfg)
|
|
|
+ self.assertEqual(main.checkpoint['id'], '002')
|
|
|
+ self.assertEqual(total(counts, 'redis_checkpoint_cache_hit'), 0)
|
|
|
+ branch, counts = measured_read(self.saver, config(ns='other'))
|
|
|
+ self.assertEqual(branch.checkpoint['id'], '999')
|
|
|
+ self.assertEqual(total(counts, 'redis_checkpoint_cache_hit'), 1)
|
|
|
+ self.assertIsNone(self.saver.get_tuple(config(ns='missing')))
|
|
|
+ self.assertIsNone(self.saver.get_tuple(config(cid='missing')))
|
|
|
+ self.assertEqual(len(list(self.saver.list(self.cfg))), 2)
|
|
|
+
|
|
|
+ def test_failure_no_stale_fallback_and_scope_cleanup(self):
|
|
|
+ with self.assertRaisesRegex(RuntimeError, 'synthetic-failure'):
|
|
|
+ with self.saver.reuse_written_checkpoint('synthetic-reuse'):
|
|
|
+ self.put()
|
|
|
+ with patch.object(self.saver.client, 'eval', side_effect=ConnectionError('private')):
|
|
|
+ with self.assertRaisesRegex(ContextError, '^redis_connection_failed$'):
|
|
|
+ self.put('002')
|
|
|
+ with self.assertRaisesRegex(ContextError, '^redis_connection_failed$'):
|
|
|
+ self.saver.get_tuple(self.cfg)
|
|
|
+ state, counts = measured_read(self.saver, self.cfg)
|
|
|
+ self.assertEqual(state.checkpoint['id'], '001')
|
|
|
+ self.assertEqual(total(counts, 'redis_checkpoint_cache_hit'), 1)
|
|
|
+ raise RuntimeError('synthetic-failure')
|
|
|
+ self.assertEqual(self.saver._reuse_scopes, {})
|
|
|
+
|
|
|
+ def test_expiry_deletion_and_reads_do_not_renew_ttl(self):
|
|
|
+ with self.saver.reuse_written_checkpoint('synthetic-reuse'):
|
|
|
+ self.put()
|
|
|
+ self.saver.client.now = 29
|
|
|
+ self.assertIsNotNone(self.saver.get_tuple(self.cfg))
|
|
|
+ self.assertEqual(self.saver.client.expiry[self.key], 30)
|
|
|
+ self.saver.client.now = 30
|
|
|
+ self.assertIsNone(self.saver.get_tuple(self.cfg))
|
|
|
+ self.put('002')
|
|
|
+ self.saver.delete_thread('synthetic-reuse')
|
|
|
+ self.assertIsNone(self.saver.get_tuple(self.cfg))
|
|
|
+
|
|
|
+ def test_small_body_clears_large_cached_body_and_close_cleans_scope(self):
|
|
|
+ with self.saver.reuse_written_checkpoint('synthetic-reuse'):
|
|
|
+ self.put()
|
|
|
+ self.put('002', text='很短')
|
|
|
+ _, counts = measured_read(self.saver, self.cfg)
|
|
|
+ self.assertEqual(total(counts, 'redis_checkpoint_cache_hit'), 0)
|
|
|
+ self.assertIsNone(self.saver._reuse_scopes[self.key][0].checkpoint)
|
|
|
+ self.saver.close()
|
|
|
+ self.assertEqual(self.saver._reuse_scopes, {})
|
|
|
+
|
|
|
+ def test_worker_thread_writes_visible_and_returned_state_is_not_shared(self):
|
|
|
+ with self.saver.reuse_written_checkpoint('synthetic-reuse'):
|
|
|
+ with ThreadPoolExecutor(1) as pool:
|
|
|
+ pool.submit(self.put).result()
|
|
|
+ first, counts = measured_read(self.saver, self.cfg)
|
|
|
+ self.assertEqual(total(counts, 'redis_checkpoint_cache_hit'), 1)
|
|
|
+ first.checkpoint['channel_values']['response'] = '客户端修改'
|
|
|
+ second = self.saver.get_tuple(self.cfg)
|
|
|
+ self.assertNotEqual(second.checkpoint['channel_values']['response'], '客户端修改')
|
|
|
+
|
|
|
+ def test_intermediate_writes_do_not_hash_bodies_until_read(self):
|
|
|
+ with self.saver.reuse_written_checkpoint('synthetic-reuse'):
|
|
|
+ with patch('src.step3_question_answer.context.hashlib.sha1', side_effect=AssertionError('unnecessary-hash')):
|
|
|
+ for cid in ('001', '002', '003'):
|
|
|
+ self.put(cid)
|
|
|
+ self.assertEqual(self.saver.get_tuple(self.cfg).checkpoint['id'], '003')
|
|
|
+
|
|
|
+ def test_parallel_sessions_overlap_fallback_and_two_agents_are_isolated(self):
|
|
|
+ barrier = threading.Barrier(2)
|
|
|
+ def request(thread):
|
|
|
+ with self.saver.reuse_written_checkpoint(thread):
|
|
|
+ self.put(cfg=config(thread), text=thread * 30000)
|
|
|
+ barrier.wait(5)
|
|
|
+ return measured_read(self.saver, config(thread))
|
|
|
+ with ThreadPoolExecutor(2) as pool:
|
|
|
+ result = list(pool.map(request, ('AAA', 'BBB')))
|
|
|
+ for thread, (state, counts) in zip(('AAA', 'BBB'), result):
|
|
|
+ self.assertEqual(state.checkpoint['channel_values']['recommendation']['sources'][0]['content'], thread * 30000)
|
|
|
+ self.assertEqual(total(counts, 'redis_checkpoint_cache_hit'), 1)
|
|
|
+ with self.saver.reuse_written_checkpoint('synthetic-reuse'):
|
|
|
+ self.put()
|
|
|
+ with self.saver.reuse_written_checkpoint('synthetic-reuse'):
|
|
|
+ self.put('002')
|
|
|
+ self.assertEqual(total(measured_read(self.saver, self.cfg)[1], 'redis_checkpoint_cache_hit'), 0)
|
|
|
+ self.put('003')
|
|
|
+ self.assertEqual(total(measured_read(self.saver, self.cfg)[1], 'redis_checkpoint_cache_hit'), 0)
|
|
|
+ from src.step4_api.talent import TalentContextSaver
|
|
|
+ talent = TalentContextSaver(ContextConfig('redis://localhost/0', 30), self.saver.client)
|
|
|
+ with talent.reuse_written_checkpoint('synthetic-reuse'):
|
|
|
+ talent.put(self.cfg, checkpoint('900'), {}, {})
|
|
|
+ self.assertEqual(talent.get_tuple(self.cfg).checkpoint['id'], '900')
|
|
|
+ self.assertEqual(self.saver.get_tuple(self.cfg).checkpoint['id'], '003')
|
|
|
+ self.assertEqual(self.saver._reuse_scopes, {})
|
|
|
+ self.assertEqual(talent._reuse_scopes, {})
|
|
|
+
|
|
|
+ def test_invalid_conditional_responses_do_not_hide_bad_state(self):
|
|
|
+ with self.saver.reuse_written_checkpoint('synthetic-reuse'):
|
|
|
+ self.put()
|
|
|
+ for response in ([1, 'bad', 1], [1, self.saver._field('c', '', '001'), None],
|
|
|
+ [1, self.saver._field('c', '', '001'), True]):
|
|
|
+ with patch.object(self.saver.client, 'eval', return_value=response):
|
|
|
+ with self.assertRaisesRegex(ContextError, '^context_data_invalid$'):
|
|
|
+ self.saver.get_tuple(self.cfg)
|
|
|
+ with patch.object(self.saver.client, 'eval', return_value=[1, self.saver._field('c', '', '001'), 1]):
|
|
|
+ with self.assertRaisesRegex(ContextError, '^context_data_invalid$'):
|
|
|
+ self.saver.get_tuple(self.cfg)
|
|
|
+
|
|
|
+ def test_bytes_transport_and_pending_updates_during_read(self):
|
|
|
+ with self.saver.reuse_written_checkpoint('synthetic-reuse'):
|
|
|
+ saved = self.put()
|
|
|
+ read = self.saver.client.eval
|
|
|
+ def server(script, *args):
|
|
|
+ if script == READ_CHECKPOINT:
|
|
|
+ self.saver.put_writes(saved, [('__interrupt__', {'value': '新补问'})], 'task')
|
|
|
+ result = read(script, *args)
|
|
|
+ return ([v.encode() if isinstance(v, str) else v for v in result]
|
|
|
+ if isinstance(result, list) else result)
|
|
|
+ with patch.object(self.saver.client, 'eval', side_effect=server):
|
|
|
+ state, counts = measured_read(self.saver, self.cfg)
|
|
|
+ self.assertEqual(total(counts, 'redis_checkpoint_cache_hit'), 1)
|
|
|
+ self.assertIn(('task', '__interrupt__', {'value': '新补问'}), state.pending_writes)
|
|
|
+
|
|
|
+
|
|
|
+class ReuseGraphTests(unittest.TestCase):
|
|
|
+ def test_runtime_failure_retry_and_stopped_requests_discard_cached_body(self):
|
|
|
+ from types import SimpleNamespace
|
|
|
+ saver = fake_saver()
|
|
|
+ fail = True
|
|
|
+ def plan(*_):
|
|
|
+ if fail:
|
|
|
+ raise RuntimeError('synthetic-plan-failure')
|
|
|
+ return {'queries': [], 'error': None}
|
|
|
+ def runtime():
|
|
|
+ return GraphRuntime(build_graph(None,
|
|
|
+ lambda _: {'conversation_kind': 'policy_overview', 'needs_company': False},
|
|
|
+ checkpointer=saver, retriever=SimpleNamespace(search=lambda *_: {'candidates': []}),
|
|
|
+ recommender=SimpleNamespace(plan=plan, answer=lambda *_: {'overview': '合成答复', 'items': [], 'sources': []})))
|
|
|
+ question = '合成问题' * 12000
|
|
|
+ with self.assertRaisesRegex(RuntimeError, 'synthetic-plan-failure'):
|
|
|
+ runtime().run('synthetic-reuse', question, lambda *_: None, threading.Event())
|
|
|
+ self.assertEqual(saver._reuse_scopes, {})
|
|
|
+ # Restore via a fresh graph as the existing /retry contract requires.
|
|
|
+ fail = False
|
|
|
+ rows = []
|
|
|
+ runtime().run('synthetic-reuse', '/retry', lambda e, d: rows.append((e, d)), threading.Event())
|
|
|
+ self.assertEqual(rows[-1], ('done', {'status': 'completed'}))
|
|
|
+ self.assertEqual(saver._reuse_scopes, {})
|
|
|
+ stopped = threading.Event()
|
|
|
+ def stop(event, data):
|
|
|
+ if event == 'progress' and data.get('stage') == 'analyze' and data.get('status') == 'completed':
|
|
|
+ stopped.set()
|
|
|
+ runtime().run('synthetic-stop', question, stop, stopped)
|
|
|
+ self.assertTrue(stopped.is_set())
|
|
|
+ self.assertEqual(saver._reuse_scopes, {})
|
|
|
+
|
|
|
+ def test_interrupt_resume_new_process_and_complete_state(self):
|
|
|
+ saver = fake_saver()
|
|
|
+ graph = build_graph(None, checkpointer=saver)
|
|
|
+ question = '合成人才培训需求' * 9000
|
|
|
+ with saver.reuse_written_checkpoint('synthetic-reuse'):
|
|
|
+ graph.invoke({'question': question}, config())
|
|
|
+ state = graph.get_state(config())
|
|
|
+ self.assertTrue(state.interrupts)
|
|
|
+ before = saver.get_tuple(config())
|
|
|
+ restored_saver = RedisContextSaver(ContextConfig('redis://localhost/0', 30), saver.client)
|
|
|
+ restored = build_graph(None, checkpointer=restored_saver)
|
|
|
+ self.assertEqual(before, restored_saver.get_tuple(config()))
|
|
|
+ with restored_saver.reuse_written_checkpoint('synthetic-reuse'):
|
|
|
+ restored.invoke(Command(resume={'needs_company': False}), config())
|
|
|
+ state = restored.get_state(config())
|
|
|
+ self.assertFalse(state.next)
|
|
|
+ self.assertFalse(state.interrupts)
|
|
|
+ self.assertEqual(state.values['original_question'], question)
|
|
|
+ cached = restored_saver.get_tuple(config())
|
|
|
+ self.assertEqual(cached, restored_saver.get_tuple(config()))
|
|
|
+
|
|
|
+ def test_api_full_recommendation_matches_unoptimized_final_result(self):
|
|
|
+ from types import SimpleNamespace
|
|
|
+ answer = checkpoint()['channel_values']['recommendation']
|
|
|
+ answer.update(items=[{'source_id': f'synthetic-card-{i}', 'title': f'合成政策{i}',
|
|
|
+ 'card': {'kind': 'policy', 'conditions': ['合成地区', '不含例外主体']}}
|
|
|
+ for i in (1, 2)], overview='完整合成结果', complete=True)
|
|
|
+ def runtime(saver):
|
|
|
+ return GraphRuntime(build_graph(None,
|
|
|
+ lambda _: {'conversation_kind': 'policy_overview', 'needs_company': False},
|
|
|
+ checkpointer=saver,
|
|
|
+ retriever=SimpleNamespace(search=lambda *_: {'candidates': []}),
|
|
|
+ recommender=SimpleNamespace(plan=lambda *_: {'queries': [], 'error': None}, answer=lambda *_: answer)))
|
|
|
+ results = []
|
|
|
+ for optimized in (False, True):
|
|
|
+ saver = fake_saver()
|
|
|
+ rows, timing = [], Timings()
|
|
|
+ # Disable only the optimization threshold for the control, keeping
|
|
|
+ # the same real graph, serialization and synthetic Redis adapter.
|
|
|
+ with patch('src.step3_question_answer.context.MIN_REUSE_BYTES', 65536 if optimized else 10**12), bind(timing):
|
|
|
+ runtime(saver).run('synthetic-reuse', '合成政策需求', lambda e, d: rows.append((e, d)), threading.Event())
|
|
|
+ results.append(next(data for event, data in rows if event == 'result'))
|
|
|
+ self.assertEqual(rows[-1], ('done', {'status': 'completed'}))
|
|
|
+ self.assertEqual(saver._reuse_scopes, {})
|
|
|
+ counts = timing.snapshot()['measurements']
|
|
|
+ self.assertEqual(total(counts, 'redis_checkpoint_cache_hit'), int(optimized))
|
|
|
+ self.assertEqual(results[0], results[1])
|
|
|
+ self.assertEqual(results[1]['recommendation']['sources'], answer['sources'])
|
|
|
+
|
|
|
+
|
|
|
+if __name__ == '__main__':
|
|
|
+ unittest.main()
|