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/parser/protocol.py b/pageindex/parser/protocol.py index 76d7b0a..939af22 100644 --- a/pageindex/parser/protocol.py +++ b/pageindex/parser/protocol.py @@ -24,5 +24,8 @@ class ParsedDocument: @runtime_checkable class DocumentParser(Protocol): - def supported_extensions(self) -> list[str]: ... - def parse(self, file_path: str, **kwargs) -> ParsedDocument: ... + def supported_extensions(self) -> list[str]: + """Return the file extensions this parser handles (e.g. ['.pdf']).""" + + def parse(self, file_path: str, **kwargs) -> ParsedDocument: + """Parse a file into a ParsedDocument (a flat list of ContentNode).""" 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/pageindex/storage/protocol.py b/pageindex/storage/protocol.py index 427021b..5d7d107 100644 --- a/pageindex/storage/protocol.py +++ b/pageindex/storage/protocol.py @@ -4,15 +4,40 @@ from typing import Protocol, runtime_checkable @runtime_checkable class StorageEngine(Protocol): - def create_collection(self, name: str) -> None: ... - def get_or_create_collection(self, name: str) -> None: ... - def list_collections(self) -> list[str]: ... - def delete_collection(self, name: str) -> None: ... - def save_document(self, collection: str, doc_id: str, doc: dict) -> None: ... - def find_document_by_hash(self, collection: str, file_hash: str) -> str | None: ... - def get_document(self, collection: str, doc_id: str) -> dict: ... - def get_document_structure(self, collection: str, doc_id: str) -> list: ... - def get_pages(self, collection: str, doc_id: str) -> list | None: ... - def list_documents(self, collection: str) -> list[dict]: ... - def delete_document(self, collection: str, doc_id: str) -> None: ... - def close(self) -> None: ... + """Persistence contract for collections and their documents.""" + + def create_collection(self, name: str) -> None: + """Create a new collection; error if it already exists.""" + + def get_or_create_collection(self, name: str) -> None: + """Create the collection if absent; no-op if it already exists.""" + + def list_collections(self) -> list[str]: + """Return all collection names.""" + + def delete_collection(self, name: str) -> None: + """Delete a collection and all its documents.""" + + def save_document(self, collection: str, doc_id: str, doc: dict) -> None: + """Persist a document under a collection.""" + + def find_document_by_hash(self, collection: str, file_hash: str) -> str | None: + """Return the doc_id with this file hash in the collection, or None.""" + + def get_document(self, collection: str, doc_id: str) -> dict: + """Return a document's metadata.""" + + def get_document_structure(self, collection: str, doc_id: str) -> list: + """Return a document's tree structure.""" + + def get_pages(self, collection: str, doc_id: str) -> list | None: + """Return cached page content, or None if not cached.""" + + def list_documents(self, collection: str) -> list[dict]: + """Return metadata for all documents in a collection.""" + + def delete_document(self, collection: str, doc_id: str) -> None: + """Delete a single document from a collection.""" + + def close(self) -> None: + """Release any underlying resources (connections, handles).""" 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]