Adds file upload feature

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

Author: ieaves

docs/ramalama-chat.1.md

@@ -29,9 +29,13 @@ Show this help message and exit
 #### **--prefix**
 Prefix for the user prompt (default: 🦭 > )
 
+#### **--rag**=path
+A file or directory of files to be loaded and provided as local context in the chat history.
+
 #### **--url**=URL
 The host to send requests to (default: http://127.0.0.1:8080)
 
+
 ## EXAMPLES
 
 Communicate with the default local OpenAI REST API. (http://127.0.0.1:8080)

ramalama/chat.py

@@ -14,6 +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.logger import logger
 
 
@@ -72,6 +73,16 @@ def __init__(self, args):
         self.prompt = args.prefix
 
         self.url = f"{args.url}/chat/completions"
+        self.prep_rag_message()
+
+    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})
 
     def handle_args(self):
         prompt = " ".join(self.args.ARGS) if self.args.ARGS else None

ramalama/cli.py

@@ -905,6 +905,7 @@ def chat_parser(subparsers):
         help='possible values are "never", "always" and "auto".',
     )
     parser.add_argument("--prefix", type=str, help="prefix for the user prompt", default=default_prefix())
+    parser.add_argument("--rag", type=str, help="a file or directory to use as context for the chat")
     parser.add_argument("--url", type=str, default="http://127.0.0.1:8080/v1", help="the url to send requests to")
     parser.add_argument("MODEL", completer=local_models)  # positional argument
     parser.add_argument(
@@ -925,6 +926,7 @@ def run_parser(subparsers):
     )
     parser.add_argument("--prefix", type=str, help="prefix for the user prompt", default=default_prefix())
     parser.add_argument("MODEL", completer=local_models)  # positional argument
+
     parser.add_argument(
         "ARGS",
         nargs="*",

ramalama/file_upload/__init__.py

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

ramalama/file_upload/file_loader.py

@@ -0,0 +1,60 @@
+import os
+from string import Template
+from typing import Type
+from warnings import warn
+
+from ramalama.file_upload.file_types import base, txt
+
+SUPPORTED_EXTENSIONS = {
+    '.txt': txt.TXTFileUpload,
+    '.sh': txt.TXTFileUpload,
+    '.md': txt.TXTFileUpload,
+    '.yaml': txt.TXTFileUpload,
+    '.yml': txt.TXTFileUpload,
+    '.json': txt.TXTFileUpload,
+    '.csv': txt.TXTFileUpload,
+    '.toml': txt.TXTFileUpload,
+}
+
+
+class BaseFileUploader:
+    """
+    Base class for file upload handlers.
+    This class should be extended by specific file type handlers.
+    """
+
+    def __init__(self, files: list[Type[base.BaseFileUpload]], delim_string: str = "<!--start_document $name-->"):
+        self.files = files
+        self.document_delimiter: Template = Template(delim_string)
+
+    def load(self) -> str:
+        """
+        Generate the output string by concatenating the processed files.
+        """
+        output = (f"\n{self.document_delimiter.substitute(name=f.file)}\n{f.load()}" for f in self.files)
+        return "".join(output)
+
+
+class FileUpLoader(BaseFileUploader):
+    def __init__(self, file_path: 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(file_path, f) for f in os.listdir(file_path) if os.path.isfile(os.path.join(file_path, f))
+            ]
+
+        extensions = [os.path.splitext(f)[1].lower() for f in files]
+
+        if set(extensions) - set(SUPPORTED_EXTENSIONS):
+            warning_message = (
+                f"Unsupported file types found: {set(extensions) - set(SUPPORTED_EXTENSIONS)}\n"
+                f"Supported types are: {set(SUPPORTED_EXTENSIONS.keys())}"
+            )
+            warn(warning_message)
+
+        files = [SUPPORTED_EXTENSIONS[ext](file=f) for ext, f in zip(extensions, files) if ext in SUPPORTED_EXTENSIONS]
+        super().__init__(files=files)

ramalama/file_upload/file_types/__init__.py

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

ramalama/file_upload/file_types/base.py

@@ -0,0 +1,19 @@
+from abc import ABC, abstractmethod
+
+
+class BaseFileUpload(ABC):
+    """
+    Base class for file upload handlers.
+    This class should be extended by specific file type handlers.
+    """
+
+    def __init__(self, file):
+        self.file = file
+
+    @abstractmethod
+    def load(self) -> str:
+        """
+        Load the content of the file.
+        This method should be implemented by subclasses to handle specific file types.
+        """
+        pass

ramalama/file_upload/file_types/pdf.py

@@ -0,0 +1,10 @@
+from ramalama.file_upload.file_types.base import BaseFileUpload
+
+
+class PDFFileUpload(BaseFileUpload):
+    def load(self) -> str:
+        """
+        Load the content of the PDF file.
+        This method should be implemented to handle PDF file reading.
+        """
+        return ""

ramalama/file_upload/file_types/txt.py

@@ -0,0 +1,12 @@
+from ramalama.file_upload.file_types.base import BaseFileUpload
+
+
+class TXTFileUpload(BaseFileUpload):
+    def load(self) -> str:
+        """
+        Load the content of the text file.
+        """
+
+        # TODO: Support for non-default encodings?
+        with open(self.file, 'r') as f:
+            return f.read()

test/unit/conftest.py

@@ -0,0 +1,6 @@
+# # tests/conftest.py
+# import pytest
+
+# @pytest.fixture(autouse=True)
+# def set_container_engine_env(monkeypatch):
+#     monkeypatch.setenv("RAMALAMA_CONTAINER_ENGINE", "docker")