MindsDB 1.2.9__py3-none-any.whl → 24.12.4.0__py3-none-any.whl

mindsdb/interfaces/knowledge_base/preprocessing/constants.py

@@ -0,0 +1,13 @@
+# Default settings for markdown header splitting
+DEFAULT_MARKDOWN_HEADERS = [
+    ("#", "Header 1"),
+    ("##", "Header 2"),
+    ("###", "Header 3"),
+]
+
+# Limits for web crawling
+DEFAULT_CRAWL_DEPTH = 1
+DEFAULT_WEB_CRAWL_LIMIT = 100
+DEFAULT_WEB_FILTERS = []
+
+DEFAULT_CONTEXT_DOCUMENT_LIMIT = 50

mindsdb/interfaces/knowledge_base/preprocessing/document_loader.py

@@ -0,0 +1,176 @@
+import os
+from typing import List, Iterator
+from langchain_core.documents import Document as LangchainDocument
+from langchain_text_splitters import MarkdownHeaderTextSplitter, RecursiveCharacterTextSplitter
+import pandas as pd
+
+from mindsdb.interfaces.file.file_controller import FileController
+from mindsdb.integrations.utilities.rag.loaders.file_loader import FileLoader
+from mindsdb.integrations.utilities.rag.splitters.file_splitter import (
+    FileSplitter,
+    DEFAULT_CHUNK_SIZE,
+    DEFAULT_CHUNK_OVERLAP
+)
+from mindsdb.integrations.handlers.web_handler.urlcrawl_helpers import get_all_websites
+from mindsdb.interfaces.knowledge_base.preprocessing.models import Document
+from mindsdb.utilities import log
+
+logger = log.getLogger(__name__)
+
+
+class DocumentLoader:
+    """Handles loading documents from various sources including SQL queries"""
+
+    def __init__(
+            self,
+            file_controller: FileController,
+            file_splitter: FileSplitter,
+            markdown_splitter: MarkdownHeaderTextSplitter,
+            file_loader_class=FileLoader,
+            mysql_proxy=None
+    ):
+        """
+        Initialize with required dependencies
+
+        Args:
+            file_controller: Controller for file operations
+            file_splitter: Splitter for file content
+            markdown_splitter: Splitter for markdown content
+            file_loader_class: Class to use for file loading
+            mysql_proxy: Proxy for executing MySQL queries
+        """
+        self.file_controller = file_controller
+        self.file_splitter = file_splitter
+        self.markdown_splitter = markdown_splitter
+        self.file_loader_class = file_loader_class
+        self.mysql_proxy = mysql_proxy
+
+        # Initialize text splitter for query results with default settings
+        self.query_splitter = RecursiveCharacterTextSplitter(
+            chunk_size=DEFAULT_CHUNK_SIZE,
+            chunk_overlap=DEFAULT_CHUNK_OVERLAP
+        )
+
+    def load_files(self, file_names: List[str]) -> Iterator[Document]:
+        """Load and split documents from files"""
+        for file_name in file_names:
+            file_path = self.file_controller.get_file_path(file_name)
+            loader = self.file_loader_class(file_path)
+
+            for doc in loader.lazy_load():
+                # Add file extension to metadata for proper splitting
+                extension = os.path.splitext(file_path)[1].lower()
+                doc.metadata['extension'] = extension
+                doc.metadata['source'] = file_name
+
+                # Use FileSplitter to handle the document based on its type
+                split_docs = self.file_splitter.split_documents([doc])
+                for split_doc in split_docs:
+                    # Preserve original metadata while adding split-specific metadata
+                    metadata = doc.metadata.copy()
+                    metadata.update(split_doc.metadata or {})
+
+                    yield Document(
+                        content=split_doc.page_content,
+                        metadata=metadata
+                    )
+
+    def load_web_pages(
+            self,
+            urls: List[str],
+            crawl_depth: int,
+            limit: int,
+            filters: List[str] = None,
+    ) -> Iterator[Document]:
+        """Load and split documents from web pages"""
+        websites_df = get_all_websites(
+            urls,
+            crawl_depth=crawl_depth,
+            limit=limit,
+            filters=filters
+        )
+
+        for _, row in websites_df.iterrows():
+            # Create a document with HTML extension for proper splitting
+            doc = LangchainDocument(
+                page_content=row['text_content'],
+                metadata={
+                    'extension': '.html',
+                    'url': row['url']
+                }
+            )
+
+            # Use FileSplitter to handle HTML content
+            split_docs = self.file_splitter.split_documents([doc])
+            for split_doc in split_docs:
+                metadata = doc.metadata.copy()
+                metadata.update(split_doc.metadata or {})
+
+                yield Document(
+                    content=split_doc.page_content,
+                    metadata=metadata
+                )
+
+    def load_query_result(self, query: str, project_name: str) -> Iterator[Document]:
+        """
+        Load documents from SQL query results
+
+        Args:
+            query: SQL query to execute
+            project_name: Name of the project context
+
+        Returns:
+            Iterator of Document objects
+
+        Raises:
+            ValueError: If mysql_proxy is not configured or query returns no data
+        """
+        if not self.mysql_proxy:
+            raise ValueError("MySQL proxy not configured")
+
+        if not query:
+            return
+
+        # Set project context and execute query
+        self.mysql_proxy.set_context({'db': project_name})
+        query_result = self.mysql_proxy.process_query(query)
+
+        if query_result.type != 'table':
+            raise ValueError('Query returned no data')
+
+        # Convert query result to DataFrame
+        df = query_result.data.to_df()
+
+        # Process each row into a Document
+        for _, row in df.iterrows():
+            # Extract content and metadata
+            content = str(row.get('content', ''))
+
+            # Convert remaining columns to metadata
+            metadata = {
+                col: str(row[col])
+                for col in df.columns
+                if col != 'content' and not pd.isna(row[col])
+            }
+            metadata['source'] = 'query'
+
+            # Split content using recursive splitter
+            if content:
+                doc = LangchainDocument(
+                    page_content=content,
+                    metadata=metadata
+                )
+                # Use FileSplitter with default recursive splitter
+                split_docs = self.file_splitter.split_documents(
+                    [doc],
+                    default_failover=True
+                )
+
+                for split_doc in split_docs:
+                    metadata = doc.metadata.copy()
+                    metadata.update(split_doc.metadata or {})
+
+                    yield Document(
+                        content=split_doc.page_content,
+                        metadata=metadata
+                    )

mindsdb/interfaces/knowledge_base/preprocessing/document_preprocessor.py

@@ -0,0 +1,401 @@
+from typing import List, Dict, Optional, Any
+import pandas as pd
+from langchain_text_splitters import RecursiveCharacterTextSplitter
+import hashlib
+import asyncio
+
+
+from mindsdb.integrations.utilities.rag.splitters.file_splitter import (
+    FileSplitter,
+    FileSplitterConfig,
+)
+
+from mindsdb.interfaces.agents.langchain_agent import create_chat_model
+
+from mindsdb.interfaces.knowledge_base.preprocessing.models import (
+    PreprocessingConfig,
+    ProcessedChunk,
+    PreprocessorType,
+    ContextualConfig,
+    Document,
+    TextChunkingConfig,
+)
+from mindsdb.utilities import log
+
+from langchain_core.documents import Document as LangchainDocument
+
+logger = log.getLogger(__name__)
+
+
+class DocumentPreprocessor:
+    """Base class for document preprocessing"""
+
+    RESERVED_METADATA_FIELDS = {
+        "content",
+        "id",
+        "embeddings",
+        "original_doc_id",
+        "chunk_index",
+    }
+
+    def __init__(self):
+        """Initialize preprocessor"""
+        self.splitter = None  # Will be set by child classes
+
+    def process_documents(self, documents: List[Document]) -> List[ProcessedChunk]:
+        """Base implementation - should be overridden by child classes"""
+        raise NotImplementedError("Subclasses must implement process_documents")
+
+    def _split_document(self, doc: Document) -> List[Document]:
+        """Split document into chunks while preserving metadata"""
+        if self.splitter is None:
+            raise ValueError("Splitter not configured")
+
+        # Convert to langchain Document for splitting
+        langchain_doc = LangchainDocument(
+            page_content=doc.content, metadata=doc.metadata or {}
+        )
+        # Split and convert back to our Document type
+        split_docs = self.splitter.split_documents([langchain_doc])
+        return [
+            Document(content=split_doc.page_content, metadata=split_doc.metadata)
+            for split_doc in split_docs
+        ]
+
+    def _get_source(self) -> str:
+        """Get the source identifier for this preprocessor"""
+        return self.__class__.__name__
+
+    def to_dataframe(self, chunks: List[ProcessedChunk]) -> pd.DataFrame:
+        """Convert processed chunks to dataframe format"""
+        return pd.DataFrame([chunk.model_dump() for chunk in chunks])
+
+    @staticmethod
+    def dict_to_document(data: Dict[str, Any]) -> Document:
+        """Convert a dictionary to a Document object"""
+        return Document(
+            id=data.get("id"),
+            content=data.get("content", ""),
+            embeddings=data.get("embeddings"),
+            metadata=data.get("metadata", {}),
+        )
+
+    def _generate_deterministic_id(
+        self, content: str, content_column: str = None, provided_id: str = None
+    ) -> str:
+        """Generate a deterministic ID based on content and column"""
+        if provided_id is not None:
+            return f"{provided_id}_{content_column}"
+
+        id_string = f"content={content}_column={content_column}"
+        return hashlib.sha256(id_string.encode()).hexdigest()
+
+    def _generate_chunk_id(
+        self,
+        content: str,
+        chunk_index: Optional[int] = None,
+        content_column: str = None,
+        provided_id: str = None,
+    ) -> str:
+        """Generate deterministic ID for a chunk"""
+        base_id = self._generate_deterministic_id(content, content_column, provided_id)
+        chunk_id = (
+            f"{base_id}_chunk_{chunk_index}" if chunk_index is not None else base_id
+        )
+        logger.debug(f"Generated chunk ID: {chunk_id} for content hash: {base_id}")
+        return chunk_id
+
+    def _prepare_chunk_metadata(
+        self,
+        doc_id: Optional[str],
+        chunk_index: Optional[int],
+        base_metadata: Optional[Dict] = None,
+    ) -> Dict:
+        """Centralized method for preparing chunk metadata"""
+        metadata = base_metadata or {}
+
+        # Always preserve original document ID
+        if doc_id is not None:
+            metadata["original_doc_id"] = doc_id
+
+        # Add chunk index only for multi-chunk cases
+        if chunk_index is not None:
+            metadata["chunk_index"] = chunk_index
+
+        # Always set source
+        metadata["source"] = self._get_source()
+
+        return metadata
+
+
+class ContextualPreprocessor(DocumentPreprocessor):
+    """Contextual preprocessing implementation that enhances document chunks with context"""
+
+    DEFAULT_CONTEXT_TEMPLATE = """
+<document>
+{{WHOLE_DOCUMENT}}
+</document>
+Here is the chunk we want to situate within the whole document
+<chunk>
+{{CHUNK_CONTENT}}
+</chunk>
+Please give a short succinct context to situate this chunk within the overall document for the purposes of improving search retrieval of the chunk. Answer only with the succinct context and nothing else."""
+
+    def __init__(self, config: ContextualConfig):
+        """Initialize with contextual configuration"""
+        super().__init__()
+        self.config = config
+        self.splitter = FileSplitter(
+            FileSplitterConfig(
+                chunk_size=config.chunk_size, chunk_overlap=config.chunk_overlap
+            )
+        )
+        self.llm = create_chat_model(
+            {
+                "model_name": self.config.llm_config.model_name,
+                "provider": self.config.llm_config.provider,
+                **self.config.llm_config.params,
+            }
+        )
+        self.context_template = config.context_template or self.DEFAULT_CONTEXT_TEMPLATE
+        self.summarize = self.config.summarize
+
+    def _prepare_prompts(
+        self, chunk_contents: list[str], full_documents: list[str]
+    ) -> list[str]:
+        prompts = [
+            self.context_template.replace("{{WHOLE_DOCUMENT}}", full_document)
+            for full_document in full_documents
+        ]
+        prompts = [
+            prompt.replace("{{CHUNK_CONTENT}}", chunk_content)
+            for prompt, chunk_content in zip(prompts, chunk_contents)
+        ]
+
+        return prompts
+
+    def _generate_context(
+        self, chunk_contents: list[str], full_documents: list[str]
+    ) -> list[str]:
+        """Generate contextual description for a chunk using LLM"""
+        prompts = self._prepare_prompts(chunk_contents, full_documents)
+
+        # Check if LLM supports async
+        if hasattr(self.llm, 'abatch'):
+            loop = asyncio.new_event_loop()
+            asyncio.set_event_loop(loop)
+            try:
+                response = loop.run_until_complete(self._async_generate(prompts))
+            finally:
+                loop.close()
+        else:
+            # Use sync batch for non-async LLMs
+            response = [resp.content for resp in self.llm.batch(prompts)]
+        return response
+
+    async def _async_generate(self, prompts: list[str]) -> list[str]:
+        """Helper for async LLM generation"""
+        return [resp.content for resp in await self.llm.abatch(prompts)]
+
+    def _split_document(self, doc: Document) -> List[Document]:
+        """Split document into chunks while preserving metadata"""
+        # Use base class implementation
+        return super()._split_document(doc)
+
+    def process_documents(self, documents: List[Document]) -> List[ProcessedChunk]:
+        chunks_list = []
+        doc_index_list = []
+        chunk_index_list = []
+        processed_chunks = []
+
+        for doc_index, doc in enumerate(documents):
+            # Get content_column from metadata if available
+            content_column = (
+                doc.metadata.get("content_column") if doc.metadata else None
+            )
+
+            # Ensure document has an ID
+            if doc.id is None:
+                doc.id = self._generate_deterministic_id(doc.content, content_column)
+
+            # Skip empty or whitespace-only content
+            if not doc.content or not doc.content.strip():
+                continue
+
+            chunk_docs = self._split_document(doc)
+
+            # Single chunk case
+            if len(chunk_docs) == 1:
+                chunk_doc = chunk_docs[0]
+                if not chunk_doc.content or not chunk_doc.content.strip():
+                    continue
+                else:
+                    chunks_list.append(chunk_doc)
+                    doc_index_list.append(doc_index)
+                    chunk_index_list.append(0)
+
+            else:
+                # Multiple chunks case
+                for i, chunk_doc in enumerate(chunk_docs):
+                    if not chunk_doc.content or not chunk_doc.content.strip():
+                        continue
+                    else:
+                        chunks_list.append(chunk_doc)
+                        doc_index_list.append(doc_index)
+                        chunk_index_list.append(i)
+
+        # Generate contexts
+        doc_contents = [documents[i].content for i in doc_index_list]
+        chunk_contents = [chunk_doc.content for chunk_doc in chunks_list]
+        contexts = self._generate_context(chunk_contents, doc_contents)
+
+        for context, chunk_doc, chunk_index, doc_index in zip(
+            contexts, chunks_list, chunk_index_list, doc_index_list
+        ):
+            processed_content = (
+                context if self.summarize else f"{context}\n\n{chunk_doc.content}"
+            )
+            doc = documents[doc_index]
+
+            # Initialize metadata
+            metadata = {}
+            if doc.metadata:
+                metadata.update(doc.metadata)
+
+            # Pass through doc.id and content_column
+            content_column = (
+                doc.metadata.get("content_column") if doc.metadata else None
+            )
+            chunk_id = self._generate_chunk_id(
+                processed_content,
+                chunk_index,
+                content_column=content_column,
+                provided_id=doc.id,
+            )
+            processed_chunks.append(
+                ProcessedChunk(
+                    id=chunk_id,
+                    content=processed_content,  # Use the content with context
+                    embeddings=doc.embeddings,
+                    metadata=self._prepare_chunk_metadata(doc.id, None, metadata),
+                )
+            )
+
+        return processed_chunks
+
+
+class TextChunkingPreprocessor(DocumentPreprocessor):
+    """Default text chunking preprocessor using RecursiveCharacterTextSplitter"""
+
+    def __init__(self, config: Optional[TextChunkingConfig] = None):
+        """Initialize with text chunking configuration"""
+        super().__init__()
+        self.config = config or TextChunkingConfig()
+        self.splitter = RecursiveCharacterTextSplitter(
+            chunk_size=self.config.chunk_size,
+            chunk_overlap=self.config.chunk_overlap,
+            length_function=self.config.length_function,
+            separators=self.config.separators,
+        )
+
+    def _split_document(self, doc: Document) -> List[Document]:
+        """Split document into chunks while preserving metadata"""
+        # Use base class implementation
+        return super()._split_document(doc)
+
+    def process_documents(self, documents: List[Document]) -> List[ProcessedChunk]:
+        processed_chunks = []
+
+        for doc in documents:
+            # Get content_column from metadata if available
+            content_column = (
+                doc.metadata.get("content_column") if doc.metadata else None
+            )
+
+            # Ensure document has an ID
+            if doc.id is None:
+                doc.id = self._generate_deterministic_id(doc.content, content_column)
+
+            # Skip empty or whitespace-only content
+            if not doc.content or not doc.content.strip():
+                continue
+
+            chunk_docs = self._split_document(doc)
+
+            # Single chunk case
+            if len(chunk_docs) == 1:
+                chunk_doc = chunk_docs[0]
+                if not chunk_doc.content or not chunk_doc.content.strip():
+                    continue
+
+                # Initialize metadata
+                metadata = {}
+                if doc.metadata:
+                    metadata.update(doc.metadata)
+
+                # Pass through doc.id and content_column
+                id = self._generate_chunk_id(
+                    chunk_doc.content, content_column=content_column, provided_id=doc.id
+                )
+                processed_chunks.append(
+                    ProcessedChunk(
+                        id=id,
+                        content=chunk_doc.content,
+                        embeddings=doc.embeddings,
+                        metadata=self._prepare_chunk_metadata(doc.id, None, metadata),
+                    )
+                )
+            else:
+                # Multiple chunks case
+                for i, chunk_doc in enumerate(chunk_docs):
+                    if not chunk_doc.content or not chunk_doc.content.strip():
+                        continue
+
+                    # Initialize metadata
+                    metadata = {}
+                    if doc.metadata:
+                        metadata.update(doc.metadata)
+
+                    # Pass through doc.id and content_column
+                    chunk_id = self._generate_chunk_id(
+                        chunk_doc.content,
+                        i,
+                        content_column=content_column,
+                        provided_id=doc.id,
+                    )
+                    processed_chunks.append(
+                        ProcessedChunk(
+                            id=chunk_id,
+                            content=chunk_doc.content,
+                            embeddings=doc.embeddings,
+                            metadata=self._prepare_chunk_metadata(doc.id, i, metadata),
+                        )
+                    )
+
+        return processed_chunks
+
+
+class PreprocessorFactory:
+    """Factory for creating preprocessors based on configuration"""
+
+    @staticmethod
+    def create_preprocessor(
+        config: Optional[PreprocessingConfig] = None,
+    ) -> DocumentPreprocessor:
+        """
+        Create appropriate preprocessor based on configuration
+        : param config: Preprocessing configuration
+        : return: Configured preprocessor instance
+        : raises ValueError: If unknown preprocessor type specified
+        """
+        if config is None:
+            return TextChunkingPreprocessor()
+
+        if config.type == PreprocessorType.CONTEXTUAL:
+            return ContextualPreprocessor(
+                config.contextual_config or ContextualConfig()
+            )
+        elif config.type == PreprocessorType.TEXT_CHUNKING:
+            return TextChunkingPreprocessor(config.text_chunking_config)
+
+        raise ValueError(f"Unknown preprocessor type: {config.type}")

mindsdb/interfaces/knowledge_base/preprocessing/models.py

@@ -0,0 +1,108 @@
+from enum import Enum
+
+from typing import List, Dict, Any, Optional, Union, Callable
+
+
+from pydantic import BaseModel, Field, model_validator
+
+
+from mindsdb.integrations.utilities.rag.settings import DEFAULT_CHUNK_OVERLAP, DEFAULT_CHUNK_SIZE
+from mindsdb.integrations.utilities.rag.settings import LLMConfig
+
+
+class PreprocessorType(Enum):
+    CONTEXTUAL = "contextual"
+    TEXT_CHUNKING = "text_chunking"
+
+
+class BasePreprocessingConfig(BaseModel):
+    """Base configuration for preprocessing"""
+    chunk_size: int = Field(default=DEFAULT_CHUNK_SIZE, description="Size of document chunks")
+    chunk_overlap: int = Field(default=DEFAULT_CHUNK_OVERLAP, description="Overlap between chunks")
+
+
+class ContextualConfig(BasePreprocessingConfig):
+    """Configuration specific to contextual preprocessing"""
+    llm_config: LLMConfig = Field(
+        default_factory=LLMConfig,
+        description="LLM configuration to use for context generation"
+    )
+    context_template: Optional[str] = Field(
+        default=None,
+        description="Custom template for context generation"
+    )
+    summarize: Optional[bool] = Field(
+        default=False,
+        description="Whether to return chunks as summarizations"
+    )
+
+
+class TextChunkingConfig(BaseModel):
+    """Configuration for text chunking preprocessor using Pydantic"""
+    chunk_size: int = Field(
+        default=1000,
+        description="The target size of each text chunk",
+        gt=0
+    )
+    chunk_overlap: int = Field(
+        default=200,
+        description="The number of characters to overlap between chunks",
+        ge=0
+    )
+    length_function: Callable = Field(
+        default=len,
+        description="Function to measure text length"
+    )
+    separators: List[str] = Field(
+        default=["\n\n", "\n", " ", ""],
+        description="List of separators to use for splitting text, in order of priority"
+    )
+
+    class Config:
+        arbitrary_types_allowed = True
+
+
+class PreprocessingConfig(BaseModel):
+    """Complete preprocessing configuration"""
+    type: PreprocessorType = Field(
+        default=PreprocessorType.TEXT_CHUNKING,
+        description="Type of preprocessing to apply"
+    )
+    contextual_config: Optional[ContextualConfig] = Field(
+        default=None,
+        description="Configuration for contextual preprocessing"
+    )
+    text_chunking_config: Optional[TextChunkingConfig] = Field(
+        default=None,
+        description="Configuration for text chunking preprocessing"
+    )
+
+    @model_validator(mode='after')
+    def validate_config_presence(self) -> 'PreprocessingConfig':
+        """Ensure the appropriate config is present for the chosen type"""
+        if self.type == PreprocessorType.CONTEXTUAL and not self.contextual_config:
+            self.contextual_config = ContextualConfig()
+        if self.type == PreprocessorType.TEXT_CHUNKING and not self.text_chunking_config:
+            self.text_chunking_config = TextChunkingConfig()
+        return self
+
+
+class Document(BaseModel):
+
+    """Document model with default metadata handling"""
+    id: Optional[Union[int, str]] = Field(default=None, description="Unique identifier for the document")
+    content: str = Field(description="The document content")
+    embeddings: Optional[List[float]] = Field(default=None, description="Vector embeddings of the content")
+    metadata: Optional[Dict[str, Any]] = Field(default=None, description="Additional document metadata")
+
+    @model_validator(mode='after')
+    def validate_metadata(self) -> 'Document':
+        """Ensure metadata is present and valid"""
+        if not self.metadata:
+            self.metadata = {'source': 'default'}
+        return self
+
+
+class ProcessedChunk(Document):
+    """Processed chunk that aligns with VectorStoreHandler schema"""
+    pass