Skip to content

Commit 2a69c76

Browse files
committed
fix: exact md page selection + restore CHATGPT_API_KEY alias
Address PR #272 review: - get_md_page_content / retrieve._get_md_page_content returned every node whose line_num fell in [min(pages), max(pages)], so a non-contiguous spec like "5,100" over-fetched everything in between. Match the exact requested line numbers instead, mirroring the PDF path. (Same bug as #280.) - Restore the CHATGPT_API_KEY -> OPENAI_API_KEY backward-compat alias dropped when pageindex/utils.py became a re-export shim; users with only CHATGPT_API_KEY set would otherwise fail auth after upgrading. It now runs in __init__.py right after load_dotenv. Adds regression tests for both. Claude-Session: https://claude.ai/code/session_01Kx5DgKbhK1N8autqXH8SmS
1 parent 2d46d68 commit 2a69c76

5 files changed

Lines changed: 91 additions & 6 deletions

File tree

pageindex/__init__.py

Lines changed: 7 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -7,6 +7,13 @@
77
from dotenv import load_dotenv as _load_dotenv
88
_load_dotenv()
99

10+
# Backward compatibility: honor CHATGPT_API_KEY as an alias for OPENAI_API_KEY
11+
# (kept from the pre-SDK pageindex.utils). Runs after load_dotenv so a value in
12+
# .env is picked up too; only fills OPENAI_API_KEY when it isn't already set.
13+
import os as _os
14+
if not _os.getenv("OPENAI_API_KEY") and _os.getenv("CHATGPT_API_KEY"):
15+
_os.environ["OPENAI_API_KEY"] = _os.getenv("CHATGPT_API_KEY")
16+
1017
# Upstream exports (backward compatibility). Import from the canonical
1118
# pageindex.index.* modules directly so `import pageindex` does NOT trip the
1219
# top-level deprecation shims (pageindex.page_index / .page_index_md / .utils).

pageindex/index/utils.py

Lines changed: 5 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -464,18 +464,20 @@ def get_pdf_page_content(file_path: str, page_nums: list[int]) -> list[dict]:
464464
def get_md_page_content(structure: list, page_nums: list[int]) -> list[dict]:
465465
"""
466466
For Markdown documents, 'pages' are line numbers.
467-
Find nodes whose line_num falls within [min(page_nums), max(page_nums)] and return their text.
467+
Return only the nodes whose line_num is one of ``page_nums`` (exact match),
468+
mirroring the PDF path. A non-contiguous spec like [5, 100] returns just
469+
those two lines, not the whole [5, 100] range.
468470
"""
469471
if not page_nums:
470472
return []
471-
min_line, max_line = min(page_nums), max(page_nums)
473+
wanted = set(page_nums)
472474
results = []
473475
seen = set()
474476

475477
def _traverse(nodes):
476478
for node in nodes:
477479
ln = node.get('line_num')
478-
if ln and min_line <= ln <= max_line and ln not in seen:
480+
if ln in wanted and ln not in seen:
479481
seen.add(ln)
480482
results.append({'page': ln, 'content': node.get('text', '')})
481483
if node.get('nodes'):

pageindex/retrieve.py

Lines changed: 6 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -56,16 +56,19 @@ def _get_pdf_page_content(doc_info: dict, page_nums: list[int]) -> list[dict]:
5656
def _get_md_page_content(doc_info: dict, page_nums: list[int]) -> list[dict]:
5757
"""
5858
For Markdown documents, 'pages' are line numbers.
59-
Find nodes whose line_num falls within [min(page_nums), max(page_nums)] and return their text.
59+
Return only the nodes whose line_num is one of ``page_nums`` (exact match),
60+
not the whole [min(page_nums), max(page_nums)] range.
6061
"""
61-
min_line, max_line = min(page_nums), max(page_nums)
62+
if not page_nums:
63+
return []
64+
wanted = set(page_nums)
6265
results = []
6366
seen = set()
6467

6568
def _traverse(nodes):
6669
for node in nodes:
6770
ln = node.get('line_num')
68-
if ln and min_line <= ln <= max_line and ln not in seen:
71+
if ln in wanted and ln not in seen:
6972
seen.add(ln)
7073
results.append({'page': ln, 'content': node.get('text', '')})
7174
if node.get('nodes'):

tests/test_env_compat.py

Lines changed: 37 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,37 @@
1+
"""CHATGPT_API_KEY must keep working as an alias for OPENAI_API_KEY (backward
2+
compat carried over from the pre-SDK pageindex.utils; PR #272 review).
3+
4+
The alias runs at import time in pageindex/__init__.py, so each case runs in a
5+
fresh subprocess with a controlled environment. cwd is a temp dir so load_dotenv
6+
can't pick up the repo's own .env and skew the result."""
7+
8+
import os
9+
import subprocess
10+
import sys
11+
from pathlib import Path
12+
13+
REPO = Path(__file__).resolve().parent.parent
14+
_PRINT_OPENAI = "import pageindex, os; print(os.environ.get('OPENAI_API_KEY', ''))"
15+
16+
17+
def _run(tmp_path, **overrides):
18+
env = {k: v for k, v in os.environ.items()
19+
if k not in ("OPENAI_API_KEY", "CHATGPT_API_KEY")}
20+
env["PYTHONPATH"] = str(REPO)
21+
env.update(overrides)
22+
r = subprocess.run(
23+
[sys.executable, "-c", _PRINT_OPENAI],
24+
env=env, cwd=str(tmp_path), capture_output=True, text=True,
25+
)
26+
assert r.returncode == 0, r.stderr
27+
return r.stdout.strip()
28+
29+
30+
def test_chatgpt_api_key_aliases_openai(tmp_path):
31+
# Only CHATGPT_API_KEY set -> OPENAI_API_KEY gets filled from it.
32+
assert _run(tmp_path, CHATGPT_API_KEY="sk-alias-123") == "sk-alias-123"
33+
34+
35+
def test_existing_openai_api_key_is_not_overwritten(tmp_path):
36+
# Both set -> the real OPENAI_API_KEY wins; the alias must not clobber it.
37+
assert _run(tmp_path, OPENAI_API_KEY="sk-real", CHATGPT_API_KEY="sk-alias") == "sk-real"

tests/test_page_content.py

Lines changed: 36 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,36 @@
1+
"""Markdown page-content selection must return exactly the requested lines,
2+
mirroring the PDF path — not the whole [min, max] range (PR #272 review / #280)."""
3+
4+
5+
def _md_structure():
6+
# line_num 40 sits *between* 5 and 100 but is NOT requested below.
7+
return [
8+
{"line_num": 5, "text": "line five", "nodes": [
9+
{"line_num": 40, "text": "line forty (should be excluded)", "nodes": []},
10+
]},
11+
{"line_num": 100, "text": "line hundred", "nodes": []},
12+
{"line_num": 101, "text": "line 101", "nodes": []},
13+
]
14+
15+
16+
def test_get_md_page_content_returns_only_requested_lines():
17+
from pageindex.index.utils import get_md_page_content
18+
19+
out = get_md_page_content(_md_structure(), [5, 100])
20+
# exactly the two requested lines — not 5, 40, 100 (the old range behavior)
21+
assert [r["page"] for r in out] == [5, 100]
22+
assert all("forty" not in r["content"] for r in out)
23+
24+
25+
def test_get_md_page_content_empty_spec():
26+
from pageindex.index.utils import get_md_page_content
27+
28+
assert get_md_page_content(_md_structure(), []) == []
29+
30+
31+
def test_retrieve_md_page_content_returns_only_requested_lines():
32+
# The legacy retrieve path has its own copy of the same logic.
33+
from pageindex.retrieve import _get_md_page_content
34+
35+
out = _get_md_page_content({"structure": _md_structure()}, [5, 100])
36+
assert [r["page"] for r in out] == [5, 100]

0 commit comments

Comments
 (0)