Add support for user-defined storage_context in llamaindex_knowledge …

…and update knowledge config structure
cmgzn · Oct 21, 2024 · 7fb1532 · 7fb1532
1 parent 2dfe7cf
commit 7fb1532
Show file tree

Hide file tree

Showing 2 changed files with 98 additions and 8 deletions.
diff --git a/examples/conversation_with_RAG_agents/configs/knowledge_config.json b/examples/conversation_with_RAG_agents/configs/knowledge_config.json
@@ -20,7 +20,7 @@
             }
           }
         },
-        "store_and_index": {
+        "data_parse": {
           "transformations": [
             {
               "create_object": true,
@@ -34,7 +34,22 @@
           ]
         }
       }
-    ]
+    ],
+    "store_and_index": {
+      "storage_context": {
+        "vector_store": {
+          "create_object": true,
+          "module": "llama_index.vector_stores.elasticsearch",
+          "class": "ElasticsearchStore",
+          "init_args": {
+            "index_name": "agentscope_code_rag",
+            "es_url": "http://localhost:9200",
+            "es_user": "elastic",
+            "es_password": ""
+          }
+        }
+      }
+    }
   },
   {
     "knowledge_id": "agentscope_api_rag",
@@ -57,7 +72,22 @@
           }
         }
       }
-    ]
+    ],
+    "store_and_index": {
+      "storage_context": {
+        "vector_store": {
+          "create_object": true,
+          "module": "llama_index.vector_stores.elasticsearch",
+          "class": "ElasticsearchStore",
+          "init_args": {
+            "index_name": "agentscope_api_rag",
+            "es_url": "http://localhost:9200",
+            "es_user": "elastic",
+            "es_password": ""
+          }
+        }
+      }
+    }
   },
   {
     "knowledge_id": "agentscope_global_rag",
@@ -95,7 +125,7 @@
             }
           }
         },
-        "store_and_index": {
+        "data_parse": {
           "transformations": [
             {
               "create_object": true,
@@ -109,6 +139,21 @@
           ]
         }
       }
-    ]
+    ],
+    "store_and_index": {
+      "storage_context": {
+        "vector_store": {
+          "create_object": true,
+          "module": "llama_index.vector_stores.elasticsearch",
+          "class": "ElasticsearchStore",
+          "init_args": {
+            "index_name": "agentscope_global_rag",
+            "es_url": "http://localhost:9200",
+            "es_user": "elastic",
+            "es_password": ""
+          }
+        }
+      }
+    }
   }
 ]
diff --git a/src/agentscope/rag/llama_index_knowledge.py b/src/agentscope/rag/llama_index_knowledge.py
@@ -252,9 +252,21 @@ def _load_index(self) -> None:
         Load the persisted index from persist_dir.
         """
         # load the storage_context
-        storage_context = StorageContext.from_defaults(
-            persist_dir=self.persist_dir,
+        # Inject the persist_dir setting
+        self.knowledge_config.get("store_and_index", {}).get(
+            "storage_context",
+            {},
+        ).update(
+            {
+                "persist_dir": self.persist_dir,
+            },
         )
+        storage_context = self._set_store(self.knowledge_config)
+        if not storage_context:
+            # if storage_context is not set, use the default
+            storage_context = StorageContext.from_defaults(
+                persist_dir=self.persist_dir,
+            )
         # construct index from
         self.index = load_index_from_storage(
             storage_context=storage_context,
@@ -287,10 +299,14 @@ def _data_to_index(self) -> None:
                 transformations=transformations,
             )
             nodes = nodes + nodes_docs
+        # set store
+        storage_context = self._set_store(self.knowledge_config)
         # convert nodes to index
+        # if storage_context is None, use the default
         self.index = VectorStoreIndex(
             nodes=nodes,
             embed_model=self.emb_model,
+            storage_context=storage_context,
         )
         logger.info("index calculation completed.")
         # persist the calculated index
@@ -401,7 +417,16 @@ def _set_transformations(self, config: dict) -> Any:
         Args:
             config (dict): a dictionary containing configurations.
         """
-        if "store_and_index" in config:
+        if "data_parse" in config:
+            temp = self._prepare_args_from_config(
+                config=config.get("data_parse", {}),
+            )
+            transformations = temp.get("transformations")
+        elif "store_and_index" in config:
+            logger.warning(
+                "The old configuration structure is deprecated, "
+                "please use data_parse instead of store_and_index.",
+            )
             temp = self._prepare_args_from_config(
                 config=config.get("store_and_index", {}),
             )
@@ -427,6 +452,26 @@ def _set_transformations(self, config: dict) -> Any:
         transformations = {"transformations": transformations}
         return transformations
 
+    def _set_store(self, config: dict) -> Any:
+        """
+        Set the store as needed, or just use the default setting.
+
+        Args:
+            config (dict): a dictionary containing configurations.
+        """
+        if "store_and_index" in config:
+            temp = self._prepare_args_from_config(
+                config=config.get("store_and_index", {}),
+            )
+            context_config = temp.get("storage_context")
+        else:
+            return None
+
+        # Create the storage context
+        storage_context = StorageContext.from_defaults(**context_config)
+        logger.info("storage_context is ready.")
+        return storage_context
+
     def _get_retriever(
         self,
         similarity_top_k: int = None,