alita-sdk 0.3.257__py3-none-any.whl → 0.3.584__py3-none-any.whl

alita_sdk/runtime/langchain/document_loaders/AlitaDocxMammothLoader.py

@@ -1,4 +1,5 @@
 import re
+import uuid
 from io import BytesIO
 
 import mammoth.images
@@ -8,6 +9,9 @@ from langchain_core.document_loaders import BaseLoader
 from langchain_core.documents import Document
 from mammoth import convert_to_html
 from markdownify import markdownify
+from docx import Document as DocxDocument
+from docx.oxml.ns import qn
+from bs4 import BeautifulSoup
 
 from alita_sdk.tools.chunkers.sematic.markdown_chunker import markdown_by_headers_chunker
 from .utils import perform_llm_prediction_for_image_bytes
@@ -17,6 +21,7 @@ class AlitaDocxMammothLoader(BaseLoader):
     """
     Loader for Docx files using Mammoth to convert to HTML, with image handling,
     and then Markdownify to convert HTML to markdown.
+    Detects bordered paragraphs and text boxes and treats them as code blocks.
     """
     def __init__(self, **kwargs):
         """
@@ -97,6 +102,295 @@ class AlitaDocxMammothLoader(BaseLoader):
         new_md = pattern.sub(replace_placeholder, original_md)
         return new_md
 
+    def __has_border(self, paragraph):
+        """
+        Check if a paragraph has border formatting.
+        
+        Args:
+            paragraph: A python-docx Paragraph object.
+            
+        Returns:
+            bool: True if paragraph has any border, False otherwise.
+        """
+        pPr = paragraph._element.pPr
+        if pPr is not None:
+            pBdr = pPr.find(qn('w:pBdr'))
+            if pBdr is not None:
+                # Check if any border side exists (top, bottom, left, right)
+                for side in ['top', 'bottom', 'left', 'right']:
+                    border = pBdr.find(qn(f'w:{side}'))
+                    if border is not None:
+                        # Check if border is not "none" or has a width
+                        val = border.get(qn('w:val'))
+                        if val and val != 'none':
+                            return True
+        return False
+
+    def __find_text_boxes(self, doc):
+        """
+        Find all text boxes in document by searching OOXML structure.
+        Text boxes are typically in w:txbxContent elements.
+        
+        Args:
+            doc: A python-docx Document object.
+            
+        Returns:
+            list: List of tuples (element, paragraphs_inside_textbox).
+        """
+        text_boxes = []
+        
+        # Iterate through document body XML to find text box content elements
+        for element in doc.element.body.iter():
+            # Look for text box content elements
+            if element.tag.endswith('txbxContent'):
+                # Collect all paragraphs inside this text box
+                txbx_paragraphs = []
+                for txbx_para_element in element.iter():
+                    if txbx_para_element.tag.endswith('p'):
+                        txbx_paragraphs.append(txbx_para_element)
+                
+                if txbx_paragraphs:
+                    text_boxes.append((element, txbx_paragraphs))
+        
+        return text_boxes
+
+    def __create_marker_paragraph(self, marker_text):
+        """
+        Create a paragraph element with marker text.
+        
+        Args:
+            marker_text (str): The marker text to insert.
+            
+        Returns:
+            Element: An OOXML paragraph element.
+        """
+        from docx.oxml import OxmlElement
+        
+        p = OxmlElement('w:p')
+        r = OxmlElement('w:r')
+        t = OxmlElement('w:t')
+        t.text = marker_text
+        r.append(t)
+        p.append(r)
+        return p
+
+    def __inject_markers_for_paragraph(self, paragraph, start_marker, end_marker):
+        """
+        Inject marker paragraphs before and after a bordered paragraph.
+        
+        Args:
+            paragraph: A python-docx Paragraph object.
+            start_marker (str): The start marker text.
+            end_marker (str): The end marker text.
+        """
+        # Insert start marker paragraph before
+        marker_p_start = self.__create_marker_paragraph(start_marker)
+        paragraph._element.addprevious(marker_p_start)
+        
+        # Insert end marker paragraph after
+        marker_p_end = self.__create_marker_paragraph(end_marker)
+        paragraph._element.addnext(marker_p_end)
+
+    def __inject_markers_for_textbox(self, textbox_element, paragraph_elements, start_marker, end_marker):
+        """
+        Inject markers around text box content.
+        
+        Args:
+            textbox_element: The w:txbxContent element.
+            paragraph_elements: List of paragraph elements inside the text box.
+            start_marker (str): The start marker text.
+            end_marker (str): The end marker text.
+        """
+        if not paragraph_elements:
+            return
+        
+        # Insert start marker before first paragraph in text box
+        first_para = paragraph_elements[0]
+        marker_p_start = self.__create_marker_paragraph(start_marker)
+        first_para.addprevious(marker_p_start)
+        
+        # Insert end marker after last paragraph in text box
+        last_para = paragraph_elements[-1]
+        marker_p_end = self.__create_marker_paragraph(end_marker)
+        last_para.addnext(marker_p_end)
+
+    def __detect_and_mark_bordered_content(self, docx_stream):
+        """
+        Detects bordered paragraphs and text boxes, injects unique markers around them.
+        Groups consecutive bordered paragraphs into single code blocks.
+        
+        Args:
+            docx_stream: A file-like object containing the DOCX document.
+            
+        Returns:
+            tuple: (modified_docx_stream, start_marker, end_marker)
+        """
+        # Load document with python-docx
+        doc = DocxDocument(docx_stream)
+        
+        # Generate unique markers to avoid conflicts with document content
+        unique_id = uuid.uuid4().hex[:8]
+        start_marker = f"<<<BORDERED_BLOCK_START_{unique_id}>>>"
+        end_marker = f"<<<BORDERED_BLOCK_END_{unique_id}>>>"
+        
+        # Group consecutive bordered paragraphs together
+        bordered_groups = []
+        current_group = []
+        
+        for para in doc.paragraphs:
+            if self.__has_border(para):
+                current_group.append(para)
+            else:
+                if current_group:
+                    # End of a bordered group
+                    bordered_groups.append(current_group)
+                    current_group = []
+        
+        # Don't forget the last group if document ends with bordered paragraphs
+        if current_group:
+            bordered_groups.append(current_group)
+        
+        # Collect all text boxes
+        # text_boxes = self.__find_text_boxes(doc)
+        
+        # Inject markers around each group of consecutive bordered paragraphs
+        for group in bordered_groups:
+            if group:
+                # Add start marker before first paragraph in group
+                first_para = group[0]
+                marker_p_start = self.__create_marker_paragraph(start_marker)
+                first_para._element.addprevious(marker_p_start)
+                
+                # Add end marker after last paragraph in group
+                last_para = group[-1]
+                marker_p_end = self.__create_marker_paragraph(end_marker)
+                last_para._element.addnext(marker_p_end)
+        
+        # Inject markers around text box content
+        # for textbox_element, para_elements in text_boxes:
+        #     self.__inject_markers_for_textbox(textbox_element, para_elements, start_marker, end_marker)
+        
+        # Save modified document to BytesIO
+        output = BytesIO()
+        doc.save(output)
+        output.seek(0)
+        
+        return output, start_marker, end_marker
+
+    def __contains_complex_structure(self, content_html):
+        """
+        Check if HTML content contains tables, lists, or other complex structures.
+        
+        Args:
+            content_html (str): HTML content to analyze.
+            
+        Returns:
+            bool: True if content contains tables/lists, False otherwise.
+        """
+        content_soup = BeautifulSoup(content_html, 'html.parser')
+        
+        # Check for tables
+        if content_soup.find('table'):
+            return True
+        
+        # Check for lists (ul, ol)
+        if content_soup.find('ul') or content_soup.find('ol'):
+            return True
+        
+        return False
+
+    def __escape_hash_symbols(self, html_content):
+        """
+        Escape hash (#) symbols at the beginning of lines in HTML to prevent
+        them from being treated as markdown headers.
+        
+        Args:
+            html_content (str): HTML content.
+            
+        Returns:
+            str: HTML with escaped hash symbols.
+        """
+        soup = BeautifulSoup(html_content, 'html.parser')
+        
+        # Process all text-containing elements
+        for element in soup.find_all(['p', 'li', 'td', 'th', 'h1', 'h2', 'h3', 'h4', 'h5', 'h6']):
+            if element.string:
+                text = element.string
+                # If line starts with #, escape it
+                if text.strip().startswith('#'):
+                    element.string = text.replace('#', '\\#', 1)
+        
+        return str(soup)
+
+    def __wrap_marked_sections_in_code_blocks(self, html, start_marker, end_marker):
+        """
+        Find content between markers and wrap appropriately:
+        - Simple text/code → <pre><code> block
+        - Tables/lists → Custom wrapper with preserved structure
+        
+        Args:
+            html (str): The HTML content from Mammoth.
+            start_marker (str): The start marker text.
+            end_marker (str): The end marker text.
+            
+        Returns:
+            str: HTML with marked sections wrapped appropriately.
+        """
+        import html as html_module
+        
+        # Mammoth escapes < and > to &lt; and &gt;, so we need to escape our markers too
+        escaped_start = html_module.escape(start_marker)
+        escaped_end = html_module.escape(end_marker)
+        
+        # Pattern to find content between HTML-escaped markers (including HTML tags)
+        # The markers will be in separate <p> tags, and content in between
+        pattern = re.compile(
+            f'<p>{re.escape(escaped_start)}</p>(.*?)<p>{re.escape(escaped_end)}</p>',
+            re.DOTALL
+        )
+        
+        def replace_with_appropriate_wrapper(match):
+            content = match.group(1)
+            
+            # Detect if content has complex structure (tables, lists)
+            has_complex_structure = self.__contains_complex_structure(content)
+            
+            if has_complex_structure:
+                # Preserve structure: keep HTML as-is, escape # symbols
+                escaped_content = self.__escape_hash_symbols(content)
+                # Wrap in a div with special class for potential custom handling
+                return f'<div class="alita-bordered-content">{escaped_content}</div>'
+            else:
+                # Simple text/code: extract as plain text and wrap in code block
+                content_soup = BeautifulSoup(content, 'html.parser')
+                
+                # Extract text from each paragraph separately to preserve line breaks
+                lines = []
+                for element in content_soup.find_all(['p', 'h1', 'h2', 'h3', 'h4', 'h5', 'h6']):
+                    # Replace <br /> within paragraphs with newlines
+                    for br in element.find_all('br'):
+                        br.replace_with('\n')
+                    text = element.get_text()
+                    # Preserve leading whitespace (indentation), only strip trailing
+                    lines.append(text.rstrip())
+                
+                # If no paragraphs found, just get all text
+                if not lines:
+                    content = content.replace('<br />', '\n').replace('<br/>', '\n').replace('<br>', '\n')
+                    content_text = content_soup.get_text()
+                    lines = [line.rstrip() for line in content_text.split('\n')]
+                
+                # Join lines, strip only leading/trailing empty lines
+                content_text = '\n'.join(lines).strip()
+                # Return as code block (need to HTML-escape the content)
+                content_escaped = html_module.escape(content_text)
+                return f'<pre><code>{content_escaped}</code></pre>'
+        
+        # Replace all marked sections with appropriate wrappers
+        result_html = pattern.sub(replace_with_appropriate_wrapper, html)
+        
+        return result_html
+
     def load(self):
         """
         Loads and converts the Docx file to markdown format.
@@ -131,6 +425,7 @@ class AlitaDocxMammothLoader(BaseLoader):
     def _convert_docx_to_markdown(self, docx_file):
         """
         Converts the content of a Docx file to markdown format.
+        Detects bordered content and treats it as code blocks.
 
         Args:
             docx_file (BinaryIO): The Docx file object.
@@ -138,11 +433,28 @@ class AlitaDocxMammothLoader(BaseLoader):
         Returns:
             str: The markdown content extracted from the Docx file.
         """
+        # Step 1: Detect and mark bordered content
+        # Reset stream position if needed
+        if hasattr(docx_file, 'seek'):
+            docx_file.seek(0)
+        
+        marked_docx, start_marker, end_marker = self.__detect_and_mark_bordered_content(docx_file)
+        
+        # Step 2: Convert marked DOCX to HTML using Mammoth
         if self.extract_images:
             # Extract images using the provided image handler
-            result = convert_to_html(docx_file, convert_image=mammoth.images.img_element(self.__handle_image))
+            result = convert_to_html(marked_docx, convert_image=mammoth.images.img_element(self.__handle_image))
         else:
             # Ignore images
-            result = convert_to_html(docx_file, convert_image=lambda image: "")
-        content = markdownify(result.value, heading_style="ATX")
+            result = convert_to_html(marked_docx, convert_image=lambda image: "")
+        
+        # Step 3: Wrap marked sections in <pre><code> tags
+        html_with_code_blocks = self.__wrap_marked_sections_in_code_blocks(
+            result.value, start_marker, end_marker
+        )
+        
+        # Step 4: Convert HTML to markdown
+        content = markdownify(html_with_code_blocks, heading_style="ATX")
+        
+        # Step 5: Post-process markdown (for image transcripts, etc.)
         return self.__postprocess_original_md(content)

alita_sdk/runtime/langchain/document_loaders/AlitaExcelLoader.py

@@ -12,61 +12,239 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 import io
+import os
 from typing import Iterator
 import pandas as pd
 from json import loads
 
-from langchain_core.tools import ToolException
+from openpyxl import load_workbook
+from xlrd import open_workbook
+from langchain_core.documents import Document
 from .AlitaTableLoader import AlitaTableLoader
-    
+from alita_sdk.runtime.langchain.constants import LOADER_MAX_TOKENS_DEFAULT
 
-class AlitaExcelLoader(AlitaTableLoader):
+cell_delimiter = " | "
 
-    excel_by_sheets: bool = False
+class AlitaExcelLoader(AlitaTableLoader):
     sheet_name: str = None
-    return_type: str = 'str'
+    file_name: str = None
+    max_tokens: int = LOADER_MAX_TOKENS_DEFAULT
+    add_header_to_chunks: bool = False
+    header_row_number: int = 1
 
     def __init__(self, **kwargs):
         if not kwargs.get('file_path'):
             file_content = kwargs.get('file_content')
             if file_content:
+                self.file_name = kwargs.get('file_name')
                 kwargs['file_path'] = io.BytesIO(file_content)
+        else:
+            self.file_name = kwargs.get('file_path')
         super().__init__(**kwargs)
-        self.excel_by_sheets = kwargs.get('excel_by_sheets')
-        self.return_type = kwargs.get('return_type')
         self.sheet_name = kwargs.get('sheet_name')
+        # Set and validate chunking parameters only once
+        self.max_tokens = int(kwargs.get('max_tokens', LOADER_MAX_TOKENS_DEFAULT))
+        self.add_header_to_chunks = bool(kwargs.get('add_header_to_chunks', False))
+        header_row_number = kwargs.get('header_row_number', 1)
+        # Validate header_row_number
+        try:
+            header_row_number = int(header_row_number)
+            if header_row_number > 0:
+                self.header_row_number = header_row_number
+            else:
+                self.header_row_number = 1
+                self.add_header_to_chunks = False
+        except (ValueError, TypeError):
+            self.header_row_number = 1
+            self.add_header_to_chunks = False
 
     def get_content(self):
         try:
-            dfs = pd.read_excel(self.file_path, sheet_name=self.sheet_name)
-
-            if self.excel_by_sheets:
-                result = {}
-                for sheet_name, df in dfs.items():
-                    df.fillna('', inplace=True)
-                    result[sheet_name] = self.parse_sheet(df)
-                return result
+            # Determine file extension
+            file_extension = os.path.splitext(self.file_name)[-1].lower()
+
+            if file_extension == '.xlsx':
+                # Use openpyxl for .xlsx files
+                return self._read_xlsx()
+            elif file_extension == '.xls':
+                # Use xlrd for .xls files
+                return self._read_xls()
             else:
-                result = []
-                for sheet_name, df in dfs.items():
-                    string_content = self.parse_sheet(df)
-                    result.append(f"====== Sheet name: {sheet_name} ======\n{string_content}")
-                return "\n\n".join(result)
+                raise ValueError(f"Unsupported file format: {file_extension}")
         except Exception as e:
-            return ToolException(f"Error reading Excel file: {e}")
+            return f"Error reading Excel file: {e}"
 
-    def parse_sheet(self, df):
-        df.fillna('', inplace=True)
+    def _read_xlsx(self):
+        """
+        Reads .xlsx files using openpyxl.
+        """
+        workbook = load_workbook(self.file_path, data_only=True)  # `data_only=True` ensures we get cell values, not formulas
+        sheets = workbook.sheetnames
+        if self.sheet_name:
+            if self.sheet_name in sheets:
+                sheet_content = self.parse_sheet(workbook[self.sheet_name])
+            else:
+                sheet_content = [f"Sheet '{self.sheet_name}' does not exist in the workbook."]
+            return {self.sheet_name: sheet_content}
+        else:
+            # Dictionary comprehension for all sheets
+            return {name: self.parse_sheet(workbook[name]) for name in sheets}
 
-        if self.return_type == 'dict':
-            return df.to_dict(orient='records')
-        elif self.return_type == 'csv':
-            return df.to_csv()
+    def _read_xls(self):
+        """
+        Reads .xls files using xlrd.
+        """
+        workbook = open_workbook(filename=self.file_name, file_contents=self.file_content)
+        sheets = workbook.sheet_names()
+        if self.sheet_name:
+            if self.sheet_name in sheets:
+                sheet = workbook.sheet_by_name(self.sheet_name)
+                return {self.sheet_name: self.parse_sheet_xls(sheet)}
+            else:
+                return {self.sheet_name: [f"Sheet '{self.sheet_name}' does not exist in the workbook."]}
         else:
-            return df.to_string(index=False)
+            # Dictionary comprehension for all sheets
+            return {name: self.parse_sheet_xls(workbook.sheet_by_name(name)) for name in sheets}
+
+    def parse_sheet(self, sheet):
+        """
+        Parses a single .xlsx sheet, extracting text and hyperlinks, and formats them.
+        """
+        sheet_content = []
+
+        for row in sheet.iter_rows():
+            row_content = []
+            for cell in row:
+                if cell.hyperlink:
+                    # If the cell has a hyperlink, format it as Markdown
+                    hyperlink = cell.hyperlink.target
+                    cell_value = cell.value or ''  # Use cell value or empty string
+                    row_content.append(f"[{cell_value}]({hyperlink})")
+                else:
+                    # If no hyperlink, use the cell value (computed value if formula)
+                    row_content.append(str(cell.value) if cell.value is not None else "")
+            # Join the row content into a single line using `|` as the delimiter
+            sheet_content.append(cell_delimiter.join(row_content))
+
+        # Format the sheet content based on the return type
+        return self._format_sheet_content(sheet_content)
+
+    def parse_sheet_xls(self, sheet):
+        """
+        Parses a single .xls sheet using xlrd, extracting text and hyperlinks, and formats them.
+        """
+        sheet_content = []
+
+        # Extract hyperlink map (if available)
+        hyperlink_map = getattr(sheet, 'hyperlink_map', {})
+
+        for row_idx in range(sheet.nrows):
+            row_content = []
+            for col_idx in range(sheet.ncols):
+                cell = sheet.cell(row_idx, col_idx)
+                cell_value = cell.value
+
+                # Check if the cell has a hyperlink
+                cell_address = (row_idx, col_idx)
+                if cell_address in hyperlink_map:
+                    hyperlink = hyperlink_map[cell_address].url_or_path
+                    if cell_value:
+                        row_content.append(f"[{cell_value}]({hyperlink})")
+                else:
+                    row_content.append(str(cell_value) if cell_value is not None else "")
+            # Join the row content into a single line using `|` as the delimiter
+            sheet_content.append(cell_delimiter.join(row_content))
+
+        # Format the sheet content based on the return type
+        return self._format_sheet_content(sheet_content)
+
+    def _format_sheet_content(self, rows):
+        """
+        Specification:
+        Formats a list of sheet rows into a list of string chunks according to the following rules:
+        1. If max_tokens < 1, returns a single chunk (list of one string) with all rows joined by a newline ('\n').
+           - If add_header_to_chunks is True and header_row_number is valid, the specified header row is prepended as the first line.
+        2. If max_tokens >= 1:
+           a. Each chunk is a string containing one or more rows, separated by newlines ('\n'), such that the total token count (as measured by tiktoken) does not exceed max_tokens.
+           b. If add_header_to_chunks is True and header_row_number is valid, the specified header row is prepended once at the top of each chunk (not before every row).
+           c. If a single row exceeds max_tokens, it is placed in its own chunk without splitting, with the header prepended if applicable.
+        3. Returns: List[str], where each string is a chunk ready for further processing.
+        """
+        import tiktoken
+        encoding = tiktoken.get_encoding('cl100k_base')
+
+        # --- Inner functions ---
+        def count_tokens(text):
+            """Count tokens in text using tiktoken encoding."""
+            return len(encoding.encode(text))
+
+        def finalize_chunk(chunk_rows):
+            """Join rows for a chunk, prepending header if needed."""
+            if self.add_header_to_chunks and header:
+                return '\n'.join([header] + chunk_rows)
+            else:
+                return '\n'.join(chunk_rows)
+        # --- End inner functions ---
+
+        # If max_tokens < 1, return all rows as a single chunk
+        if self.max_tokens < 1:
+            return ['\n'.join(rows)]
+
+        # Extract header if needed
+        header = None
+        if self.add_header_to_chunks and rows:
+            header_idx = self.header_row_number - 1
+            header = rows.pop(header_idx)
+
+        chunks = []  # List to store final chunks
+        current_chunk = []  # Accumulate rows for the current chunk
+        current_tokens = 0  # Token count for the current chunk
+
+        for row in rows:
+            row_tokens = count_tokens(row)
+            # If row itself exceeds max_tokens, flush current chunk and add row as its own chunk (with header if needed)
+            if row_tokens > self.max_tokens:
+                if current_chunk:
+                    chunks.append(finalize_chunk(current_chunk))
+                    current_chunk = []
+                    current_tokens = 0
+                # Add the large row as its own chunk, with header if needed
+                if self.add_header_to_chunks and header:
+                    chunks.append(finalize_chunk([row]))
+                else:
+                    chunks.append(row)
+                continue
+            # If adding row would exceed max_tokens, flush current chunk and start new
+            if current_tokens + row_tokens > self.max_tokens:
+                if current_chunk:
+                    chunks.append(finalize_chunk(current_chunk))
+                current_chunk = [row]
+                current_tokens = row_tokens
+            else:
+                current_chunk.append(row)
+                current_tokens += row_tokens
+        # Add any remaining rows as the last chunk
+        if current_chunk:
+            chunks.append(finalize_chunk(current_chunk))
+        return chunks
+
+    def load(self) -> list:
+        docs = []
+        content_per_sheet = self.get_content()
+        # content_per_sheet is a dict of sheet_name: list of chunk strings
+        for sheet_name, content_chunks in content_per_sheet.items():
+            metadata = {
+                "source": f'{self.file_path}:{sheet_name}',
+                "sheet_name": sheet_name,
+                "file_type": "excel",
+            }
+            # Each chunk is a separate Document
+            for chunk in content_chunks:
+                docs.append(Document(page_content=chunk, metadata=metadata))
+        return docs
 
-    def read(self):
-        df = pd.read_excel(self.file_path, sheet_name=None)
+    def read(self, lazy: bool = False):
+        df = pd.read_excel(self.file_path, sheet_name=None, engine='calamine')
         docs = []
         for key in df.keys():
             if self.raw_content:
@@ -77,7 +255,7 @@ class AlitaExcelLoader(AlitaTableLoader):
         return docs
 
     def read_lazy(self) -> Iterator[dict]:
-        df = pd.read_excel(self.file_path, sheet_name=None)
+        df = pd.read_excel(self.file_path, sheet_name=None, engine='calamine')
         for key in df.keys():
             if self.raw_content:
                 yield df[key].to_string()

alita_sdk/runtime/langchain/document_loaders/AlitaImageLoader.py

@@ -151,5 +151,5 @@ class AlitaImageLoader(BaseLoader):
         """Load text from image using OCR or LLM if llm is provided, supports SVG."""
         text_content = self.get_content()
 
-        metadata = {"source": str(self.file_path)}  # Ensure source is always a string for metadata
+        metadata = {"source": str(self.file_path if hasattr(self, 'file_path') else self.file_name)}  # Ensure source is always a string for metadata
         return [Document(page_content=text_content, metadata=metadata)]