feat: pass all toc pages to LLM (#338)

Author: Moskize91

.gitignore

@@ -3,6 +3,7 @@
 /.venv
 /analysing
 /models-cache
+/format.json
 
 # Byte-compiled / optimized / DLL files
 __pycache__/

format.template.json

@@ -0,0 +1,11 @@
+{
+  "key": "<YOUR_API_KEY>",
+  "url": "https://sample-llm.com/v1",
+  "model": "sample-llm-model",
+  "token_encoding": "o200k_base",
+  "timeout": 360.0,
+  "retry_times": 10,
+  "retry_interval_seconds": 0.75,
+  "temperature": 0.6,
+  "top_p": 0.6
+}

pdf_craft/common/reader.py

@@ -1,6 +1,6 @@
 import re
 from pathlib import Path
-from typing import Callable, Generator, Generic, TypeVar
+from typing import Callable, Container, Generator, Generic, TypeVar
 from xml.etree.ElementTree import Element
 
 from .xml import read_xml
@@ -25,11 +25,15 @@ def __init__(
             indexed_files.append((idx, p))
 
         indexed_files.sort(key=lambda t: t[0])
-        self._file_paths: list[Path] = [path for _, path in indexed_files]
+        self._indexed_files: list[tuple[int, Path]] = indexed_files
         self._decode: Callable[[Element], T] = decode
 
-    def read(self) -> Generator[T, None, None]:
-        for xml_path in self._file_paths:
+    def read(
+        self, page_indexes: Container[int] | None = None
+    ) -> Generator[T, None, None]:
+        for page_index, xml_path in self._indexed_files:
+            if page_indexes is not None and page_index not in page_indexes:
+                continue
             root = read_xml(xml_path)
             try:
                 yield self._decode(root)

pdf_craft/toc/analysing.py

@@ -88,10 +88,16 @@ def _do_analyse_toc(
 
         try:
             ref2level = analyse_toc_levels_by_llm(
-                toc_pages=toc_pages,
                 llm=llm,
+                toc_page_refs=toc_pages,
+                toc_page_contents=list(
+                    pages.read(
+                        page_indexes={toc_page.page_index for toc_page in toc_pages},
+                    )
+                ),
             )
             toc_page_indexes.extend(ref.page_index for ref in toc_pages)
+
         except LLMAnalysisError as e:
             print(f"LLM analysis failed, falling back to statistical method: {e}")
             ref2level = analyse_toc_levels(