From 42a5cbd21e2fd50d7412935f484b41f590916229 Mon Sep 17 00:00:00 2001
From: "Tuan Anh Nguyen Dang (Tadashi_Cin)" <tadashi@cinnamon.is>
Date: Thu, 7 Nov 2024 21:07:16 +0700
Subject: [PATCH] feat: include light rag build in Docker (#475) #none

* feat: update docker build for lightrag

* feat: update docker build

* feat: update docker build
---
 Dockerfile      |  6 ++++++
 flowsettings.py | 37 ++++++++++++++++++++-----------------
 2 files changed, 26 insertions(+), 17 deletions(-)

diff --git a/Dockerfile b/Dockerfile
index 10de1321..d467fccb 100644
--- a/Dockerfile
+++ b/Dockerfile
@@ -81,6 +81,12 @@ RUN --mount=type=ssh  \
     pip install -e "libs/kotaemon[adv]" \
     && pip install unstructured[all-docs]
 
+# Install lightRAG
+ENV USE_LIGHTRAG=true
+RUN --mount=type=ssh  \
+    --mount=type=cache,target=/root/.cache/pip  \
+    pip install aioboto3 nano-vectordb ollama xxhash lightrag-hku
+
 # Clean up
 RUN apt-get autoremove \
     && apt-get clean \
diff --git a/flowsettings.py b/flowsettings.py
index a2f9e4c4..c4bb74e4 100644
--- a/flowsettings.py
+++ b/flowsettings.py
@@ -289,29 +289,32 @@
 USE_NANO_GRAPHRAG = config("USE_NANO_GRAPHRAG", default=False, cast=bool)
 USE_LIGHTRAG = config("USE_LIGHTRAG", default=False, cast=bool)
 
+GRAPHRAG_INDEX_TYPES = ["ktem.index.file.graph.GraphRAGIndex"]
+
 if USE_NANO_GRAPHRAG:
-    GRAPHRAG_INDEX_TYPE = "ktem.index.file.graph.NanoGraphRAGIndex"
+    GRAPHRAG_INDEX_TYPES.append("ktem.index.file.graph.NanoGraphRAGIndex")
 elif USE_LIGHTRAG:
-    GRAPHRAG_INDEX_TYPE = "ktem.index.file.graph.LightRAGIndex"
-else:
-    GRAPHRAG_INDEX_TYPE = "ktem.index.file.graph.GraphRAGIndex"
+    GRAPHRAG_INDEX_TYPES.append("ktem.index.file.graph.LightRAGIndex")
 
 KH_INDEX_TYPES = [
     "ktem.index.file.FileIndex",
-    GRAPHRAG_INDEX_TYPE,
+    *GRAPHRAG_INDEX_TYPES,
 ]
 
-GRAPHRAG_INDEX = {
-    "name": GRAPHRAG_INDEX_TYPE.split(".")[-1].replace("Index", ""),  # get last name
-    "config": {
-        "supported_file_types": (
-            ".png, .jpeg, .jpg, .tiff, .tif, .pdf, .xls, .xlsx, .doc, .docx, "
-            ".pptx, .csv, .html, .mhtml, .txt, .md, .zip"
-        ),
-        "private": False,
-    },
-    "index_type": GRAPHRAG_INDEX_TYPE,
-}
+GRAPHRAG_INDICES = [
+    {
+        "name": graph_type.split(".")[-1].replace("Index", ""),  # get last name
+        "config": {
+            "supported_file_types": (
+                ".png, .jpeg, .jpg, .tiff, .tif, .pdf, .xls, .xlsx, .doc, .docx, "
+                ".pptx, .csv, .html, .mhtml, .txt, .md, .zip"
+            ),
+            "private": False,
+        },
+        "index_type": graph_type,
+    }
+    for graph_type in GRAPHRAG_INDEX_TYPES
+]
 
 KH_INDICES = [
     {
@@ -325,5 +328,5 @@
         },
         "index_type": "ktem.index.file.FileIndex",
     },
-    GRAPHRAG_INDEX,
+    *GRAPHRAG_INDICES,
 ]