PDF Knowledgebase MCP Server

"""Chunker using LangChain's MarkdownHeaderTextSplitter.""" import logging from datetime import datetime, timezone from typing import Any, Dict, List, Optional, Tuple from ..models import Chunk from .chunker import Chunker logger = logging.getLogger(__name__) class LangChainChunker(Chunker): """Chunker using LangChain's MarkdownHeaderTextSplitter.""" def __init__( self, headers_to_split_on: Optional[List[Tuple[str, str]]] = None, chunk_size: int = 1000, chunk_overlap: int = 200, cache_dir: str = None, min_chunk_size: int = 0, ): """Initialize the LangChain chunker. Args: headers_to_split_on: List of (header_tag, header_name) tuples. chunk_size: Maximum size of each chunk. chunk_overlap: Overlap between chunks. cache_dir: Optional cache directory. min_chunk_size: Minimum size for chunks (0 = disabled). """ super().__init__(cache_dir=cache_dir, min_chunk_size=min_chunk_size) self.headers_to_split_on = headers_to_split_on or [ ("#", "Header 1"), ("##", "Header 2"), ("###", "Header 3"), ] self.chunk_size = chunk_size self.chunk_overlap = chunk_overlap try: from langchain_text_splitters import MarkdownHeaderTextSplitter, RecursiveCharacterTextSplitter self.header_splitter = MarkdownHeaderTextSplitter(headers_to_split_on=self.headers_to_split_on) self.text_splitter = RecursiveCharacterTextSplitter( chunk_size=self.chunk_size, chunk_overlap=self.chunk_overlap, separators=["\n\n", "\n", ". ", " ", ""], ) except ImportError: raise ImportError("LangChain text splitters not available. Install with: pip install pdfkb-mcp[langchain]") def chunk(self, markdown_content: str, metadata: Dict[str, Any]) -> List[Chunk]: """Chunk markdown content using LangChain's splitter. Args: markdown_content: Markdown text to chunk. metadata: Document metadata. Returns: List of Chunk objects. """ try: if not markdown_content or not markdown_content.strip(): logger.warning("Empty markdown content provided to chunker") return [] # Split the markdown content by headers header_splits = self.header_splitter.split_text(markdown_content) # Process each header split and further split if needed final_chunks = [] for split in header_splits: # Extract the text content from the split if hasattr(split, "page_content"): # LangChain Document object text_content = split.page_content split_metadata = getattr(split, "metadata", {}) else: # Plain string text_content = str(split) split_metadata = {} # If the split is small enough, use it as is if len(text_content) <= self.chunk_size: final_chunks.append({"text": text_content, "metadata": split_metadata}) else: # Split large sections into smaller chunks sub_chunks = self.text_splitter.split_text(text_content) for sub_chunk in sub_chunks: final_chunks.append({"text": sub_chunk, "metadata": split_metadata}) # Convert to Chunk objects with metadata chunks = [] for i, chunk_data in enumerate(final_chunks): chunk_text = chunk_data["text"].strip() if not chunk_text: # Skip empty chunks continue chunk_metadata = { "chunk_strategy": "langchain_markdown_header", "headers_to_split_on": self.headers_to_split_on, "chunk_size": self.chunk_size, "chunk_overlap": self.chunk_overlap, "created_at": datetime.now(timezone.utc).isoformat(), } # Add header metadata if available if chunk_data["metadata"]: chunk_metadata.update(chunk_data["metadata"]) # Add any provided metadata chunk_metadata.update(metadata) chunk = Chunk(text=chunk_text, chunk_index=i, metadata=chunk_metadata) chunks.append(chunk) # Apply minimum chunk size filtering chunks = self._filter_small_chunks(chunks) logger.info(f"Created {len(chunks)} chunks using LangChain MarkdownHeaderTextSplitter") return chunks except Exception as e: logger.error(f"Failed to chunk markdown content with LangChain: {e}") # Fallback to simple text splitting return self._fallback_chunk(markdown_content, metadata) def _fallback_chunk(self, text: str, metadata: Dict[str, Any]) -> List[Chunk]: """Fallback chunking method using only RecursiveCharacterTextSplitter. Args: text: Text content to chunk. metadata: Document metadata. Returns: List of Chunk objects. """ try: logger.warning("Using fallback chunking method") chunks = [] text_chunks = self.text_splitter.split_text(text) for i, chunk_text in enumerate(text_chunks): if not chunk_text.strip(): continue chunk_metadata = { "chunk_strategy": "langchain_fallback", "chunk_size": self.chunk_size, "chunk_overlap": self.chunk_overlap, "created_at": datetime.now(timezone.utc).isoformat(), } chunk_metadata.update(metadata) chunk = Chunk(text=chunk_text.strip(), chunk_index=i, metadata=chunk_metadata) chunks.append(chunk) # Apply minimum chunk size filtering chunks = self._filter_small_chunks(chunks) logger.info(f"Created {len(chunks)} chunks using fallback method") return chunks except Exception as e: logger.error(f"Fallback chunking also failed: {e}") raise RuntimeError(f"Failed to chunk text: {e}") from e