Skip to content

fix: remove redundant import, use tempdir for batch, add logging #362

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
HemantSudarshan wants to merge 1 commit into
base: main
Choose a base branch
from
Open

fix: remove redundant import, use tempdir for batch, add logging #362

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
8 changes: 6 additions & 2 deletions extract_thinker/extractor.py
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -1,7 +1,6 @@
import asyncio
import base64
from typing import Any, Dict, List, Optional, IO, Type, Union, get_origin, get_type_hints, get_args, Annotated
from instructor.batch import BatchJob
import uuid
from pydantic import BaseModel
from extract_thinker.document_loader.document_loader_data import DocumentLoaderData
Expand All @@ -13,6 +12,7 @@
from extract_thinker.models.classification_response import ClassificationResponse, ClassificationResponseInternal
from extract_thinker.llm import LLM
import os
import tempfile
from extract_thinker.document_loader.loader_interceptor import LoaderInterceptor
from extract_thinker.document_loader.llm_interceptor import LlmInterceptor
from concurrent.futures import ThreadPoolExecutor, as_completed
Expand Down Expand Up @@ -83,7 +83,11 @@ def get_document_loader_for_file(self, source: Union[str, IO]) -> DocumentLoader
return self.document_loader

# Check all registered loaders
checked_loaders = set()
for loader in self.document_loaders_by_file_type.values():
if loader in checked_loaders:
continue
checked_loaders.add(loader)
if loader.can_handle(source):
return loader

Expand Down Expand Up @@ -983,7 +987,7 @@ def extract_batch(
)

# Create batch directory if it doesn't exist
batch_dir = os.path.join(os.getcwd(), "extract_thinker_batch")
batch_dir = tempfile.gettempdir()

Copilot AI Feb 14, 2026

Copy link

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

extract_batch() now uses tempfile.gettempdir() directly as batch_dir. This makes the batch input/output files live directly in the shared system temp directory, and (more importantly) BatchJob._cleanup_files() removes the parent directory when it becomes empty—so in some environments this could attempt to remove the system temp dir itself. Use a dedicated subdirectory under the temp dir (e.g., tempdir + "extract_thinker_batch" or a per-job temp folder) so cleanup only ever targets directories created by this library.

Suggested change
batch_dir = tempfile.gettempdir()
batch_dir = os.path.join(tempfile.gettempdir(), "extract_thinker_batch")

Copilot uses AI. Check for mistakes.
os.makedirs(batch_dir, exist_ok=True)

# Generate unique paths if not provided
Expand Down
43 changes: 23 additions & 20 deletions extract_thinker/markdown/markdown_converter.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -8,6 +8,9 @@
from concurrent.futures import ThreadPoolExecutor, as_completed
from extract_thinker.utils import encode_image, json_to_formatted_string, extract_thinking_json
import re
import logging

logger = logging.getLogger(__name__)

class ContentItem(BaseModel):
"""Represents a single piece of extracted content with certainty."""
Expand Down Expand Up @@ -176,7 +179,7 @@ def to_markdown_structured(self, source: Union[str, IO, List[Union[str, IO]]], p
# If setting vision mode fails, we probably can't proceed as expected.
raise ValueError(f"Failed to set vision mode on document loader: {e}") from e
else:
print("Warning: Document loader does not have set_vision_mode. Assuming it handles vision implicitly.")
logger.warning("Document loader does not have set_vision_mode. Assuming it handles vision implicitly.")

pages_data = self.document_loader.load(source)
if not isinstance(pages_data, list):
Expand Down Expand Up @@ -208,7 +211,7 @@ def to_markdown_structured(self, source: Union[str, IO, List[Union[str, IO]]], p
try:
result_strings[index] = future.result()
except Exception as exc:
print(f'Page {index + 1} processing failed: {exc}')
logger.error('Page %d processing failed: %s', index + 1, exc)
result_strings[index] = f"<!-- Error processing page {index + 1}: {exc} -->"

return result_strings
Expand Down Expand Up @@ -293,7 +296,7 @@ def _build_message_content(
placeholders_found = re.findall(placeholder_pattern, processed_text)

if placeholders_found:
print(f"Detected {len(placeholders_found)} image extraction placeholder(s). Adding specific instructions to content.")
logger.info("Detected %d image extraction placeholder(s). Adding specific instructions to content.", len(placeholders_found))
# Construct an instruction to append *after* the main text
instruction = (
"\n\n---\n"
Expand Down Expand Up @@ -354,7 +357,7 @@ def _process_page_with_llm(self, page_data: Any) -> PageContent:
raise ValueError("LLM is required for structured extraction but not configured.")

if not isinstance(page_data, dict):
print(f"Warning: Unexpected page data type: {type(page_data)}. Skipping LLM processing for this page.")
logger.warning("Unexpected page data type: %s. Skipping LLM processing for this page.", type(page_data))
return f"<!-- Error: Unexpected page data type: {type(page_data)} -->"

messages = self._build_messages(self._build_message_content(page_data, vision=True))
Expand All @@ -364,7 +367,7 @@ def _process_page_with_llm(self, page_data: Any) -> PageContent:
raw_response = self.llm.raw_completion(messages=messages)
return extract_thinking_json(raw_response, PageContent)
except Exception as e:
print(f"LLM request failed for page: {e}")
logger.error("LLM request failed for page: %s", e)
raise

def _process_content_data(
Expand Down Expand Up @@ -522,9 +525,9 @@ def _append_images(
}
})
else:
print(f"Warning: Could not get base64 for image item: {type(img)}")
logger.warning("Could not get base64 for image item: %s", type(img))
except Exception as e:
print(f"Warning: Error processing image item {type(img)}: {e}")
logger.warning("Error processing image item %s: %s", type(img), e)

# --- Copied Methods from Extractor --- END ---

Expand Down Expand Up @@ -560,7 +563,7 @@ def to_markdown(self, source: Union[str, IO, List[Union[str, IO]]], vision: bool
try:
self.document_loader.set_vision_mode(vision)
except Exception as e:
print(f"Warning: Failed to set vision mode on document loader: {e}")
logger.warning("Failed to set vision mode on document loader: %s", e)

# Load the document
pages_data = self.document_loader.load(source)
Expand All @@ -581,7 +584,7 @@ def to_markdown(self, source: Union[str, IO, List[Union[str, IO]]], vision: bool
if vision:
has_images = any(isinstance(page, dict) and (page.get('image') or page.get('images')) for page in pages_data)
if not has_images:
print("Warning: Vision processing enabled but no images found. Will process as text-only.")
logger.warning("Vision processing enabled but no images found. Will process as text-only.")

# Process pages in parallel
markdown_parts = [None] * len(pages_data) # Pre-allocate list
Expand All @@ -595,7 +598,7 @@ def to_markdown(self, source: Union[str, IO, List[Union[str, IO]]], vision: bool
try:
markdown_parts[index] = future.result()
except Exception as exc:
print(f'Page {index + 1} processing failed: {exc}')
logger.error('Page %d processing failed: %s', index + 1, exc)
markdown_parts[index] = f"<!-- Error processing page {index + 1}: {exc} -->"

# Return the list of markdown parts instead of joining them
Expand Down Expand Up @@ -634,13 +637,13 @@ def to_markdown(self, source: Union[str, IO, List[Union[str, IO]]], vision: bool
markdown_content = self.llm.request(messages=messages)
return [markdown_content]
except Exception as e:
print(f"LLM request failed: {e}")
logger.error("LLM request failed: %s", e)
raise

except Exception as e:
print(f"Error in markdown conversion: {e}")
logger.error("Error in markdown conversion: %s", e)
if self.document_loader:
print("Falling back to basic conversion.")
logger.info("Falling back to basic conversion.")
# Call basic conversion but convert its result to a list too
basic_result = self._basic_to_markdown(source, vision=vision, pages=pages)
return [basic_result]
Expand All @@ -664,7 +667,7 @@ def _process_markdown_page(self, page_data: Any) -> str:
raise ValueError("LLM is required for markdown extraction but not configured.")

if not isinstance(page_data, dict):
print(f"Warning: Unexpected page data type: {type(page_data)}. Skipping LLM processing for this page.")
logger.warning("Unexpected page data type: %s. Skipping LLM processing for this page.", type(page_data))
return f"<!-- Error: Unexpected page data type: {type(page_data)} -->"

# Use structured=False to get only Markdown content without JSON
Expand All @@ -675,7 +678,7 @@ def _process_markdown_page(self, page_data: Any) -> str:
markdown_content = self.llm.request(messages=messages)
return markdown_content.choices[0].message.content
except Exception as e:
print(f"LLM request failed for page: {e}")
logger.error("LLM request failed for page: %s", e)
raise

def _basic_to_markdown(self, source: Union[str, IO, List[Union[str, IO]]], vision: bool = False, pages: Optional[List[int]] = None) -> str:
Expand All @@ -701,7 +704,7 @@ def _basic_to_markdown(self, source: Union[str, IO, List[Union[str, IO]]], visio
try:
self.document_loader.set_vision_mode(vision)
except Exception as e:
print(f"Warning: Failed to set vision mode on document loader: {e}")
logger.warning("Failed to set vision mode on document loader: %s", e)

pages_data = self.document_loader.load(source)
if not isinstance(pages_data, list):
Expand All @@ -728,7 +731,7 @@ def _basic_to_markdown(self, source: Union[str, IO, List[Union[str, IO]]], visio
try:
markdown_parts[index] = future.result()
except Exception as exc:
print(f'Basic Page {index + 1} conversion failed: {exc}')
logger.error('Basic page %d conversion failed: %s', index + 1, exc)
markdown_parts[index] = f"\n\n<!-- Error converting page {index + 1}: {exc} -->\n\n"

return "\n\n".join(part for part in markdown_parts if part)
Expand All @@ -743,7 +746,7 @@ def _convert_page_basic(self, page_data: Any, vision: bool) -> str:
vision: If True, include the first image in Markdown format.
"""
if not isinstance(page_data, dict):
print(f"Warning: Unexpected page data type: {type(page_data)}. Converting to string.")
logger.warning("Unexpected page data type: %s. Converting to string.", type(page_data))
return str(page_data)

text_content = page_data.get("content", "")
Expand All @@ -760,9 +763,9 @@ def _convert_page_basic(self, page_data: Any, vision: bool) -> str:
# Basic image tag, assuming PNG
image_md = f"\n![Page Image](data:image/png;base64,{b64_img})\n"
else:
print(f"Warning: Image data is not in bytes format: {type(first_image_data)}")
logger.warning("Image data is not in bytes format: %s", type(first_image_data))
except Exception as e:
print(f"Error processing image on page: {e}")
logger.error("Error processing image on page: %s", e)
image_md = "\n<!-- Error processing image -->\n"

separator = "\n" if text_content and image_md else ""
Expand Down
5 changes: 4 additions & 1 deletion extract_thinker/pagination_handler.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -6,6 +6,9 @@
from extract_thinker.utils import encode_image, json_to_formatted_string, make_all_fields_optional
import yaml
from concurrent.futures import ThreadPoolExecutor, as_completed
import logging

logger = logging.getLogger(__name__)

class ConflictResolution(BaseModel):
resolved_fields: Dict[str, Dict[str, Any]] = Field(
Expand Down Expand Up @@ -58,7 +61,7 @@ def handle(self,
results.append(result)
except Exception as e:
# Log error but continue processing other pages
print(f"Error processing page: {str(e)}")
logger.error("Error processing page: %s", e)

if not results:
raise ValueError("No valid results obtained from any page")
Expand Down
Loading