From f9d665056398c0f6e6cc61ae6456d900e5bdd55e Mon Sep 17 00:00:00 2001
From: Tadashi <tadashi@cinnamon.is>
Date: Tue, 17 Dec 2024 16:40:37 +0700
Subject: [PATCH] fix: update non local mode in LightRAG

---
 flowsettings.py                               |  2 +-
 libs/ktem/ktem/assets/css/main.css            |  5 ++
 .../index/file/graph/lightrag_pipelines.py    | 54 +++++++++++------
 .../ktem/index/file/graph/nano_pipelines.py   | 59 +++++++++++++------
 libs/ktem/ktem/utils/render.py                |  8 ++-
 5 files changed, 90 insertions(+), 38 deletions(-)

diff --git a/flowsettings.py b/flowsettings.py
index 119adab12..0647b71a4 100644
--- a/flowsettings.py
+++ b/flowsettings.py
@@ -302,7 +302,7 @@
 
 if USE_NANO_GRAPHRAG:
     GRAPHRAG_INDEX_TYPES.append("ktem.index.file.graph.NanoGraphRAGIndex")
-elif USE_LIGHTRAG:
+if USE_LIGHTRAG:
     GRAPHRAG_INDEX_TYPES.append("ktem.index.file.graph.LightRAGIndex")
 
 KH_INDEX_TYPES = [
diff --git a/libs/ktem/ktem/assets/css/main.css b/libs/ktem/ktem/assets/css/main.css
index f85ad39f0..82689a597 100644
--- a/libs/ktem/ktem/assets/css/main.css
+++ b/libs/ktem/ktem/assets/css/main.css
@@ -204,6 +204,11 @@ mark {
   right: 15px;
 }
 
+/* prevent overflow of html info panel */
+#html-info-panel {
+  overflow-x: auto !important;
+}
+
 #chat-expand-button {
   position: absolute;
   top: 6px;
diff --git a/libs/ktem/ktem/index/file/graph/lightrag_pipelines.py b/libs/ktem/ktem/index/file/graph/lightrag_pipelines.py
index 49e29b767..6a374f4a8 100644
--- a/libs/ktem/ktem/index/file/graph/lightrag_pipelines.py
+++ b/libs/ktem/ktem/index/file/graph/lightrag_pipelines.py
@@ -340,7 +340,7 @@ def get_user_settings(cls) -> dict:
             "search_type": {
                 "name": "Search type",
                 "value": "local",
-                "choices": ["local", "global", "hybrid", "naive"],
+                "choices": ["local", "global", "hybrid"],
                 "component": "dropdown",
                 "info": "Whether to use local or global search in the graph.",
             }
@@ -425,20 +425,40 @@ def run(
             return []
 
         graphrag_func, query_params = self._build_graph_search()
-        entities, relationships, sources = asyncio.run(
-            lightrag_build_local_query_context(graphrag_func, text, query_params)
-        )
 
-        documents = self.format_context_records(entities, relationships, sources)
-        plot = self.plot_graph(relationships)
-
-        return documents + [
-            RetrievedDocument(
-                text="",
-                metadata={
-                    "file_name": "GraphRAG",
-                    "type": "plot",
-                    "data": plot,
-                },
-            ),
-        ]
+        # only local mode support graph visualization
+        if query_params.mode == "local":
+            entities, relationships, sources = asyncio.run(
+                lightrag_build_local_query_context(graphrag_func, text, query_params)
+            )
+            documents = self.format_context_records(entities, relationships, sources)
+            plot = self.plot_graph(relationships)
+            documents += [
+                RetrievedDocument(
+                    text="",
+                    metadata={
+                        "file_name": "GraphRAG",
+                        "type": "plot",
+                        "data": plot,
+                    },
+                ),
+            ]
+        else:
+            context = graphrag_func.query(text, query_params)
+
+            # account for missing ``` for closing code block
+            context += "\n```"
+
+            documents = [
+                RetrievedDocument(
+                    text=context,
+                    metadata={
+                        "file_name": "GraphRAG {} Search".format(
+                            query_params.mode.capitalize()
+                        ),
+                        "type": "table",
+                    },
+                )
+            ]
+
+        return documents
diff --git a/libs/ktem/ktem/index/file/graph/nano_pipelines.py b/libs/ktem/ktem/index/file/graph/nano_pipelines.py
index 01fdb4db9..bbfdf26be 100644
--- a/libs/ktem/ktem/index/file/graph/nano_pipelines.py
+++ b/libs/ktem/ktem/index/file/graph/nano_pipelines.py
@@ -335,7 +335,7 @@ def get_user_settings(cls) -> dict:
             "search_type": {
                 "name": "Search type",
                 "value": "local",
-                "choices": ["local", "global", "naive"],
+                "choices": ["local", "global"],
                 "component": "dropdown",
                 "info": "Whether to use local or global search in the graph.",
             }
@@ -428,22 +428,43 @@ def run(
             return []
 
         graphrag_func, query_params = self._build_graph_search()
-        entities, relationships, reports, sources = asyncio.run(
-            nano_graph_rag_build_local_query_context(graphrag_func, text, query_params)
-        )
 
-        documents = self.format_context_records(
-            entities, relationships, reports, sources
-        )
-        plot = self.plot_graph(relationships)
-
-        return documents + [
-            RetrievedDocument(
-                text="",
-                metadata={
-                    "file_name": "GraphRAG",
-                    "type": "plot",
-                    "data": plot,
-                },
-            ),
-        ]
+        # only local mode support graph visualization
+        if query_params.mode == "local":
+            entities, relationships, reports, sources = asyncio.run(
+                nano_graph_rag_build_local_query_context(
+                    graphrag_func, text, query_params
+                )
+            )
+
+            documents = self.format_context_records(
+                entities, relationships, reports, sources
+            )
+            plot = self.plot_graph(relationships)
+
+            documents += [
+                RetrievedDocument(
+                    text="",
+                    metadata={
+                        "file_name": "GraphRAG",
+                        "type": "plot",
+                        "data": plot,
+                    },
+                ),
+            ]
+        else:
+            context = graphrag_func.query(text, query_params)
+
+            documents = [
+                RetrievedDocument(
+                    text=context,
+                    metadata={
+                        "file_name": "GraphRAG {} Search".format(
+                            query_params.mode.capitalize()
+                        ),
+                        "type": "table",
+                    },
+                )
+            ]
+
+        return documents
diff --git a/libs/ktem/ktem/utils/render.py b/libs/ktem/ktem/utils/render.py
index 9bab73186..c42b8ad66 100644
--- a/libs/ktem/ktem/utils/render.py
+++ b/libs/ktem/ktem/utils/render.py
@@ -49,7 +49,13 @@ def collapsible(header, content, open: bool = False) -> str:
     def table(text: str) -> str:
         """Render table from markdown format into HTML"""
         text = replace_mardown_header(text)
-        return markdown.markdown(text, extensions=["markdown.extensions.tables"])
+        return markdown.markdown(
+            text,
+            extensions=[
+                "markdown.extensions.tables",
+                "markdown.extensions.fenced_code",
+            ],
+        )
 
     @staticmethod
     def preview(