diff --git a/.env.example b/.env.example index 855ee71f3..32ea96da0 100644 --- a/.env.example +++ b/.env.example @@ -82,6 +82,13 @@ JWT_SECRET=weknora-jwt-secret # Embedding并发数，出现429错误时，可调小此参数 CONCURRENCY_POOL_SIZE=5 +# Docreader 并发任务数（图片OCR/Caption等异步任务），默认 1 +# 默认使用的 paddleocr 在高并发场景下，会出现异常，请谨慎设置 +# IMAGE_MAX_CONCURRENT=1 + +# Docreader OCR后端(no_ocr, paddle, nanonets) +# OCR_BACKEND=paddle + # 如果使用ElasticSearch作为向量存储，需要配置以下参数 # ElasticSearch地址，例如 http://localhost:9200 # ELASTICSEARCH_ADDR=your_elasticsearch_addr diff --git a/docker-compose.dev.yml b/docker-compose.dev.yml index a4dc79020..7d1667488 100644 --- a/docker-compose.dev.yml +++ b/docker-compose.dev.yml @@ -120,6 +120,8 @@ services: - WEB_PROXY=${WEB_PROXY:-} # File size limit (in MB) - MAX_FILE_SIZE_MB=${MAX_FILE_SIZE_MB:-50} + - IMAGE_MAX_CONCURRENT=${IMAGE_MAX_CONCURRENT:-1} + - OCR_BACKEND=${OCR_BACKEND:-paddle} healthcheck: test: ["CMD", "grpc_health_probe", "-addr=:50051"] interval: 30s diff --git a/docker-compose.yml b/docker-compose.yml index 00a6d93e6..14b28fe51 100644 --- a/docker-compose.yml +++ b/docker-compose.yml @@ -148,6 +148,8 @@ services: - MINERU_ENDPOINT=${MINERU_ENDPOINT:-} # File size limit (in MB) - MAX_FILE_SIZE_MB=${MAX_FILE_SIZE_MB:-50} + - IMAGE_MAX_CONCURRENT=${IMAGE_MAX_CONCURRENT:-1} + - OCR_BACKEND=${OCR_BACKEND:-paddle} healthcheck: test: ["CMD", "grpc_health_probe", "-addr=:50051"] interval: 30s diff --git a/docreader/parser/base_parser.py b/docreader/parser/base_parser.py index 4a7fb19a1..781830dab 100644 --- a/docreader/parser/base_parser.py +++ b/docreader/parser/base_parser.py @@ -35,34 +35,38 @@ class BaseParser(ABC): @staticmethod def _is_safe_url(url: str) -> bool: """Validate URL to prevent SSRF attacks - + Args: url: URL to validate - + Returns: True if URL is safe to fetch, False otherwise """ try: parsed_url = urlparse(url) - + # Only allow http and https schemes - if parsed_url.scheme not in ['http', 'https']: + if parsed_url.scheme not in ["http", "https"]: logger.warning(f"Rejected URL with invalid scheme: {parsed_url.scheme}") return False - + # Extract hostname hostname = parsed_url.hostname if not hostname: logger.warning("No hostname found in URL") return False - + # Try to parse as IP address try: ip = ipaddress.ip_address(hostname) # Reject private, loopback, link-local, multicast addresses - if (ip.is_private or ip.is_loopback or - ip.is_link_local or ip.is_multicast or - ip.is_reserved): + if ( + ip.is_private + or ip.is_loopback + or ip.is_link_local + or ip.is_multicast + or ip.is_reserved + ): logger.warning(f"Rejected URL with restricted IP: {ip}") return False except ValueError: @@ -70,17 +74,19 @@ class BaseParser(ABC): # Reject localhost and common internal hostnames hostname_lower = hostname.lower() restricted_hostnames = [ - 'localhost', - '127.0.0.1', - '::1', - 'metadata.google.internal', - 'metadata.tencentyun.com', - '169.254.169.254', # AWS metadata endpoint + "localhost", + "127.0.0.1", + "::1", + "metadata.google.internal", + "metadata.tencentyun.com", + "169.254.169.254", # AWS metadata endpoint ] - if hostname_lower in restricted_hostnames or hostname_lower.endswith('.local'): + if hostname_lower in restricted_hostnames or hostname_lower.endswith( + ".local" + ): logger.warning(f"Rejected URL with restricted hostname: {hostname}") return False - + return True except Exception as e: logger.warning(f"Error validating URL: {str(e)}") @@ -121,7 +127,7 @@ class BaseParser(ABC): chunk_size: int = 1000, chunk_overlap: int = 200, separators: list[str] = ["\n\n", "\n", "。"], - ocr_backend: str = "paddle", + ocr_backend: str = "no_ocr", ocr_config: dict = {}, max_image_size: int = 1920, # Maximum image size max_concurrent_tasks: int = 5, # Max concurrent tasks @@ -151,7 +157,7 @@ class BaseParser(ABC): self.chunk_size = chunk_size self.chunk_overlap = chunk_overlap self.separators = separators - self.ocr_backend = os.getenv("OCR_BACKEND", ocr_backend) + self.ocr_backend = ocr_backend self.ocr_config = ocr_config self.max_image_size = max_image_size self.max_concurrent_tasks = max_concurrent_tasks @@ -161,13 +167,14 @@ class BaseParser(ABC): self.chunking_config.storage_config if self.chunking_config else None ) - logger.info(f"Initializing parser for file: {file_name}, type: {file_type}") logger.info( - f"Parser config: chunk_size={chunk_size}, " + f"Initializing parser for file={file_name}, type={file_type}, " + f"parser config: chunk_size={chunk_size}, " f"overlap={chunk_overlap}, " f"multimodal={enable_multimodal}, " + f"max_chunks={max_chunks}, " f"ocr_backend={ocr_backend}, " - f"max_chunks={max_chunks}" + f"max_concurrent_tasks={max_concurrent_tasks}" ) # Only initialize Caption service if multimodal is enabled vlm_config = self.chunking_config.vlm_config if self.chunking_config else None @@ -305,13 +312,8 @@ class BaseParser(ABC): logger.warning("No image data to process") return [] - # Set max concurrency, reduce concurrency to avoid resource contention - max_concurrency = min( - self.max_concurrent_tasks, 1 - ) # Reduce concurrency to prevent excessive memory usage - # Use semaphore to limit concurrency - semaphore = asyncio.Semaphore(max_concurrency) + semaphore = asyncio.Semaphore(self.max_concurrent_tasks) # Store results to avoid overall failure due to task failure results = [] @@ -741,7 +743,7 @@ class BaseParser(ABC): if not self._is_safe_url(img_url): logger.error(f"URL failed validation check: {img_url}") return img_url, img_url, None - + # Still need to get image object for OCR processing # Get proxy settings from environment variables http_proxy = os.environ.get("EXTERNAL_HTTP_PROXY") @@ -792,7 +794,7 @@ class BaseParser(ABC): if not self._is_safe_url(img_url): logger.error(f"URL failed validation check: {img_url}") return img_url, img_url, None - + # Get proxy settings from environment variables http_proxy = os.environ.get("EXTERNAL_HTTP_PROXY") https_proxy = os.environ.get("EXTERNAL_HTTPS_PROXY") @@ -952,10 +954,8 @@ class BaseParser(ABC): # Create and run all Chunk concurrent processing tasks async def process_all_chunks(): - # Set max concurrency, reduce concurrency to avoid resource contention - max_concurrency = min(self.max_concurrent_tasks, 1) # Reduce concurrency # Use semaphore to limit concurrency - semaphore = asyncio.Semaphore(max_concurrency) + semaphore = asyncio.Semaphore(self.max_concurrent_tasks) async def process_with_limit(chunk, idx, total): """Use semaphore to control concurrent processing of Chunks""" diff --git a/docreader/parser/parser.py b/docreader/parser/parser.py index 2cb0da223..204364859 100644 --- a/docreader/parser/parser.py +++ b/docreader/parser/parser.py @@ -1,4 +1,5 @@ import logging +import os from typing import Dict, Type from docreader.models.document import Document @@ -109,8 +110,9 @@ class Parser: 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 + max_concurrent_tasks=int(os.getenv("IMAGE_MAX_CONCURRENT", "1")), + ocr_backend=os.getenv("OCR_BACKEND", "paddle"), ) logger.info(f"Starting to parse file content, size: {len(content)} bytes") @@ -155,8 +157,9 @@ class Parser: 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 + max_concurrent_tasks=int(os.getenv("IMAGE_MAX_CONCURRENT", "1")), + ocr_backend=os.getenv("OCR_BACKEND", "paddle"), ) logger.info("Starting to parse URL content")