From 3e31fdeefd5e190ce2ff27296d59ddb93ffebe24 Mon Sep 17 00:00:00 2001 From: begoniezhao Date: Mon, 1 Dec 2025 17:43:26 +0800 Subject: [PATCH] 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 --- docreader/parser/caption.py | 31 +++- docreader/parser/chain_parser.py | 103 +++++++++++- docreader/parser/csv_parser.py | 51 ++++++ docreader/parser/excel_parser.py | 69 +++++++- docreader/parser/markdown_parser.py | 246 ++++++++++++++++++++++++++-- docreader/parser/mineru_parser.py | 92 +++++++++-- docreader/parser/ocr_engine.py | 25 ++- docreader/parser/parser.py | 50 +++--- docreader/parser/pdf_parser.py | 9 + docreader/parser/storage.py | 12 +- docreader/parser/web_parser.py | 52 +++++- docreader/splitter/splitter.py | 173 +++++++++++++++---- docreader/utils/endecode.py | 130 +++++++++++++-- docreader/utils/split.py | 56 ++++++- mcp-server/weknora_mcp_server.py | 159 +++++++++++------- 15 files changed, 1065 insertions(+), 193 deletions(-) diff --git a/docreader/parser/caption.py b/docreader/parser/caption.py index f3e6f69e..548381f9 100644 --- a/docreader/parser/caption.py +++ b/docreader/parser/caption.py @@ -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: diff --git a/docreader/parser/chain_parser.py b/docreader/parser/chain_parser.py index 7edac887..0b5dec48 100644 --- a/docreader/parser/chain_parser.py +++ b/docreader/parser/chain_parser.py @@ -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")) diff --git a/docreader/parser/csv_parser.py b/docreader/parser/csv_parser.py index f2a713d1..c2c445a9 100644 --- a/docreader/parser/csv_parser.py +++ b/docreader/parser/csv_parser.py @@ -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) diff --git a/docreader/parser/excel_parser.py b/docreader/parser/excel_parser.py index 6c9ac3c1..d1d20568 100644 --- a/docreader/parser/excel_parser.py +++ b/docreader/parser/excel_parser.py @@ -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 diff --git a/docreader/parser/markdown_parser.py b/docreader/parser/markdown_parser.py index 70084049..7b979b99 100644 --- a/docreader/parser/markdown_parser.py +++ b/docreader/parser/markdown_parser.py @@ -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() diff --git a/docreader/parser/mineru_parser.py b/docreader/parser/mineru_parser.py index ef1c5cae..c7eba34e 100644 --- a/docreader/parser/mineru_parser.py +++ b/docreader/parser/mineru_parser.py @@ -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) diff --git a/docreader/parser/ocr_engine.py b/docreader/parser/ocr_engine.py index 0a999b99..7ca25cb6 100644 --- a/docreader/parser/ocr_engine.py +++ b/docreader/parser/ocr_engine.py @@ -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] diff --git a/docreader/parser/parser.py b/docreader/parser/parser.py index ffba7d33..2cb0da22 100644 --- a/docreader/parser/parser.py +++ b/docreader/parser/parser.py @@ -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 diff --git a/docreader/parser/pdf_parser.py b/docreader/parser/pdf_parser.py index c17184dc..7e1e6dda 100644 --- a/docreader/parser/pdf_parser.py +++ b/docreader/parser/pdf_parser.py @@ -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) diff --git a/docreader/parser/storage.py b/docreader/parser/storage.py index 767ab0a9..87c3f2ee 100644 --- a/docreader/parser/storage.py +++ b/docreader/parser/storage.py @@ -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) diff --git a/docreader/parser/web_parser.py b/docreader/parser/web_parser.py index ce259e08..a7906902 100644 --- a/docreader/parser/web_parser.py +++ b/docreader/parser/web_parser.py @@ -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) diff --git a/docreader/splitter/splitter.py b/docreader/splitter/splitter.py index c3c7128f..a14fb731 100644 --- a/docreader/splitter/splitter.py +++ b/docreader/splitter/splitter.py @@ -1,4 +1,11 @@ -"""Token splitter.""" +"""Token splitter. + +This module provides text splitting functionality with support for: +- Configurable chunk size and overlap +- Protected regex patterns (e.g., math formulas, images, links, tables) +- Header tracking for context preservation +- Smart merging with overlap handling +""" import itertools import logging @@ -12,8 +19,9 @@ from docreader.splitter.header_hook import ( ) from docreader.utils.split import split_by_char, split_by_sep -DEFAULT_CHUNK_OVERLAP = 100 -DEFAULT_CHUNK_SIZE = 512 +# Default configuration for text chunking +DEFAULT_CHUNK_OVERLAP = 100 # Number of tokens to overlap between chunks +DEFAULT_CHUNK_SIZE = 512 # Maximum size of each chunk in tokens T = TypeVar("T") @@ -21,6 +29,15 @@ logger = logging.getLogger(__name__) class TextSplitter(BaseModel, Generic[T]): + """Text splitter with support for protected patterns and header tracking. + + This class splits text into chunks while: + - Respecting chunk size and overlap constraints + - Preserving protected patterns (formulas, tables, code blocks) + - Tracking headers for context preservation + - Maintaining text integrity with smart merging + """ + chunk_size: int = Field(description="The token chunk size for each chunk.") chunk_overlap: int = Field( description="The token overlap of each chunk when splitting." @@ -31,14 +48,18 @@ class TextSplitter(BaseModel, Generic[T]): # Try to keep the matched characters as a whole. # If it's too long, the content will be further segmented. + # 尝试将匹配的字符作为整体保留，如果太长则进一步分段 protected_regex: List[str] = Field( description="Protected regex for splitting into words" ) len_function: Callable[[str], int] = Field(description="The length function.") # Header tracking Hook related attributes + # 标题跟踪钩子相关属性 header_hook: HeaderTracker = Field(default_factory=HeaderTracker, exclude=True) + # Compiled regex patterns for protected content _protected_fns: List[Pattern] = PrivateAttr() + # Split functions for different separators _split_fns: List[Callable] = PrivateAttr() def __init__( @@ -47,22 +68,33 @@ class TextSplitter(BaseModel, Generic[T]): chunk_overlap: int = DEFAULT_CHUNK_OVERLAP, separators: List[str] = ["\n", "。", " "], protected_regex: List[str] = [ - # math formula + # math formula - LaTeX style formulas enclosed in $$ r"\$\$[\s\S]*?\$\$", - # image + # image - Markdown image syntax ![alt](url) r"!\[.*?\]\(.*?\)", - # link + # link - Markdown link syntax [text](url) r"\[.*?\]\(.*?\)", - # table header + # table header - Markdown table header with separator line r"(?:\|[^|\n]*)+\|[\r\n]+\s*(?:\|\s*:?-{3,}:?\s*)+\|[\r\n]+", - # table body + # table body - Markdown table rows r"(?:\|[^|\n]*)+\|[\r\n]+", - # code header + # code header - Code block start with language identifier r"```(?:\w+)[\r\n]+[^\r\n]*", ], length_function: Callable[[str], int] = lambda x: len(x), ): - """Initialize with parameters.""" + """Initialize with parameters. + + Args: + chunk_size: Maximum size of each chunk + chunk_overlap: Number of tokens to overlap between chunks + separators: List of separators to use for splitting (in priority order) + protected_regex: Regex patterns for content that should be kept intact + length_function: Function to calculate text length (default: character count) + + Raises: + ValueError: If chunk_overlap is larger than chunk_size + """ if chunk_overlap > chunk_size: raise ValueError( f"Got a larger chunk overlap ({chunk_overlap}) than chunk size " @@ -76,75 +108,120 @@ class TextSplitter(BaseModel, Generic[T]): protected_regex=protected_regex, len_function=length_function, ) + # Compile all protected regex patterns for efficient matching self._protected_fns = [re.compile(reg) for reg in protected_regex] + # Create split functions: one for each separator, plus character-level splitting as fallback self._split_fns = [split_by_sep(sep) for sep in separators] + [split_by_char()] def split_text(self, text: str) -> List[Tuple[int, int, str]]: - """Split text into chunks.""" + """Split text into chunks with overlap and protected pattern handling. + + Args: + text: The input text to split + + Returns: + List of tuples (start_pos, end_pos, chunk_text) representing each chunk + """ if text == "": return [] + # Step 1: Split text by separators recursively splits = self._split(text) + # Step 2: Extract protected content positions protect = self._split_protected(text) + # Step 3: Merge splits with protected content to ensure integrity splits = self._join(splits, protect) + # Verify that joining all splits reconstructs the original text assert "".join(splits) == text + # Step 4: Merge splits into final chunks with overlap chunks = self._merge(splits) return chunks def _split(self, text: str) -> List[str]: """Break text into splits that are smaller than chunk size. + + This method recursively splits text using separators in priority order. + It tries each separator until it finds one that can split the text, + then recursively processes any splits that are still too large. NOTE: the splits contain the separators. + + Args: + text: The text to split + + Returns: + List of text splits, each smaller than chunk_size """ + # If text is already small enough, return as-is if self.len_function(text) <= self.chunk_size: return [text] + # Try each split function in order until one successfully splits the text splits = [] for split_fn in self._split_fns: splits = split_fn(text) if len(splits) > 1: break + # Process each split: keep if small enough, otherwise recursively split further new_splits = [] for split in splits: split_len = self.len_function(split) if split_len <= self.chunk_size: new_splits.append(split) else: - # recursively split + # Recursively split oversized chunks new_splits.extend(self._split(split)) return new_splits def _merge(self, splits: List[str]) -> List[Tuple[int, int, str]]: - """Merge splits into chunks. + """Merge splits into chunks with overlap and header tracking. The high-level idea is to keep adding splits to a chunk until we exceed the chunk size, then we start a new chunk with overlap. When we start a new chunk, we pop off the first element of the previous chunk until the total length is less than the chunk size. + + Headers are tracked and prepended to chunks for context preservation. + + Args: + splits: List of text splits to merge + + Returns: + List of tuples (start_pos, end_pos, chunk_text) representing merged chunks """ + # Final list of chunks with their positions chunks: List[Tuple[int, int, str]] = [] + # Current chunk being built: list of (start, end, text) tuples cur_chunk: List[Tuple[int, int, str]] = [] + # Track current headers and chunk length cur_headers, cur_len = "", 0 + # Track position in original text cur_start, cur_end = 0, 0 + for split in splits: + # Calculate position of current split in original text cur_end = cur_start + len(split) split_len = self.len_function(split) + + # Warn if a single split exceeds chunk size (shouldn't happen after _split) if split_len > self.chunk_size: logger.error( f"Got a split of size {split_len}, ", f"larger than chunk size {self.chunk_size}.", ) + # Update header tracking with current split self.header_hook.update(split) cur_headers = self.header_hook.get_headers() cur_headers_len = self.len_function(cur_headers) + # If headers are too large, skip them to avoid oversized chunks if cur_headers_len > self.chunk_size: logger.error( f"Got headers of size {cur_headers_len}, ", @@ -152,31 +229,35 @@ class TextSplitter(BaseModel, Generic[T]): ) cur_headers, cur_headers_len = "", 0 - # if we exceed the chunk size after adding the new split, then - # we need to end the current chunk and start a new one + # Check if adding this split would exceed chunk size + # If so, finalize current chunk and start a new one with overlap if cur_len + split_len + cur_headers_len > self.chunk_size: - # end the previous chunk + # Finalize the previous chunk if it has content if len(cur_chunk) > 0: chunks.append( ( - cur_chunk[0][0], - cur_chunk[-1][1], - "".join([c[2] for c in cur_chunk]), + cur_chunk[0][0], # Start position of first element + cur_chunk[-1][1], # End position of last element + "".join([c[2] for c in cur_chunk]), # Concatenated text ) ) - # start a new chunk with overlap - # keep popping off the first element of the previous chunk until: + # Start a new chunk with overlap from previous chunk + # Keep popping off the first element of the previous chunk until: # 1. the current chunk length is less than chunk overlap # 2. the total length is less than chunk size while cur_chunk and ( cur_len > self.chunk_overlap or cur_len + split_len + cur_headers_len > self.chunk_size ): - # pop off the first element + # Remove the first element to reduce overlap first_chunk = cur_chunk.pop(0) cur_len -= self.len_function(first_chunk[2]) + # Prepend headers to new chunk if: + # 1. Headers exist + # 2. Headers + split fit in chunk size + # 3. Headers are not already in the split if ( cur_headers and split_len + cur_headers_len < self.chunk_size @@ -192,11 +273,12 @@ class TextSplitter(BaseModel, Generic[T]): ) cur_len += cur_headers_len + # Add current split to the chunk cur_chunk.append((cur_start, cur_end, split)) cur_len += split_len cur_start = cur_end - # handle the last chunk + # Handle the last chunk (there should always be at least one) assert cur_chunk chunks.append( ( @@ -209,29 +291,44 @@ class TextSplitter(BaseModel, Generic[T]): return chunks def _split_protected(self, text: str) -> List[Tuple[int, str]]: + """Extract protected content from text based on regex patterns. + + Args: + text: The input text to scan for protected patterns + + Returns: + List of tuples (start_position, protected_text) for each protected match + """ + # Find all matches for all protected patterns matches = [ (match.start(), match.end()) for pattern in self._protected_fns for match in pattern.finditer(text) ] + # Sort by start position (ascending), then by length (descending) to handle overlaps matches.sort(key=lambda x: (x[0], -x[1])) res = [] def fold(initial: int, current: Tuple[int, int]) -> int: + """Accumulator function to filter overlapping matches.""" + # Only process if match starts after previous match ended if current[0] >= initial: + # Only keep protected content if it fits within chunk size if current[1] - current[0] < self.chunk_size: res.append((current[0], text[current[0] : current[1]])) else: logger.warning(f"Protected text ignore: {current}") + # Return the end position of the furthest match so far return max(initial, current[1]) - # filter overlapping matches + # Filter overlapping matches using accumulate list(itertools.accumulate(matches, fold, initial=-1)) return res def _join(self, splits: List[str], protect: List[Tuple[int, str]]) -> List[str]: - """ + """Merge splits with protected content to ensure protected patterns remain intact. + Merges and splits elements in splits array based on protected substrings. The function processes the input splits to ensure all protected substrings @@ -243,45 +340,63 @@ class TextSplitter(BaseModel, Generic[T]): Key behaviors: 1. Preserves the complete structure of each protected substring 2. Separates protected substrings from any adjacent non-protected content - 3. Maintains the original sequence of all content except for necessary + 3. Maintains the original sequence of all content 4. Handles cases where protected substrings are partially concatenated + + Args: + splits: List of text splits from _split() + protect: List of (position, text) tuples for protected content + + Returns: + List of text splits with protected content properly isolated """ - j = 0 - point, start = 0, 0 - res = [] + j = 0 # Index for protected content list + point, start = 0, 0 # Track current position in original text + res = [] # Result list of merged splits for split in splits: + # Calculate end position of current split end = start + len(split) + # Get the portion of split starting from current point cur = split[point - start :] + + # Process all protected content that overlaps with current split while j < len(protect): p_start, p_content = protect[j] p_end = p_start + len(p_content) + # If protected content is beyond current split, move to next split if end <= p_start: break + # Add content before protected section if point < p_start: local_end = p_start - point res.append(cur[:local_end]) cur = cur[local_end:] point = p_start + # Add the protected content as a single unit res.append(p_content) j += 1 + # Skip content that's part of the protected section if point < p_end: local_start = p_end - point cur = cur[local_start:] point = p_end + # If no more content in current split, break if not cur: break + # Add any remaining content from current split if cur: res.append(cur) point = end + # Move to next split start = end return res diff --git a/docreader/utils/endecode.py b/docreader/utils/endecode.py index 2457d072..ec4b4641 100644 --- a/docreader/utils/endecode.py +++ b/docreader/utils/endecode.py @@ -1,3 +1,12 @@ +""" +Encoding and Decoding Utilities Module + +This module provides utilities for encoding and decoding various data types, +with a focus on image and text data conversion: +- Image encoding/decoding (base64) +- Text encoding/decoding (multiple character sets) +- Bytes conversion utilities +""" import base64 import binascii import io @@ -11,31 +20,50 @@ logger = logging.getLogger(__name__) def decode_image(image: Union[str, bytes, Image.Image, np.ndarray]) -> str: - """Convert image to base64 encoded string + """Convert image to base64 encoded string. + + This function handles multiple image input formats and converts them + to a base64 encoded string representation, which is useful for embedding + images in JSON, HTML, or other text-based formats. Args: - image: Image file path, bytes, PIL Image object, or numpy array + image: Image in one of the following formats: + - str: File path to an image file + - bytes: Raw image bytes data + - Image.Image: PIL/Pillow Image object + - np.ndarray: NumPy array representing image data Returns: - Base64 encoded image string, or empty string if conversion fails + str: Base64 encoded string representation of the image + + Raises: + ValueError: If the image type is not supported + + Example: + >>> # From file path + >>> base64_str = decode_image("/path/to/image.png") + >>> # From PIL Image + >>> from PIL import Image + >>> img = Image.open("photo.jpg") + >>> base64_str = decode_image(img) """ if isinstance(image, str): - # It's a file path + # Handle file path: read file and encode to base64 with open(image, "rb") as image_file: return base64.b64encode(image_file.read()).decode() elif isinstance(image, bytes): - # It's bytes data + # Handle raw bytes: directly encode to base64 return base64.b64encode(image).decode() elif isinstance(image, Image.Image): - # It's a PIL Image + # Handle PIL Image: save to buffer then encode buffer = io.BytesIO() image.save(buffer, format=image.format) return base64.b64encode(buffer.getvalue()).decode() elif isinstance(image, np.ndarray): - # It's a numpy array + # Handle numpy array: convert to PIL Image, then encode as PNG pil_image = Image.fromarray(image) buffer = io.BytesIO() pil_image.save(buffer, format="PNG") @@ -45,19 +73,35 @@ def decode_image(image: Union[str, bytes, Image.Image, np.ndarray]) -> str: def encode_image(image: str, errors="strict") -> bytes: - """ - Decode image bytes using base64. + """Decode a base64 encoded image string back to bytes. + + This function converts a base64 encoded string representation of an image + back into its original binary bytes format. - errors - The error handling scheme to use for the handling of decoding errors. - The default is 'strict' meaning that decoding errors raise a - UnicodeDecodeError. Other possible values are 'ignore' and '????' - as well as any other name registered with codecs.register_error that - can handle UnicodeDecodeErrors. + Args: + image: Base64 encoded string representation of an image + errors: Error handling scheme for decoding errors: + - 'strict' (default): Raise binascii.Error on decoding errors + - 'ignore': Return empty bytes on decoding errors + - Any other name registered with codecs.register_error + + Returns: + bytes: Decoded image bytes, or empty bytes if errors='ignore' and decoding fails + + Raises: + binascii.Error: If decoding fails and errors='strict' + + Example: + >>> base64_str = "iVBORw0KGgoAAAANSUhEUgAAAAUA..." + >>> image_bytes = encode_image(base64_str) + >>> # With error handling + >>> image_bytes = encode_image(base64_str, errors="ignore") """ try: + # Attempt to decode the base64 string to bytes image_bytes = base64.b64decode(image) except binascii.Error as e: + # Handle decoding errors based on the errors parameter if errors == "ignore": return b"" else: @@ -66,6 +110,20 @@ def encode_image(image: str, errors="strict") -> bytes: def encode_bytes(content: str) -> bytes: + """Convert a string to bytes using UTF-8 encoding. + + Args: + content: String to be encoded + + Returns: + bytes: UTF-8 encoded bytes representation of the string + + Example: + >>> text = "Hello, 世界" + >>> encoded = encode_bytes(text) + >>> type(encoded) + + """ return content.encode() @@ -81,15 +139,53 @@ def decode_bytes( "latin-1", ], ) -> str: - # Try decoding with each encoding format + """Decode bytes to string with automatic encoding detection. + + This function attempts to decode bytes using multiple encoding formats + in order of priority. It's particularly useful for handling text files + with unknown or mixed encodings, especially for Chinese text. + + The function tries encodings in the provided order and returns the first + successful decode. If all encodings fail, it falls back to latin-1 with + error replacement to ensure a result is always returned. + + Args: + content: Bytes content to be decoded + encodings: List of encoding formats to try, in order of priority. + Default includes common encodings for Chinese and Western text: + - utf-8: Universal encoding (tried first) + - gb18030, gb2312, gbk: Chinese encodings (Simplified) + - big5: Chinese encoding (Traditional) + - ascii, latin-1: Western encodings + + Returns: + str: Decoded string content + + Note: + - If all encodings fail, latin-1 with error='replace' is used as fallback + - The fallback may result in character replacement (�) for invalid bytes + - A warning is logged when fallback encoding is used + + Example: + >>> # Decode with default encodings + >>> text = decode_bytes(b"\\xe4\\xb8\\xad\\xe6\\x96\\x87") # UTF-8 Chinese + >>> print(text) + 中文 + >>> # Decode with custom encodings + >>> text = decode_bytes(content, encodings=["utf-8", "gbk"]) + """ + # Try decoding with each encoding format in order for encoding in encodings: try: text = content.decode(encoding) logger.debug(f"Decode content with {encoding}: {len(text)} characters") return text except UnicodeDecodeError: + # This encoding didn't work, try the next one continue + # Fallback: use latin-1 with error replacement if all encodings fail + # latin-1 can decode any byte sequence, but may produce incorrect characters text = content.decode(encoding="latin-1", errors="replace") logger.warning( "Unable to determine correct encoding, using latin-1 as fallback. " @@ -99,5 +195,7 @@ def decode_bytes( if __name__ == "__main__": + # Example: Test encode_image with error handling + # This demonstrates decoding a base64 string with 'ignore' error mode img = "test![](data:image/png;base64,iVBORw0KGgoAAAANSUhEUgAAAMgA)test" encode_image(img, errors="ignore") diff --git a/docreader/utils/split.py b/docreader/utils/split.py index 6442c4fb..4950f7ac 100644 --- a/docreader/utils/split.py +++ b/docreader/utils/split.py @@ -3,14 +3,37 @@ from typing import Callable, List def split_text_keep_separator(text: str, separator: str) -> List[str]: - """Split text with separator and keep the separator at the end of each split.""" + """Split text with separator and keep the separator at the end of each split. + + Args: + text: The input text to split + separator: The separator string to split by + + Returns: + List of text chunks with separator preserved at the start of each chunk (except first) + + Example: + >>> split_text_keep_separator("Hello\nWorld\nTest", "\n") + ["Hello", "\nWorld", "\nTest"] + """ + # Split text by separator parts = text.split(separator) + # Add separator back to the beginning of each part (except the first one) result = [separator + s if i > 0 else s for i, s in enumerate(parts)] + # Filter out empty strings return [s for s in result if s] def split_by_sep(sep: str, keep_sep: bool = True) -> Callable[[str], List[str]]: - """Split text by separator.""" + """Create a function that splits text by a given separator. + + Args: + sep: The separator string to split by + keep_sep: If True, keep the separator in the result; if False, discard it + + Returns: + A callable function that takes text and returns a list of split strings + """ if keep_sep: return lambda text: split_text_keep_separator(text, sep) else: @@ -18,17 +41,40 @@ def split_by_sep(sep: str, keep_sep: bool = True) -> Callable[[str], List[str]]: def split_by_char() -> Callable[[str], List[str]]: - """Split text by character.""" + """Create a function that splits text into individual characters. + + Returns: + A callable function that takes text and returns a list of characters + """ return lambda text: list(text) def split_by_regex(regex: str) -> Callable[[str], List[str]]: - """Split text by regex.""" + """Create a function that splits text by a regex pattern. + + Args: + regex: The regular expression pattern to split by + + Returns: + A callable function that takes text and returns a list of split strings + The regex pattern is captured, so the separators are included in the result + """ + # Compile regex with capturing group to keep separators in result pattern = re.compile(f"({regex})") + # Split by pattern and filter out None/empty values return lambda text: list(filter(None, pattern.split(text))) def match_by_regex(regex: str) -> Callable[[str], bool]: - """Split text by regex.""" + """Create a function that checks if text matches a regex pattern. + + Args: + regex: The regular expression pattern to match against + + Returns: + A callable function that takes text and returns True if it matches the pattern + """ + # Compile the regex pattern for efficient reuse pattern = re.compile(regex) + # Return a function that checks if text matches the pattern from the start return lambda text: bool(pattern.match(text)) diff --git a/mcp-server/weknora_mcp_server.py b/mcp-server/weknora_mcp_server.py index 2e2cc484..c8ca4269 100644 --- a/mcp-server/weknora_mcp_server.py +++ b/mcp-server/weknora_mcp_server.py @@ -17,11 +17,11 @@ import mcp.types as types from mcp.server import NotificationOptions, Server from mcp.server.models import InitializationOptions -# Set up logging +# Set up logging configuration for the MCP server logging.basicConfig(level=logging.INFO) logger = logging.getLogger(__name__) -# Configuration +# Configuration - Load from environment variables with defaults WEKNORA_BASE_URL = os.getenv("WEKNORA_BASE_URL", "http://localhost:8080/api/v1") WEKNORA_API_KEY = os.getenv("WEKNORA_API_KEY", "") @@ -29,33 +29,48 @@ class WeKnoraClient: """Client for interacting with WeKnora API""" def __init__(self, base_url: str, api_key: str): + """Initialize the WeKnora API client with base URL and authentication""" self.base_url = base_url self.api_key = api_key + # Create a persistent session for connection pooling and performance self.session = requests.Session() + # Set default headers for all requests self.session.headers.update({ - "X-API-Key": api_key, - "Content-Type": "application/json" + "X-API-Key": api_key, # API key for authentication + "Content-Type": "application/json" # Default content type }) def _request(self, method: str, endpoint: str, **kwargs) -> Dict[str, Any]: - """Make a request to the WeKnora API""" + """Make a request to the WeKnora API + + Args: + method: HTTP method (GET, POST, PUT, DELETE) + endpoint: API endpoint path + **kwargs: Additional arguments to pass to requests + + Returns: + JSON response as dictionary + """ url = f"{self.base_url}{endpoint}" try: + # Execute HTTP request with the specified method response = self.session.request(method, url, **kwargs) + # Raise exception for HTTP error status codes (4xx, 5xx) response.raise_for_status() + # Parse and return JSON response return response.json() except RequestException as e: logger.error(f"API request failed: {e}") raise - # Tenant Management + # Tenant Management - Methods for managing multi-tenant configurations def create_tenant(self, name: str, description: str, business: str, retriever_engines: Dict) -> Dict: - """Create a new tenant""" + """Create a new tenant with specified configuration""" data = { "name": name, "description": description, "business": business, - "retriever_engines": retriever_engines + "retriever_engines": retriever_engines # Configuration for search engines } return self._request("POST", "/tenants", json=data) @@ -67,13 +82,13 @@ class WeKnoraClient: """List all tenants""" return self._request("GET", "/tenants") - # Knowledge Base Management + # Knowledge Base Management - Methods for managing knowledge bases def create_knowledge_base(self, name: str, description: str, config: Dict) -> Dict: - """Create a new knowledge base""" + """Create a new knowledge base with chunking and model configuration""" data = { "name": name, "description": description, - **config + **config # Merge additional configuration (chunking, models, etc.) } return self._request("POST", "/knowledge-bases", json=data) @@ -94,22 +109,24 @@ class WeKnoraClient: return self._request("DELETE", f"/knowledge-bases/{kb_id}") def hybrid_search(self, kb_id: str, query: str, config: Dict) -> Dict: - """Perform hybrid search in knowledge base""" + """Perform hybrid search combining vector and keyword search""" data = { "query_text": query, - **config + **config # Include thresholds and match count } return self._request("GET", f"/knowledge-bases/{kb_id}/hybrid-search", json=data) - # Knowledge Management + # Knowledge Management - Methods for creating and managing knowledge entries def create_knowledge_from_file(self, kb_id: str, file_path: str, enable_multimodel: bool = True) -> Dict: - """Create knowledge from file""" + """Create knowledge from a local file with optional multimodal processing""" with open(file_path, 'rb') as f: files = {'file': f} data = {'enable_multimodel': str(enable_multimodel).lower()} - # Temporarily remove Content-Type for multipart request + # Temporarily remove Content-Type header for multipart/form-data request + # (requests will set it automatically with boundary) headers = self.session.headers.copy() del headers['Content-Type'] + # Use requests.post directly instead of session to avoid header conflicts response = requests.post( f"{self.base_url}/knowledge-bases/{kb_id}/knowledge/file", headers=headers, @@ -120,10 +137,10 @@ class WeKnoraClient: return response.json() def create_knowledge_from_url(self, kb_id: str, url: str, enable_multimodel: bool = True) -> Dict: - """Create knowledge from URL""" + """Create knowledge from a web URL with optional multimodal processing""" data = { - "url": url, - "enable_multimodel": enable_multimodel + "url": url, # Web URL to fetch and process + "enable_multimodel": enable_multimodel # Enable image/multimodal extraction } return self._request("POST", f"/knowledge-bases/{kb_id}/knowledge/url", json=data) @@ -140,16 +157,16 @@ class WeKnoraClient: """Delete knowledge""" return self._request("DELETE", f"/knowledge/{knowledge_id}") - # Model Management + # Model Management - Methods for managing AI models (LLM, Embedding, Rerank) def create_model(self, name: str, model_type: str, source: str, description: str, parameters: Dict, is_default: bool = False) -> Dict: - """Create a new model""" + """Create a new AI model configuration""" data = { "name": name, - "type": model_type, - "source": source, + "type": model_type, # KnowledgeQA, Embedding, or Rerank + "source": source, # local, openai, etc. "description": description, - "parameters": parameters, - "is_default": is_default + "parameters": parameters, # API keys, base URLs, etc. + "is_default": is_default # Set as default model for this type } return self._request("POST", "/models", json=data) @@ -161,12 +178,12 @@ class WeKnoraClient: """Get model details""" return self._request("GET", f"/models/{model_id}") - # Session Management + # Session Management - Methods for managing chat sessions def create_session(self, kb_id: str, strategy: Dict) -> Dict: - """Create a new chat session""" + """Create a new chat session with conversation strategy""" data = { - "knowledge_base_id": kb_id, - "session_strategy": strategy + "knowledge_base_id": kb_id, # Knowledge base to query + "session_strategy": strategy # Conversation settings (max rounds, rewrite, etc.) } return self._request("POST", "/sessions", json=data) @@ -183,16 +200,17 @@ class WeKnoraClient: """Delete session""" return self._request("DELETE", f"/sessions/{session_id}") - # Chat Functionality + # Chat Functionality - Methods for conversational interactions def chat(self, session_id: str, query: str) -> Dict: - """Send a chat message""" + """Send a chat message and get AI response""" data = {"query": query} - # Note: This returns SSE stream, simplified here + # Note: The actual API returns Server-Sent Events (SSE) stream + # This simplified version returns the complete response return self._request("POST", f"/knowledge-chat/{session_id}", json=data) - # Chunk Management + # Chunk Management - Methods for managing knowledge chunks (text segments) def list_chunks(self, knowledge_id: str, page: int = 1, page_size: int = 20) -> Dict: - """List chunks of knowledge""" + """List text chunks of a knowledge entry with pagination""" params = {"page": page, "page_size": page_size} return self._request("GET", f"/chunks/{knowledge_id}", params=params) @@ -200,14 +218,15 @@ class WeKnoraClient: """Delete a chunk""" return self._request("DELETE", f"/chunks/{knowledge_id}/{chunk_id}") -# Initialize MCP server +# Initialize MCP server instance app = Server("weknora-server") +# Initialize WeKnora API client with configuration client = WeKnoraClient(WEKNORA_BASE_URL, WEKNORA_API_KEY) -# Tool definitions +# Tool definitions - Register all available tools for the MCP protocol @app.list_tools() async def handle_list_tools() -> list[types.Tool]: - """List all available WeKnora tools""" + """List all available WeKnora tools with their schemas""" return [ # Tenant Management types.Tool( @@ -497,17 +516,27 @@ async def handle_list_tools() -> list[types.Tool]: async def handle_call_tool( name: str, arguments: dict | None ) -> list[types.TextContent | types.ImageContent | types.EmbeddedResource]: - """Handle tool execution""" + """Handle tool execution requests from MCP clients + + Args: + name: Name of the tool to execute + arguments: Tool arguments as dictionary + + Returns: + List of content items (text, image, or embedded resources) + """ try: + # Use empty dict if no arguments provided args = arguments or {} - # Tenant Management + # Tenant Management - Route tenant-related operations if name == "create_tenant": result = client.create_tenant( args["name"], args["description"], args["business"], + # Default to postgres-based keyword and vector search if not specified args.get("retriever_engines", { "engines": [ {"retriever_type": "keywords", "retriever_engine_type": "postgres"}, @@ -518,14 +547,15 @@ async def handle_call_tool( elif name == "list_tenants": result = client.list_tenants() - # Knowledge Base Management + # Knowledge Base Management - Route knowledge base operations elif name == "create_knowledge_base": + # Build configuration with defaults for chunking and models config = { "chunking_config": args.get("chunking_config", { - "chunk_size": 1000, - "chunk_overlap": 200, - "separators": ["."], - "enable_multimodal": True + "chunk_size": 1000, # Default chunk size in characters + "chunk_overlap": 200, # Default overlap between chunks + "separators": ["."], # Default text separators + "enable_multimodal": True # Enable image processing by default }), "embedding_model_id": args.get("embedding_model_id", ""), "summary_model_id": args.get("summary_model_id", "") @@ -542,10 +572,11 @@ async def handle_call_tool( elif name == "delete_knowledge_base": result = client.delete_knowledge_base(args["kb_id"]) elif name == "hybrid_search": + # Configure hybrid search with thresholds and result count config = { - "vector_threshold": args.get("vector_threshold", 0.5), - "keyword_threshold": args.get("keyword_threshold", 0.3), - "match_count": args.get("match_count", 5) + "vector_threshold": args.get("vector_threshold", 0.5), # Minimum similarity score + "keyword_threshold": args.get("keyword_threshold", 0.3), # Minimum keyword match score + "match_count": args.get("match_count", 5) # Number of results to return } result = client.hybrid_search(args["kb_id"], args["query"], config) @@ -573,11 +604,12 @@ async def handle_call_tool( elif name == "delete_knowledge": result = client.delete_knowledge(args["knowledge_id"]) - # Model Management + # Model Management - Route model configuration operations elif name == "create_model": + # Build model parameters (API credentials, endpoints, etc.) parameters = { - "base_url": args.get("base_url", ""), - "api_key": args.get("api_key", "") + "base_url": args.get("base_url", ""), # Model API endpoint + "api_key": args.get("api_key", "") # Model API key } result = client.create_model( args["name"], @@ -592,17 +624,18 @@ async def handle_call_tool( elif name == "get_model": result = client.get_model(args["model_id"]) - # Session Management + # Session Management - Route chat session operations elif name == "create_session": + # Build session strategy with conversation settings strategy = { - "max_rounds": args.get("max_rounds", 5), - "enable_rewrite": args.get("enable_rewrite", True), - "fallback_strategy": "FIXED_RESPONSE", + "max_rounds": args.get("max_rounds", 5), # Maximum conversation turns + "enable_rewrite": args.get("enable_rewrite", True), # Enable query rewriting + "fallback_strategy": "FIXED_RESPONSE", # Strategy when no answer found "fallback_response": args.get("fallback_response", "Sorry, I cannot answer this question."), - "embedding_top_k": 10, - "keyword_threshold": 0.5, - "vector_threshold": 0.7, - "summary_model_id": args.get("summary_model_id", "") + "embedding_top_k": 10, # Number of chunks to retrieve + "keyword_threshold": 0.5, # Keyword match threshold + "vector_threshold": 0.7, # Vector similarity threshold + "summary_model_id": args.get("summary_model_id", "") # Model for summarization } result = client.create_session(args["kb_id"], strategy) elif name == "get_session": @@ -630,17 +663,20 @@ async def handle_call_tool( result = client.delete_chunk(args["knowledge_id"], args["chunk_id"]) else: + # Handle unknown tool names return [types.TextContent( type="text", text=f"Unknown tool: {name}" )] + # Return successful result as formatted JSON return [types.TextContent( type="text", text=json.dumps(result, indent=2, ensure_ascii=False) )] except Exception as e: + # Log and return error message logger.error(f"Tool execution failed: {e}") return [types.TextContent( type="text", @@ -648,8 +684,10 @@ async def handle_call_tool( )] async def run(): - """Run the MCP server""" + """Run the MCP server using stdio transport""" + # Create stdio streams for communication with MCP client async with mcp.server.stdio.stdio_server() as (read_stream, write_stream): + # Run the server with initialization options await app.run( read_stream, write_stream, @@ -664,8 +702,9 @@ async def run(): ) def main(): - """主函数入口点，用于 console_scripts""" + """Main entry point for console_scripts""" import asyncio + # Run the async server asyncio.run(run()) if __name__ == "__main__":