ara-cli 0.1.9.69__py3-none-any.whl → 0.1.10.8__py3-none-any.whl

ara_cli/file_loaders/document_reader.py

@@ -0,0 +1,245 @@
+import os
+from abc import ABC, abstractmethod
+from typing import Tuple, Optional
+
+
+class DocumentReader(ABC):
+    """Abstract base class for document readers"""
+
+    def __init__(self, file_path: str):
+        self.file_path = file_path
+        self.base_dir = os.path.dirname(file_path)
+
+    @abstractmethod
+    def read(self, extract_images: bool = False) -> str:
+        """Read document and optionally extract images"""
+        pass
+
+    def create_image_data_dir(self, extension_suffix: str) -> str:
+        """
+        Create data directory for images with file extension suffix to avoid conflicts.
+
+        Returns:
+            str: Path to images directory
+        """
+        file_name_with_ext = os.path.splitext(os.path.basename(self.file_path))[
+            0] + f"_{extension_suffix}"
+        data_dir = os.path.join(self.base_dir, f"{file_name_with_ext}.data")
+        images_dir = os.path.join(data_dir, "images")
+        if not os.path.exists(images_dir):
+            os.makedirs(images_dir)
+        return images_dir
+
+    def save_and_describe_image(
+        self,
+        image_data: bytes,
+        image_format: str,
+        save_dir: str,
+        image_counter: int
+    ) -> Tuple[str, str]:
+        """
+        Save image data and get its description from LLM.
+
+        Returns:
+            tuple: (relative_image_path, description)
+        """
+        from ara_cli.prompt_handler import describe_image
+
+        # Save image
+        image_filename = f"{image_counter}.{image_format}"
+        image_path = os.path.join(save_dir, image_filename)
+
+        with open(image_path, "wb") as image_file:
+            image_file.write(image_data)
+
+        # Get image description from LLM
+        description = describe_image(image_path)
+
+        # Get relative path
+        relative_image_path = os.path.relpath(image_path, self.base_dir)
+
+        return relative_image_path, description
+
+
+class DocxReader(DocumentReader):
+    """Reader for DOCX files"""
+
+    def read(self, extract_images: bool = False) -> str:
+        import docx
+
+        doc = docx.Document(self.file_path)
+        text_content = '\n'.join(para.text for para in doc.paragraphs)
+
+        if not extract_images:
+            return text_content
+
+        from PIL import Image
+        import io
+
+        # Create data directory for images
+        images_dir = self.create_image_data_dir("docx")
+
+        # Extract and process images
+        image_descriptions = []
+        image_counter = 1
+
+        for rel in doc.part.rels.values():
+            if "image" in rel.reltype:
+                image_data = rel.target_part.blob
+
+                # Determine image format
+                image = Image.open(io.BytesIO(image_data))
+                image_format = image.format.lower()
+
+                # Save and describe image
+                relative_path, description = self.save_and_describe_image(
+                    image_data, image_format, images_dir, image_counter
+                )
+
+                # Add formatted description to list
+                image_description = f"\nImage: {relative_path}\n[{description}]\n"
+                image_descriptions.append(image_description)
+
+                image_counter += 1
+
+        # Combine text content with image descriptions
+        if image_descriptions:
+            text_content += "\n\n### Extracted Images\n" + \
+                "\n".join(image_descriptions)
+
+        return text_content
+
+
+class PdfReader(DocumentReader):
+    """Reader for PDF files"""
+
+    def read(self, extract_images: bool = False) -> str:
+        import pymupdf4llm
+
+        if not extract_images:
+            return pymupdf4llm.to_markdown(self.file_path, write_images=False)
+
+        import fitz  # PyMuPDF
+
+        # Create images directory
+        images_dir = self.create_image_data_dir("pdf")
+
+        # Extract text without images first
+        text_content = pymupdf4llm.to_markdown(
+            self.file_path, write_images=False)
+
+        # Extract and process images
+        doc = fitz.open(self.file_path)
+        image_descriptions = []
+        image_counter = 1
+
+        for page_num, page in enumerate(doc):
+            image_list = page.get_images()
+
+            for img_index, img in enumerate(image_list):
+                # Extract image
+                xref = img[0]
+                base_image = doc.extract_image(xref)
+                image_bytes = base_image["image"]
+                image_ext = base_image["ext"]
+
+                # Save and describe image
+                relative_path, description = self.save_and_describe_image(
+                    image_bytes, image_ext, images_dir, image_counter
+                )
+
+                # Add formatted description to list
+                image_description = f"\nImage: {relative_path}\n[{description}]\n"
+                image_descriptions.append(image_description)
+
+                image_counter += 1
+
+        doc.close()
+
+        # Combine text content with image descriptions
+        if image_descriptions:
+            text_content += "\n\n### Extracted Images\n" + \
+                "\n".join(image_descriptions)
+
+        return text_content
+
+
+class OdtReader(DocumentReader):
+    """Reader for ODT files"""
+
+    def read(self, extract_images: bool = False) -> str:
+        import pymupdf4llm
+
+        if not extract_images:
+            return pymupdf4llm.to_markdown(self.file_path, write_images=False)
+
+        import zipfile
+        from PIL import Image
+        import io
+
+        # Create data directory for images
+        images_dir = self.create_image_data_dir("odt")
+
+        # Get text content
+        text_content = pymupdf4llm.to_markdown(
+            self.file_path, write_images=False)
+
+        # Extract and process images from ODT
+        image_descriptions = []
+        image_counter = 1
+
+        try:
+            with zipfile.ZipFile(self.file_path, 'r') as odt_zip:
+                # List all files in the Pictures directory
+                picture_files = [
+                    f for f in odt_zip.namelist() if f.startswith('Pictures/')]
+
+                for picture_file in picture_files:
+                    # Extract image data
+                    image_data = odt_zip.read(picture_file)
+
+                    # Determine image format
+                    image = Image.open(io.BytesIO(image_data))
+                    image_format = image.format.lower()
+
+                    # Save and describe image
+                    relative_path, description = self.save_and_describe_image(
+                        image_data, image_format, images_dir, image_counter
+                    )
+
+                    # Add formatted description to list
+                    image_description = f"\nImage: {relative_path}\n[{description}]\n"
+                    image_descriptions.append(image_description)
+
+                    image_counter += 1
+        except Exception as e:
+            print(f"Warning: Could not extract images from ODT: {e}")
+
+        # Combine text content with image descriptions
+        if image_descriptions:
+            text_content += "\n\n### Extracted Images\n" + \
+                "\n".join(image_descriptions)
+
+        return text_content
+
+
+class DocumentReaderFactory:
+    """Factory for creating appropriate document readers"""
+
+    @staticmethod
+    def create_reader(file_path: str) -> Optional[DocumentReader]:
+        """Create appropriate reader based on file extension"""
+        _, ext = os.path.splitext(file_path)
+        ext = ext.lower()
+
+        readers = {
+            '.docx': DocxReader,
+            '.pdf': PdfReader,
+            '.odt': OdtReader
+        }
+
+        reader_class = readers.get(ext)
+        if reader_class:
+            return reader_class(file_path)
+
+        return None

ara_cli/file_loaders/document_readers.py

@@ -0,0 +1,233 @@
+import os
+from abc import ABC, abstractmethod
+from typing import Tuple, Optional
+
+
+class DocumentReader(ABC):
+    """Abstract base class for document readers"""
+    
+    def __init__(self, file_path: str):
+        self.file_path = file_path
+        self.base_dir = os.path.dirname(file_path)
+    
+    @abstractmethod
+    def read(self, extract_images: bool = False) -> str:
+        """Read document and optionally extract images"""
+        pass
+    
+    def create_image_data_dir(self, extension_suffix: str) -> str:
+        """
+        Create data directory for images with file extension suffix to avoid conflicts.
+        
+        Returns:
+            str: Path to images directory
+        """
+        file_name_with_ext = os.path.splitext(os.path.basename(self.file_path))[0] + f"_{extension_suffix}"
+        data_dir = os.path.join(self.base_dir, f"{file_name_with_ext}.data")
+        images_dir = os.path.join(data_dir, "images")
+        if not os.path.exists(images_dir):
+            os.makedirs(images_dir)
+        return images_dir
+    
+    def save_and_describe_image(self, image_data: bytes, image_format: str, 
+                               save_dir: str, image_counter: int) -> Tuple[str, str]:
+        """
+        Save image data and get its description from LLM.
+        
+        Returns:
+            tuple: (relative_image_path, description)
+        """
+        from ara_cli.prompt_handler import describe_image
+        
+        # Save image
+        image_filename = f"{image_counter}.{image_format}"
+        image_path = os.path.join(save_dir, image_filename)
+        
+        with open(image_path, "wb") as image_file:
+            image_file.write(image_data)
+        
+        # Get image description from LLM
+        description = describe_image(image_path)
+        
+        # Get relative path
+        relative_image_path = os.path.relpath(image_path, self.base_dir)
+        
+        return relative_image_path, description
+
+
+class DocxReader(DocumentReader):
+    """Reader for DOCX files"""
+    
+    def read(self, extract_images: bool = False) -> str:
+        import docx
+        
+        doc = docx.Document(self.file_path)
+        text_content = '\n'.join(para.text for para in doc.paragraphs)
+        
+        if not extract_images:
+            return text_content
+        
+        from PIL import Image
+        import io
+        
+        # Create data directory for images
+        images_dir = self.create_image_data_dir("docx")
+        
+        # Extract and process images
+        image_descriptions = []
+        image_counter = 1
+        
+        for rel in doc.part.rels.values():
+            if "image" in rel.reltype:
+                image_data = rel.target_part.blob
+                
+                # Determine image format
+                image = Image.open(io.BytesIO(image_data))
+                image_format = image.format.lower()
+                
+                # Save and describe image
+                relative_path, description = self.save_and_describe_image(
+                    image_data, image_format, images_dir, image_counter
+                )
+                
+                # Add formatted description to list
+                image_description = f"\nImage: {relative_path}\n[{description}]\n"
+                image_descriptions.append(image_description)
+                
+                image_counter += 1
+        
+        # Combine text content with image descriptions
+        if image_descriptions:
+            text_content += "\n\n### Extracted Images\n" + "\n".join(image_descriptions)
+        
+        return text_content
+
+
+class PdfReader(DocumentReader):
+    """Reader for PDF files"""
+    
+    def read(self, extract_images: bool = False) -> str:
+        import pymupdf4llm
+        
+        if not extract_images:
+            return pymupdf4llm.to_markdown(self.file_path, write_images=False)
+        
+        import fitz  # PyMuPDF
+        
+        # Create images directory
+        images_dir = self.create_image_data_dir("pdf")
+        
+        # Extract text without images first
+        text_content = pymupdf4llm.to_markdown(self.file_path, write_images=False)
+        
+        # Extract and process images
+        doc = fitz.open(self.file_path)
+        image_descriptions = []
+        image_counter = 1
+        
+        for page_num, page in enumerate(doc):
+            image_list = page.get_images()
+            
+            for img_index, img in enumerate(image_list):
+                # Extract image
+                xref = img[0]
+                base_image = doc.extract_image(xref)
+                image_bytes = base_image["image"]
+                image_ext = base_image["ext"]
+                
+                # Save and describe image
+                relative_path, description = self.save_and_describe_image(
+                    image_bytes, image_ext, images_dir, image_counter
+                )
+                
+                # Add formatted description to list
+                image_description = f"\nImage: {relative_path}\n[{description}]\n"
+                image_descriptions.append(image_description)
+                
+                image_counter += 1
+        
+        doc.close()
+        
+        # Combine text content with image descriptions
+        if image_descriptions:
+            text_content += "\n\n### Extracted Images\n" + "\n".join(image_descriptions)
+        
+        return text_content
+
+
+class OdtReader(DocumentReader):
+    """Reader for ODT files"""
+    
+    def read(self, extract_images: bool = False) -> str:
+        import pymupdf4llm
+        
+        if not extract_images:
+            return pymupdf4llm.to_markdown(self.file_path, write_images=False)
+        
+        import zipfile
+        from PIL import Image
+        import io
+        
+        # Create data directory for images
+        images_dir = self.create_image_data_dir("odt")
+        
+        # Get text content
+        text_content = pymupdf4llm.to_markdown(self.file_path, write_images=False)
+        
+        # Extract and process images from ODT
+        image_descriptions = []
+        image_counter = 1
+        
+        try:
+            with zipfile.ZipFile(self.file_path, 'r') as odt_zip:
+                # List all files in the Pictures directory
+                picture_files = [f for f in odt_zip.namelist() if f.startswith('Pictures/')]
+                
+                for picture_file in picture_files:
+                    # Extract image data
+                    image_data = odt_zip.read(picture_file)
+                    
+                    # Determine image format
+                    image = Image.open(io.BytesIO(image_data))
+                    image_format = image.format.lower()
+                    
+                    # Save and describe image
+                    relative_path, description = self.save_and_describe_image(
+                        image_data, image_format, images_dir, image_counter
+                    )
+                    
+                    # Add formatted description to list
+                    image_description = f"\nImage: {relative_path}\n[{description}]\n"
+                    image_descriptions.append(image_description)
+                    
+                    image_counter += 1
+        except Exception as e:
+            print(f"Warning: Could not extract images from ODT: {e}")
+        
+        # Combine text content with image descriptions
+        if image_descriptions:
+            text_content += "\n\n### Extracted Images\n" + "\n".join(image_descriptions)
+        
+        return text_content
+
+
+class DocumentReaderFactory:
+    """Factory for creating appropriate document readers"""
+    
+    @staticmethod
+    def create_reader(file_path: str) -> Optional[DocumentReader]:
+        """Create appropriate reader based on file extension"""
+        _, ext = os.path.splitext(file_path)
+        ext = ext.lower()
+        
+        readers = {
+            '.docx': DocxReader,
+            '.pdf': PdfReader,
+            '.odt': OdtReader
+        }
+        
+        reader_class = readers.get(ext)
+        if reader_class:
+            return reader_class(file_path)
+        
+        return None

ara_cli/file_loaders/file_loader.py

@@ -0,0 +1,50 @@
+from abc import ABC, abstractmethod
+from typing import Optional
+
+
+class FileLoader(ABC):
+    """Abstract base class for file loaders"""
+
+    def __init__(self, chat_instance):
+        self.chat = chat_instance
+
+    @abstractmethod
+    def load(self, file_path: str, **kwargs) -> bool:
+        """Load file with specific implementation"""
+        pass
+
+    def add_prompt_tag_if_needed(self):
+        """Add prompt tag to chat if needed"""
+        self.chat.add_prompt_tag_if_needed(self.chat.chat_name)
+
+
+class FileLoaderFactory:
+    """Factory for creating appropriate file loaders"""
+    BINARY_TYPE_MAPPING = {
+        ".png": "image/png",
+        ".jpg": "image/jpeg",
+        ".jpeg": "image/jpeg",
+    }
+
+    DOCUMENT_TYPE_EXTENSIONS = [".docx", ".doc", ".odt", ".pdf"]
+
+    @staticmethod
+    def create_loader(file_name: str, chat_instance) -> Optional[FileLoader]:
+        """Create appropriate loader based on file type"""
+        from ara_cli.file_loaders.binary_file_loader import BinaryFileLoader
+        from ara_cli.file_loaders.text_file_loader import TextFileLoader
+        from ara_cli.file_loaders.document_file_loader import DocumentFileLoader
+
+        file_name_lower = file_name.lower()
+
+        # Check if it's a binary file
+        for extension, mime_type in FileLoaderFactory.BINARY_TYPE_MAPPING.items():
+            if file_name_lower.endswith(extension):
+                return BinaryFileLoader(chat_instance)
+
+        # Check if it's a document
+        if any(file_name_lower.endswith(ext) for ext in FileLoaderFactory.DOCUMENT_TYPE_EXTENSIONS):
+            return DocumentFileLoader(chat_instance)
+
+        # Default to text file loader
+        return TextFileLoader(chat_instance)

ara_cli/file_loaders/file_loaders.py

@@ -0,0 +1,123 @@
+import os
+import base64
+from abc import ABC, abstractmethod
+from typing import Optional
+from ara_cli.file_loaders.markdown_reader import MarkdownReader
+from ara_cli.file_loaders.document_readers import DocumentReaderFactory
+
+
+class FileLoader(ABC):
+    """Abstract base class for file loaders"""
+    
+    def __init__(self, chat_instance):
+        self.chat = chat_instance
+        
+    @abstractmethod
+    def load(self, file_path: str, **kwargs) -> bool:
+        """Load file with specific implementation"""
+        pass
+        
+    def add_prompt_tag_if_needed(self):
+        """Add prompt tag to chat if needed"""
+        self.chat.add_prompt_tag_if_needed(self.chat.chat_name)
+
+
+class TextFileLoader(FileLoader):
+    """Loads text files"""
+    
+    def load(self, file_path: str, prefix: str = "", suffix: str = "", 
+             block_delimiter: str = "", extract_images: bool = False) -> bool:
+        """Load text file with optional markdown image extraction"""
+        
+        is_md_file = file_path.lower().endswith('.md')
+        
+        if is_md_file and extract_images:
+            reader = MarkdownReader(file_path)
+            file_content = reader.read(extract_images=True)
+        else:
+            with open(file_path, 'r', encoding='utf-8', errors="replace") as file:
+                file_content = file.read()
+        
+        if block_delimiter:
+            file_content = f"{block_delimiter}\n{file_content}\n{block_delimiter}"
+        
+        write_content = f"{prefix}{file_content}{suffix}\n"
+        
+        with open(self.chat.chat_name, 'a', encoding='utf-8') as chat_file:
+            chat_file.write(write_content)
+            
+        return True
+
+
+class BinaryFileLoader(FileLoader):
+    """Loads binary files (images)"""
+    
+    def load(self, file_path: str, mime_type: str, prefix: str = "", suffix: str = "") -> bool:
+        """Load binary file as base64"""
+        
+        with open(file_path, 'rb') as file:
+            file_content = file.read()
+            
+        base64_image = base64.b64encode(file_content).decode("utf-8")
+        write_content = f"{prefix}![{os.path.basename(file_path)}](data:{mime_type};base64,{base64_image}){suffix}\n"
+        
+        with open(self.chat.chat_name, 'a', encoding='utf-8') as chat_file:
+            chat_file.write(write_content)
+            
+        return True
+
+
+class DocumentFileLoader(FileLoader):
+    """Loads document files (PDF, DOCX, ODT)"""
+    
+    def load(self, file_path: str, prefix: str = "", suffix: str = "", 
+             block_delimiter: str = "```", extract_images: bool = False) -> bool:
+        """Load document file with optional image extraction"""
+        
+        reader = DocumentReaderFactory.create_reader(file_path)
+        
+        if not reader:
+            print("Unsupported document type.")
+            return False
+        
+        text_content = reader.read(extract_images=extract_images)
+        
+        if block_delimiter:
+            text_content = f"{block_delimiter}\n{text_content}\n{block_delimiter}"
+        
+        write_content = f"{prefix}{text_content}{suffix}\n"
+        
+        with open(self.chat.chat_name, 'a', encoding='utf-8') as chat_file:
+            chat_file.write(write_content)
+            
+        return True
+
+
+class FileLoaderFactory:
+    """Factory for creating appropriate file loaders"""
+    
+    BINARY_TYPE_MAPPING = {
+        ".png": "image/png",
+        ".jpg": "image/jpeg",
+        ".jpeg": "image/jpeg",
+    }
+    
+    DOCUMENT_TYPE_EXTENSIONS = [".docx", ".doc", ".odt", ".pdf"]
+    
+    @staticmethod
+    def create_loader(file_name: str, chat_instance) -> Optional[FileLoader]:
+        """Create appropriate loader based on file type"""
+        
+        file_name_lower = file_name.lower()
+        
+        # Check if it's a binary file
+        for extension, mime_type in FileLoaderFactory.BINARY_TYPE_MAPPING.items():
+            if file_name_lower.endswith(extension):
+                return BinaryFileLoader(chat_instance)
+        
+        # Check if it's a document
+        if any(file_name_lower.endswith(ext) for ext in FileLoaderFactory.DOCUMENT_TYPE_EXTENSIONS):
+            return DocumentFileLoader(chat_instance)
+        
+        # Default to text file loader
+        return TextFileLoader(chat_instance)