adds vision to chat

Signed-off-by: Ian Eaves <[email protected]>

Author: ieaves

pyproject.toml

@@ -77,6 +77,26 @@ dictionary = ".codespelldict"
 ignore-words-list = ["cann", "clos", "creat", "ro", "hastable", "shouldnot", "mountns", "passt" ,"assertin"]
 check-hidden = true
 
+[tool.ruff]
+line-length = 120
+target-version = "py312"
+include = ["\\.pyi?$"]
+exclude = [
+  "/\\.git",
+  "/\\.tox",
+  "/\\.venv",
+  "/\\.history",
+  "/build",
+  "/dist",
+  "/docs",
+  "/hack",
+  "/venv"
+]
+
+[tool.ruff.format]
+preview = true
+quote-style = "preserve"
+
 [tool.pytest.ini_options]
 testpaths = ["."]
 log_cli = true

ramalama/chat.py

@@ -14,7 +14,7 @@
 from ramalama.config import CONFIG
 from ramalama.console import EMOJI, should_colorize
 from ramalama.engine import dry_run, stop_container
-from ramalama.file_upload.file_loader import FileUpLoader
+from ramalama.file_loaders.file_manager import OpanAIChatAPIMessageBuilder
 from ramalama.logger import logger
 
 
@@ -90,10 +90,9 @@ def prep_rag_message(self):
         if (context := getattr(self.args, "rag", None)) is None:
             return
 
-        if not (message_content := FileUpLoader(context).load()):
-            return
-
-        self.conversation_history.append({"role": "system", "content": message_content})
+        builder = OpanAIChatAPIMessageBuilder()
+        messages = builder.load(context)
+        self.conversation_history.extend(messages)
 
     def handle_args(self):
         prompt = " ".join(self.args.ARGS) if self.args.ARGS else None

ramalama/file_loaders/__init__.py

@@ -0,0 +1,3 @@
+from ramalama.file_loaders import file_manager, file_types
+
+__all__ = ["file_manager", "file_types"]

ramalama/file_loaders/file_manager.py

@@ -0,0 +1,129 @@
+import os
+from abc import ABC, abstractmethod
+from string import Template
+from typing import Dict, List, Type
+from warnings import warn
+
+from ramalama.file_loaders.file_types import base, image, txt
+
+
+class BaseFileManager(ABC):
+    """
+    Base class for file upload handlers.
+    This class should be extended by specific file type handlers.
+    """
+
+    def __init__(self):
+        self.loaders = {ext.lower(): loader() for loader in self.get_loaders() for ext in loader.file_extensions()}
+
+    def _get_loader(self, file: str) -> base.BaseFileLoader:
+        return self.loaders[os.path.splitext(file)[1].lower()]
+
+    @abstractmethod
+    def load(self):
+        pass
+
+    @classmethod
+    @abstractmethod
+    def get_loaders(cls) -> List[Type[base.BaseFileLoader]]:
+        pass
+
+
+class TextFileManager(BaseFileManager):
+    def __init__(self, delim_string: str = "<!--start_document $name-->"):
+        self.document_delimiter: Template = Template(delim_string)
+        super().__init__()
+
+    @classmethod
+    def get_loaders(cls) -> List[Type[base.BaseFileLoader]]:
+        return [txt.TXTFileLoader]
+
+    def load(self, files: list[str]) -> str:
+        """
+        Generate the output string by concatenating the processed files.
+        """
+        contents = []
+        for file in files:
+            loader = self._get_loader(file)
+            content = f"\n{self.document_delimiter.substitute(name=file)}\n{loader.load(file)}"
+            contents.append(content)
+
+        return "".join(contents)
+
+
+class ImageFileManager(BaseFileManager):
+    @classmethod
+    def get_loaders(cls) -> List[Type[base.BaseFileLoader]]:
+        return [image.BasicImageFileLoader]
+
+    def load(self, files: list[str]) -> list[str]:
+        """
+        Generate the output string by concatenating the processed image files.
+        """
+        return [self._get_loader(file).load(file) for file in files]
+
+
+def unsupported_files_warning(unsupported_files: list[str], supported_extensions: list[str]):
+    supported_extensions = sorted(supported_extensions)
+    formatted_supported = ", ".join(supported_extensions)
+    formatted_unsupported = "- " + "\n- ".join(unsupported_files)
+    warn(
+        f"""
+⚠️  Unsupported file types detected!
+
+Ramalama supports the following file types:
+{formatted_supported}
+
+The following unsupported files were found and ignored:
+{formatted_unsupported}
+    """.strip()
+    )
+
+
+class OpanAIChatAPIMessageBuilder:
+    def __init__(self):
+        self.text_manager = TextFileManager()
+        self.image_manager = ImageFileManager()
+
+    def partition_files(self, file_path: str) -> tuple[list[str], list[str], list[str]]:
+        if not os.path.exists(file_path):
+            raise ValueError(f"{file_path} does not exist.")
+
+        if not os.path.isdir(file_path):
+            files = [file_path]
+        else:
+            files = [os.path.join(root, name) for root, _, files in os.walk(file_path) for name in files]
+
+        text_files = []
+        image_files = []
+        unsupported_files = []
+
+        for file in files:
+            file_type = os.path.splitext(file)[1].lower()  # Convert to lowercase for case-insensitive matching
+            if file_type in self.text_manager.loaders:
+                text_files.append(file)
+            elif file_type in self.image_manager.loaders:
+                image_files.append(file)
+            else:
+                unsupported_files.append(file)
+
+        return text_files, image_files, unsupported_files
+
+    def supported_extensions(self) -> set[str]:
+        return self.text_manager.loaders.keys() | self.image_manager.loaders.keys()
+
+    def load(self, file_path: str) -> list[dict]:
+        text_files, image_files, unsupported_files = self.partition_files(file_path)
+
+        if unsupported_files:
+            unsupported_files_warning(unsupported_files, list(self.supported_extensions()))
+
+        messages = []
+        if text_files:
+            messages.append({"role": "system", "content": self.text_manager.load(text_files)})
+        if image_files:
+            message = {"role": "system", "content": []}
+            for content in self.image_manager.load(image_files):
+                message['content'].append({"type": "image_url", "image_url": {"url": content}})
+            messages.append(message)
+        return messages

ramalama/file_loaders/file_types/__init__.py

@@ -0,0 +1,3 @@
+from ramalama.file_loaders.file_types import base, image, txt
+
+__all__ = ["base", "txt", "image"]

ramalama/file_upload/file_types/base.py renamed to ramalama/file_loaders/file_types/base.py

@@ -1,19 +1,25 @@
 from abc import ABC, abstractmethod
 
 
-class BaseFileUpload(ABC):
+class BaseFileLoader(ABC):
     """
     Base class for file upload handlers.
     This class should be extended by specific file type handlers.
     """
 
-    def __init__(self, file):
-        self.file = file
-
+    @staticmethod
     @abstractmethod
-    def load(self) -> str:
+    def load(file: str) -> str:
         """
         Load the content of the file.
         This method should be implemented by subclasses to handle specific file types.
         """
         pass
+
+    @staticmethod
+    @abstractmethod
+    def file_extensions() -> set[str]:
+        """
+        Get the file extension supported by this file type handler.
+        """
+        pass

ramalama/file_loaders/file_types/image.py

@@ -0,0 +1,32 @@
+import base64
+import mimetypes
+
+from ramalama.file_loaders.file_types.base import BaseFileLoader
+
+
+class BasicImageFileLoader(BaseFileLoader):
+    @staticmethod
+    def file_extensions() -> set[str]:
+        return {
+            ".jpg",
+            ".jpeg",
+            ".png",
+            ".gif",
+            ".bmp",
+            ".tiff",
+            ".tif",
+            ".webp",
+            ".ico",
+        }
+
+    @staticmethod
+    def load(file: str) -> str:
+        """
+        Load the content of the text file.
+        """
+
+        mime_type, _ = mimetypes.guess_type(file)
+        with open(file, "rb") as f:
+            data = base64.b64encode(f.read()).decode("utf-8")
+
+        return f"data:{mime_type};base64,{data}"

ramalama/file_loaders/file_types/pdf.py

@@ -0,0 +1,15 @@
+from ramalama.file_loaders.file_types.base import BaseFileLoader
+
+
+class PDFFileLoader(BaseFileLoader):
+    @staticmethod
+    def file_extensions() -> set[str]:
+        return {".pdf"}
+
+    @staticmethod
+    def load(file: str) -> str:
+        """
+        Load the content of the PDF file.
+        This method should be implemented to handle PDF file reading.
+        """
+        return ""

ramalama/file_loaders/file_types/txt.py

@@ -0,0 +1,26 @@
+from ramalama.file_loaders.file_types.base import BaseFileLoader
+
+
+class TXTFileLoader(BaseFileLoader):
+    @staticmethod
+    def file_extensions() -> set[str]:
+        return {
+            ".txt",
+            ".sh",
+            ".md",
+            ".yaml",
+            ".yml",
+            ".json",
+            ".csv",
+            ".toml",
+        }
+
+    @staticmethod
+    def load(file: str) -> str:
+        """
+        Load the content of the text file.
+        """
+
+        # TODO: Support for non-default encodings?
+        with open(file, "r") as f:
+            return f.read()