diff --git a/pageindex/__init__.py b/pageindex/__init__.py index 61b8958..e05c5b4 100644 --- a/pageindex/__init__.py +++ b/pageindex/__init__.py @@ -7,6 +7,13 @@ from dotenv import load_dotenv as _load_dotenv _load_dotenv() +# Backward compatibility: honor CHATGPT_API_KEY as an alias for OPENAI_API_KEY +# (kept from the pre-SDK pageindex.utils). Runs after load_dotenv so a value in +# .env is picked up too; only fills OPENAI_API_KEY when it isn't already set. +import os as _os +if not _os.getenv("OPENAI_API_KEY") and _os.getenv("CHATGPT_API_KEY"): + _os.environ["OPENAI_API_KEY"] = _os.getenv("CHATGPT_API_KEY") + # Upstream exports (backward compatibility). Import from the canonical # pageindex.index.* modules directly so `import pageindex` does NOT trip the # top-level deprecation shims (pageindex.page_index / .page_index_md / .utils). diff --git a/pageindex/index/utils.py b/pageindex/index/utils.py index e8be39e..7a10ed2 100644 --- a/pageindex/index/utils.py +++ b/pageindex/index/utils.py @@ -464,18 +464,20 @@ def get_pdf_page_content(file_path: str, page_nums: list[int]) -> list[dict]: def get_md_page_content(structure: list, page_nums: list[int]) -> list[dict]: """ For Markdown documents, 'pages' are line numbers. - Find nodes whose line_num falls within [min(page_nums), max(page_nums)] and return their text. + Return only the nodes whose line_num is one of ``page_nums`` (exact match), + mirroring the PDF path. A non-contiguous spec like [5, 100] returns just + those two lines, not the whole [5, 100] range. """ if not page_nums: return [] - min_line, max_line = min(page_nums), max(page_nums) + wanted = set(page_nums) results = [] seen = set() def _traverse(nodes): for node in nodes: ln = node.get('line_num') - if ln and min_line <= ln <= max_line and ln not in seen: + if ln in wanted and ln not in seen: seen.add(ln) results.append({'page': ln, 'content': node.get('text', '')}) if node.get('nodes'): diff --git a/pageindex/retrieve.py b/pageindex/retrieve.py index e0e1537..18a4494 100644 --- a/pageindex/retrieve.py +++ b/pageindex/retrieve.py @@ -56,16 +56,19 @@ def _get_pdf_page_content(doc_info: dict, page_nums: list[int]) -> list[dict]: def _get_md_page_content(doc_info: dict, page_nums: list[int]) -> list[dict]: """ For Markdown documents, 'pages' are line numbers. - Find nodes whose line_num falls within [min(page_nums), max(page_nums)] and return their text. + Return only the nodes whose line_num is one of ``page_nums`` (exact match), + not the whole [min(page_nums), max(page_nums)] range. """ - min_line, max_line = min(page_nums), max(page_nums) + if not page_nums: + return [] + wanted = set(page_nums) results = [] seen = set() def _traverse(nodes): for node in nodes: ln = node.get('line_num') - if ln and min_line <= ln <= max_line and ln not in seen: + if ln in wanted and ln not in seen: seen.add(ln) results.append({'page': ln, 'content': node.get('text', '')}) if node.get('nodes'): diff --git a/tests/test_env_compat.py b/tests/test_env_compat.py new file mode 100644 index 0000000..9dee3db --- /dev/null +++ b/tests/test_env_compat.py @@ -0,0 +1,37 @@ +"""CHATGPT_API_KEY must keep working as an alias for OPENAI_API_KEY (backward +compat carried over from the pre-SDK pageindex.utils; PR #272 review). + +The alias runs at import time in pageindex/__init__.py, so each case runs in a +fresh subprocess with a controlled environment. cwd is a temp dir so load_dotenv +can't pick up the repo's own .env and skew the result.""" + +import os +import subprocess +import sys +from pathlib import Path + +REPO = Path(__file__).resolve().parent.parent +_PRINT_OPENAI = "import pageindex, os; print(os.environ.get('OPENAI_API_KEY', ''))" + + +def _run(tmp_path, **overrides): + env = {k: v for k, v in os.environ.items() + if k not in ("OPENAI_API_KEY", "CHATGPT_API_KEY")} + env["PYTHONPATH"] = str(REPO) + env.update(overrides) + r = subprocess.run( + [sys.executable, "-c", _PRINT_OPENAI], + env=env, cwd=str(tmp_path), capture_output=True, text=True, + ) + assert r.returncode == 0, r.stderr + return r.stdout.strip() + + +def test_chatgpt_api_key_aliases_openai(tmp_path): + # Only CHATGPT_API_KEY set -> OPENAI_API_KEY gets filled from it. + assert _run(tmp_path, CHATGPT_API_KEY="sk-alias-123") == "sk-alias-123" + + +def test_existing_openai_api_key_is_not_overwritten(tmp_path): + # Both set -> the real OPENAI_API_KEY wins; the alias must not clobber it. + assert _run(tmp_path, OPENAI_API_KEY="sk-real", CHATGPT_API_KEY="sk-alias") == "sk-real" diff --git a/tests/test_page_content.py b/tests/test_page_content.py new file mode 100644 index 0000000..e810651 --- /dev/null +++ b/tests/test_page_content.py @@ -0,0 +1,36 @@ +"""Markdown page-content selection must return exactly the requested lines, +mirroring the PDF path — not the whole [min, max] range (PR #272 review / #280).""" + + +def _md_structure(): + # line_num 40 sits *between* 5 and 100 but is NOT requested below. + return [ + {"line_num": 5, "text": "line five", "nodes": [ + {"line_num": 40, "text": "line forty (should be excluded)", "nodes": []}, + ]}, + {"line_num": 100, "text": "line hundred", "nodes": []}, + {"line_num": 101, "text": "line 101", "nodes": []}, + ] + + +def test_get_md_page_content_returns_only_requested_lines(): + from pageindex.index.utils import get_md_page_content + + out = get_md_page_content(_md_structure(), [5, 100]) + # exactly the two requested lines — not 5, 40, 100 (the old range behavior) + assert [r["page"] for r in out] == [5, 100] + assert all("forty" not in r["content"] for r in out) + + +def test_get_md_page_content_empty_spec(): + from pageindex.index.utils import get_md_page_content + + assert get_md_page_content(_md_structure(), []) == [] + + +def test_retrieve_md_page_content_returns_only_requested_lines(): + # The legacy retrieve path has its own copy of the same logic. + from pageindex.retrieve import _get_md_page_content + + out = _get_md_page_content({"structure": _md_structure()}, [5, 100]) + assert [r["page"] for r in out] == [5, 100]