Merge pull request #41 from andrej1991/assisted-chat

Adding the ability to choose from a text (sentence) processor or a markdown processor

Author: tisnik

src/lightspeed_rag_content/document_processor.py

@@ -25,6 +25,7 @@
 import faiss
 from llama_index.core import Settings, SimpleDirectoryReader, VectorStoreIndex
 from llama_index.core.llms.utils import resolve_llm
+from llama_index.core.node_parser import MarkdownNodeParser
 from llama_index.core.readers.base import BaseReader
 from llama_index.core.schema import Document, TextNode
 from llama_index.core.storage.storage_context import StorageContext
@@ -72,6 +73,8 @@ def __init__(self, config: _Config):
                 model_name=str(self.config.embeddings_model_dir)
             )
             Settings.llm = resolve_llm(None)
+        if config.doc_type == "markdown":
+            Settings.node_parser = MarkdownNodeParser()
 
     @staticmethod
     def _got_whitespace(text: str) -> bool:
@@ -391,6 +394,7 @@ def __init__(
         vector_store_type: str = "faiss",
         table_name: Optional[str] = None,
         manual_chunking: bool = True,
+        doc_type: str = "text",
     ):
         """Initialize instance."""
         if vector_store_type == "postgres" and not table_name:
@@ -406,6 +410,7 @@ def __init__(
             vector_store_type=vector_store_type,
             table_name=table_name,
             manual_chunking=manual_chunking,
+            doc_type=doc_type,
         )
 
         self._check_config(self.config)

src/lightspeed_rag_content/utils.py

@@ -68,4 +68,11 @@ def get_common_arg_parser() -> argparse.ArgumentParser:
         help="How to do the chunking for llama-stack, manually like in "
         "llama-index or automatically using the RAG runtime tool.",
     )
+    parser.add_argument(
+        "-dt",
+        "--document-type",
+        dest="doc_type",
+        default="text",
+        help="The type of the document which is to be added to the RAG",
+    )
     return parser

tests/test_document_processor.py

@@ -52,6 +52,7 @@ def setUp(self):
             "model_name": "sentence-transformers/all-mpnet-base-v2",
             "embeddings_model_dir": "embeddings_model",
             "num_workers": 10,
+            "doc_type": "text",
         }
         self.log = self.patch_object(document_processor, "LOG")
         self.indexdb = self.patch_object(document_processor, "_LlamaIndexDB")

tests/test_document_processor_llama_stack.py

@@ -54,6 +54,7 @@ def setUp(self):
             vector_store_type="llamastack-faiss",
             embedding_dimension=None,
             manual_chunking=True,
+            doc_type="text",
         )
 
     @mock.patch.object(document_processor.tempfile, "TemporaryDirectory")