alita-sdk 0.3.263__py3-none-any.whl → 0.3.499__py3-none-any.whl

alita_sdk/runtime/langchain/document_loaders/AlitaMarkdownLoader.py

@@ -0,0 +1,66 @@
+from pathlib import Path
+from typing import Any, List, Union, Generator, Iterator
+from langchain_core.documents import Document
+
+from langchain_community.document_loaders.unstructured import (
+    UnstructuredFileLoader,
+    validate_unstructured_version,
+)
+
+class AlitaMarkdownLoader(UnstructuredFileLoader):
+
+    def __init__(
+        self,
+        file_path: Union[str, Path],
+        mode: str = "elements",
+        chunker_config: dict = None,
+        **unstructured_kwargs: Any,
+    ):
+        """
+        Args:
+            file_path: The path to the Markdown file to load.
+            mode: The mode to use when loading the file. Can be one of "single",
+                "multi", or "all". Default is "single".
+            chunker_config: Configuration dictionary for the markdown chunker.
+            **unstructured_kwargs: Any kwargs to pass to the unstructured.
+        """
+        file_path = str(file_path)
+        validate_unstructured_version("0.4.16")
+        self.chunker_config = chunker_config or {
+            "strip_header": False,
+            "return_each_line": False,
+            "headers_to_split_on": [],
+            "max_tokens": 512,
+            "token_overlap": 10,
+        }
+        super().__init__(file_path=file_path, mode=mode, **unstructured_kwargs)
+
+    def _file_content_generator(self) -> Generator[Document, None, None]:
+        """
+        Creates a generator that yields a single Document object
+        representing the entire content of the Markdown file.
+        """
+        with open(self.file_path, "r", encoding="utf-8") as file:
+            content = file.read()
+        yield Document(page_content=content, metadata={"source": self.file_path})
+
+    def _get_elements(self) -> List[Document]:
+        """
+        Processes the Markdown file using the markdown_chunker and returns the chunks.
+        """
+        from alita_sdk.tools.chunkers.sematic.markdown_chunker import markdown_chunker
+
+        # Create a generator for the file content
+        file_content_generator = self._file_content_generator()
+
+        # Use the markdown_chunker to process the content
+        chunks = markdown_chunker(file_content_generator, config=self.chunker_config)
+
+        # Convert the generator to a list of Document objects
+        return list(chunks)
+
+    def lazy_load(self) -> Iterator[Document]:
+        """Load file."""
+        elements = self._get_elements()
+        self._post_process_elements(elements)
+        yield from elements

alita_sdk/runtime/langchain/document_loaders/AlitaPDFLoader.py

@@ -1,5 +1,8 @@
 import pymupdf
-from langchain_community.document_loaders import PyPDFLoader
+import fitz
+from langchain_community.document_loaders import PyPDFium2Loader
+
+from .ImageParser import ImageParser
 from .utils import perform_llm_prediction_for_image_bytes, create_temp_file
 from langchain_core.tools import ToolException
 
@@ -20,6 +23,7 @@ class AlitaPDFLoader:
         self.headers = kwargs.get('headers', None)
         self.extraction_mode = kwargs.get('extraction_mode', "plain")
         self.extraction_kwargs = kwargs.get('extraction_kwargs', None)
+        self.images_parser=ImageParser(llm=self.llm, prompt=self.prompt)
 
     def get_content(self):
         if hasattr(self, 'file_path'):
@@ -41,8 +45,59 @@ class AlitaPDFLoader:
         return text_content
 
     def read_pdf_page(self, report, page, index):
-        text_content = f'Page: {index}\n'
-        text_content += page.get_text()
+        # Extract text in block format (to more accurately match hyperlinks to text)
+        text_blocks = page.get_text("blocks")  # Returns a list of text blocks
+        words = page.get_text("words")  # Returns words with their coordinates
+
+        # Extract hyperlinks
+        links = page.get_links()
+
+        # Create a list to store the modified text
+        modified_text = []
+
+        for block in text_blocks:
+            block_rect = fitz.Rect(block[:4])  # Coordinates of the text block
+            block_text = block[4]  # The actual text of the block
+
+            # Check if there are hyperlinks intersecting with this text block
+            for link in links:
+                if "uri" in link:  # Ensure this is a hyperlink
+                    link_rect = link["from"]  # Coordinates of the hyperlink area
+                    link_uri = link["uri"]  # The URL of the hyperlink
+
+                    # Expand the hyperlink area slightly to account for inaccuracies
+                    link_rect = fitz.Rect(
+                        link_rect.x0 - 1, link_rect.y0 - 1, link_rect.x1 + 1, link_rect.y1 + 1
+                    )
+
+                    # Find words that are inside the hyperlink area
+                    link_text = []
+                    for word in words:
+                        word_rect = fitz.Rect(word[:4])  # Coordinates of the word
+                        word_text = word[4]
+
+                        # Check if the word rectangle is fully inside the hyperlink rectangle
+                        if link_rect.contains(word_rect):
+                            link_text.append(word_text)
+                        # If the word partially intersects, check vertical alignment
+                        elif link_rect.intersects(word_rect):
+                            # Condition: The word must be on the same line as the hyperlink
+                            if abs(link_rect.y0 - word_rect.y0) < 2 and abs(link_rect.y1 - word_rect.y1) < 2:
+                                link_text.append(word_text)
+
+                    # Format the hyperlink in Markdown
+                    full_text = " ".join(link_text) if link_text else "No text"
+                    hyperlink = f"[{full_text}]({link_uri})"
+
+                    # Replace the hyperlink text in the block with the formatted hyperlink
+                    block_text = block_text.replace(full_text, hyperlink)
+
+            # Add the processed text block to the result
+            modified_text.append(block_text)
+
+        # Combine all text blocks into the final text for the page
+        text_content = f'Page: {index}\n' + "\n".join(modified_text)
+
         if self.extract_images:
             images = page.get_images(full=True)
             for i, img in enumerate(images):
@@ -54,10 +109,24 @@ class AlitaPDFLoader:
 
     def load(self):
         if not hasattr(self, 'file_path'):
-            self.file_path = create_temp_file(self.file_content)
-        return PyPDFLoader(file_path=self.file_path,
-                           password=self.password,
-                           headers=self.headers,
-                           extract_images=self.extract_images,
-                           extraction_mode=self.extraction_mode,
-                           extraction_kwargs=self.extraction_kwargs).load()
+            import tempfile
+
+            with tempfile.NamedTemporaryFile(mode='w+b', delete=True, suffix=".pdf") as temp_file:
+                temp_file.write(self.file_content)
+                temp_file.flush()
+                self.file_path = temp_file.name
+                return self._load_docs()
+        else:
+            return self._load_docs()
+
+    def _load_docs(self):
+        docs = PyPDFium2Loader(
+                file_path = self.file_path,
+                password=self.password,
+                headers=self.headers,
+                extract_images = self.extract_images,
+                images_parser = ImageParser(llm=self.llm, prompt=self.prompt),
+            ).load()
+        for doc in docs:
+            doc.metadata['chunk_id'] = doc.metadata['page']
+        return docs

alita_sdk/runtime/langchain/document_loaders/AlitaPowerPointLoader.py

@@ -1,10 +1,10 @@
 import io
 
-from langchain_community.document_loaders import UnstructuredPowerPointLoader
 from langchain_core.tools import ToolException
 from pptx import Presentation
 from .utils import perform_llm_prediction_for_image_bytes, create_temp_file
 from pptx.enum.shapes import MSO_SHAPE_TYPE
+from langchain_core.documents import Document
 
 class AlitaPowerPointLoader:
 
@@ -22,33 +22,70 @@ class AlitaPowerPointLoader:
         self.extract_images = unstructured_kwargs.get('extract_images', False)
         self.llm = unstructured_kwargs.get('llm', None)
         self.prompt = unstructured_kwargs.get('prompt', "Describe image")
+        self.pages_per_chunk = unstructured_kwargs.get('pages_per_chunk', 5)
 
     def get_content(self):
-        prs = Presentation(io.BytesIO(self.file_content))
-        text_content = ''
+        if hasattr(self, 'file_path'):
+            with open(self.file_path, 'rb') as f:
+                prs = Presentation(f)
+        elif hasattr(self, 'file_content'):
+            prs = Presentation(io.BytesIO(self.file_content))
+        pages = []
         if self.page_number is not None:
-            text_content += self.read_pptx_slide(prs.slides[self.page_number - 1], self.page_number)
+            pages.append(self.read_pptx_slide(prs.slides[self.page_number - 1], self.page_number))
         else:
             for index, slide in enumerate(prs.slides, start=1):
-                text_content += self.read_pptx_slide(slide, index)
-        return text_content
+                pages.append(self.read_pptx_slide(slide, index))
+        if self.mode == 'single':
+            return "\n".join(pages)
+        if self.mode == 'paged':
+            return pages
+        else:
+            raise ToolException(f"Unknown mode value: {self.mode}. Only 'single', 'paged' values allowed.")
 
     def read_pptx_slide(self, slide, index):
         text_content = f'Slide: {index}\n'
         for shape in slide.shapes:
-            if hasattr(shape, "text"):
-                text_content += shape.text + "\n"
+            if hasattr(shape, "text_frame") and shape.text_frame is not None:
+                for paragraph in shape.text_frame.paragraphs:
+                    for run in paragraph.runs:
+                        if run.hyperlink and run.hyperlink.address:
+                            link_text = run.text.strip() or "Link"
+                            link_url = run.hyperlink.address
+                            text_content += f" [{link_text}]({link_url}) "
+                        else:
+                            text_content += run.text
+                text_content += "\n"
             elif self.extract_images and shape.shape_type == MSO_SHAPE_TYPE.PICTURE:
                 try:
-                    caption = perform_llm_prediction_for_image_bytes(shape.image.blob, self.llm)
+                    caption = perform_llm_prediction_for_image_bytes(shape.image.blob, self.llm, self.prompt)
                 except:
                     caption = "unknown"
                 text_content += "\n**Image Transcript:**\n" + caption + "\n--------------------\n"
-        return text_content
+        return text_content + "\n"
 
     def load(self):
-        if not self.file_path:
-            self.file_path = create_temp_file(self.file_content)
-        return UnstructuredPowerPointLoader(file_path=self.file_path,
-                           mode=self.mode,
-                           **self.unstructured_kwargs).load()
+        content = self.get_content()
+        if isinstance(content, str):
+            yield Document(page_content=content, metadata={})
+        elif isinstance(content, list):
+            chunk = []
+            chunk_count = 0
+            for page_number, page in enumerate(content, start=1):
+                chunk.append(page)
+                if len(chunk) == self.pages_per_chunk:
+                    chunk_content = "\n".join(chunk)
+                    yield Document(
+                        page_content=chunk_content,
+                        metadata={"chunk_number": chunk_count + 1,
+                                  "pages_in_chunk": list(range(page_number - len(chunk) + 1, page_number + 1))}
+                    )
+                    chunk = []
+                    chunk_count += 1
+            if chunk:
+                chunk_content = "\n".join(chunk)
+                yield Document(
+                    page_content=chunk_content,
+                    metadata={"chunk_number": chunk_count + 1,
+                              "pages_in_chunk": list(range(len(content) - len(chunk) + 1, len(content) + 1))}
+                )

alita_sdk/runtime/langchain/document_loaders/AlitaPythonLoader.py

@@ -0,0 +1,9 @@
+from pathlib import Path
+from typing import Union
+from langchain_community.document_loaders.python import PythonLoader
+
+class AlitaPythonLoader(PythonLoader):
+    """Load `Python` files, respecting any non-default encoding if specified."""
+
+    def __init__(self, file_path: Union[str, Path], **kwargs):
+        super().__init__(file_path)

alita_sdk/runtime/langchain/document_loaders/AlitaTableLoader.py

@@ -17,8 +17,6 @@ from langchain_core.documents import Document
 from typing import List, Optional, Iterator
 from json import dumps
 from .utils import cleanse_data
-from ..tools.log import print_log
-
 
 class AlitaTableLoader(BaseLoader):
     def __init__(self,
@@ -65,7 +63,7 @@ class AlitaTableLoader(BaseLoader):
                 "source": f'{self.file_path}:{idx+1}',
                 "table_source": self.file_path,
             }
-            if len(docs) == 0:
+            if len(docs) == 0 and not self.raw_content:
                 header_metadata = metadata.copy()
                 header_metadata["header"] = "true"
                 header = "\t".join([str(value) for value in row.keys()])
@@ -74,7 +72,6 @@ class AlitaTableLoader(BaseLoader):
                 docs.append(Document(page_content=row, metadata=metadata))
                 continue
             if self.json_documents:
-                # print_log(row)
                 metadata['columns'] = list(row.keys())
                 metadata['og_data'] = dumps(row)
                 docs.append(Document(page_content=self.row_processor(row), metadata=metadata))

alita_sdk/runtime/langchain/document_loaders/AlitaTextLoader.py

@@ -1,4 +1,4 @@
-from typing import Iterator
+from typing import Iterator, Generator
 
 from langchain_core.documents import Document
 
@@ -6,6 +6,9 @@ from langchain_community.document_loaders.base import BaseLoader
 from langchain_community.document_loaders.helpers import detect_file_encodings
 from langchain_core.tools import ToolException
 
+from alita_sdk.tools.chunkers import markdown_chunker
+
+
 class AlitaTextLoader(BaseLoader):
 
     def __init__(self, **kwargs):
@@ -19,6 +22,8 @@ class AlitaTextLoader(BaseLoader):
             raise ToolException("'file_path' or 'file_content' parameter should be provided.")
         self.encoding = kwargs.get('encoding', 'utf-8')
         self.autodetect_encoding = kwargs.get('autodetect_encoding', False)
+        self.max_tokens=kwargs.get('max_tokens', 1024)
+        self.token_overlap = kwargs.get('token_overlap', 10)
 
     def get_content(self):
         text = ""
@@ -59,8 +64,16 @@ class AlitaTextLoader(BaseLoader):
 
         return text
 
+    def generate_document(self, text, metadata) -> Generator[Document, None, None]:
+        yield Document(page_content=text, metadata=metadata)
+
     def lazy_load(self) -> Iterator[Document]:
         """Load from file path."""
         text = self.get_content()
         metadata = {"source": str(self.file_path) if hasattr(self, 'file_path') else self.file_name}
-        yield Document(page_content=text, metadata=metadata)
+        chunks = markdown_chunker(file_content_generator=self.generate_document(text, metadata),
+                                  config={
+                                      "max_tokens": self.max_tokens,
+                                      "token_overlap": self.token_overlap
+                                  })
+        yield from chunks

alita_sdk/runtime/langchain/document_loaders/ImageParser.py

@@ -0,0 +1,30 @@
+from typing import Iterator
+
+from langchain_community.document_loaders.parsers.images import BaseImageBlobParser
+from langchain_core.documents import Document
+from langchain_core.documents.base import Blob
+
+from alita_sdk.runtime.langchain.document_loaders.AlitaImageLoader import AlitaImageLoader
+
+class ImageParser(BaseImageBlobParser):
+
+    def __init__(self, **kwargs):
+        self.llm = kwargs.get('llm')
+        self.prompt = kwargs.get('prompt')
+
+    def lazy_parse(self, blob: Blob) -> Iterator[Document]:
+        try:
+            yield from super().lazy_parse(blob)
+        except Exception:
+            yield Document(page_content="[Image: Unknown]")
+
+    def _analyze_image(self, img) -> str:
+        from io import BytesIO
+
+        byte_stream = BytesIO()
+        img.save(byte_stream, format='PNG')
+        image_bytes = byte_stream.getvalue()
+        try:
+            return AlitaImageLoader(file_content=image_bytes, file_name="image.png", prompt=self.prompt, llm=self.llm).get_content()
+        except Exception:
+            return "Image: unknown"