import copy import io import unittest from pypdf import PdfWriter from pypdf.generic import DictionaryObject, NameObject, DecodedStreamObject from context_manager import (index_document, load_document, prepare_request, prompt_tokens, rank_excerpts, estimate_tokens, MAX_DOCUMENT_BYTES) def synthetic_pdf(): writer = PdfWriter() for number in range(1, 13): page = writer.add_blank_page(width=600, height=800) font = DictionaryObject({NameObject('/Type'): NameObject('/Font'), NameObject('/Subtype'): NameObject('/Type1'), NameObject('/BaseFont'): NameObject('/Helvetica')}) page[NameObject('/Resources')] = DictionaryObject({NameObject('/Font'): DictionaryObject({NameObject('/F1'): font})}) text = ('Routine warehouse inspection and inventory notes. ' * 70 if number < 12 else 'Project Aurora release code is ZX731. Release date is 17 November 2027.') stream = DecodedStreamObject() stream.set_data(f'BT /F1 10 Tf 30 700 Td ({text}) Tj ET'.encode()) page[NameObject('/Contents')] = writer._add_object(stream) result = io.BytesIO() writer.write(result) return result.getvalue() class ContextTests(unittest.TestCase): def setUp(self): self.model = {'context_window': 4096, 'max_completion_tokens': 2048, 'rate_limits': {'tpm': 6000}} def test_real_pdf_last_page_is_retrieved_instead_of_initial_prefix(self): document = load_document(synthetic_pdf(), 'synthetic.pdf') self.assertEqual(len(document['pages']), 12) self.assertGreater(estimate_tokens(' '.join(document['pages'][:11])), 1200) ranked = rank_excerpts(document, 'What is Project Aurora release code?') self.assertEqual(ranked[0]['page'], 12) request = prepare_request([], 'What is Project Aurora release code?', self.model, 1024, document) self.assertIn('ZX731', request['messages'][0]['content']) self.assertEqual(request['sources'][0]['location'], 'page 12') self.assertFalse(request['no_document_match']) self.assertLessEqual(request['prompt_tokens'] + request['max_tokens'] + request['margin'], 4096) def test_long_history_is_trimmed_as_complete_turns_without_mutation(self): history = [] for i in range(100): history += [{'role': 'user', 'content': f'Question {i}: ' + 'example ' * 80}, {'role': 'assistant', 'content': f'Answer {i}: ' + 'response ' * 80}] original = copy.deepcopy(history) request = prepare_request(history, 'What did we discuss most recently?', self.model, 1024) self.assertGreater(request['omitted_messages'], 0) self.assertEqual(request['omitted_messages'] % 2, 0) self.assertEqual(request['messages'][0]['role'], 'user') self.assertIn('Answer 99', request['messages'][-2]['content']) self.assertEqual(history, original) self.assertLessEqual(prompt_tokens(request['messages']) + request['max_tokens'] + request['margin'], 4096) def test_tpm_and_output_and_context_limits_are_independent(self): for context, tpm, output_cap, requested in [(8192, 2000, 4000, 3000), (1024, 8000, 100, 500), (4096, None, 2048, 9000)]: model = {'context_window': context, 'max_completion_tokens': output_cap, 'rate_limits': {'tpm': tpm}} request = prepare_request([], 'Hello', model, requested) self.assertLessEqual(request['max_tokens'], output_cap) self.assertLessEqual(request['max_tokens'], requested) self.assertLessEqual(request['prompt_tokens'] + request['max_tokens'] + request['margin'], min(context, tpm or context)) def test_oversized_question_is_rejected_without_truncation(self): with self.assertRaisesRegex(ValueError, 'too large'): prepare_request([], 'question ' * 10000, self.model, 512) def test_image_allowance_is_counted_and_payload_preserved(self): content = [{'type': 'text', 'text': 'Describe this image'}, {'type': 'image_url', 'image_url': {'url': 'data:image/png;base64,abc'}}] model = {'context_window': 16000, 'max_completion_tokens': 2000} request = prepare_request([], content, model, 512) self.assertGreater(request['prompt_tokens'], 4096) self.assertEqual(request['messages'][-1]['content'], content) with self.assertRaises(ValueError): prepare_request([], content, self.model, 512) def test_no_match_is_explicit_and_does_not_inject_irrelevant_prefix(self): doc = index_document(['Warehouse inventory records.'], 'inventory.txt') request = prepare_request([], 'Meteorite tungsten isotope?', self.model, 512, doc) self.assertTrue(request['no_document_match']) self.assertEqual(request['sources'], []) def test_txt_and_markdown_keep_end_and_use_section_labels(self): for name in ['notes.txt', 'notes.md']: data = ('Routine inventory. ' * 10000 + '\nAurora release code ZX731').encode() doc = load_document(data, name) self.assertIn('ZX731', doc['pages'][0]) request = prepare_request([], 'Aurora release code?', self.model, 512, doc) self.assertTrue(any('ZX731' in s['text'] for s in request['sources'])) self.assertTrue(request['sources'][0]['location'].startswith('section')) def test_followup_can_retrieve_subject_from_previous_question(self): doc = index_document(['Aurora release code ZX731 and delivery on Tuesday.'], 'notes.txt') request = prepare_request([{'role': 'user', 'content': 'Tell me about Aurora release.'}], 'And when?', self.model, 512, doc) self.assertFalse(request['no_document_match']) def test_blank_and_oversized_documents_are_rejected(self): with self.assertRaisesRegex(ValueError, 'No readable text'): load_document(b'', 'empty.txt') with self.assertRaisesRegex(ValueError, '10 MB'): load_document(b'x' * (MAX_DOCUMENT_BYTES + 1), 'large.txt') def test_missing_metadata_uses_bounded_fallback(self): request = prepare_request([], 'Hello', {}, 100000) self.assertLessEqual(request['max_tokens'], 4000) self.assertLessEqual(request['prompt_tokens'] + request['max_tokens'] + request['margin'], 8192) if __name__ == '__main__': unittest.main()