Allow file index to be private (#45)

* Fix breaking reranker

* Allow private file index

* Avoid setting default to 1 when user management is enabled

Author: trducng

libs/ktem/ktem/index/base.py

@@ -109,11 +109,16 @@ def get_admin_settings(cls) -> dict:
         return {}
 
     @abc.abstractmethod
-    def get_indexing_pipeline(self, settings: dict) -> "BaseComponent":
+    def get_indexing_pipeline(
+        self, settings: dict, user_id: Optional[int]
+    ) -> "BaseComponent":
         """Return the indexing pipeline that populates the entities into the index
 
         Args:
             settings: the user settings of the index
+            user_id: the user id who is accessing the index
+                TODO: instead of having a user_id, should have an app_state
+                which might also contain the settings.
 
         Returns:
             BaseIndexing: the indexing pipeline

libs/ktem/ktem/index/file/index.py

@@ -31,6 +31,26 @@ class FileIndex(BaseIndex):
 
     def __init__(self, app, id: int, name: str, config: dict):
         super().__init__(app, id, name, config)
+
+        self._indexing_pipeline_cls: Type[BaseFileIndexIndexing]
+        self._retriever_pipeline_cls: list[Type[BaseFileIndexRetriever]]
+        self._selector_ui_cls: Type
+        self._selector_ui: Any = None
+        self._index_ui_cls: Type
+        self._index_ui: Any = None
+
+        self._default_settings: dict[str, dict] = {}
+        self._setting_mappings: dict[str, dict] = {}
+
+    def _setup_resources(self):
+        """Setup resources for the file index
+
+        The resources include:
+            - Database table
+            - Vector store
+            - Document store
+            - File storage path
+        """
         Base = declarative_base()
         Source = type(
             "Source",
@@ -50,6 +70,7 @@ def __init__(self, app, id: int, name: str, config: dict):
                 "date_created": Column(
                     DateTime(timezone=True), server_default=func.now()
                 ),
+                "user": Column(Integer, default=1),
             },
         )
         Index = type(
@@ -61,6 +82,7 @@ def __init__(self, app, id: int, name: str, config: dict):
                 "source_id": Column(String),
                 "target_id": Column(String),
                 "relation_type": Column(Integer),
+                "user": Column(Integer, default=1),
             },
         )
         self._vs: BaseVectorStore = get_vectorstore(f"index_{self.id}")
@@ -74,16 +96,6 @@ def __init__(self, app, id: int, name: str, config: dict):
             "FileStoragePath": self._fs_path,
         }
 
-        self._indexing_pipeline_cls: Type[BaseFileIndexIndexing]
-        self._retriever_pipeline_cls: list[Type[BaseFileIndexRetriever]]
-        self._selector_ui_cls: Type
-        self._selector_ui: Any = None
-        self._index_ui_cls: Type
-        self._index_ui: Any = None
-
-        self._default_settings: dict[str, dict] = {}
-        self._setting_mappings: dict[str, dict] = {}
-
     def _setup_indexing_cls(self):
         """Retrieve the indexing class for the file index
 
@@ -247,6 +259,7 @@ def on_create(self):
         self.config = config
 
         # create the resources
+        self._setup_resources()
         self._resources["Source"].metadata.create_all(engine)  # type: ignore
         self._resources["Index"].metadata.create_all(engine)  # type: ignore
         self._fs_path.mkdir(parents=True, exist_ok=True)
@@ -255,6 +268,7 @@ def on_delete(self):
         """Clean up the index when the user delete it"""
         import shutil
 
+        self._setup_resources()
         self._resources["Source"].__table__.drop(engine)  # type: ignore
         self._resources["Index"].__table__.drop(engine)  # type: ignore
         self._vs.drop()
@@ -263,6 +277,7 @@ def on_delete(self):
 
     def on_start(self):
         """Setup the classes and hooks"""
+        self._setup_resources()
         self._setup_indexing_cls()
         self._setup_retriever_cls()
         self._setup_file_index_ui_cls()
@@ -326,9 +341,16 @@ def get_admin_settings(cls):
                     "Set 0 to disable."
                 ),
             },
+            "private": {
+                "name": "Make private",
+                "value": False,
+                "component": "radio",
+                "choices": [("Yes", True), ("No", False)],
+                "info": "If private, files will not be accessible across users.",
+            },
         }
 
-    def get_indexing_pipeline(self, settings) -> BaseFileIndexIndexing:
+    def get_indexing_pipeline(self, settings, user_id) -> BaseFileIndexIndexing:
         """Define the interface of the indexing pipeline"""
 
         prefix = f"index.options.{self.id}."
@@ -341,6 +363,7 @@ def get_indexing_pipeline(self, settings) -> BaseFileIndexIndexing:
 
         obj = self._indexing_pipeline_cls.get_pipeline(stripped_settings, self.config)
         obj.set_resources(resources=self._resources)
+        obj._user_id = user_id
 
         return obj
 

libs/ktem/ktem/index/file/pipelines.py

@@ -13,6 +13,7 @@
 from ktem.components import filestorage_path
 from ktem.db.models import engine
 from ktem.embeddings.manager import embedding_models_manager
+from ktem.llms.manager import llms
 from llama_index.vector_stores import (
     FilterCondition,
     FilterOperator,
@@ -28,7 +29,7 @@
 from kotaemon.base import RetrievedDocument
 from kotaemon.indices import VectorIndexing, VectorRetrieval
 from kotaemon.indices.ingests import DocumentIngestor
-from kotaemon.indices.rankings import BaseReranking
+from kotaemon.indices.rankings import BaseReranking, LLMReranking
 
 from .base import BaseFileIndexIndexing, BaseFileIndexRetriever
 
@@ -72,7 +73,7 @@ class DocumentRetrievalPipeline(BaseFileIndexRetriever):
     """
 
     vector_retrieval: VectorRetrieval = VectorRetrieval.withx()
-    reranker: BaseReranking
+    reranker: BaseReranking = LLMReranking.withx()
     get_extra_table: bool = False
     mmr: bool = False
     top_k: int = 5
@@ -225,12 +226,15 @@ def get_pipeline(cls, user_settings, index_settings, selected):
         """
         retriever = cls(
             get_extra_table=user_settings["prioritize_table"],
-            reranker=user_settings["reranking_llm"],
             top_k=user_settings["num_retrieval"],
             mmr=user_settings["mmr"],
         )
         if not user_settings["use_reranking"]:
             retriever.reranker = None  # type: ignore
+        else:
+            retriever.reranker.llm = llms.get(
+                user_settings["reranking_llm"], llms.get_default()
+            )
 
         retriever.vector_retrieval.embedding = embedding_models_manager[
             index_settings.get("embedding", embedding_models_manager.get_default_name())
@@ -342,6 +346,7 @@ def run(
                 name=Path(file_path).name,
                 path=file_hash,
                 size=Path(file_path).stat().st_size,
+                user=self._user_id,  # type: ignore
             )
             file_to_source[file_path] = source