From 81002d585d41d5a0371af312063c45499d4828b6 Mon Sep 17 00:00:00 2001 From: prthm2910 Date: Fri, 3 Jul 2026 14:30:36 +0530 Subject: [PATCH 1/5] feat: add rate limiting, 429 retry, and error recovery for LLM calls --- pageindex/page_index.py | 74 ++++++++++---- pageindex/utils.py | 219 +++++++++++++++++++++++++++++++++++++--- 2 files changed, 263 insertions(+), 30 deletions(-) diff --git a/pageindex/page_index.py b/pageindex/page_index.py index 9004309fb..c8bc3c89a 100644 --- a/pageindex/page_index.py +++ b/pageindex/page_index.py @@ -119,6 +119,8 @@ def toc_detector_single_page(content, model=None): response = llm_completion(model=model, prompt=prompt) # print('response', response) json_content = extract_json(response) + if not json_content or 'toc_detected' not in json_content: + return False return json_content['toc_detected'] @@ -137,6 +139,8 @@ def check_if_toc_extraction_is_complete(content, toc, model=None): prompt = prompt + '\n Document:\n' + content + '\n Table of contents:\n' + toc response = llm_completion(model=model, prompt=prompt) json_content = extract_json(response) + if not json_content or 'completed' not in json_content: + return "no" return json_content['completed'] @@ -155,6 +159,8 @@ def check_if_toc_transformation_is_complete(content, toc, model=None): prompt = prompt + '\n Raw Table of contents:\n' + content + '\n Cleaned Table of contents:\n' + toc response = llm_completion(model=model, prompt=prompt) json_content = extract_json(response) + if not json_content or 'completed' not in json_content: + return "no" return json_content['completed'] def extract_toc_content(content, model=None): @@ -217,6 +223,8 @@ def detect_page_index(toc_content, model=None): response = llm_completion(model=model, prompt=prompt) json_content = extract_json(response) + if not json_content or 'page_index_given_in_toc' not in json_content: + return "no" return json_content['page_index_given_in_toc'] def toc_extractor(page_list, toc_page_list, model): @@ -296,7 +304,10 @@ def toc_transformer(toc_content, model=None): if_complete = check_if_toc_transformation_is_complete(toc_content, last_complete, model) if if_complete == "yes" and finish_reason == "finished": last_complete = extract_json(last_complete) - cleaned_response=convert_page_to_int(last_complete['table_of_contents']) + if not last_complete or 'table_of_contents' not in last_complete: + cleaned_response = [] + else: + cleaned_response=convert_page_to_int(last_complete['table_of_contents']) return cleaned_response last_complete = get_json_content(last_complete) @@ -332,7 +343,10 @@ def toc_transformer(toc_content, model=None): last_complete = extract_json(last_complete) - cleaned_response=convert_page_to_int(last_complete['table_of_contents']) + if not last_complete or 'table_of_contents' not in last_complete: + cleaned_response = [] + else: + cleaned_response=convert_page_to_int(last_complete['table_of_contents']) return cleaned_response @@ -485,9 +499,12 @@ def add_page_number_to_toc(part, structure, model=None): current_json_raw = llm_completion(model=model, prompt=prompt) json_result = extract_json(current_json_raw) - for item in json_result: - if 'start' in item: - del item['start'] + if isinstance(json_result, list): + for item in json_result: + if isinstance(item, dict) and 'start' in item: + del item['start'] + else: + json_result = [] return json_result @@ -532,11 +549,16 @@ def generate_toc_continue(toc_content, part, model=None): Directly return the additional part of the final JSON structure. Do not output anything else.""" prompt = prompt + '\nGiven text\n:' + part + '\nPrevious tree structure\n:' + json.dumps(toc_content, indent=2) - response, finish_reason = llm_completion(model=model, prompt=prompt, return_finish_reason=True) - if finish_reason == 'finished': - return extract_json(response) - else: - raise Exception(f'finish reason: {finish_reason}') + try: + response, finish_reason = llm_completion(model=model, prompt=prompt, return_finish_reason=True) + if finish_reason == 'finished': + return extract_json(response) + else: + print(f"Warning: generate_toc_continue failed with finish reason: {finish_reason}") + return [] + except Exception as e: + print(f"Warning: generate_toc_continue encountered an exception: {e}") + return [] ### add verify completeness def generate_toc_init(part, model=None): @@ -566,12 +588,16 @@ def generate_toc_init(part, model=None): Directly return the final JSON structure. Do not output anything else.""" prompt = prompt + '\nGiven text\n:' + part - response, finish_reason = llm_completion(model=model, prompt=prompt, return_finish_reason=True) - - if finish_reason == 'finished': - return extract_json(response) - else: - raise Exception(f'finish reason: {finish_reason}') + try: + response, finish_reason = llm_completion(model=model, prompt=prompt, return_finish_reason=True) + if finish_reason == 'finished': + return extract_json(response) + else: + print(f"Warning: generate_toc_init failed with finish reason: {finish_reason}") + return [] + except Exception as e: + print(f"Warning: generate_toc_init encountered an exception: {e}") + return [] def process_no_toc(page_list, start_index=1, model=None, logger=None): page_contents=[] @@ -753,6 +779,8 @@ async def single_toc_item_index_fixer(section_title, content, model=None): prompt = toc_extractor_prompt + '\nSection Title:\n' + str(section_title) + '\nDocument pages:\n' + content response = await llm_acompletion(model=model, prompt=prompt) json_content = extract_json(response) + if not json_content or 'physical_index' not in json_content: + return None return convert_physical_index_to_int(json_content['physical_index']) @@ -994,7 +1022,11 @@ async def meta_processor(page_list, mode=None, toc_content=None, toc_page_list=N elif mode == 'process_toc_no_page_numbers': return await meta_processor(page_list, mode='process_no_toc', start_index=start_index, opt=opt, logger=logger) else: - raise Exception('Processing failed') + if logger: + logger.error('Processing failed, returning empty list as fallback.') + else: + print('Processing failed, returning empty list as fallback.') + return [] async def process_large_node_recursively(node, page_list, opt=None, logger=None): @@ -1111,13 +1143,19 @@ async def page_index_builder(): def page_index(doc, model=None, toc_check_page_num=None, max_page_num_each_node=None, max_token_num_each_node=None, - if_add_node_id=None, if_add_node_summary=None, if_add_doc_description=None, if_add_node_text=None): + if_add_node_id=None, if_add_node_summary=None, if_add_doc_description=None, if_add_node_text=None, + max_concurrent_requests=None, rpm_limit=None): user_opt = { arg: value for arg, value in locals().items() if arg != "doc" and value is not None } opt = ConfigLoader().load(user_opt) + + from .utils import set_concurrency_limit, set_rpm_limit + set_concurrency_limit(opt.max_concurrent_requests) + set_rpm_limit(opt.rpm_limit) + return page_index_main(doc, opt) diff --git a/pageindex/utils.py b/pageindex/utils.py index f00ccf3a7..12f0105cf 100644 --- a/pageindex/utils.py +++ b/pageindex/utils.py @@ -1,3 +1,6 @@ +import random +import re +import threading import litellm import logging import os @@ -23,23 +26,196 @@ litellm.drop_params = True +# Global concurrency and Rate Limiting settings +# No defaults — callers must configure via set_concurrency_limit() / set_rpm_limit() +# or pass values to page_index(). The config.yaml file provides repo defaults. +_concurrency_limit = None +_semaphores = {} +_rpm_limit = None +_rate_limiter = None + +class SlidingWindowRateLimiter: + """Sliding-window log rate limiter. + + Tracks request timestamps across a 60-second window and enforces + a minimum spacing between consecutive requests to prevent bursts. + """ + def __init__(self, requests_per_minute): + self.requests_per_minute = requests_per_minute + self.request_times = [] + self.lock = threading.Lock() + + def _wait_time(self): + with self.lock: + now = time.time() + self.request_times = [t for t in self.request_times if now - t < 60.0] + + min_spacing = 60.0 / self.requests_per_minute if self.requests_per_minute > 0 else 0 + + if self.request_times: + last_scheduled = self.request_times[-1] + earliest_start = last_scheduled + min_spacing + else: + earliest_start = now + + if len(self.request_times) < self.requests_per_minute: + scheduled_time = max(now, earliest_start) + self.request_times.append(scheduled_time) + return max(0.0, scheduled_time - now) + + oldest = self.request_times[0] + window_wait = 60.0 - (now - oldest) + scheduled_time = max(now + window_wait, earliest_start) + self.request_times.append(scheduled_time) + return max(0.0, scheduled_time - now) + + def wait_sync(self): + if self.requests_per_minute <= 0: + return + wait_t = self._wait_time() + if wait_t > 0: + time.sleep(wait_t) + + async def wait_async(self): + if self.requests_per_minute <= 0: + return + wait_t = self._wait_time() + if wait_t > 0: + await asyncio.sleep(wait_t) + +def set_concurrency_limit(limit): + global _concurrency_limit + if limit is not None: + _concurrency_limit = limit + +def set_rpm_limit(limit): + global _rpm_limit, _rate_limiter + _rpm_limit = limit + if limit is not None and limit > 0: + _rate_limiter = SlidingWindowRateLimiter(limit) + else: + _rate_limiter = None + +def get_rate_limiter(): + global _rate_limiter + if _rate_limiter is None and _rpm_limit is not None and _rpm_limit > 0: + _rate_limiter = SlidingWindowRateLimiter(_rpm_limit) + return _rate_limiter + +def get_llm_semaphore(): + if _concurrency_limit is None or _concurrency_limit <= 0: + return None + loop = asyncio.get_running_loop() + if loop not in _semaphores: + _semaphores[loop] = asyncio.Semaphore(_concurrency_limit) + return _semaphores[loop] + def count_tokens(text, model=None): if not text: return 0 return litellm.token_counter(model=model, text=text) +def _is_rate_limit_error(exception): + """Check if an exception is caused by a rate limit (HTTP 429).""" + error_str = str(exception).lower() + # Check for 429 status code or rate limit keywords + if "429" in error_str or "rate limit" in error_str or "too many requests" in error_str: + return True + # Check for LiteLLM RateLimitError class + if hasattr(litellm, "RateLimitError") and isinstance(exception, litellm.RateLimitError): + return True + return False + + +def _extract_retry_headers(exception): + """Extract rate-limit headers from the exception if available.""" + retry_after = None + reset_requests = None + + # LiteLLM attaches headers to the exception or its response + # Check _response_headers, response_headers, or the exception args + for attr in ("_response_headers", "response_headers"): + headers = getattr(exception, attr, None) + if headers and isinstance(headers, dict): + retry_after = headers.get("retry-after") or headers.get("Retry-After") + reset_requests = headers.get("x-ratelimit-reset-requests") + break + + if not retry_after: + match = re.search(r"Retry-After:\s*(\d+)", str(exception)) + if match: + retry_after = match.group(1) + + return retry_after, reset_requests + + +def _extract_retry_delay(exception, attempt): + """ + Determine how long to wait before retrying after an error. + + Priority: + 1. Retry-After header from the provider (seconds or HTTP-date) + 2. x-ratelimit-reset-requests header (estimated reset time) + 3. Exponential backoff with jitter (fallback for non-429 errors) + """ + if not _is_rate_limit_error(exception): + # Non-rate-limit errors: standard exponential backoff with jitter + delay = min(10.0, 1.0 * (1.5 ** attempt)) + return delay + random.uniform(0.0, 0.5) + + # For rate-limit errors, try to respect provider headers + retry_after, reset_requests = _extract_retry_headers(exception) + + # 1. Retry-After header (can be seconds or HTTP-date) + if retry_after: + try: + # Try as integer seconds first (most common: "5", "30", "60") + seconds = int(retry_after) + # Cap and add jitter to spread retries + capped = min(seconds, 120) + return capped + random.uniform(0.0, 1.0) + except ValueError: + # Could be an HTTP-date (e.g. "Wed, 21 Oct 2026 07:28:00 GMT") + try: + from email.utils import parsedate_to_datetime + retry_time = parsedate_to_datetime(retry_after) + from datetime import datetime, timezone + wait = (retry_time - datetime.now(timezone.utc)).total_seconds() + if 0 < wait <= 120: + return wait + random.uniform(0.0, 1.0) + except Exception: + pass + + # 2. x-ratelimit-reset-requests (some providers send this as a duration or timestamp) + if reset_requests: + try: + seconds = float(reset_requests) + if 0 < seconds <= 120: + return seconds + random.uniform(0.0, 1.0) + except ValueError: + pass + + # 3. Fallback: exponential backoff with jitter (headers unavailable) + delay = min(30.0, 5.0 * (1.5 ** attempt)) # Longer base for rate limits + return delay + random.uniform(0.0, 2.0) + + def llm_completion(model, prompt, chat_history=None, return_finish_reason=False): if model: model = model.removeprefix("litellm/") max_retries = 10 messages = list(chat_history) + [{"role": "user", "content": prompt}] if chat_history else [{"role": "user", "content": prompt}] + limiter = get_rate_limiter() for i in range(max_retries): try: + if limiter: + limiter.wait_sync() response = litellm.completion( model=model, messages=messages, temperature=0, + num_retries=0, # disable LiteLLM's built-in retry; we handle it ourselves ) content = response.choices[0].message.content if return_finish_reason: @@ -47,10 +223,12 @@ def llm_completion(model, prompt, chat_history=None, return_finish_reason=False) return content, finish_reason return content except Exception as e: - print('************* Retrying *************') - logging.error(f"Error: {e}") + is_429 = _is_rate_limit_error(e) + logging.error(f"Error: {e}" + (" [Rate Limit Detected]" if is_429 else "")) if i < max_retries - 1: - time.sleep(1) + print('************* Retrying *************') + delay = _extract_retry_delay(e, i) + time.sleep(delay) else: logging.error('Max retries reached for prompt: ' + prompt) if return_finish_reason: @@ -58,25 +236,42 @@ def llm_completion(model, prompt, chat_history=None, return_finish_reason=False) return "" - async def llm_acompletion(model, prompt): if model: model = model.removeprefix("litellm/") max_retries = 10 messages = [{"role": "user", "content": prompt}] + + sem = get_llm_semaphore() + limiter = get_rate_limiter() + for i in range(max_retries): try: - response = await litellm.acompletion( - model=model, - messages=messages, - temperature=0, - ) + if limiter: + await limiter.wait_async() + if sem is not None: + async with sem: + response = await litellm.acompletion( + model=model, + messages=messages, + temperature=0, + num_retries=0, # disable LiteLLM's built-in retry; we handle it ourselves + ) + else: + response = await litellm.acompletion( + model=model, + messages=messages, + temperature=0, + num_retries=0, + ) return response.choices[0].message.content except Exception as e: - print('************* Retrying *************') - logging.error(f"Error: {e}") + is_429 = _is_rate_limit_error(e) + logging.error(f"Error: {e}" + (" [Rate Limit Detected]" if is_429 else "")) if i < max_retries - 1: - await asyncio.sleep(1) + print('************* Retrying *************') + delay = _extract_retry_delay(e, i) + await asyncio.sleep(delay) else: logging.error('Max retries reached for prompt: ' + prompt) return "" From c7ba971b80716801ef04192ca17c79867e980d09 Mon Sep 17 00:00:00 2001 From: prthm2910 Date: Fri, 3 Jul 2026 14:30:45 +0530 Subject: [PATCH 2/5] chore: add max_concurrent_requests and rpm_limit config defaults --- pageindex/config.yaml | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/pageindex/config.yaml b/pageindex/config.yaml index 591fe9331..7b78960aa 100644 --- a/pageindex/config.yaml +++ b/pageindex/config.yaml @@ -7,4 +7,6 @@ max_token_num_each_node: 20000 if_add_node_id: "yes" if_add_node_summary: "yes" if_add_doc_description: "no" -if_add_node_text: "no" \ No newline at end of file +if_add_node_text: "no" +max_concurrent_requests: 0 # 0 = unlimited; set >0 to cap concurrent API calls +rpm_limit: 0 # 0 = unlimited; set >0 to rate-limit requests per minute \ No newline at end of file From 0f84adc4c70a7dc5b0771d2172a77c4984d48694 Mon Sep 17 00:00:00 2001 From: prthm2910 Date: Fri, 3 Jul 2026 14:31:02 +0530 Subject: [PATCH 3/5] feat: add --rpm-limit and --max-concurrent-requests CLI flags --- run_pageindex.py | 24 +++++++++++++++++++++--- 1 file changed, 21 insertions(+), 3 deletions(-) diff --git a/run_pageindex.py b/run_pageindex.py index 673439d89..1805ac949 100644 --- a/run_pageindex.py +++ b/run_pageindex.py @@ -28,6 +28,10 @@ help='Whether to add doc description to the doc') parser.add_argument('--if-add-node-text', type=str, default=None, help='Whether to add text to the node') + parser.add_argument('--rpm-limit', type=int, default=None, + help='Rate limit in requests per minute (RPM)') + parser.add_argument('--max-concurrent-requests', type=int, default=None, + help='Max concurrent API requests (0 = unlimited)') # Markdown specific arguments parser.add_argument('--if-thinning', type=str, default='no', @@ -61,9 +65,16 @@ 'if_add_node_summary': args.if_add_node_summary, 'if_add_doc_description': args.if_add_doc_description, 'if_add_node_text': args.if_add_node_text, + 'rpm_limit': args.rpm_limit, + 'max_concurrent_requests': args.max_concurrent_requests, } opt = ConfigLoader().load({k: v for k, v in user_opt.items() if v is not None}) + # Apply rate limiting settings + from pageindex.utils import set_concurrency_limit, set_rpm_limit + set_concurrency_limit(opt.max_concurrent_requests) + set_rpm_limit(opt.rpm_limit) + # Process the PDF toc_with_page_number = page_index_main(args.pdf_path, opt) print('Parsing done, saving to file...') @@ -102,11 +113,18 @@ 'if_add_node_summary': args.if_add_node_summary, 'if_add_doc_description': args.if_add_doc_description, 'if_add_node_text': args.if_add_node_text, - 'if_add_node_id': args.if_add_node_id + 'if_add_node_id': args.if_add_node_id, + 'rpm_limit': args.rpm_limit, + 'max_concurrent_requests': args.max_concurrent_requests, } - + # Load config with defaults from config.yaml - opt = config_loader.load(user_opt) + opt = config_loader.load({k: v for k, v in user_opt.items() if v is not None}) + + # Apply rate limiting settings for markdown as well + from pageindex.utils import set_concurrency_limit, set_rpm_limit + set_concurrency_limit(opt.max_concurrent_requests) + set_rpm_limit(opt.rpm_limit) toc_with_page_number = asyncio.run(md_to_tree( md_path=args.md_path, From 7919d283be4d86acf4f9a8ccd87c157d5293bd5c Mon Sep 17 00:00:00 2001 From: prthm2910 Date: Fri, 3 Jul 2026 14:31:14 +0530 Subject: [PATCH 4/5] test: add tests for rate limiter, 429 handling, concurrency, and error recovery --- requirements-test.txt | 2 + tests/__init__.py | 0 tests/conftest.py | 15 ++++ tests/test_429_handling.py | 153 +++++++++++++++++++++++++++++++++++ tests/test_concurrency.py | 104 ++++++++++++++++++++++++ tests/test_error_recovery.py | 76 +++++++++++++++++ tests/test_rate_limiter.py | 80 ++++++++++++++++++ 7 files changed, 430 insertions(+) create mode 100644 requirements-test.txt create mode 100644 tests/__init__.py create mode 100644 tests/conftest.py create mode 100644 tests/test_429_handling.py create mode 100644 tests/test_concurrency.py create mode 100644 tests/test_error_recovery.py create mode 100644 tests/test_rate_limiter.py diff --git a/requirements-test.txt b/requirements-test.txt new file mode 100644 index 000000000..94022afa2 --- /dev/null +++ b/requirements-test.txt @@ -0,0 +1,2 @@ +pytest>=8.0 +pytest-asyncio>=0.23 diff --git a/tests/__init__.py b/tests/__init__.py new file mode 100644 index 000000000..e69de29bb diff --git a/tests/conftest.py b/tests/conftest.py new file mode 100644 index 000000000..edda5628e --- /dev/null +++ b/tests/conftest.py @@ -0,0 +1,15 @@ +"""Shared fixtures for PageIndex tests.""" +import pytest + +from pageindex.utils import set_concurrency_limit, set_rpm_limit + + +@pytest.fixture(autouse=True) +def reset_global_settings(): + """Reset global concurrency and RPM settings after each test.""" + yield + set_concurrency_limit(0) + set_rpm_limit(0) + import pageindex.utils as _utils + _utils._semaphores.clear() + _utils._rate_limiter = None diff --git a/tests/test_429_handling.py b/tests/test_429_handling.py new file mode 100644 index 000000000..51c5bc86b --- /dev/null +++ b/tests/test_429_handling.py @@ -0,0 +1,153 @@ +"""Tests for 429 rate limit detection and retry behavior.""" +import asyncio +import time +from unittest.mock import patch, MagicMock + +import pytest + +from pageindex.utils import ( + _is_rate_limit_error, + _extract_retry_headers, + _extract_retry_delay, + llm_completion, + llm_acompletion, +) + + +class MockRateLimitError(Exception): + def __init__(self, message, headers=None): + super().__init__(message) + self._response_headers = headers or {} + + +class TestRateLimitDetection: + + def test_detects_429_in_message(self): + err = Exception("HTTP 429 Too Many Requests") + assert _is_rate_limit_error(err) is True + + def test_detects_rate_limit_text(self): + err = Exception("Rate limit exceeded for this model") + assert _is_rate_limit_error(err) is True + + def test_detects_too_many_requests(self): + err = Exception("Error: too many requests, please slow down") + assert _is_rate_limit_error(err) is True + + def test_non_rate_limit_error(self): + err = Exception("Connection refused") + assert _is_rate_limit_error(err) is False + + def test_detects_litellm_rate_limit_error(self): + import litellm + if hasattr(litellm, "RateLimitError"): + err = litellm.RateLimitError( + "Rate limited", + llm_provider="openai", + model="gpt-4" + ) + assert _is_rate_limit_error(err) is True + + +class TestRetryHeaderExtraction: + + def test_extracts_retry_after_from_headers(self): + err = MockRateLimitError("429", headers={"Retry-After": "30"}) + retry_after, reset = _extract_retry_headers(err) + assert retry_after == "30" + + def test_extracts_ratelimit_reset_from_headers(self): + err = MockRateLimitError("429", headers={ + "x-ratelimit-reset-requests": "45", + }) + retry_after, reset = _extract_retry_headers(err) + assert reset == "45" + + def test_returns_none_when_no_headers(self): + err = Exception("429 Too Many Requests") + retry_after, reset = _extract_retry_headers(err) + assert retry_after is None + assert reset is None + + +class TestRetryDelayCalculation: + + def test_retry_after_header_respected(self): + err = MockRateLimitError("429", headers={"Retry-After": "10"}) + delay = _extract_retry_delay(err, 0) + assert delay >= 10 + + def test_fallback_exponential_backoff(self): + err = Exception("429 Too Many Requests") + delay_0 = _extract_retry_delay(err, 0) + delay_1 = _extract_retry_delay(err, 1) + delay_2 = _extract_retry_delay(err, 2) + assert delay_1 > delay_0 + assert delay_2 > delay_1 + + def test_non_429_uses_standard_backoff(self): + err = Exception("Connection timeout") + delay = _extract_retry_delay(err, 0) + assert delay <= 10.5 + + def test_no_thundering_herd(self): + err = Exception("429 Too Many Requests") + delays = [_extract_retry_delay(err, 0) for _ in range(20)] + unique_delays = set(round(d, 2) for d in delays) + assert len(unique_delays) > 1 + + def test_retry_after_capped(self): + err = MockRateLimitError("429", headers={"Retry-After": "3600"}) + delay = _extract_retry_delay(err, 0) + assert delay <= 121 + + +class TestLLMCompletionRetry: + + def test_llm_completion_returns_empty_on_max_retries(self): + call_count = 0 + + def mock_completion(*args, **kwargs): + nonlocal call_count + call_count += 1 + raise Exception("429 Too Many Requests") + + with patch("pageindex.utils.litellm.completion", side_effect=mock_completion): + with patch("pageindex.utils.time.sleep"): + result = llm_completion("gpt-4", "test") + assert result == "" + assert call_count == 10 + + def test_llm_completion_succeeds_on_retry(self): + call_count = 0 + + def mock_completion(*args, **kwargs): + nonlocal call_count + call_count += 1 + if call_count <= 2: + raise Exception("429 Too Many Requests") + response = MagicMock() + response.choices = [MagicMock()] + response.choices[0].message.content = '{"answer": "yes"}' + response.choices[0].finish_reason = "finished" + return response + + with patch("pageindex.utils.litellm.completion", side_effect=mock_completion): + with patch("pageindex.utils.time.sleep"): + result = llm_completion("gpt-4", "test") + assert result == '{"answer": "yes"}' + + @pytest.mark.asyncio + async def test_llm_acompletion_returns_empty_on_max_retries(self): + call_count = 0 + + async def mock_acompletion(*args, **kwargs): + nonlocal call_count + call_count += 1 + raise Exception("429 Too Many Requests") + + with patch("pageindex.utils.litellm.acompletion", side_effect=mock_acompletion): + with patch("pageindex.utils.asyncio.sleep"): + result = await llm_acompletion("gpt-4", "test") + assert result == "" + assert call_count == 10 diff --git a/tests/test_concurrency.py b/tests/test_concurrency.py new file mode 100644 index 000000000..0e72aa46c --- /dev/null +++ b/tests/test_concurrency.py @@ -0,0 +1,104 @@ +"""Tests for concurrency limiter (semaphore + settings).""" +import asyncio + +import pytest + +from pageindex.utils import ( + set_concurrency_limit, + set_rpm_limit, + get_rate_limiter, + get_llm_semaphore, + SlidingWindowRateLimiter, + _semaphores, +) + + +class TestSettingsManagement: + + def test_default_concurrency_limit(self): + import pageindex.utils as _utils + set_concurrency_limit(0) + assert _utils._concurrency_limit == 0 + + def test_set_concurrency_limit(self): + import pageindex.utils as _utils + set_concurrency_limit(10) + assert _utils._concurrency_limit == 10 + + def test_set_rpm_limit_creates_limiter(self): + set_rpm_limit(50) + limiter = get_rate_limiter() + assert limiter is not None + assert isinstance(limiter, SlidingWindowRateLimiter) + assert limiter.requests_per_minute == 50 + + def test_set_rpm_limit_zero_disables(self): + set_rpm_limit(50) + get_rate_limiter() + set_rpm_limit(0) + limiter = get_rate_limiter() + assert limiter is None + + +class TestSemaphoreBehavior: + + def test_semaphore_limits_concurrency(self): + import pageindex.utils as _utils + _semaphores.clear() + set_concurrency_limit(2) + + max_concurrent = 0 + current_concurrent = 0 + + async def task(): + nonlocal max_concurrent, current_concurrent + async with get_llm_semaphore(): + current_concurrent += 1 + max_concurrent = max(max_concurrent, current_concurrent) + await asyncio.sleep(0.05) + current_concurrent -= 1 + + async def run(): + tasks = [task() for _ in range(6)] + await asyncio.gather(*tasks) + + asyncio.run(run()) + assert max_concurrent <= 2 + + def test_semaphore_no_deadlock(self): + import pageindex.utils as _utils + _semaphores.clear() + set_concurrency_limit(1) + + completed = [] + + async def task(task_id): + async with get_llm_semaphore(): + await asyncio.sleep(0.01) + completed.append(task_id) + return task_id + + async def run(): + tasks = [task(i) for i in range(5)] + results = await asyncio.gather(*tasks) + return results + + results = asyncio.run(run()) + assert len(results) == 5 + assert len(completed) == 5 + assert set(completed) == {0, 1, 2, 3, 4} + + +class TestRateLimiterIntegration: + + def test_get_rate_limiter_lazy_init(self): + import pageindex.utils as _utils + set_rpm_limit(0) + _utils._rate_limiter = None + assert _utils._rate_limiter is None + + set_rpm_limit(42) + + limiter = get_rate_limiter() + assert limiter is not None + assert limiter.requests_per_minute == 42 diff --git a/tests/test_error_recovery.py b/tests/test_error_recovery.py new file mode 100644 index 000000000..961314595 --- /dev/null +++ b/tests/test_error_recovery.py @@ -0,0 +1,76 @@ +"""Tests for cascading error recovery (KeyError prevention).""" +import asyncio +from unittest.mock import patch, MagicMock + +import pytest + +from pageindex.page_index import ( + toc_detector_single_page, + generate_toc_init, + generate_toc_continue, + add_page_number_to_toc, +) + + +class TestTocDetectorRecovery: + + def test_returns_false_on_empty_response(self): + with patch("pageindex.page_index.llm_completion", return_value=""): + result = toc_detector_single_page("some text", model="gpt-4") + assert result is False + + def test_returns_false_on_malformed_json(self): + with patch("pageindex.page_index.llm_completion", return_value="just plain text"): + result = toc_detector_single_page("some text", model="gpt-4") + assert result is False + + def test_returns_false_on_missing_key(self): + with patch("pageindex.page_index.llm_completion", return_value='{"thinking": "test"}'): + result = toc_detector_single_page("some text", model="gpt-4") + assert result is False + + def test_returns_true_when_valid(self): + with patch("pageindex.page_index.llm_completion", return_value='{"thinking": "test", "toc_detected": "yes"}'): + result = toc_detector_single_page("some text", model="gpt-4") + assert result == "yes" + + +class TestGenerateTocRecovery: + + def test_generate_toc_init_returns_empty_on_exception(self): + with patch("pageindex.page_index.llm_completion", side_effect=Exception("API error")): + result = generate_toc_init("some text", model="gpt-4") + assert result == [] + + def test_generate_toc_continue_returns_empty_on_exception(self): + with patch("pageindex.page_index.llm_completion", side_effect=Exception("API error")): + result = generate_toc_continue([], "some text", model="gpt-4") + assert result == [] + + def test_generate_toc_init_returns_empty_on_length(self): + def mock_completion(*args, return_finish_reason=False, **kwargs): + if return_finish_reason: + return '{"structure": "1", "title": "Test"', "length" + return '{"structure": "1", "title": "Test"' + + with patch("pageindex.page_index.llm_completion", side_effect=mock_completion): + result = generate_toc_init("some text", model="gpt-4") + assert result == [] + + +class TestAddPageNumberRecovery: + + def test_handles_non_list_json(self): + with patch("pageindex.page_index.llm_completion", return_value='{"error": "something went wrong"}'): + result = add_page_number_to_toc("text", [], model="gpt-4") + assert result == [] + + def test_handles_valid_list(self): + response = ( + '[{"structure": "1", "title": "Test", "physical_index": ' + '"", "start": "yes"}]' + ) + with patch("pageindex.page_index.llm_completion", return_value=response): + result = add_page_number_to_toc("text", [], model="gpt-4") + assert isinstance(result, list) + assert "start" not in result[0] diff --git a/tests/test_rate_limiter.py b/tests/test_rate_limiter.py new file mode 100644 index 000000000..69b30f3f3 --- /dev/null +++ b/tests/test_rate_limiter.py @@ -0,0 +1,80 @@ +"""Tests for SlidingWindowRateLimiter rate limiting logic.""" +import asyncio +import time +from unittest.mock import patch + +from pageindex.utils import SlidingWindowRateLimiter + + +class TestSlidingWindowRateLimiter: + + def test_wait_time_under_limit(self): + limiter = SlidingWindowRateLimiter(requests_per_minute=60) + wait = limiter._wait_time() + assert wait == 0.0 + + def test_wait_time_over_limit(self): + limiter = SlidingWindowRateLimiter(requests_per_minute=1) + limiter._wait_time() + wait = limiter._wait_time() + assert wait > 0 + assert wait <= 60.0 + + def test_window_sliding(self): + limiter = SlidingWindowRateLimiter(requests_per_minute=2) + limiter._wait_time() + limiter._wait_time() + + limiter.request_times = [time.time() - 61.0] + wait = limiter._wait_time() + assert wait == 0.0 + + def test_minimum_spacing(self): + rpm = 30 + limiter = SlidingWindowRateLimiter(requests_per_minute=rpm) + expected_spacing = 60.0 / rpm + + limiter._wait_time() + wait = limiter._wait_time() + assert wait >= expected_spacing - 0.01 + assert wait <= expected_spacing + 0.01 + + def test_sync_wait_completes(self): + limiter = SlidingWindowRateLimiter(requests_per_minute=60) + limiter.wait_sync() + start = time.time() + limiter.wait_sync() + elapsed = time.time() - start + assert elapsed >= 0.9 + + def test_async_wait_completes(self): + limiter = SlidingWindowRateLimiter(requests_per_minute=60) + limiter.wait_sync() + + async def run_test(): + start = time.time() + await limiter.wait_async() + return time.time() - start + + elapsed = asyncio.run(run_test()) + assert elapsed >= 0.9 + + def test_disabled_limiter(self): + limiter_zero = SlidingWindowRateLimiter(requests_per_minute=0) + limiter_neg = SlidingWindowRateLimiter(requests_per_minute=-5) + + start = time.time() + limiter_zero.wait_sync() + limiter_neg.wait_sync() + assert time.time() - start < 0.1 + + def test_disabled_limiter_async(self): + limiter = SlidingWindowRateLimiter(requests_per_minute=0) + + async def run_test(): + start = time.time() + await limiter.wait_async() + return time.time() - start + + elapsed = asyncio.run(run_test()) + assert elapsed < 0.1 From 5c7999434bb34e9aef3dc970f5f397f89a77a17b Mon Sep 17 00:00:00 2001 From: prthm2910 Date: Fri, 3 Jul 2026 14:43:52 +0530 Subject: [PATCH 5/5] fix: remove duplicate import re after rebase --- pageindex/utils.py | 1 - 1 file changed, 1 deletion(-) diff --git a/pageindex/utils.py b/pageindex/utils.py index 923ab5221..12f0105cf 100644 --- a/pageindex/utils.py +++ b/pageindex/utils.py @@ -19,7 +19,6 @@ import yaml from pathlib import Path from types import SimpleNamespace as config -import re # Backward compatibility: support CHATGPT_API_KEY as alias for OPENAI_API_KEY if not os.getenv("OPENAI_API_KEY") and os.getenv("CHATGPT_API_KEY"):