style: add necessary comments to improve code quality

- Add docstrings and inline comments for key functions and complex logic
- Unify comment style, eliminate magic numbers and ambiguous variable names
- No functional changes, only improve maintainability
This commit is contained in:
begoniezhao
2025-12-01 17:43:26 +08:00
parent be411affdb
commit 3e31fdeefd
15 changed files with 1065 additions and 193 deletions
+23 -8
View File
@@ -181,17 +181,22 @@ class Caption:
from parameters or environment variables.
"""
logger.info("Initializing Caption service")
# Default prompt for image captioning in Chinese: "Briefly describe the main content of the image"
self.prompt = """简单凝炼的描述图片的主要内容"""
# API request timeout in seconds
self.timeout = 30
# Use provided VLM config if available,
# otherwise fall back to environment variables
if vlm_config and vlm_config.get("base_url") and vlm_config.get("model_name"):
# Build completion URL from provided base URL
self.completion_url = vlm_config.get("base_url", "") + "/chat/completions"
self.model = vlm_config.get("model_name", "")
self.api_key = vlm_config.get("api_key", "")
# Interface type: "ollama" or "openai" (default)
self.interface_type = vlm_config.get("interface_type", "openai").lower()
else:
# Fall back to environment variables if config not provided
base_url = os.getenv("VLM_MODEL_BASE_URL")
model_name = os.getenv("VLM_MODEL_NAME")
if not base_url or not model_name:
@@ -202,7 +207,7 @@ class Caption:
self.api_key = os.getenv("VLM_MODEL_API_KEY", "")
self.interface_type = os.getenv("VLM_INTERFACE_TYPE", "openai").lower()
# 验证接口类型
# Validate interface type - must be either "ollama" or "openai"
if self.interface_type not in ["ollama", "openai"]:
logger.warning(
f"Unknown interface type: {self.interface_type}, defaulting to openai"
@@ -227,7 +232,7 @@ class Caption:
logger.info("Calling Caption API for image captioning")
logger.info(f"Processing image data: {image_data[:50]}...")
# 根据接口类型选择调用方式
# Route to appropriate API based on interface type
if self.interface_type == "ollama":
return self._call_ollama_api(image_data)
else:
@@ -236,8 +241,10 @@ class Caption:
def _call_ollama_api(self, image_base64: str) -> Optional[CaptionChatResp]:
"""Call Ollama API for image captioning using base64 encoded image data."""
# Extract host URL by removing the chat completions endpoint
host = self.completion_url.replace("/v1/chat/completions", "")
# Initialize Ollama client with host and timeout
client = ollama.Client(
host=host,
timeout=self.timeout,
@@ -246,16 +253,17 @@ class Caption:
try:
logger.info(f"Calling Ollama API with model: {self.model}")
# 调用Ollama API,使用images参数传递base64编码的图片
# Call Ollama API with base64 encoded image
# Prompt: "Briefly describe the main content of the image"
response = client.generate(
model=self.model,
prompt="简单凝炼的描述图片的主要内容",
images=[image_base64], # image_base64是base64编码的图片数据
options={"temperature": 0.1},
images=[image_base64], # Pass base64 encoded image data
options={"temperature": 0.1}, # Low temperature for more deterministic output
stream=False,
)
# 构造响应对象
# Construct response object in standard format
caption_resp = CaptionChatResp(
id="ollama_response",
created=int(time.time()),
@@ -277,6 +285,7 @@ class Caption:
"""Call OpenAI-compatible API for image captioning."""
logger.info(f"Calling OpenAI-compatible API with model: {self.model}")
# Construct user message with text prompt and base64 encoded image
user_msg = UserMessage(
role="user",
content=[
@@ -290,20 +299,23 @@ class Caption:
],
)
# Build completion request with model parameters
gpt_req = CompletionRequest(
model=self.model,
temperature=0.3,
top_p=0.8,
temperature=0.3, # Moderate randomness for balanced output
top_p=0.8, # Nucleus sampling parameter
messages=[user_msg],
user="abc",
)
# Set up HTTP headers for the API request
headers = {
"Content-Type": "application/json",
"Accept": "text/event-stream",
"Cache-Control": "no-cache",
"Connection": "keep-alive",
}
# Add authorization header if API key is provided
if self.api_key:
headers["Authorization"] = f"Bearer {self.api_key}"
@@ -311,12 +323,14 @@ class Caption:
logger.info(
f"Sending request to OpenAI-compatible API with model: {self.model}"
)
# Send POST request to the API endpoint
response = requests.post(
self.completion_url,
data=json.dumps(gpt_req, default=lambda o: o.__dict__, indent=4),
headers=headers,
timeout=self.timeout,
)
# Check for successful response
if response.status_code != 200:
logger.error(
f"OpenAI API returned non-200 status code: {response.status_code}"
@@ -325,6 +339,7 @@ class Caption:
logger.info(f"Received from OpenAI with status: {response.status_code}")
logger.info("Converting response to CaptionChatResp object")
# Parse JSON response into structured object
caption_resp = CaptionChatResp.from_json(response.json())
if caption_resp.usage:
+100 -3
View File
@@ -1,3 +1,11 @@
"""
Chain Parser Module
This module provides two chain-of-responsibility pattern implementations for document parsing:
1. FirstParser: Tries multiple parsers sequentially until one succeeds
2. PipelineParser: Chains parsers where each parser processes the output of the previous one
"""
import logging
from typing import Dict, List, Tuple, Type
@@ -10,17 +18,43 @@ logger.setLevel(logging.INFO)
class FirstParser(BaseParser):
"""
First-success parser that tries multiple parsers in sequence.
This parser attempts to parse content using each registered parser in order.
It returns the result from the first parser that successfully produces a valid document.
If all parsers fail, it returns an empty Document.
Usage:
# Create a custom FirstParser with specific parser classes
CustomParser = FirstParser.create(MarkdownParser, HTMLParser)
parser = CustomParser()
document = parser.parse_into_text(content_bytes)
"""
# Tuple of parser classes to be instantiated
_parser_cls: Tuple[Type["BaseParser"], ...] = ()
def __init__(self, *args, **kwargs):
"""Initialize FirstParser with configured parser classes."""
super().__init__(*args, **kwargs)
# Instantiate all parser classes into parser instances
self._parsers: List[BaseParser] = []
for parser_cls in self._parser_cls:
parser = parser_cls(*args, **kwargs)
self._parsers.append(parser)
def parse_into_text(self, content: bytes) -> Document:
"""Parse content using the first parser that succeeds.
Args:
content: Raw bytes content to be parsed
Returns:
Document: Parsed document from the first successful parser,
or an empty Document if all parsers fail
"""
for p in self._parsers:
logger.info(f"FirstParser: using parser {p.__class__.__name__}")
document = p.parse_into_text(content)
@@ -31,41 +65,104 @@ class FirstParser(BaseParser):
@classmethod
def create(cls, *parser_classes: Type["BaseParser"]) -> Type["FirstParser"]:
"""Factory method to create a FirstParser subclass with specific parsers.
Args:
*parser_classes: Variable number of BaseParser subclasses to try in order
Returns:
Type[FirstParser]: A new FirstParser subclass configured with the given parsers
Example:
CustomParser = FirstParser.create(MarkdownParser, HTMLParser)
parser = CustomParser()
"""
# Generate a descriptive class name based on parser names
names = "_".join([p.__name__ for p in parser_classes])
# Dynamically create a new class with the parser configuration
return type(f"FirstParser_{names}", (cls,), {"_parser_cls": parser_classes})
class PipelineParser(BaseParser):
"""
Pipeline parser that chains multiple parsers sequentially.
This parser processes content through a series of parsers where each parser
receives the output of the previous parser as input. Images from all parsers
are accumulated and merged into the final document.
Usage:
# Create a custom PipelineParser with specific parser classes
CustomParser = PipelineParser.create(PreParser, MarkdownParser, PostParser)
parser = CustomParser()
document = parser.parse_into_text(content_bytes)
"""
# Tuple of parser classes to be instantiated and chained
_parser_cls: Tuple[Type["BaseParser"], ...] = ()
def __init__(self, *args, **kwargs):
"""Initialize PipelineParser with configured parser classes."""
super().__init__(*args, **kwargs)
# Instantiate all parser classes into parser instances
self._parsers: List[BaseParser] = []
for parser_cls in self._parser_cls:
parser = parser_cls(*args, **kwargs)
self._parsers.append(parser)
def parse_into_text(self, content: bytes) -> Document:
"""Parse content through a pipeline of parsers.
Each parser in the pipeline processes the output of the previous parser.
Images from all parsers are accumulated and merged into the final document.
Args:
content: Raw bytes content to be parsed
Returns:
Document: Final document after processing through all parsers,
with accumulated images from all stages
"""
# Accumulate images from all parsers
images: Dict[str, str] = {}
document = Document()
for p in self._parsers:
logger.info(f"PipelineParser: using parser {p.__class__.__name__}")
# Parse content with current parser
document = p.parse_into_text(content)
# Convert document content back to bytes for next parser
content = endecode.encode_bytes(document.content)
# Accumulate images from this parser
images.update(document.images)
# Merge all accumulated images into final document
document.images.update(images)
return document
@classmethod
def create(cls, *parser_classes: Type["BaseParser"]) -> Type["PipelineParser"]:
"""Factory method to create a PipelineParser subclass with specific parsers.
Args:
*parser_classes: Variable number of BaseParser subclasses to chain in order
Returns:
Type[PipelineParser]: A new PipelineParser subclass configured with the given parsers
Example:
CustomParser = PipelineParser.create(PreprocessParser, MarkdownParser)
parser = CustomParser()
"""
# Generate a descriptive class name based on parser names
names = "_".join([p.__name__ for p in parser_classes])
# Dynamically create a new class with the parser configuration
return type(f"PipelineParser_{names}", (cls,), {"_parser_cls": parser_classes})
if __name__ == "__main__":
from docreader.parser.markdown_parser import MarkdownParser
cls = FirstParser.create(MarkdownParser)
parser = cls()
print(parser.parse_into_text(b"aaa"))
# Example: Create and use a FirstParser with MarkdownParser
FpCls = FirstParser.create(MarkdownParser)
lparser = FpCls()
print(lparser.parse_into_text(b"aaa"))
+51
View File
@@ -1,3 +1,10 @@
"""
CSV Parser Module
This module provides a parser for CSV (Comma-Separated Values) files.
It converts CSV data into a Document with structured chunks, where each row
becomes a separate chunk with key-value pairs.
"""
import logging
from io import BytesIO
from typing import List
@@ -11,23 +18,64 @@ logger = logging.getLogger(__name__)
class CSVParser(BaseParser):
"""
Parser for CSV files that converts tabular data into structured text.
This parser reads CSV content and transforms each row into a formatted string
with column-value pairs. Each row is stored as a separate Chunk in the Document,
allowing for granular access to individual records.
The output format for each row is:
"column1: value1, column2: value2, column3: value3\n"
Usage:
parser = CSVParser()
with open("data.csv", "rb") as f:
document = parser.parse_into_text(f.read())
"""
def parse_into_text(self, content: bytes) -> Document:
"""Parse CSV content into a Document with structured chunks.
Each row in the CSV is converted into a formatted string and stored as
a separate Chunk. The chunks maintain sequential order and track their
position in the overall document.
Args:
content: Raw bytes content of the CSV file
Returns:
Document: A Document object containing:
- content: Full text with all rows concatenated
- chunks: List of Chunk objects, one per CSV row
Note:
Bad lines in the CSV are automatically skipped using pandas'
on_bad_lines="skip" parameter.
"""
chunks: List[Chunk] = []
text: List[str] = []
start, end = 0, 0
# Read CSV content into a pandas DataFrame, skipping malformed lines
df = pd.read_csv(BytesIO(content), on_bad_lines="skip")
# Process each row in the DataFrame
for i, (idx, row) in enumerate(df.iterrows()):
# Format row as "column: value" pairs separated by commas
content_row = (
",".join(
f"{col.strip()}: {str(row[col]).strip()}" for col in df.columns
)
+ "\n"
)
# Update end position for this chunk
end += len(content_row)
text.append(content_row)
# Create a chunk for this row with position tracking
chunks.append(Chunk(content=content_row, seq=i, start=start, end=end))
# Update start position for next chunk
start = end
return Document(
@@ -37,6 +85,7 @@ class CSVParser(BaseParser):
if __name__ == "__main__":
# Example usage: Parse a CSV file and display its content
logging.basicConfig(level=logging.DEBUG)
your_file = "/path/to/your/file.csv"
@@ -44,7 +93,9 @@ if __name__ == "__main__":
with open(your_file, "rb") as f:
content = f.read()
document = parser.parse_into_text(content)
# Display full document content
logger.error(document.content)
# Display individual chunks (rows)
for chunk in document.chunks:
logger.error(chunk.content)
+67 -2
View File
@@ -1,3 +1,10 @@
"""
Excel Parser Module
This module provides functionality to parse Excel files (.xlsx, .xls) into
structured Document objects with text content and chunks. It supports multiple
sheets and handles various Excel formats using pandas.
"""
import logging
from io import BytesIO
from typing import List
@@ -11,44 +18,102 @@ logger = logging.getLogger(__name__)
class ExcelParser(BaseParser):
"""Parser for Excel files (.xlsx, .xls).
This parser extracts text content from Excel files by processing all sheets
and converting each row into a structured text format. Each row becomes a
separate chunk with key-value pairs.
Features:
- Supports multiple sheets in a single Excel file
- Automatically removes completely empty rows
- Converts each row to "column: value" format
- Creates individual chunks for each row for better granularity
Example:
>>> parser = ExcelParser()
>>> with open("data.xlsx", "rb") as f:
... content = f.read()
... document = parser.parse_into_text(content)
>>> print(document.content)
Name: John,Age: 30,City: NYC
Name: Jane,Age: 25,City: LA
"""
def parse_into_text(self, content: bytes) -> Document:
"""Parse Excel file bytes into a Document object.
Args:
content: Raw bytes of the Excel file
Returns:
Document: Parsed document containing:
- content: Full text with all rows from all sheets
- chunks: List of Chunk objects, one per row
Note:
- Empty rows (all NaN values) are automatically skipped
- Each row is formatted as: "col1: val1,col2: val2,..."
- Chunks maintain sequential ordering across all sheets
"""
chunks: List[Chunk] = []
text: List[str] = []
start, end = 0, 0
# Load Excel file from bytes into pandas ExcelFile object
excel_file = pd.ExcelFile(BytesIO(content))
# Process each sheet in the Excel file
for excel_sheet_name in excel_file.sheet_names:
# Parse the sheet into a DataFrame
df = excel_file.parse(sheet_name=excel_sheet_name)
# Remove rows where all values are NaN (completely empty rows)
df.dropna(how="all", inplace=True)
# Process each row in the DataFrame
for _, row in df.iterrows():
page_content = []
# Build key-value pairs for non-null values
for k, v in row.items():
if pd.notna(v):
if pd.notna(v): # Skip NaN/null values
page_content.append(f"{k}: {v}")
# Skip rows with no valid content
if not page_content:
continue
# Format row as comma-separated key-value pairs
content_row = ",".join(page_content) + "\n"
end += len(content_row)
text.append(content_row)
# Create a chunk for this row with position tracking
chunks.append(
Chunk(content=content_row, seq=len(chunks), start=start, end=end)
)
start = end
# Combine all text and return as Document
return Document(content="".join(text), chunks=chunks)
if __name__ == "__main__":
# Example usage: Parse an Excel file and display results
logging.basicConfig(level=logging.DEBUG)
# Specify the path to your Excel file
your_file = "/path/to/your/file.xlsx"
parser = ExcelParser()
# Read and parse the Excel file
with open(your_file, "rb") as f:
content = f.read()
document = parser.parse_into_text(content)
# Display the full document content
logger.error(document.content)
# Display the first chunk as an example
for chunk in document.chunks:
logger.error(chunk.content)
break
break # Only show the first chunk
+229 -17
View File
@@ -1,3 +1,15 @@
"""
Markdown Parser Module
This module provides comprehensive Markdown parsing functionality including:
- Table formatting and standardization
- Base64 image extraction and conversion
- Image path replacement and URL generation
- Pipeline-based parsing with multiple stages
The parser uses a pipeline approach to process Markdown content through
multiple stages: table formatting -> image processing.
"""
import base64
import logging
import os
@@ -15,37 +27,76 @@ logger = logging.getLogger(__name__)
class MarkdownTableUtil:
"""Utility class for formatting Markdown tables.
This class standardizes Markdown table formatting by:
- Normalizing column alignment markers (e.g., :---, :---:, ---:)
- Adding consistent spacing around pipes (|)
- Preserving indentation levels
- Handling both header rows and data rows
Example:
Input: |姓名|年龄|城市|
|:---|---:|:---:|
|张三|25|北京|
Output: | 姓名 | 年龄 | 城市 |
| :--- | ---: | :---: |
| 张三 | 25 | 北京 |
"""
def __init__(self):
# Pattern to match alignment row (e.g., |:---|---:|:---:|)
self.align_pattern = re.compile(
r"^([\t ]*)\|[\t ]*[:-]+(?:[\t ]*\|[\t ]*[:-]+)*[\t ]*\|[\t ]*$",
re.MULTILINE,
)
# Pattern to match regular table rows (header or data)
self.line_pattern = re.compile(
r"^([\t ]*)\|[\t ]*[^|\r\n]*(?:[\t ]*\|[^|\r\n]*)*\|[\t ]*$",
re.MULTILINE,
)
def format_table(self, content: str) -> str:
"""Format all Markdown tables in the content.
Args:
content: Raw Markdown text containing tables
Returns:
Formatted Markdown text with standardized table formatting
"""
def process_align(match: Match[str]) -> str:
"""Process alignment row to standardize format."""
# Split by | and remove empty strings
columns = [col.strip() for col in match.group(0).split("|") if col.strip()]
processed = []
for col in columns:
# Preserve left alignment marker (:---)
left_colon = ":" if col.startswith(":") else ""
# Preserve right alignment marker (---:)
right_colon = ":" if col.endswith(":") else ""
processed.append(left_colon + "---" + right_colon)
# Preserve original indentation
prefix = match.group(1)
return prefix + "| " + " | ".join(processed) + " |"
def process_line(match: Match[str]) -> str:
"""Process regular table row to standardize format."""
# Split by | and remove empty strings
columns = [col.strip() for col in match.group(0).split("|") if col.strip()]
# Preserve original indentation
prefix = match.group(1)
return prefix + "| " + " | ".join(columns) + " |"
formatted_content = content
# First format regular rows (header and data)
formatted_content = self.line_pattern.sub(process_line, formatted_content)
# Then format alignment rows (must be done after to avoid conflicts)
formatted_content = self.align_pattern.sub(process_align, formatted_content)
return formatted_content
@@ -73,22 +124,64 @@ class MarkdownTableUtil:
class MarkdownTableFormatter(BaseParser):
"""Parser for formatting Markdown tables.
This parser standardizes the formatting of all Markdown tables in the
document to ensure consistent spacing and alignment markers.
Example:
>>> formatter = MarkdownTableFormatter()
>>> content = b"|Name|Age|\n|---|---|\n|John|30|"
>>> doc = formatter.parse_into_text(content)
>>> print(doc.content)
| Name | Age |
| --- | --- |
| John | 30 |
"""
def __init__(self, **kwargs):
super().__init__(**kwargs)
self.table_helper = MarkdownTableUtil()
def parse_into_text(self, content: bytes) -> Document:
"""Parse and format Markdown tables.
Args:
content: Raw Markdown content as bytes
Returns:
Document with formatted table content
"""
# Decode bytes to string with automatic encoding detection
text = endecode.decode_bytes(content)
# Format all tables in the content
text = self.table_helper.format_table(text)
return Document(content=text)
class MarkdownImageUtil:
"""Utility class for handling images in Markdown.
This class provides functionality to:
- Extract base64-encoded images from Markdown
- Extract image paths from Markdown
- Replace image paths with new URLs
- Convert base64 images to binary format
Supported formats:
- Base64 embedded images: ![alt](data:image/png;base64,iVBORw0...)
- Regular image links: ![alt](path/to/image.png)
"""
def __init__(self):
# Pattern to match base64 embedded images
# Captures: (1) alt text, (2) image format, (3) base64 data
self.b64_pattern = re.compile(
r"!\[([^\]]*)\]\(data:image/(\w+)\+?\w*;base64,([^\)]+)\)"
)
# Pattern to match regular image syntax
self.image_pattern = re.compile(r"!\[([^\]]*)\]\(([^)]+)\)")
# Pattern for replacing image paths
self.replace_pattern = re.compile(r"!\[([^\]]*)\]\(([^)]+)\)")
def extract_image(
@@ -97,23 +190,41 @@ class MarkdownImageUtil:
path_prefix: Optional[str] = None,
replace: bool = True,
) -> Tuple[str, List[str]]:
"""Extract base64 encoded images from Markdown content"""
# image_path => base64 bytes
"""Extract image paths from Markdown content.
Args:
content: Markdown text containing images
path_prefix: Optional prefix to add to image paths
replace: Whether to replace image syntax in content
Returns:
Tuple of (processed_text, list_of_image_paths)
Example:
>>> util = MarkdownImageUtil()
>>> text, images = util.extract_image("![logo](img/logo.png)")
>>> print(images)
['img/logo.png']
"""
# List to store extracted image paths
images: List[str] = []
def repl(match: Match[str]) -> str:
title = match.group(1)
image_path = match.group(2)
"""Replacement function for each image match."""
title = match.group(1) # Alt text
image_path = match.group(2) # Image path
# Add prefix if specified
if path_prefix:
image_path = f"{path_prefix}/{image_path}"
images.append(image_path)
# Keep original if replace is False
if not replace:
return match.group(0)
# Replace image path with URL
# Replace image path with potentially prefixed path
return f"![{title}]({image_path})"
text = self.image_pattern.sub(repl, content)
@@ -126,30 +237,55 @@ class MarkdownImageUtil:
path_prefix: Optional[str] = None,
replace: bool = True,
) -> Tuple[str, Dict[str, bytes]]:
"""Extract base64 encoded images from Markdown content"""
# image_path => base64 bytes
"""Extract and decode base64 embedded images from Markdown.
This method finds all base64-encoded images in the Markdown content,
decodes them to binary format, generates unique filenames, and
optionally replaces them with file path references.
Args:
content: Markdown text containing base64 images
path_prefix: Optional directory prefix for generated paths
replace: Whether to replace base64 syntax with file paths
Returns:
Tuple of (processed_text, dict_of_path_to_bytes)
Example:
>>> util = MarkdownImageUtil()
>>> text = "![logo](data:image/png;base64,iVBORw0KGg...)"
>>> new_text, images = util.extract_base64(text, "images")
>>> print(new_text)
![logo](images/uuid.png)
>>> print(len(images))
1
"""
# Dictionary mapping generated file paths to binary image data
images: Dict[str, bytes] = {}
def repl(match: Match[str]) -> str:
title = match.group(1)
img_ext = match.group(2)
img_b64 = match.group(3)
"""Replacement function for each base64 image match."""
title = match.group(1) # Alt text
img_ext = match.group(2) # Image format (png, jpg, etc.)
img_b64 = match.group(3) # Base64 encoded data
# Decode base64 string to bytes
image_byte = endecode.encode_image(img_b64, errors="ignore")
if not image_byte:
logger.error(f"Failed to decode base64 image skip it: {img_b64}")
return title
return title # Return just the alt text if decode fails
# Generate unique filename with original extension
image_path = f"{uuid.uuid4()}.{img_ext}"
if path_prefix:
image_path = f"{path_prefix}/{image_path}"
images[image_path] = image_byte
# Keep original base64 if replace is False
if not replace:
return match.group(0)
# Replace image path with URL
# Replace base64 data with file path reference
return f"![{title}]({image_path})"
text = self.b64_pattern.sub(repl, content)
@@ -157,15 +293,40 @@ class MarkdownImageUtil:
return text, images
def replace_path(self, content: str, images: Dict[str, str]) -> str:
"""Replace image paths in Markdown with new URLs.
This method is typically used to replace local file paths with
uploaded URLs after images have been stored.
Args:
content: Markdown text with image references
images: Mapping of old paths to new URLs
Returns:
Markdown text with updated image URLs
Example:
>>> util = MarkdownImageUtil()
>>> content = "![logo](temp/img.png)"
>>> mapping = {"temp/img.png": "https://cdn.com/img.png"}
>>> result = util.replace_path(content, mapping)
>>> print(result)
![logo](https://cdn.com/img.png)
"""
# Track which paths were actually replaced
content_replace: set = set()
def repl(match: Match[str]) -> str:
title = match.group(1)
image_path = match.group(2)
"""Replacement function for each image match."""
title = match.group(1) # Alt text
image_path = match.group(2) # Current image path
# Only replace if path exists in mapping
if image_path not in images:
return match.group(0)
return match.group(0) # Keep original
content_replace.add(image_path)
# Get new URL from mapping
image_path = images[image_path]
return f"![{title}]({image_path})"
@@ -186,43 +347,94 @@ class MarkdownImageUtil:
class MarkdownImageBase64(BaseParser):
"""Parser for extracting and uploading base64 images from Markdown.
This parser:
1. Extracts base64-encoded images from Markdown content
2. Uploads the decoded images to storage
3. Replaces base64 data with uploaded URLs
4. Returns a Document with clean Markdown and image mappings
Requires:
- self.storage: Storage backend for uploading images
Example:
>>> parser = MarkdownImageBase64(storage=my_storage)
>>> content = b"![logo](data:image/png;base64,iVBORw0...)"
>>> doc = parser.parse_into_text(content)
>>> print(doc.content)
![logo](https://storage.com/uuid.png)
"""
def __init__(self, **kwargs):
super().__init__(**kwargs)
self.image_helper = MarkdownImageUtil()
def parse_into_text(self, content: bytes) -> Document:
"""Parse Markdown and process base64 images.
Args:
content: Raw Markdown content as bytes
Returns:
Document with:
- content: Markdown with base64 replaced by URLs
- images: Dict mapping URLs to base64 strings
"""
# Convert byte content to string using universal decoding method
text = endecode.decode_bytes(content)
# Extract base64 images and replace with temporary paths
text, img_b64 = self.image_helper.extract_base64(text, path_prefix="images")
# Final image mapping: URL -> base64 string (for Document)
images: Dict[str, str] = {}
# Temporary mapping: temp_path -> uploaded_URL (for replacement)
image_replace: Dict[str, str] = {}
logger.debug(f"Uploading {len(img_b64)} images from markdown")
# Upload each extracted image to storage
for ipath, b64_bytes in img_b64.items():
# Get file extension for proper MIME type
ext = os.path.splitext(ipath)[1].lower()
# Upload binary data and get back URL
image_url = self.storage.upload_bytes(b64_bytes, ext)
# Map temp path to uploaded URL for replacement
image_replace[ipath] = image_url
# Store base64 string in final images dict
images[image_url] = base64.b64encode(b64_bytes).decode()
# Replace temporary paths with actual uploaded URLs
text = self.image_helper.replace_path(text, image_replace)
return Document(content=text, images=images)
class MarkdownParser(PipelineParser):
"""Complete Markdown parser using pipeline approach.
This parser processes Markdown content through multiple stages:
1. MarkdownTableFormatter: Standardizes table formatting
2. MarkdownImageBase64: Extracts and uploads base64 images
The pipeline ensures that content flows through each parser in sequence,
with each stage's output becoming the next stage's input.
"""
_parser_cls = (MarkdownTableFormatter, MarkdownImageBase64)
if __name__ == "__main__":
# Example usage and testing
logging.basicConfig(level=logging.DEBUG)
# Test the complete MarkdownParser pipeline
your_content = "test![](data:image/png;base64,iVBORw0KGgoAAAANSUhEUgAAAMgA)test"
parser = MarkdownParser()
# Parse content and display results
document = parser.parse_into_text(your_content.encode())
logger.info(document.content)
logger.info(f"Images: {len(document.images)}, name: {document.images.keys()}")
# Run individual utility tests
MarkdownImageUtil._self_test()
MarkdownTableUtil._self_test()
+75 -17
View File
@@ -16,20 +16,47 @@ logger = logging.getLogger(__name__)
class StdMinerUParser(BaseParser):
"""
Standard MinerU Parser for document parsing.
This parser uses MinerU API to parse documents (especially PDFs) into markdown format,
with support for tables, formulas, and images extraction.
"""
def __init__(
self,
enable_markdownify: bool = True,
mineru_endpoint: str = "",
**kwargs,
):
"""
Initialize MinerU parser.
Args:
enable_markdownify: Whether to convert HTML tables to markdown format
mineru_endpoint: MinerU API endpoint URL
**kwargs: Additional arguments passed to BaseParser
"""
super().__init__(**kwargs)
# Get MinerU endpoint from environment variable or parameter
self.minerU = os.getenv("MINERU_ENDPOINT", mineru_endpoint)
self.enable_markdownify = enable_markdownify
# Helper for processing markdown images
self.image_helper = MarkdownImageUtil()
# Pattern to match base64 encoded images
self.base64_pattern = re.compile(r"data:image/(\w+);base64,(.*)")
# Check if MinerU API is available
self.enable = self.ping()
def ping(self, timeout: int = 5) -> bool:
"""
Check if MinerU API is available.
Args:
timeout: Request timeout in seconds
Returns:
True if API is available, False otherwise
"""
try:
response = requests.get(
self.minerU + "/docs", timeout=timeout, allow_redirects=True
@@ -40,6 +67,15 @@ class StdMinerUParser(BaseParser):
return False
def parse_into_text(self, content: bytes) -> Document:
"""
Parse document content into text using MinerU API.
Args:
content: Raw document content in bytes
Returns:
Document object containing parsed text and images
"""
if not self.enable:
logger.debug("MinerU API is not enabled")
return Document()
@@ -48,22 +84,23 @@ class StdMinerUParser(BaseParser):
md_content: str = ""
images_b64: Dict[str, str] = {}
try:
# Call MinerU API to parse document
response = requests.post(
url=self.minerU + "/file_parse",
data={
"return_md": True,
"return_images": True,
"lang_list": ["ch", "en"],
"table_enable": True,
"formula_enable": True,
"parse_method": "auto",
"start_page_id": 0,
"end_page_id": 99999,
"backend": "pipeline",
"response_format_zip": False,
"return_middle_json": False,
"return_model_output": False,
"return_content_list": False,
"return_md": True, # Return markdown content
"return_images": True, # Return extracted images
"lang_list": ["ch", "en"], # Support Chinese and English
"table_enable": True, # Enable table parsing
"formula_enable": True, # Enable formula parsing
"parse_method": "auto", # Auto detect parsing method
"start_page_id": 0, # Start from first page
"end_page_id": 99999, # Parse all pages
"backend": "pipeline", # Use pipeline backend
"response_format_zip": False, # Return JSON instead of ZIP
"return_middle_json": False, # Don't return intermediate JSON
"return_model_output": False, # Don't return model output
"return_content_list": False, # Don't return content list
},
files={"files": content},
timeout=1000,
@@ -76,38 +113,47 @@ class StdMinerUParser(BaseParser):
logger.error(f"MinerU parsing failed: {e}", exc_info=True)
return Document()
# convert table(HTML) in markdown to markdown table
# Convert HTML tables in markdown to markdown table format
if self.enable_markdownify:
logger.debug("Converting HTML to Markdown")
md_content = markdownify.markdownify(md_content)
images = {}
image_replace = {}
# image in images_bs64 may not be used in md_content
# such as: table ...
# so we need to filter them
# Filter images that are actually used in markdown content
# Some images in images_b64 may not be referenced in md_content
# (e.g., images embedded in tables), so we need to filter them
for ipath, b64_str in images_b64.items():
# Skip images that are not referenced in markdown content
if f"images/{ipath}" not in md_content:
logger.debug(f"Image {ipath} not used in markdown")
continue
# Parse base64 image data
match = self.base64_pattern.match(b64_str)
if match:
# Extract image format (e.g., png, jpg)
file_ext = match.group(1)
# Extract base64 encoded data
b64_str = match.group(2)
# Decode base64 string to bytes
image_bytes = endecode.encode_image(b64_str, errors="ignore")
if not image_bytes:
logger.error("Failed to decode base64 image skip it")
continue
# Upload image to storage and get URL
image_url = self.storage.upload_bytes(
image_bytes, file_ext=f".{file_ext}"
)
# Store image mapping for later use
images[image_url] = b64_str
# Prepare replacement mapping for markdown content
image_replace[f"images/{ipath}"] = image_url
logger.info(f"Replaced {len(image_replace)} images in markdown")
# Replace image paths in markdown with uploaded URLs
text = self.image_helper.replace_path(md_content, image_replace)
logger.info(
@@ -117,15 +163,27 @@ class StdMinerUParser(BaseParser):
class MinerUParser(PipelineParser):
"""
MinerU Parser with pipeline processing.
This parser combines StdMinerUParser for document parsing and
MarkdownTableFormatter for table formatting in a pipeline.
"""
_parser_cls = (StdMinerUParser, MarkdownTableFormatter)
if __name__ == "__main__":
# Example usage for testing
logging.basicConfig(level=logging.DEBUG)
# Configure your file path and MinerU endpoint
your_file = "/path/to/your/file.pdf"
your_mineru = "http://host.docker.internal:9987"
# Create parser instance
parser = MinerUParser(mineru_endpoint=your_mineru)
# Parse PDF file
with open(your_file, "rb") as f:
content = f.read()
document = parser.parse_into_text(content)
+16 -9
View File
@@ -52,8 +52,10 @@ class PaddleOCRBackend(OCRBackend):
os.environ["CUDA_VISIBLE_DEVICES"] = ""
paddle.device.set_device("cpu")
# Try to detect if CPU supports AVX instruction set
# 尝试检测CPU是否支持AVX指令集
try:
# Detect if CPU supports AVX
# 检测CPU是否支持AVX
if platform.system() == "Linux":
try:
@@ -69,6 +71,7 @@ class PaddleOCRBackend(OCRBackend):
"CPU does not support AVX instructions, "
"using compatibility mode"
)
# Further restrict instruction set usage
# 进一步限制指令集使用
os.environ["FLAGS_use_avx2"] = "0"
os.environ["FLAGS_use_avx"] = "1"
@@ -96,9 +99,9 @@ class PaddleOCRBackend(OCRBackend):
"use_gpu": False,
"text_det_limit_type": "max",
"text_det_limit_side_len": 960,
"use_doc_orientation_classify": True, # 启用文档方向分类
"use_doc_orientation_classify": True, # Enable document orientation classification / 启用文档方向分类
"use_doc_unwarping": False,
"use_textline_orientation": True, # 启用文本行方向检测
"use_textline_orientation": True, # Enable text line orientation detection / 启用文本行方向检测
"text_recognition_model_name": "PP-OCRv4_server_rec",
"text_detection_model_name": "PP-OCRv4_server_det",
"text_det_thresh": 0.3,
@@ -174,13 +177,13 @@ class PaddleOCRBackend(OCRBackend):
if image.mode != "RGB":
image = image.convert("RGB")
# Convert to numpy array if needed
# Convert to numpy array for PaddleOCR processing
image_array = np.array(image)
# Perform OCR
# Perform OCR recognition
ocr_result = self.ocr.ocr(image_array, cls=False)
# Extract text
# Extract and concatenate text from OCR results
ocr_text = ""
if ocr_result and ocr_result[0]:
text = [
@@ -209,6 +212,7 @@ class NanonetsOCRBackend(OCRBackend):
base_url: Base URL for OpenAI API
model: Model name
"""
# Load configuration from environment variables
base_url = os.getenv("OCR_API_BASE_URL", "http://localhost:8000/v1")
api_key = os.getenv("OCR_API_KEY", "123")
timeout = 30
@@ -218,6 +222,7 @@ class NanonetsOCRBackend(OCRBackend):
logger.info(f"Nanonets OCR engine initialized with model: {self.model}")
self.temperature = 0.0
self.max_tokens = 15000
# Prompt for OCR text extraction with specific formatting requirements
self.prompt = """## 任务说明
请从上传的文档中提取文字内容,严格按自然阅读顺序(从上到下,从左到右)输出,并遵循以下格式规范。
@@ -258,12 +263,12 @@ class NanonetsOCRBackend(OCRBackend):
return ""
try:
# Encode image to base64
# Encode image to base64 format for API transmission
img_base64 = endecode.decode_image(image)
if not img_base64:
return ""
# Call Nanonets OCR API
# Call Nanonets OCR API using OpenAI-compatible format
logger.info(f"Calling Nanonets OCR API with model: {self.model}")
response = self.client.chat.completions.create(
model=self.model,
@@ -294,13 +299,14 @@ class NanonetsOCRBackend(OCRBackend):
class OCREngine:
"""OCR Engine factory class"""
"""OCR Engine factory class for managing different OCR backend instances"""
# Singleton pattern: cache instances for each backend type
_instance: Dict[str, OCRBackend] = {}
@classmethod
def get_instance(cls, backend_type: str) -> OCRBackend:
"""Get OCR engine instance
"""Get OCR engine instance using factory pattern
Args:
backend_type: OCR backend type, one of: "paddle", "nanonets"
@@ -310,6 +316,7 @@ class OCREngine:
OCR engine instance or None if initialization fails
"""
backend_type = backend_type.lower()
# Return cached instance if already initialized
if cls._instance.get(backend_type):
return cls._instance[backend_type]
+31 -19
View File
@@ -24,13 +24,15 @@ class Parser:
"""
def __init__(self):
# Initialize all parser types
# Initialize all parser types - maps file extensions to their corresponding parser classes
self.parsers: Dict[str, Type[BaseParser]] = {
# Document formats
"docx": Docx2Parser,
"doc": DocParser,
"pdf": PDFParser,
"md": MarkdownParser,
"txt": TextParser,
# Image formats - all use the same ImageParser
"jpg": ImageParser,
"jpeg": ImageParser,
"png": ImageParser,
@@ -38,7 +40,9 @@ class Parser:
"bmp": ImageParser,
"tiff": ImageParser,
"webp": ImageParser,
# Alternative markdown extension
"markdown": MarkdownParser,
# Spreadsheet formats
"csv": CSVParser,
"xlsx": ExcelParser,
"xls": ExcelParser,
@@ -59,8 +63,10 @@ class Parser:
Returns:
Parser class for the file type, or None if unsupported
"""
# Look up parser by file type (case-insensitive)
parser = self.parsers.get(file_type.lower())
if not parser:
# Raise error if file type is not supported
raise ValueError(f"Unsupported file type: {file_type}")
return parser
@@ -90,31 +96,34 @@ class Parser:
f"multimodal={config.enable_multimodal}"
)
# Get appropriate parser for file type
# Get appropriate parser class for the file type
cls = self.get_parser(file_type)
# Parse file content
# Create parser instance with configuration
logger.info(f"Creating parser instance for {file_type} file")
parser = cls(
file_name=file_name,
file_type=file_type,
chunk_size=config.chunk_size,
chunk_overlap=config.chunk_overlap,
separators=config.separators,
enable_multimodal=config.enable_multimodal,
max_image_size=1920, # Limit image size to 1920px
max_concurrent_tasks=5, # Limit concurrent tasks to 5
chunking_config=config, # Pass the entire chunking config
chunk_size=config.chunk_size, # Size of each text chunk
chunk_overlap=config.chunk_overlap, # Overlap between consecutive chunks
separators=config.separators, # Text separators for chunking
enable_multimodal=config.enable_multimodal, # Enable image/multimodal processing
max_image_size=1920, # Limit image size to 1920px for performance
max_concurrent_tasks=5, # Limit concurrent tasks to 5 to avoid resource exhaustion
chunking_config=config, # Pass the entire chunking config for advanced options
)
logger.info(f"Starting to parse file content, size: {len(content)} bytes")
# Execute the parsing process
result = parser.parse(content)
# Validate parsing results and log warnings if needed
if not result.content:
logger.warning(f"Parser returned empty content for file: {file_name}")
elif not result.chunks:
logger.warning(f"Parser returned empty chunks for file: {file_name}")
elif result.chunks[0]:
# Log first chunk size for debugging
logger.info(f"First chunk content length: {len(result.chunks[0].content)}")
logger.info(f"Parsed file {file_name}, with {len(result.chunks)} chunks")
return result
@@ -137,27 +146,30 @@ class Parser:
f"overlap={config.chunk_overlap}, multimodal={config.enable_multimodal}"
)
# Create web parser instance
# Create web parser instance with configuration
logger.info("Creating WebParser instance")
parser = WebParser(
title=title,
chunk_size=config.chunk_size,
chunk_overlap=config.chunk_overlap,
separators=config.separators,
enable_multimodal=config.enable_multimodal,
max_image_size=1920, # Limit image size
max_concurrent_tasks=5, # Limit concurrent tasks
chunking_config=config,
title=title, # Webpage title for metadata
chunk_size=config.chunk_size, # Size of each text chunk
chunk_overlap=config.chunk_overlap, # Overlap between consecutive chunks
separators=config.separators, # Text separators for chunking
enable_multimodal=config.enable_multimodal, # Enable image/multimodal processing
max_image_size=1920, # Limit image size to 1920px for performance
max_concurrent_tasks=5, # Limit concurrent tasks to avoid resource exhaustion
chunking_config=config, # Pass the entire chunking config
)
logger.info("Starting to parse URL content")
# Parse URL content (encode URL string to bytes as required by parser interface)
result = parser.parse(url.encode())
# Validate parsing results and log warnings if needed
if not result.content:
logger.warning(f"Parser returned empty content for url: {url}")
elif not result.chunks:
logger.warning(f"Parser returned empty chunks for url: {url}")
elif result.chunks[0]:
# Log first chunk size for debugging
logger.info(f"First chunk content length: {len(result.chunks[0].content)}")
logger.info(f"Parsed url {url}, with {len(result.chunks)} chunks")
return result
+9
View File
@@ -4,4 +4,13 @@ from docreader.parser.mineru_parser import MinerUParser
class PDFParser(FirstParser):
"""PDF Parser using chain of responsibility pattern
Attempts to parse PDF files using multiple parser backends in order:
1. MinerUParser - Primary parser for PDF documents
2. MarkitdownParser - Fallback parser if MinerU fails
The first successful parser result will be returned.
"""
# Parser classes to try in order (chain of responsibility pattern)
_parser_cls = (MinerUParser, MarkitdownParser)
+11 -1
View File
@@ -244,8 +244,18 @@ class MinioStorage(Storage):
found = client.bucket_exists(bucket_name)
if not found:
client.make_bucket(bucket_name)
# Set public read policy for the bucket
policy = (
'{"Version":"2012-10-17","Statement":[{"Effect":"Allow","Principal":{"AWS":["*"]},"Action":["s3:GetBucketLocation","s3:ListBucket"],"Resource":["arn:aws:s3:::%s"]},{"Effect":"Allow","Principal":{"AWS":["*"]},"Action":["s3:GetObject"],"Resource":["arn:aws:s3:::%s/*"]}]}'
'{'
'"Version":"2012-10-17",'
'"Statement":['
'{"Effect":"Allow","Principal":{"AWS":["*"]},'
'"Action":["s3:GetBucketLocation","s3:ListBucket"],'
'"Resource":["arn:aws:s3:::%s"]},'
'{"Effect":"Allow","Principal":{"AWS":["*"]},'
'"Action":["s3:GetObject"],'
'"Resource":["arn:aws:s3:::%s/*"]}'
']}'
% (bucket_name, bucket_name)
)
client.set_bucket_policy(bucket_name, policy)
+45 -7
View File
@@ -15,19 +15,40 @@ logger = logging.getLogger(__name__)
class StdWebParser(BaseParser):
"""Web page parser"""
"""Standard web page parser using Playwright and Trafilatura.
This parser scrapes web pages using Playwright's WebKit browser and extracts
clean content using Trafilatura library. It supports proxy configuration and
converts HTML content to markdown format.
"""
def __init__(self, title: str, **kwargs):
"""Initialize the web parser.
Args:
title: Title of the web page to be used as file name
**kwargs: Additional arguments passed to BaseParser
"""
self.title = title
# Get proxy configuration from environment variable if available
self.proxy = os.environ.get("WEB_PROXY", "")
super().__init__(file_name=title, **kwargs)
logger.info(f"Initialized WebParser with title: {title}")
async def scrape(self, url: str) -> str:
"""Scrape web page content using Playwright.
Args:
url: The URL of the web page to scrape
Returns:
HTML content of the web page as string, empty string on error
"""
logger.info(f"Starting web page scraping for URL: {url}")
try:
async with async_playwright() as p:
kwargs = {}
# Configure proxy if available
if self.proxy:
kwargs["proxy"] = {"server": self.proxy}
logger.info("Launching WebKit browser")
@@ -36,6 +57,7 @@ class StdWebParser(BaseParser):
logger.info(f"Navigating to URL: {url}")
try:
# Navigate to URL with 30 second timeout
await page.goto(url, timeout=30000)
logger.info("Initial page load complete")
except Exception as e:
@@ -44,35 +66,40 @@ class StdWebParser(BaseParser):
return ""
logger.info("Retrieving page HTML content")
# Get the full HTML content of the page
content = await page.content()
logger.info(f"Retrieved {len(content)} bytes of HTML content")
await browser.close()
logger.info("Browser closed")
# Parse HTML content with BeautifulSoup
# Return raw HTML content for further processing
logger.info("Parsing HTML with BeautifulSoup")
logger.info("Successfully parsed HTML content")
return content
except Exception as e:
logger.error(f"Failed to scrape web page: {str(e)}")
# Return empty BeautifulSoup object on error
# Return empty string on error
return ""
def parse_into_text(self, content: bytes) -> Document:
"""Parse web page
"""Parse web page content into a Document object.
Args:
content: Web page content
content: URL encoded as bytes
Returns:
Parse result
Document object containing the parsed markdown content
"""
# Decode bytes to get the URL string
url = endecode.decode_bytes(content)
logger.info(f"Scraping web page: {url}")
# Run async scraping in sync context
chtml = asyncio.run(self.scrape(url))
# Extract clean content from HTML using Trafilatura
# Convert to markdown format with metadata, images, tables, and links
md_text = extract(
chtml,
output_format="markdown",
@@ -80,7 +107,7 @@ class StdWebParser(BaseParser):
include_images=True,
include_tables=True,
include_links=True,
deduplicate=True,
deduplicate=True, # Remove duplicate content
)
if not md_text:
logger.error("Failed to parse web page")
@@ -89,16 +116,27 @@ class StdWebParser(BaseParser):
class WebParser(PipelineParser):
"""Web parser using pipeline pattern.
This parser chains StdWebParser (for web scraping and HTML to markdown conversion)
with MarkdownParser (for markdown processing). The pipeline processes content
sequentially through both parsers.
"""
# Parser classes to be executed in sequence
_parser_cls = (StdWebParser, MarkdownParser)
if __name__ == "__main__":
# Configure logging for debugging
logging.basicConfig(level=logging.DEBUG)
logger.setLevel(logging.DEBUG)
# Example URL to scrape
url = "https://cloud.tencent.com/document/product/457/6759"
# Create parser instance and parse the web page
parser = WebParser(title="")
cc = parser.parse_into_text(url.encode())
# Save the parsed markdown content to file
with open("./tencent.md", "w") as f:
f.write(cc.content)