Files
HKUDS-RAG-Anything/raganything/processor.py
2025-07-22 01:58:42 +08:00

725 lines
27 KiB
Python

"""
Document processing functionality for RAGAnything
Contains methods for parsing documents and processing multimodal content
"""
import os
import time
import hashlib
import json
from typing import Dict, List, Any
from pathlib import Path
from raganything.parser import MineruParser, DoclingParser
from raganything.utils import (
separate_content,
insert_text_content,
get_processor_for_type,
)
class ProcessorMixin:
"""ProcessorMixin class containing document processing functionality for RAGAnything"""
def _generate_cache_key(
self, file_path: Path, parse_method: str = None, **kwargs
) -> str:
"""
Generate cache key based on file path and parsing configuration
Args:
file_path: Path to the file
parse_method: Parse method used
**kwargs: Additional parser parameters
Returns:
str: Cache key for the file and configuration
"""
# Get file modification time
mtime = file_path.stat().st_mtime
# Create configuration dict for cache key
config_dict = {
"file_path": str(file_path.absolute()),
"mtime": mtime,
"parser": self.config.parser,
"parse_method": parse_method or self.config.parse_method,
}
# Add relevant kwargs to config
relevant_kwargs = {
k: v
for k, v in kwargs.items()
if k
in [
"lang",
"device",
"start_page",
"end_page",
"formula",
"table",
"backend",
"source",
]
}
config_dict.update(relevant_kwargs)
# Generate hash from config
config_str = json.dumps(config_dict, sort_keys=True)
cache_key = hashlib.md5(config_str.encode()).hexdigest()
return cache_key
def _generate_content_based_doc_id(self, content_list: List[Dict[str, Any]]) -> str:
"""
Generate doc_id based on document content
Args:
content_list: Parsed content list
Returns:
str: Content-based document ID with doc- prefix
"""
from lightrag.utils import compute_mdhash_id
# Extract key content for ID generation
content_hash_data = []
for item in content_list:
if isinstance(item, dict):
# For text content, use the text
if item.get("type") == "text" and item.get("text"):
content_hash_data.append(item["text"].strip())
# For other content types, use key identifiers
elif item.get("type") == "image" and item.get("img_path"):
content_hash_data.append(f"image:{item['img_path']}")
elif item.get("type") == "table" and item.get("table_body"):
content_hash_data.append(f"table:{item['table_body']}")
elif item.get("type") == "equation" and item.get("text"):
content_hash_data.append(f"equation:{item['text']}")
else:
# For other types, use string representation
content_hash_data.append(str(item))
# Create a content signature
content_signature = "\n".join(content_hash_data)
# Generate doc_id from content
doc_id = compute_mdhash_id(content_signature, prefix="doc-")
return doc_id
async def _get_cached_result(
self, cache_key: str, file_path: Path, parse_method: str = None, **kwargs
) -> tuple[List[Dict[str, Any]], str] | None:
"""
Get cached parsing result if available and valid
Args:
cache_key: Cache key to look up
file_path: Path to the file for mtime check
parse_method: Parse method used
**kwargs: Additional parser parameters
Returns:
tuple[List[Dict[str, Any]], str] | None: (content_list, doc_id) or None if not found/invalid
"""
if not hasattr(self, "parse_cache") or self.parse_cache is None:
return None
try:
cached_data = await self.parse_cache.get_by_id(cache_key)
if not cached_data:
return None
# Check file modification time
current_mtime = file_path.stat().st_mtime
cached_mtime = cached_data.get("mtime", 0)
if current_mtime != cached_mtime:
self.logger.debug(f"Cache invalid - file modified: {cache_key}")
return None
# Check parsing configuration
cached_config = cached_data.get("parse_config", {})
current_config = {
"parser": self.config.parser,
"parse_method": parse_method or self.config.parse_method,
}
# Add relevant kwargs to current config
relevant_kwargs = {
k: v
for k, v in kwargs.items()
if k
in [
"lang",
"device",
"start_page",
"end_page",
"formula",
"table",
"backend",
"source",
]
}
current_config.update(relevant_kwargs)
if cached_config != current_config:
self.logger.debug(f"Cache invalid - config changed: {cache_key}")
return None
content_list = cached_data.get("content_list", [])
doc_id = cached_data.get("doc_id")
if content_list and doc_id:
self.logger.debug(
f"Found valid cached parsing result for key: {cache_key}"
)
return content_list, doc_id
else:
self.logger.debug(
f"Cache incomplete - missing content or doc_id: {cache_key}"
)
return None
except Exception as e:
self.logger.warning(f"Error accessing parse cache: {e}")
return None
async def _store_cached_result(
self,
cache_key: str,
content_list: List[Dict[str, Any]],
doc_id: str,
file_path: Path,
parse_method: str = None,
**kwargs,
) -> None:
"""
Store parsing result in cache
Args:
cache_key: Cache key to store under
content_list: Content list to cache
doc_id: Content-based document ID
file_path: Path to the file for mtime storage
parse_method: Parse method used
**kwargs: Additional parser parameters
"""
if not hasattr(self, "parse_cache") or self.parse_cache is None:
return
try:
# Get file modification time
file_mtime = file_path.stat().st_mtime
# Create parsing configuration
parse_config = {
"parser": self.config.parser,
"parse_method": parse_method or self.config.parse_method,
}
# Add relevant kwargs to config
relevant_kwargs = {
k: v
for k, v in kwargs.items()
if k
in [
"lang",
"device",
"start_page",
"end_page",
"formula",
"table",
"backend",
"source",
]
}
parse_config.update(relevant_kwargs)
cache_data = {
cache_key: {
"content_list": content_list,
"doc_id": doc_id,
"mtime": file_mtime,
"parse_config": parse_config,
"cached_at": time.time(),
"cache_version": "1.0",
}
}
await self.parse_cache.upsert(cache_data)
# Ensure data is persisted to disk
await self.parse_cache.index_done_callback()
self.logger.info(f"Stored parsing result in cache: {cache_key}")
except Exception as e:
self.logger.warning(f"Error storing to parse cache: {e}")
async def parse_document(
self,
file_path: str,
output_dir: str = None,
parse_method: str = None,
display_stats: bool = None,
**kwargs,
) -> tuple[List[Dict[str, Any]], str]:
"""
Parse document with caching support
Args:
file_path: Path to the file to parse
output_dir: Output directory (defaults to config.parser_output_dir)
parse_method: Parse method (defaults to config.parse_method)
display_stats: Whether to display content statistics (defaults to config.display_content_stats)
**kwargs: Additional parameters for parser (e.g., lang, device, start_page, end_page, formula, table, backend, source)
Returns:
tuple[List[Dict[str, Any]], str]: (content_list, doc_id)
"""
# Use config defaults if not provided
if output_dir is None:
output_dir = self.config.parser_output_dir
if parse_method is None:
parse_method = self.config.parse_method
if display_stats is None:
display_stats = self.config.display_content_stats
self.logger.info(f"Starting document parsing: {file_path}")
file_path = Path(file_path)
if not file_path.exists():
raise FileNotFoundError(f"File not found: {file_path}")
# Generate cache key based on file and configuration
cache_key = self._generate_cache_key(file_path, parse_method, **kwargs)
# Check cache first
cached_result = await self._get_cached_result(
cache_key, file_path, parse_method, **kwargs
)
if cached_result is not None:
content_list, doc_id = cached_result
self.logger.info(f"Using cached parsing result for: {file_path}")
if display_stats:
self.logger.info(
f"* Total blocks in cached content_list: {len(content_list)}"
)
return content_list, doc_id
# Choose appropriate parsing method based on file extension
ext = file_path.suffix.lower()
try:
doc_parser = (
DoclingParser() if self.config.parser == "docling" else MineruParser()
)
# Log parser and method information
self.logger.info(
f"Using {self.config.parser} parser with method: {parse_method}"
)
if ext in [".pdf"]:
self.logger.info("Detected PDF file, using parser for PDF...")
content_list = doc_parser.parse_pdf(
pdf_path=file_path,
output_dir=output_dir,
method=parse_method,
**kwargs,
)
elif ext in [
".jpg",
".jpeg",
".png",
".bmp",
".tiff",
".tif",
".gif",
".webp",
]:
self.logger.info("Detected image file, using parser for images...")
# Use the selected parser's image parsing capability
if hasattr(doc_parser, "parse_image"):
content_list = doc_parser.parse_image(
image_path=file_path, output_dir=output_dir, **kwargs
)
else:
# Fallback to MinerU for image parsing if current parser doesn't support it
self.logger.warning(
f"{self.config.parser} parser doesn't support image parsing, falling back to MinerU"
)
content_list = MineruParser().parse_image(
image_path=file_path, output_dir=output_dir, **kwargs
)
elif ext in [
".doc",
".docx",
".ppt",
".pptx",
".xls",
".xlsx",
".html",
".htm",
".xhtml",
]:
self.logger.info(
"Detected Office or HTML document, using parser for Office/HTML..."
)
content_list = doc_parser.parse_office_doc(
doc_path=file_path, output_dir=output_dir, **kwargs
)
else:
# For other or unknown formats, use generic parser
self.logger.info(
f"Using generic parser for {ext} file (method={parse_method})..."
)
content_list = doc_parser.parse_document(
file_path=file_path,
method=parse_method,
output_dir=output_dir,
**kwargs,
)
except Exception as e:
self.logger.error(
f"Error during parsing with {self.config.parser} parser: {str(e)}"
)
self.logger.warning("Falling back to MinerU parser...")
# If specific parser fails, fall back to MinerU parser
content_list = MineruParser().parse_document(
file_path=file_path,
method=parse_method,
output_dir=output_dir,
**kwargs,
)
self.logger.info(
f"Parsing complete! Extracted {len(content_list)} content blocks"
)
# Generate doc_id based on content
doc_id = self._generate_content_based_doc_id(content_list)
# Store result in cache
await self._store_cached_result(
cache_key, content_list, doc_id, file_path, parse_method, **kwargs
)
# Display content statistics if requested
if display_stats:
self.logger.info("\nContent Information:")
self.logger.info(f"* Total blocks in content_list: {len(content_list)}")
# Count elements by type
block_types: Dict[str, int] = {}
for block in content_list:
if isinstance(block, dict):
block_type = block.get("type", "unknown")
if isinstance(block_type, str):
block_types[block_type] = block_types.get(block_type, 0) + 1
self.logger.info("* Content block types:")
for block_type, count in block_types.items():
self.logger.info(f" - {block_type}: {count}")
return content_list, doc_id
async def _process_multimodal_content(
self, multimodal_items: List[Dict[str, Any]], file_path: str, doc_id: str
):
"""
Process multimodal content (using specialized processors)
Args:
multimodal_items: List of multimodal items
file_path: File path (for reference)
doc_id: Document ID for proper chunk association
"""
if not multimodal_items:
self.logger.debug("No multimodal content to process")
return
# Check if multimodal content for this document is already processed
try:
existing_doc_status = await self.lightrag.doc_status.get_by_id(doc_id)
if existing_doc_status:
# Check if multimodal processing is already completed
multimodal_processed = existing_doc_status.get(
"multimodal_processed", False
)
existing_multimodal_chunks = existing_doc_status.get(
"multimodal_chunks_list", []
)
if multimodal_processed and len(existing_multimodal_chunks) >= len(
multimodal_items
):
self.logger.info(
f"Multimodal content already processed for document {doc_id} "
f"({len(existing_multimodal_chunks)} chunks found, {len(multimodal_items)} items to process)"
)
return
elif len(existing_multimodal_chunks) > 0:
self.logger.info(
f"Partial multimodal content found for document {doc_id} "
f"({len(existing_multimodal_chunks)} chunks exist, will reprocess all)"
)
except Exception as e:
self.logger.debug(f"Error checking multimodal cache for {doc_id}: {e}")
# Continue with processing if cache check fails
self.logger.info("Starting multimodal content processing...")
file_name = os.path.basename(file_path)
# Collect all chunk results for batch processing (similar to text content processing)
all_chunk_results = []
multimodal_chunk_ids = []
# Get current text chunks count to set proper order indexes for multimodal chunks
existing_doc_status = await self.lightrag.doc_status.get_by_id(doc_id)
existing_chunks_count = (
existing_doc_status.get("chunks_count", 0) if existing_doc_status else 0
)
for i, item in enumerate(multimodal_items):
try:
content_type = item.get("type", "unknown")
self.logger.info(
f"Processing item {i+1}/{len(multimodal_items)}: {content_type} content"
)
# Select appropriate processor
processor = get_processor_for_type(self.modal_processors, content_type)
if processor:
# Prepare item info for context extraction
item_info = {
"page_idx": item.get("page_idx", 0),
"index": i,
"type": content_type,
}
# Process content and get chunk results instead of immediately merging
(
enhanced_caption,
entity_info,
chunk_results,
) = await processor.process_multimodal_content(
modal_content=item,
content_type=content_type,
file_path=file_name,
item_info=item_info, # Pass item info for context extraction
batch_mode=True,
doc_id=doc_id, # Pass doc_id for proper association
chunk_order_index=existing_chunks_count
+ i
+ 1, # Proper order index
)
# Collect chunk results for batch processing
all_chunk_results.extend(chunk_results)
# Extract chunk ID from the entity_info (actual chunk_id created by processor)
if entity_info and "chunk_id" in entity_info:
chunk_id = entity_info["chunk_id"]
multimodal_chunk_ids.append(chunk_id)
self.logger.info(
f"{content_type} processing complete: {entity_info.get('entity_name', 'Unknown')}"
)
else:
self.logger.warning(
f"No suitable processor found for {content_type} type content"
)
except Exception as e:
self.logger.error(f"Error processing multimodal content: {str(e)}")
self.logger.debug("Exception details:", exc_info=True)
continue
# Update doc_status to include multimodal chunks
if multimodal_chunk_ids:
try:
# Get current document status
current_doc_status = await self.lightrag.doc_status.get_by_id(doc_id)
if current_doc_status:
existing_multimodal_chunks = current_doc_status.get(
"multimodal_chunks_list", []
)
# Combine existing chunks with new multimodal chunks
updated_multimodal_chunks_list = (
existing_multimodal_chunks + multimodal_chunk_ids
)
# Update document status with separated chunk lists
await self.lightrag.doc_status.upsert(
{
doc_id: {
**current_doc_status, # Keep existing fields
"multimodal_chunks_list": updated_multimodal_chunks_list, # Separated multimodal chunks
"multimodal_chunks_count": len(
updated_multimodal_chunks_list
),
"multimodal_processed": True, # Mark multimodal processing as complete
"updated_at": time.strftime("%Y-%m-%dT%H:%M:%S+00:00"),
}
}
)
# Ensure doc_status update is persisted to disk
await self.lightrag.doc_status.index_done_callback()
self.logger.info(
f"Updated doc_status with {len(multimodal_chunk_ids)} multimodal chunks"
)
except Exception as e:
self.logger.warning(
f"Error updating doc_status with multimodal chunks: {e}"
)
# Batch merge all multimodal content results (similar to text content processing)
if all_chunk_results:
from lightrag.operate import merge_nodes_and_edges
from lightrag.kg.shared_storage import (
get_namespace_data,
get_pipeline_status_lock,
)
# Get pipeline status and lock from shared storage
pipeline_status = await get_namespace_data("pipeline_status")
pipeline_status_lock = get_pipeline_status_lock()
await merge_nodes_and_edges(
chunk_results=all_chunk_results,
knowledge_graph_inst=self.lightrag.chunk_entity_relation_graph,
entity_vdb=self.lightrag.entities_vdb,
relationships_vdb=self.lightrag.relationships_vdb,
global_config=self.lightrag.__dict__,
pipeline_status=pipeline_status,
pipeline_status_lock=pipeline_status_lock,
llm_response_cache=self.lightrag.llm_response_cache,
current_file_number=1,
total_files=1,
file_path=file_name,
)
await self.lightrag._insert_done()
self.logger.info("Multimodal content processing complete")
async def process_document_complete(
self,
file_path: str,
output_dir: str = None,
parse_method: str = None,
display_stats: bool = None,
split_by_character: str | None = None,
split_by_character_only: bool = False,
doc_id: str | None = None,
**kwargs,
):
"""
Complete document processing workflow
Args:
file_path: Path to the file to process
output_dir: output directory (defaults to config.parser_output_dir)
parse_method: Parse method (defaults to config.parse_method)
display_stats: Whether to display content statistics (defaults to config.display_content_stats)
split_by_character: Optional character to split the text by
split_by_character_only: If True, split only by the specified character
doc_id: Optional document ID, if not provided will be generated from content
**kwargs: Additional parameters for parser (e.g., lang, device, start_page, end_page, formula, table, backend, source)
"""
# Ensure LightRAG is initialized
await self._ensure_lightrag_initialized()
# Use config defaults if not provided
if output_dir is None:
output_dir = self.config.parser_output_dir
if parse_method is None:
parse_method = self.config.parse_method
if display_stats is None:
display_stats = self.config.display_content_stats
self.logger.info(f"Starting complete document processing: {file_path}")
# Step 1: Parse document
content_list, content_based_doc_id = await self.parse_document(
file_path, output_dir, parse_method, display_stats, **kwargs
)
# Use provided doc_id or fall back to content-based doc_id
if doc_id is None:
doc_id = content_based_doc_id
# Step 2: Separate text and multimodal content
text_content, multimodal_items = separate_content(content_list)
# Step 2.5: Set content source for context extraction in multimodal processing
if hasattr(self, "set_content_source_for_context") and multimodal_items:
self.logger.info(
"Setting content source for context-aware multimodal processing..."
)
self.set_content_source_for_context(
content_list, self.config.content_format
)
# Step 3: Insert pure text content with all parameters
if text_content.strip():
file_name = os.path.basename(file_path)
await insert_text_content(
self.lightrag,
text_content,
file_paths=file_name,
split_by_character=split_by_character,
split_by_character_only=split_by_character_only,
ids=doc_id,
)
# Step 4: Process multimodal content (using specialized processors)
if multimodal_items:
await self._process_multimodal_content(multimodal_items, file_path, doc_id)
else:
# If no multimodal content, mark as processed to avoid future checks
try:
existing_doc_status = await self.lightrag.doc_status.get_by_id(doc_id)
if existing_doc_status and not existing_doc_status.get(
"multimodal_processed", False
):
existing_multimodal_chunks = existing_doc_status.get(
"multimodal_chunks_list", []
)
await self.lightrag.doc_status.upsert(
{
doc_id: {
**existing_doc_status,
"multimodal_chunks_list": existing_multimodal_chunks,
"multimodal_chunks_count": len(
existing_multimodal_chunks
),
"multimodal_processed": True,
"updated_at": time.strftime("%Y-%m-%dT%H:%M:%S+00:00"),
}
}
)
await self.lightrag.doc_status.index_done_callback()
self.logger.debug(
f"Marked document {doc_id[:8]}... as having no multimodal content"
)
except Exception as e:
self.logger.debug(
f"Error updating doc_status for no multimodal content: {e}"
)
self.logger.info(f"Document {file_path} processing complete!")