162 lines
8.8 KiB
Python
162 lines
8.8 KiB
Python
"""Context overflow must split inputs, never silently discard their remainder."""
|
|
import json
|
|
import os
|
|
import re
|
|
import tempfile
|
|
import unittest
|
|
from pathlib import Path
|
|
from unittest.mock import patch
|
|
|
|
import requests
|
|
|
|
import test_pipeline as fixtures
|
|
import ollama_client as client
|
|
from pdf_text import _ocr_page
|
|
|
|
|
|
def echo_request(base, model, source, timeout, style, system, schema, image, max_tokens):
|
|
return {'title': 'Einladung', 'paragraphs': [source]}
|
|
|
|
|
|
class ContextTests(unittest.TestCase):
|
|
setUp = fixtures.PipelineTests.setUp
|
|
|
|
def test_character_limit_splits_in_order_and_keeps_every_word(self):
|
|
source = '\n\n'.join(f'Absatz {i}: ' + 'Einladung ' * 35 for i in range(9))
|
|
with patch.dict(os.environ, {'OLLAMA_MAX_INPUT_CHARS': '1000'}), \
|
|
patch.object(client, '_request', side_effect=echo_request) as request:
|
|
result = client.transform(source)
|
|
self.assertGreater(request.call_count, 1)
|
|
self.assertEqual(' '.join(result.paragraphs).split(), source.split())
|
|
self.assertTrue(all(len(call.args[2]) <= 1000 for call in request.call_args_list))
|
|
self.assertEqual(result.title, 'Einladung')
|
|
|
|
def test_byte_budget_protects_unicode_even_below_character_limit(self):
|
|
source = 'Überraschung 🎉 漢字 ' * 1800
|
|
sent = []
|
|
def answer(*args):
|
|
text, system = args[2], args[5]
|
|
self.assertLessEqual(len(text.encode()) + len(system.encode()) + args[-1] + 2048, 32768)
|
|
sent.append(text)
|
|
return {'title': 'Einladung', 'paragraphs': ['Ankündigung.']}
|
|
with patch.object(client, '_request', side_effect=answer):
|
|
client.transform(source)
|
|
self.assertGreater(len(sent), 1)
|
|
self.assertEqual(' '.join(sent[:-1]).split(), source.split())
|
|
self.assertTrue(sent[-1].startswith('Einladung')) # Smaller results get one final editorial pass.
|
|
|
|
def test_only_the_oversized_chunk_is_retried(self):
|
|
source = 'A' * 700 + '\n\n' + 'B' * 700 + '\n\n' + 'C' * 700 + '\n\n' + 'D' * 700
|
|
def answer(*args):
|
|
if len(args[2]) > 2000 or (args[2].startswith('C') and len(args[2]) > 1000):
|
|
raise client.ContextLimitError('Tokenlimit')
|
|
return echo_request(*args)
|
|
with patch.object(client, '_request', side_effect=answer) as request:
|
|
result = client.transform(source)
|
|
self.assertEqual(' '.join(result.paragraphs).split(), source.split())
|
|
self.assertEqual([call.args[2] for call in request.call_args_list[:-1]],
|
|
[source, source[:1402], source[1404:], 'C' * 700, 'D' * 700])
|
|
self.assertTrue(request.call_args_list[-1].args[2].startswith('Einladung'))
|
|
|
|
def test_permanent_errors_do_not_retry_or_return_partial_content(self):
|
|
source = 'A' * 700 + '\n\n' + 'B' * 700
|
|
for error in (client.OllamaError('authentication'), client.ContentMismatchError('Fakten')):
|
|
with self.subTest(error=error), patch.object(client, '_request', side_effect=error) as request:
|
|
with self.assertRaises(type(error)):
|
|
client.transform(source)
|
|
self.assertEqual(request.call_count, 1)
|
|
with patch.dict(os.environ, {'OLLAMA_MAX_INPUT_CHARS': '1000'}), patch.object(client, '_request',
|
|
side_effect=[{'title': 'Einladung', 'paragraphs': ['Erster Teil.']}, client.OllamaError('offline')]):
|
|
with self.assertRaisesRegex(client.OllamaError, 'offline'):
|
|
client.transform(source)
|
|
|
|
def test_context_retry_has_a_bound(self):
|
|
with patch.object(client, '_request', side_effect=client.ContextLimitError('Tokenlimit')) as request:
|
|
with self.assertRaisesRegex(client.ContextLimitError, 'Tokenlimit'):
|
|
client.transform('Text ' * 800)
|
|
self.assertLessEqual(request.call_count, 4)
|
|
|
|
def test_facts_moved_into_section_titles_are_not_discarded(self):
|
|
source = 'Fest im Garten. ' * 45 + '\n\n' + 'Abschluss im Saal. ' * 40
|
|
with patch.dict(os.environ, {'OLLAMA_MAX_INPUT_CHARS': '1000'}), patch.object(client, '_request',
|
|
side_effect=[{'title': 'Fest im Garten', 'paragraphs': ['Herzlich willkommen.']},
|
|
{'title': 'Abschluss im Saal', 'paragraphs': ['Wir freuen uns auf Euch.']},
|
|
{'title': 'Fest im Garten', 'paragraphs': ['Abschluss im Saal']}]) as request:
|
|
result = client.transform(source)
|
|
self.assertEqual(result.title, 'Fest im Garten')
|
|
self.assertIn('Abschluss im Saal', result.paragraphs)
|
|
self.assertIn('Abschluss im Saal', request.call_args_list[-1].args[2])
|
|
|
|
def test_invalid_final_edit_is_not_published_as_successful_sections(self):
|
|
with patch.dict(os.environ, {'OLLAMA_MAX_INPUT_CHARS': '1000'}), patch.object(client, '_request',
|
|
side_effect=[{'title': 'Einladung', 'paragraphs': ['Beginn im Garten.']},
|
|
{'title': 'Einladung', 'paragraphs': ['Abschluss im Saal.']},
|
|
client.OllamaError('invalid final response')]):
|
|
with self.assertRaisesRegex(client.OllamaError, 'invalid final response'):
|
|
client.transform('A' * 700 + '\n\n' + 'B' * 700)
|
|
|
|
def test_only_size_related_http_errors_enable_splitting(self):
|
|
for status, message, expected in (
|
|
(400, "This model's maximum context length is 32768 tokens", client.ContextLimitError),
|
|
(422, 'input is too long', client.ContextLimitError),
|
|
(413, 'Request Entity Too Large', client.ContextLimitError),
|
|
(401, 'Unauthorized', client.OllamaError),
|
|
(429, 'Rate limit exceeded', client.OllamaError),
|
|
(500, 'out of memory', client.OllamaError),
|
|
(400, 'model does not support images', client.OllamaError),
|
|
):
|
|
response = requests.Response()
|
|
response.status_code = status
|
|
response._content = json.dumps({'error': {'message': message}}).encode()
|
|
self.assertIs(type(client._gateway_error(response)), expected)
|
|
|
|
def test_invalid_context_fails_before_any_request(self):
|
|
for value in ('oops', '4096', '999999'):
|
|
with self.subTest(value=value), patch.dict(os.environ, {'OLLAMA_CONTEXT_LENGTH': value}), \
|
|
patch.object(client, '_request') as request:
|
|
with self.assertRaises(client.OllamaError):
|
|
client.transform('Einladung zum Fest.')
|
|
request.assert_not_called()
|
|
|
|
def test_image_limit_splits_with_overlap_and_keeps_bottom(self):
|
|
from pdf_text import run
|
|
with tempfile.TemporaryDirectory() as temp:
|
|
pdf = Path(temp) / 'page.pdf'
|
|
fixtures.make_text_pdf(pdf)
|
|
for rotated in (False, True):
|
|
def metadata(args, **kwargs):
|
|
info = run(args, **kwargs)
|
|
# Landscape MediaBox rotated by 90° must use its full width as height.
|
|
if args[0] == 'pdfinfo' and rotated:
|
|
info = re.sub(r'(Page\s+1\s+MediaBox:).*', r'\1 0 0 842 595', info)
|
|
info = re.sub(r'(Page\s+1\s+rot:)\s+0', r'\1 90', info)
|
|
return info
|
|
with self.subTest(rotated=rotated), patch('pdf_text.run', side_effect=metadata) as commands, \
|
|
patch('pdf_text.transcribe_image', side_effect=[
|
|
client.ContextLimitError('Tokenlimit'), 'Anfang\nÜbergang', 'Übergang\nEnde']):
|
|
result = _ocr_page(pdf, 1)
|
|
crops = [call.args[0] for call in commands.call_args_list if call.args[0][0] == 'pdftoppm']
|
|
self.assertNotIn('-y', crops[0])
|
|
top = int(crops[1][crops[1].index('-y') + 1])
|
|
height = int(crops[1][crops[1].index('-H') + 1])
|
|
bottom_top = int(crops[2][crops[2].index('-y') + 1])
|
|
bottom_height = int(crops[2][crops[2].index('-H') + 1])
|
|
self.assertEqual(top, 0)
|
|
self.assertEqual(height - bottom_top, 96)
|
|
self.assertGreater(bottom_top + bottom_height, 2300)
|
|
self.assertEqual(result, 'Anfang\nÜbergang\n\nÜbergang\nEnde')
|
|
|
|
def test_image_retry_does_not_retry_network_errors(self):
|
|
with tempfile.TemporaryDirectory() as temp:
|
|
pdf = Path(temp) / 'page.pdf'
|
|
fixtures.make_text_pdf(pdf)
|
|
for error, max_calls in ((client.OllamaError('offline'), 1), (client.ContextLimitError('Tokenlimit'), 4)):
|
|
with self.subTest(error=error), patch('pdf_text.transcribe_image', side_effect=error) as vision:
|
|
with self.assertRaises(type(error)):
|
|
_ocr_page(pdf, 1)
|
|
self.assertEqual(vision.call_count, max_calls)
|
|
|
|
|
|
if __name__ == '__main__':
|
|
unittest.main()
|